Update agent.py

agent.py

@@ -1,7 +1,10 @@
+# ===========================
+# FINAL agent.py
+# Tool Tracking Enabled
+# ===========================
+
 import os
 import re
-import time
-import json
 import tempfile
 import subprocess
 from typing import Literal
@@ -23,16 +26,15 @@ from bs4 import BeautifulSoup
 
 load_dotenv()
 
-# ============================================================
+# ==========================================================
 # CONFIG
-# ============================================================
+# ==========================================================
 
-MAX_TOOL_LOOPS = 3
-QUESTION_TIMEOUT_SECONDS = 85
+MAX_TOOL_LOOPS = 2
 
-# ============================================================
+# ==========================================================
 # TOOLS
-# ============================================================
+# ==========================================================
 
 @tool
 def add(a: float, b: float) -> float:
@@ -42,7 +44,7 @@ def add(a: float, b: float) -> float:
 
 @tool
 def subtract(a: float, b: float) -> float:
-    """Subtract b from a."""
+    """Subtract two numbers."""
     return a - b
 
 
@@ -54,96 +56,103 @@ def multiply(a: float, b: float) -> float:
 
 @tool
 def divide(a: float, b: float) -> str:
-    """Divide a by b."""
+    """Divide two numbers."""
     if b == 0:
-        return "division by zero"
+        return "N/A"
     return str(a / b)
 
 
 @tool
 def reverse_text(text: str) -> str:
-    """Reverse text character by character."""
+    """Reverse text."""
     return text[::-1]
 
 
 @tool
 def wiki_search(query: str) -> str:
-    """Search Wikipedia."""
+    """Wikipedia search."""
     try:
         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]}"
+                f"{d.metadata.get('title','')}\n"
+                f"{d.page_content[:1800]}"
             )
-        return "\n\n---\n\n".join(out)
 
-    except Exception as e:
-        return f"wiki error: {e}"
+        return "\n\n".join(out)
+
+    except Exception:
+        return "N/A"
 
 
 @tool
 def web_search(query: str) -> str:
-    """Search web using DuckDuckGo."""
+    """DuckDuckGo search."""
     try:
-        return DuckDuckGoSearchRun(max_results=5).run(query)
-    except Exception as e:
-        return f"web error: {e}"
+        return DuckDuckGoSearchRun(max_results=4).run(query)
+    except Exception:
+        return "N/A"
 
 
 @tool
 def fetch_page(url: str) -> str:
     """Fetch webpage text."""
     try:
-        headers = {"User-Agent": "Mozilla/5.0"}
-        r = requests.get(url, timeout=15, headers=headers)
+        r = requests.get(
+            url,
+            timeout=12,
+            headers={"User-Agent": "Mozilla/5.0"}
+        )
+
         r.raise_for_status()
 
         soup = BeautifulSoup(r.text, "html.parser")
 
-        for t in soup(["script", "style", "nav", "footer", "header"]):
+        for t in soup(["script", "style", "nav", "footer"]):
             t.decompose()
 
-        txt = soup.get_text(separator="\n", strip=True)
-        return txt[:8000]
+        txt = soup.get_text(
+            separator="\n",
+            strip=True
+        )
 
-    except Exception as e:
-        return f"fetch error: {e}"
+        return txt[:7000]
+
+    except Exception:
+        return "N/A"
 
 
 @tool
 def run_python(code: str) -> str:
-    """Execute Python safely."""
+    """Run Python code."""
     try:
         with tempfile.NamedTemporaryFile(
-            mode="w", suffix=".py", delete=False
+            mode="w",
+            suffix=".py",
+            delete=False
         ) as f:
             f.write(code)
             fname = f.name
 
-        result = subprocess.run(
+        p = subprocess.run(
             ["python3", fname],
             capture_output=True,
             text=True,
-            timeout=18
+            timeout=15
         )
 
-        out = result.stdout.strip()
-        err = result.stderr.strip()
-
-        if err:
-            out += f"\nSTDERR:\n{err}"
+        out = p.stdout.strip()
 
-        return out if out else "(no output)"
+        return out if out else "N/A"
 
-    except subprocess.TimeoutExpired:
-        return "python timeout"
-    except Exception as e:
-        return f"python error: {e}"
+    except Exception:
+        return "N/A"
 
 
 TOOLS = [
@@ -158,11 +167,11 @@ TOOLS = [
     run_python,
 ]
 
-# ============================================================
-# MODEL LAYER
-# ============================================================
+# ==========================================================
+# MODELS
+# ==========================================================
 
-def build_model(model_name: str):
+def llm(model_name):
     return ChatGroq(
         model=model_name,
         api_key=os.getenv("GROQ_API_KEY"),
@@ -170,198 +179,189 @@ def build_model(model_name: str):
     )
 
 
-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")
+MODEL_FAST = llm("llama-3.1-8b-instant")
+MODEL_SMART = llm("qwen/qwen3-32b")
 
+try:
+    MODEL_BIG = llm("llama-3.3-70b-versatile")
+except Exception:
+    MODEL_BIG = MODEL_SMART
 
-# ============================================================
+# ==========================================================
 # ROUTER
-# ============================================================
+# ==========================================================
 
 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"
+        "youtube", "video", "audio",
+        "image", "xlsx", "excel",
+        "pdf", "attached file",
+        "chess", "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"
+        "what is", "who is",
+        "surname", "first name",
+        "how many", "reverse",
+        "opposite"
     ]):
         return "simple"
 
     return "medium"
 
 
