Update agent.py

agent.py

@@ -6,8 +6,14 @@ from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader
-from youtube_transcript_api import YouTubeTranscriptApi
+from youtube_transcript_api import YouTubeTranscriptApi, NoTranscriptFound
+from duckduckgo_search import DDGS
+from langchain_community.document_loaders import ArxivLoader
+from sympy import sympify
+from PIL import Image
 import re
+import requests
+from io import BytesIO
 
 # Load system prompt
 with open("system_prompt.txt", "r", encoding="utf-8") as f:
@@ -16,7 +22,7 @@ with open("system_prompt.txt", "r", encoding="utf-8") as f:
 # Tool: Wikipedia search
 @tool
 def wiki_search(query: str) -> str:
-    """Search Wikipedia for a query and return max 2 results."""
+    """Wikipedia search tool."""
     try:
         docs = WikipediaLoader(query=query, load_max_docs=2).load()
         return "\n\n---\n\n".join([doc.page_content for doc in docs])
@@ -26,30 +32,96 @@ def wiki_search(query: str) -> str:
 # Tool: Tavily web search
 @tool
 def web_search(query: str) -> str:
-    """Web search with Tavily (Google-like)."""
+    """Tavily web search tool."""
     try:
-        docs = TavilySearchResults(max_results=3).invoke(query)
-        return "\n\n---\n\n".join([doc.page_content for doc in docs])
+        results = TavilySearchResults(max_results=3).invoke(query)
+        if isinstance(results, list):
+            return "\n\n---\n\n".join([r["content"] if isinstance(r, dict) else str(r) for r in results])
+        return str(results)
     except Exception as e:
         return f"Web search failed: {e}"
 
-# Tool: YouTube transcript parser
+# Tool: DuckDuckGo search
 @tool
-def youtube_transcript(video_url: str) -> str:
-    """Extract transcript from a YouTube video URL."""
+def duckduckgo_search(query: str) -> str:
+    """DuckDuckGo search tool."""
     try:
-        video_id = re.search(r"v=([a-zA-Z0-9_-]{11})", video_url)
-        if not video_id:
-            return "Invalid YouTube URL"
-        transcript = YouTubeTranscriptApi.get_transcript(video_id.group(1))
-        text = " ".join([entry['text'] for entry in transcript])
-        return text
+        with DDGS() as ddgs:
+            results = ddgs.text(query, max_results=3)
+            return "\n\n---\n\n".join([r["body"] for r in results if "body" in r])
     except Exception as e:
-        return f"Transcript fetch failed: {e}"
+        return f"DuckDuckGo search failed: {e}"
 
-tools = [wiki_search, web_search, youtube_transcript]
+# Tool: YouTube transcript or duration extractor
+@tool
+def youtube_transcript(video_title_or_url: str) -> str:
+    """YouTube transcript or duration extractor tool."""
+    try:
+        with DDGS() as ddgs:
+            results = ddgs.videos(video_title_or_url, max_results=1)
+            if not results:
+                return "No video found by that title."
+            video = results[0]
+            video_url = video["url"]
+            duration = video.get("duration")
+            return f"Duration: {duration}"
+    except Exception as e:
+        return f"YouTube search failed: {e}"
+
+# Tool: Arxiv paper fetcher (parse arXiv.org abstract directly)
+@tool
+def arxiv_fetch(query_or_id: str) -> str:
+    """Arxiv paper fetcher tool."""
+    try:
+        if re.match(r"\d{4}\.\d{5}(v\d+)?", query_or_id):
+            abs_url = f"https://arxiv.org/abs/{query_or_id}"
+            api_url = f"http://export.arxiv.org/api/query?id_list={query_or_id}"
+            res = requests.get(api_url)
+            if res.status_code == 200:
+                return res.text[:2000] + f"\n\nFull: {abs_url}"
+            return f"Could not retrieve metadata from arXiv API"
+        else:
+            docs = ArxivLoader(query=query_or_id, load_max_docs=2).load()
+            return "\n\n---\n\n".join([doc.page_content for doc in docs])
+    except Exception as e:
+        return f"ArXiv fetch failed: {e}"
+
+@tool
+def math_solver(expression: str) -> str:
+    """Math solver tool."""
+    try:
+        result = sympify(expression).evalf()
+        return str(result)
+    except Exception as e:
+        return f"Math error: {e}"
+
+@tool
+def reverse_text(text: str) -> str:
+    """Text reversal tool."""
+    return text[::-1]
+
+@tool
+def image_info(url: str) -> str:
+    """Image dimension fetcher tool."""
+    try:
+        response = requests.get(url)
+        img = Image.open(BytesIO(response.content))
+        return f"Image size: {img.size} (width x height)"
+    except Exception as e:
+        return f"Image error: {e}"
+
+# Tools list
+tools = [
+    wiki_search,
+    web_search,
+    duckduckgo_search,
+    youtube_transcript,
+    arxiv_fetch,
+    math_solver,
+    reverse_text,
+    image_info
+]
 
-# Build LangGraph
 def build_graph():
     llm = ChatOpenAI(
         model="gpt-4o",
@@ -74,7 +146,6 @@ def build_graph():
     builder.add_edge("tools", "assistant")
     return builder.compile()
 
-# Agent for GAIA benchmark
 class BasicAgent:
     def __init__(self):
         print("GAIA LangGraph Agent Initialized")
@@ -90,3 +161,16 @@ class BasicAgent:
             return final_msg
         except Exception as e:
             return f"FINAL ANSWER: error - {str(e)}"
+
+if __name__ == "__main__":
+    agent = BasicAgent()
+    questions = [
+        "What is the zip code of the Eiffel Tower?",
+        "What is the capital city of Australia?",
+        "How long is the video titled 'The History of Time' on YouTube?",
+        "What does the arXiv paper '2303.12712' say about Transformer performance?",
+    ]
+
+    for q in questions:
+        print(f"\n[Question]: {q}")
+        print(agent(q))