agent.py

# agent.py — agent builder, trace extractor, benchmark runner
import json

from smolagents import CodeAgent, InferenceClientModel

from tools import TOOL_LIST, AUTHORIZED_IMPORTS

# ── Instructions ────────────────────────────────────────────────────────────
INSTRUCTIONS = """
You are a general AI assistant. I will ask you a question.
Think step by step, use tools as needed, then call final_answer() with your answer.

═══════════════════════════════════════════════════════
FINAL ANSWER FORMAT — HIGHEST PRIORITY
═══════════════════════════════════════════════════════
final_answer() replaces the GAIA "FINAL ANSWER: [X]" template.
Pass ONLY the bare answer — no sentences, no explanation, nothing else.

The answer must be one of:
  • A NUMBER
  • As few words as possible (a name, a place, a date, a short phrase)
  • A comma-separated list of numbers and/or strings

─── Numbers ────────────────────────────────────────────
• Do NOT use commas as thousands separator  → 1000 not 1,000
• Do NOT include units ($, %, km…) unless the question explicitly asks
• Do NOT round unless the question says to
• Pay attention to the unit the question asks for:
  "how many thousand hours" → answer 17, not 17000
  "how many millions" → answer 3.2, not 3200000
• WRONG: final_answer("$1,200.50")   RIGHT: final_answer("1200.5")
• WRONG: final_answer("42 meters")   RIGHT: final_answer("42")

─── Strings ────────────────────────────────────────────
• No articles (the, a, an)
  WRONG: final_answer("The Eiffel Tower")   RIGHT: final_answer("Eiffel Tower")
• No abbreviations — write the full form
  WRONG: final_answer("NY")   RIGHT: final_answer("New York")
• Write digits in plain text, not as words
  WRONG: final_answer("forty-two")   RIGHT: final_answer("42")
• Exact spelling as found in the authoritative source

─── Lists ──────────────────────────────────────────────
• Comma-separated, no trailing "and"
• WRONG: final_answer("Alice, Bob, and Carol")   RIGHT: final_answer("Alice, Bob, Carol")

═══════════════════════════════════════════════════════
TOOL STRATEGY
═══════════════════════════════════════════════════════
1. If a file is attached — read it FIRST before anything else.
   PDF   → read_pdf(file_path)
   CSV   → read_csv_file(file_path)
   Excel → read_excel_file(file_path)
   Image → extract_text_from_image(file_path)
   Audio → transcribe_audio(file_path)

2. For current facts, news, prices, live data → search_tool (uses DuckDuckGo, always free).
3. For reading a specific URL in full → fetch_webpage(url).
   If the answer is in a table on that page → extract_table_from_url(url) instead.
   If the URL points directly to a file (PDF/CSV/Excel) → download_and_read(url).
3b. For any YouTube URL in the question → get_youtube_transcript(url) immediately.
4. For math / calculations → calculator or write Python code.
5. For chess positions / best move questions → analyze_chess_position(fen).
   Convert board description to FEN first if needed.
6. For academic papers → arxiv_search first, then fetch the PDF URL and use read_pdf.
   Never guess numbers from academic papers — always read the source.
6. For counting words/chars/patterns in text → count_and_find.
7. For flight distance from Delhi → flight_time_from_delhi.
8. Two reliable sources agree → stop searching and answer.
9. If a tool fails, try a different source — do NOT retry the same broken call.
10. Do not import unauthorized libraries.
11. For Wikipedia discography / filmography / album lists:
    → use wikipedia_section(page_title, "Discography") NOT fetch_webpage.
    The section tool returns plain text with years already readable.
    Count only entries whose year falls in the requested range.
═══════════════════════════════════════════════════════
MEMORY RULES
═══════════════════════════════════════════════════════
• User says "remember / save / store / note that" → call save_memory.
• User refers to "my office / son / broker / trip / project / city" → call retrieve_memory first.
• Never invent personal information.
"""


# ── Agent factory ────────────────────────────────────────────────────────────

