storage/vector_store.py

import chromadb
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
from pathlib import Path
from typing import Optional

from agent.evidence import make_evidence_record

CHROMA_PATH = Path("data/chroma")
_ef = SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")

RERANK_POOL_MULTIPLIER = 5
RERANK_POOL_FLOOR = 20
_COLLECTION_SOURCES = {
    "filings": {"10-K", "10-Q"},
    "transcripts": {"transcript"},
}


def _client() -> chromadb.PersistentClient:
    CHROMA_PATH.mkdir(parents=True, exist_ok=True)
    return chromadb.PersistentClient(path=str(CHROMA_PATH))


def _collection_missing(exc: Exception) -> bool:
    """Recognise Chroma's version-dependent missing-collection errors only."""
    name = type(exc).__name__.casefold()
    message = str(exc).casefold()
    return (
        "notfound" in name
        or "invalidcollection" in name
        or "does not exist" in message
        or "not found" in message
    )


def add_chunks(
    collection_name: str,
    chunks: list[str],
    metadatas: list[dict],
    ids: list[str],
) -> None:
    if not (len(chunks) == len(metadatas) == len(ids)):
        raise ValueError("chunks, metadatas, and ids must have the same length")
    enriched: list[dict] = []
    allowed_sources = _COLLECTION_SOURCES.get(collection_name)
    if allowed_sources is None:
        raise ValueError(f"Unsupported evidence collection: {collection_name!r}")
    for chunk, metadata, chunk_id in zip(chunks, metadatas, ids):
        meta = dict(metadata)
        source = meta.get("source") or ("transcript" if collection_name == "transcripts" else None)
        if source not in allowed_sources:
            raise ValueError(
                f"Source {source!r} is invalid for collection {collection_name!r}"
            )
        document_id = str(meta.get("document_id") or ":".join(filter(None, [
            collection_name,
            str(meta.get("ticker", "")),
            str(source),
            str(meta.get("period", "")),
            str(meta.get("filing_date") or meta.get("date") or ""),
            str(meta.get("section", "")),
        ])))
        record = make_evidence_record(
            source=source,
            content=chunk,
            document_id=document_id,
            chunk_id=str(meta.get("chunk_id") or chunk_id),
            source_url=meta.get("source_url"),
            as_of=meta.get("filing_date") or meta.get("date"),
        )
        meta.update({
            "source": record.ref.source,
            "document_id": record.ref.document_id,
            "chunk_id": record.ref.chunk_id or str(chunk_id),
            "content_hash": record.ref.content_hash,
            "evidence_id": record.ref.evidence_id,
        })
        enriched.append(meta)
    col = _client().get_or_create_collection(name=collection_name, embedding_function=_ef)
    col.upsert(documents=chunks, metadatas=enriched, ids=ids)


def delete_by_ticker(collection_name: str, ticker: str) -> None:
    """Delete all chunks belonging to a ticker from the given collection."""
    try:
        col = _client().get_collection(name=collection_name, embedding_function=_ef)
    except Exception as exc:
        if _collection_missing(exc):
            return
        raise RuntimeError(
            f"Unable to open Chroma collection {collection_name!r} for deletion"
        ) from exc
    results = col.get(where={"ticker": ticker.upper()})
    if results["ids"]:
        col.delete(ids=results["ids"])


def search(
    collection_name: str,
    query: str,
    ticker: str,
    n_results: int = 3,
    min_filing_date: Optional[str] = None,
    period: Optional[str] = None,
) -> list[dict]:
    allowed_sources = _COLLECTION_SOURCES.get(collection_name)
    if allowed_sources is None:
        raise ValueError(f"Unsupported evidence collection: {collection_name!r}")
    try:
        col = _client().get_collection(name=collection_name, embedding_function=_ef)
    except Exception as exc:
        if _collection_missing(exc):
            return []
        raise RuntimeError(
            f"Unable to open Chroma collection {collection_name!r}; storage may be unavailable or corrupt"
        ) from exc

    # Build where clause — period equality is exact, date filtering is done in Python
    where: dict = {"ticker": ticker.upper()}
    if period:
        where = {"$and": [{"ticker": {"$eq": ticker.upper()}}, {"period": {"$eq": period}}]}

    # Over-fetch then re-rank: 5x always, so the cross-encoder has real choice
    # even when no date filter is applied.
    fetch_limit = max(n_results * RERANK_POOL_MULTIPLIER, RERANK_POOL_FLOOR)
    results = col.query(
        query_texts=[query],
        n_results=fetch_limit,
        where=where,
    )
    if not results["documents"][0]:
        return []

    result_list = []
    for doc, raw_meta in zip(results["documents"][0], results["metadatas"][0]):
        meta = dict(raw_meta)
        source = meta.get("source") or ("transcript" if collection_name == "transcripts" else None)
        if source not in allowed_sources:
            continue

        # Only chunks written through add_chunks have a content-addressed
        # provenance tuple.  Re-deriving a fresh ref for legacy rows would make
        # unverifiable historical content look newly verified.
        if not all(meta.get(key) for key in (
            "document_id", "chunk_id", "content_hash", "evidence_id",
        )):
            continue
        record = make_evidence_record(
            source=source,
            content=doc,
            document_id=str(meta["document_id"]),
            chunk_id=str(meta["chunk_id"]),
            source_url=meta.get("source_url"),
            as_of=meta.get("filing_date") or meta.get("date"),
            metadata=meta,
        )
        if (
            meta.get("content_hash") != record.ref.content_hash
            or meta.get("evidence_id") != record.ref.evidence_id
        ):
            continue
        result_list.append({
            "text": record.content,
            "metadata": meta,
            "evidence_ref": record.ref.model_dump(mode="json"),
        })

    # Apply date filter in Python if specified
    if min_filing_date:
        # Support both "filing_date" (for filings) and "date" (for transcripts)
        result_list = [
            r for r in result_list
            if (r["metadata"].get("filing_date", "") >= min_filing_date or
                r["metadata"].get("date", "") >= min_filing_date)
        ]

    from storage.reranker import rerank
    return rerank(query, result_list, top_k=n_results)