add duckduckgo tool and improve prompting

leehuwuj · leehuwuj · commit adff220f7954 · 2024-09-26T15:20:14.000+07:00
diff --git a/templates/types/multiagent/fastapi/app/examples/orchestrator.py b/templates/types/multiagent/fastapi/app/examples/orchestrator.py
@@ -11,14 +11,19 @@ def create_orchestrator(chat_history: Optional[List[ChatMessage]] = None):
     researcher = create_researcher(chat_history)
     writer = FunctionCallingAgent(
         name="writer",
-        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. Don't make up any information yourself. If you don't have the necessary information to write a blog post, reply "I need information about the topic to write the blog post". If you have all the information needed, write the blog post.""",
+        role="expert in writing blog posts, need information and images to write a post",
+        system_prompt="""You are an expert in writing blog posts. 
+        You are given a task to write a blog post. Don't make up any information yourself. 
+        If you don't have the necessary information to write a blog post, reply "I need information about the topic to write the blog post". 
+        If you need to use images, reply "I need images about the topic to write the blog post". Don't use any dummy images made up by you.
+        If you have all the information needed, write the blog post.""",
         chat_history=chat_history,
     )
     reviewer = FunctionCallingAgent(
         name="reviewer",
-        role="expert in reviewing blog posts",
+        role="expert in reviewing blog posts, need a written blog post to review",
         system_prompt="""You are an expert in reviewing blog posts. You are given a task to review a blog post. Review the post and fix the issues found yourself. You must output a final blog post.
+        A post must include at lease one valid image, if not, reply "I need images about the topic to write the blog post". An image URL start with example or your website is not valid.
         Especially check for logical inconsistencies and proofread the post for grammar and spelling errors.""",
         chat_history=chat_history,
     )
diff --git a/templates/types/multiagent/fastapi/app/examples/researcher.py b/templates/types/multiagent/fastapi/app/examples/researcher.py
@@ -1,10 +1,11 @@
 import os
 from typing import List
-from llama_index.core.tools import QueryEngineTool, ToolMetadata
+
 from app.agents.single import FunctionCallingAgent
 from app.engine.index import get_index
-
+from app.tools.duckduckgo import get_tools as get_duckduckgo_tools
 from llama_index.core.chat_engine.types import ChatMessage
+from llama_index.core.tools import QueryEngineTool, ToolMetadata
 
 
 def get_query_engine_tool() -> QueryEngineTool:
@@ -30,10 +31,11 @@ def get_query_engine_tool() -> QueryEngineTool:
 
 
 def create_researcher(chat_history: List[ChatMessage]):
+    duckduckgo_search_tools = get_duckduckgo_tools()
     return FunctionCallingAgent(
         name="researcher",
-        tools=[get_query_engine_tool()],
-        role="expert in retrieving any unknown content",
-        system_prompt="You are a researcher agent. You are given a researching task. You must use your tools to complete the research.",
+        tools=[get_query_engine_tool(), *duckduckgo_search_tools],
+        role="expert in retrieving any unknown content or searching for images from the internet",
+        system_prompt="You are a researcher agent. You are given a researching task. You must use tools to retrieve information from the knowledge base and search for needed images from the internet for the post.",
         chat_history=chat_history,
     )
diff --git a/templates/types/multiagent/fastapi/app/tools/artifact.py b/templates/types/multiagent/fastapi/app/tools/artifact.py
@@ -1,8 +1,11 @@
+import logging
 import os
 import re
 from enum import Enum
 from io import BytesIO
 
+from llama_index.core.tools.function_tool import FunctionTool
+
 OUTPUT_DIR = "output/tools"
 
 
@@ -11,43 +14,92 @@ class ArtifactType(Enum):
     HTML = "html"
 
 
-HTML_FILE_TEMPLATE = """
+COMMON_STYLES = """
+body {
+    font-family: Arial, sans-serif;
+    line-height: 1.3;
+    color: #333;
+}
+h1, h2, h3, h4, h5, h6 {
+    margin-top: 1em;
+    margin-bottom: 0.5em;
+}
+p {
+    margin-bottom: 0.7em;
+}
+code {
+    background-color: #f4f4f4;
+    padding: 2px 4px;
+    border-radius: 4px;
+}
+pre {
+    background-color: #f4f4f4;
+    padding: 10px;
+    border-radius: 4px;
+    overflow-x: auto;
+}
+table {
+    border-collapse: collapse;
+    width: 100%;
+    margin-bottom: 1em;
+}
+th, td {
+    border: 1px solid #ddd;
+    padding: 8px;
+    text-align: left;
+}
+th {
+    background-color: #f2f2f2;
+    font-weight: bold;
+}
+img {
+    max-width: 90%;
+    height: auto;
+    display: block;
+    margin: 1em auto;
+    border-radius: 10px;
+}
+"""
+
+HTML_SPECIFIC_STYLES = """
+body {
+    max-width: 800px;
+    margin: 0 auto;
+    padding: 20px;
+}
+"""
+
+PDF_SPECIFIC_STYLES = """
+@page {
+    size: letter;
+    margin: 2cm;
+}
+body {
+    font-size: 11pt;
+}
+h1 { font-size: 18pt; }
+h2 { font-size: 16pt; }
+h3 { font-size: 14pt; }
+h4, h5, h6 { font-size: 12pt; }
+pre, code {
+    font-family: Courier, monospace;
+    font-size: 0.9em;
+}
+"""
+
+HTML_TEMPLATE = """
 <!DOCTYPE html>
 <html lang="en">
 <head>
     <meta charset="UTF-8">
     <meta name="viewport" content="width=device-width, initial-scale=1.0">
     <style>
-        body {{
-            font-family: Arial, sans-serif;
-            line-height: 1.3;
-            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;
-        }}
+        {common_styles}
+        {specific_styles}
     </style>
 </head>
 <body>
-    {html_content}
+    {content}
 </body>
 </html>
 """
@@ -66,10 +118,23 @@ def _generate_html_content(cls, original_content: str) -> str:
                 "Failed to import required modules. Please install markdown."
             )
 
