Update agent.py

agent.py

@@ -81,7 +81,9 @@ def calculator(inputs: dict):
         return a % b
     else:
         return "Unknown operation"
-        
+
+
+
 
 @tool
 def wiki_search(query: str) -> str:
@@ -282,12 +284,16 @@ enabled_tool_names = [
 
 tools = [tool_map[name] for name in enabled_tool_names]
 
+
+# Safe version
 tools = []
 for name in enabled_tool_names:
-    if name in tool_map:
-        tools.append(tool_map[name])
-    else:
-        print(f"⚠️ Warning: Tool '{name}' not found in tool_map. Skipping.")
+    if name not in tool_map:
+        print(f"❌ Tool not found: {name}")
+        continue
+    tools.append(tool_map[name])
+
+    
 
 # -------------------------------
 # Step 2: Load the JSON file or tasks (Replace this part if you're loading tasks dynamically)
@@ -503,57 +509,61 @@ def get_llm(provider: str, config: dict):
 
 
 
-
 # ----------------------------------------------------------------
 # Planning & Execution Logic
 # ----------------------------------------------------------------
-def planner(question: str, tools: list) -> list:
+def planner(question: str, tools: list) -> tuple:
+    """
+    Select the best-matching tool(s) for a question based on keyword-based intent detection and tool metadata.
+    Returns the detected intent and matched tools.
+    """
     question = question.lower().strip()
 
-    # Define intent-based keywords (broad categories)
+    # Define intent-based keywords
     intent_keywords = {
         "math": ["calculate", "evaluate", "add", "subtract", "multiply", "divide", "modulus", "plus", "minus", "times"],
         "wiki_search": ["who is", "what is", "define", "explain", "tell me about", "overview of"],
         "web_search": ["search", "find", "look up", "google", "latest news", "current info"],
-        "arxiv": ["arxiv", "research paper", "scientific paper", "preprint"],
-        "youtube": ["youtube", "watch", "play video", "show me a video"],
+        "arxiv_search": ["arxiv", "research paper", "scientific paper", "preprint"],
+        "get_youtube_transcript": ["youtube", "watch", "play video", "show me a video"],
         "video_analysis": ["analyze video", "summarize video", "video content"],
         "data_analysis": ["analyze", "plot", "graph", "data", "visualize"],
         "wikidata_query": ["wikidata", "sparql", "run sparql", "query wikidata"],
-        "general_qa": ["why", "how", "difference between", "compare", "what happens", "reason for", "cause of", "effect of"]
+        "default": ["why", "how", "difference between", "compare", "what happens", "reason for", "cause of", "effect of"]
     }
 
-    matched_tools = []
-
-    # Try to find tools whose description matches intent
+    # Step 1: Identify intent
+    detected_intent = None
     for intent, keywords in intent_keywords.items():
         if any(keyword in question for keyword in keywords):
-            for tool in tools:
-                description = getattr(tool, "description", "").lower()
-                name = getattr(tool, "name", "").lower()
-
-                # Match based on intent keywords appearing in tool description or name
-                if intent in description or intent in name:
-                    matched_tools.append(tool)
-
-            # Break after first matching intent — you can remove this if you want to allow multi-intent matching
-            if matched_tools:
-                break
+            detected_intent = intent
+            break
 
-    # Fallback 1: try using general-purpose tools if available
+    # Step 2: Match tools by intent
+    matched_tools = []
+    if detected_intent:
+        for tool in tools:
+            name = getattr(tool, "name", "").lower()
+            description = getattr(tool, "description", "").lower()
+            if detected_intent in name or detected_intent in description:
+                matched_tools.append(tool)
+
+    # Step 3: Fallback to general-purpose/default tools if no match found
     if not matched_tools:
         matched_tools = [
-            tool for tool in tools 
-            if "default" in getattr(tool, "name", "").lower() 
+            tool for tool in tools
+            if "default" in getattr(tool, "name", "").lower()
             or "qa" in getattr(tool, "description", "").lower()
         ]
 
-    # Fallback 2: return first tool to prevent failure
-    return matched_tools if matched_tools else [tools[0]]
+    return detected_intent, matched_tools if matched_tools else [tools[0]]
 
 
 
 def task_classifier(question: str) -> str:
+    """
+    Classifies the question into one of the predefined task categories.
+    """
     question = question.lower().strip()
 
     # Context-aware intent patterns
@@ -599,6 +609,44 @@ def task_classifier(question: str) -> str:
 
     return "default"
 
+def select_tool_and_run(question: str, tools: list) -> any:
+    """
+    Select the best tool based on question intent and run it.
+    """
+    # Step 1: Use task_classifier to identify the intent of the question
+    intent = task_classifier(question)
+
+    # Step 2: Use the intent_tool_map to map the intent to the tool
+    intent_tool_map = {
+        "math": "calculator",
+        "wiki_search": "wiki_tool",
+        "web_search": "web_search_tool",
+        "arxiv": "arxiv_tool",
+        "get_youtube_transcript": "youtube_tool",
+        "video_analysis": "video_tool",
+        "analyze_tool": "analyze_tool",
+        "wikidata_query": "wikidata_query",
+        "default": "default_tool"
+    }
+
+    tool_name = intent_tool_map.get(intent, "default_tool")
+
+    # Step 3: Get the selected tool
+    tool_func = tools.get(tool_name)
+
+    if not tool_func:
+        print(f"No matching tool found for {intent}")
+        return None
+
+    # Step 4: Call the tool function and return the result
+    return tool_func.run(question)  # or tool_func(question), depending on how your tools are structured
+
+# Run the function to select and execute the tool
+result = select_tool_and_run(question, tools)
+
+# Output the result
+print(result)
+
 
 # Function to extract math operation from the question
 
@@ -640,18 +688,31 @@ def extract_math_from_question(question: str):
 
 # Example tool set (adjust these to match your actual tool names)
 tools = {
-    "math": calc_tool,          # Example tool for math tasks
-    "wiki_search": wiki_tool,   # Example tool for wiki search tasks
-    "retriever": retriever_tool, # Example tool for retriever tasks
-    "default": default_tool     # Fallback tool
+    "math": calculator,
+    "wiki_search": wiki_tool,
+    "web_search": web_search_tool,
+    "arvix_search": arvix_tool,
+    "retriever": retriever_tool,
+    "get_youtube_transcript": youtube_tool,
+    "extract_video_id": video_tool,
+    "analyze_attachment": analyze_tool,
+    "wikidata_query": wikiq_tool,
+    "default": default_tool
 }
 
+
 # The task order can also include the tools for each task
 priority_order = [
-    {"task": "math", "tool": "math"},        # Priority task and tool
+    {"task": "math", "tool": "math"},
     {"task": "wiki_search", "tool": "wiki_search"},
+    {"task": "web_search", "tool": "web_search"},
+    {"task": "arvix_search", "tool": "arvix_search"},
+    {"task": "wikidata_query", "tool": "wikidata_query"},
     {"task": "retriever", "tool": "retriever"},
-    {"task": "default", "tool": "default"}  # Fallback tool
+    {"task": "get_youtube_transcript", "tool": "get_youtube_transcript"},
+    {"task": "extract_video_id", "tool": "extract_video_id"},
+    {"task": "analyze_attachment", "tool": "analyze_attachment"},
+    {"task": "default", "tool": "default"}  # Fallback
 ]
 
 def decide_task(state: dict) -> str: