# Query rewriting: LLM transforms user query to better match document language
REWRITE_PROMPT = """Rewrite the following user question to be more likely to
match technical documentation. Make it precise and use domain terminology.
Output only the rewritten question, nothing else.

User question: {query}

Rewritten:"""

def rewrite_query(query: str) -> str:
    response = client.messages.create(
        model="claude-3-haiku-20240307",
        max_tokens=100,
        messages=[{"role": "user",
                   "content": REWRITE_PROMPT.format(query=query)}]
    )
    return response.content[0].text.strip()

# "how do I make packets go faster?" →
# "methods to improve packet processing throughput and reduce latency in DPDK"

# Two-stage retrieval:
# Stage 1 — Retrieve: fast embedding search, get top-50
# Stage 2 — Rerank:   cross-encoder scores each of the 50 precisely
# Return top-5 of the reranked 50
#
# Why not use the cross-encoder for all 50,000 chunks?
# Cross-encoders are ~100x slower — fine for 50, too slow for 50,000

pip install cohere

import cohere
co = cohere.Client()   # COHERE_API_KEY from environment

def retrieve_and_rerank(
    query: str,
    collection,
    retrieve_k: int = 50,   # retrieve many
    return_k: int = 5       # return few, best ones
) -> list[dict]:
    # Stage 1: fast vector search
    results = collection.query(
        query_texts=[query], n_results=retrieve_k,
        include=["documents", "metadatas"]
    )
    docs  = results["documents"][0]
    metas = results["metadatas"][0]

    if not docs:
        return []

    # Stage 2: Cohere reranker
    rerank_response = co.rerank(
        model="rerank-english-v3.0",
        query=query,
        documents=docs,
        top_n=return_k,
        return_documents=True
    )

    return [
        {
            "text":      hit.document.text,
            "score":     hit.relevance_score,    # 0-1, higher = more relevant
            "rank":      hit.index,              # original rank before reranking
            "meta":      metas[hit.index],
        }
        for hit in rerank_response.results
    ]

# Usage
results = retrieve_and_rerank("How does DPDK mempool work?", collection)
for r in results:
    print(f"Score: {r['score']:.3f} (was rank {r['rank']+1}) | {r['text'][:60]}")

# Standard search: embed query → find similar chunks
# Problem: "make packets go faster" != "throughput optimisation"
#
# HyDE search: generate a hypothetical document → embed that → find similar chunks
# "make packets go faster" → generates paragraph using "throughput", "mbps", "pps"
# → now embedding matches real doc language

HYDE_PROMPT = """Write a short technical document passage (2-3 sentences) that would
directly answer the following question. Write as if you are an expert
writing documentation. Use precise technical terminology.

Question: {query}

Technical passage:"""

def hyde_retrieve(query: str, collection, n_results: int = 5) -> list[dict]:
    # Step 1: generate hypothetical document
    response = client.messages.create(
        model="claude-3-haiku-20240307",
        max_tokens=200,
        messages=[{"role": "user",
                   "content": HYDE_PROMPT.format(query=query)}]
    )
    hypothetical_doc = response.content[0].text.strip()

    # Step 2: embed the hypothetical doc and search
    results = collection.query(
        query_texts=[hypothetical_doc],   # ← key: search with generated doc, not original query
        n_results=n_results,
        include=["documents", "distances", "metadatas"]
    )
    return [
        {"text": doc, "score": 1-dist, "meta": meta,
         "hypothetical": hypothetical_doc}
        for doc, dist, meta in zip(
            results["documents"][0],
            results["distances"][0],
            results["metadatas"][0]
        )
    ]

# Hybrid: retrieve with both original query and HyDE, merge
def hybrid_hyde(query: str, collection, n_results: int = 5) -> list[dict]:
    standard = collection.query(query_texts=[query], n_results=n_results,
                                include=["documents", "distances", "metadatas", "ids"])
    hyde_res  = hyde_retrieve(query, collection, n_results=n_results)

    # Merge unique results, original query results get slight preference
    seen = set()
    merged = []
    for r in hyde_res:
        if r["text"] not in seen:
            seen.add(r["text"])
            merged.append(r)
    for doc, dist, meta in zip(standard["documents"][0],
                                standard["distances"][0],
                                standard["metadatas"][0]):
        if doc not in seen:
            seen.add(doc)
            merged.append({"text": doc, "score": 1-dist, "meta": meta})
    return sorted(merged, key=lambda x: x["score"], reverse=True)[:n_results]

