msaas-rag 1.0.0__py3-none-any.whl

msaas_rag-1.0.0.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,11 @@
+rag/__init__.py,sha256=ca84grhMOHLh_cfFsJtbAU-jp2DkqEPjAjJwO65tdGo,1235
+rag/chunking.py,sha256=XpuIU44VMyHlp_YuSKbIZpEvQct1VsryCs-qHDuufgA,7538
+rag/config.py,sha256=ITny9Edku1M82KdzfeZ-SmZA9wotR5P7O8cUO6JDKgQ,1240
+rag/embeddings.py,sha256=UwK63qcJEBvkXPXkYCbljx_pSWRG-DNZRmAkrln840o,4870
+rag/models.py,sha256=NdtpR6X9I5lSnMISWOQ9c8iY4ranBg6qsME4PxvPUg8,1777
+rag/pipeline.py,sha256=DAmgmzwa_BnnJ27hgf9l0-e0rPoAfVoNN3br-1NATAc,3631
+rag/reranker.py,sha256=2a7veHSoIA5SgOGxEMzVLAdzKizY0O9jbkz8EBN7sl8,3605
+rag/vector_store.py,sha256=8IuMhv9fkexbkd2BSpyKoQQWqjzzS37gw-nJDJZJM8w,10659
+msaas_rag-1.0.0.dist-info/METADATA,sha256=79VoK8PKNFWy-xg35JoHjVbJEiSFErWcYhdjrH_YsH8,711
+msaas_rag-1.0.0.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+msaas_rag-1.0.0.dist-info/RECORD,,

msaas_rag-1.0.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.30.1
+Root-Is-Purelib: true
+Tag: py3-none-any

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",
+]

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

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

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)

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)

rag/pipeline.py

