Update agents.py

agents.py

@@ -51,42 +51,52 @@ def modulus(a: int, b: int) -> int:
 # ---- Search Tools ---- #
 @tool
 def search_wikipedia(query: str) -> str:
-    """Returns up to 2 documents related to a query from Wikipedia."""
-    docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    return {"wiki_results": "\n\n---\n\n".join(
-        f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}'
-        for doc in docs
-    )}
+    """Search Wikipedia for a query and return maximum 2 results.
+    
+    Args:
+        query: The search query."""
+    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"wiki_results": formatted_search_docs}
 
 @tool
 def search_web(query: str) -> str:
-    """Fetches up to 3 web results using Tavily."""
-    results = TavilySearchResults(max_results=3).invoke(query=query)
-    return {"web_results": "\n\n---\n\n".join(
-        f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}'
-        for doc in results
-    )}
+    """Search Tavily for a query and return maximum 3 results.
+    
+    Args:
+        query: The search query."""
+    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"web_results": formatted_search_docs}
 
 @tool
 def search_arxiv(query: str) -> str:
-    """Retrieves up to 3 papers related to the query from ArXiv."""
-    results = ArxivLoader(query=query, load_max_docs=3).load()
-    return {"arvix_results": "\n\n---\n\n".join(
-        f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}'
-        for doc in results
-    )}
-
-
-system_message = SystemMessage(content="""You are a helpful assistant tasked with answering questions using a set of tools. Now, I will ask you a question. Report your thoughts, and finish your answer with the following template:
-
-FINAL ANSWER: [YOUR FINAL ANSWER]
-
-YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma-separated list of numbers and/or strings. 
-- If you are asked for a number, don't use a comma in the number and avoid units like $ or % unless specified otherwise.
-- If you are asked for a string, avoid using articles and abbreviations (e.g. for cities), and write digits in plain text unless specified otherwise.
-- If you are asked for a comma-separated list, apply the above rules depending on whether each item is a number or string.
-
-Your answer should start only with "FINAL ANSWER: ", followed by your result.""")
+    """Search Arxiv for a query and return maximum 3 result.
+    
+    Args:
+        query: The search query."""
+    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"arvix_results": formatted_search_docs}
+
+
+system_message = SystemMessage(content="""You are a helpful assistant tasked with answering questions using a set of tools. 
+Now, I will ask you a question. Report your thoughts, and finish your answer with the following template: 
+FINAL ANSWER: [YOUR FINAL ANSWER]. 
+YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
+Your answer should only start with "FINAL ANSWER: ", then follows with the answer. """)
 
 toolset = [
     multiply,