added wiki search

agent.py

@@ -1,14 +1,13 @@
 '''
 TODO tools
-- web_search
 
 Facultative
 - wiki_search
 - arxiv_search
 '''
 import os
-from smolagents import CodeAgent, tool, DuckDuckGoSearchTool, OpenAIServerModel
-
+from smolagents import CodeAgent, tool, DuckDuckGoSearchTool, OpenAIServerModel, VisitWebpageTool, Tool
+from langchain_community import WikipediaLoader
 
 @tool
 def add(a:int, b:int) -> int:
@@ -88,13 +87,30 @@ def rounder(a:float, n:int) -> float:
 
   return round(a,n)
 
+@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}
+
 
 def get_agent() -> CodeAgent:
   search_tool = DuckDuckGoSearchTool()
+  web_page_tool = VisitWebpageTool()
+
+  api_key = os.getenv('CODESTRAL_API_KEY')
 
   model = OpenAIServerModel(
     model_id='codestral-latest',
     api_base="https://codestral.mistral.ai/v1/",
-    api_key=os.getenv('CODESTRAL_API_KEY'))
+    api_key=api_key)
 
-  return CodeAgent(tools=[add, subtract, multiply, divide, modulus, rounder, search_tool], model=model)
+  return CodeAgent(tools=[add, subtract, multiply, divide, modulus, rounder, search_tool, web_page_tool, wiki_search], model=model)

requirements.txt

@@ -1,3 +1,6 @@
 gradio
 requests
-smolagents[openai]
+smolagents
+smolagents[openai]
+langchain_community
+wikipedia