# ============================================================
# 02 — Wire Madhava into a RAG agent
#
# Colab:  !pip install winnex-madhava
#
# WHAT THIS SHOWS
#   The pattern every RAG pipeline needs and almost none implement:
#   a retrieval step that can answer "did you miss anything?"
#
#   A standard RAG agent retrieves k chunks and hands them to the LLM.
#   If a relevant chunk was dropped, nobody knows. Here the retriever
#   returns the chunks AND the proof that the rest could not belong.
#
# WHY IT MATTERS
#   When a regulated customer asks "why didn't the assistant cite the
#   addendum?", "the vector index is approximate" is not an answer.
#   "Here is the bound proving the addendum could not be in the top-k,
#   and here is the one that could" is.
# ============================================================

# instalação automática — Colab, Jupyter e script local.
# No Colab, "!pip" dentro de comentário não executa: a instalação
# precisa acontecer de dentro da célula, antes dos imports.
import importlib
import subprocess
import sys

for _pkg in ["winnex-madhava"]:
    _mod = _pkg.replace("-", "_")
    try:
        __import__(_mod)
        continue
    except ImportError:
        pass
    _r = subprocess.run(
        [sys.executable, "-m", "pip", "install", "-q", "--upgrade", _pkg],
        capture_output=True, text=True
    )
    # O interpretador do Colab já está em execução quando o pip termina.
    # Sem invalidar o cache, o import abaixo pode não enxergar o pacote.
    importlib.invalidate_caches()
    try:
        __import__(_mod)
    except ImportError:
        sys.stderr.write(_r.stdout or "")
        sys.stderr.write(_r.stderr or "")
        raise SystemExit(
            "\n\nNão foi possível importar " + _pkg + " após instalar.\n"
            "No Colab: menu 'Ambiente de execução' -> 'Reiniciar sessão',\n"
            "e execute esta célula novamente.\n"
        )

import numpy as np
import winnex_madhava as mad

DIM = 128

def embed(text: str) -> np.ndarray:
    """Deterministic demo embedding. Swap for a real encoder in production."""
    v = np.zeros(DIM, dtype=np.float32)
    for ch in text.lower():
        v[ord(ch) % DIM] += 1.0
    n = np.linalg.norm(v)
    return v / n if n else v

# ── 1. a document corpus ────────────────────────────────────

chunks = [
    "O contrato tem prazo de 24 meses, renovável por igual período.",
    "O aluguel mensal é de R$ 2.000,00, corrigido anualmente pelo IGP-M.",
    "O aditivo primeiro altera o aluguel para R$ 2.400,00 a partir de janeiro.",
    "A multa por rescisão antecipada é de três aluguéis.",
    "O foro de eleição é a comarca de Goiânia.",
    "A garantia locatícia é fiança bancária de R$ 7.200,00.",
]
corpus = np.stack([embed(c) for c in chunks]).astype(np.float32)

# ── 2. index with the audit path enabled ────────────────────
retriever = mad.build_engine(
    corpus,
    dim=DIM,
    metric="cosine",
    k=3,
    k1_fraction=1.0,        # keep every bound survivor
    postfilter=True,        # exact re-rank
    audit_exhaustive=True,  # produce the per-document record
)

# ── 3. the retrieval tool an agent would call ───────────────
def retrieve_with_proof(question: str, k: int = 3) -> dict:
    """Return chunks plus evidence that the rest were excluded correctly."""
    result = mad.retrieve(retriever, embed(question), corpus_text=chunks, k=k)

    q = embed(question)
    passagens = [
        {
            "texto": chunks[idx],
            "score": float(np.dot(q, corpus[idx])),
        }
        for idx in result["indices"]
    ]

    return {
        "passagens": passagens,
        "proof": {
            "bound_violations": result["bound_violations"],
            "excluded_proven": len(chunks) - len(result["indices"]),
            "claim": result["proof"],
            "latency_ms": result["latency_ms"],
        },
    }

# ── 4. what the agent does with it ──────────────────────────
pergunta = "qual o valor atual do aluguel"
saida = retrieve_with_proof(pergunta)

print(f"pergunta: {pergunta}\n")
print("contexto entregue ao modelo:")
for p in saida["passagens"]:
    print(f"  [{p['score']:.3f}] {p['texto']}")

print(f"\nevidence for the audit log:")
print(f"  documents excluded with proof : {saida['proof']['excluded_proven']}")
print(f"  bound violations              : {saida['proof']['bound_violations']}")
print(f"  latency                       : {saida['proof']['latency_ms']:.2f} ms")

assert saida["proof"]["bound_violations"] == 0

# ── 5. the system prompt that makes the agent honest ────────
SYSTEM_PROMPT = """You answer questions using only the retrieved passages.

If the answer is not in the passages, say so plainly — do not guess.

The retrieval step is deterministic and carries a proof: every document
that was NOT returned has been shown, by bound, to be irrelevant to this
query. If a user asks why a specific document was not cited, the system
can produce that proof. You may tell them this."""

print(f"\nsystem prompt in use:\n{SYSTEM_PROMPT}")
print("\nOK — retrieval is wired, and it can answer the audit question.")