Retrieval strategies & query rewriting

from openai import OpenAI

client = OpenAI()

def embed(text: str) -> list[float]:
    r = client.embeddings.create(model="text-embedding-3-small", input=text)
    return r.data[0].embedding

def naive_retrieve(question: str, store, k: int = 5) -> list[dict]:
    """Single embedding, single search — baseline."""
    q_vec = embed(question)
    return store.similarity_search(q_vec, k=k)

# store = your VectorStore wrapper (Pinecone, pgvector, Chroma, etc.)

HYDE_PROMPT = """Write a single paragraph that would appear in internal documentation
and answer the question. Use neutral policy language. Do not invent IDs or numbers.

Question: {question}

Passage:"""

def hyde_retrieve(question: str, store, k: int = 5) -> list[dict]:
    hypo = client.chat.completions.create(
        model="gpt-4o-mini",
        temperature=0,
        max_tokens=256,
        messages=[{"role": "user", "content": HYDE_PROMPT.format(question=question)}],
    ).choices[0].message.content.strip()

    hypo_vec = embed(hypo)
    hits = store.similarity_search(hypo_vec, k=k)
    for h in hits:
        h["meta"] = {**h.get("meta", {}), "hyde_hypothesis": hypo[:200]}
    return hits

MULTI_QUERY_PROMPT = """Generate {n} different search queries that would help retrieve
documentation to answer the user question. One per line, no numbering.

Question: {question}"""

def generate_query_variants(question: str, n: int = 3) -> list[str]:
    text = client.chat.completions.create(
        model="gpt-4o-mini",
        temperature=0.3,
        max_tokens=200,
        messages=[{"role": "user", "content": MULTI_QUERY_PROMPT.format(n=n, question=question)}],
    ).choices[0].message.content
    variants = [question] + [ln.strip() for ln in text.splitlines() if ln.strip()]
    return variants[: n + 1]

def dedupe_by_id(lists: list[list[dict]], top_n: int = 8) -> list[dict]:
    seen: dict[str, dict] = {}
    for results in lists:
        for rank, hit in enumerate(results):
            cid = hit["id"]
            if cid not in seen or hit.get("score", 0) > seen[cid].get("score", 0):
                seen[cid] = {**hit, "best_rank": min(seen.get(cid, {}).get("best_rank", 999), rank)}
    ordered = sorted(seen.values(), key=lambda h: (h.get("best_rank", 999), -h.get("score", 0)))
    return ordered[:top_n]

def multi_query_retrieve(question: str, store, k_per_query: int = 4) -> list[dict]:
    variants = generate_query_variants(question)
    lists = [naive_retrieve(q, store, k=k_per_query) for q in variants]
    return dedupe_by_id(lists)

STEP_BACK_PROMPT = """You are an expert at generalizing questions to underlying principles.
Given a specific user question, write ONE broader question that would help retrieve
high-level policy or concept documentation. Output only the question.

Specific question: {question}"""

def step_back_question(question: str) -> str:
    return client.chat.completions.create(
        model="gpt-4o-mini",
        temperature=0,
        max_tokens=120,
        messages=[{"role": "user", "content": STEP_BACK_PROMPT.format(question=question)}],
    ).choices[0].message.content.strip()

def step_back_retrieve(question: str, store) -> list[dict]:
    broad = step_back_question(question)
    principle_hits = naive_retrieve(broad, store, k=3)
    detail_hits = naive_retrieve(question, store, k=5)
    return dedupe_by_id([principle_hits, detail_hits], top_n=8)

DECOMPOSE_PROMPT = """Break the user question into independent sub-questions answerable from a knowledge base.
Return JSON: {{"sub_questions": ["...", "..."]}} Max 4 sub-questions.

Question: {question}"""

def decompose(question: str) -> list[str]:
    import json
    raw = client.chat.completions.create(
        model="gpt-4o-mini",
        temperature=0,
        response_format={"type": "json_object"},
        messages=[{"role": "user", "content": DECOMPOSE_PROMPT.format(question=question)}],
    ).choices[0].message.content
    data = json.loads(raw)
    subs = data.get("sub_questions", [])
    return subs if subs else [question]

def decomposed_retrieve(question: str, store, k_per_sub: int = 4) -> list[dict]:
    subs = decompose(question)
    lists = []
    for sq in subs:
        hits = naive_retrieve(sq, store, k=k_per_sub)
        for h in hits:
            h.setdefault("meta", {})["sub_question"] = sq
        lists.append(hits)
    merged = dedupe_by_id(lists, top_n=10)
    return merged

import cohere

co = cohere.Client()

def rerank_hits(question: str, candidates: list[dict], top_n: int = 5) -> list[dict]:
    if not candidates:
        return []
    docs = [c["text"] for c in candidates]
    response = co.rerank(
        model="rerank-v3.5",
        query=question,
        documents=docs,
        top_n=min(top_n, len(docs)),
    )
    out = []
    for r in response.results:
        hit = {**candidates[r.index], "rerank_score": r.relevance_score}
        out.append(hit)
    return out

