src/store.py

"""ChromaDB-backed vector store (persistent on disk, cosine space).

Collection-generic: the triage app keeps two collections — the policy rules and past
adjudicated cases — so every function takes a collection name. We supply our own embeddings
(no Chroma embedding function); `search()` converts cosine distance to similarity via 1 - dist.
"""
from __future__ import annotations

import os
from dataclasses import dataclass, field
from functools import lru_cache

from .config import CONFIG
from .embed import embed_query, embed_texts


@dataclass
class Retrieved:
    text: str
    score: float
    metadata: dict = field(default_factory=dict)


@lru_cache(maxsize=1)
def _client():
    # One PersistentClient per process. Re-instantiating Chroma's client against the same
    # path repeatedly corrupts its shared-system-client lifecycle (tenant / RustBindings
    # errors), which surfaced as flaky failures across the agent's many tool calls.
    import chromadb

    os.makedirs(CONFIG.chroma_dir, exist_ok=True)
    return chromadb.PersistentClient(path=CONFIG.chroma_dir)


@lru_cache(maxsize=16)
def get_collection(name: str):
    return _client().get_or_create_collection(
        name=name,
        metadata={"hnsw:space": "cosine"},
    )


def index_documents(collection: str, ids, texts, metadatas, batch_size: int = 128) -> int:
    """Embed and upsert documents into a named collection. Returns count indexed."""
    ids, texts, metadatas = list(ids), list(texts), list(metadatas)
    col = get_collection(collection)
    for i in range(0, len(texts), batch_size):
        sl = slice(i, i + batch_size)
        col.upsert(
            ids=ids[sl],
            documents=texts[sl],
            embeddings=embed_texts(texts[sl]),
            metadatas=metadatas[sl],
        )
    return len(texts)


def count(collection: str) -> int:
    try:
        return get_collection(collection).count()
    except Exception:
        return 0


def search(collection: str, query: str, top_k: int = 5) -> list[Retrieved]:
    col = get_collection(collection)
    n = col.count()
    if n == 0:
        return []
    res = col.query(
        query_embeddings=[embed_query(query)],
        n_results=min(top_k, n),
        include=["documents", "metadatas", "distances"],
    )
    out: list[Retrieved] = []
    for doc, meta, dist in zip(res["documents"][0], res["metadatas"][0], res["distances"][0]):
        out.append(Retrieved(text=doc, score=1.0 - float(dist), metadata=meta or {}))
    return out