from sentence_transformers import SentenceTransformer, util

ENCODER_MODEL = "all-MiniLM-L6-v2"
DEFAULT_INTENT = "knowledge_base"

INTENT_EXAMPLES: dict[str, list[str]] = {
    "greeting": ["Hi", "Hello", "Thanks", "Goodbye"],
    "knowledge_base": ["How do I reset my password?", "What is RAG?"],
    "structured_data": ["What's my account balance?", "Show my recent orders"],
    "real_time": ["What's the weather in SF?", "AAPL stock price"],
    "out_of_scope": ["Tell me a joke", "What's the meaning of life?"],
}


class SemanticIntentClassifier:
    """Classifies query intent by cosine similarity against few-shot exemplars."""

    def __init__(self, model_name: str = ENCODER_MODEL) -> None:
        self.model = SentenceTransformer(model_name)
        self.intent_embeddings = {
            intent: self.model.encode(examples)
            for intent, examples in INTENT_EXAMPLES.items()
        }

    def classify(self, query: str, threshold: float = 0.5) -> str:
        """Returns the best-matching intent, or DEFAULT_INTENT if below threshold."""
        if not query.strip():
            raise ValueError("query must be non-empty")
        query_embedding = self.model.encode(query)
        intent_scores = {
            intent: float(util.cos_sim(query_embedding, embs)[0].max())
            for intent, embs in self.intent_embeddings.items()
        }
        best_intent = max(intent_scores, key=intent_scores.get)
        return best_intent if intent_scores[best_intent] >= threshold else DEFAULT_INTENT

import re
from typing import Literal

IntentType = Literal["greeting", "knowledge_base", "structured_data", "real_time", "out_of_scope"]

INTENT_PATTERNS: dict[IntentType, str] = {
    "greeting": r"\b(hi|hello|hey|thanks|thank you|bye|goodbye)\b",
    "structured_data": r"\b(my account|my balance|my order|my subscription|purchase history)\b",
    "real_time": r"\b(weather|temperature|forecast|stock price|news)\b",
    "out_of_scope": r"\b(joke|poem|story|write code|meaning of life)\b",
}


def classify_intent_simple(query: str) -> IntentType:
    """Matches query against keyword patterns; falls back to knowledge_base."""
    query_lower = query.lower()
    for intent, pattern in INTENT_PATTERNS.items():
        if re.search(pattern, query_lower):
            return intent
    return "knowledge_base"

from sentence_transformers import SentenceTransformer
from sklearn.linear_model import LogisticRegression

ENCODER_MODEL = "all-MiniLM-L6-v2"
MAX_SOLVER_ITERATIONS = 1000


class MLIntentClassifier:
    """Trains a logistic regression classifier on embedded query examples."""

    def __init__(
        self,
        training_data: list[tuple[str, str]],
        model_name: str = ENCODER_MODEL,
    ) -> None:
        self.encoder = SentenceTransformer(model_name)
        queries, intents = zip(*training_data)
        embeddings = self.encoder.encode(queries)
        self.classifier = LogisticRegression(max_iter=MAX_SOLVER_ITERATIONS)
        self.classifier.fit(embeddings, intents)

    def classify(self, query: str) -> str:
        """Returns predicted intent label for a single query."""
        return self.classifier.predict([self.encoder.encode([query])[0]])[0]


training_data = [
    ("Hi", "greeting"), ("How do I reset my password?", "knowledge_base"),
    ("What's my account balance?", "structured_data"),
    ("What's the weather in SF?", "real_time"),
    ("Tell me a joke", "out_of_scope"),
    # ... 100-1000+ examples
]
classifier = MLIntentClassifier(training_data)

GREETING_RESPONSE = "Hi! How can I help?"
OUT_OF_SCOPE_RESPONSE = "I'm designed for account and product support."


def route_query(query: str, user_id: str) -> dict[str, str]:
    """Classifies query intent and dispatches to the appropriate backend."""
    intent = classify_intent(query)
    if intent == "greeting":
        return {"intent": intent, "response": GREETING_RESPONSE, "source": "direct"}
    if intent == "knowledge_base":
        chunks = hybrid_search(query, documents, k=5)
        return {"intent": intent, "response": llm_generate(query, chunks), "source": "vector_db"}
    if intent == "structured_data":
        result = database.execute(generate_sql(query, user_id))
        return {"intent": intent, "response": format_sql_result(result), "source": "sql"}
    if intent == "real_time":
        return {"intent": intent, "response": call_api(query), "source": "api"}
    return {"intent": "out_of_scope", "response": OUT_OF_SCOPE_RESPONSE, "source": "direct"}

from sentence_transformers import SentenceTransformer
import numpy as np
import faiss

ENCODER_MODEL = "sentence-transformers/all-mpnet-base-v2"
HNSW_M = 32
STAGE1_K = 50

model = SentenceTransformer(ENCODER_MODEL)

doc_embeddings = model.encode(documents, convert_to_numpy=True)
# L2 default; equivalent to cosine on normalized vectors
index = faiss.IndexHNSWFlat(doc_embeddings.shape[1], HNSW_M)
index.add(doc_embeddings)

query_embedding = model.encode("What is RAG?", convert_to_numpy=True)
distances, indices = index.search(query_embedding[np.newaxis, :], k=STAGE1_K)

