Spaces:
Runtime error
Runtime error
File size: 5,057 Bytes
d4ec2fa | 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 | """
Lightweight in-memory vector search — replaces chromadb (Ticket 7's original choice).
Why the change: chromadb pulls in a huge dependency tree (kubernetes client,
opentelemetry, onnxruntime, grpcio, aiohttp...) meant for running a full client-server
vector database. At our actual scale (2,736 vectors x 384 dims = ~4MB), that's wildly
disproportionate -- it was the direct cause of Render's free-tier 512MB build running
out of memory. Brute-force cosine similarity via numpy does the identical job in
~2 milliseconds (measured), with zero extra dependencies.
Loads once at startup from the same embeddings.npy / report_ids.json Ticket 7 already
produced -- no data or embeddings were regenerated, only how they're queried changed.
"""
import json
import logging
import sqlite3
import numpy as np
from api.config import BASE_DIR, DATABASE_URL
logger = logging.getLogger(__name__)
EMBEDDINGS_PATH = BASE_DIR / "data" / "processed" / "embeddings.npy"
REPORT_IDS_PATH = BASE_DIR / "data" / "processed" / "report_ids.json"
class VectorStore:
def __init__(self):
self.embeddings = None # (N, 384) float32, L2-normalized
self.report_ids = None # list[str], same order as embeddings rows
self.metadata = {} # report_id -> {narrative_text, ata_chapter_label, severity_label, ...}
self.loaded = False
self.load_error = None
def load(self):
try:
self.embeddings = np.load(EMBEDDINGS_PATH).astype(np.float32)
with open(REPORT_IDS_PATH) as f:
self.report_ids = [str(x) for x in json.load(f)]
db_path = DATABASE_URL.replace("sqlite:///", "")
conn = sqlite3.connect(db_path)
rows = conn.execute("""
SELECT r.asrs_report_id, r.narrative_text, g.ata_chapter_label, g.severity_label
FROM reports r LEFT JOIN gold_labels g ON r.asrs_report_id = g.asrs_report_id
""").fetchall()
conn.close()
self.metadata = {
rid: {"narrative_text": text, "ata_chapter_label": ata, "severity_label": sev}
for rid, text, ata, sev in rows
}
self.loaded = True
logger.info(f"Vector store loaded: {len(self.report_ids)} vectors")
except Exception as e:
self.loaded = False
self.load_error = str(e)
logger.error(f"Vector store load failed: {e}")
def query(self, query_embedding, system=None, severity=None, limit=20):
if not self.loaded:
raise RuntimeError(self.load_error or "Vector store not loaded")
q = np.asarray(query_embedding, dtype=np.float32)
q = q / (np.linalg.norm(q) + 1e-9)
scores = self.embeddings @ q # cosine similarity, since both sides are L2-normalized
# Apply metadata filters BEFORE ranking, same semantics as the old Chroma `where` clause
candidate_idx = np.arange(len(self.report_ids))
if system or severity:
keep = []
for i in candidate_idx:
meta = self.metadata.get(self.report_ids[i], {})
if system and meta.get("ata_chapter_label") != system:
continue
if severity and meta.get("severity_label") != severity:
continue
keep.append(i)
candidate_idx = np.array(keep, dtype=int)
if len(candidate_idx) == 0:
return []
candidate_scores = scores[candidate_idx]
top_n = min(limit, len(candidate_idx))
top_order = np.argsort(-candidate_scores)[:top_n]
top_idx = candidate_idx[top_order]
results = []
for i in top_idx:
rid = self.report_ids[i]
meta = self.metadata.get(rid, {})
results.append({
"report_id": rid,
"excerpt": (meta.get("narrative_text") or "")[:300],
"ata_chapter": meta.get("ata_chapter_label"),
"severity": meta.get("severity_label"),
"score": round(float(scores[i]), 4),
})
return results
def filter_only(self, system=None, severity=None, limit=20):
"""No query vector at all -- just metadata filtering, no ranking."""
if not self.loaded:
raise RuntimeError(self.load_error or "Vector store not loaded")
results = []
for rid, meta in self.metadata.items():
if system and meta.get("ata_chapter_label") != system:
continue
if severity and meta.get("severity_label") != severity:
continue
results.append({
"report_id": rid,
"excerpt": (meta.get("narrative_text") or "")[:300],
"ata_chapter": meta.get("ata_chapter_label"),
"severity": meta.get("severity_label"),
"score": None,
})
if len(results) >= limit:
break
return results
vector_store = VectorStore()
|