from openai import OpenAI

client = OpenAI()

async def rag_query(question: str, top_k: int = 5) -> dict:
    """Basic RAG with citations."""
    docs = await vector_db.search(question, limit=top_k)
    context = "\n\n".join([f"[{i+1}] {doc.text}" for i, doc in enumerate(docs)])

    response = await llm.chat([
        {"role": "system", "content": "Answer with inline citations [1], [2]. Use ONLY provided context."},
        {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}
    ])

    return {"answer": response.content, "sources": [d.metadata['source'] for d in docs]}

from pydantic import BaseModel, Field
from typing import Literal

class QueryRoute(BaseModel):
    strategy: Literal["direct", "hyde", "decompose", "web"] = Field(
        description="Best retrieval strategy for this query"
    )
    reasoning: str

async def route_query(question: str) -> str:
    route = await llm.with_structured_output(QueryRoute).ainvoke(
        f"Choose the best retrieval strategy for: {question}\n"
        "- direct: Simple factual queries with clear keywords\n"
        "- hyde: Abstract/conceptual queries with vocabulary mismatch\n"
        "- decompose: Multi-concept queries spanning multiple topics\n"
        "- web: Recent events or data not in knowledge base"
    )
    return route.strategy

async def adaptive_search(question: str, top_k: int = 10) -> list[dict]:
    strategy = await route_query(question)

    if strategy == "direct":
        return await retriever.search(question, top_k=top_k)
    elif strategy == "hyde":
        hyde_result = await hyde_service.generate(question)
        return await retriever.search_by_embedding(hyde_result.embedding, top_k=top_k)
    elif strategy == "decompose":
        return await decomposed_search(question, top_k=top_k)
    elif strategy == "web":
        return await web_search(question)

async def search(question: str) -> list[dict]:
    # Always uses HyDE regardless of query type
    hyde_result = await hyde_service.generate(question)
    return await retriever.search_by_embedding(hyde_result.embedding, top_k=10)

async def adaptive_search(question: str) -> list[dict]:
    strategy = await route_query(question)  # Choose best approach

    if strategy == "direct":
        return await retriever.search(question, top_k=10)  # Fast path
    elif strategy == "hyde":
        hyde_result = await hyde_service.generate(question)
        return await retriever.search_by_embedding(hyde_result.embedding, top_k=10)

def build_crag_workflow() -> StateGraph:
    workflow = StateGraph(RAGState)

    workflow.add_node("retrieve", retrieve)
    workflow.add_node("grade", grade_documents)
    workflow.add_node("generate", generate)
    workflow.add_node("web_search", web_search)
    workflow.add_node("transform_query", transform_query)

    workflow.add_edge(START, "retrieve")
    workflow.add_edge("retrieve", "grade")

    workflow.add_conditional_edges("grade", route_after_grading, {
        "generate": "generate",
        "transform_query": "transform_query",
        "web_search": "web_search"
    })

    workflow.add_edge("transform_query", "retrieve")  # Retry
    workflow.add_edge("web_search", "generate")
    workflow.add_edge("generate", END)

    return workflow.compile()

def route_after_grading(state: RAGState) -> str:
    if state["web_search_needed"]:
        if state.get("retry_count", 0) < 2:
            return "transform_query"
        return "web_search"
    return "generate"

def web_search(state: RAGState) -> dict:
    web_results = tavily_client.search(state["question"], max_results=5, search_depth="advanced")
    web_docs = [
        Document(page_content=r["content"], metadata={"source": r["url"], "type": "web"})
        for r in web_results
    ]
    return {"documents": web_docs, "web_search_needed": False}

def route_after_grading(state: RAGState) -> str:
    if state["web_search_needed"]:
        return "transform_query"  # Infinite loop possible!
    return "generate"

def route_after_grading(state: RAGState) -> str:
    if state["web_search_needed"]:
        if state.get("retry_count", 0) < 2:  # Max 2 retries
            return "transform_query"
        return "web_search"  # Fallback to web
    return "generate"

Query → [Entity Extraction] → [KG Lookup] → [Vector Search] → [Merge] → [Generate]

async def search(query: str) -> list[dict]:
    # Misses relationships between entities
    return await vector_db.search(query, limit=10)

async def graph_rag_search(query: str) -> list[dict]:
    entities = await extract_entities(query)  # Extract entities from query
    kg_results = await neo4j.lookup_entities(entities)  # KG lookup
    vector_results = await vector_db.search(query, limit=10)  # Vector search
    return merge_results(kg_results, vector_results)  # Combine both

from langgraph.graph import StateGraph, START, END
from typing import TypedDict, List, Annotated
from langchain_core.documents import Document
import operator

class RAGState(TypedDict):
    question: str
    documents: Annotated[List[Document], operator.add]
    generation: str
    web_search_needed: bool
    retry_count: int
    relevance_scores: dict[str, float]

from pydantic import BaseModel, Field

class GradeDocuments(BaseModel):
    binary_score: str = Field(description="Relevance score 'yes' or 'no'")

