webapp/agent_loop.py

"""ReAct agent loop over the uploaded-packet tools in :mod:`webapp.agent_tools`.

The model drives: each step it emits ONE JSON action — ``{"thought", "tool",
"args"}`` — we run the tool, append the observation to a scratchpad, and ask
again, until it emits the terminal ``final_answer`` tool (or a budget runs out).

Why a JSON-action ReAct loop rather than the OpenAI ``tools=`` function-calling
API: the shipped runtime is a local GGUF (Gemma) via llama.cpp, whose native
tool-calling is unreliable. A plain "emit one JSON object" protocol is robust on
local models and works identically against the remote vLLM endpoint we validate
against first — so the loop only needs a ``complete(prompt, system) -> str``
callable and stays backend-agnostic.

The loop is a generator of *frames* (``{"stage", ...}``) so the caller
(:func:`webapp.backend.agent_chat`) can stream the agent's thinking, tool calls,
tool results, and final answer to the React UI live.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Callable, Iterator

from webapp.agent_tools import (
    FINAL_ANSWER,
    ToolContext,
    observation_text,
    run_tool,
    tool_catalog,
)

# Turn budgets. Steps cap tool calls (and, on ZeroGPU, GPU-window seconds); the
# parse budget tolerates a few malformed JSON replies before we give up.
DEFAULT_MAX_STEPS = 6
_PARSE_RETRIES = 2
# How much prior conversation to carry (chars) — keeps the prompt inside an 8K ctx.
_HISTORY_CHARS = 4000


def _system_prompt(ctx: ToolContext) -> str:
    catalog = json.dumps(tool_catalog(), ensure_ascii=False, indent=2)
    return (
        "You are the Agenda Parser's research agent. You help the user understand a "
        "single uploaded public-meeting agenda packet (one compiled PDF: the agenda "
        "plus every item's backup documents).\n\n"
        "You work by calling tools, one at a time, and reasoning over what they return. "
        "Start with list_agenda_items to see the agenda and which items have backup "
        "pages.\n\n"
        "Available tools (JSON Schema):\n" + catalog + "\n\n"
        "PROTOCOL — follow exactly:\n"
        "- Respond with a SINGLE JSON object and nothing else (no prose, no code fence).\n"
        '- Shape: {"thought": "<one short sentence>", "tool": "<tool name>", "args": {<args>}}\n'
        "- Call exactly one tool per step. Read its result, then decide the next step.\n"
        "- Example step:\n"
        '  {"thought": "Find where the budget total is stated.", "tool": "find_text", '
        '"args": {"query": "$"}}\n'
        f'- When you can answer, call "{FINAL_ANSWER}" with an "answer" in markdown.\n\n'
        "CHOOSING A TOOL:\n"
        "- find_text — exact strings: dollar amounts, dates, names, acronyms, ordinance/"
        "statute numbers, \"Item 7\". Fast and literal.\n"
        "- search_packet — conceptual lookups (\"where is X discussed\"); ranks by meaning, "
        "so it can miss exact tokens that find_text would catch.\n"
        "- get_item_text — read the backup pages for a specific item; if has_more is true, "
        "call again with offset=next_offset to read the rest.\n"
        "- summarize / report — heavy: each is a multi-pass LLM run over the whole packet "
        "(slow). Use ONLY for whole-agenda summaries or thorough briefings, never for a "
        "single fact.\n\n"
        "RULES:\n"
        "- Cite the item number and page range you drew each fact from. Do not invent "
        "facts you did not see via a tool.\n"
        "- If list_agenda_items reports confidence \"poor\", the page ranges came from text, "
        "not bookmarks — verify with get_item_text or find_text before trusting them.\n"
        "- If a tool errors, change your args or switch tools; never repeat the identical "
        "failing call. If something fails twice, answer with what you have.\n"
        "- You have a limited number of steps. Be efficient and call "
        f'"{FINAL_ANSWER}" as soon as the question is answerable.'
    )


def _render_history(messages: list[dict]) -> str:
    """The prior conversation as a compact transcript (most recent kept)."""
    lines = []
    for m in messages:
        role = "User" if m.get("role") == "user" else "Assistant"
        content = (m.get("content") or "").strip()
        if content:
            lines.append(f"{role}: {content}")
    text = "\n".join(lines)
    if len(text) > _HISTORY_CHARS:
        text = "…" + text[-_HISTORY_CHARS:]
    return text or "(no prior messages)"


def _extract_json(text: str) -> dict | None:
    """Best-effort parse of a single JSON object from a model reply.

    Tolerates code fences and surrounding prose by falling back to the first
    balanced ``{...}`` span. Returns ``None`` if nothing parses.
    """
    s = (text or "").strip()
    if s.startswith("```"):
        s = s.split("```", 2)[1] if s.count("```") >= 2 else s.strip("`")
        if s.lstrip().lower().startswith("json"):
            s = s.lstrip()[4:]
    s = s.strip()
    try:
        obj = json.loads(s)
        return obj if isinstance(obj, dict) else None
    except (json.JSONDecodeError, ValueError):
        pass
    # Fall back: scan for the first balanced brace span.
    start = s.find("{")
    while start != -1:
        depth = 0
        for i in range(start, len(s)):
            if s[i] == "{":
                depth += 1
            elif s[i] == "}":
                depth -= 1
                if depth == 0:
                    try:
                        obj = json.loads(s[start : i + 1])
                        if isinstance(obj, dict):
                            return obj
                    except (json.JSONDecodeError, ValueError):
                        break
        start = s.find("{", start + 1)
    return None


def _result_summary(tool: str, result: dict) -> str:
    """A one-line, human-readable gist of a tool result for the live UI (the raw JSON
    stays available behind a disclosure). Best-effort — empty string when nothing fits."""
    if not isinstance(result, dict):
        return ""
    if result.get("error"):
        return f"error: {str(result['error'])[:80]}"
    if tool == "find_text" and isinstance(result.get("count"), int):
        n = result["count"]
        return f"{n} match{'' if n == 1 else 'es'}"
    if isinstance(result.get("count"), int):
        n = result["count"]
        noun = "item" if tool == "list_agenda_items" else "result"
        s = f"{n} {noun}{'' if n == 1 else 's'}"
        conf = result.get("confidence")
        return f"{s} · {conf}" if conf else s
    if isinstance(result.get("results"), list):
        n = len(result["results"])
        return f"{n} passage{'' if n == 1 else 's'}"
    if tool == "get_item_text":
        if result.get("sliced") and result.get("pages"):
            return f"pp. {result['pages']}"
        return str(result.get("method") or "full packet")
    if result.get("summary"):
        return "summary ready"
    if result.get("report"):
        return "report ready"
    if isinstance(result.get("items"), list):
        n = len(result["items"])
        return f"{n} item{'' if n == 1 else 's'}"
    return ""


@dataclass(frozen=True)
class Toolkit:
    """The agenda-specific pieces the otherwise-generic loop plugs in.

    Bundling these lets a second agent (e.g. the Cornell LII legal-research agent) reuse
    the whole ReAct loop with a different toolset by passing its own ``Toolkit`` — only the
    system prompt, tool dispatch, and result-summary differ; the protocol, JSON parsing,
    scratchpad, retries, and budget logic are shared.

    - ``system_prompt(ctx) -> str``: the agent's role + tool catalog + PROTOCOL.
    - ``run_tool(name, args, ctx) -> dict``: dispatch one tool call (errors as ``{"error"}``).
    - ``observation_text(result) -> str``: compact JSON of a result for the scratchpad.
    - ``result_summary(tool, result) -> str``: one-line human gist for the live UI.
    - ``final_answer``: the terminal pseudo-tool name that ends the turn.
    """

    system_prompt: Callable[..., str]
    run_tool: Callable[..., dict]
    observation_text: Callable[..., str]
    result_summary: Callable[..., str]
    final_answer: str


# The default toolkit — the uploaded-packet research agent (unchanged behavior).
AGENDA_TOOLKIT = Toolkit(
    system_prompt=_system_prompt,
    run_tool=run_tool,
    observation_text=observation_text,
    result_summary=_result_summary,
    final_answer=FINAL_ANSWER,
)


def agent_turn(
    messages: list[dict],
    ctx: object,
    complete: Callable[..., str],
    *,
    max_steps: int = DEFAULT_MAX_STEPS,
    toolkit: Toolkit | None = None,
) -> Iterator[dict]:
    """Run one user turn of the ReAct loop, yielding progress/result frames.

    Args:
        messages: the full chat so far (list of ``{"role", "content"}``), ending with
            the user's current message.
        ctx: the tool context the toolkit's tools operate on (packet ``ToolContext`` for
            the agenda agent, ``LiiContext`` for the LII agent) — passed straight through
            to ``toolkit.run_tool``.
        complete: ``complete(prompt, system=...) -> str`` the agent reasons with.
        max_steps: maximum tool calls before the loop force-stops.
        toolkit: which toolset/system-prompt to drive the loop with; defaults to the
            agenda packet agent (:data:`AGENDA_TOOLKIT`).

    Frames (``{"stage", ...}``):
        ``thinking`` (text, step) · ``tool_call`` (tool, args, step) ·
        ``tool_result`` (tool, result, step) · ``answer`` (text) · ``error`` (text).
    """
    tk = toolkit or AGENDA_TOOLKIT
    system = tk.system_prompt(ctx)
    history = _render_history(messages)
    scratchpad: list[str] = []
    parse_fails = 0
    seen_actions: set[str] = set()  # canonical (tool, args) already run this turn

    for step in range(1, max_steps + 1):
        pad = "\n".join(scratchpad) if scratchpad else "(empty — no tool calls yet)"
        prompt = (
            f"Conversation so far:\n{history}\n\n"
            f"Your scratchpad this turn (tool calls and their results):\n{pad}\n\n"
            "Respond with the next JSON action."
        )
        try:
            raw = complete(prompt, system=system)
        except Exception as e:  # noqa: BLE001
            yield {"stage": "error", "text": f"Model call failed: {type(e).__name__}: {e}"}
            return

        action = _extract_json(raw)
        if action is None or "tool" not in action:
            parse_fails += 1
            if parse_fails > _PARSE_RETRIES:
                # Out of retries — surface whatever the model said as the answer.
                yield {"stage": "answer", "text": (raw or "").strip() or
                       "I couldn't formulate a structured answer."}
                return
            # Tell the UI we're nudging the model back to valid JSON, so the work area
            # isn't silent during a retry.
            yield {"stage": "notice", "text": "Re-reading the request…", "step": step}
            scratchpad.append(
                'SYSTEM: Your last reply was not a single valid JSON action. '
                'Reply with ONLY {"thought":..., "tool":..., "args":{...}}.'
            )
            continue

        thought = str(action.get("thought") or "").strip()
        tool = str(action.get("tool") or "").strip()
        targs = action.get("args")
        if not isinstance(targs, dict):
            targs = {}

        if tool == tk.final_answer:
            if thought:
                yield {"stage": "thinking", "text": thought, "step": step}
            answer = str(targs.get("answer") or action.get("answer") or "").strip()
            yield {"stage": "answer", "text": answer or "(no answer)"}
            return

        # Skip an identical repeated tool call — a known small-model failure mode the
        # PROTOCOL warns against (it re-emits the same action despite the scratchpad). Don't
        # re-run it or stream a duplicate step to the UI; nudge the model to use the prior
        # result, vary the call, or finish, and move on to the next step.
        action_key = json.dumps({"tool": tool, "args": targs}, sort_keys=True, ensure_ascii=False)
        if action_key in seen_actions:
            scratchpad.append(
                f"SYSTEM: You already ran {tool} with those exact arguments this turn; its "
                "result is in the scratchpad above. Do NOT repeat it — use that result, try "
                f'different arguments or another tool, or call "{tk.final_answer}".'
            )
            continue
        seen_actions.add(action_key)

        if thought:
            yield {"stage": "thinking", "text": thought, "step": step}
        yield {"stage": "tool_call", "tool": tool, "args": targs, "step": step}
        result = tk.run_tool(tool, targs, ctx)
        obs = tk.observation_text(result)
        yield {
            "stage": "tool_result", "tool": tool, "args": targs, "result": obs,
            "summary": tk.result_summary(tool, result), "step": step,
        }
        scratchpad.append(
            f'ACTION: {json.dumps({"tool": tool, "args": targs}, ensure_ascii=False)}\n'
            f"OBSERVATION: {obs}"
        )

    # Budget exhausted: ask for a final answer from what we've gathered.
    pad = "\n".join(scratchpad) if scratchpad else "(no tool results)"
    closing = (
        f"Conversation so far:\n{history}\n\n"
        f"Your scratchpad this turn:\n{pad}\n\n"
        "You have reached the step limit. Write your best final answer for the user now, "
        "in markdown, using only what the tools returned above."
    )
    try:
        final = complete(closing, system=system)
    except Exception as e:  # noqa: BLE001
        yield {"stage": "error", "text": f"Model call failed: {type(e).__name__}: {e}"}
        return
    # The closing prompt asks for prose, but tolerate a stray JSON final_answer.
    parsed = _extract_json(final)
    if parsed and parsed.get("tool") == FINAL_ANSWER:
        final = str((parsed.get("args") or {}).get("answer") or final)
    yield {"stage": "answer", "text": (final or "").strip() or "(no answer)"}