agent_single_pass.py

"""
agent_single_pass.py

NFL analyst agent (single-pass planning): plan → execute (parallel) → synthesize.

The planner emits every tool call up front, the executor runs them in parallel,
and the synthesizer writes the answer. A single bounded re-plan round lets it
issue follow-up calls that depend on the first round's results.

This is the original agent, preserved alongside the newer ReAct loop in
agent.py. Both expose the same run()/_get_model() interface so the eval harness
can drive either one.
"""

import json
import re
import time
import argparse
import urllib.request
from datetime import date
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed


from tools.registry import TOOLS, DISPATCH

_PROMPTS_DIR = Path(__file__).resolve().parent / "prompts"

# Character budgets for the tool-results block fed back into each LLM call.
# These keep prompt + reserved completion within the server's --max-model-len
# (8192). The planner's system prompt embeds the full tool schemas plus a large
# rules block (~3.4k tokens) and 2048 completion tokens are reserved, so its
# results budget must stay small. Dense JSON tokenizes to more than chars/4, so
# we keep a conservative safety margin rather than filling the nominal space.
_PLAN_RESULTS_BUDGET  = 4_000
_SYNTH_RESULTS_BUDGET = 14_000

def _load_prompt(name: str) -> str:
    return (_PROMPTS_DIR / name).read_text()

# ---------------------------------------------------------------------------
# vLLM client helpers
# ---------------------------------------------------------------------------

def _post(base_url: str, model: str, messages: list[dict], max_tokens: int = 1024,
          stop: list[str] | None = None) -> str:
    """Send a chat completion request and return the assistant message text."""
    body: dict = {
        "model": model,
        "messages": messages,
        "max_tokens": max_tokens,
    }
    if stop:
        body["stop"] = stop
    payload = json.dumps(body).encode()
    req = urllib.request.Request(
        f"{base_url}/v1/chat/completions",
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req) as resp:
        data = json.loads(resp.read())
        usage = data.get("usage", {})

    return data["choices"][0]["message"]["content"], usage


def _get_model(base_url: str) -> str:
    """Auto-detect the first loaded model from the vLLM server."""
    with urllib.request.urlopen(f"{base_url}/v1/models") as resp:
        data = json.loads(resp.read())
    return data["data"][0]["id"]


def _strip_think(text: str) -> str:
    """Remove <think>...</think> chain-of-thought blocks (closed or unclosed)."""
    # Closed blocks
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    # Unclosed block trailing to end-of-string (model ran out of tokens mid-think)
    text = re.sub(r"<think>(?:(?!</think>).)*$", "", text, flags=re.DOTALL)
    return text.strip()


def _trim_results(results: list[dict]) -> list[dict]:
    """
    Shrink oversized tool payloads so prompt + completion fit the model's
    context window.

    An UNFILTERED get_game_results call (team is None) returns every game in the
    season — the whole league overflows an 8k context on its own. Those queries
    are about playoff progression, so we drop REG-season rows. Team-filtered
    results stay intact (their REG games are needed to compute records).
    """
    trimmed = []
    for r in results:
        res = r.get("result")
        if (r.get("tool") == "get_game_results"
                and isinstance(res, dict)
                and res.get("team") is None
                and isinstance(res.get("games"), list)):
            games = [g for g in res["games"] if g.get("game_type") != "REG"]
            trimmed.append({**r, "result": {**res, "games": games}})
        else:
            trimmed.append(r)
    return trimmed


def _format_results(results: list[dict], max_chars: int) -> str:
    """
    Format tool results as a compact context block for the LLM prompt.

    The block is capped at `max_chars` so it can never overflow the model's
    context window, no matter how many tools the planner decided to call. Whole
    results are kept until the budget is hit; the rest are summarized as omitted.
    """
    blocks = [
        f"[{r['tool']}]\n"
        + (json.dumps(r["result"]) if "result" in r else f"ERROR: {r['error']}")
        for r in _trim_results(results)
    ]

    kept: list[str] = []
    used = 0
    for i, block in enumerate(blocks):
        if used + len(block) > max_chars:
            remaining = len(blocks) - i
            kept.append(f"[... {remaining} more tool result(s) omitted to fit context ...]")
            break
        kept.append(block)
        used += len(block) + 1  # +1 for the joining newline

    return "\n".join(kept)


def _repair_json(text: str) -> str:
    """
    Repair structurally unbalanced JSON brackets/braces.

    Small models frequently drop a closing brace before a closing bracket, e.g.
    '[{"args": {"team": "DAL"}]' (missing the '}' for the outer object). This
    walks the text tracking an open-bracket stack (ignoring chars inside
    strings) and inserts the correct closers — both for mismatches mid-string
    and for anything still open at the end.
    """
    closer = {"{": "}", "[": "]"}
    stack: list[str] = []
    out: list[str] = []
    in_string = False
    escaped = False

    for ch in text:
        if in_string:
            out.append(ch)
            if escaped:
                escaped = False
            elif ch == "\\":
                escaped = True
            elif ch == '"':
                in_string = False
            continue

        if ch == '"':
            in_string = True
            out.append(ch)
        elif ch in "{[":
            stack.append(ch)
            out.append(ch)
        elif ch in "}]":
            # Close any open brackets whose closer doesn't match this char yet.
            while stack and closer[stack[-1]] != ch:
                out.append(closer[stack.pop()])
            if stack:
                stack.pop()
            out.append(ch)
        else:
            out.append(ch)

    # Close an unterminated string (truncated mid-value), then anything still open.
    if in_string:
        out.append('"')
    while stack:
        out.append(closer[stack.pop()])

    return "".join(out)


