Refactor artifact generation tools by introducing separate CodeGenerator and DocumentGenerator classes. Update app_writer to utilize FunctionAgent for code and document generation workflows. Remove deprecated ArtifactGenerator class. Enhance artifact transformation logic in callbacks. Improve system prompts for clarity and instruction adherence.

leehuwuj · leehuwuj · commit 630a76d6c80d · 2025-04-24T11:21:21.000+07:00
diff --git a/llama-index-server/examples/app_writer.py b/llama-index-server/examples/app_writer.py
@@ -1,30 +1,57 @@
-from typing import List
-
 from fastapi import FastAPI
 
-from llama_index.core.agent.workflow import AgentWorkflow
-from llama_index.core.tools import BaseTool
+from llama_index.core.agent.workflow import AgentWorkflow, FunctionAgent
 from llama_index.core.workflow import Workflow
 from llama_index.llms.openai import OpenAI
 from llama_index.server import LlamaIndexServer, UIConfig
 from llama_index.server.api.models import ChatRequest
 from llama_index.server.api.utils import get_last_artifact
-from llama_index.server.tools.artifact_generator import ArtifactGenerator
+from llama_index.server.tools.artifact import CodeGenerator, DocumentGenerator
 
 
 def create_workflow(chat_request: ChatRequest) -> Workflow:
