Build a Citation-Aware RAG Pipeline

The citation problem in RAG

A user asks: "What's your return policy for shoes?"

Your RAG system retrieves three chunks:

• Chunk 45 (from doc-returns): "Returns are accepted within 30 days of purchase."
• Chunk 127 (from doc-footwear): "Shoe damage from normal wear is not covered under returns."
• Chunk 89 (from doc-shipping): "We ship shoes within 2 business days."

The model generates: "We accept shoe returns within 30 days, though normal wear damage isn't covered."

The problem: Which chunk did each claim come from? Did the model cite chunk 45 and 127? Or did it invent the claim about damage coverage? Without explicit citation enforcement, you can't tell.

The solution: A citation-aware pipeline that:

1. Tags ingestion: Mark every chunk with its source document and section.
2. Carries IDs through retrieval: Each retrieved chunk knows its chunk_id and doc_id.
3. Enforces citations in generation: Instruct the model to cite source IDs after each claim.
4. Verifies citations: Check that cited IDs were actually retrieved.

This lesson builds a complete, production-ready implementation.

Step 1: Document ingestion with metadata tagging

Raw documents come from various sources: PDFs, web pages, internal wikis, databases. Each document needs metadata: title, URL, owner, publication date, access level, and section path.

import json
from datetime import datetime
from dataclasses import dataclass
from typing import List, Optional

@dataclass
class DocumentMetadata:
    """Metadata for every document."""
    doc_id: str                    # Unique identifier (e.g., "doc-returns-2024")
    title: str                     # Human-readable title
    url: str                       # Source URL or internal path
    owner: str                     # Who owns this document (for access control)
    published_date: str            # ISO 8601 date
    version: str                   # Document version (e.g., "1.0", "2.1")
    access_level: str              # "public", "internal", "confidential"
    source_type: str               # "policy", "faq", "blog", "legal", "internal"

@dataclass
class Document:
    """A full document with metadata."""
    metadata: DocumentMetadata
    text: str                      # Full document text
    sections: List[str] = None     # Optional: pre-identified sections (["intro", "policy", "exceptions"])

def parse_documents_from_disk(doc_dir):
    """Load documents from disk. Format: JSON files with 'metadata' and 'text' keys."""
    documents = []

    import os
    for filename in os.listdir(doc_dir):
        if not filename.endswith('.json'):
            continue

        with open(os.path.join(doc_dir, filename)) as f:
            raw = json.load(f)

        # Validate required fields
        assert 'doc_id' in raw, f"Missing doc_id in {filename}"
        assert 'title' in raw, f"Missing title in {filename}"
        assert 'url' in raw, f"Missing url in {filename}"
        assert 'text' in raw, f"Missing text in {filename}"

        # Set defaults for optional fields
        raw.setdefault('owner', 'Unknown')
        raw.setdefault('published_date', datetime.now().isoformat()[:10])
        raw.setdefault('version', '1.0')
        raw.setdefault('access_level', 'public')
        raw.setdefault('source_type', 'general')

        metadata = DocumentMetadata(
            doc_id=raw['doc_id'],
            title=raw['title'],
            url=raw['url'],
            owner=raw['owner'],
            published_date=raw['published_date'],
            version=raw['version'],
            access_level=raw['access_level'],
            source_type=raw['source_type']
        )

        doc = Document(
            metadata=metadata,
            text=raw['text'],
            sections=raw.get('sections', None)
        )
        documents.append(doc)

    print(f"Loaded {len(documents)} documents from {doc_dir}")
    return documents

# Example usage
documents = parse_documents_from_disk('documents/')

# Inspect
for doc in documents[:2]:
    print(f"- {doc.metadata.doc_id}: {doc.metadata.title} (published {doc.metadata.published_date})")

Step 2: Chunking with provenance preservation

Documents are chunked, but every chunk must remember its parent document and section.

from dataclasses import dataclass
import re

@dataclass
class Chunk:
    """A chunk of text with full provenance."""
    chunk_id: str                  # Unique ID (e.g., "doc-returns-chunk-3")
    doc_id: str                    # Parent document
    section: str                   # Section within document (e.g., "Return Policy")
    text: str                      # Chunk text
    url: str                        # Document URL
    metadata: dict                 # Extra: owner, published_date, access_level

