Update basic_agent.py

basic_agent.py

@@ -6,25 +6,32 @@ from tools import (
 )
 
 class BasicAgent:
+    def __init__(self):
+        self.tool_registry = {
+            "analyze_youtube_video": (analyze_youtube_video, "url"),
+            "extract_number_from_text": (extract_number_from_text, "text"),
+            "search_wikipedia": (search_wikipedia, "query"),
+        }
+
     def __call__(self, question: str) -> str:
         try:
-            tool = self.select_tool(question)
-            result = tool.invoke({"query": question})
-            return result
+            tool, input_key = self.select_tool(question)
+            return tool.invoke({input_key: question})
         except Exception as e:
             return f"Error: {e}"
 
     def select_tool(self, question: str):
         q = question.lower()
-    
+
         if "youtube.com" in q or "youtu.be" in q or "in the video" in q:
-            return analyze_youtube_video  # direct reference, not index
-    
+            return self.tool_registry["analyze_youtube_video"]
+
         elif "how many" in q or "number of" in q:
-            return extract_number_from_text
-    
+            return self.tool_registry["extract_number_from_text"]
+
         elif "wikipedia" in q or "encyclopedia" in q or "who is" in q:
-            return search_wikipedia
-    
-        return extract_number_from_text  # fallback tool
+            return self.tool_registry["search_wikipedia"]
+
+        return self.tool_registry["extract_number_from_text"]  # fallback
+