webapp/agent_tools.py

"""Tool registry for Agent Mode — uploaded-packet access as callable tools.

Agent Mode lets the model *drive* the agenda packet itself: instead of one fixed
pipeline, it picks a tool, sees the result, and decides what to do next (see
:mod:`webapp.agent_loop`). This module is the tool catalog — each tool is a thin
wrapper over an existing :mod:`webapp.backend` / :mod:`chroma` function, plus a
JSON-Schema describing its arguments so the loop can advertise it to the model and
validate what comes back.

Design notes:

* **One uploaded packet.** Every tool runs against the single uploaded agenda packet
  (``ctx.upload_id``). The model never picks a source; it explores *within* the one
  packet (list items → read an item / search → summarize / report).
* **LLM-backed tools reuse the caller's completer.** ``summarize`` / ``report`` run
  the chroma map-reduce with ``ctx.complete`` — the *same* in-window completer the
  agent loop reasons with. On ZeroGPU that means the already-loaded GGUF, not a
  nested ``@spaces.GPU`` call (which would try to grab a second GPU window).
* **Bounded outputs.** Every result is capped (text truncated, lists trimmed) so a
  tool observation can't blow the model's context window — doubly important for the
  8K-context local GGUF.
* **PDF safety.** All packet reads go through ``backend``'s cached, PDFium-lock-
  serialized helpers; this module never touches pypdfium2 directly.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Callable

# Output caps — keep any single observation small enough to feed back into an 8K ctx.
_MAX_ITEMS = 60          # agenda items returned to the model
_MAX_TEXT_CHARS = 2600   # item text / summary / report body slice (one page of reading)
_MAX_CHUNK_CHARS = 650   # per semantic-search hit
_MAX_HITS = 8            # exact-match hits returned by find_text
_FIND_CONTEXT = 120      # chars of context shown on each side of an exact match
_REPORT_MAX_CHUNKS = 12  # cap the report map-reduce hard in agent mode
_SUMMARY_CHARS = 16000   # agenda text fed to the summarizer (pre-chunked anyway)


@dataclass
class ToolContext:
    """Everything a tool needs that the *model* shouldn't have to supply.

    ``complete`` is the LLM completer the agent loop is already using (remote client
    or in-window llama.cpp); LLM-backed tools call it so nothing re-enters the GPU.
    ``_store`` memoizes this turn's ephemeral packet vector store so repeated
    ``search_packet`` calls don't re-chunk the same packet.
    """

    upload_id: str
    complete: Callable[..., str] | None = None
    _store: object | None = None


@dataclass
class Tool:
    name: str
    description: str
    parameters: dict  # JSON-Schema for the args object
    run: Callable[..., dict]


# --------------------------------------------------------------------------- #
# Small helpers
# --------------------------------------------------------------------------- #
def _clip(text: str, limit: int) -> str:
    text = (text or "").strip()
    if len(text) <= limit:
        return text
    return text[:limit].rstrip() + " …[truncated]"


def _agenda_items(ctx: ToolContext) -> list[dict]:
    """The parsed agenda items for the uploaded packet (bookmarks-first parse)."""
    from webapp import backend

    return backend.parse_agenda_outline_from_packet(ctx.upload_id).get("items", [])


def _item_titles(items: list[dict]) -> list[str]:
    """Agenda items as ordered ``"<number> <name>"`` lines — the form the packet
    slicer anchors on (mirrors what the React UI sends)."""
    return [f"{(it.get('number') or '').strip()} {(it.get('name') or '').strip()}".strip()
            for it in items]


def _agenda_portion_text(ctx: ToolContext) -> str:
    """Text of the page range the user marked as the agenda (the TOC portion)."""
    from webapp import backend

    pages = backend.cached_packet_pages(ctx.upload_id)
    if not pages:
        return ""
    a_start, a_end = backend._parse_page_range(
        backend._agenda_pages_for(ctx.upload_id), len(pages))
    return "\n\n".join(p for p in pages[a_start:a_end] if p).strip()


# --------------------------------------------------------------------------- #
# Tool implementations
# --------------------------------------------------------------------------- #
def _list_agenda_items(ctx: ToolContext, **args) -> dict:
    """The parsed agenda items, in order, each with its packet page range."""
    from webapp import backend

    parsed = backend.parse_agenda_outline_from_packet(ctx.upload_id)
    items = parsed.get("items", [])
    # How the outline was obtained: "outline" (PDF bookmarks, reliable), "text"
    # (parsed from the agenda page text — verify ranges), or "none". Surfaced so the
    # agent can hedge / double-check with find_text when confidence is "poor".
    source = parsed.get("source", "none")
    confidence = parsed.get("confidence", "empty")
    if not items:
        return {
            "count": 0,
            "source": source,
            "confidence": confidence,
            "note": "No agenda items were parsed from this packet — it may have no "
            "bookmarks and an unparseable agenda page range.",
            "items": [],
        }
    trimmed = [
        {
            "index": i,
            "number": it.get("number", ""),
            "name": it.get("name", ""),
            "is_section": it.get("is_section", False),
            "has_pages": it.get("has_pages", False),
            "pages": it.get("pages", ""),
        }
        for i, it in enumerate(items[:_MAX_ITEMS])
    ]
    out = {"count": len(items), "source": source, "confidence": confidence,
           "items": trimmed}
    if confidence == "poor":
        out["note"] = ("Page ranges were parsed from text, not bookmarks — verify them "
                       "with get_item_text / find_text before relying on them.")
    return out


def _get_item_text(ctx: ToolContext, **args) -> dict:
    """Raw packet text backing one agenda item, sliced to its section."""
    from webapp import backend

    items = _agenda_items(ctx)
    if not items:
        return {"error": "No agenda items to anchor on. Call list_agenda_items first."}
    titles = _item_titles(items)

    index = args.get("item_index")
    if index is None and args.get("item"):
        needle = str(args["item"]).strip().lower()
        index = next((i for i, t in enumerate(titles) if needle in t.lower()), None)
    if index is None:
        return {
            "error": "Specify item_index (0-based) or item (title text). "
            "Available items: " + "; ".join(f"{i}:{t}" for i, t in enumerate(titles[:40])),
        }
    index = int(index)
    if not (0 <= index < len(items)):
        return {"error": f"item_index out of range (0..{len(items) - 1})."}

    it = items[index]
    try:
        pages = backend.cached_packet_pages(ctx.upload_id)
    except Exception as e:  # noqa: BLE001
        return {"error": f"Could not load the packet: {type(e).__name__}: {e}"}
    if not pages:
        return {"error": "The packet is no longer on the server — ask the user to "
                "re-upload it."}

    start, end = int(it.get("start") or 0), int(it.get("end") or 0)
    if it.get("has_pages") and end > start:
        text = "\n\n".join(p for p in pages[start:end] if p).strip()
        # Long items are returned a page-of-text at a time; the model continues by
        # calling again with `offset = next_offset` until has_more is false.
        try:
            offset = max(0, int(args.get("offset") or 0))
        except (TypeError, ValueError):
            offset = 0
        total = len(text)
        window = text[offset : offset + _MAX_TEXT_CHARS]
        next_offset = offset + len(window)
        has_more = next_offset < total
        return {
            "item": titles[index],
            "sliced": True,
            "pages": it.get("pages", ""),
            "text": window + (" …[more — call again with this offset]" if has_more else ""),
            "offset": offset,
            "next_offset": next_offset,
            "total_chars": total,
            "has_more": has_more,
        }
    return {
        "item": titles[index],
        "sliced": False,
        "note": "This item has no backup pages in the packet. Use search_packet to "
        "find relevant content instead.",
        "text": "",
    }


def _ephemeral_store(ctx: ToolContext):
    """A per-turn vector store of the uploaded packet's full text."""
    from chroma import AgendaStore
    from webapp import backend

    if ctx._store is not None:
        return ctx._store
    docs = backend._packet_documents(ctx.upload_id)
    # Chroma collection names must be 3-512 chars of [a-zA-Z0-9._-] and start/end
    # alphanumeric; the url-safe upload_id can contain (or end in) "_"/"-".
    import re

    safe = re.sub(r"[^a-zA-Z0-9]", "", ctx.upload_id) or "packet"
    store = AgendaStore(None, collection=f"agent_{safe}")
    if docs:
        store.add_documents(docs)
    ctx._store = store
    return store