def grade_documents(state: RAGState) -> dict:
    """Grade documents for relevance — core Self-RAG pattern."""
    question = state["question"]
    documents = state["documents"]
    filtered_docs, relevance_scores = [], {}

    for doc in documents:
        score = retrieval_grader.invoke({
            "question": question, "document": doc.page_content
        })
        doc_id = doc.metadata.get("id", hash(doc.page_content))
        relevance_scores[doc_id] = 1.0 if score.binary_score == "yes" else 0.0
        if score.binary_score == "yes":
            filtered_docs.append(doc)

    web_search_needed = len(filtered_docs) < len(documents) // 2
    return {
        "documents": filtered_docs,
        "web_search_needed": web_search_needed,
        "relevance_scores": relevance_scores
    }

def generate(state: RAGState) -> dict:
    # Uses all retrieved docs without quality check
    context = "\n\n".join([d.page_content for d in state["documents"]])
    return {"generation": llm.invoke(context)}

def grade_documents(state: RAGState) -> dict:
    filtered_docs = []
    for doc in state["documents"]:
        score = grader.invoke({"question": state["question"], "document": doc.page_content})
        if score.binary_score == "yes":  # Only keep relevant docs
            filtered_docs.append(doc)

    web_search_needed = len(filtered_docs) < len(state["documents"]) // 2
    return {"documents": filtered_docs, "web_search_needed": web_search_needed}

from rank_bm25 import BM25Okapi
import numpy as np

class HybridRetriever:
    def __init__(self, chunks: list[str], embeddings: np.ndarray):
        self.chunks = chunks
        self.embeddings = embeddings
        tokenized = [c.lower().split() for c in chunks]
        self.bm25 = BM25Okapi(tokenized)

    def search(
        self, query: str, query_embedding: np.ndarray,
        top_k: int = 20, bm25_weight: float = 0.4, vector_weight: float = 0.6
    ) -> list[tuple[int, float]]:
        bm25_scores = self.bm25.get_scores(query.lower().split())
        bm25_scores = (bm25_scores - bm25_scores.min()) / (bm25_scores.max() - bm25_scores.min() + 1e-6)

        vector_scores = np.dot(self.embeddings, query_embedding)
        vector_scores = (vector_scores - vector_scores.min()) / (vector_scores.max() - vector_scores.min() + 1e-6)

        combined = bm25_weight * bm25_scores + vector_weight * vector_scores
        top_indices = np.argsort(combined)[::-1][:top_k]
        return [(i, combined[i]) for i in top_indices]

def search(query: str, query_embedding: np.ndarray, top_k: int = 20) -> list[int]:
    # Misses exact-match queries
    vector_scores = np.dot(self.embeddings, query_embedding)
    return np.argsort(vector_scores)[::-1][:top_k]

def search(query: str, query_embedding: np.ndarray, top_k: int = 20) -> list[tuple[int, float]]:
    bm25_scores = self.bm25.get_scores(query.lower().split())
    bm25_norm = (bm25_scores - bm25_scores.min()) / (bm25_scores.max() - bm25_scores.min() + 1e-6)

    vector_scores = np.dot(self.embeddings, query_embedding)
    vector_norm = (vector_scores - vector_scores.min()) / (vector_scores.max() - vector_scores.min() + 1e-6)

    combined = 0.4 * bm25_norm + 0.6 * vector_norm  # Research-backed weights
    return [(i, combined[i]) for i in np.argsort(combined)[::-1][:top_k]]

from dataclasses import dataclass

@dataclass
class ContextualChunk:
    original: str
    contextualized: str
    embedding: list[float]
    doc_id: str
    chunk_index: int

class ContextualRetriever:
    def __init__(self, embed_model, llm_client):
        self.embed_model = embed_model
        self.llm = llm_client
        self.chunks: list[ContextualChunk] = []

    def add_document(self, doc_id: str, text: str, chunk_size: int = 512):
        raw_chunks = self._chunk_text(text, chunk_size)
        contextualized = self._contextualize_batch(text, raw_chunks)
        embeddings = self.embed_model.embed(contextualized)

        for i, (raw, ctx, emb) in enumerate(zip(raw_chunks, contextualized, embeddings)):
            self.chunks.append(ContextualChunk(
                original=raw, contextualized=ctx, embedding=emb,
                doc_id=doc_id, chunk_index=i
            ))
        self._rebuild_bm25()

    def search(self, query: str, top_k: int = 10) -> list[ContextualChunk]:
        query_emb = self.embed_model.embed([query])[0]
        bm25_scores = self.bm25.get_scores(query.lower().split())
        embeddings = np.array([c.embedding for c in self.chunks])
        vector_scores = np.dot(embeddings, query_emb)
        combined = 0.4 * self._normalize(bm25_scores) + 0.6 * self._normalize(vector_scores)
        top_indices = np.argsort(combined)[::-1][:top_k]
        return [self.chunks[i] for i in top_indices]

async def contextualize_parallel(document: str, chunks: list[str]) -> list[str]:
    semaphore = asyncio.Semaphore(10)  # Max 10 concurrent
    async def process_chunk(chunk: str) -> str:
        async with semaphore:
            context = await async_generate_context(document, chunk)
            return f"{context}\n\n{chunk}"
    return await asyncio.gather(*[process_chunk(c) for c in chunks])

def add_document(self, doc_id: str, text: str):
    # No contextualization, no hybrid indexing
    raw_chunks = self._chunk_text(text, 512)
    embeddings = self.embed_model.embed(raw_chunks)
    self.chunks.extend(embeddings)