def chunk_document(doc: Document, chunk_size=300, overlap=50) -> List[Chunk]:
    """Split a document into overlapping chunks, preserving provenance."""

    chunks = []
    text = doc.text

    # Optional: use document sections if available
    if doc.sections:
        # Split by section headers
        # (simplified: assume section headers are marked as "## Section Name")
        section_pattern = r'^## (.+)$'
        parts = re.split(section_pattern, text, flags=re.MULTILINE)
        # parts will be [prefix, section_1_title, section_1_text, section_2_title, ...]

        current_section = "Preamble"
        i = 0
        while i < len(parts):
            if i % 2 == 0:
                # This is text
                section_text = parts[i]
            else:
                # This is a section title
                current_section = parts[i]
                section_text = parts[i + 1] if i + 1 < len(parts) else ""
                i += 1

            # Chunk the section text
            start = 0
            while start < len(section_text):
                end = min(start + chunk_size, len(section_text))
                chunk_text = section_text[start:end]

                chunk_id = f"{doc.metadata.doc_id}-chunk-{len(chunks)}"
                chunks.append(Chunk(
                    chunk_id=chunk_id,
                    doc_id=doc.metadata.doc_id,
                    section=current_section,
                    text=chunk_text,
                    url=doc.metadata.url,
                    metadata={
                        'owner': doc.metadata.owner,
                        'published_date': doc.metadata.published_date,
                        'access_level': doc.metadata.access_level,
                        'source_type': doc.metadata.source_type
                    }
                ))

                # Move to next chunk with overlap
                start = end - overlap if end < len(section_text) else end

            i += 1
    else:
        # Simple fixed-size chunking
        start = 0
        while start < len(text):
            end = min(start + chunk_size, len(text))
            chunk_text = text[start:end]

            chunk_id = f"{doc.metadata.doc_id}-chunk-{len(chunks)}"
            chunks.append(Chunk(
                chunk_id=chunk_id,
                doc_id=doc.metadata.doc_id,
                section="Main",
                text=chunk_text,
                url=doc.metadata.url,
                metadata={
                    'owner': doc.metadata.owner,
                    'published_date': doc.metadata.published_date,
                    'access_level': doc.metadata.access_level,
                    'source_type': doc.metadata.source_type
                }
            ))

            start = end - overlap if end < len(text) else end

    print(f"Chunked {doc.metadata.doc_id} into {len(chunks)} chunks")
    return chunks

# Chunk all documents
all_chunks = []
for doc in documents:
    chunks = chunk_document(doc, chunk_size=300, overlap=50)
    all_chunks.extend(chunks)

print(f"Total chunks: {len(all_chunks)}")

Step 3: Indexing and retrieval with ID preservation

Chunks are embedded and indexed. Crucially, we store chunk_id and doc_id alongside the embedding, so retrieval returns structured objects, not just text.

import numpy as np
from sentence_transformers import SentenceTransformer

class CitationAwareRetriever:
    def __init__(self, chunks: List[Chunk], model_name='all-mpnet-base-v2'):
        self.chunks = chunks
        self.chunk_by_id = {chunk.chunk_id: chunk for chunk in chunks}

        # Embed all chunks
        self.model = SentenceTransformer(model_name)
        chunk_texts = [chunk.text for chunk in chunks]
        self.embeddings = self.model.encode(chunk_texts, convert_to_numpy=True)
        self.embeddings = self.embeddings / np.linalg.norm(self.embeddings, axis=1, keepdims=True)

        print(f"Indexed {len(chunks)} chunks")

    def retrieve(self, query: str, top_k=10) -> List[dict]:
        """Retrieve chunks and return structured objects with metadata."""

        # Embed query
        query_embedding = self.model.encode(query, convert_to_numpy=True)
        query_embedding = query_embedding / np.linalg.norm(query_embedding)

        # Find top-k
        similarities = np.dot(self.embeddings, query_embedding)
        top_indices = np.argsort(similarities)[::-1][:top_k]

        # Build result objects with full provenance
        results = []
        for rank, idx in enumerate(top_indices, 1):
            chunk = self.chunks[idx]
            results.append({
                'rank': rank,
                'chunk_id': chunk.chunk_id,
                'doc_id': chunk.doc_id,
                'section': chunk.section,
                'text': chunk.text,
                'url': chunk.url,
                'score': similarities[idx],
                'metadata': chunk.metadata
            })

        return results

# Create retriever
retriever = CitationAwareRetriever(all_chunks)

# Test retrieval
query = "What's your return policy?"
results = retriever.retrieve(query, top_k=10)

print(f"\nRetrieved {len(results)} chunks for: {query}")
for r in results[:3]:
    print(f"\n{r['rank']}. Chunk {r['chunk_id']} (score={r['score']:.3f})")
    print(f"   Section: {r['section']}")
    print(f"   Text: {r['text'][:100]}...")

Step 4: Prompt design for enforced citations

The model must cite a source ID after every factual claim. This requires explicit instruction in the prompt.

from anthropic import Anthropic

