ssearch/query_hybrid.py

# query_hybrid.py
# Hybrid retrieval: BM25 (sparse) + vector similarity (dense) + cross-encoder
#
# Combines two retrieval strategies to catch both exact term matches and
# semantic similarity:
#   1. Retrieve top-20 via vector similarity (bi-encoder, catches meaning)
#   2. Retrieve top-20 via BM25 (term frequency, catches exact names/dates)
#   3. Merge and deduplicate candidates by node ID
#   4. Re-rank the union with a cross-encoder -> top-15
#   5. Pass re-ranked chunks to LLM for synthesis
#
# The cross-encoder doesn't care where candidates came from -- it scores
# each (query, chunk) pair on its own merits. BM25's job is just to
# nominate candidates that vector similarity might miss.
#
# E.M.F. February 2026

# Environment vars must be set before importing huggingface/transformers
# libraries, because huggingface_hub.constants evaluates HF_HUB_OFFLINE
# at import time.
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["HF_HUB_CACHE"] = "./models"  # HF hub cache for offline model loads
os.environ["HF_HUB_OFFLINE"] = "1"

from llama_index.core import (
    StorageContext,
    load_index_from_storage,
    Settings,
    get_response_synthesizer,
)
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.ollama import Ollama
from llama_index.core.prompts import PromptTemplate
from llama_index.core.postprocessor import SentenceTransformerRerank
from llama_index.core.schema import MetadataMode
from llama_index.retrievers.bm25 import BM25Retriever
import sys

#
# Globals
#

# Embedding model (must match build_store.py)
EMBED_MODEL = HuggingFaceEmbedding(cache_folder="./models", model_name="BAAI/bge-large-en-v1.5", local_files_only=True)

# LLM model for generation.
#LLM_MODEL = "command-r7b"
LLM_MODEL = "gemma4:e4b"

# Temperature. Gemma's own guidance is 1.0, but that is for open-ended chat.
# This is grounded extraction -- the prompt says to answer only from context --
# so a low temperature is the safer default. 1.0 previously; raise it back if
# answers come out flat.
LLM_TEMPERATURE = 0.3

# Ollama num_ctx. The prompt itself is only ~2500 tokens with 15 chunks, but
# thinking=True spends reasoning tokens out of this same budget, so leave room.
LLM_CONTEXT_WINDOW = 16384

# Cross-encoder model for re-ranking (cached in ./models/)
# Was cross-encoder/ms-marco-MiniLM-L-12-v2 (2021, trained on MS MARCO web
# queries). It scored journal-style queries far out of its calibrated range --
# top hits at -0.77 and -1.90 logits -- because "what would the author say
# about..." is nothing like a search-engine query. bge-reranker-v2-m3 returns
# calibrated 0-1 scores and handles the phrasing. Costs ~5s more per query.
RERANK_MODEL = "BAAI/bge-reranker-v2-m3"
RERANK_TOP_N = 15

# Retrieval parameters
VECTOR_TOP_K = 20   # candidates from vector similarity
BM25_TOP_K = 20     # candidates from BM25 term matching

#
# Custom prompt -- same as v3
#
PROMPT = PromptTemplate(
"""You are a precise research assistant analyzing excerpts from a personal journal collection.
Every excerpt below has been selected and ranked for relevance to the query.

CONTEXT (ranked by relevance):
{context_str}

QUERY:
{query_str}

Instructions:
- Answer ONLY using information explicitly present in the CONTEXT above
- Examine ALL provided excerpts, not just the top few -- each one was selected for relevance
- Be specific: quote or closely paraphrase key passages and cite their file names
- When multiple files touch on the query, note what each one contributes
- If the context doesn't contain enough information to answer fully, say so

Your response should:
1. Directly answer the query, drawing on as many relevant excerpts as possible
2. Reference specific files and their content (e.g., "In <filename>, ...")
3. End with a list of all files that contributed to your answer, with a brief note on each

If the context is insufficient, explain what's missing."""
)