def add_document(self, doc_id: str, text: str, chunk_size: int = 512):
    raw_chunks = self._chunk_text(text, chunk_size)
    contextualized = self._contextualize_batch(text, raw_chunks)  # Add context
    embeddings = self.embed_model.embed(contextualized)  # Embed with context

    for i, (raw, ctx, emb) in enumerate(zip(raw_chunks, contextualized, embeddings)):
        self.chunks.append(ContextualChunk(
            original=raw, contextualized=ctx, embedding=emb,
            doc_id=doc_id, chunk_index=i
        ))
    self._rebuild_bm25()  # Hybrid BM25 + vector

Original: "ACME Q3 2024 Earnings Report..."
Chunk: "Revenue increased 15% compared to the previous quarter."
Query: "What was ACME's Q3 2024 revenue growth?"
Result: Chunk doesn't mention "ACME" or "Q3 2024" — retrieval fails

import anthropic
client = anthropic.Anthropic()

CONTEXT_PROMPT = """
<document>
{document}
</document>

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

Please give a short, succinct context (1-2 sentences) to situate this chunk
within the overall document. Focus on information that would help retrieval.
Answer only with the context, nothing else.
"""

def contextualize_chunk(document: str, chunk: str) -> str:
    response = client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=150,
        messages=[{"role": "user",
                   "content": CONTEXT_PROMPT.format(document=document, chunk=chunk)}]
    )
    return f"{response.content[0].text}\n\n{chunk}"

def contextualize_cached(document: str, chunk: str) -> str:
    response = client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=150,
        messages=[{"role": "user", "content": [
            {"type": "text", "text": f"<document>\n{document}\n</document>",
             "cache_control": {"type": "ephemeral"}},
            {"type": "text", "text": f"Situate this chunk (1-2 sentences):\n<chunk>\n{chunk}\n</chunk>"}
        ]}]
    )
    return f"{response.content[0].text}\n\n{chunk}"

def index_chunk(chunk: str) -> str:
    # Missing document context — retrieval will fail
    embedding = embed_model.embed([chunk])[0]
    return embedding

def contextualize_chunk(document: str, chunk: str) -> str:
    context = client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=150,
        messages=[{"role": "user", "content": [
            {"type": "text", "text": f"<document>\n{document}\n</document>",
             "cache_control": {"type": "ephemeral"}},  # Cache for 90% cost reduction
            {"type": "text", "text": f"Situate this chunk (1-2 sentences):\n<chunk>\n{chunk}\n</chunk>"}
        ]}]
    )
    return f"{context.content[0].text}\n\n{chunk}"  # Prepend context

async def rag_query(question: str, top_k: int = 5) -> str:
    """Basic RAG: retrieve then generate."""
    docs = await vector_db.search(question, limit=top_k)

    context = "\n\n".join([
        f"[{i+1}] {doc.text}"
        for i, doc in enumerate(docs)
    ])

    response = await llm.chat([
        {"role": "system", "content":
            "Answer using ONLY the provided context. "
            "If not in context, say 'I don't have that information.'"},
        {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}
    ])

    return response.content

async def rag_with_citations(question: str) -> dict:
    """RAG with inline citations [1], [2], etc."""
    docs = await vector_db.search(question, limit=5)

    context = "\n\n".join([
        f"[{i+1}] {doc.text}\nSource: {doc.metadata['source']}"
        for i, doc in enumerate(docs)
    ])

    response = await llm.chat([
        {"role": "system", "content":
            "Answer with inline citations like [1], [2]. "
            "End with a Sources section."},
        {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}
    ])

    return {
        "answer": response.content,
        "sources": [doc.metadata['source'] for doc in docs]
    }

async def rag_query(question: str) -> str:
    docs = await vector_db.search(question, limit=5)
    context = "\n\n".join([doc.text for doc in docs])  # No citations
    response = await llm.chat([
        {"role": "user", "content": f"{context}\n\n{question}"}  # No grounding instruction
    ])
    return response.content  # No source tracking

async def rag_with_citations(question: str) -> dict:
    docs = await vector_db.search(question, limit=5)
    context = "\n\n".join([
        f"[{i+1}] {doc.text}\nSource: {doc.metadata['source']}"  # Numbered citations
        for i, doc in enumerate(docs)
    ])

    response = await llm.chat([
        {"role": "system", "content": "Answer with inline citations like [1], [2]. Use ONLY the provided context."},
        {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}
    ])

    return {"answer": response.content, "sources": [doc.metadata['source'] for doc in docs]}

def fit_context(docs: list, max_tokens: int = 6000) -> list:
    """Truncate context to fit token budget."""
    total_tokens = 0
    selected = []

    for doc in docs:
        doc_tokens = count_tokens(doc.text)
        if total_tokens + doc_tokens > max_tokens:
            break
        selected.append(doc)
        total_tokens += doc_tokens

    return selected

from pydantic import BaseModel

class SufficiencyCheck(BaseModel):
    is_sufficient: bool
    confidence: float  # 0.0-1.0
    missing_info: str | None = None