def _parse_tool_calls(raw: str) -> list[dict]:
    """
    Parse the planner's raw output into a list of tool-call dicts.

    Strips chain-of-thought and markdown fences, then normalizes dict-shaped
    outputs (single call object or {"...": [...]} wrapper) into a list.
    Raises ValueError if the output cannot be parsed into a JSON array.
    """
    text = _strip_think(raw)

    # Strip markdown fences if the model wraps output in ```json ... ```
    if text.startswith("```"):
        text = text.split("\n", 1)[1]
        if text.endswith("```"):
            text = text[:-3]
        text = text.strip()

    if not text:
        raise ValueError(f"Planner returned an empty response.\nRaw output: {raw!r}")

    try:
        tool_calls = json.loads(text)
    except json.JSONDecodeError:
        # Small models often drop a closing brace; try a structural repair.
        tool_calls = json.loads(_repair_json(text))

    # Normalize dict outputs to a list. The model sometimes wraps the array in
    # an object (e.g. {"tool_calls": [...]}) or returns a single call object
    # (e.g. {"tool": ..., "args": ...}) instead of a one-element array.
    if isinstance(tool_calls, dict):
        if "tool" in tool_calls:
            tool_calls = [tool_calls]
        else:
            wrapped = next(
                (v for v in tool_calls.values() if isinstance(v, list)), None
            )
            if wrapped is not None:
                tool_calls = wrapped

    if not isinstance(tool_calls, list):
        raise ValueError(
            f"Planner output is not a JSON array: {type(tool_calls).__name__}"
        )

    return tool_calls


# ---------------------------------------------------------------------------
# Phase 1: Plan
# ---------------------------------------------------------------------------

def plan(query: str, base_url: str, model: str,
         prior_results: list[dict] | None = None) -> tuple[list[dict], float]:
    """
    Ask the LLM to produce a list of tool calls for the given query.

    If prior_results is provided, the planner sees earlier tool outputs and
    can issue follow-up calls that depend on them.

    Returns:
        tool_calls: list of {"tool": str, "args": dict}
        elapsed:    planning phase latency in seconds
    """
    system = _load_prompt("planner.txt").format(
        tool_schemas=json.dumps(TOOLS, indent=2),
        current_date=date.today().isoformat(),
    )

    user_content = query
    if prior_results:
        results_block = _format_results(prior_results, _PLAN_RESULTS_BUDGET)
        user_content = (
            f"PRIOR TOOL RESULTS:\n{results_block}\n\n"
            f"Based on these results, output the next round of tool calls "
            f"needed to answer the original question.\n\n"
            f"ORIGINAL QUESTION: {query}"
        )

    messages = [
        {"role": "system", "content": system},
        {"role": "user",   "content": user_content},
    ]

    t0 = time.perf_counter()
    raw, usage = _post(base_url, model, messages, max_tokens=2048)
    try:
        tool_calls = _parse_tool_calls(raw)
    except (ValueError, json.JSONDecodeError) as first_err:
        # Retry once, disabling chain-of-thought and demanding bare JSON. This
        # recovers from truncated/malformed output (e.g. JSON cut off after a
        # long think block).
        retry_messages = [
            {"role": "system", "content": system},
            {"role": "user",
             "content": user_content
             + "\n\nOutput ONLY the JSON array of tool calls. No prose, no "
               "markdown, no thinking. /no_think"},
        ]
        raw, usage = _post(base_url, model, retry_messages, max_tokens=2048)
        try:
            tool_calls = _parse_tool_calls(raw)
        except (ValueError, json.JSONDecodeError):
            raise ValueError(
                f"Planner returned invalid JSON after retry: {first_err}\n"
                f"Raw output: {raw}"
            )
    elapsed = time.perf_counter() - t0

    return tool_calls, elapsed, usage


def needs_replan(results: list[dict]) -> bool:
    """Check whether results contain data that enables follow-up tool calls."""
    for r in results:
        if "result" not in r:
            continue
        data = r["result"]
        # game_results contain week numbers needed for get_game_breakdown
        if r["tool"] == "get_game_results" and isinstance(data, list) and len(data) > 0:
            return True
    return False


# ---------------------------------------------------------------------------
# Phase 2: Execute
# ---------------------------------------------------------------------------

def _call_tool(tool_name: str, args: dict) -> dict:
    """Dispatch a single tool call and return {"tool", "result"} or {"tool", "error"}."""
    fn = DISPATCH.get(tool_name)
    if fn is None:
        return {"tool": tool_name, "error": f"Unknown tool: {tool_name!r}"}
    try:
        result = fn(**args)
        return {"tool": tool_name, "result": result}
    except Exception as e:
        return {"tool": tool_name, "error": str(e)}