def build_agent(hf_token: str) -> CodeAgent:
    model = InferenceClientModel(token=hf_token, model="deepseek-ai/DeepSeek-V3")
    return CodeAgent(
        tools=TOOL_LIST,
        model=model,
        instructions=INSTRUCTIONS,
        max_steps=15,
        verbosity_level=1,
        additional_authorized_imports=AUTHORIZED_IMPORTS,
    )


# ── Trace extractor ──────────────────────────────────────────────────────────

def extract_trace(agent: CodeAgent) -> str:
    """Serialize agent step logs into a readable reasoning trace."""
    try:
        parts = []
        logs = getattr(agent, "logs", None) or getattr(agent, "memory", [])
        for i, step in enumerate(logs):
            chunk = []
            # model output — attribute name varies across smolagents versions
            for attr in ("model_output", "llm_output", "agent_memory"):
                val = getattr(step, attr, None)
                if val:
                    chunk.append(str(val)[:800])
                    break
            # observations
            for attr in ("observations", "observation"):
                obs = getattr(step, attr, None)
                if obs:
                    if isinstance(obs, list):
                        obs = "\n".join(str(o) for o in obs)
                    chunk.append(f"→ {str(obs)[:400]}")
                    break
            # tool calls
            tool_calls = getattr(step, "tool_calls", None)
            if tool_calls:
                chunk.append(f"tools: {str(tool_calls)[:200]}")
            # error
            err = getattr(step, "error", None)
            if err:
                chunk.append(f"✗ {str(err)[:200]}")
            if chunk:
                parts.append(f"[Step {i+1}]\n" + "\n".join(chunk))
        return "\n\n".join(parts)[:5000] if parts else ""
    except Exception:
        return ""


# ── Benchmark runner ─────────────────────────────────────────────────────────

def gaia_benchmark_iter(level_filter: str, split: str, max_q: int, hf_token: str):
    """Synchronous generator — yields (log_text, jsonl_filename_or_None)."""
    from datasets import load_dataset
    from huggingface_hub import hf_hub_download

    yield "Loading gaia-benchmark/GAIA dataset…", None

    try:
        dataset = load_dataset("gaia-benchmark/GAIA", "2023_all", split=split, token=hf_token)
    except Exception as e:
        yield f"Dataset load failed: {e}", None
        return

    if level_filter != "All":
        lvl = level_filter.split()[-1]
        dataset = dataset.filter(lambda x: str(x["Level"]) == lvl)

    total = len(dataset)
    if max_q > 0:
        total = min(max_q, total)
        dataset = dataset.select(range(total))

    log = [f"GAIA — Level: {level_filter} | Split: {split} | Questions: {total}", "─" * 60]
    results = []

    for i, task in enumerate(dataset):
        task_id  = task.get("task_id", f"task_{i}")
        question = task.get("Question", "")
        ground_truth = task.get("Final answer", "")
        level        = task.get("Level", "")

        fname    = (task.get("file_name") or "").strip()

        attached = None
        if fname:
            try:
                attached = hf_hub_download(
                    repo_id="gaia-benchmark/GAIA",
                    filename=f"2023/{split}/{fname}",
                    repo_type="dataset",
                    token=hf_token,
                )
            except Exception as fe:
                log.append(f"  ⚠ file download failed ({fname}): {fe}")

        full_q = question.strip()
        if attached:
            full_q += f"\n\nAttached file path: {attached}"

        agent = build_agent(hf_token)
        try:
            answer = str(agent.run(full_q))
            trace  = extract_trace(agent)
        except Exception as e:
            answer = f"ERROR: {e}"
            trace  = str(e)

        results.append({
    "task_id":         task_id,
    "question":        question,
    "level":           level,
    "model_answer":    answer,
    "ground_truth":    ground_truth,
    "reasoning_trace": trace,
})

        icon = "✓" if not answer.startswith("ERROR") else "✗"
        log.append(f"[{i+1:>3}/{total}] {icon}  {task_id}  →  {answer[:70]}")
        yield "\n".join(log), None

    log += ["─" * 60, f"✅  Done — {total} questions"]
    yield "\n".join(log), results