Advanced RAG on Milvus: Hybrid Search, HyDE, Query Rewriting and Reranking

A proof-of-concept RAG system takes minutes: chunk a PDF, push the chunks through an embedding model, hand the matches to an LLM. Production is where that vector-search-only setup stalls. Real users don't phrase queries like tidy data engines. They write vague wording, leave out keywords and bundle several questions into one, and a single semantic lookup misses much of that. When retrieval fails, the LLM has to summarize irrelevant context, which yields subtle hallucinations or answers that don't help.

Closing the gap between basic semantic search and a multi-stage retrieval pipeline takes layers. The build below uses Milvus. It starts with dense vector lookups, adds hybrid search, then stacks four techniques (metadata filtering, query rewriting, HyDE and query decomposition) and finishes with cross-encoder reranking. Each step comes with working code.

Why RAG exists

Models like GPT-4 learn from general internet data up to a knowledge cutoff. Private documents, internal wikis and company policies are out of reach. Ask a raw model what the company's leave policy is and it either says it doesn't know or invents a plausible, wrong answer.

RAG (Retrieval-Augmented Generation) splits the job into two steps:

1. Retrieve the most relevant passages from your documents with semantic search.
2. Generate an answer by passing those passages to the LLM as context.

The overall flow looks like this:

User Question
     │
     ▼
[Embed the question]  →  query vector
     │
     ▼
[Search Vector DB]  →  top-K relevant document chunks
     │
     ▼
[LLM prompt: "Given these passages, answer: {question}"]
     │
     ▼
  Accurate, Grounded Answer

Final answer quality depends almost entirely on retrieval quality, so getting retrieval right is the whole game.

Where a vector database comes in

Conventional databases compare exact values, so a WHERE text LIKE '%leave%' query skips documents that say "vacation entitlement" or "medical absence" despite the same meaning. A vector database instead stores each document as a high-dimensional numeric vector (an embedding) and retrieves by semantic similarity, meaning how close the query's meaning is to each chunk's meaning. Milvus suits that role for these reasons:

• It's open source and can be self-hosted, which keeps data private.
• Milvus Lite runs from a single .db file, so development needs no server.
• Hybrid search is native, letting dense (semantic) and sparse (keyword) vectors combine in one query.
• It scales from a laptop to billions of vectors with no application code changes.

Dense and sparse embeddings

The pipeline uses two kinds of vectors, so it pays to know them before writing code. A dense embedding is a fixed-length vector of, for instance, 1024 numbers produced by a neural network. All dimensions carry values, and training teaches the model to put sentences with similar meaning close together:

"sick leave policy"           →  [0.12, -0.87, 0.34, 0.56, ...]  (1024 numbers)
"medical absence entitlement" →  [0.13, -0.85, 0.31, 0.54, ...]  ← very close
"quarterly revenue target"    →  [0.91,  0.23, -0.67, 0.02, ...] ← far away

Dense vectors capture meaning and synonyms well. Models such as OpenAI's text-embedding-3-large produce them. A sparse embedding is a huge vector where most entries are zero. Each non-zero position stands for one vocabulary word, and its value is that word's importance, a TF-IDF-style weight:

"sick leave policy" → {word_index_for_"sick": 0.82, word_index_for_"leave": 0.91, ...}

BM25 is the classic sparse algorithm and excels at exact keyword matching. It catches proper nouns, product names and technical terms that dense search can miss because they lack a close semantic neighbor.

Hybrid search uses both, taking the strengths of each approach.

Setup and configuration

Install the dependencies first:

pip install --upgrade pymilvus
pip install "pymilvus[model]"
pip install sentence-transformers
pip install langchain-text-splitters
pip install langchain-openai
pip install langchain-community
pip install scipy
pip install nltk

Then the imports:

import uuid
from tqdm import tqdm
from pymilvus import (
    MilvusClient, DataType,
    AnnSearchRequest, RRFRanker
)
from pymilvus.model.sparse import BM25EmbeddingFunction
from pymilvus.model.sparse.bm25.tokenizers import build_default_analyzer
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
import scipy.sparse as sp
import re, json
import nltk
nltk.download('stopwords')

Keep every tunable parameter together at the top of the file:

PDF_PATH        = "./data/sample_employee_handbook.pdf" # path of you document
COLLECTION_NAME = "rag_documents_hybrid"
MILVUS_DB_PATH  = "./db/milvus_demo.db"
API_KEY         = "sk-..."
EMBEDDING_MODEL = "text-embedding-3-large"
EMBEDDING_DIM   = 1024
CHUNK_SIZE      = 500
CHUNK_OVERLAP   = 100
TOP_K           = 5

On chunk size: 500 characters with a 100-character overlap suits simple documents as a starting point. Narrative text does better with larger chunks of 800 to 1000, and dense technical documentation with smaller ones of 200 to 300. Overlap makes sure sentences cut by a chunk boundary can still be found.

The indexing pipeline

Indexing runs once, or again when documents change. It turns raw documents into vectors and loads them into Milvus. The path is:

PDF  →  Pages  →  Chunks  →  Dense Embeddings
                           →  Sparse Embeddings
                           →  Milvus Collection

First initialize the models and connect:

# Dense embedding model here we'll be using OpenAI's embedding model
embedding_obj = OpenAIEmbeddings(
    model=EMBEDDING_MODEL,
    api_key=API_KEY,
    dimensions=EMBEDDING_DIM
)
# Milvus Lite - single file, no server needed
client = MilvusClient(MILVUS_DB_PATH)
print("Models and DB connection ready.")

Then load the PDF into one document object per page:

# Load PDF — one Document object per page
loader = PyPDFLoader(PDF_PATH)
documents = loader.load()
print(f"Loaded {len(documents)} pages.")

# Split into overlapping chunks
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=CHUNK_SIZE,
    chunk_overlap=CHUNK_OVERLAP,
    separators=["\n\n", "\n", ".", " ", ""]
)
chunks = text_splitter.split_documents(documents)
print(f"Created {len(chunks)} chunks.")

RecursiveCharacterTextSplitter splits on paragraph breaks first, then newlines, then sentences, which keeps related sentences together in a chunk as far as possible. Next come the embeddings for both vector types:

texts = [doc.page_content for doc in chunks]

# Dense embeddings - one API call for the entire corpus
print("Generating dense embeddings...")
dense_embeddings = embedding_obj.embed_documents(texts)
print(f"Dense dimension: {len(dense_embeddings[0])}")

# Sparse embeddings - BM25 must be fit on YOUR corpus first
print("Fitting BM25 on corpus...")
analyzer = build_default_analyzer(language="en") # for this you will require nltk-stopwords
bm25_ef = BM25EmbeddingFunction(analyzer)
bm25_ef.fit(texts)  # Builds vocabulary from your documents
sparse_embeddings = bm25_ef.encode_documents(texts)
print("Sparse embeddings generated.")

The bm25_ef.fit(texts) call is critical. BM25 learns vocabulary and term frequencies from your own corpus, and that's why it works. A generic BM25 model in its place drops keyword matching quality noticeably. Now the collection, with its schema and indexes:

# Drop and recreate for a clean state
if COLLECTION_NAME in client.list_collections():
    client.drop_collection(COLLECTION_NAME)

# Define schema
schema = client.create_schema()
schema.add_field("id",            DataType.VARCHAR,           is_primary=True, max_length=100)
schema.add_field("vector",        DataType.FLOAT_VECTOR,      dim=EMBEDDING_DIM)
schema.add_field("sparse_vector", DataType.SPARSE_FLOAT_VECTOR)
schema.add_field("text",          DataType.VARCHAR,           max_length=65535)
schema.add_field("page_number",   DataType.INT64)
schema.add_field("source",        DataType.VARCHAR,           max_length=500)
schema.add_field("chunk_id",      DataType.INT64)

