Retrieval Foundations for LLM Systems

# Minimal RAG pipeline skeleton
import openai

def simple_rag(query, retriever, client, top_k=5):
    """Retrieve, then generate a grounded answer."""

    # Step 1: Retrieve relevant passages
    passages = retriever.search(query, top_k=top_k)

    # Step 2: Assemble context for the generator
    context = "\n\n".join(
        f"[Source {i+1}]: {p.text}"
        for i, p in enumerate(passages)
    )

    # Step 3: Generate with retrieved context
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": """Answer based ONLY on the
provided context. If the context does not contain
the answer, say so. Cite [Source N] for each claim."""},
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}"},
        ],
    )
    return response.choices[0].message.content

from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer
import numpy as np

docs = [
    "Error E-4021: deployment pipeline timeout after 300s",
    "Resolving CI/CD failures in production environments",
    "How to configure deployment retry settings",
]

# --- Lexical retrieval with BM25 ---
tokenized = [d.lower().split() for d in docs]
bm25 = BM25Okapi(tokenized)
query = "fix error E-4021 deploying"
lex_scores = bm25.get_scores(query.lower().split())
print("BM25 scores:", lex_scores)  # high for doc with "E-4021"

# --- Dense retrieval with embeddings ---
model = SentenceTransformer("all-MiniLM-L6-v2")
doc_embs = model.encode(docs)
q_emb = model.encode([query])
dense_scores = np.dot(doc_embs, q_emb.T).flatten()
print("Dense scores:", dense_scores)  # high for semantically similar

def reciprocal_rank_fusion(result_lists, k=60):
    """Merge multiple ranked lists using RRF.

    Args:
        result_lists: list of lists of (doc_id, score) tuples
        k: smoothing constant (default 60)
    Returns:
        sorted list of (doc_id, rrf_score) tuples
    """
    scores = {}
    for results in result_lists:
        for rank, (doc_id, _) in enumerate(results):
            if doc_id not in scores:
                scores[doc_id] = 0.0
            # RRF formula: 1 / (k + rank)
            scores[doc_id] += 1.0 / (k + rank + 1)

    return sorted(scores.items(), key=lambda x: x[1], reverse=True)

# Usage: fuse BM25 and dense retrieval results
bm25_results = [("doc_1", 8.2), ("doc_3", 6.1), ("doc_5", 4.0)]
dense_results = [("doc_3", 0.92), ("doc_2", 0.88), ("doc_1", 0.85)]

fused = reciprocal_rank_fusion([bm25_results, dense_results])
print("Fused ranking:", fused[:5])

def chunk_text(tokens, chunk_size=400, overlap=60):
    """Split a token list into overlapping chunks.

    Overlap ensures information at chunk boundaries
    is captured by at least one chunk, improving recall.
    """
    chunks = []
    start = 0
    while start < len(tokens):
        end = min(start + chunk_size, len(tokens))
        chunks.append(tokens[start:end])
        if end == len(tokens):
            break
        # Step forward by (chunk_size - overlap)
        start = end - overlap
    return chunks

# Example: chunk a document of 1200 tokens
tokens = list(range(1200))  # simulated token IDs
chunks = chunk_text(tokens, chunk_size=400, overlap=60)
print(f"Created {len(chunks)} chunks")
for i, c in enumerate(chunks):
    print(f"  Chunk {i}: tokens {c[0]}-{c[-1]} ({len(c)} tokens)")
# Output: chunks with 60-token overlap at boundaries

# Metadata-filtered search with a vector database (Qdrant example)
from qdrant_client import QdrantClient
from qdrant_client.models import Filter, FieldCondition, MatchValue, Range

client = QdrantClient("localhost", port=6333)

# Search with metadata filters: right tenant, recent docs only
results = client.search(
    collection_name="knowledge_base",
    query_vector=query_embedding,
    query_filter=Filter(
        must=[
            # Only this tenant's documents
            FieldCondition(
                key="tenant_id",
                match=MatchValue(value="acme-corp")
            ),
            # Only documents updated in last 90 days
            FieldCondition(
                key="updated_at",
                range=Range(gte="2025-02-01")
            ),
            # Only English-language content
            FieldCondition(
                key="language",
                match=MatchValue(value="en")
            ),
        ]
    ),
    limit=10,
)
# Filters run BEFORE vector search, so all 10 results
# are guaranteed to match the constraints

# Example: indexing and searching with Chroma (lightweight)
import chromadb

# Create a local collection
client = chromadb.Client()
collection = client.create_collection(
    name="knowledge_base",
    metadata={"hnsw:space": "cosine"}  # cosine similarity
)

# Add documents with embeddings and metadata
collection.add(
    documents=[
        "Refunds must be requested within 14 days.",
        "Shipping is free for orders over $50.",
        "Contact support at support@example.com.",
    ],
    metadatas=[
        {"category": "refunds", "region": "US"},
        {"category": "shipping", "region": "US"},
        {"category": "support", "region": "global"},
    ],
    ids=["doc1", "doc2", "doc3"],
)

# Search with optional metadata filter
results = collection.query(
    query_texts=["How do I get a refund?"],
    n_results=2,
    where={"region": "US"},  # metadata filter
)
print(results["documents"])

import faiss
import numpy as np

# Generate 100K random vectors (simulating embeddings)
d = 384                       # embedding dimension
n = 100_000                   # number of documents
vectors = np.random.randn(n, d).astype('float32')

# --- Exact search (brute force) ---
exact_index = faiss.IndexFlatIP(d)  # inner product
exact_index.add(vectors)

# --- ANN search (HNSW) ---
hnsw_index = faiss.IndexHNSWFlat(d, 32)  # 32 neighbors per node
hnsw_index.hnsw.efConstruction = 200      # build quality
hnsw_index.hnsw.efSearch = 64             # query quality (tune this)
hnsw_index.add(vectors)

# Compare: query with 5 random vectors
queries = np.random.randn(5, d).astype('float32')
import time

t0 = time.time()
D_exact, I_exact = exact_index.search(queries, 10)
print(f"Exact:  {(time.time()-t0)*1000:.1f}ms")

t0 = time.time()
D_ann, I_ann = hnsw_index.search(queries, 10)
print(f"HNSW:   {(time.time()-t0)*1000:.1f}ms")

# Measure recall: how many ANN results match exact results?
recall = np.mean([
    len(set(I_exact[i]) & set(I_ann[i])) / 10.0
    for i in range(5)
])
print(f"Recall@10: {recall:.1%}")

from sentence_transformers import CrossEncoder

# Load a cross-encoder reranker
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

query = "What is the refund policy for international orders?"

# Candidates from first-stage retrieval (bi-encoder)
candidates = [
    "Shipping costs are non-refundable for all orders.",
    "International refunds take 10-15 business days.",
    "Our refund policy allows returns within 14 days.",
    "Contact support for order tracking information.",
    "International orders may incur customs duties.",
]

# Score each candidate against the query
pairs = [[query, doc] for doc in candidates]
scores = reranker.predict(pairs)

# Rerank by cross-encoder score
ranked = sorted(
    zip(candidates, scores),
    key=lambda x: x[1],
    reverse=True
)
for doc, score in ranked:
    print(f"  {score:>6.3f}  {doc}")
# "International refunds take 10-15 business days" ranks #1

import openai

def rewrite_query(user_query, client, chat_history=None):
    """Rewrite a user query for better retrieval."""
    context = ""
    if chat_history:
        context = "\n".join(
            f"{m['role']}: {m['content']}"
            for m in chat_history[-3:]
        )

    response = client.chat.completions.create(
        model="gpt-4o-mini",  # cheap and fast for rewrites
        messages=[{
            "role": "system",
            "content": """Rewrite the user's question to improve