def build_citation_prompt(query: str, retrieved_chunks: List[dict]) -> str:
    """Build a prompt that enforces citations."""

    # Format retrieved chunks with IDs
    context_text = ""
    for r in retrieved_chunks:
        context_text += f"\n[{r['chunk_id']}] {r['section']}:\n{r['text']}\n"

    prompt = f"""You are a helpful support assistant. Answer the user's question using ONLY the provided sources.

Instructions:
1. Answer using only information from the sources below.
2. After every factual claim, cite the source ID in square brackets, e.g.: "Returns are accepted within 30 days [chunk-123]."
3. If the sources don't contain enough information to fully answer, say: "I can only partially answer this based on available sources: [explain what you can answer]"
4. If the sources don't address the question at all, say: "I don't have information about that in the available sources."
5. Never cite a source ID that doesn't appear in the sources below.
6. Never invent source IDs.

Available sources:
{context_text}

User question: {query}

Answer:"""

    return prompt

# Test
query = "What's your return policy for shoes?"
retrieved = retriever.retrieve(query, top_k=5)
prompt = build_citation_prompt(query, retrieved)

print("Generated prompt (first 500 chars):")
print(prompt[:500])

Step 5: Generation with citation enforcement

Call the model with the citation-enforcing prompt, extract citations, and prepare for verification.

from anthropic import Anthropic

def generate_answer_with_citations(query: str, retriever: CitationAwareRetriever):
    """Generate an answer and track which source IDs the model cited."""

    client = Anthropic()

    # Retrieve
    retrieved = retriever.retrieve(query, top_k=10)
    retrieved_ids = set(r['chunk_id'] for r in retrieved)

    # Build prompt
    prompt = build_citation_prompt(query, retrieved)

    # Generate answer
    response = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=500,
        messages=[
            {"role": "user", "content": prompt}
        ]
    )

    answer_text = response.content[0].text

    # Extract cited source IDs from the answer (pattern: [chunk-id])
    import re
    cited_ids = set(re.findall(r'\[([a-z0-9\-]+)\]', answer_text.lower()))

    result = {
        'query': query,
        'answer': answer_text,
        'retrieved_chunks': retrieved,
        'retrieved_ids': retrieved_ids,
        'cited_ids': cited_ids,
    }

    return result

# Test end-to-end
query = "What's your return policy for shoes?"
result = generate_answer_with_citations(query, retriever)

print(f"Query: {result['query']}")
print(f"\nAnswer:\n{result['answer']}")
print(f"\nRetrieved chunk IDs: {result['retrieved_ids']}")
print(f"Cited chunk IDs: {result['cited_ids']}")

Step 6: Citation verification

The critical step: verify that every cited source ID was actually retrieved.

def verify_citations(result: dict) -> dict:
    """Verify that cited sources are in the retrieved set."""

    cited = result['cited_ids']
    retrieved = result['retrieved_ids']

    # Check: is every cited ID in the retrieved set?
    hallucinated = cited - retrieved  # Cited but not retrieved
    verified = cited & retrieved       # Cited and retrieved (good)

    verification = {
        'cited_ids': cited,
        'retrieved_ids': retrieved,
        'verified_citations': verified,
        'hallucinated_citations': hallucinated,
        'is_valid': len(hallucinated) == 0,
        'validity_score': len(verified) / max(1, len(cited))  # Fraction of citations that are valid
    }

    return verification

# Verify
verification = verify_citations(result)

print("\n=== CITATION VERIFICATION ===")
print(f"Valid citations: {len(verification['verified_citations'])}")
print(f"Hallucinated citations: {len(verification['hallucinated_citations'])}")

if hallucinated := verification['hallucinated_citations']:
    print(f"\nWARNING: Model cited sources that were not retrieved: {hallucinated}")
    print("-> This suggests the model is making up sources.")
else:
    print("\nSUCCESS: All citations are from the retrieved set.")

print(f"Citation validity score: {verification['validity_score']:.1%}")

Step 7: Complete end-to-end workflow

Putting it all together:

def citation_aware_rag(query: str, retriever: CitationAwareRetriever, client: Anthropic) -> dict:
    """Full citation-aware RAG pipeline."""

    # Step 1: Retrieve
    retrieved = retriever.retrieve(query, top_k=10)
    retrieved_ids = set(r['chunk_id'] for r in retrieved)

    # Step 2: Generate with citations
    prompt = build_citation_prompt(query, retrieved)

    response = client.messages.create(
        model="claude-3-5-sonnet-20241022",
        max_tokens=500,
        messages=[{"role": "user", "content": prompt}]
    )

    answer_text = response.content[0].text

    # Step 3: Extract citations
    import re
    cited_ids = set(re.findall(r'\[([a-z0-9\-]+)\]', answer_text.lower()))

    # Step 4: Verify
    hallucinated = cited_ids - retrieved_ids
    verified = cited_ids & retrieved_ids

    # Step 5: Build enriched answer object
    if hallucinated:
        # Regenerate without the hallucinated citations (optional)
        # For now, just mark as unverified
        confidence = "LOW"
    else:
        confidence = "HIGH" if cited_ids else "MEDIUM"

    result = {
        'query': query,
        'answer': answer_text,
        'retrieved_chunks': [
            {
                'id': r['chunk_id'],
                'section': r['section'],
                'url': r['url'],
                'text': r['text']
            }
            for r in retrieved
        ],
        'citations': {
            'verified': list(verified),
            'hallucinated': list(hallucinated),
            'confidence': confidence
        },
        'metrics': {
            'retrieval_count': len(retrieved),
            'citation_count': len(cited_ids),
            'verification_rate': len(verified) / max(1, len(cited_ids))
        }
    }

    return result