from transformers import DPRQuestionEncoder, DPRContextEncoder, DPRQuestionEncoderTokenizer, DPRContextEncoderTokenizer
import torch
import numpy as np

DPR_QUESTION_MODEL = "facebook/dpr-question_encoder-single-nq-base"
DPR_CONTEXT_MODEL = "facebook/dpr-ctx_encoder-single-nq-base"
STAGE1_K = 50

q_encoder = DPRQuestionEncoder.from_pretrained(DPR_QUESTION_MODEL)
c_encoder = DPRContextEncoder.from_pretrained(DPR_CONTEXT_MODEL)
q_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained(DPR_QUESTION_MODEL)
c_tokenizer = DPRContextEncoderTokenizer.from_pretrained(DPR_CONTEXT_MODEL)

doc_tokens = c_tokenizer(documents, return_tensors="pt", padding=True, truncation=True)
with torch.no_grad():
    doc_embeddings = c_encoder(**doc_tokens).pooler_output.numpy()

query_tokens = q_tokenizer([query], return_tensors="pt")
with torch.no_grad():
    query_embedding = q_encoder(**query_tokens).pooler_output.numpy()

similarities = np.dot(doc_embeddings, query_embedding.T).squeeze()
top_indices = np.argsort(similarities)[::-1][:STAGE1_K]

from openai import OpenAI

EMBEDDING_MODEL = "text-embedding-3-large"
TRUNCATED_DIMS = 256

client = OpenAI()

full = client.embeddings.create(input="What is RAG?", model=EMBEDDING_MODEL)

compact = client.embeddings.create(
    input="What is RAG?", model=EMBEDDING_MODEL, dimensions=TRUNCATED_DIMS
)
# For models without native truncation: embedding[:TRUNCATED_DIMS] then L2-normalize

from rank_bm25 import BM25Okapi
import numpy as np

STAGE1_K = 50

tokenized_docs = [doc.lower().split() for doc in documents]
bm25 = BM25Okapi(tokenized_docs)

query = "HIPAA compliance requirements"
scores = bm25.get_scores(query.lower().split())
top_indices = np.argsort(scores)[::-1][:STAGE1_K]

from transformers import AutoModelForMaskedLM, AutoTokenizer
import torch

SPLADE_MODEL = "naver/splade-cocondenser-ensembledistil"
MAX_TOKEN_LENGTH = 256


class SpladeEncoder:
    """Encodes text into SPLADE sparse vectors using learned term expansion."""

    def __init__(self, model_name: str = SPLADE_MODEL) -> None:
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForMaskedLM.from_pretrained(model_name)

    def encode(self, text: str) -> dict[str, float]:
        """Returns sparse vector mapping vocabulary terms to importance weights."""
        tokens = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=MAX_TOKEN_LENGTH)
        with torch.no_grad():
            logits = self.model(**tokens).logits
        # Max pooling over sequence length, then ReLU + log(1+x) for sparsity
        weights = torch.max(
            torch.log1p(torch.relu(logits)) * tokens["attention_mask"].unsqueeze(-1), dim=1
        ).values
        non_zero = weights.squeeze().nonzero().squeeze()
        return {self.tokenizer.decode(idx.item()): weights[0, idx].item() for idx in non_zero}

EPSILON = 1e-8


def _min_max_normalize(scores: np.ndarray) -> np.ndarray:
    """Normalizes scores to [0, 1] range."""
    return (scores - scores.min()) / (scores.max() - scores.min() + EPSILON)


def hybrid_search(
    query: str,
    documents: list[str],
    doc_embeddings: np.ndarray,
    encoder: SentenceTransformer,
    bm25: BM25Okapi,
    k: int = 50,
    alpha: float = 0.5,
) -> list[tuple[int, float]]:
    """Fuses dense cosine and BM25 sparse scores with min-max normalization."""
    if not documents:
        raise ValueError("documents must be non-empty")

    query_embedding = encoder.encode(query, convert_to_numpy=True)
    dense_scores = np.dot(doc_embeddings, query_embedding)
    sparse_scores = bm25.get_scores(query.lower().split())

    # Normalize both to [0, 1] — critical for fair combination
    dense_norm = _min_max_normalize(dense_scores)
    sparse_norm = _min_max_normalize(sparse_scores)

    hybrid_scores = alpha * dense_norm + (1 - alpha) * sparse_norm
    top_indices = np.argsort(hybrid_scores)[::-1][:k]
    return [(int(i), float(hybrid_scores[i])) for i in top_indices]

def reciprocal_rank_fusion(
    dense_ranks: list[int],
    sparse_ranks: list[int],
    k: int = 60,
) -> list[tuple[int, float]]:
    """Merges ranked lists using RRF: score(doc) = sum(1 / (k + rank))."""
    scores: dict[int, float] = {}
    for rank, doc_id in enumerate(dense_ranks):
        scores[doc_id] = scores.get(doc_id, 0) + 1 / (k + rank + 1)
    for rank, doc_id in enumerate(sparse_ranks):
        scores[doc_id] = scores.get(doc_id, 0) + 1 / (k + rank + 1)
    return sorted(scores.items(), key=lambda x: x[1], reverse=True)

from sentence_transformers import CrossEncoder
import numpy as np

RERANKER_MODEL = "cross-encoder/ms-marco-MiniLM-L6-v2"
STAGE1_K = 50
FINAL_K = 5

reranker = CrossEncoder(RERANKER_MODEL)

