File size: 6,607 Bytes
35676b4 7880373 35676b4 7880373 35676b4 7880373 35676b4 7880373 35676b4 7880373 35676b4 7880373 35676b4 7880373 35676b4 7880373 35676b4 7880373 35676b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 | 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)
|