# Create the collection
client.create_collection(collection_name=COLLECTION_NAME, schema=schema)

# Build indexes separately
index_params = client.prepare_index_params()

index_params.add_index(
    field_name="vector",
    index_type="FLAT",          # Exact search - swap to HNSW for production
    metric_type="COSINE"
)

index_params.add_index(
    field_name="sparse_vector",
    index_type="SPARSE_INVERTED_INDEX",
    metric_type="IP"            # Inner Product is the only valid metric for sparse
)

client.create_index(collection_name=COLLECTION_NAME, index_params=index_params)

print("Collection and indexes created.")

Two index choices matter here:

• FLAT for dense vectors. It performs exact nearest-neighbor search, fully accurate but slower at scale. It's fine in development and for datasets below roughly 100K chunks, and HNSW should replace it in production.
• SPARSE_INVERTED_INDEX for sparse vectors. It's the only valid index type for sparse data in Milvus and works like a classic inverted index, the way search engines index documents by word.

Finally, prepare records and insert them:

def sparse_to_dict(s_emb) -> dict:
    """Convert a scipy sparse row into Milvus-compatible {index: value} dict."""
    if sp.issparse(s_emb):
        coo = s_emb.tocoo()
        return {int(col): float(val) for col, val in zip(coo.col, coo.data)}
    elif isinstance(s_emb, dict):
        return s_emb
    else:
        return {int(i): float(v) for i, v in enumerate(s_emb) if v != 0.0}

# Build the records list
data = []
for idx, (chunk, d_emb) in enumerate(tqdm(zip(chunks, dense_embeddings), total=len(chunks))):
    sparse_dict = sparse_to_dict(sparse_embeddings[idx])
    if not sparse_dict:
        print(f"Warning: empty sparse vector at chunk {idx}, skipping.")
        continue
    data.append({
        "id":            str(uuid.uuid4()),
        "vector":        d_emb,
        "sparse_vector": sparse_dict,
        "text":          chunk.page_content,
        "page_number":   int(chunk.metadata.get("page", -1)),
        "source":        PDF_PATH,
        "chunk_id":      idx
    })

# Insert into Milvus
res = client.insert(collection_name=COLLECTION_NAME, data=data)
print(f"Inserted {res['insert_count']} records.")

# Load into memory - required before any search operation
client.load_collection(COLLECTION_NAME)
print(f"Load state: {client.get_load_state(COLLECTION_NAME)}")

The sparse_to_dict helper is needed because BM25 produces a scipy sparse matrix, yet Milvus wants a plain Python {int: float} dictionary. It copes with 2D sparse rows, 1D arrays and dictionaries that are already converted. One common gotcha: always call client.load_collection() before searching. Indexes live on disk in Milvus and are loaded into RAM on demand, so without that call a search errors out or comes back empty.

Dense search as the baseline

With data indexed, the simplest search embeds the query and finds the nearest document vectors by cosine similarity:

# ════════════════════════════════════════════════════════════
# Dense Vector Search
# ════════════════════════════════════════════════════════════

query = "What is the leave policy?"

# Step 1: Embed the query using the same model used at index time
query_dense_embedding = embedding_obj.embed_query(query)

# Step 2: Search
results = client.search(
    collection_name=COLLECTION_NAME,
    data=[query_dense_embedding],
    anns_field="vector",
    search_param={"metric_type": "COSINE"},
    limit=TOP_K,
    output_fields=["text", "page_number", "source"]
)

# Step 3: Display results
for idx, hit in enumerate(results[0], start=1):
    entity = hit["entity"]
    print(f"Rank {idx} | Cosine Score: {hit['distance']:.4f} | Page: {entity['page_number']}")
    print(f"  {entity['text'][:300]}\n")

