rag.py

"""
The chatbot's brain: embed the question, retrieve a broad set of relevant
chunks (hybrid vector + keyword search) from Supabase, and ask the LLM to
write a genuinely generative answer grounded in that material — not a
templated or copy-pasted response — citing which document each part came
from.
"""
import logging
import re
from typing import List, Optional

import requests

import db
from config import GROQ_API_KEY, GROQ_CHAT_MODEL, TOP_K_CHUNKS, MAX_CONTEXT_CHARS
from embeddings import embed_query

logger = logging.getLogger("janmat.rag")

GROQ_URL = "https://api.groq.com/openai/v1/chat/completions"

SYSTEM_PROMPT = """You are JanMat, an assistant that answers questions about Indian \
legislative bills, their PRS Legislative Research briefs, related news articles, and \
public opinion submissions.

Rules:
- Answer using the numbered SOURCE excerpts provided below. Do not invent facts that \
aren't supported by them.
- Write a genuinely thought-out, generative answer in your own words — synthesise \
across the sources, explain context and implications, and structure the answer to \
actually help the reader understand the topic. Do not just copy or lightly rephrase a \
single source's sentence.
- If the sources give a fuller picture together than any one of them alone, weave that \
together rather than picking one source and ignoring the rest.
- After each claim, cite the source number(s) it came from, like [1] or [2,3].
- Format for readability: a short opening paragraph that answers the question \
directly, then bullet points or short paragraphs for the details. Use **bold** \
sparingly for key terms. Avoid wide multi-column tables — they're hard to read on \
a phone; prefer bullets. Don't add a heading above a single short paragraph.
- Explain the substance in your own words. Short quoted phrases are fine where \
the exact legal wording matters, but don't string together long verbatim passages \
from the source text — the point is to make the material understandable, not to \
reproduce it.
- Only cite a source if you actually used it for that claim. Don't add citations \
to look thorough.
- If the sources don't contain enough information to answer, say so plainly instead of \
guessing.
- Keep the answer in plain, accessible language — assume the reader is a citizen, not \
a lawyer or policy expert.
- If sources disagree (e.g. a bill's official text vs. a news article's framing vs. a \
public comment), point that out rather than picking one silently.
"""


def _format_context(chunks: List[dict]) -> str:
    parts = []
    total = 0
    for i, c in enumerate(chunks, start=1):
        header = f"[{i}] {c['title']} ({c['kind']}, {c['url']}" + (
            f", page {c['page_number']})" if c.get("page_number") else ")"
        )
        block = f"{header}\n{c['content']}\n"
        if total + len(block) > MAX_CONTEXT_CHARS:
            break
        parts.append(block)
        total += len(block)
    return "\n---\n".join(parts)


def retrieve(question: str, top_k: int = TOP_K_CHUNKS, filter_kind: Optional[str] = None) -> List[dict]:
    """Hybrid retrieval: embeds the question for vector search, and passes the
    raw text too so Postgres can blend in trigram keyword matching — catches
    exact bill names / section numbers that a pure embedding search can miss,
    while still ranking mostly on semantic similarity."""
    query_vec = embed_query(question)
    return db.match_chunks(query_vec, query_text=question, match_count=top_k, filter_kind=filter_kind)


def _call_groq(messages: list, temperature: float = 0.3) -> str:
    if not GROQ_API_KEY:
        raise RuntimeError("GROQ_API_KEY is not set")
    headers = {"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"}
    payload = {"model": GROQ_CHAT_MODEL, "temperature": temperature, "messages": messages}
    resp = requests.post(GROQ_URL, headers=headers, json=payload, timeout=60)
    resp.raise_for_status()
    return resp.json()["choices"][0]["message"]["content"]


def _extract_cited_numbers(answer: str, max_n: int) -> List[int]:
    """
    Find which source numbers the model actually cited, e.g. [1], [2,3],
    [4, 5]. Returns them sorted and de-duplicated, ignoring any number
    outside the range of sources we actually supplied.
    """
    cited = set()
    for group in re.findall(r"\[([\d\s,]+)\]", answer):
        for part in group.split(","):
            part = part.strip()
            if part.isdigit():
                n = int(part)
                if 1 <= n <= max_n:
                    cited.add(n)
    return sorted(cited)


def _renumber_answer(answer: str, old_to_new: dict) -> str:
    """
    After dropping uncited sources, the surviving ones are renumbered 1..N.
    Rewrite the markers in the answer text to match, so [12] becomes [3] if
    source 12 is now the third in the list.
    """
    def replace(match):
        seen = []
        for part in match.group(1).split(","):
            part = part.strip()
            if part.isdigit() and int(part) in old_to_new:
                new = str(old_to_new[int(part)])
                # Two chunks from the same document map to the same new
                # number — don't emit "[1,1]".
                if new not in seen:
                    seen.append(new)
        return f"[{','.join(seen)}]" if seen else ""

    return re.sub(r"\[([\d\s,]+)\]", replace, answer)


def answer_question(session_id: str, question: str, filter_kind: Optional[str] = None) -> dict:
    """
    Returns:
      {"answer": str, "citations": [{"n": 1, "title":.., "url":.., "kind":.., "document_id":..}], "chunks_used": int}

    Note: `citations` contains only the sources the model actually cited in
    its answer — not every chunk that was retrieved. Retrieval deliberately
    pulls a broad set of candidates (TOP_K_CHUNKS) so the model has enough
    context to work with, but listing all of them under the answer would be
    misleading: most retrieved chunks don't end up informing the response.
    """
    chunks = retrieve(question, filter_kind=filter_kind)
    if not chunks:
        return {
            "answer": (
                "I don't have any ingested bill text, PDFs, articles, or public "
                "opinions that match this question yet. Try asking about a bill "
                "that's already been ingested, or trigger an ingestion run."
            ),
            "citations": [],
            "chunks_used": 0,
        }

    context = _format_context(chunks)

    history = db.get_recent_messages(session_id, limit=6)
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for h in history:
        messages.append({"role": h["role"], "content": h["content"]})
    messages.append(
        {
            "role": "user",
            "content": f"SOURCES:\n{context}\n\nQUESTION: {question}",
        }
    )

    answer = _call_groq(messages)

    # Keep only the sources the model actually cited. Multiple chunks often
    # come from the same document, so also de-duplicate by URL — citing the
    # same PDF three times as [2], [7], [13] is noise, not precision.
    cited_numbers = _extract_cited_numbers(answer, len(chunks))

    citations = []
    old_to_new = {}
    seen_urls = {}
    for n in cited_numbers:
        c = chunks[n - 1]
        url = c["url"]
        if url in seen_urls:
            # Already citing this document — point this marker at the
            # existing entry instead of adding a duplicate.
            old_to_new[n] = seen_urls[url]
            continue
        new_n = len(citations) + 1
        seen_urls[url] = new_n
        old_to_new[n] = new_n
        citations.append(
            {
                "n": new_n,
                "document_id": c["document_id"],
                "title": c["title"],
                "url": url,
                "kind": c["kind"],
                "page_number": c.get("page_number"),
            }
        )

    answer = _renumber_answer(answer, old_to_new)

    return {"answer": answer, "citations": citations, "chunks_used": len(chunks)}