ULTRON V.10.4

# How many contract passages to retrieve per question. Deliberately higher than
# a generic document agent: clause-heavy questions (termination, fees, renewal)
# often span several sections, and current local defaults (8K+ token context)
# leave room for the richer grounding.
RAG_TOP_K = [[[RAG_TOP_K | 6]]]

# Session context set from the page: what kind of contract, and whose side the
# explanation should take. Both optional.
_CONTEXT = {"contract_type": "", "role": "", "sources": []}


def _set_context(contract_type, role):
    """(JS-callable) Remember the contract type and the user's side.

    Underscore-prefixed so it is never exposed to the LLM as a tool.
    """
    _CONTEXT["contract_type"] = (contract_type or "").strip()
    _CONTEXT["role"] = (role or "").strip()
    return "ok"


def _extract_pdf_text(pdf_bytes):
    """(JS-callable) Extract plain text from a PDF's raw bytes via pypdf."""
    import io
    from pypdf import PdfReader

    reader = PdfReader(io.BytesIO(pdf_bytes))
    return "\n\n".join((page.extract_text() or "") for page in reader.pages)


async def _ingest_contract(source, text):
    """(JS-callable) Chunk + embed + index one contract document or amendment.

    Returns the number of passages stored, as a string (for the UI).
    """
    count = await agentop_rag.add_document(source, text)
    _CONTEXT["sources"].append(source)
    return str(count)


def _framing():
    bits = []
    if _CONTEXT["contract_type"]:
        bits.append(f"The document is a {_CONTEXT['contract_type']}.")
    if _CONTEXT["role"]:
        bits.append(
            f"Explain everything from the perspective of the {_CONTEXT['role']}."
        )
    return " ".join(bits)


async def process_user_query(query):
    """Plain-language contract answering, grounded in retrieved clauses.

    Overrides the default router so retrieval is deterministic (small local
    models are unreliable at deciding to call a search tool themselves).
    """
    hits = await agentop_rag.search(query, RAG_TOP_K)
    if not hits:
        return (
            "I don't have a contract to read yet — upload it above (PDF or "
            "text) and I'll explain it in plain language."
        )
    excerpts = "\n\n".join(
        f"[{i + 1}] (from {h['source']}) {h['text']}" for i, h in enumerate(hits)
    )
    grounded_prompt = (
        "You are explaining a contract to someone who is not a lawyer. "
        f"{_framing()} Use ONLY the CONTRACT EXCERPTS below.\n"
        "Answer in plain everyday language: give the short answer first, then "
        "what the contract actually says, citing excerpt numbers like [1]. "
        "If the excerpts mention deadlines, fees, penalties, or automatic "
        "renewal, point them out. If the excerpts do not contain the answer, "
        "say so plainly instead of guessing.\n\n"
        f"CONTRACT EXCERPTS:\n{excerpts}\n\nQUESTION: {query}"
    )
    # TEMPLATE_SYSTEM_PROMPT only exists as a JS global; the query bridge falls
    # back to it when custom_prompt is empty (same idiom as the base template).
    return await process_user_query_wllama(
        grounded_prompt, globals().get("TEMPLATE_SYSTEM_PROMPT", "")
    )