Retrievers

from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.vectorstores import FAISS
from langchain_community.retrievers import BM25Retriever
from langchain.retrievers import EnsembleRetriever, ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain_core.documents import Document

# 1) The contract — every retriever exposes a single method
class Retriever:
    def get_relevant_documents(self, query: str) -> list[Document]: ...
    async def aget_relevant_documents(self, query: str) -> list[Document]: ...

# 2) Build a small corpus
docs = [
    Document(page_content='Kubernetes uses ReplicaSets to keep N pods running.', metadata={'src': 'k8s.md'}),
    Document(page_content='Postgres MVCC means writers never block readers.',     metadata={'src': 'pg.md'}),
    Document(page_content='Redis is single-threaded; long-running commands block.', metadata={'src': 'redis.md'}),
    Document(page_content='Django ORM uses lazy querysets — evaluated on iteration.', metadata={'src': 'django.md'}),
    # ... pretend there are 10,000 of these
]

# 3) Vector retriever — semantic similarity
emb = OpenAIEmbeddings(model='text-embedding-3-small')
vs  = FAISS.from_documents(docs, emb)
vector_retriever = vs.as_retriever(search_kwargs={'k': 4})

vector_retriever.invoke('How does Postgres handle concurrent writes?')

# 4) Keyword retriever — BM25, no embeddings, no GPU
bm25 = BM25Retriever.from_documents(docs)
bm25.k = 4
bm25.invoke('replicaset pod')

# 5) MMR — diverse results, not just nearest neighbours
mmr = vs.as_retriever(
    search_type='mmr',
    search_kwargs={'k': 4, 'fetch_k': 20, 'lambda_mult': 0.5},
)
# fetch_k candidates → re-rank for diversity → return k.
# lambda_mult 1.0 = pure relevance, 0.0 = pure diversity.

# 6) Threshold filter — drop weak matches
threshold = vs.as_retriever(
    search_type='similarity_score_threshold',
    search_kwargs={'score_threshold': 0.78, 'k': 4},
)
# Returns at most k docs whose similarity >= threshold; empty list is allowed.

# 7) Metadata filter — combine semantic + structured
meta_filter = vs.as_retriever(
    search_kwargs={'k': 4, 'filter': {'src': 'pg.md'}},
)
# Pre-filter the vector index by metadata before ANN search. Saves cost and improves relevance.

# 8) Ensemble (hybrid) — combine vector + BM25
ensemble = EnsembleRetriever(retrievers=[vector_retriever, bm25], weights=[0.7, 0.3])
ensemble.invoke('does Redis block?')
# Hybrid retrieval reliably beats either alone for production search.

# 9) Contextual compression — re-rank or summarise with an LLM after fetch
llm = ChatOpenAI(model='gpt-4o-mini', temperature=0)
extractor = LLMChainExtractor.from_llm(llm)
compressed = ContextualCompressionRetriever(
    base_compressor=extractor,
    base_retriever=ensemble,
)
compressed.invoke('How does Postgres handle concurrent writes?')
# Returns only the SENTENCES from each retrieved doc that actually answer the question.
# Slower (extra LLM call), much higher signal-to-noise in the prompt.

# 10) Cross-encoder reranker — Cohere, Voyage, or a local bge-reranker
from langchain.retrievers.document_compressors import CrossEncoderReranker
from langchain_community.cross_encoders import HuggingFaceCrossEncoder

ce = HuggingFaceCrossEncoder(model_name='BAAI/bge-reranker-v2-m3')
reranked = ContextualCompressionRetriever(
    base_compressor=CrossEncoderReranker(model=ce, top_n=4),
    base_retriever=ensemble.with_config(search_kwargs={'k': 20}),
)
# Strategy: fetch 20 candidates from cheap retrievers, rerank to 4 with a strong model.
# Big quality gains; tiny extra latency.

# 11) MultiQueryRetriever — let the LLM generate query variations
from langchain.retrievers import MultiQueryRetriever
mq = MultiQueryRetriever.from_llm(retriever=vector_retriever, llm=llm)
mq.invoke('How does Postgres handle simultaneous writers?')
# Generates 3 paraphrases, retrieves for each, dedupes — catches phrasing mismatches.

# 12) ParentDocumentRetriever — small chunks for matching, large chunks for context
from langchain_community.storage import InMemoryStore
from langchain.retrievers import ParentDocumentRetriever
from langchain_text_splitters import RecursiveCharacterTextSplitter

parent = ParentDocumentRetriever(
    vectorstore=FAISS.from_documents([], emb),
    docstore=InMemoryStore(),
    child_splitter=RecursiveCharacterTextSplitter(chunk_size=400, chunk_overlap=40),
    parent_splitter=RecursiveCharacterTextSplitter(chunk_size=2000),
)
parent.add_documents(docs)
# Search uses tight 400-char windows; downstream gets the 2000-char parent for context.

# 13) Self-query retriever — LLM writes its own metadata filter
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain.chains.query_constructor.base import AttributeInfo

field_info = [
    AttributeInfo(name='src', description='Source file', type='string'),
    AttributeInfo(name='year', description='Year of publication', type='integer'),
]
self_q = SelfQueryRetriever.from_llm(
    llm=llm,
    vectorstore=vs,
    document_contents='technical docs',
    metadata_field_info=field_info,
)
self_q.invoke('Postgres concurrency, from 2023')   # generates filter: { year: 2023 } automatically

# 14) Build it into a RAG chain
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

prompt = ChatPromptTemplate.from_messages([
    ('system', 'Answer based ONLY on the context.\nIf the context does not contain the answer, say so.\n\nContext:\n{context}'),
    ('user', '{question}'),
])

def fmt(docs):
    return '\n\n'.join(f'[{i+1}] {d.page_content}' for i, d in enumerate(docs))

rag = (
    {'context': reranked | fmt, 'question': RunnablePassthrough()}
    | prompt | llm | StrOutputParser()
)
rag.invoke('How does Postgres handle concurrent writes?')

# 15) Evaluation
#   Metrics:
#     • Hit@K — does the gold doc appear in the top K?
#     • MRR / nDCG — ranking quality
#     • Latency p50 / p95 — what users feel
#     • End-to-end answer quality with eval LLM (ragas, evals frameworks)
#   Always evaluate against a curated query set, not synthetic data.

# 16) Common bugs
#   • k=4 default is often too small; production tuned values are 8-20
#   • Pure vector retrieval misses exact-keyword queries (model numbers, error codes)
#   • No metadata filter → retrieving last year's deprecated docs
#   • Chunk size too small → answers fragmented; too big → retrieval mixes topics
#   • Forgetting to dedupe across retrievers in an ensemble → repeated context wastes prompt budget
#   • Reranker fed only 4 results → no headroom to improve ordering

retriever = store.as_retriever(
    search_type='mmr', search_kwargs={'k': 4, 'fetch_k': 20},
)
hits = retriever.invoke('return policy')