Update app.py

app.py

@@ -312,40 +312,54 @@ class BasicAgent:
     def __init__(self):
         self.llm = GeminiLLM()
         self.chain = LLMChain(prompt=PromptTemplate.from_template("{question}"), llm=self.llm)
-        print("BasicAgent con GeminiLLM inicializado.")
-
+    
+    def call_tool(self, tool_name: str, *args, **kwargs):
+        # Simple mapping from tool name to function
+        tools = {
+            "wiki_search": wiki_search,
+            "load_file": load_file,
+            "web_search": web_search,
+            "arxiv_search": arxiv_search,
+            "multiply": multiply,
+            "add": add,
+            "subtract": subtract,
+            "divide": divide,
+            "modulus": modulus,
+            "power": power,
+            "square_root": square_root,
+            "extract_numbers": extract_numbers,
+            "extract_keywords": extract_keywords,
+            "extract_names": extract_names,
+            "find_non_commutative_pairs": find_non_commutative_pairs,
+        }
+        func = tools.get(tool_name)
+        if func:
+            return func(*args, **kwargs)
+        else:
+            return f"Tool '{tool_name}' not found."
+    
     def __call__(self, question: str) -> str:
-        print(f"Pregunta recibida: {question[:50]}...")
-
-        max_retries = 5
-        backoff = 2  # segundos
-
-        for attempt in range(max_retries):
-            try:
-                respuesta = self.chain.run({"question": question})
-                print(f"Respuesta: {respuesta}")
-                return respuesta
-            except Exception as e:
-                error_str = str(e).lower()
-                print(f"Error en intento {attempt + 1}: {e}")
-
-                # Manejar error 503 específico
-                if "503" in error_str or "model is overloaded" in error_str:
-                    if attempt < max_retries - 1:
-                        print(f"Modelo saturado, reintentando en {backoff} segundos...")
-                        time.sleep(backoff)
-                        backoff *= 2
-                        continue
-                    else:
-                        return "El modelo está saturado. Por favor, inténtalo más tarde."
-                # Otros errores que pueden ser manejados
-                elif "quota" in error_str:
-                    return "Se ha superado la cuota de la API. Intenta nuevamente más tarde."
-                elif "timeout" in error_str:
-                    return "La solicitud tardó demasiado. Intenta nuevamente."
-                else:
-                    return "Error procesando la pregunta."
+        question = question.lower()
+
+        # Very simple logic to decide which tool to call based on question keywords
+        if "wikipedia" in question or "who is" in question or "what is" in question:
+            # Extract query for Wikipedia (here just remove the word "wikipedia")
+            query = question.replace("wikipedia", "").strip()
+            docs = self.call_tool("wiki_search", query=query, load_max_docs=2)
+            # Return titles of found documents
+            return "\n".join([doc.metadata.get("title", "No title") for doc in docs])
+
+        elif "add" in question or "plus" in question:
+            # Extract numbers and add (simple demo)
+            nums = self.call_tool("extract_numbers", question)
+            if len(nums) >= 2:
+                return str(self.call_tool("add", nums[0], nums[1]))
+            else:
+                return "I couldn’t find enough numbers to add."
 
+        else:
+            # If no tool matched, fallback to the LLM
+            return self.chain.run({"question": question})
 
 
 def run_and_submit_all(profile: gr.OAuthProfile | None):