Update agent.py

agent.py

@@ -1,11 +1,11 @@
 # ==========================================================
-# UPDATED agent.py
-# Observability + Adaptive Search + Better Recursion Control
+# FINAL agent.py
+# GAIA Optimized | Multi-Provider | Strong Prompt | Tools
 # ==========================================================
 
 import os
 import re
-import time
+import io
 import tempfile
 import subprocess
 from typing import Literal
@@ -15,25 +15,28 @@ from dotenv import load_dotenv
 from langgraph.graph import START, StateGraph, MessagesState
 from langgraph.prebuilt import ToolNode, tools_condition
 
-from langchain_groq import ChatGroq
 from langchain_core.tools import tool
 from langchain_core.messages import SystemMessage
 
+from langchain_groq import ChatGroq
+from langchain_google_genai import ChatGoogleGenerativeAI
+
 from langchain_community.document_loaders import WikipediaLoader
 from langchain_community.tools import DuckDuckGoSearchRun
 
 import requests
+import pandas as pd
 from bs4 import BeautifulSoup
 
 load_dotenv()
 
 # ==========================================================
-# GLOBAL TRACE
+# TRACE VARIABLES
 # ==========================================================
 
 LAST_MODEL_USED = "N/A"
-LAST_MODEL_FALLBACK = "No"
-LAST_MODEL_ERROR = "None"
+LAST_FALLBACK = "No"
+LAST_ERROR = "None"
 
 # ==========================================================
 # TOOLS
@@ -43,21 +46,17 @@ LAST_MODEL_ERROR = "None"
 def wiki_search(query: str) -> str:
     """Search Wikipedia."""
     try:
-        docs = WikipediaLoader(
-            query=query,
-            load_max_docs=2
-        ).load()
+        docs = WikipediaLoader(query=query, load_max_docs=2).load()
 
         if not docs:
             return "No results"
 
         return "\n\n".join(
-            d.page_content[:1800]
-            for d in docs
+            d.page_content[:1800] for d in docs
         )
 
-    except Exception as e:
-        return f"N/A {e}"
+    except Exception:
+        return "N/A"
 
 
 @tool
@@ -67,13 +66,14 @@ def web_search(query: str) -> str:
         return DuckDuckGoSearchRun(
             max_results=4
         ).run(query)
-    except Exception as e:
-        return f"N/A {e}"
+
+    except Exception:
+        return "N/A"
 
 
 @tool
 def fetch_page(url: str) -> str:
-    """Fetch page text."""
+    """Fetch webpage text."""
     try:
         r = requests.get(
             url,
@@ -94,10 +94,12 @@ def fetch_page(url: str) -> str:
         ):
             t.decompose()
 
-        return soup.get_text(
+        text = soup.get_text(
             separator="\n",
             strip=True
-        )[:7000]
+        )
+
+        return text[:7000]
 
     except Exception:
         return "N/A"
@@ -105,7 +107,7 @@ def fetch_page(url: str) -> str:
 
 @tool
 def run_python(code: str) -> str:
-    """Run Python."""
+    """Execute Python code."""
     try:
         with tempfile.NamedTemporaryFile(
             mode="w",
@@ -113,13 +115,13 @@ def run_python(code: str) -> str:
             delete=False
         ) as f:
             f.write(code)
-            fname = f.name
+            path = f.name
 
         p = subprocess.run(
-            ["python3", fname],
+            ["python3", path],
             capture_output=True,
             text=True,
-            timeout=15
+            timeout=20
         )
 
         out = p.stdout.strip()
@@ -130,53 +132,111 @@ def run_python(code: str) -> str:
         return "N/A"
 
 
+@tool
+def add(a: float, b: float) -> str:
+    """Add numbers."""
+    return str(a + b)
+
+
+@tool
+def subtract(a: float, b: float) -> str:
+    """Subtract numbers."""
+    return str(a - b)
+
+
+@tool
+def multiply(a: float, b: float) -> str:
+    """Multiply numbers."""
+    return str(a * b)
+
+
+@tool
+def divide(a: float, b: float) -> str:
+    """Divide numbers."""
+    if b == 0:
+        return "N/A"
+    return str(a / b)
+
+
 @tool
 def reverse_text(text: str) -> str:
     """Reverse text."""
     return text[::-1]
 
 
+@tool
+def read_excel(url: str) -> str:
+    """Read Excel file from URL."""
+    try:
+        r = requests.get(url, timeout=20)
+        r.raise_for_status()
+
+        bio = io.BytesIO(r.content)
+
+        df = pd.read_excel(bio)
+
+        return df.head(25).to_csv(index=False)[:7000]
+
+    except Exception:
+        return "N/A"
+
+
+@tool
+def youtube_transcript(url: str) -> str:
+    """
+    Placeholder transcript tool.
+    Use web_search if transcript unavailable.
+    """
+    return "Transcript unavailable"
+
+
 TOOLS = [
     wiki_search,
     web_search,
     fetch_page,
     run_python,
-    reverse_text
+    add,
+    subtract,
+    multiply,
+    divide,
+    reverse_text,
+    read_excel,
+    youtube_transcript,
 ]
 
 # ==========================================================
 # MODELS
 # ==========================================================
 
