microrag 0.1.0__py3-none-any.whl

microrag/__init__.py

@@ -0,0 +1,57 @@
+"""MicroRAG - A feature-rich, universal RAG library for Python.
+
+MicroRAG provides:
+- ONNX-backed embeddings (CPU-only, no PyTorch at runtime)
+- DuckDB storage with HNSW vector indexes
+- Three-tier hybrid search (semantic + BM25 + FTS) with RRF fusion
+- Query preprocessing with abbreviation expansion
+
+Example:
+    ```python
+    from microrag import MicroRAG, RAGConfig
+
+    config = RAGConfig(
+        model_path="/path/to/all-MiniLM-L6-v2",
+        db_path="./rag.duckdb",
+    )
+
+    with MicroRAG(config) as rag:
+        rag.add_documents(["Document 1", "Document 2"])
+        rag.build_index()
+        results = rag.search("query")
+    ```
+"""
+
+from microrag.config import RAGConfig
+from microrag.core import MicroRAG
+from microrag.exceptions import (
+    ConfigurationError,
+    DocumentError,
+    EmbeddingError,
+    IndexError,
+    MicroRAGError,
+    SearchError,
+    StorageError,
+)
+from microrag.models import Document, SearchResult
+
+__version__ = "0.1.0"
+
+__all__ = [
+    # Main classes
+    "MicroRAG",
+    "RAGConfig",
+    # Data models
+    "Document",
+    "SearchResult",
+    # Exceptions
+    "MicroRAGError",
+    "ConfigurationError",
+    "DocumentError",
+    "EmbeddingError",
+    "IndexError",
+    "SearchError",
+    "StorageError",
+    # Version
+    "__version__",
+]

microrag/config.py

@@ -0,0 +1,86 @@
+"""Configuration for MicroRAG."""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from microrag.stopwords import ENGLISH_STOPWORDS
+
+
+@dataclass(frozen=True)
+class RAGConfig:
+    """Configuration for MicroRAG instance.
+
+    Attributes:
+        model_path: Path to sentence-transformer model directory.
+        model_file: ONNX model filename within model_path (for quantized models).
+        db_path: DuckDB database path. Use ":memory:" for in-memory database.
+        embedding_dim: Dimension of embedding vectors.
+        chunk_size: Maximum size of text chunks in characters.
+        chunk_overlap: Overlap between consecutive chunks.
+        hybrid_enabled: Enable hybrid search (semantic + BM25 + FTS).
+        hybrid_alpha: Weight for semantic search in fusion (0-1).
+        similarity_threshold: Minimum score threshold for results.
+        abbreviations: Mapping of abbreviations to expansions for query processing.
+        stopwords: Set of stopwords for BM25 tokenization.
+        remove_stopwords: Enable stopword removal in BM25 tokenization.
+        hnsw_ef_construction: HNSW index build-time parameter.
+        hnsw_ef_search: HNSW index search-time parameter.
+        hnsw_m: HNSW index M parameter (connections per layer).
+        batch_size: Batch size for embedding generation.
+    """
+
+    model_path: str
+    model_file: str | None = None
+    db_path: str = ":memory:"
+    embedding_dim: int = 384
+    chunk_size: int = 1000
+    chunk_overlap: int = 200
+    hybrid_enabled: bool = True
+    hybrid_alpha: float = 0.7
+    similarity_threshold: float = 0.4
+    abbreviations: dict[str, str] = field(default_factory=dict)
+    stopwords: frozenset[str] = field(default_factory=lambda: ENGLISH_STOPWORDS)
+    remove_stopwords: bool = True
+    hnsw_ef_construction: int = 200
+    hnsw_ef_search: int = 100
+    hnsw_m: int = 16
+    batch_size: int = 32
+
+    def __post_init__(self) -> None:
+        if not self.model_path:
+            raise ValueError("model_path is required")
+        if self.embedding_dim <= 0:
+            raise ValueError("embedding_dim must be positive")
+        if self.chunk_size <= 0:
+            raise ValueError("chunk_size must be positive")
+        if self.chunk_overlap < 0:
+            raise ValueError("chunk_overlap cannot be negative")
+        if self.chunk_overlap >= self.chunk_size:
+            raise ValueError("chunk_overlap must be less than chunk_size")
+        if not 0 <= self.hybrid_alpha <= 1:
+            raise ValueError("hybrid_alpha must be between 0 and 1")
+        if not 0 <= self.similarity_threshold <= 1:
+            raise ValueError("similarity_threshold must be between 0 and 1")
+
+    def with_updates(self, **kwargs: Any) -> "RAGConfig":
+        """Create a new config with updated values."""
+        current: dict[str, Any] = {
+            "model_path": self.model_path,
+            "model_file": self.model_file,
+            "db_path": self.db_path,
+            "embedding_dim": self.embedding_dim,
+            "chunk_size": self.chunk_size,
+            "chunk_overlap": self.chunk_overlap,
+            "hybrid_enabled": self.hybrid_enabled,
+            "hybrid_alpha": self.hybrid_alpha,
+            "similarity_threshold": self.similarity_threshold,
+            "abbreviations": dict(self.abbreviations),
+            "stopwords": self.stopwords,
+            "remove_stopwords": self.remove_stopwords,
+            "hnsw_ef_construction": self.hnsw_ef_construction,
+            "hnsw_ef_search": self.hnsw_ef_search,
+            "hnsw_m": self.hnsw_m,
+            "batch_size": self.batch_size,
+        }
+        current.update(kwargs)
+        return RAGConfig(**current)

