Upload agent.py

agent.py

@@ -15,7 +15,7 @@ from langgraph.prebuilt import ToolNode, tools_condition
 
 from langchain_groq import ChatGroq
 from langchain_core.tools import tool
-from langchain_core.messages import SystemMessage, AIMessage
+from langchain_core.messages import SystemMessage
 
 from langchain_community.document_loaders import WikipediaLoader
 from langchain_community.tools import DuckDuckGoSearchRun
@@ -76,22 +76,7 @@ def fetch_page(url: str) -> str:
 
 @tool
 def run_python(code: str) -> str:
-    """Execute Python code and return stdout.
-
-    Use for: arithmetic, counting, sorting, data transformation, and
-    processing Excel/CSV files with pandas.
-
-    pandas (3.x) and openpyxl are available. To read an Excel file use:
-        import pandas as pd
-        df = pd.read_excel('/tmp/<task_id>.xlsx')
-
-    The attached file path is provided in the question as:
-        [ATTACHED FILE: /tmp/<filename>]
-    Always use that exact path when reading files.
-    Print your final answer with print() — only stdout is returned.
-    """
-    if len(code) > 8000:
-        return "Error: code too long (>8000 chars). Rewrite with a shorter, focused script."
+    """Execute Python code and return stdout. Use for calculations, counting, data processing."""
     try:
         with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
             f.write(code)
@@ -124,7 +109,7 @@ def reverse_text(text: str) -> str:
 TOOLS = [wiki_search, web_search, fetch_page, run_python, reverse_text]
 
 # ==========================================================
-# MODELS — two chains based on question type
+# MODELS — primary + ordered fallback chain
 # ==========================================================
 
 def _llm(name: str) -> ChatGroq:
@@ -135,30 +120,13 @@ def _llm(name: str) -> ChatGroq:
     )
 
 
-MODEL_PRIMARY  = _llm("qwen/qwen3-32b")        # strong reasoning, search tasks
-MODEL_FALLBACK = _llm("llama-3.3-70b-versatile") # better at code, numeric, structured
-MODEL_LAST     = _llm("llama-3.1-8b-instant")   # last resort
+# All questions use the same primary model.
+# Fallback chain kicks in only on errors (rate limits, timeouts, etc.)
+MODEL_PRIMARY  = _llm("qwen/qwen3-32b")
+MODEL_FALLBACK = _llm("llama-3.3-70b-versatile")
+MODEL_LAST     = _llm("llama-3.1-8b-instant")
 
-# General questions: Qwen leads, 70b fallback
-GENERAL_CHAIN = [MODEL_PRIMARY, MODEL_FALLBACK, MODEL_LAST]
-
-# Code/numeric questions: still Qwen first to preserve 70b daily token budget (TPD=100k).
-# 70b fires only when Qwen returns N/A — which is when it proved superior (Q10, Q12).
-CODE_CHAIN = [MODEL_PRIMARY, MODEL_FALLBACK, MODEL_LAST]
-
-# Keywords that indicate a code/numeric/structured-data question
-_CODE_SIGNALS = [
-    "python code", "attached python", "numeric output", "final output",
-    "excel", "xlsx", "spreadsheet", "csv", "total sales", "sum of",
-    "how much", "calculate", "computation", "menu items", "sales from",
-    "grocery list", "shopping list", "pie", "filling", "recipe",
-    "attached file", "attached excel",
-]
-
-def _is_code_question(question: str) -> bool:
-    """Return True if the question is best handled by the code-oriented chain (70B first)."""
-    q = question.lower()
-    return any(sig in q for sig in _CODE_SIGNALS)
+FALLBACK_CHAIN = [MODEL_PRIMARY, MODEL_FALLBACK, MODEL_LAST]
 
 # ==========================================================
 # SYSTEM PROMPT — single prompt for all question types
@@ -174,43 +142,16 @@ Produce the exact correct answer — nothing more, nothing less.
 - Use web_search for recent events, specific articles, prices, or anything time-sensitive.
 - Use fetch_page when a URL is provided or a search result points to a relevant page.
 - Use run_python for any arithmetic, counting, sorting, or data transformation.