document retrieval. Expand acronyms, add relevant
keywords, and resolve ambiguous references using
conversation context if provided. Output ONLY the
rewritten query, nothing else."""
        }, {
            "role": "user",
            "content": f"Context:\n{context}\n\nOriginal: {user_query}"
        }],
        temperature=0.0,
        max_tokens=100,
    )
    return response.choices[0].message.content.strip()

# Example: "fix the SSO thing on mobile"
# -> "Troubleshoot Single Sign-On authentication failures on mobile"

import numpy as np

def recall_at_k(retrieved_ids, relevant_ids, k):
    """Fraction of relevant docs found in top-k retrieved."""
    top_k = set(retrieved_ids[:k])
    relevant = set(relevant_ids)
    return len(top_k & relevant) / max(len(relevant), 1)

def mrr(retrieved_ids, relevant_ids):
    """Mean Reciprocal Rank: 1/position of first relevant result."""
    relevant = set(relevant_ids)
    for i, doc_id in enumerate(retrieved_ids):
        if doc_id in relevant:
            return 1.0 / (i + 1)
    return 0.0

def ndcg_at_k(retrieved_ids, relevant_ids, k):
    """Normalized Discounted Cumulative Gain at k."""
    relevant = set(relevant_ids)
    dcg = sum(
        (1.0 if retrieved_ids[i] in relevant else 0.0) / np.log2(i + 2)
        for i in range(min(k, len(retrieved_ids)))
    )
    ideal = sum(
        1.0 / np.log2(i + 2)
        for i in range(min(k, len(relevant_ids)))
    )
    return dcg / max(ideal, 1e-10)

# Example evaluation
retrieved = ["d3", "d7", "d1", "d5", "d2"]
relevant = ["d1", "d3"]

print(f"Recall@5: {recall_at_k(retrieved, relevant, 5):.2f}")
print(f"MRR:      {mrr(retrieved, relevant):.2f}")
print(f"nDCG@5:   {ndcg_at_k(retrieved, relevant, 5):.2f}")