# rag/chunking.py

from abc import ABC, abstractmethod

from dataclasses import dataclass



@dataclass

class Chunk:

    text: str

    metadata: dict

    token_count: int

    chunk_id: str



class ChunkingStrategy(ABC):

    @abstractmethod

    def chunk(self, document: str, metadata: dict) -> list[Chunk]: ...



class FixedTokenChunker(ChunkingStrategy):

    """Baseline. Good for homogeneous content."""

    def __init__(self, chunk_size=512, overlap=50):

        self.chunk_size = chunk_size

        self.overlap = overlap



class RecursiveChunker(ChunkingStrategy):

    """Respects structure: markdown headers, code blocks, paragraphs."""

    def __init__(self, separators=["\n## ", "\n### ", "\n\n", "\n", " "], chunk_size=512):

        self.separators = separators

        self.chunk_size = chunk_size



class SemanticChunker(ChunkingStrategy):

    """Uses embedding similarity to find natural boundaries."""

    def __init__(self, model="text-embedding-3-small", threshold=0.7):

        self.model = model

        self.threshold = threshold



class AgenticChunker(ChunkingStrategy):

    """LLM decides boundaries. Expensive but highest quality for complex docs."""

    def __init__(self, model="gpt-4o-mini"):

        self.model = model

# rag/retrieval.py

class HybridRetriever:

    def __init__(self, vector_store, bm25_index, reranker, weights=(0.4, 0.3, 0.3)):

        self.vector = vector_store

        self.bm25 = bm25_index

        self.reranker = reranker

        self.weights = weights  # vector, bm25, reranker



    async def retrieve(self, query: str, k=20, final_k=5):

        # Stage 1: Parallel retrieval

        vector_results = await self.vector.search(query, k=k)

        bm25_results = await self.bm25.search(query, k=k)



        # Stage 2: Reciprocal Rank Fusion

        fused = self._rrf(vector_results, bm25_results, k=50)



        # Stage 3: Cross-encoder rerank (top 50 → top 5)

        reranked = await self.reranker.rerank(query, fused[:50])



        return reranked[:final_k]



    def _rrf(self, *result_lists, k=60):

        """Reciprocal Rank Fusion — no score calibration needed."""

        scores = defaultdict(float)

        for results in result_lists:

            for rank, doc in enumerate(results):

                scores[doc.id] += 1 / (k + rank + 1)

        return sorted(scores.items(), key=lambda x: -x[1])

# rag/query_transform.py

class QueryTransformer:

    def __init__(self, llm_model="gpt-4o-mini"):

        self.llm = instructor.from_openai(AsyncOpenAI())



    async def expand(self, query: str, context: dict = None) -> list[str]:

        """Generate multiple search queries from one user question."""



        class QuerySet(BaseModel):

            queries: list[str] = Field(min_length=3, max_length=5)

            reasoning: str



        result = await self.llm.chat.completions.create(

            model=self.model,

            response_model=QuerySet,

            messages=[

                {"role": "system", "content": """

Generate diverse search queries that collectively cover the user's intent.

Include: exact phrasing, synonyms, broader/narrower, hypothetical answer.

"""},

                {"role": "user", "content": f"Original: {query}\nContext: {context}"}

            ],

            temperature=0.3,

        )

        return result.queries



    async def decompose(self, query: str) -> list[str]:

        """Break multi-hop questions into sub-questions."""



        class SubQuestions(BaseModel):

            questions: list[str]

            needs_synthesis: bool



        return await self.llm.chat.completions.create(

            model=self.model,

            response_model=SubQuestions,

            messages=[...],

        )

# rag/optimization.py

import optuna

from dataclasses import dataclass



@dataclass

class RetrievalConfig:

    chunk_size: int

    overlap: int

    top_k: int

    vector_weight: float

    bm25_weight: float

    rerank_top_k: int



def objective(trial: optuna.Trial) -> tuple[float, float]:

    config = RetrievalConfig(

        chunk_size=trial.suggest_categorical("chunk_size", [256, 512, 768, 1024, 1536]),

        overlap=trial.suggest_int("overlap", 0, 200, step=25),

        top_k=trial.suggest_int("top_k", 5, 50, step=5),

        vector_weight=trial.suggest_float("vector_weight", 0.1, 0.8),

        bm25_weight=trial.suggest_float("bm25_weight", 0.1, 0.8),

        rerank_top_k=trial.suggest_int("rerank_top_k", 10, 100, step=10),

    )



    # Evaluate on golden set (200 queries)

    recall, latency = evaluate_config(config, golden_set)



    # Multi-objective: maximize recall, minimize latency

    return recall, latency / 1000  # seconds



study = optuna.create_study(

    directions=["maximize", "minimize"],

    sampler=optuna.samplers.TPESampler(multivariate=True),

)

study.optimize(objective, n_trials=100, timeout=3600)  # 1 hour



# Pareto frontier gives you the tradeoff curve

pareto = [t for t in study.trials if t.state == TrialState.COMPLETE]

# rag/metrics.py

from prometheus_client import Histogram, Counter, Gauge



RETRIEVAL_LATENCY = Histogram("rag_retrieval_latency_seconds", "End-to-end retrieval time")

RECALL_AT_K = Gauge("rag_recall_at_k", "Recall@k on golden set", ["k"])

QUERY_EXPANSION_COUNT = Counter("rag_query_expansions_total", "Number of expanded queries")

RERANKER_LATENCY = Histogram("rag_reranker_latency_seconds", "Cross-encoder rerank time")



class InstrumentedRetriever(HybridRetriever):

    async def retrieve(self, query, k=20, final_k=5):

        with RETRIEVAL_LATENCY.time():

            expanded = await self.transformer.expand(query)

            QUERY_EXPANSION_COUNT.inc(len(expanded))



            results = await super().retrieve(expanded, k, final_k)



            # Track recall on sampled golden queries (1% of traffic)

            if random.random() < 0.01:

                RECALL_AT_K.labels(k=10).set(self._eval_recall(query, results))



            return results