async def rag_with_sufficiency(question: str, top_k: int = 5) -> str:
    """RAG with hallucination prevention via sufficiency check."""
    docs = await vector_db.search(question, limit=top_k)
    context = "\n\n".join([f"[{i+1}] {doc.text}" for i, doc in enumerate(docs)])

    check = await llm.with_structured_output(SufficiencyCheck).ainvoke(
        f"Does this context contain sufficient information to answer?\n"
        f"Question: {question}\nContext:\n{context}"
    )

    if not check.is_sufficient and check.confidence > 0.7:
        return f"I don't have enough information. Missing: {check.missing_info}"

    if not check.is_sufficient and check.confidence <= 0.7:
        more_docs = await vector_db.search(question, limit=top_k * 2)
        context = "\n\n".join([f"[{i+1}] {doc.text}" for i, doc in enumerate(more_docs)])

    return await generate_with_context(question, context)

async def rag_query(question: str) -> str:
    docs = await vector_db.search(question, limit=100)  # No limit!
    context = "\n\n".join([doc.text for doc in docs])  # Could exceed context window
    return await generate_with_context(question, context)  # No sufficiency check

async def rag_with_sufficiency(question: str, top_k: int = 5) -> str:
    docs = await vector_db.search(question, limit=top_k)
    fitted = fit_context(docs, max_tokens=6000)  # Budget enforcement
    context = "\n\n".join([f"[{i+1}] {doc.text}" for i, doc in enumerate(fitted)])

    check = await llm.with_structured_output(SufficiencyCheck).ainvoke(
        f"Does this context contain sufficient information?\nQuestion: {question}\nContext:\n{context}"
    )

    if not check.is_sufficient and check.confidence > 0.7:
        return f"I don't have enough information. Missing: {check.missing_info}"

    return await generate_with_context(question, context)

def reciprocal_rank_fusion(
    semantic_results: list,
    keyword_results: list,
    k: int = 60
) -> list:
    """Combine semantic and keyword search with RRF."""
    scores = {}

    for rank, doc in enumerate(semantic_results):
        scores[doc.id] = scores.get(doc.id, 0) + 1 / (k + rank + 1)

    for rank, doc in enumerate(keyword_results):
        scores[doc.id] = scores.get(doc.id, 0) + 1 / (k + rank + 1)

    ranked_ids = sorted(scores.keys(), key=lambda x: scores[x], reverse=True)
    return [get_doc(id) for id in ranked_ids]

from collections import defaultdict

def multi_rrf(result_lists: list[list[dict]], k: int = 60) -> list[dict]:
    """Combine multiple ranked lists using RRF."""
    scores: defaultdict[str, float] = defaultdict(float)
    docs: dict[str, dict] = {}

    for results in result_lists:
        for rank, doc in enumerate(results, start=1):
            doc_id = doc["id"]
            scores[doc_id] += 1.0 / (k + rank)
            docs[doc_id] = doc

    ranked_ids = sorted(scores.keys(), key=lambda x: scores[x], reverse=True)
    return [docs[doc_id] for doc_id in ranked_ids]

def hybrid_search(query: str, top_k: int = 10) -> list:
    semantic = vector_search(query, top_k)
    keyword = bm25_search(query, top_k)
    # Naive merge without RRF
    return semantic[:5] + keyword[:5]

def reciprocal_rank_fusion(semantic_results: list, keyword_results: list, k: int = 60) -> list:
    scores = {}
    for rank, doc in enumerate(semantic_results):
        scores[doc.id] = scores.get(doc.id, 0) + 1 / (k + rank + 1)

    for rank, doc in enumerate(keyword_results):
        scores[doc.id] = scores.get(doc.id, 0) + 1 / (k + rank + 1)

    ranked_ids = sorted(scores.keys(), key=lambda x: scores[x], reverse=True)
    return [get_doc(id) for id in ranked_ids]

Query → [Decompose?] → [HyDE?] → [Retrieve] → [Rerank] → [Context Fit] → [Generate]

class RAGPipeline:
    """Composable RAG pipeline with optional stages."""

    def __init__(self, retriever, reranker=None, hyde_service=None, decomposer=None):
        self.retriever = retriever
        self.reranker = reranker
        self.hyde = hyde_service
        self.decomposer = decomposer

    async def query(self, question: str, top_k: int = 10) -> list[dict]:
        # Stage 1: Query enhancement (optional)
        queries = [question]
        if self.decomposer:
            concepts = await self.decomposer.decompose(question)
            if len(concepts) > 1:
                queries = concepts

        # Stage 2: Retrieve (with optional HyDE)
        all_results = []
        for q in queries:
            if self.hyde:
                hyde_result = await self.hyde.generate(q)
                results = await self.retriever.search_by_embedding(hyde_result.embedding, top_k=top_k * 3)
            else:
                results = await self.retriever.search(q, top_k=top_k * 3)
            all_results.append(results)

        # Stage 3: Fuse if multiple queries
        if len(all_results) > 1:
            merged = reciprocal_rank_fusion(all_results)
        else:
            merged = all_results[0]

        # Stage 4: Rerank (optional)
        if self.reranker:
            merged = await self.reranker.rerank(question, merged, top_k=top_k)

        return merged[:top_k]

async def query(question: str) -> list[dict]:
    # No optional stages, fixed pipeline
    docs = await retriever.search(question, top_k=10)
    return docs

