updated with max results and wikipedia serach tool

app.py

@@ -1,6 +1,7 @@
+import dataclasses
 import os
 from math import sqrt
-from typing import Dict
+from typing import Dict, List
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader
 from langchain_community.document_loaders import ArxivLoader
@@ -15,34 +16,39 @@ from smolagents import CodeAgent, tool, InferenceClientModel
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 
+@dataclasses.dataclass
+class WikiSourceDocument:
+    source: str
+    page: str
+    page_content: str
+
 @tool
-def wiki_search(query: str) -> Dict[str, str]:
+def wiki_search(query: str, load_max_docs: int=3) -> List[WikiSourceDocument]:
     """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}
+    search_docs = WikipediaLoader(query=query, load_max_docs=load_max_docs).load()
+
+    return [WikiSourceDocument(
+        source=doc.metadata["Source"],
+        page=doc.metadata["Page"],
+        page_content=doc.page_content
+    ) for doc in search_docs]
 
+@tool
+def load_file(file_id: str) -> str:
+    """Load a file from the Hugging Face Hub. It returns the content in bytes.
+    Args:
+        file_id: The file ID to load."""
+    return requests.get(f"https://agents-course-unit4-scoring.hf.space/files/{file_id}").content
 
 @tool
-def web_search(query: str) -> Dict[str, str]:
+def web_search(query: str, max_results: int) -> Dict[str, str]:
     """Search Tavily for a query and return maximum 3 results.
     Args:
         query: The search query."""
-    search_docs = TavilySearchResults(max_results=3).invoke(input=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}
+    search_docs = TavilySearchResults(max_results=max_results).invoke(input=query)
+    return {"web_results": search_docs}
 
 
 @tool
@@ -53,7 +59,7 @@ def arxiv_search(query: str) -> Dict[str, str]:
     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>'
+            f'<Document Title="{doc.metadata["Title"]}" Published="{doc.metadata['Published']}" Authors="{doc.metadata['Authors']} Summary={doc.metadata['Summary']}"/>\n{doc.page_content}\n</Document>'
             for doc in search_docs
         ]
     )
@@ -154,6 +160,7 @@ class BasicAgent:
         self.agent = CodeAgent(
             tools=[multiply, add, subtract, power, square_root, modulus, wiki_search, web_search, arxiv_search],
             model=InferenceClientModel(model_id=model_id, token=os.getenv("HF_TOKEN")),
+            max_steps=10,
         )
         print("BasicAgent initialized.")
 

requirements.txt

@@ -3,4 +3,6 @@ requests
 smolagents
 pandas
 langchain-community
-wikipedia
+wikipedia
+arxiv
+pymupdf