def two_stage_retrieve(question: str, store, retrieve_n: int = 50, final_k: int = 5) -> list[dict]:
    coarse = naive_retrieve(question, store, k=retrieve_n)
    return rerank_hits(question, coarse, top_n=final_k)

from sentence_transformers import CrossEncoder

_reranker = CrossEncoder("BAAI/bge-reranker-v2-m3")

def bge_rerank(question: str, candidates: list[dict], top_n: int = 5) -> list[dict]:
    pairs = [(question, c["text"]) for c in candidates]
    scores = _reranker.predict(pairs)
    ranked = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
    return [{**c, "rerank_score": float(s)} for c, s in ranked[:top_n]]

CONTEXTUALIZE_PROMPT = """<document>
{whole_document}
</document>

Here is the chunk we want to situate within the whole document:
<chunk>
{chunk}
</chunk>

Give a short succinct context (50-100 tokens) to situate this chunk in the document
for search retrieval. Answer only with the context, nothing else."""

def contextualize_chunk(whole_document: str, chunk: str) -> str:
    msg = CONTEXTUALIZE_PROMPT.format(whole_document=whole_document[:120_000], chunk=chunk)
    prefix = client.chat.completions.create(
        model="gpt-4o-mini",
        temperature=0,
        max_tokens=150,
        messages=[{"role": "user", "content": msg}],
    ).choices[0].message.content.strip()
    return f"{prefix}\n\n{chunk}"

def index_with_context(chunks: list[dict], doc_text: str, store) -> None:
    for ch in chunks:
        embedded_text = contextualize_chunk(doc_text, ch["text"])
        vec = embed(embedded_text)
        store.upsert(id=ch["id"], vector=vec, text=ch["text"], meta={
            **ch.get("meta", {}),
            "embedded_text": embedded_text[:500],  # debug only; redact in prod logs
        })

from rank_bm25 import BM25Okapi

class Bm25Index:
    def __init__(self, chunks: list[dict]):
        self.chunks = chunks
        corpus = [c["text"].lower().split() for c in chunks]
        self.bm25 = BM25Okapi(corpus)

    def search(self, query: str, k: int = 10) -> list[dict]:
        tokens = query.lower().split()
        scores = self.bm25.get_scores(tokens)
        ranked = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:k]
        return [
            {"id": self.chunks[i]["id"], "text": self.chunks[i]["text"],
             "score": float(scores[i]), "channel": "bm25"}
            for i in ranked if scores[i] > 0
        ]

# OpenSearch 2.x — combine knn and match in one query (then RRF in app or search pipeline)
hybrid_query = {
    "query": {
        "hybrid": {
            "queries": [
                {"match": {"text": {"query": question}}},
                {"knn": {"embedding": {"vector": q_vec, "k": 50}}},
            ]
        }
    }
}
# Many teams still run two queries + RRF in application code for tunability

from collections import defaultdict

def rrf_score(rank: int, k: int = 60) -> float:
    """rank is 1-based (first result = 1)."""
    return 1.0 / (k + rank)

def rrf_merge(lists: list[list[dict]], k: int = 60, top_n: int = 10) -> list[dict]:
    """
    lists: each inner list is ordered best-first with unique chunk 'id' keys.
    Returns merged list sorted by RRF score descending.
    """
    scores: dict[str, float] = defaultdict(float)
    payload: dict[str, dict] = {}

    for results in lists:
        for rank_0, hit in enumerate(results):
            cid = hit["id"]
            scores[cid] += rrf_score(rank_0 + 1, k=k)
            # keep richest payload (prefer dense metadata if duplicate)
            payload[cid] = {**payload.get(cid, {}), **hit}

    ordered_ids = sorted(scores.keys(), key=lambda c: scores[c], reverse=True)[:top_n]
    return [{**payload[cid], "rrf_score": scores[cid]} for cid in ordered_ids]

def hybrid_retrieve(question: str, dense_store, bm25: Bm25Index, k_each: int = 25, final_n: int = 10) -> list[dict]:
    dense_hits = naive_retrieve(question, dense_store, k=k_each)
    sparse_hits = bm25.search(question, k=k_each)
    return rrf_merge([dense_hits, sparse_hits], k=60, top_n=final_n)

def retrieve_for_production(question: str, store, bm25: Bm25Index, cfg) -> tuple[list[dict], dict]:
    meta = {"path": "fast"}
    coarse = hybrid_retrieve(question, store, bm25, k_each=25, final_n=50)

    if coarse and coarse[0].get("score", 0) >= cfg.confidence_skip_rewrite:
        meta["path"] = "hybrid_only"
        final = rerank_hits(question, coarse, top_n=5) if cfg.use_rerank else coarse[:5]
        return final, meta

    if cfg.use_multi_query:
        coarse = multi_query_retrieve(question, store)
        meta["path"] = "multi_query+hybrid"

    if cfg.use_rerank:
        final = rerank_hits(question, coarse, top_n=5)
        meta["path"] += "+rerank"
    else:
        final = coarse[:5]

    return final, meta