-- Use run_python with pandas for Excel/CSV questions. When the question contains
-  [ATTACHED FILE DATA], that IS the file content as CSV text. Parse it with:
-    import pandas as pd, io
-    data = '<copy the CSV lines from the question here>'
-    df = pd.read_csv(io.StringIO(data))
-  Then compute and print() the answer.
-- If the question says [ATTACHED FILE CONTENT], that is Python code. Run it with run_python directly
-  — copy the code exactly as given into the run_python tool and print the final output.
-- When given an operation table on a set S with a specific subset question,
-  use run_python to systematically check the required property rather than
-  reasoning about it manually. Manual reasoning on table operations is error-prone.
-- If a question mentions an attached file but the file is not available in the question text
-  (marked as [ATTACHED FILE DATA...] or [ATTACHED FILE CONTENT...]), return N/A.
-  Do NOT invent or hallucinate file contents.
-- NEVER call the same tool with the same query twice. If a tool returns no useful result,
-  try a different tool or a different query — do not repeat.
+- Use reverse_text only when asked to reverse a string.
+- You may use up to 5 tool calls. Stop as soon as you have a confident answer.
 
 ## Answer format rules
-1. Output the raw value only — no explanation, no preamble, no surrounding quotes.
-2. If asked for a first name, output ONLY the first/given name — not the full name, not the surname.
-3. If asked for a surname or last name, output ONLY the family name — not the full name.
-4. Numbers: digits only. Never include currency symbols ($, €, £) even if the question
-   mentions USD — just output the number e.g. 300.00 not $300.00.
-5. Lists: comma-separated with a space after each comma on one line
-   (e.g. "a, b, e" not "a,b,e"). Follow the model's natural formatting.
-6. For subset/set questions (e.g. "subset of S involving..."), output only the elements, comma-separated.
-7. If you cannot find the answer after searching, output: N/A
-8. For discography questions: count studio albums only. Live albums, compilations, box sets do not count.
-9. For botany/classification questions: use strict scientific categories.
-   - Fruits (botanical): tomato, pepper, cucumber, avocado, squash, beans, corn kernels.
-   - Vegetables (botanical): true vegetables are leaves (lettuce, spinach), stems (celery), roots
-     (carrot, sweet potato), bulbs (onion), or flowers (broccoli, cauliflower).
-   - Do NOT confuse culinary and botanical definitions. A tomato is a fruit botanically.
-10. Never abbreviate UNLESS the question explicitly asks for an abbreviation or code:
-    - If asked for an IOC country code: return the 3-letter code (e.g. CUB, GBR, EGY).
-    - If asked for an ISO code, airport code, or similar: return the code as-is.
-    - Otherwise write full words: "United States" not "US", "Saint Petersburg" not "St. Petersburg".
+1. Output the raw value only — no explanation, no preamble.
+2. If asked for a first name, output only the first name.
+3. If asked for a surname, output only the surname.
+4. Numbers: digits only unless units were explicitly requested.
+5. Lists: comma-separated on one line.
+6. If you cannot find the answer after searching, output: N/A
 
 ## Required final line
 Always end your response with exactly:
@@ -218,88 +159,19 @@ FINAL ANSWER: <your answer>
 """
 
 # ==========================================================
-# DIRECT ANSWER PRECHECKS
-# ==========================================================
-
-def maybe_answer_direct(question: str) -> str | None:
-    """
-    Return a direct answer string for questions that don't need tool calls,
-    or None if normal graph processing should continue.
-    """
-    q = question.strip()
-    ql = q.lower()
-
-    # Reversed sentence asking for the opposite of "left"
-    if "etisoppo" in ql and "tfel" in ql:
-        return "right"
-    if 'write the opposite of the word "left"' in ql:
-        return "right"
-
-    # Algebraic subset question — verified correct answer is b, e (with space after comma)
-    # The question involves a 5-element set {a,b,c,d,e} with operation table
-    if 'set s = {a, b, c, d, e}' in ql and 'subset' in ql:
-        return "b, e"
-
-    return None
-
-
-# ==========================================================
-# ANSWER QUALITY CHECK — used by invoke to decide fallback
-# ==========================================================
-
-def _answer_looks_weak(result) -> bool:
-    """
-    Return True if the model's response content does not contain a usable
-    FINAL ANSWER — meaning we should try the next model in the chain.
-    Only applies when the model made NO tool calls (pure text response).
-    If tool calls are present, we let the graph continue normally.
-    """
-    # If the model wants to call tools, don't short-circuit — let the graph run
-    tool_calls = getattr(result, "tool_calls", None)
-    if tool_calls:
-        return False
-
-    content = getattr(result, "content", "") or ""
-    if not isinstance(content, str):
-        return False
-
-    # Check if a FINAL ANSWER line is present and non-empty
-    match = re.search(r"FINAL ANSWER:\s*(.+)", content, re.I | re.S)
-    if not match:
-        return True
-
-    answer = match.group(1).strip()
-    # Treat explicit N/A or empty as weak
-    if not answer or answer.lower() in ("n/a", "none", "unknown", ""):
-        return True
-
-    return False
-
-
-# ==========================================================
-# INVOKE — exception fallback + content-quality fallback
+# INVOKE — with fallback chain
 # ==========================================================
 
-def invoke(messages: list, is_final: bool = False, chain: list = None) -> object:
-    """
-    Call models in chain order with quality-based fallback.
-
-    chain       — ordered list of models to try (GENERAL_CHAIN or CODE_CHAIN)
-    is_final    — True when tool results are already in history (terminal answer turn)
-    """
+def invoke(messages: list) -> object:
     global LAST_MODEL_USED, LAST_MODEL_FALLBACK, LAST_MODEL_ERROR
 
-    if chain is None:
-        chain = GENERAL_CHAIN
-
     LAST_MODEL_FALLBACK = "No"
     LAST_MODEL_ERROR = "None"
 
     seen: set[str] = set()
     first = True
-    last_result = None
 
-    for model in chain:
+    for model in FALLBACK_CHAIN:
         key = model.model_name
         if key in seen:
             continue
@@ -310,52 +182,11 @@ def invoke(messages: list, is_final: bool = False, chain: list = None) -> object
             if not first:
                 LAST_MODEL_FALLBACK = "Yes"
             first = False
-
-            result = model.bind_tools(TOOLS).invoke(messages)
-            last_result = result
-
-            # Model wants to call a tool — return immediately regardless of turn
-            tool_calls = getattr(result, "tool_calls", None)
-            if tool_calls:
-                return result
-
-            # Not a final turn — allow ONE fallback if answer is weak, then stop
-            if not is_final:
-                if _answer_looks_weak(result):
-                    LAST_MODEL_FALLBACK = "Yes"
-                    LAST_MODEL_ERROR = f"weak first-turn answer from {key}"
-                    time.sleep(0.5)
-                    for next_model in chain:
-                        nkey = next_model.model_name
-                        if nkey in seen:
-                            continue
-                        seen.add(nkey)
-                        try:
-                            LAST_MODEL_USED = nkey
-                            r2 = next_model.bind_tools(TOOLS).invoke(messages)
-                            return r2
-                        except Exception as e2:
-                            LAST_MODEL_ERROR = str(e2)
-                            continue
-                return result
-
-            # Final turn — try all remaining models until a non-weak answer
-            if not _answer_looks_weak(result):
-                return result
-
-            LAST_MODEL_FALLBACK = "Yes"
-            LAST_MODEL_ERROR = f"weak answer from {key}"
-            time.sleep(0.5)
-            continue
-
+            return model.bind_tools(TOOLS).invoke(messages)
         except Exception as e:
             LAST_MODEL_ERROR = str(e)
-            time.sleep(0.5)
             continue
 
-    if last_result is not None:
-        return last_result
-
     raise RuntimeError(f"All models failed. Last error: {LAST_MODEL_ERROR}")
 
 # ==========================================================
@@ -363,28 +194,8 @@ def invoke(messages: list, is_final: bool = False, chain: list = None) -> object
 # ==========================================================
 
 def assistant(state: MessagesState) -> dict:
-    user_q = state["messages"][-1].content if state.get("messages") else ""
-
-    # Short-circuit for questions answerable without any tool calls
-    direct = maybe_answer_direct(user_q)
-    if direct is not None:
-        return {"messages": [AIMessage(content=f"FINAL ANSWER: {direct}")]}
-
-    # Pick model chain based on question type:
-    # code/numeric/structured-data → 70B leads (observed better accuracy)
-    # everything else → Qwen leads
-    # Extract the original user question (not injected file content) for routing
-    original_q = user_q.split("\n[ATTACHED FILE")[0]
-    chain = CODE_CHAIN if _is_code_question(original_q) else GENERAL_CHAIN
-
-    # Detect whether this is a final turn (tool results already in history)
-    is_final = any(
-        getattr(m, "type", "") == "tool"
-        for m in state["messages"]
-    )
-
     messages = [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]
-    result = invoke(messages, is_final=is_final, chain=chain)
+    result = invoke(messages)
     return {"messages": [result]}
 
 
@@ -404,25 +215,10 @@ def build_graph():
 def _clean_answer(raw: str) -> str:
     """Normalise the extracted answer string."""
     answer = raw.strip()
-
-    # Strip surrounding quotes the model sometimes wraps answers in
-    # e.g. '"Extremely"' → 'Extremely', '"No, it\'s a sarcophagus."' → stripped later
-    if len(answer) >= 2 and answer[0] in ('"', "'", "\u201c", "\u2018") and answer[-1] in ('"', "'", "\u201d", "\u2019"):
-        answer = answer[1:-1].strip()
-
-    # Strip trailing punctuation
+    # Strip trailing punctuation that the model sometimes adds
     answer = answer.rstrip(".,;:")
-
-    # Strip leading currency symbols — benchmark expects raw numbers, not formatted currency
-    # e.g. "$300.00" → "300.00", "£1,234.56" → "1,234.56"
-    # Exception: if question explicitly asks for USD/currency format, the system prompt
-    # instructs the model accordingly — but _clean_answer always strips symbols here
-    # because the scorer does exact-match and won't accept "$"
-    answer = re.sub(r'^[$€£¥₹]\s*', '', answer)
-
     # Collapse internal whitespace / newlines
     answer = " ".join(answer.split())
-
     # Remove common LLM filler prefixes the regex sometimes captures
     for prefix in (
         "the answer is",
@@ -435,76 +231,6 @@ def _clean_answer(raw: str) -> str:
     ):
         if answer.lower().startswith(prefix):
             answer = answer[len(prefix):].strip()
-
-    # Expand common abbreviations
-    answer = _expand_abbreviation(answer)
-    return answer
-
-
-# Whole-answer abbreviation expansion table.
-# Keys are lowercase stripped versions; values are the canonical full forms.
-_ABBREV_MAP: dict[str, str] = {
-    # Countries / regions
-    "us":    "United States",
-    "usa":   "United States",
-    "u.s.":  "United States",
-    "u.s.a.":"United States",
-    "uk":    "United Kingdom",
-    "u.k.":  "United Kingdom",
-    "uae":   "United Arab Emirates",
-    "ussr":  "Soviet Union",
-    "drc":   "Democratic Republic of the Congo",
-    "dprk":  "North Korea",
-    "rok":   "South Korea",
-    # Cities that frequently get abbreviated
-    "hcmc":  "Ho Chi Minh City",
-    "nyc":   "New York City",
-    "la":    "Los Angeles",
-    "dc":    "Washington, D.C.",
-    # Saint / Sainte prefixes
-    "st.":   "Saint",
-    # Common institution abbreviations
-    "mit":   "Massachusetts Institute of Technology",
-    "un":    "United Nations",
-    "eu":    "European Union",
-    "nato":  "North Atlantic Treaty Organization",
-    "who":   "World Health Organization",
-    "imf":   "International Monetary Fund",
-    "nba":   "National Basketball Association",
-    "nfl":   "National Football League",
-    "fifa":  "Fédération Internationale de Football Association",
-}
-
-
-def _expand_abbreviation(answer: str) -> str:
-    """
-    If the entire answer is a known abbreviation, replace it with the full form.
-    Also expands 'St. <Name>' → 'Saint <Name>' for city/place names.
-
-    Does NOT expand:
-    - IOC country codes (3 uppercase letters like CUB, EGY, GBR) — these are
-      intentional when the question asks for the IOC code.
-    - Answers longer than 5 words — too risky to mutate longer text.
-    """
-    # Only act on short answers (≤ 5 words) to stay safe
-    if len(answer.split()) > 5:
-        return answer
-
-    # Never expand 3-letter ALL-CAPS strings — these are almost certainly
-    # IOC codes, ISO codes, or other intentional abbreviations the question asked for
-    if re.fullmatch(r'[A-Z]{3}', answer.strip()):
-        return answer
-
-    # Whole-answer lookup (case-insensitive)
-    lookup = answer.lower().strip(".")
-    if lookup in _ABBREV_MAP:
-        return _ABBREV_MAP[lookup]
-
-    # Expand "St. <Word>" → "Saint <Word>" (e.g. "St. Petersburg" → "Saint Petersburg")
-    expanded = re.sub(r'\bSt\.\s+', 'Saint ', answer)
-    if expanded != answer:
-        return expanded
-
     return answer
 
 
@@ -544,4 +270,4 @@ def get_last_trace() -> dict:
         "model": LAST_MODEL_USED,
         "fallback": LAST_MODEL_FALLBACK,
         "model_error": LAST_MODEL_ERROR,
-    }
+    }