McQbis/document-intelligence-rag

1from __future__ import annotations2 3from enum import Enum4from typing import List, Optional, Tuple5 6from rag.ingestion.base import TextChunk7from rag.cache.query_cache import QueryCache8from rag.retrieval.retriever import HybridRetriever9 10 11class RouteMode(str, Enum):12    FAST = "fast"13    DEEP = "deep"14    AUTO = "auto"15 16 17_DEEP_KEYWORDS = {18    "explain", "summarize", "compare", "analyze", "describe",19    "elaborate", "detail", "comprehensive", "thorough", "research",20    "differences", "similarities",21}22 23_SHORT_QUERY_THRESHOLD = 624 25 26class QueryRouter:27    """Routes queries to FAST or DEEP retrieval strategies."""28 29    def __init__(30        self,31        retriever: HybridRetriever,32        cache: Optional[QueryCache] = None,33        top_k: int = 10,34        candidate_k: int = 30,35    ):36        self.retriever = retriever37        self.cache = cache38        self.top_k = top_k39        self.candidate_k = candidate_k40 41    def search(42        self,43        query: str,44        mode: RouteMode = RouteMode.AUTO,45        top_k: Optional[int] = None,46        candidate_k: Optional[int] = None,47    ) -> List[Tuple[TextChunk, float]]:48        """Execute routed retrieval with optional caching."""49 50        resolved_mode = self._resolve_mode(query, mode)51 52        _top_k = top_k or self.top_k53        _cand_k = candidate_k or self.candidate_k54        use_rerank = resolved_mode == RouteMode.DEEP55 56        if resolved_mode == RouteMode.FAST and self.cache:57            results, _ = self.cache.search(query, top_k=_top_k, candidate_k=_cand_k, rerank=False)58            return results59 60        return self.retriever.search(61            query,62            top_k=_top_k,63            candidate_k=_cand_k,64            rerank=use_rerank,65        )66 67    def classify(self, query: str) -> RouteMode:68        """Classify query into FAST or DEEP mode."""69        return self._resolve_mode(query, RouteMode.AUTO)70 71    def _resolve_mode(self, query: str, mode: RouteMode) -> RouteMode:72        if mode != RouteMode.AUTO:73            return mode74 75        words = query.lower().split()76 77        if len(words) > _SHORT_QUERY_THRESHOLD:78            return RouteMode.DEEP79 80        if any(w in _DEEP_KEYWORDS for w in words):81            return RouteMode.DEEP82 83        return RouteMode.FAST