ragforge-sdk 1.0.0__py3-none-any.whl

__init__.py

@@ -0,0 +1,10 @@
+from .ragforge.pipeline import SearchPipeline
+from .ragforge.models import FastembedEmbedder, FastembedReranker
+from .ragforge.retrieval import BM25Retriever, VectorRetriever, HybridRetriever
+from .ragforge.fusion import RRFFusion, PositionAwareBlend, AdaptiveFusion
+from .ragforge.llm import LLMClient
+from .ragforge.query import QueryPlanner
+from .ragforge.evaluation import LLMJudge, Evaluator
+from .ragforge.cache import SemanticCache
+from .ragforge.dedup import Deduplicator
+from .ragforge.config import ModelConfig, PipelineConfig, LLMConfig, QueryTransformStrategy

ragforge/__init__.py

@@ -0,0 +1,21 @@
+from .cache.semantic_cache import SemanticCache
+from .config.llm_config import LLMConfig
+from .config.model_config import ModelConfig
+from .config.pipeline_config import PipelineConfig, QueryTransformStrategy
+from .dedup.deduplicator import Deduplicator
+from .evaluation.evaluator import Evaluator
+from .evaluation.judge import LLMJudge
+from .fusion.adaptive import AdaptiveFusion
+from .fusion.blend import PositionAwareBlend
+from .fusion.rrf import RRFFusion
+from .llm.llm_client import LLMClient
+from .models.embedding import FastembedEmbedder
+from .models.reranker import FastembedReranker
+from .query.planner import QueryPlanner
+from .retrieval.bm25 import BM25Retriever
+from .retrieval.hybrid import HybridRetriever
+from .retrieval.vector import VectorRetriever
+from .tracing.trace import Tracer
+
+__version__ = "1.0.0"
+__author__ = "jiangnanboy"

ragforge/cache/__init__.py

@@ -0,0 +1,4 @@
+"""Cache exports."""
+from .semantic_cache import SemanticCache
+__all__ = ["SemanticCache"]
+

ragforge/cache/semantic_cache.py

@@ -0,0 +1,191 @@
+"""
+Semantic Cache
+==============
+Cache search results keyed by query embedding similarity.
+
+Features:
+    - Embeds queries and finds cached entries by cosine similarity
+    - Configurable similarity threshold
+    - TTL-based expiration
+    - Thread-safe (basic)
+    - Uses any ``Embedder`` implementation
+"""
+
+from __future__ import annotations
+
+import time
+
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
+
+from ..protocols import Embedder
+
+class _CacheEntry:
+    """Internal cache entry."""
+
+    __slots__ = ("query", "query_embedding", "results", "timestamp")
+
+    def __init__(
+        self,
+        query: str,
+        query_embedding: np.ndarray,
+        results: list[tuple[str, float]],
+    ) -> None:
+        self.query = query
+        self.query_embedding = query_embedding
+        self.results = results
+        self.timestamp = time.time()
+
+
+class SemanticCache:
+    """Semantic cache for RAG search results.
+
+    Caches pipeline results keyed by query embedding. When a new query
+    is similar enough to a cached one, returns the cached results instead
+    of re-running the pipeline.
+
+    Args:
+        embedder:
+            Any ``Embedder`` for computing query embeddings.
+        similarity_threshold:
+            Minimum cosine similarity to consider a cache hit (default 0.95).
+        ttl_seconds:
+            Time-to-live for cache entries in seconds (default 3600).
+            Set to 0 for no expiration.
+        max_size:
+            Maximum number of entries to keep in cache (default 1000).
+
+    Example::
+
+        from ragforge import SemanticCache, FastembedEmbedder, ModelConfig
+
+        cache = SemanticCache(
+            embedder=FastembedEmbedder(ModelConfig()),
+            similarity_threshold=0.95,
+        )
+
+        # First call: runs pipeline, stores result
+        result = cache.get_or_search(
+            "苹果手机价格",
+            search_fn=lambda q, d: pipeline.search(q, d),
+            documents=docs,
+        )
+
+        # Similar query: returns cached result instantly
+        result = cache.get_or_search(
+            "苹果手机多少钱",
+            search_fn=lambda q, d: pipeline.search(q, d),
+            documents=docs,
+        )
+    """
+
+    def __init__(
+        self,
+        embedder: Embedder,
+        similarity_threshold: float = 0.95,
+        ttl_seconds: float = 3600,
+        max_size: int = 1000,
+    ) -> None:
+        self._embedder = embedder
+        self._threshold = similarity_threshold
+        self._ttl = ttl_seconds
+        self._max_size = max_size
+        self._entries: list[_CacheEntry] = []
+        self._hits = 0
+        self._misses = 0
+
+    def get(
+        self, query: str
+    ) -> list[tuple[str, float]] | None:
+        """Look up a cached result for the query.
+
+        Args:
+            query: Search query.
+
+        Returns:
+            Cached results if a hit is found (similarity > threshold),
+            ``None`` otherwise.
+        """
+        now = time.time()
+        query_emb = self._embedder.embed(query).reshape(1, -1)
+
+        for entry in self._entries:
+            # Check TTL
+            if self._ttl > 0 and (now - entry.timestamp) > self._ttl:
+                continue
+            # Check similarity
+            sim = cosine_similarity(query_emb, entry.query_embedding.reshape(1, -1))[0][0]
+            if sim >= self._threshold:
+                self._hits += 1
+                return entry.results
+
+        self._misses += 1
+        return None
+
+    def put(
+        self,
+        query: str,
+        results: list[tuple[str, float]],
+    ) -> None:
+        """Store a search result in the cache.
+
+        Args:
+            query: The query that produced these results.
+            results: Pipeline output to cache.
+        """
+        query_emb = self._embedder.embed(query)
+
+        # Evict oldest if at capacity
+        if len(self._entries) >= self._max_size:
+            self._entries.pop(0)
+
+        self._entries.append(_CacheEntry(query, query_emb, results))
+
+    def get_or_search(
+        self,
+        query: str,
+        search_fn: object,
+        documents: list[str],
+    ) -> list[tuple[str, float]]:
+        """Get cached result or execute search function.
+
+        Args:
+            query: Search query.
+            search_fn:
+                Callable ``search_fn(query, documents) -> results``.
+            documents: Candidate documents.
+
+        Returns:
+            Either cached or fresh search results.
+        """
+        cached = self.get(query)
+        if cached is not None:
+            return cached
+
+        results = search_fn(query, documents)
+        self.put(query, results)
+        return results
+
+    def clear(self) -> int:
+        """Clear all cache entries.
+
+        Returns:
+            Number of entries cleared.
+        """
+        count = len(self._entries)
+        self._entries.clear()
+        self._hits = 0
+        self._misses = 0
+        return count
+
+    @property
+    def stats(self) -> dict[str, int | float]:
+        """Cache statistics."""
+        total = self._hits + self._misses
+        hit_rate = self._hits / total if total > 0 else 0.0
+        return {
+            "entries": len(self._entries),
+            "hits": self._hits,
+            "misses": self._misses,
+            "hit_rate": round(hit_rate, 4),
+        }

ragforge/config/__init__.py

@@ -0,0 +1,8 @@
+"""Configuration exports."""
+from .llm_config import LLMConfig
+from .model_config import ModelConfig
+from .pipeline_config import PipelineConfig, QueryTransformStrategy
+
+__all__ = ["ModelConfig", "PipelineConfig", "LLMConfig", "QueryTransformStrategy"]
+
+

ragforge/config/llm_config.py

@@ -0,0 +1,38 @@
+"""
+LLM Configuration
+=================
+Configuration for LLM backends (DeepSeek, etc.).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class LLMConfig:
+    """Configuration for LLM API access.
+
+    Attributes:
+        api_key:
+            API key for the LLM service.
+        base_url:
+            Base URL of the LLM API endpoint.
+            Defaults to DeepSeek API.
+        model:
+            Model name to use.
+            Defaults to ``deepseek-chat``.
+        temperature:
+            Sampling temperature for generation (0.0 = deterministic).
+        max_tokens:
+            Maximum number of tokens in the response.
+        timeout:
+            Request timeout in seconds.
+    """
+
+    api_key: str = ""
+    base_url: str = "https://api.deepseek.com"
+    model: str = "deepseek-v4-flash"
+    temperature: float = 0.1
+    max_tokens: int = 1024
+    timeout: int = 60

ragforge/config/model_config.py

@@ -0,0 +1,45 @@
+"""
+Model Configuration
+===================
+Dataclass holding model paths and loading parameters.
+
+Separation of Concerns:
+    - *ModelConfig* → model paths, file names, dimensions
+    - *PipelineConfig* → algorithm hyper-parameters
+
+This makes it easy to swap models without touching pipeline logic.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class ModelConfig:
+    """Configuration for embedding and reranker models.
+
+    Attributes:
+        embedding_model_path:
+            Local directory containing the embedding ONNX model files.
+            If ``None``, the model will be auto-downloaded on first use.
+        embedding_onnx_file:
+            Name of the embedding ONNX weight file (default ``model.onnx``).
+        embedding_dim:
+            Dimensionality of the embedding vectors.
+        rerank_model_path:
+            Local directory containing the reranker ONNX model files.
+            If ``None``, the model will be auto-downloaded on first use.
+        rerank_onnx_file:
+            Name of the reranker ONNX weight file
+            (default ``model_int8.onnx``).
+    """
+
+    # --- Embedding model ---
+    embedding_model_path: str | None = None
+    embedding_onnx_file: str = "model.onnx"
+    embedding_dim: int = 384
+
+    # --- Reranker model ---
+    rerank_model_path: str | None = None
+    rerank_onnx_file: str = "model_int8.onnx"

ragforge/config/pipeline_config.py

