""" Adaptive RAG System — The LLM Decides Its Own Retrieval Strategy The Problem with Fixed Pipelines: Every other system in this benchmark uses the SAME retrieval strategy for every question, regardless of what the question actually needs. - Naive RAG uses FAISS for a simple factual question ("What is RLHF?") AND a complex multi-part question ("Compare FAISS and Pinecone across 5 dimensions") — even though these need very different amounts of context. - Advanced RAG runs the full pipeline (rewrite + hybrid + rerank) even for a trivial question where a single chunk would have been enough. The Fix — Route First, Retrieve Second: Before doing any retrieval, we ask the LLM to classify the query into one of three categories: "simple" → A focused factual question. One clear answer exists in the knowledge base. Naive RAG (fast FAISS search) is enough. Example: "What is LoRA?" / "Who introduced the Transformer?" "complex" → A multi-part, comparative, or reasoning-heavy question. Needs broad, high-quality context. Advanced RAG is used: query rewriting + hybrid search + cross-encoder reranking. Example: "Compare BM25 and FAISS across recall, precision, and latency" / "What are the production challenges in LLMOps?" "ambiguous" → The question is vague, uses jargon, or could be interpreted multiple ways. Hybrid RAG (BM25 + FAISS) gives balanced coverage without the full cost of Advanced RAG. Example: "How does attention work?" / "What's the best RAG?" After classifying, the system runs the appropriate retrieval pipeline and logs which strategy was chosen — so you can inspect routing decisions in the benchmark results. Why is this better than always running Advanced RAG? - Simple questions get answers faster (lower latency) and cheaper (less cost) - Complex questions still get the full pipeline — no quality sacrifice - The routing itself is a form of reasoning — the LLM understands what kind of retrieval its own question requires Trade-offs: - One extra LLM call per question (the classification step) - If classification is wrong, the wrong pipeline runs (mitigated by the fallback: ambiguous → hybrid, which is middle-ground) - Average cost is between naive-rag and advanced-rag, not the cheapest """ import asyncio import logging import time from langchain_groq import ChatGroq from langchain_core.prompts import ChatPromptTemplate from sentence_transformers import CrossEncoder from eval_framework.config import get_settings from eval_framework.systems.shared import SharedIndex from eval_framework.types import QAPair, SystemOutput # We reuse the existing system implementations rather than duplicating code. # Adaptive RAG is a router ON TOP of the other systems — not a new retrieval # method itself. from eval_framework.systems.naive_rag import NaiveRAGSystem from eval_framework.systems.hybrid_rag import HybridRAGSystem from eval_framework.systems.advanced_rag import AdvancedRAGSystem logger = logging.getLogger(__name__) _COST_PER_OUTPUT_TOKEN = 0.59 / 1_000_000 # The three routing categories the LLM can choose from. # Kept lowercase so we can do a simple string match on the LLM response. _ROUTE_SIMPLE = "simple" _ROUTE_COMPLEX = "complex" _ROUTE_AMBIGUOUS = "ambiguous" class AdaptiveRAGSystem: """ A meta-RAG system that routes each query to the right retrieval pipeline. Instead of applying the same strategy to every question, this system: 1. Classifies the query complexity with one LLM call 2. Delegates to the appropriate sub-system: simple -> NaiveRAGSystem (FAISS only, fast) complex -> AdvancedRAGSystem (rewrite + hybrid + rerank, best quality) ambiguous -> HybridRAGSystem (BM25 + FAISS, balanced) """ def __init__( self, index: SharedIndex, model_name: str = "llama-3.3-70b-versatile", reranker: CrossEncoder | None = None, ): """ Args: index: The shared FAISS + BM25 index, passed down to sub-systems. model_name: Groq model used for both routing classification and answers. reranker: Pre-loaded CrossEncoder shared from compare_systems.py. Avoids downloading the model multiple times. """ self.model_name = model_name self._index = index settings = get_settings() # ── Router LLM ───────────────────────────────────────────────────────── # Low temperature (0.0) for deterministic routing decisions. # We don't want the classifier to be creative — just consistent. self._router_llm = ChatGroq( api_key=settings.groq_api_key, model_name=model_name, temperature=0.0, max_tokens=10, # We only need one word back: "simple", "complex", or "ambiguous" ) # ── Routing Prompt ───────────────────────────────────────────────────── # The prompt is deliberately strict: # - Exactly one word response # - Three options spelled out with clear definitions # - Examples given for each category to reduce misclassification self._router_prompt = ChatPromptTemplate.from_template( "Classify the following question into exactly ONE category.\n\n" "Categories:\n" "- simple: A direct factual question with one clear answer.\n" " Examples: 'What is FAISS?', 'Who introduced the Transformer?'\n\n" "- complex: A multi-part, comparative, or reasoning-heavy question.\n" " Examples: 'Compare BM25 and FAISS', 'What are the production challenges of LLMOps?'\n\n" "- ambiguous: A vague or open-ended question that could be interpreted multiple ways.\n" " Examples: 'How does attention work?', 'What is the best RAG approach?'\n\n" "Respond with ONLY the category name (simple / complex / ambiguous).\n\n" "Question: {question}\n\n" "Category:" ) # ── Sub-Systems ──────────────────────────────────────────────────────── # Instantiate all three sub-systems upfront. # They share the same index and reranker — no duplicate downloads. self._naive = NaiveRAGSystem(index=index, model_name=model_name) self._hybrid = HybridRAGSystem(index=index, model_name=model_name) self._advanced = AdvancedRAGSystem( index=index, model_name=model_name, reranker=reranker, ) async def _classify(self, question: str) -> str: """ Ask the router LLM to classify the question. Returns one of: "simple", "complex", "ambiguous". Falls back to "ambiguous" if the response is unexpected — this is the safest middle-ground (hybrid retrieval) when we're unsure. """ messages = await self._router_prompt.ainvoke({"question": question}) response = await self._router_llm.ainvoke(messages) # Normalize: strip whitespace, lowercase, take first word only. # LLMs sometimes return "Simple." or "COMPLEX\n" — this handles that. raw = response.content.strip().lower().split()[0].rstrip(".,:") if raw in (_ROUTE_SIMPLE, _ROUTE_COMPLEX, _ROUTE_AMBIGUOUS): return raw # If we got something unexpected (e.g. "moderate", "medium"), fall back # to ambiguous — hybrid search is a safe default. logger.warning(f"Unexpected route '{raw}' for: '{question}' — defaulting to ambiguous") return _ROUTE_AMBIGUOUS async def query(self, qa_pair: QAPair) -> SystemOutput: """ Route the question to the right sub-system and return its output. The metadata field logs: - which route was chosen - how long the routing classification itself took so you can inspect routing decisions in the benchmark results. """ start = time.time() # ── Step 1: Classify ─────────────────────────────────────────────────── route = await self._classify(qa_pair.question) route_latency_ms = (time.time() - start) * 1000 logger.info(f"AdaptiveRAG routed '{qa_pair.question[:60]}...' → {route.upper()} " f"(classification took {route_latency_ms:.0f}ms)") # ── Step 2: Delegate to sub-system ──────────────────────────────────── # Each sub-system returns a SystemOutput with its own latency/cost. # We add the routing overhead on top of that. if route == _ROUTE_SIMPLE: result = await self._naive.query(qa_pair) elif route == _ROUTE_COMPLEX: result = await self._advanced.query(qa_pair) else: # ambiguous (and any unexpected fallback) result = await self._hybrid.query(qa_pair) # ── Step 3: Patch metadata ───────────────────────────────────────────── # Add routing info to the result so it shows up in the benchmark DB. # The sub-system's own latency is already in result.latency_ms — # we add the classification overhead separately so both are visible. total_latency_ms = (time.time() - start) * 1000 routing_cost = len(qa_pair.question.split()) * 1.3 * _COST_PER_OUTPUT_TOKEN result.latency_ms = total_latency_ms result.cost_usd = result.cost_usd + routing_cost result.model = self.model_name result.metadata = { **result.metadata, "system": "adaptive_rag", "route": route, # which pipeline was chosen "route_latency_ms": round(route_latency_ms, 1), # cost of classification step } return result