def main():
    # Configure LLM and embedding model
    # for local model using ollama
    # Note: Ollama temperature defaults to 0.8
    Settings.llm = Ollama(
        model=LLM_MODEL,
        temperature=LLM_TEMPERATURE,
        thinking=True,          # enable
        request_timeout=360.0,
        context_window=LLM_CONTEXT_WINDOW,
    )
    
    # Use OpenAI API:
    # from llama_index.llms.openai import OpenAI
    # Settings.llm = OpenAI(
    #     model="gpt-4o-mini",   # or "gpt-4o" for higher quality
    #     temperature=0.3,
    # )

    Settings.embed_model = EMBED_MODEL


    # Load persisted vector store
    storage_context = StorageContext.from_defaults(persist_dir="./store")
    index = load_index_from_storage(storage_context)

    # --- Retrievers ---

    # Vector retriever (dense: cosine similarity over embeddings)
    vector_retriever = index.as_retriever(similarity_top_k=VECTOR_TOP_K)

    # BM25 retriever (sparse: term frequency scoring)
    bm25_retriever = BM25Retriever.from_defaults(
        index=index,
        similarity_top_k=BM25_TOP_K,
    )

    # Cross-encoder re-ranker
    reranker = SentenceTransformerRerank(
        model=RERANK_MODEL,
        top_n=RERANK_TOP_N,
    )

    # --- Query ---

    if len(sys.argv) < 2:
        print("Usage: python query_hybrid_bm25_v4.py QUERY_TEXT")
        sys.exit(1)
    q = " ".join(sys.argv[1:])

    # Retrieve from both sources
    vector_nodes = vector_retriever.retrieve(q)
    bm25_nodes = bm25_retriever.retrieve(q)

    # Merge and deduplicate by node ID
    seen_ids = set()
    merged = []
    for node in vector_nodes + bm25_nodes:
        node_id = node.node.node_id
        if node_id not in seen_ids:
            seen_ids.add(node_id)
            merged.append(node)

    # Re-rank the merged candidates with cross-encoder
    reranked = reranker.postprocess_nodes(merged, query_str=q)

    # Report retrieval stats
    n_vector_only = len([n for n in vector_nodes if n.node.node_id not in {b.node.node_id for b in bm25_nodes}])
    n_bm25_only = len([n for n in bm25_nodes if n.node.node_id not in {v.node.node_id for v in vector_nodes}])
    n_both = len(vector_nodes) + len(bm25_nodes) - len(merged)

    # Estimate context length (prompt + node text)
    # MetadataMode.LLM to match what the synthesizer actually sends
    # (response_synthesizers/base.py uses MetadataMode.LLM). The default mode
    # strips the "file_name: ..." prefix and understates the prompt.
    context_text = "\n\n".join(
        n.node.get_content(metadata_mode=MetadataMode.LLM) for n in reranked)
    prompt_text = PROMPT.format(context_str=context_text, query_str=q)
    n_context_tokens = len(prompt_text.split())  # rough word count; ~1.3 tokens/word

    print(f"\nQuery: {q}")
    print(f"Vector: {len(vector_nodes)} ({n_vector_only} unique), "
          f"BM25: {len(bm25_nodes)} ({n_bm25_only} unique), "
          f"overlap: {n_both}, merged: {len(merged)}, re-ranked to: {len(reranked)}")
    
    # The token estimate uses a ~1.3 tokens/word ratio, which is a rough approximation. 
    # For an exact count you'd need the model's tokenizer, but this gives a useful ballpark 
    # for gauging how much of the context window we use.
    print(f"Context: ~{n_context_tokens} words (~{int(n_context_tokens * 1.3)} tokens)")

    # Synthesize response with LLM
    synthesizer = get_response_synthesizer(text_qa_template=PROMPT)
    response = synthesizer.synthesize(q, nodes=reranked)

    # Output
    print("\nResponse:\n")
    print(response.response)

    print("\nSource documents:")
    for node in response.source_nodes:
        meta = getattr(node, "metadata", None) or node.node.metadata
        score = getattr(node, "score", None)
        print(f"data/{meta.get('file_name')}  {score:.3f}")


if __name__ == "__main__":
    main()