-def choose_model(question: str):
+def pick_model(question: str):
     level = classify_question(question)
 
     if level == "simple":
-        return MODEL_8B
+        return MODEL_FAST
 
     if level == "medium":
-        return MODEL_QWEN
-
-    # hard
-    try:
-        return MODEL_70B
-    except Exception:
-        return MODEL_QWEN
+        return MODEL_SMART
 
+    return MODEL_BIG
 
-# ============================================================
-# SYSTEM PROMPT
-# ============================================================
+# ==========================================================
+# PROMPT
+# ==========================================================
 
 SYSTEM_PROMPT = """
-You are an exact-answer benchmark agent.
-
-STRICT RULES:
-
-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.
+You solve benchmark tasks.
 
-FINAL OUTPUT MUST END EXACTLY AS:
+RULES:
+1. Exact answer only.
+2. Never explain.
+3. Use tools when needed.
+4. If asked first name / surname, return only requested part.
+5. If list requested, give comma separated list only.
+6. Do not over-answer.
+7. If unsure, search once then answer.
 
-FINAL ANSWER: answer_here
+Always end with:
 
-Rules:
-- no explanation after final answer
-- plain text only
-- list answers comma-separated
+FINAL ANSWER: value
 """
 
+# ==========================================================
+# INVOKER
+# ==========================================================
 
-# ============================================================
-# FALLBACK INVOKER
-# ============================================================
+def invoke(messages, question):
+    models = [
+        pick_model(question),
+        MODEL_SMART,
+        MODEL_FAST
+    ]
 
-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:
+    for m in models:
         key = str(m)
-        if key not in seen:
-            unique.append(m)
-            seen.add(key)
 
-    last_err = None
+        if key in seen:
+            continue
 
-    for model in unique:
-        try:
-            return model.bind_tools(TOOLS).invoke(messages)
+        seen.add(key)
 
-        except Exception as e:
-            last_err = e
-            time.sleep(1)
-
-    raise RuntimeError(f"all models failed: {last_err}")
+        try:
+            return m.bind_tools(TOOLS).invoke(messages)
+        except Exception:
+            continue
 
+    raise RuntimeError("All models failed")
 
-# ============================================================
-# GRAPH STATE
-# ============================================================
+# ==========================================================
+# GRAPH NODE
+# ==========================================================
 
 def assistant(state: MessagesState):
     msgs = state["messages"]
 
-    user_question = ""
+    question = ""
+
     for m in reversed(msgs):
         if hasattr(m, "content"):
-            user_question = str(m.content)
+            question = str(m.content)
             break
 
     messages = [SystemMessage(content=SYSTEM_PROMPT)] + msgs
 
-    result = invoke_with_escalation(messages, user_question)
+    result = invoke(messages, question)
 
     return {"messages": [result]}
 
 
 def build_graph():
-    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")
+    g = StateGraph(MessagesState)
 
-    return graph.compile()
+    g.add_node("assistant", assistant)
+    g.add_node("tools", ToolNode(TOOLS))
 
+    g.add_edge(START, "assistant")
+    g.add_conditional_edges("assistant", tools_condition)
+    g.add_edge("tools", "assistant")
 
-agent = build_graph()
+    return g.compile()
 
-# ============================================================
+# ==========================================================
 # ANSWER EXTRACTION
-# ============================================================
+# ==========================================================
 
 def extract_final_answer(output: dict) -> str:
     msgs = output.get("messages", [])
 
     for m in reversed(msgs):
         content = getattr(m, "content", "")
+
         if isinstance(content, str):
             match = re.search(
                 r"FINAL ANSWER:\s*(.+)",
                 content,
                 re.I | re.S
             )
+
             if match:
-                return match.group(1).strip()
+                ans = match.group(1).strip()
+                ans = ans.replace("\n", " ").strip()
 
-    return "N/A"
+                return ans[:250]
 
+    return "N/A"
 
-# ============================================================
-# PUBLIC SOLVER
-# ============================================================
+# ==========================================================
+# TOOL TRACKING
+# ==========================================================
 
-def solve_question(question: str) -> str:
-    start = time.time()
+def extract_tools_used(output: dict):
+    msgs = output.get("messages", [])
 
-    try:
-        result = agent.invoke(
-            {"messages": [("user", question)]},
-            {"recursion_limit": MAX_TOOL_LOOPS * 2}
-        )
+    tools = []
 
-        answer = extract_final_answer(result)
+    for m in msgs:
+        name = getattr(m, "name", None)
 
-        if time.time() - start > QUESTION_TIMEOUT_SECONDS:
-            return "N/A"
+        if name:
+            tools.append(str(name))
 
-        return answer
+        tc = getattr(m, "tool_calls", None)
 
-    except Exception:
-        return "N/A"
+        if tc:
+            for item in tc:
+                if isinstance(item, dict):
+                    nm = item.get("name")
+                    if nm:
+                        tools.append(nm)
 
+    # deduplicate
+    seen = set()
+    final = []
 
-# ============================================================
-# TEST
-# ============================================================
+    for t in tools:
+        if t not in seen:
+            final.append(t)
+            seen.add(t)
 
-if __name__ == "__main__":
-    q = "Write the opposite of left."
-    print(solve_question(q))
+    return final