-def llm(name):
+def groq(model_name):
     return ChatGroq(
-        model=name,
+        model=model_name,
         api_key=os.getenv("GROQ_API_KEY"),
         temperature=0
     )
 
 
-MODEL_FAST = llm("llama-3.1-8b-instant")
-MODEL_SMART = llm("qwen/qwen3-32b")
-
-try:
-    MODEL_BIG = llm(
-        "llama-3.3-70b-versatile"
+def gemini(model_name):
+    return ChatGoogleGenerativeAI(
+        model=model_name,
+        google_api_key=os.getenv("GOOGLE_API_KEY"),
+        temperature=0
     )
-except Exception:
-    MODEL_BIG = MODEL_SMART
+
+
+MODEL_FAST = groq("llama-3.1-8b-instant")
+MODEL_MED = groq("qwen/qwen3-32b")
+MODEL_HARD = gemini("gemini-3.1-flash-lite")
+MODEL_REASON = gemini("gemma-4-26b")
 
 # ==========================================================
-# ROUTER
+# QUESTION CLASSIFIER
 # ==========================================================
 
 def classify_question(
     question: str
-) -> Literal[
-    "simple",
-    "medium",
-    "hard"
-]:
+) -> Literal["simple", "medium", "hard"]:
+
     q = question.lower()
 
     if any(x in q for x in [
@@ -186,132 +246,160 @@ def classify_question(
         "image",
         "excel",
         "xlsx",
+        "csv",
         "pdf",
+        "attached",
         "chess",
-        "attached"
     ]):
         return "hard"
 
     if any(x in q for x in [
-        "how many",
-        "who is",
         "what is",
+        "who is",
+        "how many",
         "surname",
         "first name",
-        "reverse"
+        "reverse",
     ]):
         return "simple"
 
     return "medium"
 
+# ==========================================================
+# MODEL ROUTING
+# ==========================================================
 
-def preferred_search(question: str):
-    q = question.lower()
-
-    if any(x in q for x in [
-        "paper",
-        "article",
-        "competition",
-        "recipient",
-        "licensed",
-        "2023",
-        "july",
-        "pdf"
-    ]):
-        return "web"
-
-    return "wiki"
-
-
-def pick_model(question: str):
+def choose_models(question: str):
     level = classify_question(question)
 
     if level == "simple":
-        return MODEL_FAST
+        return [
+            ("llama-3.1-8b", MODEL_FAST),
+            ("qwen3-32b", MODEL_MED),
+        ]
 
     if level == "medium":
-        return MODEL_SMART
-
-    return MODEL_BIG
+        return [
+            ("qwen3-32b", MODEL_MED),
+            ("gemma-4-26b", MODEL_REASON),
+            ("gemini-3.1-lite", MODEL_HARD),
+        ]
+
+    return [
+        ("gemini-3.1-lite", MODEL_HARD),
+        ("gemma-4-26b", MODEL_REASON),
+        ("qwen3-32b", MODEL_MED),
+    ]
 
 # ==========================================================
-# PROMPT
+# STRONG SYSTEM PROMPT
 # ==========================================================
 
-def make_prompt(question: str):
-    search_pref = preferred_search(question)
+SYSTEM_PROMPT = """
+You are a precise AI assistant for GAIA benchmark evaluation.
+Your objective is to maximize exact-match score.
+
+AVAILABLE TOOLS:
+wiki_search, web_search, fetch_page, run_python,
+add, subtract, multiply, divide, reverse_text,
+read_excel, youtube_transcript
+
+DECISION RULES:
+
+1. Use tools for factual questions requiring external knowledge:
+names, dates, rankings, counts, locations, webpages, files.
+
+2. Do NOT use tools for trivial math, simple logic,
+direct text transforms, or obvious reasoning.
+
+3. If a URL is present -> use fetch_page first.
+
+4. If code / python file is present -> use run_python.
+
+5. If reversed / encoded text -> decode first.
+
+6. For math -> use math tools or direct arithmetic.
+
+7. If spreadsheet / xlsx / csv appears -> use read_excel.
+
+8. For YouTube/video/audio -> use youtube_transcript first,
+otherwise web_search.
+
+9. Use at most 2 search attempts total.
+Do not loop endlessly.
 
-    return f"""
-You solve benchmark tasks.
+ACCURACY RULES:
 
-Preferred search tool:
-{search_pref}_search
+10. Read qualifiers carefully:
+only, least, highest, lowest, included,
+before, after, between, as of date,
+first alphabetically.
 
-Rules:
-1. Exact answer only.
-2. Never explain.
-3. If asked first name / surname,
-return only requested part.
-4. Use at most 2 tool calls.
-5. If unsure, short best guess.
-6. No extra words.
+11. If asked first name -> first name only.
+If asked surname -> surname only.
 
-Always end with:
+12. If asked for list:
+return only valid items, comma-separated.
 
-FINAL ANSWER: value
+13. Do not over-answer.
+Do not include related but unasked facts.
+
+14. If uncertain, give the most likely concise answer.
+Do not apologize.
+
+OUTPUT RULES:
+
+15. Final output must be exactly one clean line:
+
+FINAL ANSWER: <answer>
+
+16. Use digits for numbers unless explicitly asked otherwise.
+
+17. No quotes, bullets, markdown, commentary, or explanation.
+"""
+
+def build_prompt(question: str):
+    return f"""{SYSTEM_PROMPT}
+
+QUESTION:
+{question}
 """
 
 # ==========================================================
-# INVOKE
+# INVOKER
 # ==========================================================
 
 def invoke(messages, question):
     global LAST_MODEL_USED
-    global LAST_MODEL_FALLBACK
-    global LAST_MODEL_ERROR
-
-    LAST_MODEL_FALLBACK = "No"
-    LAST_MODEL_ERROR = "None"
-
-    models = [
-        pick_model(question),
-        MODEL_SMART,
-        MODEL_FAST
-    ]
+    global LAST_FALLBACK
+    global LAST_ERROR
 
-    used = set()
+    LAST_MODEL_USED = "N/A"
+    LAST_FALLBACK = "No"
+    LAST_ERROR = "None"
 
     first = True
 
-    for m in models:
-        key = str(m)
-
-        if key in used:
-            continue
-
-        used.add(key)
-
+    for name, model in choose_models(question):
         try:
-            LAST_MODEL_USED = key
+            LAST_MODEL_USED = name
 
             if not first:
-                LAST_MODEL_FALLBACK = "Yes"
+                LAST_FALLBACK = "Yes"
 
             first = False
 
-            return m.bind_tools(
+            return model.bind_tools(
                 TOOLS
             ).invoke(messages)
 
         except Exception as e:
-            LAST_MODEL_ERROR = str(e)
+            LAST_ERROR = str(e)
 
-    raise RuntimeError(
-        "All models failed"
-    )
+    raise RuntimeError("All models failed")
 
 # ==========================================================
-# GRAPH
+# GRAPH NODE
 # ==========================================================
 
 def assistant(state: MessagesState):
@@ -324,7 +412,7 @@ def assistant(state: MessagesState):
             question = str(m.content)
             break
 
-    prompt = make_prompt(question)
+    prompt = build_prompt(question)
 
     messages = [
         SystemMessage(content=prompt)
@@ -371,23 +459,14 @@ def build_graph():
     return g.compile()
 
 # ==========================================================
-# EXTRACTION
+# HELPERS
 # ==========================================================
 
-def extract_final_answer(
-    output: dict
-):
-    msgs = output.get(
-        "messages",
-        []
-    )
+def extract_final_answer(output: dict):
+    msgs = output.get("messages", [])
 
     for m in reversed(msgs):
-        txt = getattr(
-            m,
-            "content",
-            ""
-        )
+        txt = getattr(m, "content", "")
 
         if isinstance(txt, str):
             match = re.search(
@@ -397,45 +476,30 @@ def extract_final_answer(
             )
 
             if match:
-                return match.group(
-                    1
-                ).strip()[:250]
+                ans = match.group(1).strip()
+                ans = ans.replace("\n", " ")
+
+                return ans[:250]
 
     return "N/A"
 
 
-def extract_tools_used(
-    output: dict
-):
-    msgs = output.get(
-        "messages",
-        []
-    )
+def extract_tools_used(output: dict):
+    msgs = output.get("messages", [])
 
     tools = []
 
     for m in msgs:
-        nm = getattr(
-            m,
-            "name",
-            None
-        )
+        nm = getattr(m, "name", None)
 
         if nm:
             tools.append(str(nm))
 
-        tc = getattr(
-            m,
-            "tool_calls",
-            None
-        )
+        tc = getattr(m, "tool_calls", None)
 
         if tc:
             for x in tc:
-                if isinstance(
-                    x,
-                    dict
-                ):
+                if isinstance(x, dict):
                     n = x.get("name")
 
                     if n:
@@ -455,6 +519,6 @@ def extract_tools_used(
 def get_last_trace():
     return {
         "model": LAST_MODEL_USED,
-        "fallback": LAST_MODEL_FALLBACK,
-        "model_error": LAST_MODEL_ERROR
+        "fallback": LAST_FALLBACK,
+        "model_error": LAST_ERROR,
     }