""" Self-RAG: adaptive retrieval with query classification and grounding feedback. Standard RAG retrieves for every query -- wasteful for greetings ("Hello!") and harmful for general knowledge questions (retrieved chunks dilute the model's built-in knowledge with irrelevant context). Self-RAG adds two decision points: 1. **Before retrieval** -- classify the query to decide whether retrieval is needed at all. Greetings and chitchat skip retrieval entirely. 2. **After generation** -- score how well the answer is grounded in the retrieved chunks (token-overlap F1). If grounding is low, the query may have been too vague or off-target -- reformulate it by extracting salient keywords from the best-scoring chunks and re-retrieve. This is a lightweight, heuristic Self-RAG -- no separate "critic" model. Classification uses keyword/pattern matching, and grounding reuses the same token F1 metric from the hallucination guardrail and citation engine. """ from __future__ import annotations import re import string from collections import Counter from dataclasses import dataclass, field from enum import Enum # --------------------------------------------------------------------------- # Query intent classification # --------------------------------------------------------------------------- class QueryIntent(Enum): """High-level intent categories that drive retrieval decisions.""" GREETING = "greeting" # "Hello!", "Good morning" CHITCHAT = "chitchat" # "How are you?", "Tell me a joke" KNOWLEDGE = "knowledge" # Requires document retrieval # NOTE: we intentionally don't have a GENERAL_KNOWLEDGE bucket. # If documents are loaded, even general-sounding questions might be # answerable from the corpus. The classifier errs toward retrieval. # Greeting patterns — short, social utterances that never need retrieval. _GREETING_PATTERNS = re.compile( r"^\s*(" r"h(?:ello|i|ey|owdy)" r"|good\s+(?:morning|afternoon|evening|day)" r"|what'?s?\s+up" r"|yo+\b" r"|sup\b" r"|greetings" r"|thanks?(?:\s+you)?" r"|thank\s+you" r"|bye|goodbye|see\s+you" r")\s*[!?.,]*\s*$", re.IGNORECASE, ) # Chitchat — questions about the assistant itself, not about documents. _CHITCHAT_PATTERNS = re.compile( r"^\s*(" r"(?:how\s+are\s+you|who\s+are\s+you|what\s+(?:are|is)\s+you)" r"|(?:tell\s+me\s+(?:a\s+)?(?:joke|story|something\s+funny))" r"|(?:what\s+(?:can|do)\s+you\s+do)" r"|(?:are\s+you\s+(?:an?\s+)?(?:ai|bot|human|real))" r")\s*[!?.,]*\s*$", re.IGNORECASE, ) def classify_query(query: str) -> QueryIntent: """Classify a user query into an intent category. Uses pattern matching for greetings and chitchat. Anything else is classified as KNOWLEDGE (needs retrieval). This intentionally errs on the side of retrieval -- it's better to retrieve and find nothing useful than to skip retrieval and miss relevant context. """ text = query.strip() if not text: return QueryIntent.GREETING if _GREETING_PATTERNS.match(text): return QueryIntent.GREETING if _CHITCHAT_PATTERNS.match(text): return QueryIntent.CHITCHAT return QueryIntent.KNOWLEDGE # --------------------------------------------------------------------------- # Grounding scorer # --------------------------------------------------------------------------- def _normalize(text: str) -> str: """Lowercase, strip articles, punctuation, and extra whitespace.""" text = text.lower() text = re.sub(r"\b(a|an|the)\b", " ", text) text = text.translate(str.maketrans("", "", string.punctuation)) return " ".join(text.split()) def _token_f1(candidate: str, reference: str) -> float: """Token-level F1 between two strings.""" cand_tokens = _normalize(candidate).split() ref_tokens = _normalize(reference).split() if not ref_tokens: return 1.0 if not cand_tokens else 0.0 if not cand_tokens: return 0.0 common = sum((Counter(cand_tokens) & Counter(ref_tokens)).values()) if common == 0: return 0.0 precision = common / len(cand_tokens) recall = common / len(ref_tokens) return 2 * precision * recall / (precision + recall) def _split_sentences(text: str) -> list[str]: """Split on sentence-ending punctuation followed by space or EOL.""" parts = re.split(r"(?<=[.!?])\s+", text.strip()) return [s.strip() for s in parts if s.strip()] def score_grounding(answer: str, chunks: list[str]) -> float: """Score how well an answer is grounded in retrieved chunks. Returns the mean of the best token F1 each sentence achieves against any chunk. Sentences shorter than 4 words (headers, transitions) are excluded. A score of 0.0 means nothing is grounded; 1.0 means perfect overlap. Args: answer: The model's generated answer. chunks: Retrieved source chunk texts. Returns: Mean grounding score in [0.0, 1.0]. """ if not answer or not chunks: return 0.0 sentences = _split_sentences(answer) scores: list[float] = [] for sent in sentences: if len(sent.split()) < 4: continue best_f1 = max((_token_f1(sent, chunk) for chunk in chunks), default=0.0) scores.append(best_f1) if not scores: return 0.0 return sum(scores) / len(scores) # --------------------------------------------------------------------------- # Query reformulation # --------------------------------------------------------------------------- def _extract_keywords(chunks: list[str], top_n: int = 5) -> list[str]: """Extract the most frequent meaningful tokens from chunk texts. Filters out stopwords and very short tokens to keep only content words that might help refine a vague query. """ _STOPWORDS = frozenset({ "a", "an", "the", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with", "by", "from", "is", "was", "are", "were", "be", "been", "being", "have", "has", "had", "do", "does", "did", "will", "would", "could", "should", "may", "might", "can", "shall", "not", "no", "nor", "so", "yet", "both", "either", "neither", "each", "every", "all", "any", "few", "more", "most", "other", "some", "such", "than", "too", "very", "just", "also", "now", "that", "this", "these", "those", "it", "its", "they", "them", "their", "we", "our", "you", "your", "he", "she", "his", "her", "which", "what", "who", "whom", "when", "where", "how", "why", "about", "into", "through", "during", "before", "after", "above", "below", "between", "under", "over", "again", "then", "once", "here", "there", "if", "as", "up", "out", "off", }) counter: Counter[str] = Counter() for chunk in chunks: tokens = _normalize(chunk).split() for tok in tokens: if len(tok) > 2 and tok not in _STOPWORDS: counter[tok] += 1 return [word for word, _ in counter.most_common(top_n)] def reformulate_query( original_query: str, chunks: list[str], top_n_keywords: int = 5, ) -> str: """Reformulate a query by enriching it with keywords from top chunks. When the original query is too vague to retrieve well-grounded results, this function extracts frequent content words from the best chunks and appends them to the query. This shifts the embedding closer to the relevant document cluster. Args: original_query: The user's original query text. chunks: Chunk texts from the first retrieval attempt. top_n_keywords: Number of keywords to extract and append. Returns: An enriched query string. """ if not chunks: return original_query keywords = _extract_keywords(chunks, top_n=top_n_keywords) # Only add keywords not already in the query query_tokens = set(_normalize(original_query).split()) new_keywords = [kw for kw in keywords if kw not in query_tokens] if not new_keywords: return original_query return "{} {}".format(original_query.rstrip("?").strip(), " ".join(new_keywords)) # --------------------------------------------------------------------------- # Adaptive retrieval result # --------------------------------------------------------------------------- @dataclass class AdaptiveResult: """Result of adaptive retrieval with metadata about the decision process.""" intent: QueryIntent chunks: list[tuple[str, float]] # final retrieved chunks skipped: bool # True if retrieval was skipped reformulated: bool = False # True if query was reformulated original_query: str = "" reformulated_query: str = "" grounding_score: float = 0.0 # grounding of answer vs chunks attempts: int = 0 # number of retrieval attempts # --------------------------------------------------------------------------- # Adaptive retriever # --------------------------------------------------------------------------- class AdaptiveRetriever: """Self-RAG style adaptive retriever wrapping a RAGPipeline. Adds two decision layers around standard retrieval: 1. **Pre-retrieval**: classifies the query intent and skips retrieval for greetings and chitchat. 2. **Post-generation**: given a generated answer, evaluates grounding against retrieved chunks. If grounding is below threshold, reformulates the query with keywords from the top chunks and re-retrieves once. Usage:: adaptive = AdaptiveRetriever(pipeline) # Phase 1: decide whether to retrieve result = adaptive.retrieve(query) if result.skipped: # generate without context answer = engine.generate_answer(messages) else: # generate with retrieved context answer = engine.generate_answer(rag_messages) # Phase 2: check grounding and optionally re-retrieve result = adaptive.check_and_requery(query, answer, result) if result.reformulated: # re-generate with new chunks answer = engine.generate_answer(new_rag_messages) """ def __init__( self, pipeline, grounding_threshold: float = 0.15, ) -> None: self.pipeline = pipeline self.grounding_threshold = grounding_threshold def retrieve( self, query: str, top_k: int = 3, metadata_filter: dict | None = None, ) -> AdaptiveResult: """Classify query and retrieve if needed. Returns an AdaptiveResult. If ``result.skipped`` is True, no chunks were retrieved and the caller should generate without RAG context. """ intent = classify_query(query) if intent in (QueryIntent.GREETING, QueryIntent.CHITCHAT): return AdaptiveResult( intent=intent, chunks=[], skipped=True, original_query=query, attempts=0, ) # Retrieve from pipeline (supports decomposed, hybrid, etc.) chunks = self.pipeline.retrieve(query, top_k=top_k, metadata_filter=metadata_filter) return AdaptiveResult( intent=intent, chunks=chunks, skipped=False, original_query=query, attempts=1, ) def check_and_requery( self, query: str, answer: str, previous_result: AdaptiveResult, top_k: int = 3, metadata_filter: dict | None = None, ) -> AdaptiveResult: """Check grounding and re-retrieve with a reformulated query if needed. Call this after generating an answer. If the answer is poorly grounded in the retrieved chunks, the query is reformulated with keywords from the top chunks and retrieval is retried once. Args: query: The original user query. answer: The generated answer to check. previous_result: The AdaptiveResult from the initial ``retrieve()``. top_k: Number of chunks to retrieve on retry. metadata_filter: Optional metadata filter. Returns: Updated AdaptiveResult. If re-retrieval happened, ``result.reformulated`` is True and ``result.chunks`` contains the new chunks. """ if previous_result.skipped: return previous_result chunk_texts = [text for text, _score in previous_result.chunks] grounding = score_grounding(answer, chunk_texts) previous_result.grounding_score = grounding if grounding >= self.grounding_threshold: return previous_result # Grounding is low — reformulate and re-retrieve new_query = reformulate_query(query, chunk_texts) if new_query == query: # Reformulation didn't change anything — no point re-retrieving return previous_result new_chunks = self.pipeline.retrieve( new_query, top_k=top_k, metadata_filter=metadata_filter, ) return AdaptiveResult( intent=previous_result.intent, chunks=new_chunks, skipped=False, reformulated=True, original_query=query, reformulated_query=new_query, grounding_score=grounding, attempts=previous_result.attempts + 1, )