def _search_packet(ctx: ToolContext, **args) -> dict:
    query = str(args.get("query") or "").strip()
    if not query:
        return {"error": "Provide a 'query' string to search the packet for."}
    n = max(1, min(int(args.get("n_results") or 5), 10))
    try:
        store = _ephemeral_store(ctx)
        hits = store.query(query, n_results=n)
    except Exception as e:  # noqa: BLE001
        return {"error": f"Packet search failed: {type(e).__name__}: {e}"}
    results = [
        {
            "file": (h.get("metadata") or {}).get("file_name", ""),
            "text": _clip(h.get("text", ""), _MAX_CHUNK_CHARS),
        }
        for h in hits
    ]
    return {"query": query, "results": results}


def _find_text(ctx: ToolContext, **args) -> dict:
    """Exact (case-insensitive) substring scan over the packet's per-page text.

    The deterministic complement to semantic ``search_packet`` — finds literal strings
    (dollar amounts, dates, acronyms, "Item 7") that embeddings can miss, and reports the
    1-indexed page each match lands on.
    """
    from webapp import backend

    query = str(args.get("query") or "").strip()
    if not query:
        return {"error": "Provide a 'query' string to find in the packet."}
    try:
        max_hits = max(1, min(int(args.get("max_hits") or _MAX_HITS), _MAX_HITS))
    except (TypeError, ValueError):
        max_hits = _MAX_HITS
    try:
        pages = backend.cached_packet_pages(ctx.upload_id)
    except Exception as e:  # noqa: BLE001
        return {"error": f"Could not load the packet: {type(e).__name__}: {e}"}
    if not pages:
        return {"error": "The packet is no longer on the server — ask the user to "
                "re-upload it."}

    needle = query.lower()
    hits: list[dict] = []
    total = 0
    for pno, page in enumerate(pages, start=1):
        text = page or ""
        low = text.lower()
        pos = low.find(needle)
        while pos != -1:
            total += 1
            if len(hits) < max_hits:
                a = max(0, pos - _FIND_CONTEXT)
                b = min(len(text), pos + len(query) + _FIND_CONTEXT)
                snippet = ("…" if a > 0 else "") + text[a:b].strip() + ("…" if b < len(text) else "")
                hits.append({"page": pno, "snippet": _clip(snippet, _MAX_CHUNK_CHARS)})
            pos = low.find(needle, pos + len(needle))
    return {
        "query": query,
        "count": total,
        "hits": hits,
        "note": "" if total <= max_hits else f"Showing first {max_hits} of {total} matches.",
    }