-        # Convert markdown to HTML
-        html_content = markdown.markdown(original_content)
+        # Convert markdown to HTML with fenced code and table extensions
+        html_content = markdown.markdown(
+            original_content, extensions=["fenced_code", "tables"]
+        )
         return html_content
 
+    @classmethod
+    def _generate_html_document(cls, html_content: str) -> str:
+        """
+        Generate a complete HTML document with the given HTML content.
+        """
+        return HTML_TEMPLATE.format(
+            common_styles=COMMON_STYLES,
+            specific_styles=HTML_SPECIFIC_STYLES,
+            content=html_content,
+        )
+
     @classmethod
     def _generate_pdf(cls, html_content: str) -> BytesIO:
         """
@@ -82,35 +147,24 @@ def _generate_pdf(cls, html_content: str) -> BytesIO:
                 "Failed to import required modules. Please install xhtml2pdf."
             )
 
+        pdf_html = HTML_TEMPLATE.format(
+            common_styles=COMMON_STYLES,
+            specific_styles=PDF_SPECIFIC_STYLES,
+            content=html_content,
+        )
+
         buffer = BytesIO()
         pdf = pisa.pisaDocument(
-            BytesIO(html_content.encode("UTF-8")),
-            buffer,
-            encoding="UTF-8",
-            path=".",
-            link_callback=None,
-            debug=0,
-            default_css=None,
-            xhtml=False,
-            xml_output=None,
-            ident=0,
-            show_error_as_pdf=False,
-            quiet=True,
-            capacity=100 * 1024 * 1024,
-            raise_exception=True,
+            BytesIO(pdf_html.encode("UTF-8")), buffer, encoding="UTF-8"
         )
+
         if pdf.err:
+            logging.error(f"PDF generation failed: {pdf.err}")
             raise ValueError("PDF generation failed")
+
         buffer.seek(0)
         return buffer
 
-    @classmethod
-    def _generate_html(cls, html_content: str) -> str:
-        """
-        Generate a complete HTML document with the given HTML content.
-        """
-        return HTML_FILE_TEMPLATE.format(html_content=html_content)
-
     @classmethod
     def generate_artifact(
         cls, original_content: str, artifact_type: str, file_name: str
@@ -119,10 +173,11 @@ def generate_artifact(
         To generate artifact as PDF or HTML file.
         Parameters:
             original_content: str (markdown style)
-            artifact_type: str (pdf or html) specify the type of the file format based on the use case
+            artifact_type: str (pdf or html or image) specify the type of the file format based on the use case
             file_name: str (name of the artifact file) must be a valid file name, no extensions needed
         Returns:
-            str (URL to the artifact file): A file URL ready to serve.
+            str (URL to the artifact file): A file URL ready to serve (pdf or html).
+            list[str] (URLs to the artifact files): A list of image URLs ready to serve (png).
         """
         try:
             artifact_type = ArtifactType(artifact_type.lower())
