agentlib/rag

The rag package handles everything related to building and querying the vector index: embedding generation, text chunking, batch caching, and the TuringDB vector index operations.

Overview

Two distinct phases use this package:

Model registry — models.py

MODELS is a plain dict[str, ModelSpec] that maps a model name to its embedding properties:

class ModelSpec(TypedDict):
    dimensions: int        # Output vector size
    context_window: int    # Max tokens the model accepts
    provider: Literal["ollama", "sentence_transformers"]

Registered models

Adding a model

Add an entry to MODELS in models.py. No other changes are required — the factory and builder pick it up automatically.

Embeddings — embeddings.py

EmbeddingsHandler is an abstract base class for all embedding providers:

class EmbeddingsHandler(ABC):
    batch_size: int = 1024          # Max texts per call

    def embed(self, texts: list[str]) -> np.ndarray: ...
    # Returns float32 array of shape (len(texts), dim)
    # Raises ValueError if len(texts) > batch_size

Factory — make_handler(model: str) -> EmbeddingsHandler

Looks up model in MODELS and returns the appropriate handler. Raises ValueError for unknown models or providers.

from agentlib.rag.embeddings import make_handler

handler = make_handler("minishlab/potion-base-32M")
vectors = handler.embed(["ATP synthesis", "cell signaling"])  # shape (2, 512)

Providers

Ollama — ollama.py

OllamaEmbeddingsHandler wraps the ollama Python client.

OllamaEmbeddingsHandler(model: str, host: str = "http://localhost:11434")

Requires Ollama running locally with the target model pulled (ollama pull <model>).

sentence-transformers — sentence_transformers.py

SentenceTransformersEmbeddingsHandler wraps SentenceTransformer. The model is downloaded automatically from Hugging Face on first use.

SentenceTransformersEmbeddingsHandler(model: str)

Text chunker — chunker.py

def chunk_text(text: str, context_window: int, fill_rate: float = 0.8) -> list[str]

Splits text into overlapping word-based chunks that fit within the model's context window.

Tokenisation

Tokens are approximated as whitespace-separated words (text.split()). This is intentionally simple — no subword tokeniser is used. The approximation is conservative enough in practice because real subword tokenisers produce at least as many tokens as words, so a word-count budget always stays within the model's hard limit.

Algorithm

Given n words and a context_window of W:

1. Compute chunk size: chunk_size = max(1, int(W * fill_rate)) At the default fill_rate=0.8, a 512-token window yields chunks of at most 409 words. The headroom prevents the embedding model from silently truncating sequences near its limit.

2. Early exit: if n <= chunk_size, the text already fits — return [text] unchanged (no chunking, no copy).

3. Compute overlap: overlap = max(1, chunk_size // 10) 10% of the chunk size. For a 409-word chunk that is ~40 words of shared context between consecutive chunks.

4. Compute step: step = chunk_size - overlap The cursor advances by step words after each chunk, so consecutive chunks share overlap words at the boundary.

5. Sliding window loop: starting at i = 0, emit words[i : i + chunk_size] as a space-joined string, then advance i += step. Continue until i >= n. The last chunk may be shorter than chunk_size if fewer than chunk_size words remain.

Example

For context_window=10, fill_rate=0.8:

chunk_size = 8
overlap    = 0  (8 // 10 = 0, but max(1, 0) = 1)
step       = 7

words = [w0, w1, w2, w3, w4, w5, w6, w7, w8, w9, w10, w11, w12, w13]

chunk 0: words[0:8]   = w0 w1 w2 w3 w4 w5 w6 w7
chunk 1: words[7:15]  = w7 w8 w9 w10 w11 w12 w13      ← w7 is the shared overlap word

Trade-offs

The overlap ensures that a concept split across a chunk boundary appears fully in at least one chunk, improving recall for phrases that straddle the cut point.

Embeddings cache — cache.py

EmbeddingsCache tracks which embedding batch CSVs have already been written to data_dir, so that interrupted builds can resume without recomputing done batches.

cache = EmbeddingsCache(data_dir)          # scans data_dir for existing CSVs on init
cache.is_cached(index_name, batch_idx)     # True if CSV already exists
cache.register(index_name, batch_idx)      # mark a batch as written
cache.expected_path(index_name, batch_idx) # canonical Path for the CSV

CSV filenames follow the pattern {index_name}_{batch_idx:06d}.csv. All mutations are protected by a threading.Lock.

Vector index manager — index.py

VectorIndexManager is a stateless class of classmethods that wrap TuringDB vector operations.

vector_search signature

VectorIndexManager.vector_search(
    client: TuringDB,
    index_name: str,
    vector: list[float],
    top_k: int = 10,
    id_pt: str = "dbId",
    name_pt: str = "displayName",
    prop_pt: str = "text",
    return_labels: bool = False,
) -> pd.DataFrame

Index builder — builder.py

build_vector_index() is the top-level function that orchestrates the full build pipeline.

def build_vector_index(
    client: TuringDB,
    graph_name: str,
    model: str,
    data_dir: Path,
) -> BuildResult

Pipeline steps

1. Ensure index — calls VectorIndexManager.ensure_index() (drops and recreates for a clean build)
2. Fetch text nodes — MATCH (n) WHERE n.text IS NOT NULL RETURN n.dbId, n.text
3. Build chunk list — applies chunk_text() to every node; produces (node_id, chunk) pairs
4. Process batches — for each batch of size handler.batch_size:
5. Skip if cached (EmbeddingsCache.is_cached)
6. Otherwise embed, write CSV, register in cache
7. Load CSV into TuringDB (VectorIndexManager.add_vectors)
8. Return BuildResult

BuildResult

@dataclass
class BuildResult:
    index_name: str
    total_chunks: int
    batches_total: int
    batches_computed: int   # batches that required embedding
    batches_skipped: int    # batches loaded from cache

The build is idempotent: re-running it skips already-computed batches. A completely cached build still reloads all vectors into TuringDB (the index is recreated fresh each time).