mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-16 03:20:46 +02:00
Add ability for Stirling engine to reason across large documents (#6314)
# Description of Changes Adds storage in the database for full document content alongside the RAG content (and changes the service to `DocumentService` instead of `RagService`). Then adds a generic capability that should be usable by any agent (currently just used by the Question Agent) which allows the agent to pull out the full contents of the doc, chunks it into various sections that will fit in the context window, and then processes them in parallel to create an intermediate result, and then processes the intermediate result into a final answer. It will re-chunk as many times as necessary to get the content small enough for the actual answer to be analysed (I've tested on PDFs ~3500 pages long, which is well above the context limit and requires maybe 3 rounds of compression to get an answer). The new full doc analysis stuff is heavier than the RAG lookup so both remain. The agents should use RAG for targeted info and the chunked reasoner for info that requires reading the full doc.
This commit is contained in:
James Brunton
@@ -0,0 +1,102 @@
|
||||
# Document Storage
|
||||
|
||||
The `documents` package owns all stored content for a document under a single
|
||||
`collection` (file id):
|
||||
|
||||
* **Vector chunks** — small, embedded chunks for RAG-style retrieval.
|
||||
* **Ordered pages** — the original page text retained in document order, used
|
||||
for whole-document reading.
|
||||
|
||||
Both representations are populated by a single `ingest()` call and removed
|
||||
together by `delete_collection()`.
|
||||
|
||||
## 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.documents import RagCapability
|
||||
|
||||
# Only searches docs indexed under "company-docs"
|
||||
scoped = RagCapability(runtime.documents, collections=["company-docs"], top_k=3)
|
||||
|
||||
# Searches multiple collections
|
||||
multi = RagCapability(runtime.documents, collections=["company-docs", "product-specs"])
|
||||
|
||||
# No collections arg = searches all collections in the store
|
||||
everything = RagCapability(runtime.documents)
|
||||
```
|
||||
|
||||
## Config
|
||||
|
||||
Non-secret defaults live in the committed `engine/.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
|
||||
```
|
||||
|
||||
Provider credentials (and any local overrides) go in the uncommitted
|
||||
`engine/.env.local`:
|
||||
|
||||
```
|
||||
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 `DocumentStore` interface, so agents and the
|
||||
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 |
|
||||
|--------|----------|---------|
|
||||
| POST | `/api/v1/documents` | Replace-ingest a document's pages |
|
||||
| DELETE | `/api/v1/documents/{document_id}` | Delete a document's stored content |
|
||||
@@ -0,0 +1,20 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from stirling.documents.embedder import EmbeddingService
|
||||
from stirling.documents.pgvector_store import PgVectorStore
|
||||
from stirling.documents.rag_capability import RagCapability
|
||||
from stirling.documents.service import DocumentService
|
||||
from stirling.documents.sqlite_vec_store import SqliteVecStore
|
||||
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
|
||||
|
||||
__all__ = [
|
||||
"Document",
|
||||
"DocumentService",
|
||||
"DocumentStore",
|
||||
"EmbeddingService",
|
||||
"PgVectorStore",
|
||||
"RagCapability",
|
||||
"SearchResult",
|
||||
"SqliteVecStore",
|
||||
"StoredPage",
|
||||
]
|
||||
@@ -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,68 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic_ai import Embedder
|
||||
|
||||
from stirling.documents.chunker import chunk_text
|
||||
from stirling.documents.store import Document
|
||||
|
||||
# Keep each upstream embed request under every major provider's per-call limit while
|
||||
# still batching large enough that a book-sized document ingests in a reasonable number
|
||||
# of round trips. VoyageAI caps at 1000, OpenAI at 2048, Cohere at 96; 256 is a good
|
||||
# default for Voyage/OpenAI. Cohere users should pass a lower value via construction.
|
||||
DEFAULT_EMBED_BATCH_SIZE = 256
|
||||
|
||||
|
||||
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,
|
||||
embed_batch_size: int = DEFAULT_EMBED_BATCH_SIZE,
|
||||
) -> None:
|
||||
self._embedder = Embedder(model_name)
|
||||
self._chunk_size = chunk_size
|
||||
self._chunk_overlap = chunk_overlap
|
||||
self._embed_batch_size = embed_batch_size
|
||||
|
||||
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.
|
||||
|
||||
Splits the input into batches of ``embed_batch_size`` so callers can hand us
|
||||
any number of chunks without hitting provider per-request limits.
|
||||
"""
|
||||
if not texts:
|
||||
return []
|
||||
all_embeddings: list[list[float]] = []
|
||||
for start in range(0, len(texts), self._embed_batch_size):
|
||||
batch = texts[start : start + self._embed_batch_size]
|
||||
result = await self._embedder.embed_documents(batch)
|
||||
all_embeddings.extend(list(emb) for emb in result.embeddings)
|
||||
return all_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,219 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import psycopg
|
||||
from pgvector.psycopg import register_vector_async
|
||||
|
||||
from stirling.contracts.documents import Page, PageRange
|
||||
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
|
||||
|
||||
|
||||
class PgVectorStore(DocumentStore):
|
||||
"""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.
|
||||
|
||||
Holds two tables under the same connection:
|
||||
|
||||
* ``rag_documents`` - vector chunks for RAG search.
|
||||
* ``document_pages`` - ordered page text for whole-document reading.
|
||||
"""
|
||||
|
||||
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 documents_meta (
|
||||
collection TEXT PRIMARY KEY,
|
||||
source TEXT NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
await cur.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS rag_documents (
|
||||
id TEXT NOT NULL,
|
||||
collection TEXT NOT NULL
|
||||
REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||
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 cur.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS document_pages (
|
||||
collection TEXT NOT NULL
|
||||
REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||
page_number INTEGER NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
char_count INTEGER NOT NULL,
|
||||
PRIMARY KEY (collection, page_number)
|
||||
)
|
||||
"""
|
||||
)
|
||||
await cur.execute("CREATE INDEX IF NOT EXISTS idx_pages_collection ON document_pages(collection)")
|
||||
await conn.commit()
|
||||
self._initialized = True
|
||||
|
||||
async def ensure_collection(self, collection: str, source: str) -> None:
|
||||
await self._ensure_schema()
|
||||
async with await self._connect() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"""
|
||||
INSERT INTO documents_meta (collection, source)
|
||||
VALUES (%s, %s)
|
||||
ON CONFLICT (collection) DO UPDATE SET source = EXCLUDED.source
|
||||
""",
|
||||
(collection, source),
|
||||
)
|
||||
await conn.commit()
|
||||
|
||||
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 add_pages(self, collection: str, pages: list[StoredPage]) -> None:
|
||||
await self._ensure_schema()
|
||||
async with await self._connect() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute("DELETE FROM document_pages WHERE collection = %s", (collection,))
|
||||
if pages:
|
||||
await cur.executemany(
|
||||
"""
|
||||
INSERT INTO document_pages (collection, page_number, text, char_count)
|
||||
VALUES (%s, %s, %s, %s)
|
||||
""",
|
||||
[(collection, p.page_number, p.text, p.char_count) for p in pages],
|
||||
)
|
||||
await conn.commit()
|
||||
|
||||
async def read_pages(
|
||||
self,
|
||||
collection: str,
|
||||
page_range: PageRange | None = None,
|
||||
) -> list[Page]:
|
||||
await self._ensure_schema()
|
||||
async with await self._connect() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
if page_range is None:
|
||||
await cur.execute(
|
||||
"SELECT page_number, text, char_count FROM document_pages "
|
||||
"WHERE collection = %s ORDER BY page_number",
|
||||
(collection,),
|
||||
)
|
||||
else:
|
||||
await cur.execute(
|
||||
"SELECT page_number, text, char_count FROM document_pages "
|
||||
"WHERE collection = %s AND page_number BETWEEN %s AND %s "
|
||||
"ORDER BY page_number",
|
||||
(collection, page_range.start, page_range.end),
|
||||
)
|
||||
rows = await cur.fetchall()
|
||||
return [Page(page_number=r[0], text=r[1], char_count=r[2]) 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:
|
||||
# Cascade FKs handle rag_documents and document_pages.
|
||||
await cur.execute("DELETE FROM documents_meta 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 collection FROM documents_meta 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 documents_meta WHERE collection = %s",
|
||||
(collection,),
|
||||
)
|
||||
row = await cur.fetchone()
|
||||
return row is not None
|
||||
|
||||
async def close(self) -> None:
|
||||
# Connections are opened and closed per call, so nothing persistent to release.
|
||||
return None
|
||||
@@ -0,0 +1,149 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
from pydantic_ai import FunctionToolset, RunContext, ToolDefinition
|
||||
from pydantic_ai.toolsets import AbstractToolset
|
||||
|
||||
from stirling.documents.service import DocumentService
|
||||
from stirling.documents.store import SearchResult
|
||||
from stirling.models import FileId
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
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.
|
||||
|
||||
Lifecycle: a ``RagCapability`` instance is intended to live for the duration of a
|
||||
single agent run.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
documents: DocumentService,
|
||||
collections: list[FileId] | None = None,
|
||||
top_k: int = 5,
|
||||
max_searches: int = 5,
|
||||
) -> None:
|
||||
self._documents = documents
|
||||
self._collections = collections
|
||||
self._top_k = top_k
|
||||
self._max_searches = max_searches
|
||||
self._search_count = 0
|
||||
toolset: FunctionToolset[None] = FunctionToolset()
|
||||
toolset.add_function(
|
||||
self._search_knowledge,
|
||||
name="search_knowledge",
|
||||
prepare=self._prepare_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[FileId]) -> 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._documents.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 _prepare_search_knowledge(
|
||||
self,
|
||||
ctx: RunContext[None],
|
||||
tool_def: ToolDefinition,
|
||||
) -> ToolDefinition | None:
|
||||
"""Remove the search tool from the agent's toolset once the per-run search
|
||||
budget is exhausted. The agent then has no choice but to answer from what it
|
||||
has already retrieved, which prevents runaway search loops."""
|
||||
if self._search_count >= self._max_searches:
|
||||
return None
|
||||
return tool_def
|
||||
|
||||
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.
|
||||
"""
|
||||
self._search_count += 1
|
||||
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._documents.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._documents.search(query, top_k=k)
|
||||
|
||||
if not results:
|
||||
logger.info("[rag] search_knowledge query=%r -> 0 results", query)
|
||||
return "No relevant results found in the knowledge base."
|
||||
|
||||
formatted = self._format_results(results)
|
||||
logger.info(
|
||||
"[rag] search_knowledge query=%r -> %d results, %d chars",
|
||||
query,
|
||||
len(results),
|
||||
len(formatted),
|
||||
)
|
||||
logger.debug("[rag] search_knowledge query=%r returned:\n%s", query, formatted)
|
||||
return formatted
|
||||
|
||||
@staticmethod
|
||||
def _format_results(results: list[SearchResult]) -> str:
|
||||
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,133 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from stirling.contracts.documents import Page, PageRange, PageText
|
||||
from stirling.documents.embedder import EmbeddingService
|
||||
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
|
||||
from stirling.models import FileId
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PAGE_NUMBER_METADATA_KEY = "page_number"
|
||||
CONTENT_TYPE_METADATA_KEY = "content_type"
|
||||
PAGE_TEXT_CONTENT_TYPE = "page_text"
|
||||
|
||||
|
||||
class DocumentService:
|
||||
"""Top-level facade for stored document content.
|
||||
|
||||
Holds two representations of every document under a single ``collection``:
|
||||
|
||||
* **Vector chunks** for RAG-style semantic retrieval (``search``).
|
||||
* **Ordered pages** for whole-document reading (``read_pages``).
|
||||
|
||||
Both are populated by :meth:`ingest` from a single ``pages`` payload. Agents
|
||||
pick the strategy that fits the question; they don't need to know which
|
||||
storage they're hitting.
|
||||
"""
|
||||
|
||||
def __init__(self, embedder: EmbeddingService, store: DocumentStore, default_top_k: int = 5) -> None:
|
||||
self._embedder = embedder
|
||||
self._store = store
|
||||
self._default_top_k = default_top_k
|
||||
|
||||
async def ingest(
|
||||
self,
|
||||
collection: FileId,
|
||||
pages: list[PageText],
|
||||
source: str,
|
||||
) -> int:
|
||||
"""Replace-ingest a document. Returns the number of vector chunks indexed.
|
||||
|
||||
This wipes any previously-stored content for ``collection`` and writes
|
||||
both the vector-chunk and page-text representations from the same
|
||||
``pages`` payload. Pages with empty/whitespace-only text are skipped
|
||||
for chunking but still written to the page store so page numbering is
|
||||
preserved end-to-end.
|
||||
"""
|
||||
await self._store.delete_collection(collection)
|
||||
await self._store.ensure_collection(collection, source)
|
||||
|
||||
stored_pages = [StoredPage(page_number=p.page_number, text=p.text, char_count=len(p.text)) for p in pages]
|
||||
await self._store.add_pages(collection, stored_pages)
|
||||
|
||||
chunks: list[Document] = []
|
||||
for page in pages:
|
||||
if not page.text.strip():
|
||||
continue
|
||||
chunks.extend(
|
||||
self._embedder.chunk_and_prepare(
|
||||
text=page.text,
|
||||
source=f"{source}:page:{page.page_number}",
|
||||
base_metadata={
|
||||
PAGE_NUMBER_METADATA_KEY: str(page.page_number),
|
||||
CONTENT_TYPE_METADATA_KEY: PAGE_TEXT_CONTENT_TYPE,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
if not chunks:
|
||||
return 0
|
||||
embeddings = await self._embedder.embed_documents([doc.text for doc in chunks])
|
||||
await self._store.add_documents(collection, chunks, embeddings)
|
||||
return len(chunks)
|
||||
|
||||
async def search(
|
||||
self,
|
||||
query: str,
|
||||
collection: FileId | 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 read_pages(
|
||||
self,
|
||||
collection: FileId,
|
||||
page_range: PageRange | None = None,
|
||||
) -> list[Page]:
|
||||
"""Return ordered page text for ``collection``.
|
||||
|
||||
Empty list if the collection has no stored pages.
|
||||
"""
|
||||
return await self._store.read_pages(collection, page_range)
|
||||
|
||||
async def delete_collection(self, collection: FileId) -> None:
|
||||
"""Remove a collection's chunks and pages."""
|
||||
await self._store.delete_collection(collection)
|
||||
|
||||
async def has_collection(self, collection: FileId) -> bool:
|
||||
"""Check whether a collection exists."""
|
||||
return await self._store.has_collection(collection)
|
||||
|
||||
async def list_collections(self) -> list[FileId]:
|
||||
"""List all available collections."""
|
||||
return [FileId(name) for name in await self._store.list_collections()]
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Release the underlying store's resources."""
|
||||
await self._store.close()
|
||||
@@ -0,0 +1,321 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
|
||||
import sqlite_vec
|
||||
|
||||
from stirling.contracts.documents import Page, PageRange
|
||||
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
|
||||
|
||||
|
||||
class SqliteVecStore(DocumentStore):
|
||||
"""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)
|
||||
# Required so cascade deletes from documents_meta clean up child tables.
|
||||
conn.execute("PRAGMA foreign_keys=ON")
|
||||
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 documents_meta (
|
||||
collection TEXT PRIMARY KEY,
|
||||
source TEXT NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
self._conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS collections (
|
||||
name TEXT PRIMARY KEY REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||
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 REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||
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.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS document_pages (
|
||||
collection TEXT NOT NULL REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||
page_number INTEGER NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
char_count INTEGER NOT NULL,
|
||||
PRIMARY KEY (collection, page_number)
|
||||
)
|
||||
"""
|
||||
)
|
||||
self._conn.execute("CREATE INDEX IF NOT EXISTS idx_pages_collection ON document_pages(collection)")
|
||||
self._conn.commit()
|
||||
|
||||
async def ensure_collection(self, collection: str, source: str) -> None:
|
||||
async with self._lock:
|
||||
await asyncio.to_thread(self._sync_ensure_collection, collection, source)
|
||||
|
||||
def _sync_ensure_collection(self, collection: str, source: str) -> None:
|
||||
self._conn.execute(
|
||||
"""
|
||||
INSERT INTO documents_meta(collection, source) VALUES (?, ?)
|
||||
ON CONFLICT(collection) DO UPDATE SET source = excluded.source
|
||||
""",
|
||||
(collection, source),
|
||||
)
|
||||
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 add_pages(self, collection: str, pages: list[StoredPage]) -> None:
|
||||
async with self._lock:
|
||||
await asyncio.to_thread(self._sync_add_pages, collection, pages)
|
||||
|
||||
def _sync_add_pages(self, collection: str, pages: list[StoredPage]) -> None:
|
||||
self._conn.execute("DELETE FROM document_pages WHERE collection = ?", (collection,))
|
||||
if pages:
|
||||
self._conn.executemany(
|
||||
"INSERT INTO document_pages(collection, page_number, text, char_count) VALUES (?, ?, ?, ?)",
|
||||
[(collection, p.page_number, p.text, p.char_count) for p in pages],
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
async def read_pages(
|
||||
self,
|
||||
collection: str,
|
||||
page_range: PageRange | None = None,
|
||||
) -> list[Page]:
|
||||
async with self._lock:
|
||||
return await asyncio.to_thread(self._sync_read_pages, collection, page_range)
|
||||
|
||||
def _sync_read_pages(
|
||||
self,
|
||||
collection: str,
|
||||
page_range: PageRange | None,
|
||||
) -> list[Page]:
|
||||
if page_range is None:
|
||||
rows = self._conn.execute(
|
||||
"SELECT page_number, text, char_count FROM document_pages WHERE collection = ? ORDER BY page_number",
|
||||
(collection,),
|
||||
).fetchall()
|
||||
else:
|
||||
rows = self._conn.execute(
|
||||
"SELECT page_number, text, char_count FROM document_pages "
|
||||
"WHERE collection = ? AND page_number BETWEEN ? AND ? ORDER BY page_number",
|
||||
(collection, page_range.start, page_range.end),
|
||||
).fetchall()
|
||||
return [Page(page_number=r[0], text=r[1], char_count=r[2]) for r in rows]
|
||||
|
||||
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:
|
||||
# Drop the sqlite-vec virtual table first; FK cascade handles the regular tables
|
||||
# (collections, documents, document_pages) when documents_meta is deleted.
|
||||
row = self._conn.execute("SELECT table_name FROM collections WHERE name = ?", (collection,)).fetchone()
|
||||
if row is not None:
|
||||
self._conn.execute(f"DROP TABLE IF EXISTS {row[0]}")
|
||||
self._conn.execute("DELETE FROM documents_meta WHERE collection = ?", (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 collection FROM documents_meta ORDER BY collection").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 documents_meta WHERE collection = ?",
|
||||
(collection,),
|
||||
).fetchone()
|
||||
return row is not None
|
||||
|
||||
async def close(self) -> None:
|
||||
async with self._lock:
|
||||
await asyncio.to_thread(self._sync_close)
|
||||
|
||||
def _sync_close(self) -> None:
|
||||
"""Checkpoint the WAL into the main database file and close the connection so
|
||||
the .db-shm and .db-wal files are cleaned up on graceful shutdown."""
|
||||
if self._db_path is not None:
|
||||
try:
|
||||
self._conn.execute("PRAGMA wal_checkpoint(TRUNCATE)")
|
||||
self._conn.commit()
|
||||
except sqlite3.Error:
|
||||
# Best effort: if checkpointing fails we still want to close the connection.
|
||||
pass
|
||||
self._conn.close()
|
||||
@@ -0,0 +1,113 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from stirling.contracts.documents import Page, PageRange
|
||||
|
||||
|
||||
@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
|
||||
|
||||
|
||||
@dataclass
|
||||
class StoredPage:
|
||||
"""A page as written to the store. ``char_count`` is precomputed at ingest."""
|
||||
|
||||
page_number: int
|
||||
text: str
|
||||
char_count: int
|
||||
|
||||
|
||||
class DocumentStore(ABC):
|
||||
"""Abstract interface for document storage backends.
|
||||
|
||||
Backends hold two representations of every document:
|
||||
|
||||
* **Vector chunks** - small, embedded chunks used for RAG search.
|
||||
* **Ordered pages** - the original page text retained in document order,
|
||||
used for whole-document reading.
|
||||
|
||||
Both representations live under the same ``collection`` (file id) and are
|
||||
rooted at a single parent row in ``documents_meta``. Removing that parent
|
||||
row cascades to both child representations, so :meth:`delete_collection`
|
||||
is one logical delete.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def ensure_collection(self, collection: str, source: str) -> None:
|
||||
"""Upsert the top-level ``documents_meta`` row for this collection.
|
||||
|
||||
Must be called before :meth:`add_pages` or :meth:`add_documents`. Both
|
||||
of those write into child tables that hold a foreign key to the parent
|
||||
row, so it must exist first.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def add_documents(
|
||||
self,
|
||||
collection: str,
|
||||
documents: list[Document],
|
||||
embeddings: list[list[float]],
|
||||
) -> None:
|
||||
"""Store vector chunks 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 vector chunks from the collection."""
|
||||
|
||||
@abstractmethod
|
||||
async def add_pages(self, collection: str, pages: list[StoredPage]) -> None:
|
||||
"""Replace the stored pages for ``collection`` with the supplied pages.
|
||||
|
||||
Implementations must remove any previously-stored pages for the
|
||||
collection before writing, so callers can re-ingest by calling this
|
||||
method again.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def read_pages(
|
||||
self,
|
||||
collection: str,
|
||||
page_range: PageRange | None = None,
|
||||
) -> list[Page]:
|
||||
"""Return ordered pages for ``collection``.
|
||||
|
||||
If ``page_range`` is ``None`` all pages are returned. Otherwise only
|
||||
pages whose ``page_number`` falls within the inclusive range are
|
||||
returned. Pages are always ordered by ``page_number`` ascending.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def delete_collection(self, collection: str) -> None:
|
||||
"""Remove a collection's chunks and pages."""
|
||||
|
||||
@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."""
|
||||
|
||||
@abstractmethod
|
||||
async def close(self) -> None:
|
||||
"""Release any resources held by the store (connections, handles, etc.)."""
|
||||
Reference in New Issue
Block a user