candidates = hybrid_search(query, documents, k=STAGE1_K)
candidate_docs = [documents[i] for i, score in candidates]

pairs = [[query, doc] for doc in candidate_docs]
rerank_scores = reranker.predict(pairs)
top_docs = [candidate_docs[i] for i in np.argsort(rerank_scores)[::-1][:FINAL_K]]

from ragatouille import RAGPretrainedModel

COLBERT_MODEL = "colbert-ir/colbertv2.0"
STAGE1_K = 100
STAGE2_K = 50
FINAL_K = 5

colbert_model = RAGPretrainedModel.from_pretrained(COLBERT_MODEL)

stage1_results = dense_retrieval(query, k=STAGE1_K)
stage2_results = colbert_model.rerank(query=query, documents=stage1_results, k=STAGE2_K)
stage3_results = cross_encoder.rerank(query, stage2_results, k=FINAL_K)

EFSEARCH_HIGH_RECALL = 50  # Default 16; increase for better recall at higher latency

index = faiss.read_index("index.faiss")
index.hnsw.efSearch = EFSEARCH_HIGH_RECALL

COLLECTION_NAME = "documents"
STAGE1_K = 50

results = client.search(
    collection_name=COLLECTION_NAME,
    query_vector=query_embedding.tolist(),
    query_filter=Filter(must=[
        FieldCondition(key="category", match=MatchValue(value="healthcare")),
        FieldCondition(key="created_at", range=DatetimeRange(gte="2024-01-01T00:00:00Z")),
    ]),
    limit=STAGE1_K,
)

COLLECTION_NAME = "documents"
JWT_ALGORITHM = "HS256"


def secure_retrieval(
    query: str,
    user_token: str,
    encoder: SentenceTransformer,
    client: QdrantClient,
    secret_key: str,
    k: int = 10,
) -> list[dict[str, str | float]]:
    """Retrieves chunks filtered by JWT-derived RBAC roles."""
    if not query.strip():
        raise ValueError("query must be non-empty")

    payload = jwt.decode(user_token, secret_key, algorithms=[JWT_ALGORITHM])
    user_roles: list[str] = payload.get("roles", [])

    rbac_filter = Filter(must=[
        FieldCondition(key="roles", match=MatchAny(any=user_roles))
    ])

    results = client.search(
        collection_name=COLLECTION_NAME,
        query_vector=encoder.encode(query).tolist(),
        query_filter=rbac_filter,
        limit=k,
    )
    return [{"text": hit.payload["text"], "score": hit.score} for hit in results]

# Department-level
metadata = {"roles": ["legal"], "departments": ["legal"]}
# Team-level
metadata = {"roles": ["legal", "compliance"], "team": "privacy-team"}
# Individual-level
metadata = {"roles": ["legal"], "owner_id": "user_456"}

from datetime import datetime, UTC

COLLECTION_NAME = "documents"
JWT_ALGORITHM = "HS256"


def secure_retrieval_with_audit(
    query: str,
    user_token: str,
    encoder: SentenceTransformer,
    client: QdrantClient,
    secret_key: str,
    k: int = 10,
) -> list:
    """Retrieves RBAC-filtered chunks and logs an audit trail to SIEM."""
    if not query.strip():
        raise ValueError("query must be non-empty")

    payload = jwt.decode(user_token, secret_key, algorithms=[JWT_ALGORITHM])
    user_id: str = payload["user_id"]
    user_roles: list[str] = payload.get("roles", [])

    results = client.search(
        collection_name=COLLECTION_NAME,
        query_vector=encoder.encode(query).tolist(),
        query_filter=Filter(must=[FieldCondition(key="roles", match=MatchAny(any=user_roles))]),
        limit=k,
    )

    audit_log = {
        "timestamp": datetime.now(UTC).isoformat(),
        "user_id": user_id, "user_roles": user_roles,
        "query": query, "chunks_returned": len(results),
    }
    log_to_siem(audit_log)
    return results

from datetime import datetime, UTC

EPSILON = 1e-8
RECENCY_BOOST = 0.2
DAYS_PER_YEAR = 365
AUTHORITY_BOOST = 1.1
DEPARTMENT_BOOST = 1.15
TRUSTED_AUTHORS = frozenset({"legal-team", "compliance-team"})


def metadata_boosted_search(
    query: str,
    documents: list[dict],
    doc_embeddings: np.ndarray,
    encoder: SentenceTransformer,
    user_department: str,
    k: int = 10,
) -> list[dict]:
    """Applies recency, authority, and department boosts to base retrieval scores."""
    base_scores = np.dot(doc_embeddings, encoder.encode(query))
    base_scores = (base_scores - base_scores.min()) / (base_scores.max() - base_scores.min() + EPSILON)

    for i, doc in enumerate(documents):
        age_days = (datetime.now(UTC) - datetime.fromisoformat(doc["created_at"])).days
        base_scores[i] *= (1 + RECENCY_BOOST * np.exp(-age_days / DAYS_PER_YEAR))
        if doc.get("author") in TRUSTED_AUTHORS:
            base_scores[i] *= AUTHORITY_BOOST
        if doc.get("category") == user_department:
            base_scores[i] *= DEPARTMENT_BOOST

    return [documents[i] for i in np.argsort(base_scores)[::-1][:k]]