@@ -0,0 +1,110 @@
+"""RAG pipeline orchestrating chunking, embedding, storage, and retrieval."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from rag.chunking import MarkdownChunker, TextChunker
+from rag.embeddings import LocalEmbeddings, _build_provider
+from rag.models import Chunk, Document, RAGConfig, SearchResult
+from rag.reranker import Reranker
+from rag.vector_store import InMemoryVectorStore, VectorStore
+
+if TYPE_CHECKING:
+    from rag.embeddings import EmbeddingProvider
+
+
+class RAGPipeline:
+    """Orchestrates the full RAG flow: chunk -> embed -> store -> retrieve.
+
+    Can be configured via RAGConfig or by injecting components directly.
+    """
+
+    def __init__(
+        self,
+        config: RAGConfig | None = None,
+        chunker: TextChunker | MarkdownChunker | None = None,
+        embedding_provider: EmbeddingProvider | None = None,
+        vector_store: VectorStore | None = None,
+    ) -> None:
+        cfg = config or RAGConfig()
+
+        self.chunker = chunker or TextChunker(
+            strategy=cfg.chunking.strategy,
+            chunk_size=cfg.chunking.chunk_size,
+            chunk_overlap=cfg.chunking.chunk_overlap,
+        )
+        self.embedding_provider: EmbeddingProvider = embedding_provider or self._safe_provider(cfg)
+        self.vector_store = vector_store or InMemoryVectorStore()
+
+    @staticmethod
+    def _safe_provider(cfg: RAGConfig) -> EmbeddingProvider:
+        """Try to build the configured provider, fall back to local."""
+        try:
+            return _build_provider(cfg.embedding)
+        except Exception:
+            return LocalEmbeddings(dims=cfg.embedding.dimensions)
+
+    async def ingest(self, documents: list[Document]) -> int:
+        """Process documents through the full pipeline: chunk -> embed -> store.
+
+        Returns the total number of chunks stored.
+        """
+        all_chunks: list[Chunk] = []
+
+        for doc in documents:
+            chunks = self.chunker.chunk(
+                text=doc.text,
+                document_id=doc.id,
+                metadata=doc.metadata,
+            )
+            all_chunks.extend(chunks)
+
+        if not all_chunks:
+            return 0
+
+        # Embed all chunks in one batch
+        texts = [c.text for c in all_chunks]
+        embeddings = await self.embedding_provider.embed(texts)
+
+        for chunk, embedding in zip(all_chunks, embeddings):
+            chunk.embedding = embedding
+
+        # Store
+        count = await self.vector_store.upsert(all_chunks)
+        return count
+
+    async def retrieve(
+        self,
+        query: str,
+        top_k: int = 10,
+        filters: dict | None = None,
+    ) -> list[SearchResult]:
+        """Embed the query and search the vector store."""
+        query_embeddings = await self.embedding_provider.embed([query])
+        query_embedding = query_embeddings[0]
+        return await self.vector_store.search(
+            query_embedding=query_embedding,
+            top_k=top_k,
+            filters=filters,
+        )
+
+    async def retrieve_with_rerank(
+        self,
+        query: str,
+        top_k: int = 10,
+        reranker: Reranker | None = None,
+        initial_k: int | None = None,
+        filters: dict | None = None,
+    ) -> list[SearchResult]:
+        """Retrieve a larger set, then rerank down to top_k.
+
+        If no reranker is provided, behaves like plain retrieve.
+        """
+        fetch_k = initial_k or top_k * 3
+        results = await self.retrieve(query, top_k=fetch_k, filters=filters)
+
+        if reranker is None:
+            return results[:top_k]
+
+        return await reranker.rerank(query, results, top_k=top_k)

rag/reranker.py

@@ -0,0 +1,113 @@
+"""Reranking strategies for improving retrieval quality."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from rag.models import SearchResult
+
+
+class Reranker(ABC):
+    """Abstract base for reranking retrieved results."""
+
+    @abstractmethod
+    async def rerank(
+        self, query: str, results: list[SearchResult], top_k: int = 10
+    ) -> list[SearchResult]:
+        """Rerank results by relevance to *query*. Returns top_k results."""
+
+
+class CrossEncoderReranker(Reranker):
+    """Placeholder for cross-encoder model reranking.
+
+    In production, load a cross-encoder model (e.g., ms-marco-MiniLM) and score
+    each (query, chunk.text) pair. This placeholder passes results through with
+    a simple length-based heuristic for testing.
+    """
+
+    async def rerank(
+        self, query: str, results: list[SearchResult], top_k: int = 10
+    ) -> list[SearchResult]:
+        # Placeholder: score by keyword overlap ratio
+        query_terms = set(query.lower().split())
+        scored: list[tuple[float, SearchResult]] = []
+        for result in results:
+            chunk_terms = set(result.chunk.text.lower().split())
+            overlap = len(query_terms & chunk_terms)
+            total = len(query_terms) if query_terms else 1
+            relevance = overlap / total
+            scored.append((relevance, result))
+
+        scored.sort(key=lambda x: x[0], reverse=True)
+        return [
+            SearchResult(
+                chunk=r.chunk,
+                score=score,
+                metadata=r.metadata,
+            )
+            for score, r in scored[:top_k]
+        ]
+
+
+class LLMReranker(Reranker):
+    """Reranker that uses an LLM to score relevance.
+
+    Accepts a callable ``score_fn`` that takes (query, text) and returns a
+    float relevance score. This allows integration with willian-ai or any
+    LLM client.
+    """
+
+    def __init__(
+        self,
+        score_fn: callable | None = None,
+    ) -> None:
+        self._score_fn = score_fn
+
+    async def rerank(
+        self, query: str, results: list[SearchResult], top_k: int = 10
+    ) -> list[SearchResult]:
+        if self._score_fn is None:
+            # No scoring function — return as-is
+            return results[:top_k]
+
+        scored: list[tuple[float, SearchResult]] = []
+        for result in results:
+            score = await self._score_fn(query, result.chunk.text)
+            scored.append((float(score), result))
+
+        scored.sort(key=lambda x: x[0], reverse=True)
+        return [
+            SearchResult(chunk=r.chunk, score=score, metadata=r.metadata)
+            for score, r in scored[:top_k]
+        ]
+
+
+def reciprocal_rank_fusion(
+    *result_lists: list[SearchResult],
+    k: int = 60,
+    top_n: int = 10,
+) -> list[SearchResult]:
+    """Merge multiple ranked result lists using Reciprocal Rank Fusion (RRF).
+
+    RRF score for document d = sum over lists of 1 / (k + rank_in_list).
+    Default k=60 following the original RRF paper.
+    """
+    scores: dict[str, float] = {}
+    chunk_map: dict[str, SearchResult] = {}
+
+    for results in result_lists:
+        for rank, result in enumerate(results):
+            cid = result.chunk.id
+            scores[cid] = scores.get(cid, 0.0) + 1.0 / (k + rank + 1)
+            if cid not in chunk_map:
+                chunk_map[cid] = result
+
+    ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
+    return [
+        SearchResult(
+            chunk=chunk_map[cid].chunk,
+            score=score,
+            metadata=chunk_map[cid].metadata,
+        )
+        for cid, score in ranked[:top_n]
+    ]