async def query(self, question: str, top_k: int = 10) -> list[dict]:
    queries = [question]
    if self.decomposer:  # Optional decomposition
        concepts = await self.decomposer.decompose(question)
        if len(concepts) > 1:
            queries = concepts

    all_results = []
    for q in queries:
        if self.hyde:  # Optional HyDE
            hyde_result = await self.hyde.generate(q)
            results = await self.retriever.search_by_embedding(hyde_result.embedding, top_k * 3)
        else:
            results = await self.retriever.search(q, top_k * 3)
        all_results.append(results)

    merged = reciprocal_rank_fusion(all_results) if len(all_results) > 1 else all_results[0]

    if self.reranker:  # Optional reranking
        merged = await self.reranker.rerank(question, merged, top_k)

    return merged[:top_k]

import hashlib
import json
import redis

class EmbeddingCache:
    def __init__(self, redis_client: redis.Redis, ttl: int = 86400):
        self.redis = redis_client
        self.ttl = ttl

    def _key(self, text: str, model: str) -> str:
        h = hashlib.md5(f"{model}:{text}".encode()).hexdigest()
        return f"emb:{h}"

    async def get_or_embed(self, text: str, model: str, embed_fn) -> list[float]:
        key = self._key(text, model)
        cached = self.redis.get(key)
        if cached:
            return json.loads(cached)
        embedding = await embed_fn(text)
        self.redis.setex(key, self.ttl, json.dumps(embedding))
        return embedding

import asyncio

async def batch_embed(texts: list[str], embed_fn, batch_size: int = 100) -> list[list[float]]:
    """Embed texts in batches with rate limiting."""
    results = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        embeddings = await embed_fn(batch)
        results.extend(embeddings)
        if i + batch_size < len(texts):
            await asyncio.sleep(0.1)  # Rate limit courtesy
    return results

# text-embedding-3 models support Matryoshka embeddings
# Truncate to fewer dimensions with minimal quality loss
response = client.embeddings.create(
    model="text-embedding-3-large",
    input="Your text",
    dimensions=1536  # Reduce from 3072 to 1536 (saves 50% storage)
)

async def embed_texts(texts: list[str]) -> list[list[float]]:
    results = []
    for text in texts:  # One API call per text!
        embedding = await client.embeddings.create(
            model="text-embedding-3-large",
            input=text
        )
        results.append(embedding.data[0].embedding)
    return results

async def batch_embed(texts: list[str], batch_size: int = 100) -> list[list[float]]:
    results = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        # Check cache first
        cached_keys = [cache.get(text) for text in batch]
        uncached = [t for t, c in zip(batch, cached_keys) if not c]

        if uncached:
            embeddings = await client.embeddings.create(
                model="text-embedding-3-large",
                input=uncached,
                dimensions=1536  # Matryoshka reduction
            )
            for text, emb in zip(uncached, embeddings.data):
                cache.set(text, emb.embedding)  # Cache for reuse

        results.extend([cached or cache.get(t) for t, cached in zip(batch, cached_keys)])
        await asyncio.sleep(0.1)  # Rate limiting
    return results

def chunk_text(text: str, chunk_size: int = 512, overlap: int = 50) -> list[str]:
    """Split text into overlapping chunks for embedding."""
    words = text.split()
    chunks = []

    for i in range(0, len(words), chunk_size - overlap):
        chunk = " ".join(words[i:i + chunk_size])
        if chunk:
            chunks.append(chunk)

    return chunks

CHUNK_CONFIG = {
    "target_tokens": 500,      # ~400-600 tokens per chunk
    "max_tokens": 800,         # Hard limit
    "overlap_tokens": 75,      # ~15% overlap
    "boundary_markers": [      # Prefer splitting at:
        "\n## ",               # H2 headers
        "\n### ",              # H3 headers
        "\n\n",                # Paragraphs
        ". ",                  # Sentences (last resort)
    ]
}

def chunk_by_sentences(text: str, chunk_size: int = 512) -> list[str]:
    sentences = text.split('. ')
    chunks, current, current_len = [], [], 0

    for sent in sentences:
        if current_len + len(sent) > chunk_size and current:
            chunks.append('. '.join(current) + '.')
            current, current_len = [sent], len(sent)
        else:
            current.append(sent)
            current_len += len(sent)

    if current:
        chunks.append('. '.join(current))
    return chunks

def chunk_text(text: str) -> list[str]:
    # Arbitrary splits, no overlap, breaks mid-sentence
    return [text[i:i+500] for i in range(0, len(text), 500)]

def chunk_by_sentences(text: str, chunk_size: int = 512, overlap: int = 75) -> list[str]:
    sentences = text.split('. ')
    chunks, current, current_len = [], [], 0

    for sent in sentences:
        if current_len + len(sent) > chunk_size and current:
            chunk_text = '. '.join(current) + '.'
            chunks.append(chunk_text)
            # Keep last few sentences for overlap
            overlap_sents = current[-2:] if len(current) > 2 else current
            current, current_len = overlap_sents, sum(len(s) for s in overlap_sents)
        else:
            current.append(sent)
            current_len += len(sent)

    if current:
        chunks.append('. '.join(current))
    return chunks

from openai import OpenAI
client = OpenAI()

# Single text embedding
response = client.embeddings.create(
    model="text-embedding-3-small",
    input="Your text here"
)
vector = response.data[0].embedding  # 1536 dimensions

# Batch embedding (efficient)
texts = ["text1", "text2", "text3"]
response = client.embeddings.create(
    model="text-embedding-3-small",
    input=texts
)
vectors = [item.embedding for item in response.data]

import numpy as np

def cosine_similarity(a: list[float], b: list[float]) -> float:
    a, b = np.array(a), np.array(b)
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
# 1.0 = identical, 0.0 = orthogonal

# Index with one model
docs_embeddings = client.embeddings.create(
    model="text-embedding-3-large",  # 3072 dims
    input=documents
)

# Query with different model
query_embedding = client.embeddings.create(
    model="text-embedding-3-small",  # 1536 dims - MISMATCH!
    input=query
)
# Results will be nonsensical due to dimension mismatch

MODEL = "text-embedding-3-small"  # Use same model everywhere

# Index
docs_embeddings = client.embeddings.create(model=MODEL, input=documents)

# Query
query_embedding = client.embeddings.create(model=MODEL, input=query)

# Now cosine similarity is meaningful
similarity = cosine_similarity(query_embedding, docs_embeddings[0])

async def hyde_with_fallback(
    query: str,
    hyde_service: HyDEService,
    embed_fn: callable,
    timeout: float = 3.0,
) -> list[float]:
    """HyDE with fallback to direct embedding on timeout."""
    try:
        async with asyncio.timeout(timeout):
            result = await hyde_service.generate(query)
            return result.embedding
    except TimeoutError:
        # Fallback to direct query embedding
        return await embed_fn(query)

async def hyde_search(query: str) -> list[float]:
    # No timeout! May hang indefinitely
    result = await hyde_service.generate(query)
    return result.embedding

async def hyde_with_fallback(query: str, timeout: float = 3.0) -> list[float]:
    try:
        async with asyncio.timeout(timeout):
            result = await hyde_service.generate(query)
            return result.embedding
    except TimeoutError:
        # Fallback to direct query embedding
        return await embed_fn(query)

Query: "scaling async data pipelines"
Docs use: "event-driven messaging", "Apache Kafka", "message brokers"
-> Low similarity scores despite high relevance

from openai import AsyncOpenAI
from pydantic import BaseModel

class HyDEResult(BaseModel):
    original_query: str
    hypothetical_doc: str
    embedding: list[float]

async def generate_hyde(
    query: str, llm: AsyncOpenAI, embed_fn: callable, max_tokens: int = 150
) -> HyDEResult:
    """Generate hypothetical document and embed it."""
    response = await llm.chat.completions.create(
        model="gpt-5.2-mini",
        messages=[
            {"role": "system", "content":
                "Write a short paragraph that would answer this query. "
                "Use technical terminology that documentation would use."},
            {"role": "user", "content": query}
        ],
        max_tokens=max_tokens,
        temperature=0.3,
    )

    hypothetical_doc = response.choices[0].message.content
    embedding = await embed_fn(hypothetical_doc)  # Embed the hypothetical doc, not the query

    return HyDEResult(
        original_query=query,
        hypothetical_doc=hypothetical_doc,
        embedding=embedding,
    )

async def generate_hyde(query: str) -> HyDEResult:
    response = await llm.chat.completions.create(
        model="gpt-5.2-mini",
        messages=[{"role": "user", "content": query}],
        max_tokens=150
    )
    hypothetical_doc = response.choices[0].message.content
    embedding = await embed_fn(query)  # WRONG: Embeds query, not hypothetical doc!
    return HyDEResult(query, hypothetical_doc, embedding)

async def generate_hyde(query: str) -> HyDEResult:
    response = await llm.chat.completions.create(
        model="gpt-5.2-mini",
        messages=[
            {"role": "system", "content": "Write a short paragraph that would answer this query."},
            {"role": "user", "content": query}
        ],
        max_tokens=150,
        temperature=0.3
    )

    hypothetical_doc = response.choices[0].message.content
    embedding = await embed_fn(hypothetical_doc)  # Embed the hypothetical doc!

    return HyDEResult(query, hypothetical_doc, embedding)

async def batch_hyde(
    concepts: list[str], hyde_service: HyDEService
) -> list[HyDEResult]:
    """Generate HyDE embeddings for multiple concepts in parallel."""
    tasks = [hyde_service.generate(concept) for concept in concepts]
    return await asyncio.gather(*tasks)

class HyDEService:
    def __init__(self, llm, embed_fn):
        self.llm = llm
        self.embed_fn = embed_fn
        self._cache: dict[str, HyDEResult] = {}

    def _cache_key(self, query: str) -> str:
        return hashlib.md5(query.lower().strip().encode()).hexdigest()

    async def generate(self, query: str) -> HyDEResult:
        key = self._cache_key(query)
        if key in self._cache:
            return self._cache[key]
        result = await generate_hyde(query, self.llm, self.embed_fn)
        self._cache[key] = result
        return result

async def batch_hyde(concepts: list[str]) -> list[HyDEResult]:
    results = []
    for concept in concepts:  # Sequential! Slow for many concepts
        result = await hyde_service.generate(concept)
        results.append(result)
    return results

async def batch_hyde(concepts: list[str]) -> list[HyDEResult]:
    # Parallel generation for all concepts simultaneously
    tasks = [hyde_service.generate(concept) for concept in concepts]
    return await asyncio.gather(*tasks)

from dataclasses import dataclass
from typing import Literal, Optional