def reorder_for_llm(chunks: list[dict]) -> list[dict]:
    """Places most relevant chunks at start and end to mitigate lost-in-the-middle effect."""
    front = chunks[::2]
    back = chunks[1::2][::-1]
    return front + back

EPSILON = 1e-8
FIRST_CHUNK_BOOST = 1.2
LAST_CHUNK_BOOST = 1.15
HEADER_BOOST = 1.1
MIDDLE_PENALTY_FACTOR = 0.15
LONG_DOC_THRESHOLD = 10

SECTION_WEIGHTS: dict[str, dict[str, float]] = {
    "research_paper": {"abstract": 1.15, "methodology": 1.25, "results": 1.2, "references": 0.8},
    "technical_doc": {"implementation": 1.3, "api_reference": 1.2, "troubleshooting": 1.15},
    "legal_doc": {"definitions": 1.3, "requirements": 1.25, "obligations": 1.25, "preamble": 0.9},
}


def positional_boosted_search(
    query: str,
    documents: list[dict],
    doc_embeddings: np.ndarray,
    encoder: SentenceTransformer,
    k: int = 10,
) -> list[dict]:
    """Boosts retrieval scores based on chunk position within the source document."""
    base_scores = np.dot(doc_embeddings, encoder.encode(query))
    base_scores = (base_scores - base_scores.min()) / (base_scores.max() - base_scores.min() + EPSILON)
    boosted_scores = base_scores.copy()

    for i, doc in enumerate(documents):
        chunk_position = doc.get("chunk_index", 0)
        total_chunks = doc.get("total_chunks", 1)

        if chunk_position == 0:
            boosted_scores[i] *= FIRST_CHUNK_BOOST
        elif chunk_position == total_chunks - 1:
            boosted_scores[i] *= LAST_CHUNK_BOOST
        if doc.get("is_section_header", False):
            boosted_scores[i] *= HEADER_BOOST
        if total_chunks > LONG_DOC_THRESHOLD:
            normalized_position = chunk_position / (total_chunks - 1)
            position_weight = 1 - MIDDLE_PENALTY_FACTOR * (1 - 4 * (normalized_position - 0.5) ** 2)
            boosted_scores[i] *= position_weight

    return [documents[i] for i in np.argsort(boosted_scores)[::-1][:k]]

MAX_PARAPHRASE_LENGTH = 50
NUM_PARAPHRASES = 3
NUM_BEAMS = 5


def expand_query(query: str, tokenizer: T5Tokenizer, model: T5ForConditionalGeneration) -> str:
    """Generates paraphrases and concatenates them with the original query for broader recall."""
    input_ids = tokenizer(f"paraphrase: {query}", return_tensors="pt").input_ids
    outputs = model.generate(
        input_ids, max_length=MAX_PARAPHRASE_LENGTH,
        num_return_sequences=NUM_PARAPHRASES, num_beams=NUM_BEAMS,
    )
    paraphrases = [tokenizer.decode(o, skip_special_tokens=True) for o in outputs]
    return query + " " + " ".join(paraphrases)

HISTORY_WINDOW = 5


def rewrite_conversational_query(query: str, chat_history: list[dict[str, str]], llm: LLM) -> str:
    """Rewrites a follow-up query into a standalone question using recent chat history."""
    recent_history = chat_history[-HISTORY_WINDOW:]
    history_text = "\n".join([f"{m['role']}: {m['content']}" for m in recent_history])
    prompt = f"""Rewrite the user's latest query into a standalone question.

Conversation history:
{history_text}

Latest query: {query}

Standalone query:"""
    return llm.invoke(prompt).content.strip()

import hashlib
import json
import re

CONTEXTUAL_REFERENCE_PATTERNS = [
    r"\b(it|its|this|that|these|those|they|them|their)\b",
    r"\b(also|too|as well|additionally|furthermore)\b",
    r"\b(instead|however)\b",
]


def needs_rewriting(query: str) -> bool:
    """Cheap heuristic gate: detects pronouns and contextual references."""
    return any(re.search(p, query, re.IGNORECASE) for p in CONTEXTUAL_REFERENCE_PATTERNS)


def _chat_history_hash(chat_history: list[dict[str, str]]) -> str:
    """Deterministic hash of chat history for cache keying."""
    return hashlib.md5(
        json.dumps(chat_history, sort_keys=True).encode(), usedforsecurity=False
    ).hexdigest()


def rewrite_with_cache(
    query: str,
    chat_history: list[dict[str, str]],
    llm: LLM,
    cache: dict[str, str],
) -> str:
    """Returns cached rewrite if available; otherwise rewrites and caches."""
    cache_key = f"{query}:{_chat_history_hash(chat_history)}"
    if cache_key in cache:
        return cache[cache_key]
    rewritten = rewrite_conversational_query(query, chat_history, llm)
    cache[cache_key] = rewritten
    return rewritten

from transformers import T5ForConditionalGeneration, T5Tokenizer

REWRITER_MODEL = "your-org/t5-query-rewriter"
HISTORY_WINDOW = 3
MAX_REWRITE_LENGTH = 100