rag/vector_store.py

@@ -0,0 +1,293 @@
+"""Vector store backends for chunk storage and similarity search."""
+
+from __future__ import annotations
+
+import json
+from abc import ABC, abstractmethod
+from typing import Any
+
+from rag.models import Chunk, SearchResult
+
+
+class VectorStore(ABC):
+    """Abstract base for vector store backends."""
+
+    @abstractmethod
+    async def upsert(self, chunks: list[Chunk]) -> int:
+        """Insert or update chunks. Returns the number of affected rows."""
+
+    @abstractmethod
+    async def search(
+        self,
+        query_embedding: list[float],
+        top_k: int = 10,
+        filters: dict[str, Any] | None = None,
+    ) -> list[SearchResult]:
+        """Search by cosine similarity. Returns results ordered by score descending."""
+
+    @abstractmethod
+    async def delete(self, chunk_ids: list[str]) -> int:
+        """Delete chunks by ID. Returns the number of deleted rows."""
+
+
+class PgVectorStore(VectorStore):
+    """PostgreSQL + pgvector backend using asyncpg.
+
+    Supports both pure vector search and hybrid (vector + full-text) search
+    via Reciprocal Rank Fusion.
+    """
+
+    def __init__(self, dsn: str, table_name: str = "rag_chunks", dimensions: int = 1536) -> None:
+        self.dsn = dsn
+        self.table_name = table_name
+        self.dimensions = dimensions
+        self._pool: Any = None
+
+    async def _get_pool(self) -> Any:
+        if self._pool is None:
+            import asyncpg
+
+            self._pool = await asyncpg.create_pool(self.dsn, min_size=2, max_size=10)
+        return self._pool
+
+    async def create_table(self) -> None:
+        """Create the chunks table with vector and tsvector columns."""
+        pool = await self._get_pool()
+        async with pool.acquire() as conn:
+            await conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
+            await conn.execute(f"""
+                CREATE TABLE IF NOT EXISTS {self.table_name} (
+                    id TEXT PRIMARY KEY,
+                    text TEXT NOT NULL,
+                    document_id TEXT NOT NULL,
+                    start_idx INTEGER NOT NULL,
+                    end_idx INTEGER NOT NULL,
+                    metadata JSONB DEFAULT '{{}}',
+                    embedding vector({self.dimensions}),
+                    tsv tsvector GENERATED ALWAYS AS (to_tsvector('english', text)) STORED
+                )
+            """)
+            await conn.execute(
+                f"CREATE INDEX IF NOT EXISTS idx_{self.table_name}_embedding "
+                f"ON {self.table_name} USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100)"
+            )
+            await conn.execute(
+                f"CREATE INDEX IF NOT EXISTS idx_{self.table_name}_tsv "
+                f"ON {self.table_name} USING gin (tsv)"
+            )
+
+    async def upsert(self, chunks: list[Chunk]) -> int:
+        """Bulk upsert chunks with ON CONFLICT."""
+        if not chunks:
+            return 0
+        pool = await self._get_pool()
+        async with pool.acquire() as conn:
+            values = [
+                (
+                    c.id,
+                    c.text,
+                    c.document_id,
+                    c.start_idx,
+                    c.end_idx,
+                    json.dumps(c.metadata),
+                    str(c.embedding) if c.embedding else None,
+                )
+                for c in chunks
+            ]
+            result = await conn.executemany(
+                f"""
+                INSERT INTO {self.table_name} (id, text, document_id, start_idx, end_idx, metadata, embedding)
+                VALUES ($1, $2, $3, $4, $5, $6::jsonb, $7::vector)
+                ON CONFLICT (id) DO UPDATE SET
+                    text = EXCLUDED.text,
+                    document_id = EXCLUDED.document_id,
+                    metadata = EXCLUDED.metadata,
+                    embedding = EXCLUDED.embedding
+                """,
+                values,
+            )
+            return len(chunks) if result is None else len(chunks)
+
+    async def search(
+        self,
+        query_embedding: list[float],
+        top_k: int = 10,
+        filters: dict[str, Any] | None = None,
+    ) -> list[SearchResult]:
+        """Cosine similarity search."""
+        pool = await self._get_pool()
+        filter_clause = self._build_filter_clause(filters)
+        embedding_str = str(query_embedding)
+
+        async with pool.acquire() as conn:
+            rows = await conn.fetch(
+                f"""
+                SELECT id, text, document_id, start_idx, end_idx, metadata,
+                       1 - (embedding <=> $1::vector) AS score
+                FROM {self.table_name}
+                {filter_clause}
+                ORDER BY embedding <=> $1::vector
+                LIMIT $2
+                """,
+                embedding_str,
+                top_k,
+            )
+            return [self._row_to_result(row) for row in rows]
+
+    async def hybrid_search(
+        self,
+        query_text: str,
+        query_embedding: list[float],
+        top_k: int = 10,
+        vector_weight: float = 0.7,
+        filters: dict[str, Any] | None = None,
+    ) -> list[SearchResult]:
+        """Hybrid search combining vector similarity and full-text search via RRF."""
+        pool = await self._get_pool()
+        filter_clause = self._build_filter_clause(filters)
+        embedding_str = str(query_embedding)
+
+        async with pool.acquire() as conn:
+            rows = await conn.fetch(
+                f"""
+                WITH vector_results AS (
+                    SELECT id, 1 - (embedding <=> $1::vector) AS vec_score,
+                           ROW_NUMBER() OVER (ORDER BY embedding <=> $1::vector) AS vec_rank
+                    FROM {self.table_name} {filter_clause}
+                    LIMIT $3
+                ),
+                text_results AS (
+                    SELECT id, ts_rank_cd(tsv, plainto_tsquery('english', $2)) AS text_score,
+                           ROW_NUMBER() OVER (ORDER BY ts_rank_cd(tsv, plainto_tsquery('english', $2)) DESC) AS text_rank
+                    FROM {self.table_name}
+                    WHERE tsv @@ plainto_tsquery('english', $2) {filter_clause.replace("WHERE", "AND") if filter_clause else ""}
+                    LIMIT $3
+                ),
+                combined AS (
+                    SELECT COALESCE(v.id, t.id) AS id,
+                           $4 * COALESCE(1.0 / (60 + v.vec_rank), 0) +
+                           (1 - $4) * COALESCE(1.0 / (60 + t.text_rank), 0) AS rrf_score
+                    FROM vector_results v
+                    FULL OUTER JOIN text_results t ON v.id = t.id
+                    ORDER BY rrf_score DESC
+                    LIMIT $3
+                )
+                SELECT c.id, c.text, c.document_id, c.start_idx, c.end_idx,
+                       c.metadata, combined.rrf_score AS score
+                FROM combined
+                JOIN {self.table_name} c ON c.id = combined.id
+                ORDER BY combined.rrf_score DESC
+                """,
+                embedding_str,
+                query_text,
+                top_k,
+                vector_weight,
+            )
+            return [self._row_to_result(row) for row in rows]
+
+    async def delete(self, chunk_ids: list[str]) -> int:
+        if not chunk_ids:
+            return 0
+        pool = await self._get_pool()
+        async with pool.acquire() as conn:
+            result = await conn.execute(
+                f"DELETE FROM {self.table_name} WHERE id = ANY($1)", chunk_ids
+            )
+            # asyncpg returns "DELETE N"
+            return int(result.split()[-1]) if result else 0
+
+    async def close(self) -> None:
+        """Close the connection pool."""
+        if self._pool:
+            await self._pool.close()
+            self._pool = None
+
+    @staticmethod
+    def _build_filter_clause(filters: dict[str, Any] | None) -> str:
+        if not filters:
+            return ""
+        conditions = []
+        for key, value in filters.items():
+            escaped = json.dumps(value)
+            conditions.append(f"metadata->>'{key}' = {escaped}")
+        return "WHERE " + " AND ".join(conditions)
+
+    @staticmethod
+    def _row_to_result(row: Any) -> SearchResult:
+        metadata = (
+            row["metadata"] if isinstance(row["metadata"], dict) else json.loads(row["metadata"])
+        )
+        chunk = Chunk(
+            id=row["id"],
+            text=row["text"],
+            document_id=row["document_id"],
+            start_idx=row["start_idx"],
+            end_idx=row["end_idx"],
+            metadata=metadata,
+        )
+        return SearchResult(chunk=chunk, score=float(row["score"]), metadata=metadata)
+
+
+class InMemoryVectorStore(VectorStore):
+    """In-memory vector store using numpy for cosine similarity. For testing."""
+
+    def __init__(self) -> None:
+        self._chunks: dict[str, Chunk] = {}
+
+    async def upsert(self, chunks: list[Chunk]) -> int:
+        count = 0
+        for chunk in chunks:
+            self._chunks[chunk.id] = chunk
+            count += 1
+        return count
+
+    async def search(
+        self,
+        query_embedding: list[float],
+        top_k: int = 10,
+        filters: dict[str, Any] | None = None,
+    ) -> list[SearchResult]:
+        import numpy as np
+
+        candidates = list(self._chunks.values())
+        if filters:
+            candidates = [
+                c for c in candidates if all(c.metadata.get(k) == v for k, v in filters.items())
+            ]
+
+        candidates_with_emb = [c for c in candidates if c.embedding is not None]
+        if not candidates_with_emb:
+            return []
+
+        query_vec = np.array(query_embedding, dtype=np.float64)
+        query_norm = np.linalg.norm(query_vec)
+        if query_norm == 0:
+            return []
+        query_vec = query_vec / query_norm
+
+        scored: list[tuple[float, Chunk]] = []
+        for chunk in candidates_with_emb:
+            chunk_vec = np.array(chunk.embedding, dtype=np.float64)
+            chunk_norm = np.linalg.norm(chunk_vec)
+            if chunk_norm == 0:
+                continue
+            score = float(np.dot(query_vec, chunk_vec / chunk_norm))
+            scored.append((score, chunk))
+
+        scored.sort(key=lambda x: x[0], reverse=True)
+        return [
+            SearchResult(chunk=chunk, score=score, metadata=chunk.metadata)
+            for score, chunk in scored[:top_k]
+        ]
+
+    async def delete(self, chunk_ids: list[str]) -> int:
+        count = 0
+        for cid in chunk_ids:
+            if cid in self._chunks:
+                del self._chunks[cid]
+                count += 1
+        return count
+
+    @property
+    def size(self) -> int:
+        return len(self._chunks)