Update src/agent.py

src/agent.py

@@ -1,40 +1,84 @@
 from src.settings import Settings
-from smolagents import LiteLLMModel, ToolCallingAgent
+from smolagents import LiteLLMModel, ToolCallingAgent, CodeAgent, InferenceClientModel, DuckDuckGoSearchTool, Tool
 from tools.FinalAnswerTool import FinalAnswerTool
 from tools.ReadAudioTool import ReadAudioTool
 from tools.ReadImageTool import ReadImageTool
 from tools.ReadTextTool import ReadTextTool
 from tools.ReadVideoTool import ReadVideoTool
-from tools.WebSearchTool import WebSearchTool
+from tools.WebSearchTool import DuckDuckGoSearchTool
 from tools.YouTubeTool import YouTubeTool
 from tools.PythonRunnerTool import PythonRunnerTool
-from tools.WikipediaTool import WikipediaTool
+from tools.PythonCalcTool import PythonCalcTool
+from tools.SemanticScholar import AcademicPaperSearchTool
 from src.utils import InputTokenRateLimiter
+import wikipedia as wiki
+from markdownify import markdownify as to_markdown
 import time
 import random
-from litellm import completion
 
 settings = Settings()
+import litellm
+from litellm import completion 
+
+#litellm._turn_on_debug()
+class WikiTitleFinder(Tool):
+    name = "wiki_titles"
+    description = "Search for related Wikipedia page titles."
+    inputs = {"query": {"type": "string", "description": "Search query."}}
+    output_type = "string"
+
+    def forward(self, query: str) -> str:
+        results = wiki.search(query)
+        return ", ".join(results) if results else "No results."
+
+class WikiContentFetcher(Tool):
+    name = "wiki_page"
+    description = "Fetch Wikipedia page content."
+    inputs = {"page_title": {"type": "string", "description": "Wikipedia page title."}}
+    output_type = "string"
+
+    def forward(self, page_title: str) -> str:
+        try:
+            return to_markdown(wiki.page(page_title).html())
+        except wiki.exceptions.PageError:
+            return f"'{page_title}' not found."
+
+class MathSolver(Tool):
+    name = "math_solver"
+    description = "Safely evaluate basic math expressions."
+    inputs = {"input": {"type": "string", "description": "Math expression to evaluate."}}
+    output_type = "string"
+
+    def forward(self, input: str) -> str:
+        try:
+            return str(eval(input, {"__builtins__": {}}))
+        except Exception as e:
+            return f"Math error: {e}"
+        
 class BasicAgent():
     def __init__(self):
         self.model = LiteLLMModel(
             model_id=settings.llm_model_id, 
             api_key=settings.llm_api_key
         )
-        self.agent = ToolCallingAgent(
+        self.agent = CodeAgent(
+            model=self.model,
             tools=[
                 FinalAnswerTool(),
                 ReadAudioTool(),
                 ReadImageTool(),
                 ReadTextTool(),
                 ReadVideoTool(),
-                WebSearchTool(),
-                WikipediaTool(),
+                DuckDuckGoSearchTool(),   
+                WikiTitleFinder(),
+                WikiContentFetcher(),     
                 YouTubeTool(),
-                PythonRunnerTool()
-                # TODO: MP3 interpretation
+                PythonRunnerTool(),
+                PythonCalcTool(), 
+                AcademicPaperSearchTool(),
+                MathSolver()
             ],
-            max_steps=3,
+            max_steps=5,
             planning_interval=3,
             model=self.model
         )