-    tools: List[BaseTool] = [
-        ArtifactGenerator(
-            last_artifact=get_last_artifact(chat_request),
-            llm=OpenAI(model="gpt-4.1"),
-        ).to_tool()
-    ]
-    agent = AgentWorkflow.from_tools_or_functions(
-        tools,  # type: ignore
+    app_writer_agent = FunctionAgent(
+        name="Coder",
+        description="A skilled full-stack developer.",
+        system_prompt="""
+        You are an skilled full-stack developer that can help user update the code by using the code generator tool.
+        Follow these instructions:
+         + Thinking and provide a correct requirement to the code generator tool. 
+         + Always use the tool to update the code.
+         + Don't need to response the code just summarize the code and the changes you made.
+        """,
+        tools=[
+            CodeGenerator(
+                last_artifact=get_last_artifact(chat_request),
+                llm=OpenAI(model="gpt-4.1"),
+            ).to_tool()
+        ],  # type: ignore
+        llm=OpenAI(model="gpt-4.1"),
+    )
+    doc_writer_agent = FunctionAgent(
+        name="Writer",
+        description="A skilled document writer.",
+        system_prompt="""
+        You are an skilled document writer that can help user update the document by using the document generator tool.
+        Follow these instructions:
+         + Thinking and provide a correct requirement to the document generator tool.
+         + Always use the tool to update the document.
+         + Don't need to response the document just summarize the document and the changes you made.
+        """,
+        tools=[
+            DocumentGenerator(
+                last_artifact=get_last_artifact(chat_request),
+                llm=OpenAI(model="gpt-4.1"),
+            ).to_tool()
+        ],  # type: ignore
         llm=OpenAI(model="gpt-4.1-mini"),
-        system_prompt="You are a helpful assistant that can generate artifacts (code or markdown document), use the provided tools to respond to the user's request.",
     )
-    return agent
+    workflow = AgentWorkflow(
+        agents=[app_writer_agent, doc_writer_agent],
+        root_agent="Coder",
+        verbose=True,
+    )
+    return workflow
 
 
 def create_app() -> FastAPI:
diff --git a/llama-index-server/llama_index/server/api/callbacks/artifact.py b/llama-index-server/llama_index/server/api/callbacks/artifact.py
@@ -1,5 +1,4 @@
-import time
-from typing import Any
+from typing import Any, Dict, Optional
 
 from llama_index.core.agent.workflow.workflow_events import ToolCallResult
 from llama_index.server.api.callbacks.base import EventCallback
@@ -15,20 +14,23 @@ class ArtifactFromToolCall(EventCallback):
                          default is "artifact"
     """
 
-    def __init__(self, tool_name: str = "artifact_generator"):
-        self.tool_name = tool_name
+    def __init__(self, tool_prefix: str = "artifact_"):
+        self.tool_prefix = tool_prefix
 
-    def transform_tool_call_result(self, event: ToolCallResult) -> Artifact:
-        artifact = event.tool_output.raw_output
-        return Artifact(
-            created_at=int(time.time()),
-            type=artifact.get("type"),
-            data=artifact.get("data"),
-        )
+    def transform_tool_call_result(
+        self, event: ToolCallResult
+    ) -> Optional[Dict[str, Any]]:
+        artifact: Artifact = event.tool_output.raw_output
+        if isinstance(artifact, str):  # Error tool output
+            return None
+        return {
+            "type": "artifact",
+            "data": artifact.model_dump(),
+        }
 
     async def run(self, event: Any) -> Any:
         if isinstance(event, ToolCallResult):
-            if event.tool_name == self.tool_name:
+            if event.tool_name.startswith(self.tool_prefix):
                 return event, self.transform_tool_call_result(event)
         return event
 
diff --git a/llama-index-server/llama_index/server/tools/artifact/__init__.py b/llama-index-server/llama_index/server/tools/artifact/__init__.py
@@ -0,0 +1,4 @@
+from llama_index.server.tools.artifact.code_generator import CodeGenerator
+from llama_index.server.tools.artifact.document_generator import DocumentGenerator
+
+__all__ = ["CodeGenerator", "DocumentGenerator"]
diff --git a/llama-index-server/llama_index/server/tools/artifact/code_generator.py b/llama-index-server/llama_index/server/tools/artifact/code_generator.py
@@ -0,0 +1,145 @@
+import logging
+import re
+import time
+from typing import List, Optional
+
+from llama_index.core.llms import LLM
+from llama_index.core.llms.llm import ChatMessage, MessageRole
+from llama_index.core.settings import Settings
+from llama_index.core.tools.function_tool import FunctionTool
+from llama_index.server.api.models import Artifact, ArtifactType, CodeArtifactData
+
+logger = logging.getLogger(__name__)
+
+CODE_GENERATION_PROMPT = """
+You are a highly skilled content creator and software engineer. 
+Your task is to generate code to resolve the user's request.
+
+Follow these instructions exactly:
+
+1. Carefully read the user's requirements. 
+   If any details are ambiguous or missing, make reasonable assumptions and clearly reflect those in your output.
+   If the previous code is provided, carefully analyze the code with the request to make the right changes.
+2. For code requests:
+   - If the user does not specify a framework or language, default to a React component using the Next.js framework.
+   - For Next.js, use Shadcn UI components, Typescript, @types/node, @types/react, @types/react-dom, PostCSS, and TailwindCSS.
+     The import pattern should be:
+     ```
+     import { ComponentName } from "@/components/ui/component-name"
+     import { Markdown } from "@llamaindex/chat-ui"
+     import { cn } from "@/lib/utils"
+     ```
+   - Ensure the code is idiomatic, production-ready, and includes necessary imports.
+   - Only generate code relevant to the user's request—do not add extra boilerplate.
+3. Don't be verbose on response, no other text or comments only return the code which wrapped by ```language``` block.
+Example:
+```typescript
+import React from "react";
+
+export default function MyComponent() {
+  return <div>Hello World</div>;
+}
+```
+"""
+
+
+class CodeGenerator:
+    def __init__(
+        self,
+        llm: Optional[LLM] = None,
+        last_artifact: Optional[Artifact] = None,
+    ) -> None:
+        if llm is None:
+            if Settings.llm is None:
+                raise ValueError(
+                    "Missing llm. Please provide a valid LLM or set the LLM using Settings.llm."
+                )
+            llm = Settings.llm
+        self.llm = llm
+        self.last_artifact = last_artifact
+
+    def prepare_chat_messages(
+        self, requirement: str, language: str, previous_code: Optional[str] = None
+    ) -> List[ChatMessage]:
+        user_messages: List[ChatMessage] = []
+        user_messages.append(ChatMessage(role=MessageRole.USER, content=requirement))
+        if previous_code:
+            user_messages.append(
+                ChatMessage(
+                    role=MessageRole.USER,
+                    content=f"```{language}\n{previous_code}\n```",
+                )
+            )
+        else:
+            user_messages.append(
+                ChatMessage(
+                    role=MessageRole.USER,
+                    content=f"Write code in {language}. Wrap the code in ```{language}``` block.",
+                )
+            )
+        return user_messages
+
+    async def generate_code(
+        self,
+        file_name: str,
+        language: str,
+        requirement: str,
+        previous_code: Optional[str] = None,
+    ) -> Artifact:
+        """
+        Generate code based on the provided requirement.
+
+        Args:
+            file_name (str): The name of the file to generate.
+            language (str): The language of the code to generate (Only "typescript" and "python" is supported now)
+            requirement (str): Provide a detailed requirement for the code to be generated/updated.
+            old_content (Optional[str]): Existing code content to be modified or referenced. Defaults to None.
+
+        Returns:
+            Artifact: A dictionary containing the generated artifact details
+                           (type, data).
+        """
+        user_messages = self.prepare_chat_messages(requirement, language, previous_code)
+
+        messages: List[ChatMessage] = [
+            ChatMessage(role=MessageRole.SYSTEM, content=CODE_GENERATION_PROMPT),
+            *user_messages,
+        ]
+
+        try:
+            response = await self.llm.achat(messages)
+            raw_content = response.message.content
+            if not raw_content:
+                raise ValueError(
+                    "Empty response. Try with a clearer requirement or provide previous code."
+                )
+
+            # Extract code from code block in raw content
+            code_block = re.search(r"```(.*?)\n(.*?)```", raw_content, re.DOTALL)
+            if not code_block:
+                raise ValueError("Couldn't parse code from the response.")
+            code = code_block.group(2).strip()
+            return Artifact(
+                created_at=int(time.time()),
+                type=ArtifactType.CODE,
+                data=CodeArtifactData(
+                    file_name=file_name,
+                    code=code,
+                    language=language,
+                ),
+            )
+        except Exception as e:
+            raise ValueError(f"Couldn't generate code. {e}")
+
+    def to_tool(self) -> FunctionTool:
+        """
+        Converts the CodeGenerator instance into a FunctionTool.
+
+        Returns:
+            FunctionTool: A tool that can be used by agents.
+        """
+        return FunctionTool.from_defaults(
+            self.generate_code,
+            name="artifact_code_generator",
+            description="Generate/update code based on a requirement.",
+        )
diff --git a/llama-index-server/llama_index/server/tools/artifact/document_generator.py b/llama-index-server/llama_index/server/tools/artifact/document_generator.py
@@ -0,0 +1,129 @@
+import logging
+import re
+import time
+from typing import List, Literal, Optional
+
+from llama_index.core.llms import LLM
+from llama_index.core.llms.llm import ChatMessage, MessageRole
+from llama_index.core.settings import Settings
+from llama_index.core.tools.function_tool import FunctionTool
+from llama_index.server.api.models import Artifact, ArtifactType, DocumentArtifactData
+
+logger = logging.getLogger(__name__)
+
+DOCUMENT_GENERATION_PROMPT = """
+You are a highly skilled writer and content creator.
+Your task is to generate documents based on the user's request.
+
+Follow these instructions exactly:
+
+1. Carefully read the user's requirements.
+   If any details are ambiguous or missing, make reasonable assumptions and clearly reflect those in your output.
+   If previous content is provided, carefully analyze it with the request to make the right changes.
+2. For document creation:
+   - Create well-structured documents with clear headings, paragraphs, and formatting.
+   - Use concise and professional language.
+   - Ensure content is accurate, well-researched, and relevant to the topic.
+   - Organize information logically with a proper introduction, body, and conclusion when appropriate.
+   - Add citations or references when necessary.
+3. For document editing:
+   - Maintain the original document's structure unless requested otherwise.
+   - Improve clarity, flow, and grammar while preserving the original message.
+   - Remove redundancies and strengthen weak points.
+4. Answer in appropriate format which wrapped by ```<format>``` block with a file name for the content, no other text or comments.
+Example:
+```markdown
+# Title
+
+Content
+```
+"""
+
+
+class DocumentGenerator:
+    def __init__(
+        self,
+        llm: Optional[LLM] = None,
+        last_artifact: Optional[Artifact] = None,
+    ) -> None:
+        if llm is None:
+            if Settings.llm is None:
+                raise ValueError(
+                    "Missing llm. Please provide a valid LLM or set the LLM using Settings.llm."
+                )
+            llm = Settings.llm
+        self.llm = llm
+        self.last_artifact = last_artifact
+
+    def prepare_chat_messages(
+        self, requirement: str, previous_content: Optional[str] = None
+    ) -> List[ChatMessage]:
+        user_messages: List[ChatMessage] = []
+        user_messages.append(ChatMessage(role=MessageRole.USER, content=requirement))
+        if previous_content:
+            user_messages.append(
+                ChatMessage(role=MessageRole.USER, content=previous_content)
+            )
+        return user_messages
+
+    async def generate_document(
+        self,
+        file_name: str,
+        document_format: Literal["markdown", "html"],
+        requirement: str,
+        previous_content: Optional[str] = None,
+    ) -> Artifact:
+        """
+        Generate document content based on the provided requirement.
+
+        Args:
+            file_name (str): The name of the file to generate.
+            document_format (str): The format of the document to generate. (Only "markdown" and "html" are supported now)
+            requirement (str): A detailed requirement for the document to be generated/updated.
+            previous_content (Optional[str]): Existing document content to be modified or referenced. Defaults to None.
+
+        Returns:
+            Artifact: The generated document.
+        """
+        user_messages = self.prepare_chat_messages(requirement, previous_content)
+
+        messages: List[ChatMessage] = [
+            ChatMessage(role=MessageRole.SYSTEM, content=DOCUMENT_GENERATION_PROMPT),
+            *user_messages,
+        ]
+
+        try:
+            response = await self.llm.achat(messages)
+            raw_content = response.message.content
+            if not raw_content:
+                raise ValueError(
+                    "Empty response. Try with a clearer requirement or provide previous content."
+                )
+            # Extract content from the response
+            content = re.search(r"```(.*?)\n(.*?)```", raw_content, re.DOTALL)
+            if not content:
+                raise ValueError("Couldn't parse content from the response.")
+            return Artifact(
+                created_at=int(time.time()),
+                type=ArtifactType.DOCUMENT,
+                data=DocumentArtifactData(
+                    title=file_name,
+                    content=content.group(2).strip(),
+                    type=document_format,
+                ),
+            )
+        except Exception as e:
+            raise ValueError(f"Couldn't generate document. {e}")
+
+    def to_tool(self) -> FunctionTool:
+        """
+        Converts the DocumentGenerator instance into a FunctionTool.
+
+        Returns:
+            FunctionTool: A tool that can be used by agents.
+        """
+        return FunctionTool.from_defaults(
+            self.generate_document,
+            name="artifact_document_generator",
+            description="Generate/update documents based on a requirement.",
+        )
diff --git a/llama-index-server/llama_index/server/tools/artifact_generator.py b/llama-index-server/llama_index/server/tools/artifact_generator.py
