add artifact generator agent

Author: leehuwuj

templates/types/multiagent/fastapi/app/agents/multi.py

@@ -8,7 +8,7 @@
 )
 from llama_index.core.tools.types import ToolMetadata, ToolOutput
 from llama_index.core.tools.utils import create_schema_from_function
-from llama_index.core.workflow import Context, Workflow
+from llama_index.core.workflow import Context, StopEvent, Workflow
 
 
 class AgentCallTool(ContextAwareTool):
@@ -35,7 +35,8 @@ async def acall(self, ctx: Context, input: str) -> ToolOutput:
         handler = self.agent.run(input=input)
         # bubble all events while running the agent to the calling agent
         async for ev in handler.stream_events():
-            ctx.write_event_to_stream(ev)
+            if type(ev) is not StopEvent:
+                ctx.write_event_to_stream(ev)
         ret: AgentRunResult = await handler
         response = ret.response.message.content
         return ToolOutput(

templates/types/multiagent/fastapi/app/examples/artifact_generator.py

@@ -0,0 +1,19 @@
+from typing import List
+
+from app.agents.single import FunctionCallingAgent
+from app.tools.artifact import ArtifactGenerator
+from llama_index.core.chat_engine.types import ChatMessage
+from llama_index.core.tools import FunctionTool
+
+
+def create_artifact_generator(chat_history: List[ChatMessage]):
+    artifact_tool = FunctionTool.from_defaults(ArtifactGenerator.generate_artifact)
+
+    return FunctionCallingAgent(
+        name="ArtifactGenerator",
+        tools=[artifact_tool],
+        role="expert in generating artifacts (pdf, html)",
+        system_prompt="You are generator that help generate artifacts (pdf, html) from a given content.",
+        chat_history=chat_history,
+        verbose=True,
+    )

templates/types/multiagent/fastapi/app/examples/choreography.py

@@ -1,25 +1,31 @@
 from typing import List, Optional
-from app.agents.single import FunctionCallingAgent
+
 from app.agents.multi import AgentCallingAgent
+from app.agents.single import FunctionCallingAgent
+from app.examples.artifact_generator import create_artifact_generator
 from app.examples.researcher import create_researcher
 from llama_index.core.chat_engine.types import ChatMessage
 
 
 def create_choreography(chat_history: Optional[List[ChatMessage]] = None):
     researcher = create_researcher(chat_history)
+    artifact_generator = create_artifact_generator(chat_history)
     reviewer = FunctionCallingAgent(
         name="reviewer",
         role="expert in reviewing blog posts",
         system_prompt="You are an expert in reviewing blog posts. You are given a task to review a blog post. Review the post for logical inconsistencies, ask critical questions, and provide suggestions for improvement. Furthermore, proofread the post for grammar and spelling errors. If the post is good, you can say 'The post is good.'",
         chat_history=chat_history,
+        verbose=True,
     )
     return AgentCallingAgent(
         name="writer",
-        agents=[researcher, reviewer],
+        agents=[researcher, reviewer, artifact_generator],
         role="expert in writing blog posts",
         system_prompt="""You are an expert in writing blog posts. You are given a task to write a blog post. Before starting to write the post, consult the researcher agent to get the information you need. Don't make up any information yourself.
         After creating a draft for the post, send it to the reviewer agent to receive some feedback and make sure to incorporate the feedback from the reviewer.
-        You can consult the reviewer and researcher maximal two times. Your output should just contain the blog post.""",
+        You can consult the reviewer and researcher maximal two times. Your output should just contain the blog post. 
+        Finally, always request the artifact generator to create an artifact (pdf, html) that user can download and use for publishing the blog post.""",
         # TODO: add chat_history support to AgentCallingAgent
         # chat_history=chat_history,
+        verbose=True,
     )

templates/types/multiagent/fastapi/app/examples/orchestrator.py

@@ -1,8 +1,9 @@
 from typing import List, Optional
-from app.agents.single import FunctionCallingAgent
+
 from app.agents.multi import AgentOrchestrator
+from app.agents.single import FunctionCallingAgent
+from app.examples.artifact_generator import create_artifact_generator
 from app.examples.researcher import create_researcher
-
 from llama_index.core.chat_engine.types import ChatMessage
 
 
@@ -21,7 +22,8 @@ def create_orchestrator(chat_history: Optional[List[ChatMessage]] = None):
         Especially check for logical inconsistencies and proofread the post for grammar and spelling errors.""",
         chat_history=chat_history,
     )
+    artifact_generator = create_artifact_generator(chat_history)
     return AgentOrchestrator(
-        agents=[writer, reviewer, researcher],
+        agents=[writer, reviewer, researcher, artifact_generator],
         refine_plan=False,
     )

templates/types/multiagent/fastapi/app/examples/workflow.py

@@ -1,6 +1,7 @@
 from typing import AsyncGenerator, List, Optional
 
 from app.agents.single import AgentRunEvent, AgentRunResult, FunctionCallingAgent
+from app.examples.artifact_generator import create_artifact_generator
 from app.examples.researcher import create_researcher
 from llama_index.core.chat_engine.types import ChatMessage
 from llama_index.core.workflow import (
@@ -17,6 +18,9 @@ def create_workflow(chat_history: Optional[List[ChatMessage]] = None):
     researcher = create_researcher(
         chat_history=chat_history,
     )
+    artifact_generator = create_artifact_generator(
+        chat_history=chat_history,
+    )
     writer = FunctionCallingAgent(
         name="writer",
         role="expert in writing blog posts",
@@ -30,7 +34,12 @@ def create_workflow(chat_history: Optional[List[ChatMessage]] = None):
         chat_history=chat_history,
     )
     workflow = BlogPostWorkflow(timeout=360)
-    workflow.add_workflows(researcher=researcher, writer=writer, reviewer=reviewer)
+    workflow.add_workflows(
+        researcher=researcher,
+        writer=writer,
+        reviewer=reviewer,
+        artifact_generator=artifact_generator,
+    )
     return workflow
 
 
@@ -40,13 +49,16 @@ class ResearchEvent(Event):
 
 class WriteEvent(Event):
     input: str
-    is_good: bool = False
 
 
 class ReviewEvent(Event):
     input: str
 
 
+class GenerateArtifactEvent(Event):
+    input: str
+
+
 class BlogPostWorkflow(Workflow):
     @step()
     async def start(self, ctx: Context, ev: StartEvent) -> ResearchEvent:
@@ -80,7 +92,7 @@ async def write(
                     msg=f"Too many attempts ({MAX_ATTEMPTS}) to write the blog post. Proceeding with the current version.",
                 )
             )
-        if ev.is_good or too_many_attempts:
+        if too_many_attempts:
             # too many attempts or the blog post is good - stream final response if requested
             result = await self.run_agent(
                 ctx, writer, ev.input, streaming=ctx.data["streaming"]
@@ -93,7 +105,7 @@ async def write(
     @step()
     async def review(
         self, ctx: Context, ev: ReviewEvent, reviewer: FunctionCallingAgent
-    ) -> WriteEvent:
+    ) -> WriteEvent | GenerateArtifactEvent:
         result: AgentRunResult = await self.run_agent(ctx, reviewer, ev.input)
         review = result.response.message.content
         old_content = ctx.data["result"].response.message.content
@@ -105,9 +117,8 @@ async def review(
             )
         )
         if post_is_good:
-            return WriteEvent(
+            return GenerateArtifactEvent(
                 input=f"You're blog post is ready for publication. Please respond with just the blog post. Blog post: ```{old_content}```",
-                is_good=True,
             )
         else:
             return WriteEvent(
@@ -123,6 +134,16 @@ async def review(
 ```"""
             )
 
+    @step()
+    async def generate_artifact(
+        self,
+        ctx: Context,
+        ev: GenerateArtifactEvent,
+        artifact_generator: FunctionCallingAgent,
+    ) -> StopEvent:
+        result: AgentRunResult = await self.run_agent(ctx, artifact_generator, ev.input)
+        return StopEvent(result=result)
+
     async def run_agent(
         self,
         ctx: Context,

templates/types/multiagent/fastapi/app/tools/artifact.py

@@ -0,0 +1,154 @@
+import os
+from enum import Enum
+from io import BytesIO
+
+OUTPUT_DIR = "output/tools"
+
+
+class ArtifactType(Enum):
+    PDF = "pdf"
+    HTML = "html"
+
+
+class ArtifactGenerator:
+    @classmethod
+    def _generate_pdf(cls, original_content: str) -> BytesIO:
+        """
+        Generate a PDF from the original content (markdown).
+        """
+        try:
+            import markdown
+            from reportlab.lib.pagesizes import letter
+            from reportlab.lib.styles import getSampleStyleSheet
+            from reportlab.platypus import Paragraph, SimpleDocTemplate
+        except ImportError:
+            raise ImportError(
+                "Failed to import required modules. Please install reportlab and markdown."
+            )
+
+        # Convert markdown to HTML
+        html = markdown.markdown(original_content)
+
+        buffer = BytesIO()
+
+        doc = SimpleDocTemplate(buffer, pagesize=letter)
+
+        # Create a list to store the flowables (content elements)
+        elements = []
+        styles = getSampleStyleSheet()
+        # TODO: Make the format nicer
+        for paragraph in html.split("<p>"):
+            if paragraph:
+                clean_text = paragraph.replace("</p>", "").strip()
+                elements.append(Paragraph(clean_text, styles["Normal"]))
+
+        # Build the PDF document
+        doc.build(elements)
+
+        # Reset the buffer position to the beginning
+        buffer.seek(0)
+
+        return buffer
+
+    @classmethod
+    def _generate_html(cls, original_content: str) -> str:
+        """
+        Generate an HTML from the original content (markdown).
+        """
+        try:
+            import markdown
+        except ImportError:
+            raise ImportError(
+                "Failed to import required modules. Please install markdown."
+            )
+
+        # Convert markdown to HTML
+        html_content = markdown.markdown(original_content)
+
+        # Create a complete HTML document with basic styling
+        html_document = f"""
+        <!DOCTYPE html>
+        <html lang="en">
+        <head>
+            <meta charset="UTF-8">
+            <meta name="viewport" content="width=device-width, initial-scale=1.0">
+            <title>Generated HTML Document</title>
+            <style>
+                body {{
+                    font-family: Arial, sans-serif;
+                    line-height: 1.6;
+                    color: #333;
+                    max-width: 800px;
+                    margin: 0 auto;
+                    padding: 20px;
+                }}
+                h1, h2, h3, h4, h5, h6 {{
+                    margin-top: 1.5em;
+                    margin-bottom: 0.5em;
+                }}
+                p {{
+                    margin-bottom: 1em;
+                }}
+                code {{
+                    background-color: #f4f4f4;
+                    padding: 2px 4px;
+                    border-radius: 4px;
+                }}
+                pre {{
+                    background-color: #f4f4f4;
+                    padding: 10px;
+                    border-radius: 4px;
+                    overflow-x: auto;
+                }}
+            </style>
+        </head>
+        <body>
+            {html_content}
+        </body>
+        </html>
+        """
+
+        return html_document
+
+    @classmethod
+    def _write_to_file(cls, content: BytesIO, file_path: str):
+        """
+        Write the content to a file.
+        """
+        os.makedirs(os.path.dirname(file_path), exist_ok=True)
+        with open(file_path, "wb") as file:
+            file.write(content.getvalue())
+
+    @classmethod
+    def generate_artifact(
+        cls, original_content: str, artifact_type: str, file_name: str
+    ) -> str:
+        """
+        Generate an artifact from the original content and write it to a file.
+        Parameters:
+            original_content: str (markdown style)
+            artifact_type: str (pdf or html). Use pdf for report, html for blog post.
+            file_name: str (name of the artifact file), don't need to include the file extension. It will be added automatically based on the artifact type.
+        Returns:
+            str (URL to the artifact file): the url that already available to the file server. No need to change the path anymore.
+        """
+        try:
+            artifact_type = ArtifactType(artifact_type.lower())
+        except ValueError:
+            raise ValueError(
+                f"Invalid artifact type: {artifact_type}. Must be 'pdf' or 'html'."
+            )
+
+        if artifact_type == ArtifactType.PDF:
+            content = cls._generate_pdf(original_content)
+            file_extension = "pdf"
+        elif artifact_type == ArtifactType.HTML:
+            content = BytesIO(cls._generate_html(original_content).encode("utf-8"))
+            file_extension = "html"
+        else:
+            raise ValueError(f"Unexpected artifact type: {artifact_type}")
+
+        file_path = os.path.join(OUTPUT_DIR, f"{file_name}.{file_extension}")
+        cls._write_to_file(content, file_path)
+        file_url = f"{os.getenv('FILESERVER_URL_PREFIX')}/{file_path}"
+        return file_url

templates/types/multiagent/fastapi/pyproject.toml

@@ -18,6 +18,8 @@ python-dotenv = "^1.0.0"
 uvicorn = { extras = ["standard"], version = "^0.23.2" }
 cachetools = "^5.3.3"
 aiostream = "^0.5.2"
+markdown = "^3.7"
+reportlab = "^4.2.2"
 
 [tool.poetry.dependencies.docx2txt]
 version = "^0.8"