msaas-rag 1.0.0__tar.gz

msaas_rag-1.0.0/.gitignore

@@ -0,0 +1,21 @@
+node_modules/
+dist/
+.next/
+.turbo/
+*.pyc
+__pycache__/
+.venv/
+*.egg-info/
+.pytest_cache/
+.ruff_cache/
+.env
+.env.local
+.env.*.local
+.DS_Store
+coverage/
+
+# Runtime artifacts
+logs_llm/
+vectors.db
+vectors.db-shm
+vectors.db-wal

msaas_rag-1.0.0/PKG-INFO

@@ -0,0 +1,20 @@
+Metadata-Version: 2.4
+Name: msaas-rag
+Version: 1.0.0
+Summary: RAG pipeline library — chunking, embeddings, vector search, and retrieval for the Willian SaaS platform
+License: MIT
+Requires-Python: >=3.12
+Requires-Dist: asyncpg>=0.30.0
+Requires-Dist: pydantic>=2.0
+Provides-Extra: all
+Requires-Dist: numpy>=2.0; extra == 'all'
+Requires-Dist: openai>=1.50.0; extra == 'all'
+Provides-Extra: dev
+Requires-Dist: numpy>=2.0; extra == 'dev'
+Requires-Dist: pytest-asyncio>=0.24.0; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Requires-Dist: ruff>=0.8; extra == 'dev'
+Provides-Extra: numpy
+Requires-Dist: numpy>=2.0; extra == 'numpy'
+Provides-Extra: openai
+Requires-Dist: openai>=1.50.0; extra == 'openai'

msaas_rag-1.0.0/pyproject.toml