class FastQueryRewriter:
    """Fine-tuned T5 model for sub-millisecond conversational query rewriting."""

    def __init__(self, model_name: str = REWRITER_MODEL) -> None:
        self.model = T5ForConditionalGeneration.from_pretrained(model_name)
        self.tokenizer = T5Tokenizer.from_pretrained(model_name)

    def rewrite(self, query: str, chat_history: list[dict[str, str]]) -> str:
        """Rewrites a follow-up query into a standalone question using recent context."""
        history_text = " ".join([msg["content"] for msg in chat_history[-HISTORY_WINDOW:]])
        input_ids = self.tokenizer(
            f"rewrite query with context: {history_text} [SEP] {query}", return_tensors="pt"
        ).input_ids
        outputs = self.model.generate(input_ids, max_length=MAX_REWRITE_LENGTH)
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)

SUB_QUERY_K = 20
FINAL_K = 5


def decompose_query(query: str, llm: LLM) -> list[str]:
    """Breaks a multi-part question into 2-4 independent sub-queries via LLM."""
    prompt = f"Break this into 2-4 simple sub-questions:\n\nQuestion: {query}\n\nSub-questions:"
    return [q.strip() for q in llm.invoke(prompt).content.split("\n") if q.strip()]


# "Compare HIPAA and GDPR for healthcare data" →
# 1. What are HIPAA's privacy requirements?
# 2. What are GDPR's privacy requirements?
# 3. How do they differ for healthcare data?

all_results: dict[int, float] = {}
for sub_q in decompose_query(query, llm):
    for doc_id, score in retrieve(sub_q, documents, k=SUB_QUERY_K):
        all_results[doc_id] = max(all_results.get(doc_id, 0), score)
top_chunks = rerank(query, list(all_results.keys()), k=FINAL_K)

from langchain.chains.query_constructor.base import AttributeInfo
from langchain.retrievers.self_query.base import SelfQueryRetriever

metadata_field_info = [
    AttributeInfo(name="category", description="Document category", type="string"),
    AttributeInfo(name="year", description="Year published", type="integer"),
    AttributeInfo(name="author", description="Author name", type="string"),
]

retriever = SelfQueryRetriever.from_llm(
    llm=llm,
    vectorstore=vectorstore,
    document_contents="Company policy documents",
    metadata_field_info=metadata_field_info,
)

# "healthcare policies from 2023" →
#   query="healthcare policies", filter=gte("year",2023)

def hyde_retrieval(
    query: str, documents: list[str], encoder: SentenceTransformer, llm: LLM, k: int = 5,
) -> list[str]:
    """Generates a hypothetical answer and uses its embedding for retrieval."""
    hypothetical_answer = llm.invoke(
        f'Given the question: "{query}"\nWrite a detailed, factual answer.'
    ).content
    hyde_embedding = encoder.encode(hypothetical_answer)
    similarities = np.dot(encoder.encode(documents), hyde_embedding)
    return [documents[i] for i in np.argsort(similarities)[::-1][:k]]

NUM_VARIATIONS = 3
STAGE1_K = 50
RRF_K = 60


def multi_query_retrieval(
    query: str, documents: list[str], encoder: SentenceTransformer, llm: LLM, k: int = 5,
) -> list[str]:
    """Generates query variations, retrieves for each, and merges via RRF."""
    variations = llm.invoke(
        f"Generate {NUM_VARIATIONS} different versions of: {query}"
    ).content.split("\n")[:NUM_VARIATIONS]

    doc_embeddings = encoder.encode(documents)
    scores: dict[int, float] = {}
    for q in [query] + variations:
        sims = np.dot(doc_embeddings, encoder.encode(q))
        ranked = np.argsort(sims)[::-1][:STAGE1_K]
        for rank, idx in enumerate(ranked):
            scores[idx] = scores.get(idx, 0) + 1 / (RRF_K + rank + 1)

    return [documents[i] for i in sorted(scores, key=scores.get, reverse=True)[:k]]

GENERAL_FRACTION = 1 / 3


def step_back_retrieval(query: str, documents: list[str], llm: LLM, k: int = 10) -> list:
    """Retrieves with both the original and a generalized query, then reranks the union."""
    step_back_query = llm.invoke(
        f'Generate a more general question for: "{query}"'
    ).content.strip()

    specific = retrieve(query, documents, k=k)
    general = retrieve(step_back_query, documents, k=k)
    n_general = int(k * GENERAL_FRACTION)
    n_specific = k - n_general
    seen: set[int] = set()
    combined: list[int] = []
    for doc_id, score in specific[:n_specific] + general[:n_general]:
        if doc_id not in seen:
            seen.add(doc_id)
            combined.append(doc_id)
    return rerank(query, combined, k=k)

def mmr_rerank(
    query: str,
    candidates: list[str],
    encoder: SentenceTransformer,
    k: int = 5,
    lambda_param: float = 0.5,
) -> list[str]:
    """Selects k candidates balancing relevance to query and diversity among selections."""
    if k > len(candidates):
        raise ValueError(f"k={k} exceeds candidate count {len(candidates)}")

    query_emb = encoder.encode(query)
    cand_embs = encoder.encode(candidates)
    relevance = np.dot(cand_embs, query_emb)

    selected = [int(np.argmax(relevance))]
    sel_embs = [cand_embs[selected[0]]]

    while len(selected) < k:
        remaining = [i for i in range(len(candidates)) if i not in selected]
        mmr_scores = []
        for i in remaining:
            diversity_penalty = max(np.dot(cand_embs[i], e) for e in sel_embs)
            mmr_scores.append((i, lambda_param * relevance[i] - (1 - lambda_param) * diversity_penalty))
        best = max(mmr_scores, key=lambda x: x[1])[0]
        selected.append(best)
        sel_embs.append(cand_embs[best])

    return [candidates[i] for i in selected]