import numpy as np

def cosine_sim(a, b):
    a, b = np.array(a), np.array(b)
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-8))

def mmr(
    query_vec: list[float],
    candidate_vecs: list[list[float]],
    candidate_docs: list[str],
    top_k: int = 5,
    lambda_param: float = 0.5   # 0=max diversity, 1=max relevance
) -> list[str]:
    """
    Maximum Marginal Relevance selection.
    Iteratively picks the candidate that maximises:
        lambda * similarity(query, doc) - (1-lambda) * max_similarity(doc, selected)
    """
    selected_idx   = []
    selected_vecs  = []
    remaining_idx  = list(range(len(candidate_docs)))

    for _ in range(min(top_k, len(candidate_docs))):
        best_score, best_idx = -1, -1

        for idx in remaining_idx:
            relevance = cosine_sim(query_vec, candidate_vecs[idx])

            if not selected_vecs:
                redundancy = 0
            else:
                redundancy = max(cosine_sim(candidate_vecs[idx], sv)
                                 for sv in selected_vecs)

            score = lambda_param * relevance - (1 - lambda_param) * redundancy
            if score > best_score:
                best_score, best_idx = score, idx

        selected_idx.append(best_idx)
        selected_vecs.append(candidate_vecs[best_idx])
        remaining_idx.remove(best_idx)

    return [candidate_docs[i] for i in selected_idx]

# Practical example
# 1. Retrieve top-20 with embeddings
# 2. Apply MMR to select 5 diverse chunks
# 3. Pass to LLM — it now has diverse context, not 5 copies of the same info

# Build a test set: queries + expected source chunks
test_set = [
    {"query": "How does DPDK mempool initialisation work?",
     "expected_source": "dpdk_guide.pdf",
     "expected_section": "Memory Management"},
    {"query": "What is the rte_ring burst size limit?",
     "expected_source": "dpdk_guide.pdf",
     "expected_section": "Ring Library"},
    # ... 20+ test cases
]

def hit_rate(results: list[dict], expected_source: str, k: int = 5) -> float:
    """1 if expected source appears in top-k, else 0."""
    top_k = results[:k]
    return 1.0 if any(expected_source in r["meta"].get("source", "")
                         for r in top_k) else 0.0

def mrr(results: list[dict], expected_source: str) -> float:
    """Mean Reciprocal Rank — higher rank = higher score."""
    for i, r in enumerate(results):
        if expected_source in r["meta"].get("source", ""):
            return 1.0 / (i + 1)
    return 0.0

def evaluate_pipeline(retrieval_fn, test_set: list[dict], k: int = 5) -> dict:
    hit_rates, mrrs = [], []
    for test in test_set:
        results = retrieval_fn(test["query"])
        hit_rates.append(hit_rate(results, test["expected_source"], k))
        mrrs.append(mrr(results, test["expected_source"]))

    return {
        f"hit_rate@{k}": round(sum(hit_rates) / len(hit_rates), 3),
        "mrr":         round(sum(mrrs) / len(mrrs), 3),
        "n_queries":   len(test_set),
    }

# Compare pipelines
baseline = evaluate_pipeline(lambda q: basic_retrieve(q), test_set)
reranked = evaluate_pipeline(lambda q: retrieve_and_rerank(q, collection), test_set)
hyde_res = evaluate_pipeline(lambda q: hyde_retrieve(q, collection), test_set)

print(f"Baseline:  {baseline}")  # {"hit_rate@5": 0.65, "mrr": 0.48}
print(f"Reranked:  {reranked}")  # {"hit_rate@5": 0.82, "mrr": 0.67}
print(f"HyDE:      {hyde_res}")  # {"hit_rate@5": 0.74, "mrr": 0.55}