raglib-py 0.1.3__py3-none-any.whl

raglib/components/planner.py

@@ -0,0 +1,44 @@
+"""Query decomposition planner for multi-step retrieval."""
+
+from __future__ import annotations
+
+import logging
+import re
+from typing import List
+
+from raglib.llm.base_client import BaseLLMClient
+
+logger = logging.getLogger(__name__)
+
+
+class Planner:
+    """Breaks complex user queries into simpler sub-queries."""
+
+    def __init__(self, llm_client: BaseLLMClient):
+        """Initialize planner with an LLM client."""
+
+        self.llm_client = llm_client
+
+    def plan(self, query: str) -> List[str]:
+        """Decompose a query into ordered sub-queries."""
+
+        prompt = (
+            "Decompose this user query into concise sub-queries for retrieval.\n"
+            f"Query: {query}\n"
+            "Return one sub-query per line."
+        )
+        raw = self.llm_client.complete(prompt=prompt)
+        lines = [line.strip() for line in raw.splitlines() if line.strip()]
+
+        sub_queries: List[str] = []
+        for line in lines:
+            cleaned = re.sub(r"^\d+[\).:-]\s*", "", line)
+            if cleaned:
+                sub_queries.append(cleaned)
+
+        if not sub_queries:
+            parts = [p.strip() for p in re.split(r"\band\b|\bthen\b|\b,\b", query) if p.strip()]
+            sub_queries = parts if parts else [query]
+
+        logger.info("Planner produced %d sub-queries", len(sub_queries))
+        return sub_queries

raglib/components/query_expander.py

@@ -0,0 +1,52 @@
+"""Query expansion component for multi-query retrieval."""
+
+from __future__ import annotations
+
+import logging
+from typing import List, Optional
+
+from raglib.llm.base_client import BaseLLMClient
+
+logger = logging.getLogger(__name__)
+
+
+class QueryExpander:
+    """Generates semantic variants of a user query."""
+
+    def __init__(self, llm_client: Optional[BaseLLMClient] = None, max_variants: int = 4):
+        """Initialize expansion settings and optional LLM client."""
+
+        self.llm_client = llm_client
+        self.max_variants = max_variants
+
+    def expand(self, query: str) -> List[str]:
+        """Expand a query into several retrieval variants."""
+
+        if self.llm_client is None:
+            variants = [
+                query,
+                f"{query} best practices",
+                f"{query} implementation",
+                f"{query} examples",
+            ]
+        else:
+            prompt = (
+                "Expand the query into semantic variants for retrieval.\n"
+                f"Query: {query}\n"
+                f"Return up to {self.max_variants} variants, one per line."
+            )
+            raw = self.llm_client.complete(prompt=prompt)
+            variants = [line.strip(" -") for line in raw.splitlines() if line.strip()]
+            variants.insert(0, query)
+
+        deduped: List[str] = []
+        seen = set()
+        for variant in variants:
+            key = variant.lower().strip()
+            if key and key not in seen:
+                deduped.append(variant.strip())
+                seen.add(key)
+
+        selected = deduped[: self.max_variants]
+        logger.info("QueryExpander produced %d variants", len(selected))
+        return selected

raglib/components/query_rewriter.py

@@ -0,0 +1,38 @@
+"""Pre-retrieval query rewriting component."""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+from raglib.llm.base_client import BaseLLMClient
+
+logger = logging.getLogger(__name__)
+
+
+class QueryRewriter:
+    """Rewrites incoming queries into retrieval-friendly forms."""
+
+    def __init__(self, llm_client: Optional[BaseLLMClient] = None):
+        """Initialize an optional LLM-backed query rewriter."""
+
+        self.llm_client = llm_client
+
+    def rewrite(self, query: str) -> str:
+        """Rewrite a query while preserving user intent."""
+
+        if self.llm_client is None:
+            rewritten = query.strip()
+            logger.debug("QueryRewriter fallback used original query")
+            return rewritten
+
+        prompt = (
+            "Rewrite the user query for search relevance while preserving intent.\n"
+            f"Query: {query}\n"
+            "Return one rewritten query."
+        )
+        rewritten = self.llm_client.complete(prompt=prompt).strip()
+        if rewritten.lower().startswith("refined query:"):
+            rewritten = rewritten.split(":", 1)[1].strip()
+        logger.info("QueryRewriter output: %s", rewritten)
+        return rewritten or query

raglib/components/refiner.py

@@ -0,0 +1,37 @@
+"""Query refinement component for corrective retrieval loops."""
+
+from __future__ import annotations
+
+import logging
+from typing import List
+
+from raglib.llm.base_client import BaseLLMClient
+from raglib.schemas import Document
+
+logger = logging.getLogger(__name__)
+
+
+class Refiner:
+    """Refines user queries after poor retrieval quality."""
+
+    def __init__(self, llm_client: BaseLLMClient, max_retries: int = 2):
+        """Initialize the refiner with an LLM client and retry policy."""
+
+        self.llm_client = llm_client
+        self.max_retries = max_retries
+
+    def refine(self, query: str, failed_docs: List[Document]) -> str:
+        """Generate a refined query from failed retrieval context."""
+
+        snippets = "\n".join(doc.content[:180] for doc in failed_docs[:3])
+        prompt = (
+            "Rewrite the query to improve retrieval quality.\n"
+            f"Original query: {query}\n"
+            f"Low-quality snippets:\n{snippets}\n"
+            "Return only the refined query."
+        )
+        refined = self.llm_client.complete(prompt=prompt).strip()
+        if refined.lower().startswith("refined query:"):
+            refined = refined.split(":", 1)[1].strip()
+        logger.info("Refiner transformed query from '%s' to '%s'", query, refined)
+        return refined or query

raglib/components/reflection.py

@@ -0,0 +1,48 @@
+"""Reflection module for deciding next retrieval action."""
+
+from __future__ import annotations
+
+import logging
+from typing import List
+
+from raglib.components.evaluator import Evaluator
+from raglib.schemas import Document
+
+logger = logging.getLogger(__name__)
+
+ACTION_SUFFICIENT: str = "sufficient"
+ACTION_RETRY: str = "retry"
+ACTION_WEB_FALLBACK: str = "web_fallback"
+ACTION_GIVE_UP: str = "give_up"
+
+
+class ReflectionModule:
+    """Inspects retrieval quality and recommends next action."""
+
+    def __init__(self, evaluator: Evaluator):
+        """Initialize reflection with an evaluator component."""
+
+        self.evaluator = evaluator
+
+    def reflect(self, query: str, documents: List[Document]) -> str:
+        """Return the next action based on current retrieval quality."""
+
+        if not documents:
+            logger.info("ReflectionModule found no documents; suggesting web fallback")
+            return ACTION_WEB_FALLBACK
+
+        scores = [self.evaluator.score_document(query, doc) for doc in documents]
+        average = sum(scores) / len(scores)
+        best = max(scores)
+
+        if average >= self.evaluator.relevance_threshold + 0.15:
+            action = ACTION_SUFFICIENT
+        elif best >= self.evaluator.relevance_threshold:
+            action = ACTION_RETRY
+        elif len(documents) >= 3:
+            action = ACTION_WEB_FALLBACK
+        else:
+            action = ACTION_GIVE_UP
+
+        logger.info("ReflectionModule action=%s avg=%.3f best=%.3f", action, average, best)
+        return action

raglib/components/reranker.py

@@ -0,0 +1,54 @@
+"""Document reranking utilities for post-retrieval relevance sorting."""
+
+from __future__ import annotations
+
+import logging
+import re
+from dataclasses import replace
+from typing import List, Set
+
+from raglib.schemas import Document
+
+logger = logging.getLogger(__name__)
+
+
+class Reranker:
+    """Reranks candidate documents using lightweight lexical scoring."""
+
+    def __init__(self, top_k: int = 3):
+        """Initialize reranker settings."""
+
+        self.top_k = top_k
+
+    def rerank(self, query: str, documents: List[Document]) -> List[Document]:
+        """Return top-k reranked documents for a query."""
+
+        if not documents:
+            return []
+
+        query_terms = self._tokenize(query)
+        rescored: List[Document] = []
+        for doc in documents:
+            score = self._score(query_terms, doc)
+            rescored.append(replace(doc, score=score))
+
+        rescored.sort(key=lambda doc: doc.score, reverse=True)
+        logger.info("Reranker reduced %d docs to top %d", len(documents), self.top_k)
+        return rescored[: self.top_k]
+
+    def _score(self, query_terms: Set[str], doc: Document) -> float:
+        """Score a document using overlap and length normalization."""
+
+        doc_terms = self._tokenize(doc.content)
+        if not doc_terms:
+            return 0.0
+
+        overlap = len(query_terms & doc_terms)
+        overlap_ratio = overlap / max(len(query_terms), 1)
+        length_norm = min(len(doc_terms), 300) / 300.0
+        return 0.75 * overlap_ratio + 0.25 * length_norm
+
+    def _tokenize(self, text: str) -> Set[str]:
+        """Tokenize text into a set of lowercase lexical units."""
+
+        return set(re.findall(r"[a-zA-Z0-9]+", text.lower()))

raglib/components/retriever.py

@@ -0,0 +1,342 @@
+"""In-memory vector retriever with optional FAISS acceleration."""
+
+from __future__ import annotations
+
+import logging
+import math
+import re
+from collections import Counter
+from dataclasses import replace
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+
+from raglib.embedding.base_embedding import BaseEmbedding
+from raglib.schemas import Document
+from raglib.vectorstores.base_store import BaseVectorStore, VectorRecord
+from raglib.vectorstores.factory import VectorStoreFactory
+
+logger = logging.getLogger(__name__)
+
+
+class Retriever:
+    """Retrieves relevant documents from an in-memory corpus."""
+
+    def __init__(
+        self,
+        documents: List[Document],
+        embedding: Optional[BaseEmbedding] = None,
+        top_k: int = 5,
+        vector_db: Optional[Union[str, BaseVectorStore]] = None,
+        vector_db_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        """Initialize retriever state and optional FAISS backend."""
+
+        self.top_k = top_k
+        self._embedding = embedding
+        self._vector_store: Optional[BaseVectorStore] = None
+        self._documents: List[Document] = []
+        self._vector_document_lookup: Dict[str, Document] = {}
+        self._doc_tokens: List[List[str]] = []
+        self._doc_tfidf: List[Dict[str, float]] = []
+        self._document_frequency: Counter[str] = Counter()
+        self._faiss_index = None
+        self._faiss_vocabulary: List[str] = []
+        self._faiss_ready = False
+        self._faiss_import_error: Optional[Exception] = None
+
+        if self._embedding is not None:
+            logger.info("Retriever using embedding model=%s", type(self._embedding).__name__)
+            self._vector_store = VectorStoreFactory.build(
+                vector_db=vector_db,
+                vector_db_kwargs=vector_db_kwargs,
+            )
+            logger.info("Retriever using vector store=%s", type(self._vector_store).__name__)
+
+        self._detect_faiss()
+        self.add_documents(documents)
+
+    def add_documents(self, documents: List[Document]) -> None:
+        """Add documents to the corpus and refresh vector indices."""
+
+        if not documents:
+            return
+        self._documents.extend(documents)
+        logger.info("Retriever added %d documents (total=%d)", len(documents), len(self._documents))
+        self._reindex()
+
+    def retrieve(self, query: str) -> List[Document]:
+        """Retrieve top-k documents ranked by query similarity."""
+
+        logger.debug("Retriever received query=%s", query)
+        if not self._documents:
+            logger.warning("Retriever called with empty corpus")
+            return []
+
+        if not query or not query.strip():
+            logger.info("Retriever returning empty result for empty query")
+            return []
+
+        if not self._tokenize(query):
+            logger.info("Retriever returning empty result for query with no indexable tokens")
+            return []
+
+        if self._embedding is not None and self._vector_store is not None:
+            embedding_docs = self._retrieve_with_embedding(query)
+            if embedding_docs:
+                logger.info("Retriever returned %d documents via embedding search", len(embedding_docs))
+                return embedding_docs
+            logger.warning("Embedding retrieval produced no results; falling back to sparse search")
+
+        if self._faiss_ready:
+            faiss_docs = self._retrieve_with_faiss(query)
+            if faiss_docs:
+                logger.info("Retriever returned %d documents", len(faiss_docs))
+                return faiss_docs
+
+        docs = self._retrieve_with_sparse_cosine(query)
+        logger.info("Retriever returned %d documents", len(docs))
+        return docs
+
+    def _detect_faiss(self) -> None:
+        """Detect FAISS availability without requiring it at runtime."""
+
+        try:
+            import faiss  # type: ignore # noqa: F401
+            import numpy  # type: ignore # noqa: F401
+
+            logger.info("FAISS backend detected and will be used when possible")
+        except Exception as exc:  # noqa: BLE001
+            self._faiss_import_error = exc
+            logger.info("FAISS backend unavailable; using pure Python retriever")
+
+    def _reindex(self) -> None:
+        """Recompute sparse vectors and optional FAISS index."""
+
+        self._doc_tokens = [self._tokenize(doc.content) for doc in self._documents]
+        self._document_frequency = Counter()
+
+        for tokens in self._doc_tokens:
+            for token in set(tokens):
+                self._document_frequency[token] += 1
+
+        self._doc_tfidf = []
+        for tokens in self._doc_tokens:
+            tf = Counter(tokens)
+            self._doc_tfidf.append(self._counter_to_tfidf(tf, len(tokens)))
+
+        if self._embedding is not None and self._vector_store is not None and self._documents:
+            try:
+                texts = [document.content for document in self._documents]
+                vectors = self._embedding.embed_documents(texts)
+                records: List[VectorRecord] = []
+                self._vector_document_lookup = {}
+
+                limit = min(len(self._documents), len(vectors))
+                for index in range(limit):
+                    document = self._documents[index]
+                    vector = vectors[index]
+                    vector_id = self._build_vector_id(document=document, index=index)
+                    records.append(
+                        VectorRecord(
+                            document_id=vector_id,
+                            embedding=vector,
+                            document_text=document.content,
+                            metadata={
+                                **document.metadata,
+                                "source": document.source,
+                                "doc_id": document.id,
+                            },
+                        )
+                    )
+                    self._vector_document_lookup[vector_id] = document
+
+                self._vector_store.clear()
+                self._vector_store.upsert(records)
+                logger.info(
+                    "Retriever indexed %d embedding vectors in %s (dim=%d)",
+                    len(records),
+                    type(self._vector_store).__name__,
+                    self._embedding.dimension,
+                )
+            except Exception as exc:  # noqa: BLE001
+                logger.warning("Embedding indexing failed; using sparse fallback only: %s", exc)
+                self._vector_document_lookup = {}
+
+        self._build_faiss_index()
+
+    def _build_faiss_index(self) -> None:
+        """Build FAISS index if optional dependencies are available."""
+
+        if self._embedding is not None:
+            self._faiss_ready = False
+            self._faiss_index = None
+            return
+
+        try:
+            import faiss  # type: ignore
+            import numpy as np  # type: ignore
+        except Exception:
+            self._faiss_ready = False
+            self._faiss_index = None
+            return
+
+        if not self._documents:
+            self._faiss_ready = False
+            self._faiss_index = None
+            return
+
+        vocabulary = sorted(self._document_frequency.keys())
+        if not vocabulary:
+            self._faiss_ready = False
+            self._faiss_index = None
+            return
+
+        dim = len(vocabulary)
+        matrix = np.zeros((len(self._documents), dim), dtype="float32")
+        for row_index, doc_vector in enumerate(self._doc_tfidf):
+            for term_index, term in enumerate(vocabulary):
+                matrix[row_index, term_index] = float(doc_vector.get(term, 0.0))
+
+        faiss.normalize_L2(matrix)
+        index = faiss.IndexFlatIP(dim)
+        index.add(matrix)
+
+        self._faiss_index = index
+        self._faiss_vocabulary = vocabulary
+        self._faiss_ready = True
+        logger.debug("FAISS index built with %d docs and dim=%d", len(self._documents), dim)
+
+    def _retrieve_with_faiss(self, query: str) -> List[Document]:
+        """Retrieve documents using FAISS inner-product search."""
+
+        if not self._faiss_ready or self._faiss_index is None:
+            return []
+
+        try:
+            import faiss  # type: ignore
+            import numpy as np  # type: ignore
+        except Exception:
+            return []
+
+        query_tokens = self._tokenize(query)
+        tf = Counter(query_tokens)
+        query_vector = self._counter_to_tfidf(tf, len(query_tokens))
+        dense = np.zeros((1, len(self._faiss_vocabulary)), dtype="float32")
+        for term_index, term in enumerate(self._faiss_vocabulary):
+            dense[0, term_index] = float(query_vector.get(term, 0.0))
+
+        faiss.normalize_L2(dense)
+        distances, indices = self._faiss_index.search(dense, min(self.top_k, len(self._documents)))
+
+        ranked: List[Tuple[float, int]] = []
+        for rank_idx, doc_idx in enumerate(indices[0].tolist()):
+            if doc_idx < 0:
+                continue
+            ranked.append((float(distances[0][rank_idx]), int(doc_idx)))
+
+        return self._to_scored_documents(ranked)
+
+    def _retrieve_with_sparse_cosine(self, query: str) -> List[Document]:
+        """Retrieve documents using sparse cosine similarity."""
+
+        query_tokens = self._tokenize(query)
+        query_tf = Counter(query_tokens)
+        query_vector = self._counter_to_tfidf(query_tf, len(query_tokens))
+
+        scored: List[Tuple[float, int]] = []
+        for idx, doc_vector in enumerate(self._doc_tfidf):
+            score = self._cosine_similarity(query_vector, doc_vector)
+            scored.append((score, idx))
+
+        scored.sort(key=lambda item: item[0], reverse=True)
+        top = scored[: min(self.top_k, len(scored))]
+        return self._to_scored_documents(top)
+
+    def _retrieve_with_embedding(self, query: str) -> List[Document]:
+        """Retrieve documents using dense embedding cosine similarity."""
+
+        if self._embedding is None or self._vector_store is None:
+            return []
+        if not self._vector_document_lookup:
+            return []
+
+        try:
+            query_vector = self._embedding.embed_query(query)
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("Embedding query failed; using sparse fallback: %s", exc)
+            return []
+
+        hits = self._vector_store.query(
+            query_embedding=query_vector,
+            top_k=min(self.top_k, len(self._vector_document_lookup)),
+        )
+        documents: List[Document] = []
+        for hit in hits:
+            base_document = self._vector_document_lookup.get(hit.document_id)
+            if base_document is None:
+                continue
+            logger.debug("Embedding score doc_id=%s score=%.4f", base_document.id, hit.score)
+            documents.append(
+                replace(
+                    base_document,
+                    score=float(hit.score),
+                    source=base_document.source or "vector",
+                    metadata={
+                        **base_document.metadata,
+                        "retriever": "vector_store",
+                        "vector_db": type(self._vector_store).__name__,
+                    },
+                )
+            )
+        return documents
+
+    def _to_scored_documents(self, scored_indices: Sequence[Tuple[float, int]]) -> List[Document]:
+        """Convert score/index pairs into copied scored Documents."""
+
+        docs: List[Document] = []
+        for score, doc_idx in scored_indices:
+            base = self._documents[doc_idx]
+            docs.append(
+                replace(
+                    base,
+                    score=float(score),
+                    source=base.source or "vector",
+                    metadata={**base.metadata, "retriever": "in_memory"},
+                )
+            )
+        return docs
+
+    def _counter_to_tfidf(self, term_counter: Counter[str], token_count: int) -> Dict[str, float]:
+        """Convert term counts to tf-idf sparse weights."""
+
+        if token_count <= 0:
+            return {}
+        vector: Dict[str, float] = {}
+        doc_total = max(len(self._documents), 1)
+        for term, count in term_counter.items():
+            tf = count / token_count
+            idf = math.log((1 + doc_total) / (1 + self._document_frequency.get(term, 0))) + 1.0
+            vector[term] = tf * idf
+        return vector
+
+    def _cosine_similarity(self, a: Dict[str, float], b: Dict[str, float]) -> float:
+        """Compute cosine similarity for sparse vectors."""
+
+        if not a or not b:
+            return 0.0
+        shared_terms = set(a.keys()) & set(b.keys())
+        dot = sum(a[t] * b[t] for t in shared_terms)
+        norm_a = math.sqrt(sum(v * v for v in a.values()))
+        norm_b = math.sqrt(sum(v * v for v in b.values()))
+        if norm_a == 0.0 or norm_b == 0.0:
+            return 0.0
+        return dot / (norm_a * norm_b)
+
+    def _build_vector_id(self, document: Document, index: int) -> str:
+        """Build a unique vector record id for a document entry."""
+
+        return f"{document.id}__{index}"
+
+    def _tokenize(self, text: str) -> List[str]:
+        """Tokenize text into lowercase alphanumeric terms."""
+
+        return re.findall(r"[a-zA-Z0-9]+", text.lower())

raglib/components/router_retriever.py

@@ -0,0 +1,70 @@
+"""Query router that dispatches to vector, web, or hybrid retrieval."""
+
+from __future__ import annotations
+
+import logging
+from typing import List, Optional
+
+from raglib.components.hybrid_retriever import HybridRetriever
+from raglib.components.retriever import Retriever
+from raglib.components.web_retriever import WebRetriever
+from raglib.llm.base_client import BaseLLMClient
+from raglib.schemas import Document
+
+logger = logging.getLogger(__name__)
+
+ROUTE_VECTOR: str = "vector"
+ROUTE_WEB: str = "web"
+ROUTE_HYBRID: str = "hybrid"
+
+
+class RouterRetriever:
+    """Routes retrieval requests to the most suitable retriever."""
+
+    def __init__(
+        self,
+        vector_retriever: Retriever,
+        web_retriever: WebRetriever,
+        hybrid_retriever: Optional[HybridRetriever] = None,
+        llm_client: Optional[BaseLLMClient] = None,
+    ):
+        """Initialize routing dependencies and optional LLM policy."""
+
+        self.vector_retriever = vector_retriever
+        self.web_retriever = web_retriever
+        self.hybrid_retriever = hybrid_retriever
+        self.llm_client = llm_client
+
+    def route(self, query: str) -> str:
+        """Return the best retrieval route for the query."""
+
+        lower = query.lower()
+        if self.llm_client is not None:
+            prompt = (
+                "Choose one retrieval mode: vector, web, hybrid.\n"
+                f"Query: {query}"
+            )
+            candidate = self.llm_client.complete(prompt=prompt).strip().lower()
+            if candidate in {ROUTE_VECTOR, ROUTE_WEB, ROUTE_HYBRID}:
+                logger.info("RouterRetriever route selected by LLM: %s", candidate)
+                return candidate
+
+        if any(token in lower for token in ("today", "latest", "news", "current", "web")):
+            route = ROUTE_WEB
+        elif any(token in lower for token in ("compare", "trend", "benchmark", "vs")):
+            route = ROUTE_HYBRID
+        else:
+            route = ROUTE_VECTOR
+
+        logger.info("RouterRetriever route selected by heuristic: %s", route)
+        return route
+
+    def retrieve(self, query: str) -> List[Document]:
+        """Retrieve documents based on selected route."""
+
+        route = self.route(query)
+        if route == ROUTE_WEB:
+            return self.web_retriever.retrieve(query)
+        if route == ROUTE_HYBRID and self.hybrid_retriever is not None:
+            return self.hybrid_retriever.retrieve(query)
+        return self.vector_retriever.retrieve(query)