def embedding_compress(
    query: str, chunks: list[str], encoder: SentenceTransformer, threshold: float = 0.5,
) -> list[str]:
    """Keeps only sentences whose embedding similarity to the query exceeds threshold."""
    query_emb = encoder.encode(query)
    compressed = []
    for chunk in chunks:
        sentences = chunk.split(". ")
        relevant = [
            s for s, sim in zip(sentences, np.dot(encoder.encode(sentences), query_emb))
            if sim >= threshold
        ]
        if relevant:
            compressed.append(". ".join(relevant))
    return compressed

EXTRACTION_PROMPT = """Extract ONLY sentences relevant to the query. Return empty if none.
Query: {query}
Chunk: {chunk}
Relevant sentences:"""


def llm_compress(query: str, chunks: list[str], llm: LLM) -> list[str]:
    """Extracts query-relevant sentences from each chunk via LLM."""
    return [
        r for chunk in chunks
        if (r := llm.invoke(EXTRACTION_PROMPT.format(query=query, chunk=chunk)).content.strip())
    ]


def pipeline_compress(
    query: str,
    chunks: list[str],
    encoder: SentenceTransformer,
    redundancy_threshold: float = 0.85,
    relevance_threshold: float = 0.5,
) -> list[str]:
    """Deduplicates sentences by embedding similarity, then filters by query relevance."""
    query_emb = encoder.encode(query)
    all_sentences = [s.strip() for chunk in chunks for s in chunk.split(". ") if s.strip()]
    sent_embs = encoder.encode(all_sentences)

    unique_sents: list[str] = []
    unique_embs: list[np.ndarray] = []
    for sent, emb in zip(all_sentences, sent_embs):
        if not unique_embs or np.max(np.dot(unique_embs, emb)) < redundancy_threshold:
            unique_sents.append(sent)
            unique_embs.append(emb)

    scores = np.dot(unique_embs, query_emb)
    return [s for s, score in zip(unique_sents, scores) if score >= relevance_threshold]

class SentenceWindowRetriever:
    """Embeds individual sentences and expands to surrounding context at search time."""

    def __init__(self, documents: list[str], encoder: SentenceTransformer, window_size: int = 3) -> None:
        self.window_size = window_size
        self.encoder = encoder
        self.sentences: list[str] = []
        self.metadata: list[dict] = []
        for doc_idx, doc in enumerate(documents):
            doc_sents = [s.strip() + "." for s in doc.split(".") if s.strip()]
            for sent_idx, sent in enumerate(doc_sents):
                self.sentences.append(sent)
                self.metadata.append({"doc_idx": doc_idx, "sent_idx": sent_idx, "doc_sents": doc_sents})
        self.embeddings = self.encoder.encode(self.sentences)

    def search(self, query: str, k: int = 5) -> list[dict[str, str]]:
        """Returns k sentence windows most similar to the query."""
        sims = np.dot(self.embeddings, self.encoder.encode(query))
        windows: list[dict[str, str]] = []
        for idx in np.argsort(sims)[::-1][:k]:
            m = self.metadata[idx]
            start = max(0, m["sent_idx"] - self.window_size)
            end = min(len(m["doc_sents"]), m["sent_idx"] + self.window_size + 1)
            windows.append({"window": " ".join(m["doc_sents"][start:end]), "target": self.sentences[idx]})
        return windows

from pydantic import BaseModel

LEAF_LEVEL = 2


class HierarchicalNode(BaseModel):
    """A node in a parent-intermediate-leaf chunking hierarchy."""

    chunk_id: str
    text: str
    level: int
    parent_id: str | None
    children_ids: list[str]


def auto_merge_retrieve(
    query: str,
    nodes: dict[str, HierarchicalNode],
    encoder: SentenceTransformer,
    k: int = 12,
    merge_threshold: float = 0.5,
) -> list[str]:
    """Retrieves leaf nodes and merges into parent text when enough siblings match."""
    leaf_nodes = [n for n in nodes.values() if n.level == LEAF_LEVEL]
    leaf_embs = encoder.encode([n.text for n in leaf_nodes])
    top_leaves = [leaf_nodes[i] for i in np.argsort(np.dot(leaf_embs, encoder.encode(query)))[::-1][:k]]

    parent_to_children: dict[str | None, list[str]] = {}
    for leaf in top_leaves:
        parent_to_children.setdefault(leaf.parent_id, []).append(leaf.chunk_id)

    final: list[str] = []
    for parent_id, child_ids in parent_to_children.items():
        parent = nodes[parent_id]
        if len(child_ids) / len(parent.children_ids) >= merge_threshold:
            final.append(parent.text)
        else:
            final.extend(nodes[cid].text for cid in child_ids)
    return final

CHILD_CHUNK_SIZE = 400
CHILD_CHUNK_OVERLAP = 50
PARENT_CHUNK_SIZE = 2000
PARENT_CHUNK_OVERLAP = 200

retriever = ParentDocumentRetriever(
    vectorstore=vectorstore, docstore=docstore,
    child_splitter=RecursiveCharacterTextSplitter.from_tiktoken_encoder(
        chunk_size=CHILD_CHUNK_SIZE, chunk_overlap=CHILD_CHUNK_OVERLAP,
    ),
    parent_splitter=RecursiveCharacterTextSplitter.from_tiktoken_encoder(
        chunk_size=PARENT_CHUNK_SIZE, chunk_overlap=PARENT_CHUNK_OVERLAP,
    ),
)

