Most RAG systems retrieve documents and generate answers — but never connect the two in a way users can actually trust or verify. This lesson builds a complete citation pipeline that maps every LLM claim back to a specific source chunk, complete with composite confidence scores and hallucination detection.
Picture this: your company has deployed a RAG-powered internal knowledge base chatbot. An analyst asks it about your data retention policy, and it confidently returns a three-paragraph answer. The answer sounds authoritative. But which policy document did it pull from? The one from 2019 that's since been superseded, or the updated 2024 version? Was it a 90% relevant match, or a desperate stretch from a loosely related HR memo? Without source attribution, you have no idea — and neither does the analyst.
This is the silent liability of unattributed RAG systems. The retrieval step works, the generation step works, but the trust layer is completely missing. In production environments — legal, compliance, healthcare, finance, enterprise knowledge management — "the LLM said so" is not an acceptable citation. You need to be able to trace every claim in a generated response back to a specific document chunk, with a meaningful confidence signal attached to that attribution.
By the end of this lesson, you'll have built a complete citation and source attribution pipeline that sits on top of any standard RAG architecture. You'll go from raw LLM output to structured, verifiable responses where every statement has a traceable source and a confidence score that actually means something.
What you'll learn:
You should be comfortable with the fundamentals of RAG — embeddings, vector stores, chunk retrieval, and LLM prompting. Familiarity with Python is assumed throughout. We'll use LangChain for orchestration, OpenAI for embeddings and generation, Cohere for reranking, and Pydantic for structured output validation. You don't need to have used all of these together before, but you should have touched at least one retrieval system in a real project.
Before building the solution, it's worth being precise about what breaks in a standard RAG pipeline when you try to add attribution after the fact.
The typical RAG loop looks like this: embed the query, retrieve the top-k chunks by cosine similarity, concatenate them into a context window, and ask the LLM to answer based on that context. The retrieval scores exist during the retrieval step, but they're almost never passed into the prompt or preserved in the output. The LLM synthesizes a response from the concatenated chunks and returns a string — with zero indication of which parts of its answer came from which chunk.
Here's a simplified version of what this looks like in practice:
from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA
vectorstore = Chroma(
persist_directory="./policy_docs",
embedding_function=OpenAIEmbeddings()
)
qa_chain = RetrievalQA.from_chain_type(
llm=ChatOpenAI(model="gpt-4o"),
retriever=vectorstore.as_retriever(search_kwargs={"k": 5}),
return_source_documents=True
)
result = qa_chain.invoke({"query": "What is our data retention policy for customer records?"})
print(result["result"])
print(result["source_documents"]) # You get back Document objects, but...
Running this, you'll get source_documents — a list of LangChain Document objects with their content and metadata. That seems like attribution! But look closer. You have five documents returned, and a single blended answer. There's no mapping from any sentence in the answer to any specific document. The cosine similarity scores used to retrieve those documents are gone — as_retriever() silently discards them. And there's no signal about whether the LLM actually used document 3 at all, or just padded its answer from documents 1 and 2.
This is the attribution gap. Closing it requires changes at four levels: the retrieval layer, the prompt layer, the generation layer, and the output structure layer.
Before writing retrieval code, get your data model right. A citation system that collapses under schema changes is worse than no system at all. We'll use Pydantic to define the structures that will flow through the entire pipeline.
from pydantic import BaseModel, Field
from typing import Optional
from datetime import datetime
class SourceChunk(BaseModel):
"""Represents a single retrieved document chunk with full provenance."""
chunk_id: str
document_id: str
document_title: str
document_url: Optional[str] = None
document_date: Optional[datetime] = None
page_number: Optional[int] = None
section_heading: Optional[str] = None
chunk_text: str
vector_similarity_score: float = Field(ge=0.0, le=1.0)
reranker_score: Optional[float] = None
composite_confidence: Optional[float] = None
class CitedClaim(BaseModel):
"""A single sentence or claim in the LLM response, with its source attribution."""
claim_text: str
source_chunk_ids: list[str] # Can reference multiple chunks
attribution_confidence: float = Field(ge=0.0, le=1.0)
is_synthesized: bool = False # True if claim spans multiple sources
class CitedResponse(BaseModel):
"""The complete structured output of an attributed RAG query."""
query: str
answer_text: str # The full LLM-generated answer
cited_claims: list[CitedClaim]
source_chunks: list[SourceChunk]
overall_confidence: float = Field(ge=0.0, le=1.0)
retrieval_timestamp: datetime = Field(default_factory=datetime.utcnow)
model_id: str
warning_flags: list[str] = Field(default_factory=list)
A few design decisions here worth explaining. CitedClaim allows a claim to reference multiple chunk IDs — this matters because LLMs frequently synthesize across sources. A sentence like "Our 30-day retention policy applies globally, with regional exceptions in the EU" might draw from two separate policy documents. Forcing a one-to-one claim-to-source mapping would either be inaccurate or require splitting claims unnaturally.
The is_synthesized flag is important for downstream consumers. A legal team needs to know the difference between a claim that was directly quoted from a document versus one the LLM constructed by combining information from multiple sources. The latter requires more scrutiny.
warning_flags is a list of strings that will catch things like: low overall confidence, retrieval scores below threshold, mismatches between what the LLM cited and what was actually retrieved. It's your safety net.
The first concrete change from a standard RAG pipeline: stop using as_retriever() for anything serious. It's convenient but lossy. Instead, call the underlying vector store's similarity search directly so you can capture scores.
import hashlib
from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
vectorstore = Chroma(
persist_directory="./policy_docs",
embedding_function=embeddings
)
def retrieve_with_scores(query: str, k: int = 8) -> list[SourceChunk]:
"""
Retrieve document chunks with cosine similarity scores preserved.
We over-retrieve (k=8) to give the reranker more to work with.
"""
results = vectorstore.similarity_search_with_relevance_scores(
query=query,
k=k
)
source_chunks = []
for doc, score in results:
metadata = doc.metadata
# Generate a stable chunk ID from content hash if not stored
chunk_id = metadata.get(
"chunk_id",
hashlib.md5(doc.page_content.encode()).hexdigest()[:12]
)
source_chunks.append(SourceChunk(
chunk_id=chunk_id,
document_id=metadata.get("document_id", "unknown"),
document_title=metadata.get("title", "Untitled Document"),
document_url=metadata.get("url"),
document_date=metadata.get("created_at"),
page_number=metadata.get("page_number"),
section_heading=metadata.get("section"),
chunk_text=doc.page_content,
vector_similarity_score=round(score, 4)
))
return source_chunks
Notice we're retrieving k=8 instead of the standard k=4 or k=5. This is intentional — we're about to rerank these results, and reranking is much more effective when it has a larger candidate pool to choose from. The top-4 by cosine similarity is often not the same as the top-4 by actual relevance, especially with long or ambiguous queries.
Tip
similarity_search_with_relevance_scores() returns normalized scores between 0 and 1, unlike similarity_search_with_score() which returns raw L2 distances. Always use the relevance score variant unless you specifically need distances.
Vector similarity scores are useful but noisy. A chunk can have high cosine similarity to a query just because it shares vocabulary, even if it doesn't actually answer the question. A cross-encoder reranker reads the query and each chunk together, producing a much more reliable relevance signal.
import cohere
import os
cohere_client = cohere.Client(os.environ["COHERE_API_KEY"])
def rerank_chunks(
query: str,
chunks: list[SourceChunk],
top_n: int = 5
) -> list[SourceChunk]:
"""
Rerank retrieved chunks using Cohere's cross-encoder.
Updates reranker_score and computes composite_confidence.
"""
documents = [chunk.chunk_text for chunk in chunks]
response = cohere_client.rerank(
model="rerank-english-v3.0",
query=query,
documents=documents,
top_n=top_n,
return_documents=True
)
reranked = []
for result in response.results:
original_chunk = chunks[result.index]
# Reranker relevance score is already 0-1
reranker_score = result.relevance_score
# Composite confidence: weighted blend of vector and reranker signals
# We weight reranker more heavily — it's a stronger signal
composite = (
0.3 * original_chunk.vector_similarity_score +
0.7 * reranker_score
)
reranked.append(original_chunk.model_copy(update={
"reranker_score": round(reranker_score, 4),
"composite_confidence": round(composite, 4)
}))
# Sort by composite confidence descending
return sorted(reranked, key=lambda c: c.composite_confidence, reverse=True)
The composite confidence formula — 30% vector similarity, 70% reranker score — isn't arbitrary. The reranker is a cross-encoder that reads both the query and the document simultaneously, making it far more context-aware than the bi-encoder used for vector similarity. In practice, the reranker catches cases where a chunk is topically adjacent but doesn't actually answer the question, which is exactly the false-confidence scenario you need to avoid in an attribution system.
You can tune these weights. If your document corpus is highly domain-specific and your embeddings are fine-tuned on that domain, the vector similarity signal becomes more trustworthy and you might shift to 40/60. For general-purpose embeddings against specialized corpora (say, OpenAI embeddings on proprietary legal documents), stick closer to 20/80.
Warning
Don't skip the reranking step just to save an API call. In attribution systems, the composite confidence score is what downstream consumers actually use to decide whether to trust a response. A score built only on vector similarity will mislead them.
This is where most teams get stuck. The temptation is to ask the LLM to "cite your sources" in free text — and then try to parse those citations out afterward. This creates a parsing nightmare and produces inconsistent results across different queries.
The better approach: give the LLM a structured schema and tell it to return JSON. More importantly, give each source chunk a reference ID in the prompt itself, so the LLM has something concrete to cite.
from langchain.chat_models import ChatOpenAI
from langchain.schema import HumanMessage, SystemMessage
import json
SYSTEM_PROMPT = """You are a precise research assistant that answers questions using only the provided source documents.
Your response MUST be a valid JSON object matching this exact schema:
{
"answer_text": "Your complete, well-written answer here",
"cited_claims": [
{
"claim_text": "A single sentence or clause from your answer",
"source_chunk_ids": ["chunk_id_1", "chunk_id_2"],
"attribution_confidence": 0.95,
"is_synthesized": false
}
]
}
Rules:
1. Break your answer into individual claims — one per sentence or distinct factual statement.
2. Each claim MUST cite at least one source chunk ID from the provided documents.
3. If a claim combines information from multiple chunks, list all relevant chunk IDs and set is_synthesized to true.
4. If you cannot support a claim with the provided documents, do not make it.
5. attribution_confidence should reflect how directly the source chunk supports the claim:
- 0.9-1.0: The source directly states this
- 0.7-0.9: The source strongly implies this
- 0.5-0.7: The source partially supports this, with some inference required
- Below 0.5: Do not include this claim
6. answer_text should be a fluent, complete answer assembled from your cited_claims.
"""
def build_context_block(chunks: list[SourceChunk]) -> str:
"""Format retrieved chunks with explicit IDs for the LLM to reference."""
lines = ["### SOURCE DOCUMENTS\n"]
for chunk in chunks:
lines.append(f"[SOURCE: {chunk.chunk_id}]")
lines.append(f"Title: {chunk.document_title}")
if chunk.section_heading:
lines.append(f"Section: {chunk.section_heading}")
if chunk.page_number:
lines.append(f"Page: {chunk.page_number}")
lines.append(f"Content: {chunk.chunk_text}")
lines.append("---")
return "\n".join(lines)
def generate_cited_response(
query: str,
chunks: list[SourceChunk],
model_id: str = "gpt-4o"
) -> dict:
"""
Call the LLM with attribution-aware prompt and return structured JSON.
"""
llm = ChatOpenAI(
model=model_id,
temperature=0, # Determinism matters for structured output
response_format={"type": "json_object"}
)
context_block = build_context_block(chunks)
user_message = f"""Answer the following question using only the provided source documents.
QUESTION: {query}
{context_block}
Return your response as JSON following the schema in your instructions."""
messages = [
SystemMessage(content=SYSTEM_PROMPT),
HumanMessage(content=user_message)
]
response = llm.invoke(messages)
return json.loads(response.content)
A few things worth examining here. Setting temperature=0 is non-negotiable for structured output. Any temperature above zero introduces JSON formatting variation that will cause parse errors in production. Using response_format={"type": "json_object"} with OpenAI's models forces JSON-only output — this is much more reliable than telling the LLM "output JSON" in the prompt alone.
The context block format ([SOURCE: chunk_id]) gives the LLM a natural reference syntax. When you see GPT-4o output "source_chunk_ids": ["a3f8b2c1"], you know exactly which chunk to look up. This is the bridge between generation and retrieval that the naive RAG pipeline completely lacks.
Tip
If you're using a model that doesn't support response_format, add this to the end of your system prompt: "Your response must start with { and end with }. Output nothing other than the JSON object." It's less reliable but often sufficient.
Now we wire everything together into the full CitedResponse object, with validation, warning flag generation, and overall confidence computation.
from datetime import datetime
def validate_chunk_references(
cited_claims: list[dict],
available_chunk_ids: set[str]
) -> tuple[list[dict], list[str]]:
"""
Validate that the LLM only cited chunks that were actually retrieved.
Returns cleaned claims and any warning flags generated.
"""
warnings = []
cleaned_claims = []
for claim in cited_claims:
valid_ids = [
cid for cid in claim["source_chunk_ids"]
if cid in available_chunk_ids
]
if not valid_ids:
warnings.append(
f"Claim dropped — cited non-existent chunk IDs: "
f"{claim['source_chunk_ids']}"
)
continue
if len(valid_ids) < len(claim["source_chunk_ids"]):
warnings.append(
f"Claim partially attributed — some cited IDs not in retrieved set"
)
cleaned_claims.append({**claim, "source_chunk_ids": valid_ids})
return cleaned_claims, warnings
def compute_overall_confidence(
cited_claims: list[dict],
source_chunks: list[SourceChunk]
) -> float:
"""
Overall confidence is a weighted combination of:
- Average attribution confidence across claims
- Average composite confidence of cited source chunks
- Coverage: what fraction of retrieved chunks were actually cited
"""
if not cited_claims:
return 0.0
# Mean attribution confidence across all claims
avg_attribution = sum(
c["attribution_confidence"] for c in cited_claims
) / len(cited_claims)
# Mean composite confidence of source chunks that were actually cited
cited_ids = set()
for claim in cited_claims:
cited_ids.update(claim["source_chunk_ids"])
cited_chunks = [c for c in source_chunks if c.chunk_id in cited_ids]
if cited_chunks:
avg_chunk_confidence = sum(
c.composite_confidence for c in cited_chunks
if c.composite_confidence is not None
) / len(cited_chunks)
else:
avg_chunk_confidence = 0.0
# Weighted blend: attribution quality matters more than retrieval quality
overall = 0.6 * avg_attribution + 0.4 * avg_chunk_confidence
return round(overall, 4)
def run_attributed_rag(
query: str,
model_id: str = "gpt-4o"
) -> CitedResponse:
"""
Full pipeline: retrieve → rerank → generate → attribute → validate.
"""
# Step 1: Retrieve with score preservation
raw_chunks = retrieve_with_scores(query, k=8)
if not raw_chunks:
return CitedResponse(
query=query,
answer_text="No relevant documents found.",
cited_claims=[],
source_chunks=[],
overall_confidence=0.0,
model_id=model_id,
warning_flags=["no_documents_retrieved"]
)
# Step 2: Rerank and select top-5
ranked_chunks = rerank_chunks(query, raw_chunks, top_n=5)
# Step 3: Generate response with citation instructions
llm_output = generate_cited_response(query, ranked_chunks, model_id)
# Step 4: Validate chunk references
available_ids = {chunk.chunk_id for chunk in ranked_chunks}
clean_claims_dicts, warnings = validate_chunk_references(
llm_output.get("cited_claims", []),
available_ids
)
# Step 5: Convert claim dicts to CitedClaim objects
cited_claims = [CitedClaim(**claim) for claim in clean_claims_dicts]
# Step 6: Compute overall confidence
overall_confidence = compute_overall_confidence(
clean_claims_dicts, ranked_chunks
)
# Step 7: Add confidence-based warnings
if overall_confidence < 0.5:
warnings.append("low_overall_confidence")
low_confidence_claims = [
c for c in cited_claims if c.attribution_confidence < 0.7
]
if len(low_confidence_claims) > len(cited_claims) * 0.3:
warnings.append("majority_claims_low_confidence")
return CitedResponse(
query=query,
answer_text=llm_output.get("answer_text", ""),
cited_claims=cited_claims,
source_chunks=ranked_chunks,
overall_confidence=overall_confidence,
model_id=model_id,
warning_flags=warnings
)
The validation step in validate_chunk_references is critical and easy to skip. LLMs occasionally hallucinate chunk IDs — they'll cite a chunk_id that doesn't exist in the retrieved set. Without this check, you'd surface fake citations to your users, which is arguably worse than no citations at all. Any claim that cites only nonexistent IDs gets dropped entirely. Any claim with a mix of valid and invalid IDs gets the invalid ones stripped and a warning flag added.
The CitedResponse object is your internal representation. What end users actually see depends on your interface, but here's a function that renders it in a readable format for API responses or terminal output.
def render_cited_response(response: CitedResponse) -> dict:
"""
Render a CitedResponse into a clean API-friendly dict.
"""
chunk_lookup = {c.chunk_id: c for c in response.source_chunks}
rendered_claims = []
for claim in response.cited_claims:
sources = []
for chunk_id in claim.source_chunk_ids:
chunk = chunk_lookup.get(chunk_id)
if chunk:
sources.append({
"document_title": chunk.document_title,
"section": chunk.section_heading,
"page": chunk.page_number,
"url": chunk.document_url,
"chunk_confidence": chunk.composite_confidence
})
rendered_claims.append({
"text": claim.claim_text,
"sources": sources,
"confidence": claim.attribution_confidence,
"synthesized": claim.is_synthesized
})
return {
"query": response.query,
"answer": response.answer_text,
"overall_confidence": response.overall_confidence,
"confidence_label": _confidence_label(response.overall_confidence),
"claims": rendered_claims,
"warnings": response.warning_flags,
"retrieved_at": response.retrieval_timestamp.isoformat()
}
def _confidence_label(score: float) -> str:
if score >= 0.85:
return "HIGH"
elif score >= 0.65:
return "MEDIUM"
elif score >= 0.45:
return "LOW"
else:
return "VERY_LOW"
This gives you a clean JSON payload that a React frontend, a Slack bot, or a downstream API consumer can actually work with. The confidence_label converts the numerical score into something a non-technical stakeholder can understand. A compliance officer doesn't know what 0.71 means, but they understand "MEDIUM confidence."
Build an end-to-end attributed RAG system for a policy document corpus. Here's the scenario: you have a set of HR policy PDFs (you can use your company's actual policies, or download public ones). Your goal is to query this corpus and return fully attributed, confidence-scored responses.
Step 1: Ingest your documents with chunk metadata
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
import uuid
def ingest_policy_document(pdf_path: str, document_title: str, document_url: str = None):
loader = PyPDFLoader(pdf_path)
pages = loader.load()
splitter = RecursiveCharacterTextSplitter(
chunk_size=800,
chunk_overlap=100,
separators=["\n\n", "\n", ". ", " "]
)
chunks = splitter.split_documents(pages)
# Enrich each chunk with the metadata our attribution system needs
for i, chunk in enumerate(chunks):
chunk.metadata.update({
"chunk_id": str(uuid.uuid4())[:12],
"document_id": hashlib.md5(pdf_path.encode()).hexdigest()[:8],
"title": document_title,
"url": document_url,
"page_number": chunk.metadata.get("page", None),
})
vectorstore.add_documents(chunks)
print(f"Ingested {len(chunks)} chunks from {document_title}")
# Usage:
ingest_policy_document("./policies/data_retention_2024.pdf", "Data Retention Policy v2024")
ingest_policy_document("./policies/remote_work_policy.pdf", "Remote Work Guidelines")
ingest_policy_document("./policies/information_security.pdf", "Information Security Standard")
Step 2: Run attributed queries and inspect the output
response = run_attributed_rag(
"What are our data retention requirements for customer PII, and how do they differ between the EU and US?"
)
rendered = render_cited_response(response)
import json
print(json.dumps(rendered, indent=2))
Step 3: Test your warning flags. Run a query that can't be answered from your documents — something like "What is our policy on employees keeping pets in the office?" — and verify that overall_confidence drops below 0.5 and warning_flags contains "low_overall_confidence". If the LLM starts fabricating policies, your citation validation should catch the hallucinated chunk IDs.
Step 4: Tune your composite confidence weights. Try adjusting the 0.3/0.7 vector-to-reranker split in rerank_chunks(). Run 10 queries with both weighting schemes and compare the warning flag rates. You're looking for the calibration that best separates genuinely well-supported answers from uncertain ones.
The LLM ignores your JSON schema and returns prose.
This almost always happens because the system prompt is too long and the schema instructions get buried. Move the schema to the beginning of the system prompt, before any other instructions. For production, use OpenAI's structured outputs feature (available on gpt-4o and later) which enforces JSON schemas at the model level.
All chunk IDs in citations are invalid (hallucinated).
The LLM is struggling to remember the chunk IDs from the context. Double-check that your build_context_block() function is clearly labeling each chunk's ID. Try making the format more visually distinct: <<<CHUNK_ID: a3f8b2c1>>> instead of [SOURCE: a3f8b2c1]. Also verify your context window isn't getting truncated — if the chunk IDs appear near the end of a very long context and the model is cutting off, it literally never sees them.
Confidence scores are consistently near 1.0 for everything. The LLM is being sycophantic about its confidence. Add this to your system prompt: "Be conservative with attribution_confidence scores. It is better to report a lower confidence than to overstate certainty." You can also do a calibration pass: sample 50 queries, manually judge attribution quality, and compare against the model's self-reported scores to see if they're inflated.
The reranker returns different top results every run. Reranking is deterministic given the same inputs. If you're seeing variation, check whether your vector retrieval is returning different chunks between runs — this can happen with approximate nearest neighbor indexes if they haven't been fully built. Also check that your query text isn't being preprocessed differently (lowercasing, stripping punctuation) inconsistently.
Warning: majority_claims_low_confidence fires constantly.
Either your retrieval corpus doesn't have good coverage of the queries you're running (corpus gap), or your chunk size is too large and chunks are too general to support specific claims. Try reducing chunk size to 400-600 tokens and re-ingesting. Overly large chunks score reasonably well in retrieval but poorly in attribution because they contain many topics simultaneously.
Pydantic validation errors on CitedClaim objects.
The most common cause is the LLM returning attribution_confidence as a string ("0.85") instead of a float. Add a Pydantic validator: @field_validator('attribution_confidence', mode='before') def coerce_float(cls, v): return float(v). This handles both string and numeric inputs robustly.
Adding reranking and structured generation adds latency compared to naive RAG. A typical breakdown on a gpt-4o + Cohere setup:
For most internal knowledge base use cases, this is acceptable. For real-time customer-facing applications, consider these tradeoffs:
Skip reranking for low-stakes queries. Implement a fast path where you bypass Cohere reranking for queries where vector similarity scores are already high and consistent (e.g., all top-5 scores above 0.85). You lose some attribution quality but halve your latency.
Cache responses by query hash. Policy documents change infrequently. A Redis cache keyed on sha256(query + corpus_version) with a 24-hour TTL will eliminate repeated computation for common questions.
Parallelize retrieval and any pre-processing. If you're running multiple vector stores (e.g., policies in one store, contracts in another), retrieve from both simultaneously with asyncio.gather().
import asyncio
async def retrieve_multi_corpus(query: str) -> list[SourceChunk]:
policy_chunks, contract_chunks = await asyncio.gather(
asyncio.to_thread(retrieve_with_scores, query, vectorstore_policies),
asyncio.to_thread(retrieve_with_scores, query, vectorstore_contracts)
)
return policy_chunks + contract_chunks
You've built a citation system that closes the attribution gap in RAG pipelines. The key architectural moves were: preserving retrieval scores by bypassing the standard retriever abstraction, adding cross-encoder reranking to produce more meaningful confidence signals, constructing a prompt that gives the LLM concrete chunk IDs to cite, validating those citations against the actually-retrieved set to catch hallucinations, and assembling everything into a structured CitedResponse object with composite confidence scoring.
This system is directly deployable for internal knowledge bases, compliance Q&A tools, and any domain where "trust but verify" is the operating requirement.
Where to go from here:
last_verified timestamp to your SourceChunk model and flag citations that reference documents not updated in N months.