@@ -0,0 +1,39 @@
+[project]
+name = "msaas-rag"
+version = "1.0.0"
+description = "RAG pipeline library — chunking, embeddings, vector search, and retrieval for the Willian SaaS platform"
+requires-python = ">=3.12"
+license = { text = "MIT" }
+dependencies = [
+    "pydantic>=2.0",
+    "asyncpg>=0.30.0",
+]
+
+[project.optional-dependencies]
+openai = ["openai>=1.50.0"]
+numpy = ["numpy>=2.0"]
+all = ["openai>=1.50.0", "numpy>=2.0"]
+dev = [
+    "pytest>=8.0",
+    "pytest-asyncio>=0.24.0",
+    "numpy>=2.0",
+    "ruff>=0.8",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/rag"]
+
+[tool.ruff]
+target-version = "py312"
+line-length = 100
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "N", "W", "UP", "B", "SIM", "TCH"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+asyncio_mode = "auto"

msaas_rag-1.0.0/src/rag/__init__.py

@@ -0,0 +1,54 @@
+"""Willian RAG — retrieval-augmented generation pipeline library."""
+
+from rag.chunking import MarkdownChunker, TextChunker
+from rag.config import get_config, get_pipeline, init_rag
+from rag.embeddings import EmbeddingProvider, LocalEmbeddings, OpenAIEmbeddings
+from rag.models import (
+    Chunk,
+    ChunkingConfig,
+    Document,
+    EmbeddingConfig,
+    RAGConfig,
+    SearchResult,
+    VectorStoreConfig,
+)
+from rag.pipeline import RAGPipeline
+from rag.reranker import (
+    CrossEncoderReranker,
+    LLMReranker,
+    Reranker,
+    reciprocal_rank_fusion,
+)
+from rag.vector_store import InMemoryVectorStore, PgVectorStore, VectorStore
+
+__all__ = [
+    # Pipeline
+    "RAGPipeline",
+    "init_rag",
+    "get_config",
+    "get_pipeline",
+    # Models
+    "Document",
+    "Chunk",
+    "SearchResult",
+    "RAGConfig",
+    "ChunkingConfig",
+    "EmbeddingConfig",
+    "VectorStoreConfig",
+    # Chunking
+    "TextChunker",
+    "MarkdownChunker",
+    # Embeddings
+    "EmbeddingProvider",
+    "OpenAIEmbeddings",
+    "LocalEmbeddings",
+    # Vector stores
+    "VectorStore",
+    "PgVectorStore",
+    "InMemoryVectorStore",
+    # Reranking
+    "Reranker",
+    "CrossEncoderReranker",
+    "LLMReranker",
+    "reciprocal_rank_fusion",
+]

msaas_rag-1.0.0/src/rag/chunking.py

@@ -0,0 +1,202 @@
+"""Text chunking strategies for document processing."""
+
+from __future__ import annotations
+
+import re
+import uuid
+from typing import Any
+
+from rag.models import Chunk
+
+
+class TextChunker:
+    """Split plain text into chunks using configurable strategies.
+
+    Supported strategies:
+    - ``fixed_size``: split at exact character boundaries
+    - ``recursive``: split by paragraph > sentence > character, preferring natural breaks
+    - ``semantic``: sentence-based splitting that groups sentences up to chunk_size
+    """
+
+    SEPARATORS = ["\n\n", "\n", ". ", " "]
+
+    def __init__(
+        self,
+        strategy: str = "recursive",
+        chunk_size: int = 512,
+        chunk_overlap: int = 64,
+    ) -> None:
+        if strategy not in ("fixed_size", "recursive", "semantic"):
+            msg = f"Unknown strategy: {strategy!r}. Use fixed_size, recursive, or semantic."
+            raise ValueError(msg)
+        self.strategy = strategy
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+
+    def chunk(
+        self,
+        text: str,
+        document_id: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> list[Chunk]:
+        """Split *text* into Chunk objects."""
+        meta = metadata or {}
+        if self.strategy == "fixed_size":
+            return self._fixed_size(text, document_id, meta)
+        if self.strategy == "semantic":
+            return self._semantic(text, document_id, meta)
+        return self._recursive(text, document_id, meta)
+
+    # -- strategies ----------------------------------------------------------
+
+    def _fixed_size(self, text: str, document_id: str, meta: dict[str, Any]) -> list[Chunk]:
+        chunks: list[Chunk] = []
+        step = max(1, self.chunk_size - self.chunk_overlap)
+        for start in range(0, len(text), step):
+            end = min(start + self.chunk_size, len(text))
+            chunks.append(self._make(text[start:end], document_id, start, end, meta))
+            if end == len(text):
+                break
+        return chunks
+
+    def _recursive(self, text: str, document_id: str, meta: dict[str, Any]) -> list[Chunk]:
+        return self._recursive_split(text, document_id, meta, self.SEPARATORS)
+
+    def _recursive_split(
+        self,
+        text: str,
+        document_id: str,
+        meta: dict[str, Any],
+        separators: list[str],
+    ) -> list[Chunk]:
+        if len(text) <= self.chunk_size:
+            return [self._make(text, document_id, 0, len(text), meta)]
+
+        sep = separators[0] if separators else ""
+        parts = text.split(sep) if sep else list(text)
+        remaining_seps = separators[1:] if separators else []
+
+        chunks: list[Chunk] = []
+        current = ""
+        offset = 0
+
+        for i, part in enumerate(parts):
+            candidate = current + (sep if current else "") + part
+            if len(candidate) > self.chunk_size and current:
+                start = offset
+                end = offset + len(current)
+                chunks.append(self._make(current, document_id, start, end, meta))
+                # Compute overlap start
+                overlap_start = max(0, len(current) - self.chunk_overlap)
+                offset = offset + len(current) - (len(current) - overlap_start) + len(sep)
+                current = (
+                    current[overlap_start:] + sep + part if overlap_start < len(current) else part
+                )
+            elif len(candidate) > self.chunk_size and remaining_seps:
+                sub = self._recursive_split(part, document_id, meta, remaining_seps)
+                chunks.extend(sub)
+                current = ""
+                offset += len(part) + len(sep)
+            else:
+                current = candidate
+                if not current.startswith(part) and i == 0:
+                    pass  # offset stays
+        if current.strip():
+            chunks.append(self._make(current, document_id, offset, offset + len(current), meta))
+        return chunks
+
+    def _semantic(self, text: str, document_id: str, meta: dict[str, Any]) -> list[Chunk]:
+        """Sentence-based chunking: group sentences up to chunk_size."""
+        sentences = re.split(r"(?<=[.!?])\s+", text)
+        chunks: list[Chunk] = []
+        current = ""
+        offset = 0
+
+        for sentence in sentences:
+            candidate = (current + " " + sentence).strip() if current else sentence
+            if len(candidate) > self.chunk_size and current:
+                end = offset + len(current)
+                chunks.append(self._make(current, document_id, offset, end, meta))
+                offset = end + 1
+                current = sentence
+            else:
+                current = candidate
+        if current.strip():
+            chunks.append(self._make(current, document_id, offset, offset + len(current), meta))
+        return chunks
+
+    @staticmethod
+    def _make(text: str, document_id: str, start: int, end: int, meta: dict[str, Any]) -> Chunk:
+        return Chunk(
+            id=uuid.uuid4().hex,
+            text=text,
+            document_id=document_id,
+            start_idx=start,
+            end_idx=end,
+            metadata=meta,
+        )
+
+
+class MarkdownChunker:
+    """Split Markdown documents by headers, preserving structure.
+
+    Each header section becomes a chunk. Sections exceeding ``chunk_size`` are
+    further split by the given ``fallback`` TextChunker strategy.
+    """
+
+    HEADER_RE = re.compile(r"^(#{1,6})\s+(.*)", re.MULTILINE)
+
+    def __init__(self, chunk_size: int = 1024, chunk_overlap: int = 64) -> None:
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+        self._fallback = TextChunker(
+            strategy="recursive", chunk_size=chunk_size, chunk_overlap=chunk_overlap
+        )
+
+    def chunk(
+        self,
+        text: str,
+        document_id: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> list[Chunk]:
+        meta = metadata or {}
+        sections = self._split_by_headers(text)
+        chunks: list[Chunk] = []
+
+        for header, body, start in sections:
+            section_meta = {**meta}
+            if header:
+                section_meta["header"] = header
+            full_text = f"{header}\n{body}".strip() if header else body.strip()
+            if len(full_text) <= self.chunk_size:
+                chunks.append(
+                    TextChunker._make(
+                        full_text, document_id, start, start + len(full_text), section_meta
+                    )
+                )
+            else:
+                sub = self._fallback.chunk(full_text, document_id, section_meta)
+                chunks.extend(sub)
+        return chunks
+
+    def _split_by_headers(self, text: str) -> list[tuple[str, str, int]]:
+        """Return list of (header_line, body_text, start_index)."""
+        matches = list(self.HEADER_RE.finditer(text))
+        if not matches:
+            return [("", text, 0)]
+
+        sections: list[tuple[str, str, int]] = []
+        # Text before first header
+        if matches[0].start() > 0:
+            preamble = text[: matches[0].start()]
+            if preamble.strip():
+                sections.append(("", preamble.strip(), 0))
+
+        for i, match in enumerate(matches):
+            header_line = match.group(0)
+            body_start = match.end()
+            body_end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
+            body = text[body_start:body_end].strip()
+            sections.append((header_line, body, match.start()))
+
+        return sections

msaas_rag-1.0.0/src/rag/config.py

@@ -0,0 +1,46 @@
+"""Global configuration and pipeline singleton management."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from rag.models import RAGConfig
+
+if TYPE_CHECKING:
+    from rag.pipeline import RAGPipeline
+
+_pipeline: RAGPipeline | None = None
+_config: RAGConfig | None = None
+
+
+def init_rag(config: RAGConfig | None = None) -> RAGConfig:
+    """Initialize global RAG configuration.
+
+    If no config is provided, a default configuration is created.
+    Returns the active config for further customization.
+    """
+    global _config, _pipeline
+    _config = config or RAGConfig()
+    _pipeline = None  # Reset pipeline so it picks up new config
+    return _config
+
+
+def get_config() -> RAGConfig:
+    """Return the current global RAG config, initializing defaults if needed."""
+    global _config
+    if _config is None:
+        _config = RAGConfig()
+    return _config
+
+
+def get_pipeline() -> RAGPipeline:
+    """Return (or create) the global RAGPipeline singleton.
+
+    Uses the current global config. Call ``init_rag()`` first to customize.
+    """
+    global _pipeline
+    if _pipeline is None:
+        from rag.pipeline import RAGPipeline
+
+        _pipeline = RAGPipeline(config=get_config())
+    return _pipeline

msaas_rag-1.0.0/src/rag/embeddings.py

@@ -0,0 +1,136 @@
+"""Embedding providers for converting text to vector representations."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class EmbeddingProvider(ABC):
+    """Abstract base for embedding providers."""
+
+    @abstractmethod
+    async def embed(self, texts: list[str]) -> list[list[float]]:
+        """Return embedding vectors for each text in *texts*."""
+
+    @abstractmethod
+    def dimensions(self) -> int:
+        """Return the dimensionality of the embedding vectors."""
+
+
+class OpenAIEmbeddings(EmbeddingProvider):
+    """OpenAI text-embedding-3-small/large provider with caching and batching.
+
+    Requires the ``openai`` extra: ``pip install willian-rag[openai]``
+    """
+
+    MODEL_DIMENSIONS = {
+        "text-embedding-3-small": 1536,
+        "text-embedding-3-large": 3072,
+        "text-embedding-ada-002": 1536,
+    }
+
+    def __init__(
+        self,
+        model: str = "text-embedding-3-small",
+        api_key: str | None = None,
+        batch_size: int = 100,
+        max_retries: int = 3,
+        retry_delay: float = 1.0,
+    ) -> None:
+        try:
+            import openai  # noqa: F811
+        except ImportError as exc:
+            msg = "Install the openai extra: pip install willian-rag[openai]"
+            raise ImportError(msg) from exc
+
+        self.model = model
+        self.batch_size = batch_size
+        self.max_retries = max_retries
+        self.retry_delay = retry_delay
+        self._client = openai.AsyncOpenAI(api_key=api_key)
+        self._cache: dict[str, list[float]] = {}
+
+    def dimensions(self) -> int:
+        return self.MODEL_DIMENSIONS.get(self.model, 1536)
+
+    async def embed(self, texts: list[str]) -> list[list[float]]:
+        """Embed texts with batching, rate limiting, and caching."""
+        results: dict[int, list[float]] = {}
+        uncached: list[tuple[int, str]] = []
+
+        for i, text in enumerate(texts):
+            key = self._cache_key(text)
+            if key in self._cache:
+                results[i] = self._cache[key]
+            else:
+                uncached.append((i, text))
+
+        # Process uncached in batches
+        for batch_start in range(0, len(uncached), self.batch_size):
+            batch = uncached[batch_start : batch_start + self.batch_size]
+            batch_texts = [t for _, t in batch]
+            embeddings = await self._embed_with_retry(batch_texts)
+            for (idx, text), emb in zip(batch, embeddings):
+                self._cache[self._cache_key(text)] = emb
+                results[idx] = emb
+
+        return [results[i] for i in range(len(texts))]
+
+    async def _embed_with_retry(self, texts: list[str]) -> list[list[float]]:
+        last_error: Exception | None = None
+        for attempt in range(self.max_retries):
+            try:
+                response = await self._client.embeddings.create(model=self.model, input=texts)
+                return [item.embedding for item in response.data]
+            except Exception as exc:
+                last_error = exc
+                if attempt < self.max_retries - 1:
+                    await asyncio.sleep(self.retry_delay * (attempt + 1))
+        msg = f"Embedding failed after {self.max_retries} retries"
+        raise RuntimeError(msg) from last_error
+
+    @staticmethod
+    def _cache_key(text: str) -> str:
+        return hashlib.sha256(text.encode()).hexdigest()
+
+
+class LocalEmbeddings(EmbeddingProvider):
+    """Placeholder for sentence-transformers local embedding.
+
+    Returns deterministic pseudo-embeddings based on text hash for testing.
+    Replace the ``embed`` method with a real model for production use.
+    """
+
+    def __init__(self, dims: int = 384) -> None:
+        self._dims = dims
+
+    def dimensions(self) -> int:
+        return self._dims
+
+    async def embed(self, texts: list[str]) -> list[list[float]]:
+        return [self._pseudo_embedding(t) for t in texts]
+
+    def _pseudo_embedding(self, text: str) -> list[float]:
+        """Generate a deterministic pseudo-embedding from text hash."""
+        h = hashlib.sha256(text.encode()).digest()
+        raw = [b / 255.0 for b in h]
+        # Extend or truncate to match dimensions
+        while len(raw) < self._dims:
+            raw = raw + raw
+        raw = raw[: self._dims]
+        # Normalize to unit vector
+        norm = sum(x * x for x in raw) ** 0.5
+        return [x / norm if norm > 0 else 0.0 for x in raw]
+
+
+def _build_provider(config: Any) -> EmbeddingProvider:
+    """Factory: create an EmbeddingProvider from an EmbeddingConfig."""
+    if config.provider == "openai":
+        return OpenAIEmbeddings(model=config.model, batch_size=config.batch_size)
+    if config.provider == "local":
+        return LocalEmbeddings(dims=config.dimensions)
+    msg = f"Unknown embedding provider: {config.provider!r}"
+    raise ValueError(msg)

msaas_rag-1.0.0/src/rag/models.py

@@ -0,0 +1,73 @@
+"""Core domain models for the RAG pipeline."""
+
+from __future__ import annotations
+
+import uuid
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+def _default_id() -> str:
+    return uuid.uuid4().hex
+
+
+class Document(BaseModel):
+    """A source document to be ingested into the RAG pipeline."""
+
+    id: str = Field(default_factory=_default_id)
+    text: str
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class Chunk(BaseModel):
+    """A text chunk derived from a document."""
+
+    id: str = Field(default_factory=_default_id)
+    text: str
+    document_id: str
+    start_idx: int
+    end_idx: int
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    embedding: list[float] | None = None
+
+
+class SearchResult(BaseModel):
+    """A single result returned from vector search."""
+
+    chunk: Chunk
+    score: float
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class ChunkingConfig(BaseModel):
+    """Configuration for text chunking."""
+
+    strategy: str = "recursive"
+    chunk_size: int = 512
+    chunk_overlap: int = 64
+
+
+class EmbeddingConfig(BaseModel):
+    """Configuration for the embedding provider."""
+
+    provider: str = "openai"
+    model: str = "text-embedding-3-small"
+    dimensions: int = 1536
+    batch_size: int = 100
+
+
+class VectorStoreConfig(BaseModel):
+    """Configuration for the vector store backend."""
+
+    backend: str = "pgvector"
+    table_name: str = "rag_chunks"
+    dsn: str = ""
+
+
+class RAGConfig(BaseModel):
+    """Top-level configuration for the entire RAG pipeline."""
+
+    chunking: ChunkingConfig = Field(default_factory=ChunkingConfig)
+    embedding: EmbeddingConfig = Field(default_factory=EmbeddingConfig)
+    vector_store: VectorStoreConfig = Field(default_factory=VectorStoreConfig)