@@ -0,0 +1,78 @@
+"""
+Pipeline Configuration
+======================
+Dataclass holding algorithm hyper-parameters for the search pipeline.
+
+Separation of Concerns:
+    - *ModelConfig* → model paths, file names, dimensions
+    - *PipelineConfig* → algorithm hyper-parameters
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+
+
+class QueryTransformStrategy(str, Enum):
+    """Strategy for how query transform results are used in the pipeline.
+
+    Attributes:
+        REPLACE:
+            Replace the original query with the transformed version.
+            This is the default and original behavior.
+        RETRIEVE_AND_FUSE:
+            Retrieve with both the original query AND the transformed
+            query, then fuse all result sets together. This implements
+            the **Multi-Query Retrieval** pattern for higher recall.
+
+            - If ``transform()`` returns a single string, two retrievals
+              are performed (original + rewritten) and results are fused.
+            - If ``transform()`` returns a list of sub-queries, N+1
+              retrievals are performed (original + each sub-query) and
+              all results are fused via RRF.
+    """
+
+    REPLACE = "replace"
+    RETRIEVE_AND_FUSE = "retrieve_and_fuse"
+
+
+@dataclass
+class PipelineConfig:
+    """Hyper-parameters for the retrieval & fusion pipeline.
+
+    Attributes:
+        rrf_k:
+            Reciprocal Rank Fusion constant. Higher values flatten
+            score differences across ranks.
+        top_k_recall:
+            Number of candidates to keep after fusion (before reranking).
+        query_weight:
+            Multiplier applied to RRF base scores.
+        bonus_rank1:
+            Score bonus for the top-1 document after RRF fusion.
+        bonus_rank2_3:
+            Score bonus for documents ranked 2-3 after RRF fusion.
+        blend_weights:
+            Position-aware blending weights mapping rank buckets
+            to ``(retrieval_weight, reranker_weight)`` dicts.
+        query_transform_strategy:
+            How to apply query transformation results.
+            ``REPLACE`` (default): use rewritten query directly.
+            ``RETRIEVE_AND_FUSE``: retrieve with both original and
+            rewritten queries, then fuse results.
+    """
+
+    rrf_k: int = 60
+    top_k_recall: int = 30
+    query_weight: float = 2.0
+    bonus_rank1: float = 0.05
+    bonus_rank2_3: float = 0.02
+    query_transform_strategy: QueryTransformStrategy = QueryTransformStrategy.REPLACE
+    blend_weights: dict[str, dict[str, float]] = field(
+        default_factory=lambda: {
+            "top1-3": {"retrieval": 0.75, "reranker": 0.25},
+            "top4-10": {"retrieval": 0.60, "reranker": 0.40},
+            "top11+": {"retrieval": 0.40, "reranker": 0.60},
+        }
+    )

ragforge/dedup/__init__.py

@@ -0,0 +1,5 @@
+"""Deduplication exports."""
+from .deduplicator import Deduplicator
+
+__all__ = ["Deduplicator"]
+

ragforge/dedup/deduplicator.py

@@ -0,0 +1,122 @@
+"""
+Document Deduplicator
+=====================
+Identifies and removes near-duplicate documents using embedding similarity.
+
+Features:
+    - Embedding-based near-duplicate detection
+    - Configurable similarity threshold
+    - Returns deduplicated document list
+    - Optionally returns duplicate clusters for inspection
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
+
+from ..protocols import Embedder
+
+
+class Deduplicator:
+    """Document deduplicator using embedding similarity.
+
+    Args:
+        embedder:
+            Any ``Embedder`` for computing document embeddings.
+        threshold:
+            Cosine similarity threshold above which documents
+            are considered duplicates (default 0.95).
+
+    Example::
+
+        from ragforge import Deduplicator, FastembedEmbedder, ModelConfig
+
+        dedup = Deduplicator(
+            embedder=FastembedEmbedder(ModelConfig()),
+            threshold=0.95,
+        )
+
+        documents = ["iPhone 价格", "iPhone价格", "华为手机报价"]
+        unique = dedup.deduplicate(documents)
+        print(f"{len(documents)} → {len(unique)} unique")
+    """
+
+    def __init__(
+        self,
+        embedder: Embedder,
+        threshold: float = 0.95,
+    ) -> None:
+        self._embedder = embedder
+        self._threshold = threshold
+
+    def deduplicate(
+        self,
+        documents: list[str],
+    ) -> list[str]:
+        """Remove near-duplicate documents.
+
+        Keeps the first occurrence of each duplicate cluster.
+
+        Args:
+            documents: List of document strings.
+
+        Returns:
+            Deduplicated list preserving original order.
+        """
+        if not documents:
+            return []
+
+        embeddings = self._embedder.embed_batch(documents)
+        emb_matrix = np.array(embeddings)
+        sim_matrix = cosine_similarity(emb_matrix)
+
+        keep: list[bool] = [True] * len(documents)
+        for i in range(len(documents)):
+            if not keep[i]:
+                continue
+            for j in range(i + 1, len(documents)):
+                if not keep[j]:
+                    continue
+                if sim_matrix[i][j] >= self._threshold:
+                    keep[j] = False  # j is a duplicate of i
+
+        return [doc for doc, k in zip(documents, keep) if k]
+
+    def find_clusters(
+        self,
+        documents: list[str],
+    ) -> list[list[int]]:
+        """Find groups of near-duplicate document indices.
+
+        Args:
+            documents: List of document strings.
+
+        Returns:
+            List of clusters, where each cluster is a list of
+            document indices. Singletons (non-duplicates) are
+            not included in the output.
+        """
+        if not documents:
+            return []
+
+        embeddings = self._embedder.embed_batch(documents)
+        emb_matrix = np.array(embeddings)
+        sim_matrix = cosine_similarity(emb_matrix)
+
+        visited: set[int] = set()
+        clusters: list[list[int]] = []
+
+        for i in range(len(documents)):
+            if i in visited:
+                continue
+            cluster = [i]
+            for j in range(i + 1, len(documents)):
+                if j not in visited and sim_matrix[i][j] >= self._threshold:
+                    cluster.append(j)
+                    visited.add(j)
+            if len(cluster) > 1:
+                clusters.append(sorted(cluster))
+            visited.add(i)
+
+        return clusters

ragforge/evaluation/__init__.py

@@ -0,0 +1,7 @@
+"""Evaluation exports."""
+from .evaluator import Evaluator
+from .judge import LLMJudge
+
+__all__ = ["LLMJudge", "Evaluator"]
+
+