Update agent.py

agent.py

@@ -58,7 +58,10 @@ from langchain.agents import AgentType
 from typing import Union, List
 from functools import reduce
 import operator
-
+from typing import Union
+from functools import reduce
+from youtube_transcript_api import YouTubeTranscriptApi
+from youtube_transcript_api._errors import TranscriptsDisabled, VideoUnavailable
 
 
 load_dotenv()
@@ -122,18 +125,17 @@ def calculator(inputs: Union[str, dict]):
 
 @tool
 def wiki_search(query: str) -> str:
-    """Search Wikipedia for a query and return maximum 2 results.
-    
-    Args:
-        query: The search query."""
+    """Search Wikipedia for a query and return up to 2 results."""
     search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+
     formatted_search_docs = "\n\n---\n\n".join(
         [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            f'<Document source="{doc.metadata.get("source", "Wikipedia")}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
             for doc in search_docs
-        ])
+        ]
+    )
     return formatted_search_docs
-
+    
     
 @tool
 def wikidata_query(query: str) -> str:
@@ -258,6 +260,8 @@ def get_youtube_transcript(url: str) -> str:
     except Exception as e:
         raise ValueError(f"Failed to fetch transcript: {e}")
 
+
+
 @tool
 def extract_video_id(url: str) -> str:
     """
@@ -302,7 +306,7 @@ tool_map = {
     "math": calculator,
     "wiki_search": wiki_search,
     "web_search": web_search,
-    "arvix_search": arvix_search,
+    "arxiv_search": arxiv_search,
     "get_youtube_transcript": get_youtube_transcript,
     "extract_video_id": extract_video_id,
     "analyze_attachment": analyze_attachment,
@@ -314,7 +318,7 @@ enabled_tool_names = [
     "math",
     "wiki_search",
     "web_search",
-    "arvix_search",
+    "arxiv_search",
     "get_youtube_transcript",
     "extract_video_id",
     "analyze_attachment",
@@ -480,9 +484,9 @@ def process_question(question):
     # Step 5: Execute task (with error handling)
     try:
         if task_type == "wiki_search":
-            response = wiki_tool(question)
+            response = wiki_search(question)
         elif task_type == "math":
-            response = calc_tool(question)
+            response = calculator(question)
         else:
             response = "Default answer logic"
         
@@ -566,12 +570,14 @@ def planner(question: str, tools: list) -> tuple:
         "web_search": ["search", "find", "look up", "google", "latest news", "current info"],
         "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"],
+        "extract_video_id": ["analyze video", "summarize video", "video content"],
         "data_analysis": ["analyze", "plot", "graph", "data", "visualize"],
         "wikidata_query": ["wikidata", "sparql", "run sparql", "query wikidata"],
         "default": ["why", "how", "difference between", "compare", "what happens", "reason for", "cause of", "effect of"]
     }
 
+
+
     # Step 1: Identify intent
     detected_intent = None
     for intent, keywords in intent_keywords.items():
@@ -599,6 +605,7 @@ def planner(question: str, tools: list) -> tuple:
     return detected_intent, matched_tools if matched_tools else [tools[0]]
 
 
+import json
 
 def task_classifier(question: str) -> str:
     """
@@ -625,12 +632,12 @@ def task_classifier(question: str) -> str:
     elif any(phrase in question for phrase in [
         "arxiv", "latest research", "scientific paper", "research paper", "preprint"
     ]):
-        return "arxiv"
+        return "arxiv_search"
 
     elif any(phrase in question for phrase in [
         "youtube", "watch", "play the video", "show me a video"
     ]):
