src/serving/adaptive.py

"""
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,
        )