Under the COSINE metric, scores run from -1 to 1. Above 0.7 usually means a strong semantic match, and top results typically land between 0.6 and 0.95. Dense search suffices for natural-language questions whose answers sit in nearby semantic space, like asking to explain parental leave benefits or what the code of conduct says. It struggles when a query holds specific codes, acronyms or proper nouns that the embedding model rarely saw.

Hybrid search with RRF

Hybrid search executes a dense search and a sparse BM25 search at once, then fuses both ranked lists through Reciprocal Rank Fusion (RRF). The idea is simple. A document's RRF score adds up 1 / (k + rank) from each list, with k=60 as a smoothing constant, so documents ranked well in both lists end up at the top.

# ════════════════════════════════════════════════════════════
# Hybrid Search (Dense + Sparse)
# ════════════════════════════════════════════════════════════

query = "leave policy?"

# Dense query vector
query_dense = embedding_obj.embed_query(query)

# Sparse query vector - uses the same BM25 model fitted on the corpus
sparse_raw  = bm25_ef.encode_queries([query])
sparse_dict = sparse_to_dict(sparse_raw[0])
print(f"Sparse query terms: {len(sparse_dict)}")  # Should be > 0

# Build two separate ANN search requests
dense_req = AnnSearchRequest(
    data=[query_dense],
    anns_field="vector",
    param={"metric_type": "COSINE"},
    limit=TOP_K
)
sparse_req = AnnSearchRequest(
    data=[sparse_dict],
    anns_field="sparse_vector",
    param={"metric_type": "IP"},
    limit=TOP_K
)

# Execute hybrid search with RRF fusion
results = client.hybrid_search(
    collection_name=COLLECTION_NAME,
    reqs=[dense_req, sparse_req],
    ranker=RRFRanker(k=60),
    limit=TOP_K,
    output_fields=["text", "page_number", "source"]
)
for idx, hit in enumerate(results[0], start=1):
    entity = hit["entity"]
    print(f"Rank {idx} | RRF Score: {hit['distance']:.4f} | Page: {entity['page_number']}")
    print(f"  {entity['text'][:300]}\n")

About scores: the hit['distance'] value from hybrid_search is the fused RRF ranking score and not raw cosine similarity. Larger is better, though the range isn't limited to 0 through 1, which means it shouldn't be set against cosine scores from a regular search(). For production RAG, hybrid search is the recommended default. BM25 adds barely any latency and visibly lifts recall on queries heavy with keywords.

Even hybrid retrieval can fail in particular situations, and four techniques address the most common failure modes.

Metadata filtering

The problem. When a collection spans several documents, periods or departments, vector search may return relevant-sounding chunks from the wrong source, such as an outdated policy version, another department's handbook or pages outside the question's scope.

How it works. Milvus applies a scalar filter expression ahead of the vector search, so only rows that match reach the ANN stage. It's no post-filter. The filter trims the candidate pool first, which preserves recall within the filtered subset.

# ════════════════════════════════════════════════════════════
# METADATA FILTERING
# ════════════════════════════════════════════════════════════

def search_with_metadata_filter(
    client, collection_name, embedding_obj, bm25_ef,
    query: str,
    page_range: tuple = None,
    source_file: str = None,
    top_k: int = 5
):
    filter_parts = []

    if page_range:
        lo, hi = page_range
        filter_parts.append(f"page_number >= {lo} && page_number <= {hi}")

    if source_file:
        filter_parts.append(f'source == "{source_file}"')

    filter_expr = " && ".join(filter_parts) if filter_parts else None

    print(f"\n[Metadata Filter] Query : '{query}'")
    print(f"[Metadata Filter] Filter: {filter_expr or 'None (unfiltered)'}")

    results = hybrid_search(
        client, collection_name, embedding_obj, bm25_ef,
        query_text=query,
        top_k=top_k,
        filters=filter_expr
    )
    return results