COLLECTION_NAME = "documents"
CITATION_PROMPT = "Answer using only these sources. Cite by number."


def retrieve_with_metadata(
    query: str, encoder: SentenceTransformer, client: QdrantClient, k: int = 5,
) -> str:
    """Retrieves chunks with source metadata and formats a citation-ready LLM prompt."""
    results = client.search(
        collection_name=COLLECTION_NAME, query_vector=encoder.encode(query).tolist(), limit=k,
    )
    context = []
    for i, r in enumerate(results):
        m = r.payload
        context.append(
            f"[Source {i+1}] {m.get('source', '?')} "
            f"(p.{m.get('page_number', '?')}, {m.get('created_at', '?')})\n{r.payload['text']}"
        )
    separator = "\n\n"
    return f"{CITATION_PROMPT}\n\n{separator.join(context)}\n\nQuestion: {query}"

DEFAULT_SIMILARITY_THRESHOLD = 0.92
DEFAULT_TTL_SECONDS = 3600


class SemanticCache:
    """Caches retrieval results keyed by query embedding similarity."""

    def __init__(
        self, redis_client: Redis, encoder: SentenceTransformer,
        threshold: float = DEFAULT_SIMILARITY_THRESHOLD, ttl: int = DEFAULT_TTL_SECONDS,
    ) -> None:
        self.redis = redis_client
        self.encoder = encoder
        self.threshold = threshold
        self.ttl = ttl

    def get(self, query: str) -> list | None:
        """Returns cached results if a similar query exists, else None."""
        ...

    def set(self, query: str, results: list) -> None:
        """Stores query embedding and results with TTL."""
        ...


def retrieve_with_cache(
    query: str, cache: SemanticCache, retriever: HybridRetriever, **kwargs,
) -> dict[str, list | bool]:
    """Checks semantic cache before falling back to live retrieval."""
    cached = cache.get(query)
    if cached:
        return {"results": cached, "from_cache": True}
    response = retriever.search(query, **kwargs)
    cache.set(query, response["results"])
    return {"results": response["results"], "from_cache": False}

import hashlib
import json

import numpy as np
from redis import Redis

DEFAULT_SIMILARITY_THRESHOLD = 0.92
DEFAULT_TTL_SECONDS = 3600
QUERY_ID_LENGTH = 16
CACHE_INDEX_KEY = "cache:query_index"


class SemanticCache:
    """Redis-backed semantic cache with cosine similarity matching."""

    def __init__(
        self, redis_client: Redis, encoder: SentenceTransformer,
        similarity_threshold: float = DEFAULT_SIMILARITY_THRESHOLD,
        ttl_seconds: int = DEFAULT_TTL_SECONDS,
    ) -> None:
        self.redis = redis_client
        self.encoder = encoder
        self.threshold = similarity_threshold
        self.ttl = ttl_seconds

    def _query_id(self, embedding: np.ndarray) -> str:
        """Derives a short deterministic ID from an embedding vector."""
        return hashlib.sha256(embedding.astype(np.float32).tobytes()).hexdigest()[:QUERY_ID_LENGTH]

    def get(self, query: str) -> list | None:
        """Returns cached results for the most similar query, or None on miss."""
        query_emb = self.encoder.encode(query)
        cached_ids = self.redis.smembers(CACHE_INDEX_KEY)
        if not cached_ids:
            return None

        # NOTE: This linear scan is O(n) per lookup — pedagogical only.
        # Production systems should use Redis Stack's vector search (HNSW indexing)
        # or a dedicated vector store for sub-linear lookup.
        max_sim, best_id = -1.0, None
        for qid in cached_ids:
            qid = qid.decode("utf-8")
            cached_emb = json.loads(self.redis.get(f"cache:query:{qid}") or "null")
            if cached_emb is None:
                self.redis.srem(CACHE_INDEX_KEY, qid)
                continue
            sim = np.dot(query_emb, cached_emb) / (np.linalg.norm(query_emb) * np.linalg.norm(cached_emb))
            if sim > max_sim:
                max_sim, best_id = sim, qid

        if max_sim >= self.threshold and best_id:
            results = self.redis.get(f"cache:results:{best_id}")
            return json.loads(results) if results else None
        return None

    def set(self, query: str, results: list) -> None:
        """Stores query embedding and results with TTL."""
        emb = self.encoder.encode(query)
        qid = self._query_id(emb)
        # NOTE: These three calls are non-atomic — production code should use
        # a Redis pipeline or transaction to avoid partial state on failure.
        self.redis.setex(f"cache:query:{qid}", self.ttl, json.dumps(emb.tolist()))
        self.redis.setex(f"cache:results:{qid}", self.ttl, json.dumps(results))
        self.redis.sadd(CACHE_INDEX_KEY, qid)

ENCODER_MODEL = "all-mpnet-base-v2"

model = SentenceTransformer(ENCODER_MODEL)
index = faiss.read_index("index.faiss")


def search(query: str, k: int = 5) -> list[dict[str, str | float]]:
    """Retrieves k nearest chunks from a FAISS index."""
    query_embedding = model.encode(query).reshape(1, -1).astype("float32")
    distances, indices = index.search(query_embedding, k=k)
    return [{"chunk": chunks[i], "distance": float(distances[0][j])} for j, i in enumerate(indices[0])]

