Refactor agent tools: add WikipediaSearchTool for querying Wikipedia and update imports

agent/agents.py

@@ -5,16 +5,17 @@ from google import genai
 from google.genai import types
 from google.genai.types import GenerateContentConfig, ThinkingConfig
 from pydantic import BaseModel
-from smolagents import CodeAgent, VisitWebpageTool, WebSearchTool
+from smolagents import CodeAgent, Tool, VisitWebpageTool, WebSearchTool
 
 from agent.prompts import system_prompt
 from agent.tools import (
+    CodeExecutionTool,
     DownloadFile,
     ReadExcelFileBytes,
     TranscribeAudioBytes,
     TranscribeYoutubeVideo,
     UnderstandImageBytes,
-    CodeExecutionTool,
+    WikipediaSearchTool,
 )
 from agent.utils import gemini_client, gemini_model_liteLLM
 
@@ -54,6 +55,7 @@ class BasicAgent:
         read_excel_file_tool = ReadExcelFileBytes()
         understand_image_bytes = UnderstandImageBytes()
         code_execution_tool = CodeExecutionTool()
+        wiki_retriever = WikipediaSearchTool()
 
         model = gemini_model_liteLLM(self.model)
 
@@ -71,6 +73,7 @@ class BasicAgent:
                 read_excel_file_tool,
                 understand_image_bytes,
                 code_execution_tool,
+                wiki_retriever,
             ],
             model=model,
             step_callbacks=STEP_CALLBACKS,

agent/tools.py

@@ -3,10 +3,33 @@ import io
 import requests
 from google import genai
 from google.genai import types
+from langchain_community.retrievers import WikipediaRetriever
 from PIL import Image
 from smolagents import Tool
 
 
+class WikipediaSearchTool(Tool):
+    name = "wikipedia_search"
+    description = """Search Wikipedia for a given query and return the summary of the first result.
+    This tool uses the WikipediaRetriever to fetch the summary of the first search result.
+    """
+    inputs = {
+        "query": {
+            "type": "string",
+            "description": "The search query to look up on Wikipedia",
+        },
+    }
+    output_type = "string"
+
+    def forward(self, query: str):
+        retriever = WikipediaRetriever()
+        docs = retriever.invoke(query)
+        if docs:
+            return docs
+        else:
+            return "No results found."
+
+
 class CodeExecutionTool(Tool):
     name = "execute_code"
     description = """Execute Python code and answer the question if provided.