# ── Run ──────────────────────────────────────────────────────
meta_results = search_with_metadata_filter(
    client, COLLECTION_NAME, embedding_obj, bm25_ef,
    query      = "What is the leave policy?",
    page_range = (1, 30),
    source_file= None,
    top_k      = 5
)

# ── Print Results ─────────────────────────────────────────────
print("\nMETADATA-FILTERED RESULTS")
print("=" * 55)
if not meta_results or not meta_results[0]:
    print("No results returned.")
else:
    for idx, hit in enumerate(meta_results[0], start=1):
        entity = hit["entity"]
        print(f"\nRank  : {idx}")
        print(f"Score : {hit['distance']:.4f}")
        print(f"Page  : {entity['page_number']}")
        print(f"Text  :\n{entity['text'][:400]}")

Filter expressions follow this syntax:

'page_number >= 1 && page_number <= 30'      # page range
'source == "hr_policy.pdf"'                  # exact source
'category in ["leave", "performance"]'       # in a list
'source like "hr%"'                          # prefix match

Any multi-document RAG system should include metadata filtering. If users work within a scope, such as only the 2024 handbook or only finance files, derive the filter from that scope automatically.

Query rewriting

The problem. Users write casually, as in "how many days off do I get?", whereas documents use formal wording like "annual leave entitlement" or "vacation days accrual policy." Dense search bridges part of that gap, and a good rewrite lifts sparse retrieval and overall recall considerably.

How it works. An LLM drafts several formal, keyword-heavy rewrites of the original query, each variant gets its own search, and RRF merges and deduplicates what comes back.

# ════════════════════════════════════════════════════════════
# QUERY REWRITING
# ════════════════════════════════════════════════════════════

import re, json

REWRITE_PROMPT = """You are an expert at reformulating search queries to improve document retrieval.

Given a user query, produce {n} alternative search queries that:
- Use formal, document-style language
- Include relevant keywords and synonyms
- Cover different angles of the same question

User query: {query}

Respond ONLY with a JSON array of strings. Example:
["rewritten query 1", "rewritten query 2", "rewritten query 3"]"""