microrag/core.py

@@ -0,0 +1,293 @@
+"""Core MicroRAG class."""
+
+from collections.abc import Sequence
+from types import TracebackType
+from typing import Any
+
+from microrag.config import RAGConfig
+from microrag.embedding import EmbeddingModel
+from microrag.exceptions import DocumentError, MicroRAGError
+from microrag.models import Document, SearchResult
+from microrag.query_processor import QueryProcessor
+from microrag.search.hybrid import HybridSearcher
+from microrag.storage import DuckDBStorage
+from microrag.utils import chunk_text, generate_id, normalize_document_input
+
+
+class MicroRAG:
+    """Main MicroRAG class for document indexing and retrieval.
+
+    MicroRAG provides a simple API for:
+    - Adding documents with automatic embedding generation
+    - Building HNSW, BM25, and FTS indexes
+    - Hybrid search with RRF fusion
+
+    Example:
+        ```python
+        config = RAGConfig(
+            model_path="/path/to/model",
+            db_path="./rag.duckdb",
+        )
+
+        with MicroRAG(config) as rag:
+            rag.add_documents(["Doc 1", "Doc 2"])
+            rag.build_index()
+            results = rag.search("query")
+        ```
+
+    Args:
+        config: RAGConfig instance with all settings.
+    """
+
+    def __init__(self, config: RAGConfig) -> None:
+        self._config = config
+        self._storage: DuckDBStorage | None = None
+        self._embedding_model: EmbeddingModel | None = None
+        self._query_processor: QueryProcessor | None = None
+        self._searcher: HybridSearcher | None = None
+        self._documents: list[Document] = []
+        self._index_built = False
+
+    def __enter__(self) -> "MicroRAG":
+        """Enter context manager."""
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exit context manager and close resources."""
+        self.close()
+
+    @property
+    def config(self) -> RAGConfig:
+        """Get the configuration."""
+        return self._config
+
+    @property
+    def storage(self) -> DuckDBStorage:
+        """Get or create the storage adapter."""
+        if self._storage is None:
+            self._storage = DuckDBStorage(
+                db_path=self._config.db_path,
+                embedding_dim=self._config.embedding_dim,
+            )
+        return self._storage
+
+    @property
+    def embedding_model(self) -> EmbeddingModel:
+        """Get or create the embedding model."""
+        if self._embedding_model is None:
+            self._embedding_model = EmbeddingModel(
+                model_path=self._config.model_path,
+                model_file=self._config.model_file,
+                batch_size=self._config.batch_size,
+            )
+        return self._embedding_model
+
+    @property
+    def query_processor(self) -> QueryProcessor:
+        """Get or create the query processor."""
+        if self._query_processor is None:
+            self._query_processor = QueryProcessor(
+                abbreviations=dict(self._config.abbreviations),
+                stopwords=self._config.stopwords,
+                remove_stopwords=self._config.remove_stopwords,
+            )
+        return self._query_processor
+
+    @property
+    def searcher(self) -> HybridSearcher:
+        """Get or create the hybrid searcher."""
+        if self._searcher is None:
+            self._searcher = HybridSearcher(
+                storage=self.storage,
+                embedding_model=self.embedding_model,
+                query_processor=self.query_processor,
+                alpha=self._config.hybrid_alpha,
+            )
+        return self._searcher
+
+    def add_documents(
+        self,
+        documents: Sequence[str | dict[str, Any] | Document],
+        chunk: bool = True,
+    ) -> list[str]:
+        """Add documents to the RAG system.
+
+        Documents can be provided in various formats:
+        - str: Plain text content
+        - dict: {"content": str, "metadata": dict, "id": str}
+        - Document: Document dataclass
+
+        Large documents are automatically chunked if chunk=True.
+
+        Args:
+            documents: Sequence of documents in any supported format.
+            chunk: Whether to chunk large documents.
+
+        Returns:
+            List of document IDs that were added.
+        """
+        doc_ids = []
+
+        for doc in documents:
+            try:
+                doc_id, content, metadata = normalize_document_input(doc)
+            except ValueError as e:
+                raise DocumentError(str(e)) from e
+
+            # Chunk if needed
+            if chunk and len(content) > self._config.chunk_size:
+                chunks = chunk_text(
+                    content,
+                    chunk_size=self._config.chunk_size,
+                    chunk_overlap=self._config.chunk_overlap,
+                )
+
+                for i, chunk_content in enumerate(chunks):
+                    chunk_id = doc_id or generate_id(chunk_content)
+                    if len(chunks) > 1:
+                        chunk_id = f"{chunk_id}_{i}"
+
+                    chunk_metadata = {**metadata, "_chunk_index": i, "_total_chunks": len(chunks)}
+
+                    doc_obj = Document(
+                        id=chunk_id,
+                        content=chunk_content,
+                        metadata=chunk_metadata,
+                    )
+                    self._documents.append(doc_obj)
+                    doc_ids.append(chunk_id)
+            else:
+                final_id = doc_id or generate_id(content)
+                doc_obj = Document(
+                    id=final_id,
+                    content=content,
+                    metadata=metadata,
+                )
+                self._documents.append(doc_obj)
+                doc_ids.append(final_id)
+
+        self._index_built = False
+        return doc_ids
+
+    def build_index(self) -> None:
+        """Build all indexes (HNSW, BM25, FTS).
+
+        This method:
+        1. Generates embeddings for all documents
+        2. Stores documents in DuckDB
+        3. Builds HNSW vector index
+        4. Builds BM25 index
+        5. Builds FTS index
+
+        Must be called after add_documents() and before search().
+        """
+        if not self._documents:
+            return
+
+        # Generate embeddings for documents without them
+        docs_needing_embeddings = [d for d in self._documents if d.embedding is None]
+
+        if docs_needing_embeddings:
+            contents = [d.content for d in docs_needing_embeddings]
+            embeddings = self.embedding_model.encode(contents)
+
+            for doc, embedding in zip(docs_needing_embeddings, embeddings, strict=True):
+                doc.embedding = embedding
+
+        # Store documents
+        self.storage.add_documents(self._documents)
+
+        # Build HNSW index
+        self.storage.build_vector_index(
+            ef_construction=self._config.hnsw_ef_construction,
+            ef_search=self._config.hnsw_ef_search,
+            m=self._config.hnsw_m,
+        )
+
+        # Build FTS index
+        self.storage.build_fts_index()
+
+        # Build BM25 index
+        self.searcher.build_index(self._documents)
+
+        self._index_built = True
+
+    def search(
+        self,
+        query: str,
+        top_k: int = 10,
+        hybrid: bool | None = None,
+        threshold: float | None = None,
+    ) -> list[SearchResult]:
+        """Search for documents matching the query.
+
+        Args:
+            query: Search query string.
+            top_k: Maximum number of results to return.
+            hybrid: Override hybrid search setting from config.
+            threshold: Override similarity threshold from config.
+
+        Returns:
+            List of SearchResult objects sorted by relevance.
+        """
+        if not self._index_built:
+            raise MicroRAGError("Index not built. Call build_index() first.")
+
+        hybrid_enabled = hybrid if hybrid is not None else self._config.hybrid_enabled
+        similarity_threshold = (
+            threshold if threshold is not None else self._config.similarity_threshold
+        )
+
+        return self.searcher.search(
+            query=query,
+            top_k=top_k,
+            hybrid_enabled=hybrid_enabled,
+            similarity_threshold=similarity_threshold,
+        )
+
+    def get_document(self, doc_id: str) -> Document | None:
+        """Retrieve a document by ID.
+
+        Args:
+            doc_id: Document ID.
+
+        Returns:
+            Document if found, None otherwise.
+        """
+        return self.storage.get_document(doc_id)
+
+    def get_all_documents(self) -> list[Document]:
+        """Retrieve all documents from storage.
+
+        Returns:
+            List of all documents.
+        """
+        return self.storage.get_all_documents()
+
+    def count(self) -> int:
+        """Get the number of documents in storage.
+
+        Returns:
+            Document count.
+        """
+        return self.storage.count()
+
+    def clear(self) -> None:
+        """Remove all documents and reset indexes."""
+        self.storage.clear()
+        self._documents.clear()
+        self._index_built = False
+
+    def close(self) -> None:
+        """Close all resources."""
+        if self._storage is not None:
+            self._storage.close()
+            self._storage = None
+        self._embedding_model = None
+        self._query_processor = None
+        self._searcher = None

microrag/embedding/__init__.py

@@ -0,0 +1,5 @@
+"""Embedding module for MicroRAG."""
+
+from microrag.embedding.model import EmbeddingModel
+
+__all__ = ["EmbeddingModel"]

microrag/embedding/model.py

@@ -0,0 +1,118 @@
+"""ONNX-backed embedding model using sentence-transformers."""
+
+from collections.abc import Sequence
+from pathlib import Path
+
+import numpy as np
+from numpy.typing import NDArray
+from sentence_transformers import SentenceTransformer
+
+from microrag.exceptions import EmbeddingError
+
+
+class EmbeddingModel:
+    """Embedding model using sentence-transformers with ONNX backend.
+
+    This class wraps sentence-transformers to provide ONNX-optimized
+    embeddings for CPU inference.
+
+    Args:
+        model_path: Path to the sentence-transformer model directory.
+        model_file: Optional ONNX model filename for quantized models.
+        batch_size: Batch size for encoding.
+    """
+
+    def __init__(
+        self,
+        model_path: str,
+        model_file: str | None = None,
+        batch_size: int = 32,
+    ) -> None:
+        self._model_path = Path(model_path)
+        self._model_file = model_file
+        self._batch_size = batch_size
+        self._model: SentenceTransformer | None = None
+
+    def _load_model(self) -> SentenceTransformer:
+        """Load the sentence-transformer model with ONNX backend."""
+        if self._model is not None:
+            return self._model
+
+        model_path = str(self._model_path)
+
+        try:
+            # Load with ONNX backend
+            model_kwargs: dict[str, dict[str, str]] = {}
+
+            if self._model_file:
+                model_kwargs["model_kwargs"] = {"file_name": self._model_file}
+
+            self._model = SentenceTransformer(
+                model_path,
+                backend="onnx",
+                **model_kwargs,
+            )
+            return self._model
+        except Exception as e:
+            raise EmbeddingError(f"Failed to load model from {model_path}: {e}") from e
+
+    @property
+    def model(self) -> SentenceTransformer:
+        """Get the loaded model, loading it if necessary."""
+        return self._load_model()
+
+    @property
+    def embedding_dim(self) -> int:
+        """Get the embedding dimension."""
+        return self.model.get_sentence_embedding_dimension() or 384
+
+    def encode(
+        self,
+        texts: str | Sequence[str],
+        normalize: bool = True,
+    ) -> NDArray[np.float32]:
+        """Encode texts into embeddings.
+
+        Args:
+            texts: Single text or sequence of texts to encode.
+            normalize: Whether to L2-normalize embeddings.
+
+        Returns:
+            Array of shape (n_texts, embedding_dim) with float32 embeddings.
+        """
+        if isinstance(texts, str):
+            texts = [texts]
+
+        if not texts:
+            return np.array([], dtype=np.float32).reshape(0, self.embedding_dim)
+
+        try:
+            embeddings = self.model.encode(
+                list(texts),
+                batch_size=self._batch_size,
+                normalize_embeddings=normalize,
+                convert_to_numpy=True,
+                show_progress_bar=False,
+            )
+            result: NDArray[np.float32] = np.asarray(embeddings, dtype=np.float32)
+            return result
+        except Exception as e:
+            raise EmbeddingError(f"Failed to encode texts: {e}") from e
+
+    def encode_single(
+        self,
+        text: str,
+        normalize: bool = True,
+    ) -> NDArray[np.float32]:
+        """Encode a single text into an embedding.
+
+        Args:
+            text: Text to encode.
+            normalize: Whether to L2-normalize the embedding.
+
+        Returns:
+            Array of shape (embedding_dim,) with float32 embedding.
+        """
+        embeddings = self.encode([text], normalize=normalize)
+        result: NDArray[np.float32] = embeddings[0]
+        return result

microrag/exceptions.py

@@ -0,0 +1,29 @@
+"""MicroRAG exception hierarchy."""
+
+
+class MicroRAGError(Exception):
+    """Base exception for MicroRAG errors."""
+
+
+class ConfigurationError(MicroRAGError):
+    """Raised when configuration is invalid."""
+
+
+class EmbeddingError(MicroRAGError):
+    """Raised when embedding generation fails."""
+
+
+class StorageError(MicroRAGError):
+    """Raised when storage operations fail."""
+
+
+class IndexError(MicroRAGError):
+    """Raised when index operations fail."""
+
+
+class SearchError(MicroRAGError):
+    """Raised when search operations fail."""
+
+
+class DocumentError(MicroRAGError):
+    """Raised when document operations fail."""

microrag/models.py

@@ -0,0 +1,45 @@
+"""Data models for MicroRAG."""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+@dataclass
+class Document:
+    """A document with content, metadata, and optional embedding."""
+
+    id: str
+    content: str
+    metadata: dict[str, Any] = field(default_factory=dict)
+    embedding: NDArray[np.float32] | None = None
+
+    def __post_init__(self) -> None:
+        if not self.content:
+            raise ValueError("Document content cannot be empty")
+
+
+@dataclass
+class SearchResult:
+    """A search result containing document, score, and rank."""
+
+    document: Document
+    score: float
+    rank: int
+
+    @property
+    def id(self) -> str:
+        """Convenience accessor for document ID."""
+        return self.document.id
+
+    @property
+    def content(self) -> str:
+        """Convenience accessor for document content."""
+        return self.document.content
+
+    @property
+    def metadata(self) -> dict[str, Any]:
+        """Convenience accessor for document metadata."""
+        return self.document.metadata