Citation RAG - Grounded Answers with Source Attribution

from typing import List, Dict, Tuple
import re


def number_sources(documents: List[Dict]) -> Tuple[List[Dict], str]:
    """
    Assign citation IDs [1], [2], ... to retrieved chunks.
    Returns annotated docs and a formatted sources block for the prompt.
    """
    numbered = []
    source_block_lines = []

    for i, doc in enumerate(documents, start=1):
        numbered.append({**doc, "citation_id": i})
        source_block_lines.append(f"[{i}] {doc['content']}")

    sources_block = "\n\n".join(source_block_lines)
    return numbered, sources_block

CITATION_SYSTEM_PROMPT = """You are a precise assistant that always cites sources.

Rules:
1. After every factual claim, add [N] where N is the source number that supports it.
2. Only state facts that are directly supported by the provided sources.
3. If the sources do not contain information needed to answer, say so — do not infer.
4. Every sentence with a factual claim must have at least one citation.
5. Do not add citations to general framing or transition sentences."""

CITATION_USER_TEMPLATE = """Sources:
{sources_block}

Question: {question}

Answer (cite sources inline as [N] after each claim):"""


def citation_prompt(question: str, sources_block: str) -> str:
    return CITATION_USER_TEMPLATE.format(
        sources_block=sources_block,
        question=question,
    )

def parse_citations(response: str) -> Dict:
    """
    Extract citation markers from the generated response.
    Returns the full response and a per-sentence citation map.
    """
    all_citations = re.findall(r'\[(\d+)\]', response)
    cited_ids = sorted(set(int(c) for c in all_citations))

    sentence_citations = []
    sentences = re.split(r'(?<=[.!?])\s+', response.strip())
    for sentence in sentences:
        markers = re.findall(r'\[(\d+)\]', sentence)
        clean = re.sub(r'\[\d+\]', '', sentence).strip()
        sentence_citations.append({
            "sentence": clean,
            "citations": [int(m) for m in markers],
            "has_citation": len(markers) > 0,
        })

    return {
        "full_response": response,
        "sentence_citations": sentence_citations,
        "all_cited_ids": cited_ids,
    }


def validate_citations(
    parsed: Dict,
    numbered_docs: List[Dict],
    llm,
    validate_with_llm: bool = False,
) -> Dict:
    """
    Validate that each cited source actually supports the claim.

    validate_with_llm=True:  use an LLM judge for each (sentence, source) pair
    validate_with_llm=False: check only that the cited ID exists (fast, lightweight)
    """
    doc_lookup = {d["citation_id"]: d for d in numbered_docs}
    valid_ids = set(doc_lookup.keys())

    validation_results = []

    for sc in parsed["sentence_citations"]:
        for cid in sc["citations"]:
            if cid not in valid_ids:
                validation_results.append({
                    "sentence": sc["sentence"],
                    "citation_id": cid,
                    "valid": False,
                    "reason": "citation ID does not exist",
                })
                continue

            if validate_with_llm:
                source_text = doc_lookup[cid]["content"]
                check_prompt = f"""Does the following source support the claim?

Source: {source_text}

Claim: {sc['sentence']}

Return JSON: {{"supports": true or false, "reason": "brief explanation"}}"""
                result = llm.invoke(check_prompt, response_format="json")
                validation_results.append({
                    "sentence": sc["sentence"],
                    "citation_id": cid,
                    "valid": result.get("supports", False),
                    "reason": result.get("reason", ""),
                })
            else:
                validation_results.append({
                    "sentence": sc["sentence"],
                    "citation_id": cid,
                    "valid": True,
                    "reason": "citation ID exists",
                })

    # Compute grounding metrics
    total_sentences = len(parsed["sentence_citations"])
    cited_sentences = sum(1 for sc in parsed["sentence_citations"] if sc["has_citation"])
    valid_count = sum(1 for v in validation_results if v["valid"])
    total_citations = len(validation_results)

    confidence = valid_count / total_citations if total_citations > 0 else 0.0
    citation_coverage = cited_sentences / total_sentences if total_sentences > 0 else 0.0

    return {
        "confidence_score": confidence,
        "citation_coverage": citation_coverage,
        "validation_results": validation_results,
        "grounded": confidence >= 0.9 and citation_coverage >= 0.8,
    }

def citation_rag(
    query: str,
    retriever,
    llm,
    top_k: int = 5,
    validate_with_llm: bool = False,
) -> Dict:
    # Retrieve
    documents = retriever.search(query, top_k=top_k)

    # Number sources
    numbered_docs, sources_block = number_sources(documents)

    # Generate with citation instruction
    prompt = citation_prompt(query, sources_block)
    response = llm.invoke(
        system=CITATION_SYSTEM_PROMPT,
        prompt=prompt,
    )

    # Parse citations from response
    parsed = parse_citations(response)

    # Validate grounding
    validation = validate_citations(parsed, numbered_docs, llm, validate_with_llm)

    return {
        "answer": response,
        "parsed": parsed,
        "validation": validation,
        "sources": numbered_docs,
        "confidence_score": validation["confidence_score"],
        "grounded": validation["grounded"],
    }


# Usage
result = citation_rag(
    query="What file types and size limits apply to uploads?",
    retriever=retriever,
    llm=llm,
    validate_with_llm=True,
)

print(result["answer"])
# "You can upload files up to 25 MB [1] in PDF, DOCX, or TXT format [2].
#  Uploaded files are retained for 30 days [3]."

print(f"Confidence: {result['confidence_score']:.0%}")
# Confidence: 100%