# Full pipeline test
client = Anthropic()
query = "What's your return policy for shoes?"
output = citation_aware_rag(query, retriever, client)

print("=== FINAL OUTPUT ===")
print(f"Query: {output['query']}")
print(f"\nAnswer:\n{output['answer']}")
print(f"\nVerified citations: {output['citations']['verified']}")
if output['citations']['hallucinated']:
    print(f"Hallucinated citations: {output['citations']['hallucinated']}")
print(f"\nConfidence: {output['citations']['confidence']}")
print(f"Citation verification rate: {output['metrics']['verification_rate']:.1%}")

Handling edge cases

Case 1: Model doesn't cite anything

If the model generates an answer but includes no citations, the system should flag confidence as LOW:

if not cited_ids:
    answer_status = "UNCITED"
    # Optionally, regenerate with stronger citation enforcement

Case 2: Unsupported question

If no retrieved chunks are relevant, the model should refuse:

# Check: do top-retrieved chunks have any overlap with the query?
max_score = retrieved[0]['score']
if max_score < 0.5:  # Threshold: if best match is weak, probably unanswerable
    print("WARNING: Retrieved chunks may not be relevant. Recommend asking for clarification.")

Case 3: Access control

Enforce that answers only cite public sources:

def verify_citations_with_access_control(result: dict, access_level='public'):
    """Verify citations and check access level."""

    cited = result['cited_ids']
    retrieved = result['retrieved_chunks']

    for chunk in retrieved:
        if chunk['chunk_id'] in cited:
            chunk_access = chunk['metadata']['access_level']
            if chunk_access not in ['public', access_level]:
                print(f"Access violation: model cited {chunk['chunk_id']} which is {chunk_access} level")
                # Optionally remove from cited

Measuring citation accuracy

Create a labeled evaluation set and measure precision:

def evaluate_citation_accuracy(test_cases, retriever, client):
    """Measure what fraction of citations are actually correct."""

    correct_citations = 0
    total_citations = 0
    hallucinated_count = 0

    for query, expected_relevant_chunks in test_cases:
        result = citation_aware_rag(query, retriever, client)

        expected_relevant = set(expected_relevant_chunks)
        cited = result['citations']['verified'] + result['citations']['hallucinated']

        total_citations += len(cited)

        for citation_id in cited:
            if citation_id in expected_relevant:
                correct_citations += 1
            else:
                hallucinated_count += 1

    citation_precision = correct_citations / max(1, total_citations)
    hallucination_rate = hallucinated_count / max(1, total_citations)

    print(f"Citation Precision: {citation_precision:.1%}")
    print(f"Hallucination Rate: {hallucination_rate:.1%}")

# Test
test_cases = [
    ("What's your return policy?", ["chunk-returns-1", "chunk-returns-2"]),
    ("How long does shipping take?", ["chunk-shipping-1"]),
    # ... more test cases
]

evaluate_citation_accuracy(test_cases, retriever, client)

Production checklist

Before deploying a citation-aware RAG system:

1. Metadata completeness: Every chunk has doc_id, section, url, published_date.
2. Citation format consistency: Model is trained to cite in format [chunk-id], not (ref 1) or other formats.
3. Verification enabled: All answers are checked for hallucinated citations before showing users.
4. Access control: Answers only cite sources the user is allowed to see.
5. Logging: Store every query, answer, citations, and verification result for auditing.
6. Refresh policy: Re-index documents on a schedule (weekly, daily, etc.) to keep content fresh.

Common mistake

Assuming that because you retrieve documents, the model will automatically cite them. Without explicit instruction ("cite the source ID after each claim"), the model often generates plausible-sounding answers without citations. Worse, it sometimes cites source IDs that don't exist.

Also common: Not verifying citations before showing users. A hallucinated citation ("This is policy is from our official handbook [chunk-999]" when chunk-999 doesn't exist) destroys trust. Always verify before display.

Final common mistake: Forgetting to log everything. If a user later says "You told me X, but it's wrong," you won't be able to debug without logs of the query, retrieved chunks, citations, and verification result. Log comprehensively.