mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-16 11:23:10 +02:00
setup RAG (#6146)
This commit is contained in:
Anthony Stirling
@@ -0,0 +1,81 @@
|
||||
# RAG Integration Guide
|
||||
|
||||
## Adding RAG to an Agent
|
||||
|
||||
```python
|
||||
from pydantic_ai import Agent
|
||||
|
||||
from stirling.services import AppRuntime
|
||||
|
||||
class MyAgent:
|
||||
def __init__(self, runtime: AppRuntime) -> None:
|
||||
rag = runtime.rag_capability
|
||||
self.agent = Agent(
|
||||
model=runtime.smart_model,
|
||||
system_prompt="Your prompt here...",
|
||||
instructions=rag.instructions,
|
||||
toolsets=[rag.toolset],
|
||||
)
|
||||
```
|
||||
|
||||
That's it. The agent gets a `search_knowledge` tool it can call autonomously.
|
||||
|
||||
## Scoping to Specific Collections
|
||||
|
||||
Collections are named buckets of indexed documents — think folders. By default an agent searches everything in the store. Pass `collections=` to restrict it to only the docs indexed under those names.
|
||||
|
||||
```python
|
||||
from stirling.rag import RagCapability
|
||||
|
||||
# Only searches docs indexed under "company-docs" — ignores everything else
|
||||
scoped = RagCapability(runtime.rag_service, collections=["company-docs"], top_k=3)
|
||||
|
||||
# Searches multiple collections
|
||||
multi = RagCapability(runtime.rag_service, collections=["company-docs", "product-specs"])
|
||||
|
||||
# No collections arg = searches all collections in the store
|
||||
everything = RagCapability(runtime.rag_service)
|
||||
```
|
||||
|
||||
## Config (.env)
|
||||
|
||||
```
|
||||
STIRLING_RAG_BACKEND=sqlite # or "pgvector"
|
||||
STIRLING_RAG_EMBEDDING_MODEL=voyageai:voyage-4
|
||||
STIRLING_RAG_STORE_PATH=data/rag.db # used when backend=sqlite
|
||||
STIRLING_RAG_PGVECTOR_DSN= # used when backend=pgvector
|
||||
STIRLING_RAG_CHUNK_SIZE=512
|
||||
STIRLING_RAG_CHUNK_OVERLAP=64
|
||||
STIRLING_RAG_TOP_K=5
|
||||
VOYAGE_API_KEY=your-key
|
||||
```
|
||||
|
||||
## Backends
|
||||
|
||||
**`sqlite`** — Embedded sqlite-vec. Single `.db` file, zero ops. Ideal for dev and self-hosted deployments.
|
||||
|
||||
**`pgvector`** — External PostgreSQL with the `vector` extension. Point `STIRLING_RAG_PGVECTOR_DSN` at your Postgres instance.
|
||||
|
||||
Both backends implement the same `VectorStore` interface, so agents and the RAG service work identically regardless of which you pick.
|
||||
|
||||
For a self-hosted embedding server (e.g. Ollama, TEI, vLLM) set the model string accordingly and point at the server via its native env var:
|
||||
|
||||
```
|
||||
# Ollama running on another machine
|
||||
STIRLING_RAG_EMBEDDING_MODEL=ollama:nomic-embed-text
|
||||
OLLAMA_HOST=http://192.168.1.50:11434
|
||||
|
||||
# Any OpenAI-compatible embedding server
|
||||
STIRLING_RAG_EMBEDDING_MODEL=openai:my-model
|
||||
OPENAI_BASE_URL=http://192.168.1.50:8080/v1
|
||||
```
|
||||
|
||||
## API Endpoints
|
||||
|
||||
| Method | Endpoint | Purpose |
|
||||
|--------|----------|---------|
|
||||
| GET | `/api/v1/rag/status` | Report embedding model and existing collections |
|
||||
| POST | `/api/v1/rag/index` | Index text into a collection |
|
||||
| POST | `/api/v1/rag/search` | Search a collection |
|
||||
| GET | `/api/v1/rag/collections` | List collections |
|
||||
| DELETE | `/api/v1/rag/collections/{name}` | Delete a collection |
|
||||
@@ -0,0 +1,19 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from stirling.rag.capability import RagCapability
|
||||
from stirling.rag.embedder import EmbeddingService
|
||||
from stirling.rag.pgvector_store import PgVectorStore
|
||||
from stirling.rag.service import RagService
|
||||
from stirling.rag.sqlite_vec_store import SqliteVecStore
|
||||
from stirling.rag.store import Document, SearchResult, VectorStore
|
||||
|
||||
__all__ = [
|
||||
"Document",
|
||||
"EmbeddingService",
|
||||
"PgVectorStore",
|
||||
"RagCapability",
|
||||
"RagService",
|
||||
"SearchResult",
|
||||
"SqliteVecStore",
|
||||
"VectorStore",
|
||||
]
|
||||
@@ -0,0 +1,108 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
from pydantic_ai import FunctionToolset
|
||||
from pydantic_ai.toolsets import AbstractToolset
|
||||
|
||||
from stirling.rag.service import RagService
|
||||
|
||||
|
||||
class RagCapability:
|
||||
"""Bundles RAG instructions and the ``search_knowledge`` toolset for agent injection.
|
||||
|
||||
Agents consume this as::
|
||||
|
||||
rag = runtime.rag_capability
|
||||
Agent(
|
||||
...,
|
||||
instructions=rag.instructions,
|
||||
toolsets=[rag.toolset],
|
||||
)
|
||||
|
||||
When no collections are pinned, the instructions are generated dynamically at
|
||||
run time so the agent sees the current list of collections in the store.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rag_service: RagService,
|
||||
collections: list[str] | None = None,
|
||||
top_k: int = 5,
|
||||
) -> None:
|
||||
self._rag_service = rag_service
|
||||
self._collections = collections
|
||||
self._top_k = top_k
|
||||
toolset: FunctionToolset[None] = FunctionToolset()
|
||||
toolset.add_function(self._search_knowledge, name="search_knowledge")
|
||||
self._toolset = toolset
|
||||
|
||||
@property
|
||||
def instructions(self) -> str | Callable[[], Awaitable[str]]:
|
||||
if self._collections:
|
||||
return self._static_instructions_text(self._collections)
|
||||
return self._dynamic_instructions
|
||||
|
||||
@property
|
||||
def toolset(self) -> AbstractToolset[None]:
|
||||
return self._toolset
|
||||
|
||||
@staticmethod
|
||||
def _static_instructions_text(collections: list[str]) -> str:
|
||||
collection_desc = f"collections: {', '.join(collections)}"
|
||||
return (
|
||||
"You have access to a knowledge base search tool called 'search_knowledge'. "
|
||||
f"It searches {collection_desc} for relevant information. "
|
||||
"Use it when the provided context is insufficient to answer the user's question, "
|
||||
"or when you think additional background information would improve your answer. "
|
||||
"You do not have to use it if the answer is already clear from the provided text."
|
||||
)
|
||||
|
||||
async def _dynamic_instructions(self) -> str:
|
||||
collections = await self._rag_service.list_collections()
|
||||
if collections:
|
||||
names = ", ".join(collections)
|
||||
collection_desc = f"the following knowledge base collections: {names}"
|
||||
else:
|
||||
collection_desc = "the knowledge base (currently empty — no collections indexed yet)"
|
||||
return (
|
||||
"You have access to a knowledge base search tool called 'search_knowledge'. "
|
||||
f"It searches {collection_desc} for relevant information. "
|
||||
"Use it when the provided context is insufficient to answer the user's question, "
|
||||
"or when you think additional background information would improve your answer. "
|
||||
"You do not have to use it if the answer is already clear from the provided text."
|
||||
)
|
||||
|
||||
async def _search_knowledge(self, query: str, max_results: int | None = None) -> str:
|
||||
"""Search the knowledge base for information relevant to the query.
|
||||
|
||||
Args:
|
||||
query: The search query describing what information you need.
|
||||
max_results: Maximum number of results to return.
|
||||
|
||||
Returns:
|
||||
Formatted text with the most relevant knowledge base excerpts.
|
||||
"""
|
||||
k = max_results if max_results is not None else self._top_k
|
||||
if self._collections:
|
||||
all_results = []
|
||||
for col in self._collections:
|
||||
col_results = await self._rag_service.search(query, collection=col, top_k=k)
|
||||
all_results.extend(col_results)
|
||||
all_results.sort(key=lambda r: r.score, reverse=True)
|
||||
results = all_results[:k]
|
||||
else:
|
||||
results = await self._rag_service.search(query, top_k=k)
|
||||
|
||||
if not results:
|
||||
return "No relevant results found in the knowledge base."
|
||||
|
||||
sections = []
|
||||
for i, result in enumerate(results, 1):
|
||||
source = result.document.metadata.get("source", "unknown")
|
||||
chunk_idx = result.document.metadata.get("chunk_index", "?")
|
||||
score = f"{result.score:.3f}"
|
||||
sections.append(
|
||||
f"[Result {i} | source: {source}, chunk: {chunk_idx}, relevance: {score}]\n{result.document.text}"
|
||||
)
|
||||
return "\n\n---\n\n".join(sections)
|
||||
@@ -0,0 +1,118 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
# TODO: replace with pydantic-ai's built-in chunking once
|
||||
# https://github.com/pydantic/pydantic-ai/issues/3962 lands.
|
||||
|
||||
|
||||
def chunk_text(text: str, chunk_size: int = 512, overlap: int = 64) -> list[str]:
|
||||
"""Split text into chunks of approximately chunk_size characters with overlap.
|
||||
|
||||
Splits on paragraph then sentence boundaries to avoid cutting mid-thought.
|
||||
Returns an empty list for empty/whitespace-only input.
|
||||
"""
|
||||
text = text.strip()
|
||||
if not text:
|
||||
return []
|
||||
|
||||
paragraphs = _split_paragraphs(text)
|
||||
chunks: list[str] = []
|
||||
current: list[str] = []
|
||||
current_len = 0
|
||||
|
||||
for para in paragraphs:
|
||||
para_len = len(para)
|
||||
|
||||
if current_len + para_len <= chunk_size:
|
||||
current.append(para)
|
||||
current_len += para_len
|
||||
continue
|
||||
|
||||
# If the current buffer has content, flush it
|
||||
if current:
|
||||
chunks.append("\n\n".join(current))
|
||||
|
||||
# If this paragraph alone exceeds chunk_size, split it by sentences
|
||||
if para_len > chunk_size:
|
||||
sentence_chunks = _split_long_paragraph(para, chunk_size, overlap)
|
||||
chunks.extend(sentence_chunks)
|
||||
current = []
|
||||
current_len = 0
|
||||
else:
|
||||
# Start new chunk with overlap from previous chunk
|
||||
overlap_text = _get_overlap(chunks, overlap) if chunks else ""
|
||||
if overlap_text:
|
||||
current = [overlap_text, para]
|
||||
current_len = len(overlap_text) + para_len
|
||||
else:
|
||||
current = [para]
|
||||
current_len = para_len
|
||||
|
||||
if current:
|
||||
chunks.append("\n\n".join(current))
|
||||
|
||||
return [c.strip() for c in chunks if c.strip()]
|
||||
|
||||
|
||||
def _split_paragraphs(text: str) -> list[str]:
|
||||
"""Split text into paragraphs on double newlines."""
|
||||
return [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
|
||||
|
||||
|
||||
def _split_sentences(text: str) -> list[str]:
|
||||
"""Split text into sentences, keeping the delimiter attached."""
|
||||
parts = re.split(r"(?<=[.!?])\s+", text)
|
||||
return [s.strip() for s in parts if s.strip()]
|
||||
|
||||
|
||||
def _split_long_paragraph(paragraph: str, chunk_size: int, overlap: int) -> list[str]:
|
||||
"""Split a single long paragraph into sentence-boundary chunks."""
|
||||
sentences = _split_sentences(paragraph)
|
||||
chunks: list[str] = []
|
||||
current: list[str] = []
|
||||
current_len = 0
|
||||
|
||||
for sentence in sentences:
|
||||
sent_len = len(sentence)
|
||||
|
||||
if current_len + sent_len <= chunk_size:
|
||||
current.append(sentence)
|
||||
current_len += sent_len + 1 # +1 for space
|
||||
continue
|
||||
|
||||
if current:
|
||||
chunks.append(" ".join(current))
|
||||
|
||||
# If a single sentence exceeds chunk_size, force-split it
|
||||
if sent_len > chunk_size:
|
||||
for i in range(0, sent_len, chunk_size - overlap):
|
||||
chunks.append(sentence[i : i + chunk_size])
|
||||
current = []
|
||||
current_len = 0
|
||||
else:
|
||||
overlap_text = _get_overlap(chunks, overlap) if chunks else ""
|
||||
if overlap_text:
|
||||
current = [overlap_text, sentence]
|
||||
current_len = len(overlap_text) + sent_len + 1
|
||||
else:
|
||||
current = [sentence]
|
||||
current_len = sent_len
|
||||
|
||||
if current:
|
||||
chunks.append(" ".join(current))
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
def _get_overlap(chunks: list[str], overlap: int) -> str:
|
||||
"""Extract the last ~`overlap` characters from the most recent chunk, snapped to a word boundary."""
|
||||
if not chunks or overlap <= 0:
|
||||
return ""
|
||||
last = chunks[-1]
|
||||
tail = last[-overlap:] if len(last) > overlap else last
|
||||
# Snap to the nearest word boundary to avoid starting mid-word
|
||||
space_idx = tail.find(" ")
|
||||
if space_idx > 0:
|
||||
tail = tail[space_idx + 1 :]
|
||||
return tail
|
||||
@@ -0,0 +1,47 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic_ai import Embedder
|
||||
|
||||
from stirling.rag.chunker import chunk_text
|
||||
from stirling.rag.store import Document
|
||||
|
||||
|
||||
class EmbeddingService:
|
||||
"""Wraps Pydantic AI's Embedder to provide document chunking and embedding."""
|
||||
|
||||
def __init__(self, model_name: str, chunk_size: int = 512, chunk_overlap: int = 64) -> None:
|
||||
self._embedder = Embedder(model_name)
|
||||
self._chunk_size = chunk_size
|
||||
self._chunk_overlap = chunk_overlap
|
||||
|
||||
async def embed_query(self, text: str) -> list[float]:
|
||||
"""Embed a search query, optimised for retrieval."""
|
||||
result = await self._embedder.embed_query(text)
|
||||
return list(result.embeddings[0])
|
||||
|
||||
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Embed multiple document texts for indexing."""
|
||||
if not texts:
|
||||
return []
|
||||
result = await self._embedder.embed_documents(texts)
|
||||
return [list(emb) for emb in result.embeddings]
|
||||
|
||||
def chunk_and_prepare(
|
||||
self,
|
||||
text: str,
|
||||
source: str = "",
|
||||
base_metadata: dict[str, str] | None = None,
|
||||
) -> list[Document]:
|
||||
"""Chunk text and return Document objects ready for embedding.
|
||||
|
||||
Each chunk gets a unique ID based on source and chunk index.
|
||||
"""
|
||||
chunks = chunk_text(text, self._chunk_size, self._chunk_overlap)
|
||||
documents: list[Document] = []
|
||||
for i, chunk in enumerate(chunks):
|
||||
meta = dict(base_metadata) if base_metadata else {}
|
||||
meta["source"] = source
|
||||
meta["chunk_index"] = str(i)
|
||||
doc_id = f"{source}:chunk:{i}" if source else f"chunk:{i}"
|
||||
documents.append(Document(id=doc_id, text=chunk, metadata=meta))
|
||||
return documents
|
||||
@@ -0,0 +1,133 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import psycopg
|
||||
from pgvector.psycopg import register_vector_async
|
||||
|
||||
from stirling.rag.store import Document, SearchResult, VectorStore
|
||||
|
||||
|
||||
class PgVectorStore(VectorStore):
|
||||
"""PostgreSQL + pgvector backed store.
|
||||
|
||||
Connects to an external Postgres instance (DSN provided via config) and uses the
|
||||
`vector` extension for similarity search. The schema is created on first use.
|
||||
"""
|
||||
|
||||
def __init__(self, dsn: str) -> None:
|
||||
if not dsn:
|
||||
raise ValueError("pgvector backend requires a non-empty DSN (STIRLING_RAG_PGVECTOR_DSN)")
|
||||
self._dsn = dsn
|
||||
self._initialized = False
|
||||
|
||||
async def _connect(self) -> psycopg.AsyncConnection:
|
||||
conn = await psycopg.AsyncConnection.connect(self._dsn)
|
||||
await register_vector_async(conn)
|
||||
return conn
|
||||
|
||||
async def _ensure_schema(self) -> None:
|
||||
if self._initialized:
|
||||
return
|
||||
async with await self._connect() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
|
||||
await cur.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS rag_documents (
|
||||
id TEXT NOT NULL,
|
||||
collection TEXT NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
metadata JSONB NOT NULL DEFAULT '{}'::jsonb,
|
||||
embedding vector NOT NULL,
|
||||
PRIMARY KEY (id, collection)
|
||||
)
|
||||
"""
|
||||
)
|
||||
await cur.execute("CREATE INDEX IF NOT EXISTS idx_rag_collection ON rag_documents(collection)")
|
||||
await conn.commit()
|
||||
self._initialized = True
|
||||
|
||||
async def add_documents(
|
||||
self,
|
||||
collection: str,
|
||||
documents: list[Document],
|
||||
embeddings: list[list[float]],
|
||||
) -> None:
|
||||
if len(documents) != len(embeddings):
|
||||
raise ValueError(f"Got {len(documents)} documents but {len(embeddings)} embeddings")
|
||||
if not documents:
|
||||
return
|
||||
|
||||
await self._ensure_schema()
|
||||
async with await self._connect() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
for doc, emb in zip(documents, embeddings):
|
||||
await cur.execute(
|
||||
"""
|
||||
INSERT INTO rag_documents (id, collection, text, metadata, embedding)
|
||||
VALUES (%s, %s, %s, %s::jsonb, %s)
|
||||
ON CONFLICT (id, collection)
|
||||
DO UPDATE SET
|
||||
text = EXCLUDED.text,
|
||||
metadata = EXCLUDED.metadata,
|
||||
embedding = EXCLUDED.embedding
|
||||
""",
|
||||
(doc.id, collection, doc.text, json.dumps(doc.metadata), emb),
|
||||
)
|
||||
await conn.commit()
|
||||
|
||||
async def search(
|
||||
self,
|
||||
collection: str,
|
||||
query_embedding: list[float],
|
||||
top_k: int = 5,
|
||||
) -> list[SearchResult]:
|
||||
await self._ensure_schema()
|
||||
async with await self._connect() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"""
|
||||
SELECT id, text, metadata, 1 - (embedding <=> %s) AS score
|
||||
FROM rag_documents
|
||||
WHERE collection = %s
|
||||
ORDER BY embedding <=> %s
|
||||
LIMIT %s
|
||||
""",
|
||||
(query_embedding, collection, query_embedding, top_k),
|
||||
)
|
||||
rows = await cur.fetchall()
|
||||
|
||||
return [
|
||||
SearchResult(
|
||||
document=Document(id=r[0], text=r[1], metadata=r[2] or {}),
|
||||
score=float(r[3]),
|
||||
)
|
||||
for r in rows
|
||||
]
|
||||
|
||||
async def delete_collection(self, collection: str) -> None:
|
||||
await self._ensure_schema()
|
||||
async with await self._connect() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute("DELETE FROM rag_documents WHERE collection = %s", (collection,))
|
||||
await conn.commit()
|
||||
|
||||
async def list_collections(self) -> list[str]:
|
||||
await self._ensure_schema()
|
||||
async with await self._connect() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute("SELECT DISTINCT collection FROM rag_documents ORDER BY collection")
|
||||
rows = await cur.fetchall()
|
||||
return [r[0] for r in rows]
|
||||
|
||||
async def has_collection(self, collection: str) -> bool:
|
||||
await self._ensure_schema()
|
||||
async with await self._connect() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"SELECT 1 FROM rag_documents WHERE collection = %s LIMIT 1",
|
||||
(collection,),
|
||||
)
|
||||
row = await cur.fetchone()
|
||||
return row is not None
|
||||
@@ -0,0 +1,80 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from stirling.rag.embedder import EmbeddingService
|
||||
from stirling.rag.store import Document, SearchResult, VectorStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RagService:
|
||||
"""Orchestrates embedding and vector storage for RAG workflows."""
|
||||
|
||||
def __init__(self, embedder: EmbeddingService, store: VectorStore, default_top_k: int = 5) -> None:
|
||||
self._embedder = embedder
|
||||
self._store = store
|
||||
self._default_top_k = default_top_k
|
||||
|
||||
async def index_text(
|
||||
self,
|
||||
collection: str,
|
||||
text: str,
|
||||
source: str = "",
|
||||
metadata: dict[str, str] | None = None,
|
||||
) -> int:
|
||||
"""Chunk, embed, and store text. Returns the number of chunks indexed."""
|
||||
documents = self._embedder.chunk_and_prepare(text, source=source, base_metadata=metadata)
|
||||
if not documents:
|
||||
return 0
|
||||
embeddings = await self._embedder.embed_documents([doc.text for doc in documents])
|
||||
await self._store.add_documents(collection, documents, embeddings)
|
||||
return len(documents)
|
||||
|
||||
async def index_documents(self, collection: str, documents: list[Document]) -> int:
|
||||
"""Embed and store pre-chunked documents. Returns the number stored."""
|
||||
if not documents:
|
||||
return 0
|
||||
embeddings = await self._embedder.embed_documents([doc.text for doc in documents])
|
||||
await self._store.add_documents(collection, documents, embeddings)
|
||||
return len(documents)
|
||||
|
||||
async def search(
|
||||
self,
|
||||
query: str,
|
||||
collection: str | None = None,
|
||||
top_k: int | None = None,
|
||||
) -> list[SearchResult]:
|
||||
"""Embed query and search across one or all collections.
|
||||
|
||||
If collection is None, searches all available collections and merges results.
|
||||
"""
|
||||
k = top_k if top_k is not None else self._default_top_k
|
||||
query_embedding = await self._embedder.embed_query(query)
|
||||
|
||||
if collection is not None:
|
||||
if not await self._store.has_collection(collection):
|
||||
return []
|
||||
return await self._store.search(collection, query_embedding, k)
|
||||
|
||||
# Search all collections, skipping any that error (e.g. dimension mismatch)
|
||||
collections = await self._store.list_collections()
|
||||
all_results: list[SearchResult] = []
|
||||
for col_name in collections:
|
||||
try:
|
||||
results = await self._store.search(col_name, query_embedding, k)
|
||||
all_results.extend(results)
|
||||
except Exception: # noqa: BLE001 — any backend error on one collection should not stop the others
|
||||
logger.warning("Skipping collection %s during cross-collection search", col_name, exc_info=True)
|
||||
|
||||
# Sort by score descending, return top_k across all collections
|
||||
all_results.sort(key=lambda r: r.score, reverse=True)
|
||||
return all_results[:k]
|
||||
|
||||
async def delete_collection(self, collection: str) -> None:
|
||||
"""Remove a collection and all its documents."""
|
||||
await self._store.delete_collection(collection)
|
||||
|
||||
async def list_collections(self) -> list[str]:
|
||||
"""List all available collections."""
|
||||
return await self._store.list_collections()
|
||||
@@ -0,0 +1,227 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
|
||||
import sqlite_vec
|
||||
|
||||
from stirling.rag.store import Document, SearchResult, VectorStore
|
||||
|
||||
|
||||
class SqliteVecStore(VectorStore):
|
||||
"""sqlite-vec backed vector store. Single-file SQLite database, embedded, no server.
|
||||
|
||||
Each collection gets its own `vec0` virtual table with a fixed embedding dimension
|
||||
(detected on first insert). Document metadata lives in a regular table joined by rowid.
|
||||
"""
|
||||
|
||||
def __init__(self, db_path: str | Path) -> None:
|
||||
is_memory = str(db_path) == ":memory:"
|
||||
self._db_path: Path | None = None if is_memory else Path(db_path)
|
||||
|
||||
if self._db_path is not None:
|
||||
self._db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
conn = sqlite3.connect(str(self._db_path), check_same_thread=False)
|
||||
else:
|
||||
conn = sqlite3.connect(":memory:", check_same_thread=False)
|
||||
|
||||
conn.enable_load_extension(True)
|
||||
sqlite_vec.load(conn)
|
||||
conn.enable_load_extension(False)
|
||||
if self._db_path is not None:
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
|
||||
self._conn = conn
|
||||
self._lock = asyncio.Lock()
|
||||
self._init_schema()
|
||||
|
||||
@classmethod
|
||||
def ephemeral(cls) -> SqliteVecStore:
|
||||
"""In-memory store for testing."""
|
||||
return cls(":memory:")
|
||||
|
||||
def _init_schema(self) -> None:
|
||||
self._conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS collections (
|
||||
name TEXT PRIMARY KEY,
|
||||
dim INTEGER NOT NULL,
|
||||
table_name TEXT NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
self._conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS documents (
|
||||
id TEXT NOT NULL,
|
||||
collection TEXT NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
metadata TEXT NOT NULL DEFAULT '{}',
|
||||
vec_rowid INTEGER NOT NULL,
|
||||
PRIMARY KEY (id, collection)
|
||||
)
|
||||
"""
|
||||
)
|
||||
self._conn.execute("CREATE INDEX IF NOT EXISTS idx_doc_collection ON documents(collection)")
|
||||
self._conn.commit()
|
||||
|
||||
@staticmethod
|
||||
def _sanitize_table_name(collection: str) -> str:
|
||||
safe = re.sub(r"[^a-zA-Z0-9_]", "_", collection)
|
||||
return f"vec_{safe}"
|
||||
|
||||
@staticmethod
|
||||
def _normalize(vector: list[float]) -> list[float]:
|
||||
norm = math.sqrt(sum(x * x for x in vector))
|
||||
if norm == 0:
|
||||
return list(vector)
|
||||
return [x / norm for x in vector]
|
||||
|
||||
async def add_documents(
|
||||
self,
|
||||
collection: str,
|
||||
documents: list[Document],
|
||||
embeddings: list[list[float]],
|
||||
) -> None:
|
||||
if len(documents) != len(embeddings):
|
||||
raise ValueError(f"Got {len(documents)} documents but {len(embeddings)} embeddings")
|
||||
if not documents:
|
||||
return
|
||||
|
||||
async with self._lock:
|
||||
await asyncio.to_thread(self._sync_add, collection, documents, embeddings)
|
||||
|
||||
def _sync_add(
|
||||
self,
|
||||
collection: str,
|
||||
documents: list[Document],
|
||||
embeddings: list[list[float]],
|
||||
) -> None:
|
||||
dim = len(embeddings[0])
|
||||
row = self._conn.execute("SELECT dim, table_name FROM collections WHERE name = ?", (collection,)).fetchone()
|
||||
if row is None:
|
||||
table_name = self._sanitize_table_name(collection)
|
||||
self._conn.execute(f"CREATE VIRTUAL TABLE IF NOT EXISTS {table_name} USING vec0(embedding float[{dim}])")
|
||||
self._conn.execute(
|
||||
"INSERT INTO collections(name, dim, table_name) VALUES (?, ?, ?)",
|
||||
(collection, dim, table_name),
|
||||
)
|
||||
else:
|
||||
existing_dim, table_name = row
|
||||
if existing_dim != dim:
|
||||
raise ValueError(f"Collection {collection} has dim {existing_dim}, got embedding of dim {dim}")
|
||||
|
||||
# Upsert: delete existing docs with matching IDs first
|
||||
ids = [doc.id for doc in documents]
|
||||
placeholders = ",".join("?" * len(ids))
|
||||
existing = self._conn.execute(
|
||||
f"SELECT vec_rowid FROM documents WHERE collection = ? AND id IN ({placeholders})",
|
||||
(collection, *ids),
|
||||
).fetchall()
|
||||
if existing:
|
||||
vec_rowids = [r[0] for r in existing]
|
||||
row_placeholders = ",".join("?" * len(vec_rowids))
|
||||
self._conn.execute(
|
||||
f"DELETE FROM {table_name} WHERE rowid IN ({row_placeholders})",
|
||||
vec_rowids,
|
||||
)
|
||||
self._conn.execute(
|
||||
f"DELETE FROM documents WHERE collection = ? AND id IN ({placeholders})",
|
||||
(collection, *ids),
|
||||
)
|
||||
|
||||
for doc, emb in zip(documents, embeddings):
|
||||
normalized = self._normalize(list(emb))
|
||||
cursor = self._conn.execute(
|
||||
f"INSERT INTO {table_name}(embedding) VALUES (?)",
|
||||
(sqlite_vec.serialize_float32(normalized),),
|
||||
)
|
||||
vec_rowid = cursor.lastrowid
|
||||
self._conn.execute(
|
||||
"INSERT INTO documents(id, collection, text, metadata, vec_rowid) VALUES (?, ?, ?, ?, ?)",
|
||||
(doc.id, collection, doc.text, json.dumps(doc.metadata), vec_rowid),
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
async def search(
|
||||
self,
|
||||
collection: str,
|
||||
query_embedding: list[float],
|
||||
top_k: int = 5,
|
||||
) -> list[SearchResult]:
|
||||
async with self._lock:
|
||||
return await asyncio.to_thread(self._sync_search, collection, query_embedding, top_k)
|
||||
|
||||
def _sync_search(
|
||||
self,
|
||||
collection: str,
|
||||
query_embedding: list[float],
|
||||
top_k: int,
|
||||
) -> list[SearchResult]:
|
||||
row = self._conn.execute("SELECT table_name, dim FROM collections WHERE name = ?", (collection,)).fetchone()
|
||||
if row is None:
|
||||
return []
|
||||
table_name, dim = row
|
||||
if len(query_embedding) != dim:
|
||||
raise ValueError(f"Query embedding dim {len(query_embedding)} does not match collection dim {dim}")
|
||||
|
||||
normalized = self._normalize(list(query_embedding))
|
||||
query_blob = sqlite_vec.serialize_float32(normalized)
|
||||
|
||||
results = self._conn.execute(
|
||||
f"""
|
||||
SELECT d.id, d.text, d.metadata, v.distance
|
||||
FROM {table_name} v
|
||||
JOIN documents d ON d.vec_rowid = v.rowid AND d.collection = ?
|
||||
WHERE v.embedding MATCH ? AND k = ?
|
||||
ORDER BY v.distance
|
||||
""",
|
||||
(collection, query_blob, top_k),
|
||||
).fetchall()
|
||||
|
||||
return [
|
||||
SearchResult(
|
||||
document=Document(
|
||||
id=r[0],
|
||||
text=r[1],
|
||||
metadata=json.loads(r[2]) if r[2] else {},
|
||||
),
|
||||
# For normalized vectors: cosine_sim = 1 - (L2^2 / 2)
|
||||
score=max(0.0, 1.0 - (r[3] ** 2) / 2.0),
|
||||
)
|
||||
for r in results
|
||||
]
|
||||
|
||||
async def delete_collection(self, collection: str) -> None:
|
||||
async with self._lock:
|
||||
await asyncio.to_thread(self._sync_delete_collection, collection)
|
||||
|
||||
def _sync_delete_collection(self, collection: str) -> None:
|
||||
row = self._conn.execute("SELECT table_name FROM collections WHERE name = ?", (collection,)).fetchone()
|
||||
if row is None:
|
||||
return
|
||||
table_name = row[0]
|
||||
self._conn.execute(f"DROP TABLE IF EXISTS {table_name}")
|
||||
self._conn.execute("DELETE FROM documents WHERE collection = ?", (collection,))
|
||||
self._conn.execute("DELETE FROM collections WHERE name = ?", (collection,))
|
||||
self._conn.commit()
|
||||
|
||||
async def list_collections(self) -> list[str]:
|
||||
async with self._lock:
|
||||
return await asyncio.to_thread(self._sync_list_collections)
|
||||
|
||||
def _sync_list_collections(self) -> list[str]:
|
||||
rows = self._conn.execute("SELECT name FROM collections ORDER BY name").fetchall()
|
||||
return [r[0] for r in rows]
|
||||
|
||||
async def has_collection(self, collection: str) -> bool:
|
||||
async with self._lock:
|
||||
return await asyncio.to_thread(self._sync_has_collection, collection)
|
||||
|
||||
def _sync_has_collection(self, collection: str) -> bool:
|
||||
row = self._conn.execute("SELECT 1 FROM collections WHERE name = ?", (collection,)).fetchone()
|
||||
return row is not None
|
||||
@@ -0,0 +1,59 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass
|
||||
class Document:
|
||||
"""A chunk of text with metadata, ready for embedding and storage."""
|
||||
|
||||
id: str
|
||||
text: str
|
||||
metadata: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SearchResult:
|
||||
"""A document returned from a vector search with its relevance score."""
|
||||
|
||||
document: Document
|
||||
score: float
|
||||
|
||||
|
||||
class VectorStore(ABC):
|
||||
"""Abstract interface for vector storage backends.
|
||||
|
||||
Implementations must handle persistence, collection management,
|
||||
and nearest-neighbor search over pre-computed embeddings.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def add_documents(
|
||||
self,
|
||||
collection: str,
|
||||
documents: list[Document],
|
||||
embeddings: list[list[float]],
|
||||
) -> None:
|
||||
"""Store documents with their embeddings in the named collection."""
|
||||
|
||||
@abstractmethod
|
||||
async def search(
|
||||
self,
|
||||
collection: str,
|
||||
query_embedding: list[float],
|
||||
top_k: int = 5,
|
||||
) -> list[SearchResult]:
|
||||
"""Return the top_k most similar documents from the collection."""
|
||||
|
||||
@abstractmethod
|
||||
async def delete_collection(self, collection: str) -> None:
|
||||
"""Remove a collection and all its documents."""
|
||||
|
||||
@abstractmethod
|
||||
async def list_collections(self) -> list[str]:
|
||||
"""Return names of all existing collections."""
|
||||
|
||||
@abstractmethod
|
||||
async def has_collection(self, collection: str) -> bool:
|
||||
"""Check whether a collection exists."""
|
||||
Reference in New Issue
Block a user