Update agent.py

agent.py

@@ -1,353 +1,367 @@
 import os
 import re
 import time
-import subprocess
+import json
 import tempfile
+import subprocess
+from typing import Literal
+
 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_community.document_loaders import WikipediaLoader
 from langchain_community.tools import DuckDuckGoSearchRun
-from langchain_core.messages import SystemMessage, HumanMessage
+
 import requests
 from bs4 import BeautifulSoup
 
 load_dotenv()
 
 # ============================================================
-# QUESTION DIFFICULTY CLASSIFIER
+# CONFIG
 # ============================================================
 
-# Keywords that indicate a SIMPLE, fast question
-SIMPLE_KEYWORDS = [
-    "what is", "who is", "when was", "where is", "how many",
-    "reverse", "opposite", "calculate", "what does",
-    "reversed text", "what year", "first name", "surname",
-    "capital of", "which country", "what color",
-]
-
-# Keywords that indicate a HARD / time-consuming question
-HARD_KEYWORDS = [
-    "youtube.com", "watch?v=", "video", "audio", "mp3", "wav",
-    "image", "excel", "xlsx", "pdf", "attached", "file",
-    "chess", "position", "board",
-    "wikipedia featured article", "libretexts", "research paper",
-    "nominated", "who awarded",
-]
-
-# Max seconds allowed per question before we skip it
-QUESTION_TIMEOUT_SECONDS = 90  # 1.5 minutes max per question
-
-
-def classify_difficulty(question: str) -> str:
-    """Returns 'simple', 'medium', or 'hard'."""
-    q = question.lower()
-
-    # Hard signals
-    hard_score = sum(1 for kw in HARD_KEYWORDS if kw in q)
-    if hard_score >= 1:
-        return "hard"
-
-    # Simple signals
-    simple_score = sum(1 for kw in SIMPLE_KEYWORDS if kw in q)
-    if simple_score >= 1 or len(question) < 120:
-        return "simple"
-
-    return "medium"
-
-
-def sort_questions_by_difficulty(questions: list) -> list:
-    """Sort questions: simple first, medium next, hard last."""
-    order = {"simple": 0, "medium": 1, "hard": 2}
-
-    def score(item):
-        q = item.get("question", "")
-        return order[classify_difficulty(q)]
-
-    return sorted(questions, key=score)
-
+MAX_TOOL_LOOPS = 3
+QUESTION_TIMEOUT_SECONDS = 85
 
 # ============================================================
-# 1. TOOLS
+# TOOLS
 # ============================================================
 
 @tool
-def add(a: int, b: int) -> int:
+def add(a: float, b: float) -> float:
     """Add two numbers."""
     return a + b
 
 
 @tool
-def multiply(a: int, b: int) -> int:
-    """Multiply two numbers."""
-    return a * b
+def subtract(a: float, b: float) -> float:
+    """Subtract b from a."""
+    return a - b
 
 
 @tool
-def subtract(a: int, b: int) -> int:
-    """Subtract b from a."""
-    return a - b
+def multiply(a: float, b: float) -> float:
+    """Multiply two numbers."""
+    return a * b
 
 
 @tool
-def divide(a: float, b: float) -> float:
+def divide(a: float, b: float) -> str:
     """Divide a by b."""
     if b == 0:
-        return "Error: Division by zero."
-    return a / b
+        return "division by zero"
+    return str(a / b)
 
 
 @tool
 def reverse_text(text: str) -> str:
-    """Reverse a string character by character."""
+    """Reverse text character by character."""
     return text[::-1]
 
 
 @tool
 def wiki_search(query: str) -> str:
-    """Search Wikipedia for factual information."""
+    """Search Wikipedia."""
     try:
-        search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-        if not search_docs:
-            return "No Wikipedia results found."
-        formatted = []
-        for doc in search_docs:
-            content = doc.page_content[:2000]
-            formatted.append(f'Source: {doc.metadata.get("source", "Unknown")}\n{content}')
-        return "\n\n---\n\n".join(formatted)
+        docs = WikipediaLoader(query=query, load_max_docs=2).load()
+        if not docs:
+            return "No results"
+
+        out = []
+        for d in docs:
+            out.append(
+                f"TITLE: {d.metadata.get('title','')}\n"
+                f"{d.page_content[:2200]}"
+            )
+        return "\n\n---\n\n".join(out)
+
     except Exception as e:
-        return f"Wiki search error: {e}"
+        return f"wiki error: {e}"
 
 
 @tool
 def web_search(query: str) -> str:
-    """Search the web using DuckDuckGo."""
+    """Search web using DuckDuckGo."""
     try:
         return DuckDuckGoSearchRun(max_results=5).run(query)
     except Exception as e:
-        return f"Web search error: {e}"
+        return f"web error: {e}"
 
 
 @tool
 def fetch_page(url: str) -> str:
-    """Fetch and read the full content of a webpage."""
+    """Fetch webpage text."""
     try:
-        headers = {"User-Agent": "Mozilla/5.0 (compatible; research-agent/1.0)"}
-        response = requests.get(url, headers=headers, timeout=15)
-        response.raise_for_status()
-        soup = BeautifulSoup(response.text, "html.parser")
-        for tag in soup(["script", "style", "nav", "footer", "header"]):
-            tag.decompose()
-        text = soup.get_text(separator="\n", strip=True)
-        if len(text) > 8000:
-            text = text[:8000] + "\n... [truncated]"
-        return text
+        headers = {"User-Agent": "Mozilla/5.0"}
+        r = requests.get(url, timeout=15, headers=headers)
+        r.raise_for_status()
+
+        soup = BeautifulSoup(r.text, "html.parser")
+
+        for t in soup(["script", "style", "nav", "footer", "header"]):
+            t.decompose()
+
+        txt = soup.get_text(separator="\n", strip=True)
+        return txt[:8000]
+
     except Exception as e:
-        return f"Error fetching page: {e}"
+        return f"fetch error: {e}"
 
 
 @tool
 def run_python(code: str) -> str:
-    """Execute Python code and return output."""
+    """Execute Python safely."""
     try:
-        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
+        with tempfile.NamedTemporaryFile(
+            mode="w", suffix=".py", delete=False
+        ) as f:
             f.write(code)
             fname = f.name
+
         result = subprocess.run(
-            ['python3', fname],
+            ["python3", fname],
             capture_output=True,
             text=True,
-            timeout=20
+            timeout=18
         )
-        output = result.stdout.strip()
-        if result.stderr.strip():
-            output += f"\nSTDERR: {result.stderr.strip()}"
-        return output or "(no output)"
+
+        out = result.stdout.strip()
+        err = result.stderr.strip()
+
+        if err:
+            out += f"\nSTDERR:\n{err}"
+
+        return out if out else "(no output)"
+
     except subprocess.TimeoutExpired:
-        return "Code execution timed out."
+        return "python timeout"
     except Exception as e:
-        return f"Error running code: {e}"
+        return f"python error: {e}"
 
 
-tools = [
-    add, subtract, multiply, divide,
+TOOLS = [
+    add,
+    subtract,
+    multiply,
+    divide,
     reverse_text,
-    wiki_search, web_search, fetch_page,
+    wiki_search,
+    web_search,
+    fetch_page,
     run_python,
 ]
 
 # ============================================================
-# 2. MULTI-MODEL FALLBACK
-# Updated: Qwen3-32B replaces non-working Gemma4
-# Order: Qwen3 (best reasoning) → Llama 3.3 70B → Llama 3.1 8B
+# MODEL LAYER
 # ============================================================
 
-def get_fallback_models():
-    """
-    Updated model stack (April 2026):
-    - qwen/qwen3-32b: Best reasoning + agentic capabilities on Groq
-    - llama-3.3-70b-versatile: Strong general fallback
-    - llama-3.1-8b-instant: Fast last-resort
-    """
-    models = []
-
-    if os.getenv("GROQ_API_KEY"):
-        groq_key = os.getenv("GROQ_API_KEY")
-
-        # Llama 3.1 8B - fast last-resort
-        models.append(("Llama 3.1 8B", ChatGroq(
-            model="llama-3.1-8b-instant",
-            api_key=groq_key,
-            temperature=0
-        )))
-        
-        
-        # Qwen3 32B - best reasoning, replaces non-working Gemma4
-        models.append(("Qwen3 32B", ChatGroq(
-            model="qwen/qwen3-32b",
-            api_key=groq_key,
-            temperature=0,
-        )))
-
-        # Llama 3.3 70B - strong general-purpose fallback
-        models.append(("Llama 3.3 70B", ChatGroq(
-            model="llama-3.3-70b-versatile",
-            api_key=groq_key,
-            temperature=0
-        )))
-
-        
-
-    if not models:
-        raise RuntimeError("No API keys found! Set GROQ_API_KEY.")
-
-    print(f"✅ Loaded {len(models)} model(s): {[n for n, _ in models]}")
-    return models
-
-
-FALLBACK_MODELS = get_fallback_models()
-
-
-def invoke_with_fallback(messages):
-    """Try each model in order with retries on rate limit and tool failures."""
-    last_error = None
-
-    for name, m in FALLBACK_MODELS:
-        for attempt in range(2):  # Max 2 attempts per model
-            try:
-                print(f"🤖 [{name}] Attempt {attempt + 1}...")
-                result = m.bind_tools(tools).invoke(messages)
-                print(f"✅ [{name}] Success.")
-                return result
-
-            except Exception as e:
-                err = str(e).lower()
-                last_error = e
-
-                is_rate_limit = any(x in err for x in [
-                    "rate_limit", "429", "quota",
-                    "too many", "resource exhausted", "token"
-                ])
-                is_tool_failed = any(x in err for x in [
-                    "tool_use_failed", "failed_generation",
-                    "tool call validation failed"
-                ])
-
-                if "decommissioned" in err or "invalid_argument" in err:
-                    print(f"❌ [{name}] Invalid/decommissioned. Skipping.")
-                    break
-
-                elif is_rate_limit:
-                    wait = (attempt + 1) * 8
-                    print(f"⚠️ [{name}] Rate limited. Waiting {wait}s...")
-                    time.sleep(wait)
-                    last_error = e
-                    break  # move to next model
-
-                elif is_tool_failed:
-                    wait = (attempt + 1) * 3
-                    print(f"⚠️ [{name}] Tool call failed. Retrying in {wait}s...")
-                    time.sleep(wait)
-                    last_error = e
-                    continue  # retry same model
-
-                else:
-                    print(f"❌ [{name}] Error: {e}")
-                    break  # move to next model
-
-    raise RuntimeError(f"All models exhausted. Last error: {last_error}")
+def build_model(model_name: str):
+    return ChatGroq(
+        model=model_name,
+        api_key=os.getenv("GROQ_API_KEY"),
+        temperature=0
+    )
+
+
+MODEL_8B = build_model("llama-3.1-8b-instant")
+MODEL_QWEN = build_model("qwen/qwen3-32b")
+
+# optional
+MODEL_70B = build_model("llama-3.3-70b-versatile")
 
 
 # ============================================================
-# 3. ASSISTANT NODE
+# ROUTER
 # ============================================================
 
-SYSTEM_PROMPT = """You are a precise AI assistant for the GAIA benchmark evaluation.
+def classify_question(question: str) -> Literal["simple", "medium", "hard"]:
+    q = question.lower()
+
+    if any(x in q for x in [
+        "youtube.com", "watch?v=", "video", "audio",
+        "image", "xlsx", "excel", "pdf",
+        "attached file", "chess", "board position"
+    ]):
+        return "hard"
+
+    if any(x in q for x in [
+        "calculate", "reverse", "opposite",
+        "who is", "what is", "where is",
+        "how many", "first name", "surname"
+    ]):
+        return "simple"
+
+    return "medium"
+
 
-TOOLS AVAILABLE: wiki_search, web_search, fetch_page, run_python,
-                 add, subtract, multiply, divide, reverse_text.
+def choose_model(question: str):
+    level = classify_question(question)
+
+    if level == "simple":
+        return MODEL_8B
+
+    if level == "medium":
+        return MODEL_QWEN
+
+    # hard
+    try:
+        return MODEL_70B
+    except Exception:
+        return MODEL_QWEN
+
+
+# ============================================================
+# SYSTEM PROMPT
+# ============================================================
+
+SYSTEM_PROMPT = """
+You are an exact-answer benchmark agent.
 
 STRICT RULES:
-1. ALWAYS use tools for factual questions. Never rely on memory alone.
-2. Search before answering any question about names, dates, numbers, or facts.
-3. If wiki_search is unclear → follow up with web_search.
-4. If a question references a URL → use fetch_page immediately.
-5. If a question includes or references a Python file or code → use run_python to execute it.
-6. For reversed/encoded text → use reverse_text first, then answer what it asks.
-7. For math → always use the math tools (add, subtract, multiply, divide).
-8. For YouTube or audio/video questions → search for transcripts with web_search.
-9. Never give up — try multiple search strategies with different keywords if needed.
-10. Be concise in your reasoning. Do not over-think simple questions.
-
-OUTPUT FORMAT — you MUST end EVERY response with exactly this line:
-FINAL ANSWER: [your answer here]
-
-Rules for the answer content:
-- Numbers: no commas, no units unless specifically asked (write 3 not $3 or 3,000)
-- Strings: no articles (a/an/the), no abbreviations, write out digits in plain text
-- Lists: comma-separated, apply above rules per element
-- Write ONLY the bare answer after "FINAL ANSWER:" — no extra explanation
+
+1. Use tools whenever facts, names, dates, counts, URLs, files or code are involved.
+2. If URL appears -> use fetch_page first.
+3. If code appears -> use run_python.
+4. If text seems reversed -> use reverse_text.
+5. For factual lookup: wiki_search first, then web_search if needed.
+6. Be concise.
+7. Never hallucinate.
+8. If unsure, search again once.
+
+FINAL OUTPUT MUST END EXACTLY AS:
+
+FINAL ANSWER: answer_here
+
+Rules:
+- no explanation after final answer
+- plain text only
+- list answers comma-separated
 """
 
 
-def assistant(state: MessagesState):
-    sys_msg = SystemMessage(content=SYSTEM_PROMPT)
-    messages = [sys_msg] + state["messages"]
-    return {"messages": [invoke_with_fallback(messages)]}
+# ============================================================
+# FALLBACK INVOKER
+# ============================================================
+
+def invoke_with_escalation(messages, question: str):
+    primary = choose_model(question)
+
+    order = [primary, MODEL_QWEN, MODEL_8B]
+
+    # keep unique order
+    unique = []
+    seen = set()
+
+    for m in order:
+        key = str(m)
+        if key not in seen:
+            unique.append(m)
+            seen.add(key)
+
+    last_err = None
+
+    for model in unique:
+        try:
+            return model.bind_tools(TOOLS).invoke(messages)
+
+        except Exception as e:
+            last_err = e
+            time.sleep(1)
+
+    raise RuntimeError(f"all models failed: {last_err}")
 
 
 # ============================================================
-# 4. BUILD GRAPH
+# GRAPH STATE
 # ============================================================
 
+def assistant(state: MessagesState):
+    msgs = state["messages"]
+
+    user_question = ""
+    for m in reversed(msgs):
+        if hasattr(m, "content"):
+            user_question = str(m.content)
+            break
+
+    messages = [SystemMessage(content=SYSTEM_PROMPT)] + msgs
+
+    result = invoke_with_escalation(messages, user_question)
+
+    return {"messages": [result]}
+
+
 def build_graph():
-    builder = StateGraph(MessagesState)
-    builder.add_node("assistant", assistant)
-    builder.add_node("tools", ToolNode(tools))
-    builder.add_edge(START, "assistant")
-    builder.add_conditional_edges("assistant", tools_condition)
-    builder.add_edge("tools", "assistant")
-    return builder.compile()
+    graph = StateGraph(MessagesState)
+
+    graph.add_node("assistant", assistant)
+    graph.add_node("tools", ToolNode(TOOLS))
+
+    graph.add_edge(START, "assistant")
+    graph.add_conditional_edges("assistant", tools_condition)
+    graph.add_edge("tools", "assistant")
+
+    return graph.compile()
 
 
+agent = build_graph()
+
 # ============================================================
-# 5. ANSWER EXTRACTION
+# ANSWER EXTRACTION
 # ============================================================
 
-def extract_final_answer(agent_output: dict) -> str:
-    """Extract FINAL ANSWER from agent output."""
-    messages = agent_output.get("messages", [])
-    if not messages:
-        return "N/A"
+def extract_final_answer(output: dict) -> str:
+    msgs = output.get("messages", [])
 
-    for msg in reversed(messages):
-        content = getattr(msg, "content", "")
+    for m in reversed(msgs):
+        content = getattr(m, "content", "")
         if isinstance(content, str):
-            match = re.search(r"FINAL ANSWER:\s*(.+)", content, re.IGNORECASE | re.DOTALL)
+            match = re.search(
+                r"FINAL ANSWER:\s*(.+)",
+                content,
+                re.I | re.S
+            )
             if match:
-                answer = match.group(1).strip()
-                return answer if answer else "N/A"
+                return match.group(1).strip()
+
+    return "N/A"
+
+
+# ============================================================
+# PUBLIC SOLVER
+# ============================================================
+
+def solve_question(question: str) -> str:
+    start = time.time()
+
+    try:
+        result = agent.invoke(
+            {"messages": [("user", question)]},
+            {"recursion_limit": MAX_TOOL_LOOPS * 2}
+        )
+
+        answer = extract_final_answer(result)
+
+        if time.time() - start > QUESTION_TIMEOUT_SECONDS:
+            return "N/A"
+
+        return answer
+
+    except Exception:
+        return "N/A"
+
+
+# ============================================================
+# TEST
+# ============================================================
 
-    return "N/A"
+if __name__ == "__main__":
+    q = "Write the opposite of left."
+    print(solve_question(q))