def two_stage_search(query: str, stage1_k: int = 50, final_k: int = 5) -> list[str]:
    """Fast ANN retrieval followed by cross-encoder reranking."""
    query_embedding = retriever.encode(query).reshape(1, -1).astype("float32")
    _, indices = index.search(query_embedding, k=stage1_k)
    stage1_chunks = [chunks[i] for i in indices[0]]

    rerank_scores = reranker.predict([[query, c] for c in stage1_chunks])
    return [stage1_chunks[i] for i in np.argsort(rerank_scores)[::-1][:final_k]]

ENCODER_MODEL = "all-mpnet-base-v2"
RERANKER_MODEL = "cross-encoder/ms-marco-MiniLM-L6-v2"
EPSILON = 1e-8


class HybridRetriever:
    """Two-stage retriever: hybrid dense+sparse search followed by cross-encoder reranking."""

    def __init__(self, index_path: str, chunks_path: str) -> None:
        self.retriever = SentenceTransformer(ENCODER_MODEL)
        self.reranker = CrossEncoder(RERANKER_MODEL)
        self.index = faiss.read_index(index_path)
        self.chunks = np.load(chunks_path, allow_pickle=True)
        tokenized_chunks = [chunk.lower().split() for chunk in self.chunks]
        self.bm25 = BM25Okapi(tokenized_chunks)

    def search(
        self, query: str, stage1_k: int = 100, final_k: int = 5, alpha: float = 0.5,
    ) -> list[dict[str, str | float]]:
        """Runs hybrid retrieval (dense + BM25) then cross-encoder reranking."""
        query_embedding = self.retriever.encode(query)
        dense_distances, dense_indices = self.index.search(
            query_embedding.reshape(1, -1).astype("float32"), k=stage1_k)
        # FAISS returns squared L2 distances; sqrt converts to true L2 before normalization
        dense_scores = 1 / (1 + np.sqrt(dense_distances[0]))

        sparse_scores = self.bm25.get_scores(query.lower().split())

        dense_norm = (dense_scores - dense_scores.min()) / (dense_scores.max() - dense_scores.min() + EPSILON)
        sparse_norm = (sparse_scores - sparse_scores.min()) / (sparse_scores.max() - sparse_scores.min() + EPSILON)

        hybrid_scores = np.zeros(len(self.chunks))
        for i, idx in enumerate(dense_indices[0]):
            hybrid_scores[idx] = alpha * dense_norm[i]
        hybrid_scores += (1 - alpha) * sparse_norm

        top_indices = np.argsort(hybrid_scores)[::-1][:stage1_k]
        stage1_chunks = [self.chunks[i] for i in top_indices]

        rerank_scores = self.reranker.predict([[query, c] for c in stage1_chunks])
        reranked = np.argsort(rerank_scores)[::-1][:final_k]
        return [{"chunk": stage1_chunks[i], "rerank_score": float(rerank_scores[i])} for i in reranked]

import logging
import time
from collections import deque

from pydantic import BaseModel

logger = logging.getLogger(__name__)

METRICS_WINDOW = 1000
SUMMARY_WINDOW = 100
MS_PER_SECOND = 1000


class RetrievalMetrics(BaseModel):
    """Latency telemetry for a single retrieval request."""

    query: str
    stage1_latency_ms: float
    stage2_latency_ms: float
    total_latency_ms: float


class ProductionRetriever:
    """Wraps HybridRetriever with query caching, latency tracking, and logging."""

    def __init__(self, index_path: str, chunks_path: str) -> None:
        self._retriever = HybridRetriever(index_path, chunks_path)
        self._metrics: deque[RetrievalMetrics] = deque(maxlen=METRICS_WINDOW)
        self._embed_cache: dict[str, np.ndarray] = {}

    def _embed_query(self, query: str) -> np.ndarray:
        if query not in self._embed_cache:
            self._embed_cache[query] = self._retriever.retriever.encode(query)
        return self._embed_cache[query]

    def search(
        self, query: str, stage1_k: int = 100, final_k: int = 5, alpha: float = 0.5,
    ) -> dict[str, list | dict]:
        """Runs hybrid retrieval and records latency metrics."""
        start = time.perf_counter()

        s1_start = time.perf_counter()
        self._embed_query(query)
        results = self._retriever.search(query, stage1_k=stage1_k, final_k=final_k, alpha=alpha)
        s1_ms = (time.perf_counter() - s1_start) * MS_PER_SECOND

        total_ms = (time.perf_counter() - start) * MS_PER_SECOND
        self._metrics.append(RetrievalMetrics(
            query=query, stage1_latency_ms=s1_ms,
            stage2_latency_ms=0, total_latency_ms=total_ms))
        logger.info("Retrieved %d chunks in %.1fms", final_k, total_ms)
        return {"results": results, "metrics": {"total_latency_ms": total_ms}}

    def get_performance_summary(self) -> dict[str, float]:
        """Returns average and p95 latency over the most recent requests."""
        recent = list(self._metrics)[-SUMMARY_WINDOW:]
        latencies = [m.total_latency_ms for m in recent]
        return {
            "avg_total_latency_ms": float(np.mean(latencies)),
            "p95_total_latency_ms": float(np.percentile(latencies, 95)),
        }