def rewrite_query(query: str, n: int = 3) -> list[str]:
    prompt   = REWRITE_PROMPT.format(query=query, n=n)
    response = llm.invoke(prompt)
    raw      = re.sub(r"^```json|^```|```quot;, "", response.content.strip(), flags=re.MULTILINE).strip()
    try:
        variants = json.loads(raw)
        return [query] + variants       # always keep the original
    except json.JSONDecodeError:
        print("Warning: Could not parse rewrites, using original query only.")
        return [query]


def search_with_query_rewriting(
    client, collection_name, embedding_obj, bm25_ef,
    query: str,
    n_rewrites: int = 3,
    top_k: int = 5
):
    variants = rewrite_query(query, n=n_rewrites)

    print(f"\n[Query Rewriting] Original  : '{query}'")
    for i, v in enumerate(variants[1:], 1):
        print(f"[Query Rewriting] Variant {i} : '{v}'")

    seen_ids    = {}
    rank_scores = {}

    for variant in variants:
        results = hybrid_search(
            client, collection_name, embedding_obj, bm25_ef,
            query_text=variant,
            top_k=top_k
        )
        if not results or not results[0]:
            continue
        for rank, hit in enumerate(results[0], start=1):
            hit_id = hit["id"]
            rank_scores[hit_id] = rank_scores.get(hit_id, 0) + 1.0 / (60 + rank)
            if hit_id not in seen_ids:
                seen_ids[hit_id] = hit

    merged = sorted(seen_ids.values(), key=lambda h: rank_scores[h["id"]], reverse=True)[:top_k]
    return [merged]


# ── Run ──────────────────────────────────────────────────────
rewrite_results = search_with_query_rewriting(
    client, COLLECTION_NAME, embedding_obj, bm25_ef,
    query      = "What is the leave policy?",
    n_rewrites = 3,
    top_k      = 5
)

# ── Print Results ─────────────────────────────────────────────
print("\nQUERY-REWRITTEN RESULTS")
print("=" * 55)
if not rewrite_results or not rewrite_results[0]:
    print("No results returned.")
else:
    for idx, hit in enumerate(rewrite_results[0], start=1):
        entity = hit["entity"]
        print(f"\nRank  : {idx}")
        print(f"Score : {hit['distance']:.4f}")
        print(f"Page  : {entity['page_number']}")
        print(f"Text  :\n{entity['text'][:400]}")

An example transformation:

Input:  "how many days off do I get?"

Output variants:
  1. "annual leave entitlement number of days employee handbook"
  2. "vacation days accrual policy full-time employee"
  3. "paid time off PTO allowance per calendar year"

Every variant widens the retrieval surface. A chunk outside the top 5 for the original query might rank second for the third variant, and RRF surfaces documents that show up consistently across queries.

HyDE (Hypothetical Document Embeddings)

The problem. A short question and its answer sit in different spots in embedding space. "How many sick days?" is a question and embeds near other questions. A policy line saying employees get 7 days of paid sick leave a year embeds near other policy statements. That gap reduces precision.

How it works. Instead of embedding the question, the LLM writes a hypothetical answer document, and that gets embedded. The made-up answer lands near real answer documents even if its specific facts are wrong. Only its embedding is used and its text is thrown away, so the invented details never reach the user.

# ════════════════════════════════════════════════════════════
# HyDE (Hypothetical Document Embeddings)
# ════════════════════════════════════════════════════════════

HYDE_PROMPT = """You are a corporate policy document writer.

Write a 2-3 paragraph excerpt from an official HR policy or company document
that would DIRECTLY ANSWER the following question.
Write in formal document style. Do not mention the question itself.

Question: {query}

Document excerpt:"""


def generate_hypothetical_document(query: str) -> str:
    response = llm.invoke(HYDE_PROMPT.format(query=query))
    return response.content.strip()


def search_with_hyde(
    client, collection_name, embedding_obj, bm25_ef,
    query: str,
    top_k: int = 5
):
    hypothetical_doc = generate_hypothetical_document(query)

    print(f"\n[HyDE] Query            : '{query}'")
    print(f"[HyDE] Hypothetical doc :\n  {hypothetical_doc[:300]}...\n")

    # Search using the hypothetical document's embedding
    hyde_results = hybrid_search(
        client, collection_name, embedding_obj, bm25_ef,
        query_text=hypothetical_doc,    # embed the answer, not the question
        top_k=top_k
    )

    # Also search with the original query and merge both via RRF
    original_results = hybrid_search(
        client, collection_name, embedding_obj, bm25_ef,
        query_text=query,
        top_k=top_k
    )

    seen_ids    = {}
    rank_scores = {}

    for result_set in [hyde_results, original_results]:
        if not result_set or not result_set[0]:
            continue
        for rank, hit in enumerate(result_set[0], start=1):
            hit_id = hit["id"]
            rank_scores[hit_id] = rank_scores.get(hit_id, 0) + 1.0 / (60 + rank)
            if hit_id not in seen_ids:
                seen_ids[hit_id] = hit

    merged = sorted(seen_ids.values(), key=lambda h: rank_scores[h["id"]], reverse=True)[:top_k]
    return [merged]


# ── Run ──────────────────────────────────────────────────────
hyde_results = search_with_hyde(
    client, COLLECTION_NAME, embedding_obj, bm25_ef,
    query = "What is the leave policy?",
    top_k = 5
)

# ── Print Results ─────────────────────────────────────────────
print("\nHyDE RESULTS")
print("=" * 55)
if not hyde_results or not hyde_results[0]:
    print("No results returned.")
else:
    for idx, hit in enumerate(hyde_results[0], start=1):
        entity = hit["entity"]
        print(f"\nRank  : {idx}")
        print(f"Score : {hit['distance']:.4f}")
        print(f"Page  : {entity['page_number']}")
        print(f"Text  :\n{entity['text'][:400]}")

HyDE helps most with short, vague queries such as a salary review process, with vocabulary gaps between questions and documents, and with very informal phrasing. It's overkill for precise, keyword-rich queries where hybrid search already does well, and when LLM latency matters, since HyDE adds one extra LLM call per query.

Query decomposition

The problem. A complex, multi-part question can't be answered from one passage. Asking about the leave policy and how performance review affects salary touches two topics in two sections. A single vector search either fixates on one topic and misses the other, or returns mediocre results for both.

How it works. An LLM splits the question into atomic sub-questions, each is searched alone, RRF merges the results, and the final LLM call gets all the relevant chunks.

# ════════════════════════════════════════════════════════════
# QUERY DECOMPOSITION
# ════════════════════════════════════════════════════════════

DECOMPOSE_PROMPT = """You are an expert at breaking down complex questions for document retrieval.

Decompose the following question into 2-4 simple, self-contained sub-questions.
Each sub-question should target a single distinct piece of information.

Complex question: {query}

Respond ONLY with a JSON array of strings. Example:
["sub-question 1", "sub-question 2", "sub-question 3"]"""


def decompose_query(query: str) -> list[str]:
    response = llm.invoke(DECOMPOSE_PROMPT.format(query=query))
    raw      = re.sub(r"^```json|^```|```quot;, "", response.content.strip(), flags=re.MULTILINE).strip()
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        print("Warning: Could not parse decomposition, using original query.")
        return [query]


def search_with_decomposition(
    client, collection_name, embedding_obj, bm25_ef,
    query: str,
    top_k: int = 5
):
    sub_questions = decompose_query(query)

    print(f"\n[Decomposition] Original query : '{query}'")
    for i, sq in enumerate(sub_questions, 1):
        print(f"[Decomposition] Sub-question {i} : '{sq}'")

    per_subquery_results = {}
    seen_ids             = {}
    rank_scores          = {}

    for sq in sub_questions:
        results = hybrid_search(
            client, collection_name, embedding_obj, bm25_ef,
            query_text=sq,
            top_k=top_k
        )
        per_subquery_results[sq] = results

        if not results or not results[0]:
            continue
        for rank, hit in enumerate(results[0], start=1):
            hit_id = hit["id"]
            rank_scores[hit_id] = rank_scores.get(hit_id, 0) + 1.0 / (60 + rank)
            if hit_id not in seen_ids:
                seen_ids[hit_id] = hit

    merged = sorted(seen_ids.values(), key=lambda h: rank_scores[h["id"]], reverse=True)[:top_k]

    # Per sub-question breakdown
    print("\n── Per Sub-question Results ──")
    for sq, res in per_subquery_results.items():
        print(f"\n  SUB-QUERY: '{sq[:60]}'")
        if res and res[0]:
            for i, hit in enumerate(res[0], start=1):
                print(f"    {i}. Page {hit['entity']['page_number']} | Score {hit['distance']:.4f} | {hit['entity']['text'][:150]}")

    return {"per_subquery": per_subquery_results, "merged": [merged]}


# ── Run ──────────────────────────────────────────────────────
decomp_results = search_with_decomposition(
    client, COLLECTION_NAME, embedding_obj, bm25_ef,
    query = "What is the leave policy and how does it affect salary deductions?",
    top_k = 5
)

# ── Print Merged Results ──────────────────────────────────────
print("\nDECOMPOSED — MERGED FINAL RESULTS")
print("=" * 55)
merged_hits = decomp_results["merged"]
if not merged_hits or not merged_hits[0]:
    print("No results returned.")
else:
    for idx, hit in enumerate(merged_hits[0], start=1):
        entity = hit["entity"]
        print(f"\nRank  : {idx}")
        print(f"Score : {hit['distance']:.4f}")
        print(f"Page  : {entity['page_number']}")
        print(f"Text  :\n{entity['text'][:400]}")

An example decomposition:

Input:  "What is the leave policy and how does performance review affect salary?"
Sub-questions:
  1. "What is the annual leave policy?"
  2. "How many sick days are employees entitled to?"
  3. "How does performance review affect salary?"
  4. "What is the performance review schedule?"

Each sub-question pulls the passages it needs on its own, and the merged set covers every part of the original question, which no single search could do.

Cross-encoder reranking

Everything so far retrieves with bi-encoders: query and documents are embedded independently, and similarity comes from a dot product or cosine distance. They're quick but not the sharpest way to score relevance. Cross-encoders read a query and one document together as a pair and score how well they match, seeing their full interaction, which makes them far more accurate. They cost speed, being too slow for thousands of documents yet fast over a small candidate pool of 20 to 50. Hence the pattern: retrieve wide with a bi-encoder, rerank narrow with a cross-encoder.

# ============================================================
# RERANKING WITH CROSS-ENCODER
# ============================================================

from sentence_transformers import CrossEncoder

# Huggingface: cross-encoder/ms-marco-MiniLM-L12-v2
cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L12-v2")
query       = "What is the leave policy?"
RETRIEVAL_K = 20   # fetch more than you need
FINAL_K     = 5    # rerank down to this

# Step 1: Broad retrieval - fetch 20 candidates
query_dense = embedding_obj.embed_query(query)

results = client.search(
    collection_name=COLLECTION_NAME,
    data=[query_dense],
    anns_field="vector",
    search_param={"metric_type": "COSINE"},
    limit=RETRIEVAL_K,
    output_fields=["text", "page_number", "source"]
)

hits = results[0]
print(f"Retrieved {len(hits)} candidates for reranking.")

# Step 2: Score each (query, chunk) pair with the cross-encoder
pairs        = [[query, hit["entity"]["text"]] for hit in hits]
rerank_scores = cross_encoder.predict(pairs)

# Step 3: Sort by cross-encoder score
for hit, score in zip(hits, rerank_scores):
    hit["rerank_score"] = float(score)
reranked = sorted(hits, key=lambda x: x["rerank_score"], reverse=True)[:FINAL_K]

# Step 4: Display
for idx, hit in enumerate(reranked, start=1):
    entity = hit["entity"]
    print(f"Rank {idx} | Rerank: {hit['rerank_score']:.4f} | Vector: {hit['distance']:.4f}")
    print(f"  Page {entity['page_number']}: {entity['text'][:300]}\n")

Cross-encoder scores are unnormalized logits, not probabilities. A higher value signals stronger relevance, so 8.0 beats 2.0 by a wide margin, but the absolute scale differs between models. The model used, ms-marco-MiniLM-L12-v2, is a strong, light choice trained on the MS MARCO passage ranking dataset, and it deals well with question-style queries over document passages.

How the techniques compare

A sensible production default stack is:

1. Hybrid search as the always-on base retrieval method.
2. Metadata filtering any time the user's scope can be worked out.
3. Cross-encoder reranking over the top 20 candidates before they reach the LLM.
4. Query rewriting or HyDE added for domains with a wide vocabulary gap between user language and document language.

The progression runs like this:

• Dense search is the starting point: fast, free of keyword dependence and good with synonyms.
• Hybrid search is the upgrade, since adding BM25 sparse vectors reliably improves recall at little extra complexity.
• Filtering, rewriting, HyDE and decomposition each target one failure mode. Don't apply them all blindly. Choose by the documents and the query patterns users actually produce.
• Cross-encoder reranking is nearly always worth adding. It runs locally, adds little latency over a small candidate set and reliably improves the final context.

In the end retrieval quality decides answer quality. An LLM can only be as good as the context it's given, and each technique above exists to make that context more relevant, complete and dependable. Everything runs locally on Milvus Lite, with the OpenAI API as the sole external dependency for embeddings and generation.