raglib-py 0.1.3__py3-none-any.whl

raglib/__init__.py

@@ -0,0 +1,17 @@
+"""Public package exports for raglib."""
+
+from __future__ import annotations
+
+from raglib.facade import RAG
+from raglib.llm.mock_client import MockLLMClient
+from raglib.schemas import Document, GenerationResult, QueryResult
+
+__all__ = [
+    "RAG",
+    "Document",
+    "QueryResult",
+    "GenerationResult",
+    "MockLLMClient",
+]
+
+__version__ = "0.1.3"

raglib/components/__init__.py

@@ -0,0 +1,41 @@
+"""Component exports for raglib retrieval and generation building blocks."""
+
+from raglib.components.context_reducer import ContextReducer
+from raglib.components.decision import DecisionEngine
+from raglib.components.deduplicator import Deduplicator
+from raglib.components.evaluator import Evaluator
+from raglib.components.generator import Generator
+from raglib.components.hybrid_retriever import HybridRetriever
+from raglib.components.memory import MemoryModule
+from raglib.components.multi_hop_retriever import MultiHopRetriever
+from raglib.components.multi_query_retriever import MultiQueryRetriever
+from raglib.components.planner import Planner
+from raglib.components.query_expander import QueryExpander
+from raglib.components.query_rewriter import QueryRewriter
+from raglib.components.refiner import Refiner
+from raglib.components.reflection import ReflectionModule
+from raglib.components.reranker import Reranker
+from raglib.components.retriever import Retriever
+from raglib.components.router_retriever import RouterRetriever
+from raglib.components.web_retriever import WebRetriever
+
+__all__ = [
+    "Retriever",
+    "WebRetriever",
+    "HybridRetriever",
+    "MultiQueryRetriever",
+    "RouterRetriever",
+    "MultiHopRetriever",
+    "Reranker",
+    "Evaluator",
+    "Refiner",
+    "QueryRewriter",
+    "QueryExpander",
+    "Deduplicator",
+    "ContextReducer",
+    "MemoryModule",
+    "Planner",
+    "ReflectionModule",
+    "DecisionEngine",
+    "Generator",
+]

raglib/components/context_reducer.py

@@ -0,0 +1,65 @@
+"""Context reduction utilities for prompt token budgeting."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import replace
+from typing import List, Optional
+
+from raglib.schemas import Document
+
+logger = logging.getLogger(__name__)
+
+
+class ContextReducer:
+    """Trims retrieved context to fit a target token budget."""
+
+    def __init__(self, max_context_tokens: int = 3000):
+        """Initialize context token budget."""
+
+        self.max_context_tokens = max_context_tokens
+
+    def reduce(self, documents: List[Document], max_tokens: Optional[int] = None) -> List[Document]:
+        """Reduce documents until the token budget is satisfied."""
+
+        budget = max_tokens if max_tokens is not None else self.max_context_tokens
+        reduced: List[Document] = []
+        used_tokens = 0
+
+        for doc in documents:
+            doc_tokens = self._estimate_tokens(doc.content)
+            if used_tokens + doc_tokens <= budget:
+                reduced.append(doc)
+                used_tokens += doc_tokens
+                continue
+
+            remaining = budget - used_tokens
+            if remaining <= 0:
+                break
+
+            trimmed_content = self._truncate_to_tokens(doc.content, remaining)
+            if trimmed_content.strip():
+                reduced.append(
+                    replace(
+                        doc,
+                        content=trimmed_content,
+                        metadata={**doc.metadata, "trimmed": True},
+                    )
+                )
+                used_tokens = budget
+                break
+
+        logger.info("ContextReducer used %d/%d tokens", used_tokens, budget)
+        return reduced
+
+    def _estimate_tokens(self, text: str) -> int:
+        """Estimate token count from whitespace-separated words."""
+
+        return max(1, int(len(text.split()) * 1.2))
+
+    def _truncate_to_tokens(self, text: str, max_tokens: int) -> str:
+        """Truncate text to an approximate token count."""
+
+        words = text.split()
+        max_words = max(1, int(max_tokens / 1.2))
+        return " ".join(words[:max_words])

raglib/components/decision.py

@@ -0,0 +1,63 @@
+"""Decision engine for retrieval gating and mode selection."""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+from raglib.llm.base_client import BaseLLMClient
+
+logger = logging.getLogger(__name__)
+
+RETRIEVAL_VECTOR: str = "vector"
+RETRIEVAL_WEB: str = "web"
+RETRIEVAL_HYBRID: str = "hybrid"
+RETRIEVAL_NONE: str = "none"
+ALLOWED_RETRIEVAL_TYPES = {
+    RETRIEVAL_VECTOR,
+    RETRIEVAL_WEB,
+    RETRIEVAL_HYBRID,
+    RETRIEVAL_NONE,
+}
+
+
+class DecisionEngine:
+    """Decides whether retrieval is needed and what retrieval mode to use."""
+
+    def __init__(self, llm_client: Optional[BaseLLMClient] = None):
+        """Initialize decision policy with an optional LLM client."""
+
+        self.llm_client = llm_client
+
+    def should_retrieve(self, query: str) -> bool:
+        """Return whether retrieval should run for this query."""
+
+        route = self.retrieval_type(query)
+        decision = route != RETRIEVAL_NONE
+        logger.info("DecisionEngine should_retrieve=%s route=%s", decision, route)
+        return decision
+
+    def retrieval_type(self, query: str) -> str:
+        """Return one of: vector, web, hybrid, none."""
+
+        if self.llm_client is not None:
+            prompt = (
+                "Decide retrieval_type for this query. Allowed: vector, web, hybrid, none.\n"
+                f"Query: {query}\n"
+                "Return one token only."
+            )
+            response = self.llm_client.complete(prompt=prompt).strip().lower()
+            if response in ALLOWED_RETRIEVAL_TYPES:
+                logger.info("DecisionEngine selected via LLM: %s", response)
+                return response
+
+        lower = query.lower().strip()
+        if not lower:
+            return RETRIEVAL_NONE
+        if any(token in lower for token in ("hello", "thanks", "hi", "good morning")):
+            return RETRIEVAL_NONE
+        if any(token in lower for token in ("latest", "today", "news", "current", "web")):
+            return RETRIEVAL_WEB
+        if any(token in lower for token in ("compare", "versus", "vs", "trend")):
+            return RETRIEVAL_HYBRID
+        return RETRIEVAL_VECTOR

raglib/components/deduplicator.py

@@ -0,0 +1,58 @@
+"""Near-duplicate removal for retrieved document sets."""
+
+from __future__ import annotations
+
+import logging
+import re
+from typing import List, Set
+
+from raglib.schemas import Document
+
+logger = logging.getLogger(__name__)
+
+
+class Deduplicator:
+    """Removes semantically similar duplicate document chunks."""
+
+    def __init__(self, similarity_threshold: float = 0.85):
+        """Initialize deduplication threshold."""
+
+        self.similarity_threshold = similarity_threshold
+
+    def deduplicate(self, documents: List[Document]) -> List[Document]:
+        """Return documents after near-duplicate filtering."""
+
+        unique_docs: List[Document] = []
+        signatures: List[Set[str]] = []
+
+        for doc in documents:
+            signature = self._signature(doc.content)
+            if not signatures:
+                unique_docs.append(doc)
+                signatures.append(signature)
+                continue
+
+            is_duplicate = any(
+                self._jaccard(signature, existing) >= self.similarity_threshold
+                for existing in signatures
+            )
+            if not is_duplicate:
+                unique_docs.append(doc)
+                signatures.append(signature)
+
+        logger.info("Deduplicator kept %d/%d documents", len(unique_docs), len(documents))
+        return unique_docs
+
+    def _signature(self, text: str) -> Set[str]:
+        """Build a token-set signature for similarity checks."""
+
+        return set(re.findall(r"[a-zA-Z0-9]+", text.lower()))
+
+    def _jaccard(self, a: Set[str], b: Set[str]) -> float:
+        """Compute Jaccard similarity between token signatures."""
+
+        if not a and not b:
+            return 1.0
+        if not a or not b:
+            return 0.0
+        return len(a & b) / len(a | b)

raglib/components/evaluator.py

@@ -0,0 +1,57 @@
+"""Retrieval quality evaluator used for corrective feedback loops."""
+
+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 Evaluator:
+    """Scores relevance and filters low-quality retrieval results."""
+
+    def __init__(self, relevance_threshold: float = 0.3):
+        """Initialize evaluator with a minimum relevance threshold."""
+
+        self.relevance_threshold = relevance_threshold
+
+    def evaluate(self, query: str, documents: List[Document]) -> List[Document]:
+        """Filter documents by relevance score threshold."""
+
+        passed: List[Document] = []
+        for doc in documents:
+            score = self.score_document(query, doc)
+            if score >= self.relevance_threshold:
+                passed.append(replace(doc, score=score))
+
+        logger.info(
+            "Evaluator kept %d/%d documents with threshold %.2f",
+            len(passed),
+            len(documents),
+            self.relevance_threshold,
+        )
+        return passed
+
+    def score_document(self, query: str, doc: Document) -> float:
+        """Compute a relevance score for a single document."""
+
+        query_terms = self._tokenize(query)
+        doc_terms = self._tokenize(doc.content)
+        if not query_terms or not doc_terms:
+            return 0.0
+
+        overlap = len(query_terms & doc_terms) / len(query_terms)
+        density = len(query_terms & doc_terms) / max(len(doc_terms), 1)
+        phrase_bonus = 0.2 if query.lower() in doc.content.lower() else 0.0
+        score = min(1.0, overlap * 0.7 + density * 0.3 + phrase_bonus)
+        return score
+
+    def _tokenize(self, text: str) -> Set[str]:
+        """Tokenize text into a normalized set of terms."""
+
+        return set(re.findall(r"[a-zA-Z0-9]+", text.lower()))

raglib/components/generator.py

@@ -0,0 +1,91 @@
+"""Answer generation component with prompt template support."""
+
+from __future__ import annotations
+
+import logging
+from typing import List, Optional
+
+from raglib.llm.base_client import BaseLLMClient
+from raglib.schemas import Document, GenerationResult
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_PROMPT_TEMPLATE: str = (
+    "You are a precise assistant. Use only the provided context when possible.\n"
+    "Query: {query}\n"
+    "Memory:\n{memory_context}\n"
+    "Context:\n{context}\n"
+    "Answer:"
+)
+
+
+class Generator:
+    """Builds prompts and generates final answers from retrieved context."""
+
+    def __init__(
+        self,
+        llm_client: BaseLLMClient,
+        prompt_template: Optional[str] = None,
+        max_context_tokens: int = 3000,
+    ):
+        """Initialize generator with LLM client and prompt settings."""
+
+        self.llm_client = llm_client
+        self.prompt_template = prompt_template or DEFAULT_PROMPT_TEMPLATE
+        self.max_context_tokens = max_context_tokens
+
+    def generate(
+        self,
+        query: str,
+        documents: List[Document],
+        memory_context: str = "",
+        reasoning_trace: List[str] = [],
+    ) -> GenerationResult:
+        """Generate an answer from query, documents, and optional memory context."""
+
+        trace = list(reasoning_trace)
+        context_docs = self._truncate_documents(documents)
+        context_block = self._build_context(context_docs)
+
+        prompt = self.prompt_template.format(
+            query=query,
+            memory_context=memory_context,
+            context=context_block,
+        )
+
+        context_tokens = sum(self._estimate_tokens(doc.content) for doc in context_docs)
+        logger.info(
+            "Generator sending prompt with %d context docs and %d context tokens",
+            len(context_docs),
+            context_tokens,
+        )
+        answer = self.llm_client.complete(prompt=prompt, system="Answer with grounded reasoning.")
+        trace.append(f"generated_with_docs={len(context_docs)}")
+
+        return GenerationResult(answer=answer, sources=context_docs, reasoning_trace=trace)
+
+    def _truncate_documents(self, documents: List[Document]) -> List[Document]:
+        """Trim document set to fit max context token budget."""
+
+        selected: List[Document] = []
+        used = 0
+        for doc in documents:
+            tokens = self._estimate_tokens(doc.content)
+            if used + tokens > self.max_context_tokens:
+                break
+            selected.append(doc)
+            used += tokens
+        return selected
+
+    def _build_context(self, documents: List[Document]) -> str:
+        """Build a source-tagged context block from documents."""
+
+        lines: List[str] = []
+        for idx, doc in enumerate(documents, start=1):
+            lines.append(f"[{idx}] ({doc.source}) {doc.content}")
+        return "\n\n".join(lines)
+
+    def _estimate_tokens(self, text: str) -> int:
+        """Estimate token usage from text length."""
+
+        return max(1, len(text) // 4)

raglib/components/hybrid_retriever.py

@@ -0,0 +1,99 @@
+"""Hybrid retrieval that blends vector and web search signals."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import replace
+from typing import Dict, List
+
+from raglib.components.retriever import Retriever
+from raglib.components.web_retriever import WebRetriever
+from raglib.schemas import Document
+
+logger = logging.getLogger(__name__)
+
+
+class HybridRetriever:
+    """Combines vector and web retrieval with weighted re-scoring."""
+
+    def __init__(
+        self,
+        vector_retriever: Retriever,
+        web_retriever: WebRetriever,
+        vector_weight: float = 0.6,
+        web_weight: float = 0.4,
+    ):
+        """Initialize weighted retriever dependencies."""
+
+        self.vector_retriever = vector_retriever
+        self.web_retriever = web_retriever
+        total = max(vector_weight + web_weight, 1e-9)
+        self.vector_weight = vector_weight / total
+        self.web_weight = web_weight / total
+
+    def retrieve(self, query: str) -> List[Document]:
+        """Retrieve and merge vector/web results with weighted scores."""
+
+        vector_docs = self.vector_retriever.retrieve(query) if self.vector_weight > 0 else []
+        web_docs = self.web_retriever.retrieve(query) if self.web_weight > 0 else []
+
+        normalized_vector = self._normalize_scores(vector_docs)
+        normalized_web = self._normalize_scores(web_docs)
+
+        merged: Dict[str, Document] = {}
+
+        for doc in normalized_vector:
+            key = self._doc_key(doc)
+            weighted = doc.score * self.vector_weight
+            merged[key] = replace(
+                doc,
+                score=weighted,
+                source=doc.source or "vector",
+                metadata={**doc.metadata, "hybrid_components": ["vector"]},
+            )
+
+        for doc in normalized_web:
+            key = self._doc_key(doc)
+            weighted = doc.score * self.web_weight
+            if key in merged:
+                base = merged[key]
+                components = list(base.metadata.get("hybrid_components", []))
+                components.append("web")
+                merged[key] = replace(
+                    base,
+                    score=base.score + weighted,
+                    source="hybrid",
+                    metadata={**base.metadata, "hybrid_components": components},
+                )
+            else:
+                merged[key] = replace(
+                    doc,
+                    score=weighted,
+                    source=doc.source or "web",
+                    metadata={**doc.metadata, "hybrid_components": ["web"]},
+                )
+
+        ranked = sorted(merged.values(), key=lambda d: d.score, reverse=True)
+        top_k = max(self.vector_retriever.top_k, self.web_retriever.top_k)
+        logger.info("HybridRetriever merged %d documents", len(ranked))
+        return ranked[:top_k]
+
+    def _normalize_scores(self, documents: List[Document]) -> List[Document]:
+        """Normalize scores into [0, 1] for stable weighting."""
+
+        if not documents:
+            return []
+        scores = [doc.score for doc in documents]
+        max_score = max(scores)
+        min_score = min(scores)
+        span = max(max_score - min_score, 1e-9)
+        normalized: List[Document] = []
+        for doc in documents:
+            scaled = (doc.score - min_score) / span if max_score != min_score else 1.0
+            normalized.append(replace(doc, score=scaled))
+        return normalized
+
+    def _doc_key(self, doc: Document) -> str:
+        """Build a stable merge key for duplicate detection."""
+
+        return doc.id or f"doc-{abs(hash(doc.content.lower().strip()))}"

raglib/components/memory.py

@@ -0,0 +1,58 @@
+"""Conversation memory store for RAG turn history."""
+
+from __future__ import annotations
+
+import logging
+from collections import deque
+from dataclasses import dataclass
+from typing import Deque, List
+
+from raglib.schemas import Document
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MemoryTurn:
+    """Represents one conversational turn in memory."""
+
+    query: str
+    answer: str
+    documents: List[Document]
+
+
+class MemoryModule:
+    """Stores and exposes bounded conversational memory context."""
+
+    def __init__(self, max_turns: int = 10):
+        """Initialize bounded memory storage."""
+
+        self.max_turns = max_turns
+        self._turns: Deque[MemoryTurn] = deque(maxlen=max_turns)
+
+    def add(self, query: str, answer: str, documents: List[Document]) -> None:
+        """Add a turn to memory."""
+
+        self._turns.append(MemoryTurn(query=query, answer=answer, documents=documents))
+        logger.info("MemoryModule stored turn count=%d", len(self._turns))
+
+    def get_context(self) -> str:
+        """Return serialized memory context for prompt injection."""
+
+        if not self._turns:
+            return ""
+
+        lines: List[str] = []
+        for idx, turn in enumerate(self._turns, start=1):
+            lines.append(f"Turn {idx} Query: {turn.query}")
+            lines.append(f"Turn {idx} Answer: {turn.answer}")
+            if turn.documents:
+                sources = ", ".join(doc.id for doc in turn.documents[:3])
+                lines.append(f"Turn {idx} Sources: {sources}")
+        return "\n".join(lines)
+
+    def clear(self) -> None:
+        """Clear all memory turns."""
+
+        self._turns.clear()
+        logger.info("MemoryModule cleared")

raglib/components/multi_hop_retriever.py

@@ -0,0 +1,40 @@
+"""Iterative multi-hop retrieval over planned sub-queries."""
+
+from __future__ import annotations
+
+import logging
+from typing import Dict
+
+from raglib.components.planner import Planner
+from raglib.components.retriever import Retriever
+from raglib.schemas import Document, QueryResult
+
+logger = logging.getLogger(__name__)
+
+
+class MultiHopRetriever:
+    """Performs staged retrieval across multiple reasoning hops."""
+
+    def __init__(self, retriever: Retriever, planner: Planner, max_hops: int = 3):
+        """Initialize multi-hop retrieval dependencies."""
+
+        self.retriever = retriever
+        self.planner = planner
+        self.max_hops = max_hops
+
+    def retrieve(self, query: str) -> QueryResult:
+        """Retrieve documents iteratively over planned hops."""
+
+        sub_queries = self.planner.plan(query)[: self.max_hops]
+        logger.info("MultiHopRetriever planned %d hops", len(sub_queries))
+
+        merged: Dict[str, Document] = {}
+        for hop_index, sub_query in enumerate(sub_queries, start=1):
+            logger.info("MultiHopRetriever hop=%d query=%s", hop_index, sub_query)
+            for doc in self.retriever.retrieve(sub_query):
+                key = doc.id or f"doc-{abs(hash(doc.content.lower().strip()))}"
+                if key not in merged or doc.score > merged[key].score:
+                    merged[key] = doc
+
+        ranked = sorted(merged.values(), key=lambda doc: doc.score, reverse=True)
+        return QueryResult(query=query, documents=ranked, hop_count=len(sub_queries))

raglib/components/multi_query_retriever.py

@@ -0,0 +1,63 @@
+"""Multi-query retrieval via semantic expansion and result merging."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import replace
+from typing import Dict, List, Optional
+
+from raglib.components.deduplicator import Deduplicator
+from raglib.components.query_expander import QueryExpander
+from raglib.components.retriever import Retriever
+from raglib.schemas import Document
+
+logger = logging.getLogger(__name__)
+
+
+class MultiQueryRetriever:
+    """Expands a query into variants and merges retrieval results."""
+
+    def __init__(
+        self,
+        retriever: Retriever,
+        query_expander: QueryExpander,
+        deduplicator: Optional[Deduplicator] = None,
+        top_k: int = 5,
+    ):
+        """Initialize variant retrieval dependencies."""
+
+        self.retriever = retriever
+        self.query_expander = query_expander
+        self.deduplicator = deduplicator
+        self.top_k = top_k
+
+    def retrieve(self, query: str) -> List[Document]:
+        """Retrieve documents across query variants and merge scores."""
+
+        variants = self.query_expander.expand(query)
+        logger.info("MultiQueryRetriever generated %d variants", len(variants))
+
+        merged: Dict[str, Document] = {}
+        counts: Dict[str, int] = {}
+
+        for variant in variants:
+            docs = self.retriever.retrieve(variant)
+            for doc in docs:
+                key = doc.id or f"doc-{abs(hash(doc.content.lower().strip()))}"
+                if key in merged:
+                    previous = merged[key]
+                    merged[key] = replace(previous, score=max(previous.score, doc.score))
+                    counts[key] += 1
+                else:
+                    merged[key] = doc
+                    counts[key] = 1
+
+        blended = [
+            replace(doc, score=doc.score * (1.0 + 0.1 * (counts[key] - 1)))
+            for key, doc in merged.items()
+        ]
+
+        ranked = sorted(blended, key=lambda d: d.score, reverse=True)
+        if self.deduplicator:
+            ranked = self.deduplicator.deduplicate(ranked)
+        return ranked[: self.top_k]