@dataclass
class Chunk:
    content: str
    chunk_type: Literal["text", "image", "table", "chart"]
    page: int
    image_path: Optional[str] = None
    embedding: Optional[list[float]] = None

def chunk_multimodal_document(pdf_path: str) -> list[Chunk]:
    import fitz  # PyMuPDF
    doc = fitz.open(pdf_path)
    chunks = []

    for page_num, page in enumerate(doc):
        text_blocks = page.get_text("blocks")
        current_text = ""

        for block in text_blocks:
            if block[6] == 0:  # Text block
                current_text += block[4] + "\n"
            else:  # Image block
                if current_text.strip():
                    chunks.append(Chunk(content=current_text.strip(), chunk_type="text", page=page_num))
                    current_text = ""
                xref = block[7]
                img = doc.extract_image(xref)
                img_path = f"/tmp/page{page_num}_img{xref}.{img['ext']}"
                with open(img_path, "wb") as f:
                    f.write(img["image"])
                caption = generate_image_caption(img_path)
                chunks.append(Chunk(content=caption, chunk_type="image", page=page_num, image_path=img_path))

        if current_text.strip():
            chunks.append(Chunk(content=current_text.strip(), chunk_type="text", page=page_num))

    return chunks

def chunk_pdf(pdf_path: str) -> list[str]:
    import fitz
    doc = fitz.open(pdf_path)
    chunks = []
    for page in doc:
        chunks.append(page.get_text())  # Text only, images lost!
    return chunks

def chunk_multimodal_document(pdf_path: str) -> list[Chunk]:
    import fitz
    doc = fitz.open(pdf_path)
    chunks = []

    for page_num, page in enumerate(doc):
        text_blocks = page.get_text("blocks")
        for block in text_blocks:
            if block[6] == 0:  # Text block
                chunks.append(Chunk(content=block[4], chunk_type="text", page=page_num))
            else:  # Image block
                xref = block[7]
                img = doc.extract_image(xref)
                img_path = f"/tmp/page{page_num}_img{xref}.{img['ext']}"
                with open(img_path, "wb") as f:
                    f.write(img["image"])
                caption = generate_image_caption(img_path)
                chunks.append(Chunk(content=caption, chunk_type="image", page=page_num, image_path=img_path))

    return chunks

from transformers import CLIPProcessor, CLIPModel
import torch

model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")

def embed_image(image_path: str) -> list[float]:
    image = Image.open(image_path)
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        embeddings = model.get_image_features(**inputs)
    embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)
    return embeddings[0].tolist()

def embed_text(text: str) -> list[float]:
    inputs = processor(text=[text], return_tensors="pt", padding=True)
    with torch.no_grad():
        embeddings = model.get_text_features(**inputs)
    embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)
    return embeddings[0].tolist()

import voyageai
client = voyageai.Client()

def embed_multimodal(texts=None, images=None) -> list[list[float]]:
    inputs = []
    if texts:
        inputs.extend([{"type": "text", "content": t} for t in texts])
    if images:
        for path in images:
            with open(path, "rb") as f:
                b64 = base64.b64encode(f.read()).decode()
                inputs.append({"type": "image", "content": f"data:image/png;base64,{b64}"})
    return client.multimodal_embed(inputs=inputs, model="voyage-multimodal-3").embeddings

def embed_image(image_path: str) -> list[float]:
    # Using text embedding model for images - wrong modality!
    caption = generate_caption(image_path)
    return text_embed_model.embed([caption])[0]  # Loses visual features

def embed_image(image_path: str) -> list[float]:
    image = Image.open(image_path)
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        embeddings = model.get_image_features(**inputs)
    embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)  # Normalize
    return embeddings[0].tolist()

class MultimodalRAG:
    def __init__(self, vector_db, vision_model, text_model):
        self.vector_db = vector_db
        self.vision_model = vision_model
        self.text_model = text_model

    async def retrieve(self, query: str, query_image: str = None, top_k: int = 10) -> list[dict]:
        results = []
        text_emb = embed_text(query)
        text_results = await self.vector_db.search(embedding=text_emb, top_k=top_k)
        results.extend(text_results)

        if query_image:
            img_emb = embed_image(query_image)
            img_results = await self.vector_db.search(embedding=img_emb, top_k=top_k)
            results.extend(img_results)

        # Dedupe by doc_id, keep highest score
        seen = {}
        for r in results:
            doc_id = r["metadata"]["doc_id"]
            if doc_id not in seen or r["score"] > seen[doc_id]["score"]:
                seen[doc_id] = r
        return sorted(seen.values(), key=lambda x: x["score"], reverse=True)[:top_k]

async def generate_with_context(query: str, chunks: list[Chunk], model: str = "claude-opus-4-6") -> str:
    content = []
    # Add images first (attention positioning)
    for chunk in chunks:
        if chunk.chunk_type == "image" and chunk.image_path:
            b64, media_type = encode_image_base64(chunk.image_path)
            content.append({"type": "image", "source": {"type": "base64", "media_type": media_type, "data": b64}})
    # Add text context
    text_context = "\n\n".join([f"[Page {c.page}]: {c.content}" for c in chunks if c.chunk_type == "text"])
    content.append({"type": "text", "text": f"Context:\n{text_context}\n\nQuestion: {query}"})

    response = client.messages.create(model=model, max_tokens=4096, messages=[{"role": "user", "content": content}])
    return response.content[0].text

