Update app.py

app.py

@@ -29,27 +29,35 @@ def wikipedia_search(query: str) -> str:
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WHERE YOU CAN BUILD WHAT YOU WANT ------
-class BasicAgent:
-    def __init__(self):
-        print('BasicAgent initialized.')
-
-    def __call__(self, question: str) -> str:
-        print(f"Agent received question: {question[:50]}...")
+def web_search(query: str):
+    """Search the web for information."""
+    return DuckDuckGoSearchRun().run(query)
 
-        web_search_tool = ToolNode([DuckDuckGoSearchRun().run])
-        open_webpage_tool = ToolNode([wikipedia_search])
+def wikipedia_search(query: str):
+    """Search Wikipedia for information."""
+    return wikipedia.run(query)
 
-        agent = ToolCallingAgent(
-            tools=[web_search_tool, open_webpage_tool],
+class BasicAgent:
+    def __init__(self, model_id="ollama/qwen2:7b", api_base="http://192.168.1.77:11434"):
+        print('BasicAgent initialized.')
+        self.tool_node = ToolNode([web_search, wikipedia_search])
+        self.agent = ToolCallingAgent(
+            tools=[self.tool_node],
             model=LiteLLMModel(
-                model_id="ollama/qwen2:7b",
-                api_base='http://192.168.1.77:11434'
+                model_id=model_id,
+                api_base=api_base
             ),
             max_steps=15,
             verbosity_level=2,
         )
 
-        return agent.run(question)
+    def __call__(self, question: str) -> str:
+        print(f"Agent received question: {question[:50]}...")
+        try:
+            return self.agent.run(question)
+        except Exception as e:
+            print(f"Error in agent: {e}")
+            return f"AGENT ERROR: {e}"
     
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """