persian-rag 0.2.0__py3-none-any.whl

persian_rag/__init__.py

@@ -0,0 +1,121 @@
+"""Persian-first retrieval-augmented generation.
+
+Normalize Farsi text, split it into chunks, and retrieve passages with a
+hybrid of BM25 and dense vectors, optionally rerank them, and evaluate the
+results. Generation is delegated to any callable.
+
+Server-backed vector stores live in :mod:`persian_rag.vectorstores`, PDF and
+DOCX loaders in :mod:`persian_rag.loaders`, and the HTTP API in
+:mod:`persian_rag.server`; their dependencies are optional extras.
+"""
+
+from persian_rag.bm25 import BM25Index
+from persian_rag.cache import (
+    CachedEmbedder,
+    InMemoryEmbeddingCache,
+    LRUCache,
+    SQLiteEmbeddingCache,
+)
+from persian_rag.chunk import Chunk, ChunkingConfig, chunk_text
+from persian_rag.documents import SourceDocument, load_documents, strip_markdown
+from persian_rag.embeddings import (
+    Embedder,
+    EmbedderSpec,
+    HashingEmbedder,
+    OpenAICompatibleEmbedder,
+    SentenceTransformerEmbedder,
+    create_embedder,
+)
+from persian_rag.errors import (
+    ConfigError,
+    IndexFormatError,
+    IndexNotFoundError,
+    MissingDependencyError,
+    PersianRagError,
+)
+from persian_rag.evaluation import (
+    EvalExample,
+    LLMFaithfulnessJudge,
+    RetrievalReport,
+    TokenOverlapFaithfulness,
+    evaluate_answers,
+    evaluate_retrieval,
+    load_eval_examples,
+)
+from persian_rag.loaders import clean_extracted_text, load_docx, load_pdf
+from persian_rag.normalize import NormalizeConfig, normalize
+from persian_rag.persist import SavedIndex, load_index, retriever_from_saved, save_index
+from persian_rag.pipeline import (
+    DEFAULT_PROMPT_EN,
+    DEFAULT_PROMPT_FA,
+    RAGAnswer,
+    RAGPipeline,
+    build_retriever,
+)
+from persian_rag.remote import make_openai_llm
+from persian_rag.rerank import CrossEncoderReranker, LLMReranker, Reranker
+from persian_rag.retrieval import HybridRetriever, RetrievalResult, reciprocal_rank_fusion
+from persian_rag.store import FaissVectorStore, InMemoryVectorStore, VectorHit, VectorStore
+from persian_rag.tokenize import split_sentence_spans, split_sentences, tokenize
+from persian_rag.version import __version__
+
+__all__ = [
+    "DEFAULT_PROMPT_EN",
+    "DEFAULT_PROMPT_FA",
+    "BM25Index",
+    "CachedEmbedder",
+    "Chunk",
+    "ChunkingConfig",
+    "ConfigError",
+    "CrossEncoderReranker",
+    "Embedder",
+    "EmbedderSpec",
+    "EvalExample",
+    "FaissVectorStore",
+    "HashingEmbedder",
+    "HybridRetriever",
+    "InMemoryEmbeddingCache",
+    "InMemoryVectorStore",
+    "IndexFormatError",
+    "IndexNotFoundError",
+    "LLMFaithfulnessJudge",
+    "LLMReranker",
+    "LRUCache",
+    "MissingDependencyError",
+    "NormalizeConfig",
+    "OpenAICompatibleEmbedder",
+    "PersianRagError",
+    "RAGAnswer",
+    "RAGPipeline",
+    "Reranker",
+    "RetrievalReport",
+    "RetrievalResult",
+    "SQLiteEmbeddingCache",
+    "SavedIndex",
+    "SentenceTransformerEmbedder",
+    "SourceDocument",
+    "TokenOverlapFaithfulness",
+    "VectorHit",
+    "VectorStore",
+    "__version__",
+    "build_retriever",
+    "chunk_text",
+    "clean_extracted_text",
+    "create_embedder",
+    "evaluate_answers",
+    "evaluate_retrieval",
+    "load_documents",
+    "load_docx",
+    "load_eval_examples",
+    "load_index",
+    "load_pdf",
+    "make_openai_llm",
+    "normalize",
+    "reciprocal_rank_fusion",
+    "retriever_from_saved",
+    "save_index",
+    "split_sentence_spans",
+    "split_sentences",
+    "strip_markdown",
+    "tokenize",
+]

