Update agent.py

agent.py

@@ -1,4 +1,6 @@
-from smolagents import CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, LiteLLMModel, tool, load_tool # HfApiModel, OpenAIServerModel
+from smolagents import CodeAgent,  LiteLLMModel, tool, load_tool # DuckDuckGoSearchTool, WikipediaSearchTool, HfApiModel, OpenAIServerModel
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader
 import asyncio
 import os
 import yaml
@@ -32,7 +34,35 @@ from io import BytesIO
 #        return "\n".join([f"{item['title']}: {item['link']}" for item in results]) or "No results found."
 #    except Exception as e:
 #        return f"Error performing Google search: {str(e)}"
-        
+
+@tool
+def wiki_search(query: str) -> str:
+    """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 web_search(query: str) -> str:
+    """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 ImageAnalysisTool(image_path: str) -> str:
     """Tool for analyzing images using computer vision
@@ -102,9 +132,11 @@ class MagAgent:
         self.agent = CodeAgent(
             model= model,
             tools=[
+                wiki_search,
+                web_search,
 #                GoogleSearchTool,
-                DuckDuckGoSearchTool(),
-                WikipediaSearchTool(),
+#                DuckDuckGoSearchTool(),
+#                WikipediaSearchTool(),
                 ImageAnalysisTool,
 #                LocalFileAudioTool()
             ]