Update app.py

app.py

@@ -25,68 +25,29 @@ load_dotenv()
 
 
 
-@tool
-def wiki_search(query: str) -> str:
-    """Search Wikipedia for a query and return up to 2 results."""
-    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata.get("source", "Wikipedia")}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
-            for doc in search_docs
-        ]
-    )
-    return formatted_search_docs
-    
-    
-@tool
-def wikidata_query(query: str) -> str:
-    """
-    Run a SPARQL query on Wikidata and return results.
-    """
-    endpoint_url = "https://query.wikidata.org/sparql"
-    headers = {
-        "Accept": "application/sparql-results+json"
-    }
-    response = requests.get(endpoint_url, headers=headers, params={"query": query})
-    data = response.json()
-    return json.dumps(data, indent=2)
-
+from smolagents import Tool
+from langchain_community.document_loaders import WikipediaLoader
 
-@tool
-def web_search(query: str) -> str:
-    """Search Tavily for a query and return up to 3 results."""
-    tavily_key = os.getenv("TAVILY_API_KEY")
-    
-    if not tavily_key:
-        return "Error: Tavily API key not set."
+class WikiSearchTool(Tool):
+    name = "wiki_search"
+    description = "Search Wikipedia for a query and return up to 2 results."
+    inputs = {
+        "query": {"type": "string", "description": "The search term for Wikipedia."}
+    }
+    output_type = "string"
 
-    search_tool = TavilySearchResults(tavily_api_key=tavily_key, max_results=3)
-    search_docs = search_tool.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 formatted_search_docs
+    def forward(self, query: str) -> str:
+        search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
 
+        formatted_search_docs = "\n\n---\n\n".join(
+            [
+                f'<Document source="{doc.metadata.get("source", "Wikipedia")}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+                for doc in search_docs
+            ]
+        )
+        return formatted_search_docs
 
-@tool
-def arxiv_search(query: str) -> str:
-    """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 formatted_search_docs
-
 
 
 class StringReverseTool(Tool):
@@ -227,7 +188,7 @@ class BasicAgent:
         )
 
         search_tool = DuckDuckGoSearchTool()
-        wiki_search_tool = wiki_search()
+        wiki_search_tool = WikiSearchTool()
         str_reverse_tool = StringReverseTool()
         keywords_extract_tool = KeywordsExtractorTool()
         speech_to_text_tool = SpeechToTextTool()