def execute(tool_calls: list[dict]) -> tuple[list[dict], dict]:
    """
    Execute tool calls in parallel using a thread pool.

    Returns:
        results:  list of {"tool": str, "result": dict} or {"tool": str, "error": str}
        timings:  {"total": float, "per_tool": {tool_name: float, ...}}
    """
    per_tool_times: dict[str, float] = {}
    results: list[dict] = []

    def _timed_call(tc: dict) -> tuple[dict, str, float]:
        t0 = time.perf_counter()
        result = _call_tool(tc["tool"], tc.get("args", {}))
        elapsed = time.perf_counter() - t0
        return result, tc["tool"], elapsed

    t_total_start = time.perf_counter()
    with ThreadPoolExecutor() as pool:
        futures = [pool.submit(_timed_call, tc) for tc in tool_calls]
        for future in as_completed(futures):
            result, tool_name, elapsed = future.result()
            results.append(result)
            per_tool_times[tool_name] = elapsed
    total_elapsed = time.perf_counter() - t_total_start

    return results, {"total": total_elapsed, "per_tool": per_tool_times}


# ---------------------------------------------------------------------------
# Phase 3: Synthesize
# ---------------------------------------------------------------------------

def synthesize(query: str, results: list[dict], base_url: str, model: str) -> tuple[str, float]:
    """
    Send the query + tool results to the LLM and return a natural language answer.

    Returns:
        answer:  the LLM's response string
        elapsed: synthesis phase latency in seconds
    """
    # Format tool results as context for the prompt
    results_block = _format_results(results, _SYNTH_RESULTS_BUDGET)
    user_content = f"TOOL RESULTS:\n{results_block}\n\nQUESTION: {query}"

    messages = [
        {"role": "system", "content": _load_prompt("synthesizer.txt")},
        {"role": "user",   "content": user_content + " /no_think"},
    ]

    t0 = time.perf_counter()
    answer, usage = _post(base_url, model, messages, max_tokens=2048)
    elapsed = time.perf_counter() - t0

    answer = _strip_think(answer)

    return answer, elapsed, usage


# ---------------------------------------------------------------------------
# Top-level agent runner
# ---------------------------------------------------------------------------

def run(query: str, base_url: str, model: str) -> dict:
    """
    Run the full plan → execute → synthesize pipeline.

    Returns:
        {
            "query":    str,
            "answer":   str,
            "timings":  {
                "plan_s":          float,
                "execute_total_s": float,
                "execute_per_tool": {tool_name: float, ...},
                "synthesize_s":    float,
                "total_s":         float,
            }
        }
    """
    t_start = time.perf_counter()

    tool_calls, plan_time, plan_usage       = plan(query, base_url, model)
    results,    execute_times               = execute(tool_calls)

    # Re-plan if results enable dependent follow-up calls
    if needs_replan(results):
        tool_calls2, plan_time2, plan_usage2 = plan(query, base_url, model, prior_results=results)
        results2, execute_times2 = execute(tool_calls2)
        results += results2
        plan_time += plan_time2
        for k in ("prompt_tokens", "completion_tokens", "total_tokens"):
            if k in plan_usage and k in plan_usage2:
                plan_usage[k] += plan_usage2[k]
        execute_times["total"] += execute_times2["total"]
        execute_times["per_tool"].update(execute_times2["per_tool"])

    answer, synth_time, synth_usage = synthesize(query, results, base_url, model)

    total = time.perf_counter() - t_start

    return {
        "query":  query,
        "answer": answer,
        "timings": {
            "plan_s":           plan_time,
            "plan_prompt_tokens":       plan_usage.get("prompt_tokens"),
            "plan_completion_tokens":   plan_usage.get("completion_tokens"),
            "execute_total_s":  execute_times["total"],
            "execute_per_tool": execute_times["per_tool"],
            "synthesize_s":     synth_time,
            "synth_prompt_tokens":      synth_usage.get("prompt_tokens"),
            "synth_completion_tokens":  synth_usage.get("completion_tokens"),
            "total_s":          total,
        },
    }


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="NFL analyst agent")
    parser.add_argument("--base-url", default="http://localhost:8000")
    parser.add_argument("--model",    default=None)
    parser.add_argument("--query",    required=True)
    args = parser.parse_args()

    model = args.model or _get_model(args.base_url)
    output = run(args.query, args.base_url, model)

    print(output["answer"])
    print("\n--- Timings ---")
    t = output["timings"]
    print(f"  Plan:       {t['plan_s']:.2f}s  "
          f"({t['plan_prompt_tokens']} prompt / {t['plan_completion_tokens']} completion tokens)")
    print(f"  Execute:    {t['execute_total_s']:.2f}s")
    for tool, secs in t["execute_per_tool"].items():
        print(f"    {tool}: {secs:.2f}s")
    print(f"  Synthesize: {t['synthesize_s']:.2f}s  "
          f"({t['synth_prompt_tokens']} prompt / {t['synth_completion_tokens']} completion tokens)")
    print(f"  Total:      {t['total_s']:.2f}s")


if __name__ == "__main__":
    main()