persian_rag/bm25.py

@@ -0,0 +1,92 @@
+"""Okapi BM25 over in-memory tokenized documents.
+
+The IDF term follows the non-negative Lucene variant:
+
+``log(1 + (N - df + 0.5) / (df + 0.5))``
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+
+
+class BM25Index:
+    """Sparse lexical index. Duplicate document ids are rejected."""
+
+    def __init__(self, *, k1: float = 1.5, b: float = 0.75) -> None:
+        if k1 <= 0:
+            raise ValueError("k1 must be > 0")
+        if not 0.0 <= b <= 1.0:
+            raise ValueError("b must be between 0 and 1")
+        self.k1 = k1
+        self.b = b
+        self._ids: list[str] = []
+        self._positions: dict[str, int] = {}
+        self._term_freqs: list[dict[str, int]] = []
+        self._lengths: list[int] = []
+        self._df: dict[str, int] = {}
+        self._avgdl = 0.0
+
+    def __len__(self) -> int:
+        return len(self._ids)
+
+    def add(self, ids: Sequence[str], documents: Sequence[Sequence[str]]) -> None:
+        """Index tokenized documents aligned with ``ids``."""
+
+        if len(ids) != len(documents):
+            raise ValueError("ids and documents must have the same length")
+        if len(set(ids)) != len(list(ids)):
+            raise ValueError("ids must be unique")
+        for doc_id, tokens in zip(ids, documents, strict=True):
+            if doc_id in self._positions:
+                raise ValueError(f"Duplicate document id: {doc_id}")
+            frequencies: dict[str, int] = {}
+            for token in tokens:
+                frequencies[token] = frequencies.get(token, 0) + 1
+            position = len(self._ids)
+            self._positions[doc_id] = position
+            self._ids.append(doc_id)
+            self._term_freqs.append(frequencies)
+            self._lengths.append(len(list(tokens)))
+            for token in frequencies:
+                self._df[token] = self._df.get(token, 0) + 1
+        self._recompute_avgdl()
+
+    def search(self, query_tokens: Sequence[str], k: int) -> list[tuple[str, float]]:
+        """Return up to ``k`` ``(doc_id, score)`` pairs, best score first."""
+
+        if k < 1:
+            raise ValueError("k must be >= 1")
+        if not self._ids or not query_tokens:
+            return []
+
+        total = len(self._ids)
+        scores = [0.0] * total
+        average = self._avgdl or 1.0
+        for token in dict.fromkeys(query_tokens):
+            df = self._df.get(token, 0)
+            if df == 0:
+                continue
+            idf = math.log(1.0 + (total - df + 0.5) / (df + 0.5))
+            for position, frequencies in enumerate(self._term_freqs):
+                freq = frequencies.get(token, 0)
+                if freq == 0:
+                    continue
+                length_norm = 1.0 - self.b + self.b * self._lengths[position] / average
+                denominator = freq + self.k1 * length_norm
+                scores[position] += idf * (freq * (self.k1 + 1.0)) / denominator
+
+        ranked = [
+            (self._ids[position], scores[position])
+            for position in range(total)
+            if scores[position] > 0.0
+        ]
+        ranked.sort(key=lambda item: (-item[1], item[0]))
+        return ranked[:k]
+
+    def _recompute_avgdl(self) -> None:
+        if not self._lengths:
+            self._avgdl = 0.0
+            return
+        self._avgdl = sum(self._lengths) / len(self._lengths)

persian_rag/cache.py

@@ -0,0 +1,252 @@
+"""Embedding and retrieval caches.
+
+* :class:`CachedEmbedder` wraps any :class:`~persian_rag.embeddings.Embedder`
+  and reuses vectors for texts it has already embedded.
+* :class:`LRUCache` is a small thread-safe in-memory cache with optional
+  expiry. :class:`HybridRetriever` uses it for query results.
+* :class:`SQLiteEmbeddingCache` persists vectors on disk with only the
+  standard library, so repeated ingestion runs skip paid embedding calls.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import sqlite3
+import struct
+import threading
+import time
+from collections import OrderedDict
+from collections.abc import Callable, Sequence
+from pathlib import Path
+from typing import Any, Generic, Protocol, TypeVar, runtime_checkable
+
+from persian_rag.embeddings import Embedder, EmbedderSpec
+
+V = TypeVar("V")
+
+
+class LRUCache(Generic[V]):
+    """Thread-safe least-recently-used cache with an optional time to live.
+
+    ``ttl`` is in seconds; ``None`` keeps entries until they are evicted.
+    """
+
+    def __init__(
+        self,
+        max_size: int = 1024,
+        ttl: float | None = None,
+        *,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> None:
+        if max_size < 1:
+            raise ValueError("max_size must be >= 1")
+        if ttl is not None and ttl <= 0:
+            raise ValueError("ttl must be > 0")
+        self.max_size = max_size
+        self.ttl = ttl
+        self._clock = clock
+        self._data: OrderedDict[str, tuple[float, V]] = OrderedDict()
+        self._lock = threading.Lock()
+        self.hits = 0
+        self.misses = 0
+
+    def get(self, key: str) -> V | None:
+        with self._lock:
+            entry = self._data.get(key)
+            if entry is None:
+                self.misses += 1
+                return None
+            stored_at, value = entry
+            if self.ttl is not None and self._clock() - stored_at > self.ttl:
+                del self._data[key]
+                self.misses += 1
+                return None
+            self._data.move_to_end(key)
+            self.hits += 1
+            return value
+
+    def set(self, key: str, value: V) -> None:
+        with self._lock:
+            self._data[key] = (self._clock(), value)
+            self._data.move_to_end(key)
+            while len(self._data) > self.max_size:
+                self._data.popitem(last=False)
+
+    def clear(self) -> None:
+        with self._lock:
+            self._data.clear()
+
+    def __len__(self) -> int:
+        return len(self._data)
+
+    def __contains__(self, key: object) -> bool:
+        return key in self._data
+
+
+@runtime_checkable
+class EmbeddingCache(Protocol):
+    """Key-value storage for embedding vectors."""
+
+    def get(self, key: str) -> list[float] | None: ...
+
+    def set(self, key: str, value: list[float]) -> None: ...
+
+
+class InMemoryEmbeddingCache(LRUCache[list[float]]):
+    """In-process embedding cache. Defaults to 100,000 vectors."""
+
+    def __init__(self, max_size: int = 100_000, ttl: float | None = None) -> None:
+        super().__init__(max_size=max_size, ttl=ttl)
+
+
+class SQLiteEmbeddingCache:
+    """Persistent embedding cache in a single SQLite file.
+
+    Vectors are stored as packed float64 values. The connection is shared
+    between threads and guarded by a lock.
+    """
+
+    def __init__(self, path: str | Path) -> None:
+        self.path = Path(path)
+        if str(path) != ":memory:":
+            self.path.parent.mkdir(parents=True, exist_ok=True)
+        self._lock = threading.Lock()
+        self._conn = sqlite3.connect(str(path), check_same_thread=False)
+        with self._lock:
+            self._conn.execute(
+                "CREATE TABLE IF NOT EXISTS embeddings (key TEXT PRIMARY KEY, vector BLOB NOT NULL)"
+            )
+            self._conn.commit()
+
+    def get(self, key: str) -> list[float] | None:
+        with self._lock:
+            row = self._conn.execute(
+                "SELECT vector FROM embeddings WHERE key = ?", (key,)
+            ).fetchone()
+        if row is None:
+            return None
+        blob = bytes(row[0])
+        return list(struct.unpack(f"<{len(blob) // 8}d", blob))
+
+    def set(self, key: str, value: list[float]) -> None:
+        blob = struct.pack(f"<{len(value)}d", *value)
+        with self._lock:
+            self._conn.execute(
+                "INSERT OR REPLACE INTO embeddings (key, vector) VALUES (?, ?)", (key, blob)
+            )
+            self._conn.commit()
+
+    def __len__(self) -> int:
+        with self._lock:
+            row = self._conn.execute("SELECT count(*) FROM embeddings").fetchone()
+        return int(row[0])
+
+    def clear(self) -> None:
+        with self._lock:
+            self._conn.execute("DELETE FROM embeddings")
+            self._conn.commit()
+
+    def close(self) -> None:
+        with self._lock:
+            self._conn.close()
+
+    def __enter__(self) -> SQLiteEmbeddingCache:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+
+def embedder_namespace(embedder: Embedder) -> str:
+    """A stable cache namespace for an embedder.
+
+    Uses the embedder's :class:`EmbedderSpec` when it has one, so changing the
+    model, dimension, or hashing settings never returns stale vectors.
+    """
+
+    to_spec = getattr(embedder, "to_spec", None)
+    if callable(to_spec):
+        spec = to_spec()
+        if isinstance(spec, EmbedderSpec):
+            return json.dumps(spec.to_dict(), sort_keys=True, ensure_ascii=False)
+    return f"{type(embedder).__module__}.{type(embedder).__qualname__}:{embedder.dimension}"
+
+
+class CachedEmbedder:
+    """Wrap an embedder and cache its vectors.
+
+    Document and query vectors are cached separately because some models
+    (for example E5) embed them differently. Only missing texts are sent to
+    the wrapped embedder, in one batch.
+    """
+
+    def __init__(
+        self,
+        embedder: Embedder,
+        cache: EmbeddingCache | None = None,
+        *,
+        namespace: str | None = None,
+    ) -> None:
+        self.embedder = embedder
+        self.cache: EmbeddingCache = cache if cache is not None else InMemoryEmbeddingCache()
+        self.namespace = namespace if namespace is not None else embedder_namespace(embedder)
+
+    @property
+    def dimension(self) -> int:
+        return self.embedder.dimension
+
+    def to_spec(self) -> EmbedderSpec:
+        """Return the wrapped embedder's spec so indexes stay portable."""
+
+        to_spec = getattr(self.embedder, "to_spec", None)
+        if not callable(to_spec):
+            raise AttributeError(f"{type(self.embedder).__name__} has no to_spec()")
+        spec = to_spec()
+        if not isinstance(spec, EmbedderSpec):
+            raise TypeError("to_spec() must return an EmbedderSpec")
+        return spec
+
+    def key(self, text: str, kind: str = "doc") -> str:
+        digest = hashlib.sha256()
+        for part in (self.namespace, kind, text):
+            digest.update(part.encode("utf-8"))
+            digest.update(b"\x00")
+        return digest.hexdigest()
+
+    def embed_documents(self, texts: Sequence[str]) -> list[list[float]]:
+        keys = [self.key(text) for text in texts]
+        found: dict[int, list[float]] = {}
+        missing: dict[str, list[int]] = {}
+        for index, (text, key) in enumerate(zip(texts, keys, strict=True)):
+            cached = self.cache.get(key)
+            if cached is not None:
+                found[index] = list(cached)
+            else:
+                missing.setdefault(text, []).append(index)
+        if missing:
+            todo = list(missing)
+            vectors = self.embedder.embed_documents(todo)
+            if len(vectors) != len(todo):
+                raise ValueError("embedder returned an unexpected number of vectors")
+            for text, vector in zip(todo, vectors, strict=True):
+                values = [float(value) for value in vector]
+                self.cache.set(self.key(text), values)
+                for index in missing[text]:
+                    found[index] = list(values)
+        return [found[index] for index in range(len(texts))]
+
+    def embed_query(self, text: str) -> list[float]:
+        key = self.key(text, "query")
+        cached = self.cache.get(key)
+        if cached is not None:
+            return list(cached)
+        vector = [float(value) for value in self.embedder.embed_query(text)]
+        self.cache.set(key, vector)
+        return vector
+
+
+def result_cache_key(*parts: Any) -> str:
+    """Build a cache key from JSON-serializable parts."""
+
+    return json.dumps(parts, ensure_ascii=False, sort_keys=True, default=str)