@@ -133,15 +188,12 @@ def generate_artifact(
         # Always generate html content first
         html_content = cls._generate_html_content(original_content)
 
-        # Based on the type of artifact, generate the corresponding file
         if artifact_type == ArtifactType.PDF:
-            content = cls._generate_pdf(cls._generate_html(html_content))
+            content = cls._generate_pdf(html_content)
             file_extension = "pdf"
         elif artifact_type == ArtifactType.HTML:
-            content = BytesIO(cls._generate_html(html_content).encode("utf-8"))
+            content = BytesIO(cls._generate_html_document(html_content).encode("utf-8"))
             file_extension = "html"
-        else:
-            raise ValueError(f"Unexpected artifact type: {artifact_type}")
 
         file_name = cls._validate_file_name(file_name)
         file_path = os.path.join(OUTPUT_DIR, f"{file_name}.{file_extension}")
@@ -176,3 +228,7 @@ def _validate_file_name(file_name: str) -> str:
             return file_name
         else:
             raise ValueError("File name is not allowed to contain special characters.")
+
+
+def get_tools(**kwargs):
+    return [FunctionTool.from_defaults(ArtifactGenerator.generate_artifact)]
diff --git a/templates/types/multiagent/fastapi/app/tools/duckduckgo.py b/templates/types/multiagent/fastapi/app/tools/duckduckgo.py
@@ -0,0 +1,68 @@
+from llama_index.core.tools.function_tool import FunctionTool
+
+
+def duckduckgo_search(
+    query: str,
+    region: str = "wt-wt",
+    max_results: int = 10,
+):
+    """
+    Use this function to search for any query in DuckDuckGo.
+    Args:
+        query (str): The query to search in DuckDuckGo.
+        region Optional(str): The region to be used for the search in [country-language] convention, ex us-en, uk-en, ru-ru, etc...
+        max_results Optional(int): The maximum number of results to be returned. Default is 10.
+    """
+    try:
+        from duckduckgo_search import DDGS
+    except ImportError:
+        raise ImportError(
+            "duckduckgo_search package is required to use this function."
+            "Please install it by running: `poetry add duckduckgo_search` or `pip install duckduckgo_search`"
+        )
+
+    params = {
+        "keywords": query,
+        "region": region,
+        "max_results": max_results,
+    }
+    results = []
+    with DDGS() as ddg:
+        results = list(ddg.text(**params))
+    return results
+
+
+def duckduckgo_image_search(
+    query: str,
+    region: str = "wt-wt",
+    max_results: int = 10,
+):
+    """
+    Use this function to search for images in DuckDuckGo.
+    Args:
+        query (str): The query to search in DuckDuckGo.
+        region Optional(str): The region to be used for the search in [country-language] convention, ex us-en, uk-en, ru-ru, etc...
+        max_results Optional(int): The maximum number of results to be returned. Default is 10.
+    """
+    try:
+        from duckduckgo_search import DDGS
+    except ImportError:
+        raise ImportError(
+            "duckduckgo_search package is required to use this function."
+            "Please install it by running: `poetry add duckduckgo_search` or `pip install duckduckgo_search`"
+        )
+    params = {
+        "keywords": query,
+        "region": region,
+        "max_results": max_results,
+    }
+    with DDGS() as ddg:
+        results = list(ddg.images(**params))
+    return results
+
+
+def get_tools(**kwargs):
+    return [
+        FunctionTool.from_defaults(duckduckgo_search),
+        FunctionTool.from_defaults(duckduckgo_image_search),
+    ]