async def retrieve(self, query: str, top_k: int = 10) -> list[dict]:
    text_emb = embed_text(query)
    text_results = await self.vector_db.search(embedding=text_emb, top_k=top_k)

    img_emb = embed_image(query_image)
    img_results = await self.vector_db.search(embedding=img_emb, top_k=top_k)

    # No deduplication! Same doc may appear twice
    return text_results + img_results

async def retrieve(self, query: str, query_image: str = None, top_k: int = 10) -> list[dict]:
    results = []
    text_emb = embed_text(query)
    text_results = await self.vector_db.search(embedding=text_emb, top_k=top_k)
    results.extend(text_results)

    if query_image:
        img_emb = embed_image(query_image)
        img_results = await self.vector_db.search(embedding=img_emb, top_k=top_k)
        results.extend(img_results)

    # Dedupe by doc_id, keep highest score
    seen = {}
    for r in results:
        doc_id = r["metadata"]["doc_id"]
        if doc_id not in seen or r["score"] > seen[doc_id]["score"]:
            seen[doc_id] = r
    return sorted(seen.values(), key=lambda x: x["score"], reverse=True)[:top_k]

async def hybrid_search(query: str, query_embedding: list[float], top_k: int = 10) -> list[Chunk]:
    FETCH_MULTIPLIER = 3
    K = 60  # RRF smoothing constant

    # Vector search subquery
    vector_subq = (
        select(Chunk.id,
            func.row_number().over(
                order_by=Chunk.embedding.cosine_distance(query_embedding)
            ).label("vector_rank"))
        .limit(top_k * FETCH_MULTIPLIER)
        .subquery()
    )

    # Keyword search subquery
    ts_query = func.plainto_tsquery("english", query)
    keyword_subq = (
        select(Chunk.id,
            func.row_number().over(
                order_by=func.ts_rank_cd(Chunk.content_tsvector, ts_query).desc()
            ).label("keyword_rank"))
        .where(Chunk.content_tsvector.op("@@")(ts_query))
        .limit(top_k * FETCH_MULTIPLIER)
        .subquery()
    )

    # RRF fusion with FULL OUTER JOIN
    rrf_subq = (
        select(
            func.coalesce(vector_subq.c.id, keyword_subq.c.id).label("chunk_id"),
            (func.coalesce(1.0 / (K + vector_subq.c.vector_rank), 0.0) +
             func.coalesce(1.0 / (K + keyword_subq.c.keyword_rank), 0.0)
            ).label("rrf_score"))
        .select_from(vector_subq.outerjoin(keyword_subq, ..., full=True))
        .order_by("rrf_score DESC")
        .limit(top_k)
        .subquery()
    )

    return await session.execute(
        select(Chunk).join(rrf_subq, Chunk.id == rrf_subq.c.chunk_id)
    )

rrf_score = 1/(k + vector_rank) + 1/(k + keyword_rank)  # k=60 (standard)

async def hybrid_search(query: str, embedding: list[float], top_k: int = 10):
    # Separate queries, no fusion!
    vector_results = await session.execute(
        select(Chunk).order_by(Chunk.embedding.cosine_distance(embedding)).limit(top_k)
    )
    keyword_results = await session.execute(
        select(Chunk).where(Chunk.content_tsvector.op("@@")(plainto_tsquery(query))).limit(top_k)
    )
    # Naive merge, no RRF
    return list(vector_results) + list(keyword_results)

async def hybrid_search(query: str, query_embedding: list[float], top_k: int = 10):
    K = 60  # RRF smoothing constant

    # Vector search subquery
    vector_subq = select(Chunk.id, func.row_number().over(
        order_by=Chunk.embedding.cosine_distance(query_embedding)
    ).label("vector_rank")).limit(top_k * 3).subquery()

    # Keyword search subquery
    keyword_subq = select(Chunk.id, func.row_number().over(
        order_by=func.ts_rank_cd(Chunk.content_tsvector, plainto_tsquery(query)).desc()
    ).label("keyword_rank")).limit(top_k * 3).subquery()

    # RRF fusion with FULL OUTER JOIN
    rrf_subq = select(
        func.coalesce(vector_subq.c.id, keyword_subq.c.id).label("chunk_id"),
        (func.coalesce(1.0 / (K + vector_subq.c.vector_rank), 0.0) +
         func.coalesce(1.0 / (K + keyword_subq.c.keyword_rank), 0.0)).label("rrf_score")
    ).select_from(vector_subq.outerjoin(keyword_subq, ..., full=True)).order_by("rrf_score DESC").limit(top_k)

    return await session.execute(select(Chunk).join(rrf_subq, Chunk.id == rrf_subq.c.chunk_id))

-- Create HNSW index
CREATE INDEX idx_chunks_embedding_hnsw ON chunks
    USING hnsw (embedding vector_cosine_ops)
    WITH (m = 16, ef_construction = 64);

-- Query-time tuning
SET hnsw.ef_search = 40;  -- Higher = better recall, slower

-- Iterative scan for filtered queries (pgvector 0.8+)
SET hnsw.iterative_scan = 'relaxed_order';

-- No index! Sequential scan is 17x slower
SELECT * FROM chunks
ORDER BY embedding <=> '[0.1, 0.2, ...]'::vector
LIMIT 10;