def _summarize(ctx: ToolContext, **args) -> dict:
    from chroma import summarize_text

    if ctx.complete is None:
        return {"error": "No LLM completer available for summarize."}
    try:
        text = _agenda_portion_text(ctx)
        if not text:
            return {"error": "No extractable agenda text (likely a scanned PDF, or the "
                    "packet is no longer on the server)."}
        summary = summarize_text(text[:_SUMMARY_CHARS], complete=ctx.complete)
    except Exception as e:  # noqa: BLE001
        return {"error": f"Summarize failed: {type(e).__name__}: {e}"}
    return {"summary": _clip(summary, _MAX_TEXT_CHARS)}


def _report(ctx: ToolContext, **args) -> dict:
    from chroma import generate_report
    from webapp import backend

    question = str(args.get("question") or "").strip()
    if not question:
        return {"error": "Provide a 'question' to frame the report."}
    if ctx.complete is None:
        return {"error": "No LLM completer available for report."}
    try:
        docs = backend._packet_documents(ctx.upload_id)
        if not docs:
            return {"error": "No packet text available (re-upload the packet)."}
        report = generate_report(
            docs, question, complete=ctx.complete, max_chunks=_REPORT_MAX_CHUNKS
        )
    except Exception as e:  # noqa: BLE001
        return {"error": f"Report failed: {type(e).__name__}: {e}"}
    return {"question": question, "report": _clip(report, _MAX_TEXT_CHARS)}