-        return "youtube"
+        return "get_youtube_transcript"
 
     elif any(phrase in question for phrase in [
         "analyze video", "summarize video", "what happens in the video", "video content"
@@ -650,42 +657,45 @@ def task_classifier(question: str) -> str:
     return "default"
 
 
+def select_tool_and_run(question: str, tools: dict):
+    # Step 1: Classify intent
+    intent = task_classifier(question)  # assuming task_classifier maps the question to intent
 
-def select_tool_and_run(question: str, tools: list):
-    # Classify intent
-    intent = task_classifier(question)
-
-    # Map intent to expected tool names
+    # Map intent to tool names
     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",
-        "data_analysis": "analyze_tool",         # FIXED key from 'analyze_tool'
-        "wikidata_query": "wikidata_query",
-        "default": "default_tool"
+        "math": "calculator",                                # maps to tools["math"] → calculator
+        "wiki_search": "wiki_search",                        # → wiki_search
+        "web_search": "web_search",                          # → web_search
+        "arxiv_search": "arxiv_search",                      # → arxiv_search (spelling fixed)
+        "get_youtube_transcript": "get_youtube_transcript",  # → get_youtube_transcript
+        "extract_video_id": "extract_video_id",              # adjust based on your tools
+        "analyze_attachment": "analyze_attachment",          # assuming analyze_attachment handles this
+        "wikidata_query": "wikidata_query",                  # → wikidata_query
+        "default": "default"                                 # → default_tool
     }
 
-    # Resolve the tool name for the detected intent
-    tool_name = intent_tool_map.get(intent, "default_tool")
+    # Get the corresponding tool name
+    tool_name = intent_tool_map.get(intent, "default")  # Default to "default" if no match
 
-    # Try to find the matching tool object from the list
-    tool_func = next((tool for tool in tools if getattr(tool, "name", "") == tool_name), None)
+    # Retrieve the tool from the tools dictionary
+    tool_func = tools.get(tool_name)
 
     if not tool_func:
-        print(f"No matching tool found for intent: {intent} (tool_name: {tool_name})")
-        return None
+        return f"Tool not found for intent '{intent}'"
 
-    # Try to parse question into JSON for tools expecting structured input
+    # Step 2: Run the tool
     try:
-        parsed_input = json.loads(question)
-    except json.JSONDecodeError:
-        parsed_input = question  # Fallback to raw string if not JSON
+        # If the tool needs JSON or structured data
+        try:
+            parsed_input = json.loads(question)
+        except json.JSONDecodeError:
+            parsed_input = question  # fallback to raw input if not JSON
 
-    # Run the tool and return result
-    return tool_func.run(parsed_input)
+        # Run the selected tool
+        print(f"Running tool: {tool_name} with input: {parsed_input}")  # log the tool name and input
+        return tool_func(parsed_input)
+    except Exception as e:
+        return f"Error while running tool '{tool_name}': {str(e)}"
 
 
 
@@ -728,26 +738,26 @@ def extract_math_from_question(question: str):
 
 
 # Example tool set (adjust these to match your actual tool names)
-tools = {
-    "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
+intent_tool_map = {
+    "math": "math",                                      # maps to tools["math"] → calculator
+    "wiki_search": "wiki_search",                        # → wiki_search
+    "web_search": "web_search",                          # → web_search
+    "arxiv_search": "arxiv_search",                      # → arxiv_search (spelling fixed)
+    "get_youtube_transcript": "get_youtube_transcript",  # → get_youtube_transcript
+    "extract_video_id": "extract_video_id",              # adjust based on your tools
+    "analyze_attachment": "analyze_attachment",          # assuming analyze_attachment handles this
+    "wikidata_query": "wikidata_query",                  # → wikidata_query
+    "default": "default"                                 # → default_tool
 }
 
 
+
 # The task order can also include the tools for each task
 priority_order = [
     {"task": "math", "tool": "math"},
     {"task": "wiki_search", "tool": "wiki_search"},
     {"task": "web_search", "tool": "web_search"},
-    {"task": "arvix_search", "tool": "arvix_search"},
+    {"task": "arxiv_search", "tool": "arxiv_search"},
     {"task": "wikidata_query", "tool": "wikidata_query"},
     {"task": "retriever", "tool": "retriever"},
     {"task": "get_youtube_transcript", "tool": "get_youtube_transcript"},
@@ -976,7 +986,7 @@ def build_graph(provider, model_config):
     calc_tool = calculator  # Math operations tool
     web_tool = web_search    # Web search tool
     wiki_tool = wiki_search  # Wikipedia search tool
-    arvix_tool = arvix_search  # Arxiv search tool
+    arxiv_tool = arxiv_search # Arxiv search tool
     youtube_tool = get_youtube_transcript  # YouTube transcript extraction
     video_tool = extract_video_id  # Video ID extraction tool
     analyze_tool = analyze_attachment  # File analysis tool
@@ -990,7 +1000,7 @@ def build_graph(provider, model_config):
         wiki_tool, 
         calc_tool, 
         web_tool, 
-        arvix_tool,
+        arxiv_tool,
         youtube_tool, 
         video_tool, 
         analyze_tool,