# --------------------------------------------------------------------------- #
# Registry
# --------------------------------------------------------------------------- #
TOOLS: list[Tool] = [
    Tool(
        name="list_agenda_items",
        description="List the uploaded agenda's items, in order — each with an index, "
        "number, name, whether it has backup pages, and its packet page range. Start "
        "here to orient yourself.",
        parameters={"type": "object", "properties": {}},
        run=_list_agenda_items,
    ),
    Tool(
        name="get_item_text",
        description="The raw packet text backing ONE agenda item, sliced to its backup "
        "pages. Identify the item by item_index (from list_agenda_items) or item (title "
        "text). Long items return one page of text at a time; if has_more is true, call "
        "again with offset=next_offset to read the rest.",
        parameters={
            "type": "object",
            "properties": {
                "item_index": {"type": "integer", "description": "0-based index from list_agenda_items."},
                "item": {"type": "string", "description": "Title text to match instead of an index."},
                "offset": {"type": "integer", "description": "Char offset to resume reading a long item (use next_offset from the prior call; default 0)."},
            },
        },
        run=_get_item_text,
    ),
    Tool(
        name="search_packet",
        description="Semantic search over the whole uploaded agenda packet. Returns the "
        "most relevant passages for a query — use for 'find/where/does the packet mention X'.",
        parameters={
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "What to look for."},
                "n_results": {"type": "integer", "description": "How many passages (default 5)."},
            },
            "required": ["query"],
        },
        run=_search_packet,
    ),
    Tool(
        name="find_text",
        description="Exact, case-insensitive text search over the packet — returns the "
        "page number and surrounding snippet for each literal match. Use this for exact "
        "strings: dollar amounts, dates, names, acronyms, statute/ordinance numbers, "
        "'Item 7'. Use search_packet instead for conceptual 'where is X discussed' queries.",
        parameters={
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "The exact text to find."},
                "max_hits": {"type": "integer", "description": "Max matches to return (default 8)."},
            },
            "required": ["query"],
        },
        run=_find_text,
    ),
    Tool(
        name="summarize",
        description="Map-reduce summary of the agenda (its table-of-contents portion). "
        "Heavier (an LLM pass) — use when the user wants the whole agenda summarized, not "
        "for a single fact.",
        parameters={"type": "object", "properties": {}},
        run=_summarize,
    ),
    Tool(
        name="report",
        description="Query-framed report mined from the whole agenda packet (heavy "
        "map-reduce). Use for thorough briefings; prefer search_packet for quick lookups.",
        parameters={
            "type": "object",
            "properties": {
                "question": {"type": "string", "description": "The question framing the report."},
            },
            "required": ["question"],
        },
        run=_report,
    ),
]

_BY_NAME = {t.name: t for t in TOOLS}

# final_answer is special-cased by the loop (it ends the turn), but advertised here
# so the model sees it in the catalog alongside the real tools.
FINAL_ANSWER = "final_answer"


def tool_catalog() -> list[dict]:
    """OpenAI-style ``{name, description, parameters}`` schema for every tool, plus
    the terminal ``final_answer`` — what the agent loop shows the model."""
    cat = [
        {"name": t.name, "description": t.description, "parameters": t.parameters}
        for t in TOOLS
    ]
    cat.append(
        {
            "name": FINAL_ANSWER,
            "description": "Give the user your final answer and end the turn. Use this as "
            "soon as you have enough to answer.",
            "parameters": {
                "type": "object",
                "properties": {
                    "answer": {"type": "string", "description": "The answer, in markdown."}
                },
                "required": ["answer"],
            },
        }
    )
    return cat


def run_tool(name: str, args: dict, ctx: ToolContext) -> dict:
    """Dispatch one tool call. Returns a JSON-serializable result dict (always — a
    failed/unknown call comes back as ``{"error": ...}`` so the model can recover)."""
    tool = _BY_NAME.get(name)
    if tool is None:
        return {"error": f"Unknown tool '{name}'. Available: {', '.join(_BY_NAME)}."}
    if not isinstance(args, dict):
        return {"error": "Tool args must be a JSON object."}
    try:
        return tool.run(ctx, **args)
    except TypeError as e:
        return {"error": f"Bad arguments for '{name}': {e}"}
    except Exception as e:  # noqa: BLE001
        return {"error": f"Tool '{name}' failed: {type(e).__name__}: {e}"}


def observation_text(result: dict, limit: int = 2800) -> str:
    """Compact JSON rendering of a tool result for the model's scratchpad."""
    try:
        text = json.dumps(result, ensure_ascii=False, separators=(",", ":"))
    except (TypeError, ValueError):
        text = str(result)
    return _clip(text, limit)