File size: 2,589 Bytes
06e54e9 0c7de6c 06e54e9 0c7de6c 06e54e9 0c7de6c 06e54e9 | 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 | """
Manages the ChromaDB vector store for incident logs.
Converts CSV rows into embedded documents for semantic search.
"""
import chromadb
from chromadb.utils import embedding_functions
from pathlib import Path
import pandas as pd
import os
CHROMA_PATH = "/tmp/chroma_db" # Use /tmp for writable storage on HF Spaces
COLLECTION_NAME = "parking_incidents"
# Use a lightweight local embedding model
EMBED_MODEL = "all-MiniLM-L6-v2"
def _get_collection():
os.makedirs(CHROMA_PATH, exist_ok=True)
client = chromadb.PersistentClient(path=CHROMA_PATH)
ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=EMBED_MODEL)
return client.get_or_create_collection(name=COLLECTION_NAME, embedding_function=ef)
def row_to_text(row: dict | pd.Series) -> str:
"""Convert one incident row into a human-readable sentence for embedding."""
return (
f"At {row['timestamp']}, a {row['vehicle_class']} with plate {row['plate']} "
f"was detected at zone {row['zone']}. Status: {row['status']}. Notes: {row['notes']}."
)
def ingest_incident(row: dict):
"""Add a single new incident row to the vector store."""
col = _get_collection()
doc_id = str(row["id"])
text = row_to_text(row)
col.add(
documents=[text],
ids=[doc_id],
metadatas=[{k: str(v) for k, v in row.items()}],
)
def rebuild_from_csv(csv_path: str = "data/incidents.csv"):
"""
Re-embed the entire CSV into ChromaDB.
Call this once on startup or after bulk imports.
"""
if not Path(csv_path).exists():
print("[VectorStore] No CSV found, skipping rebuild.")
return
df = pd.read_csv(csv_path)
if df.empty:
return
col = _get_collection()
# Clear old data
existing = col.get()
if existing["ids"]:
col.delete(ids=existing["ids"])
docs, ids, metas = [], [], []
for _, row in df.iterrows():
docs.append(row_to_text(row))
ids.append(str(row["id"]))
metas.append({k: str(v) for k, v in row.items()})
col.add(documents=docs, ids=ids, metadatas=metas)
print(f"[VectorStore] Ingested {len(docs)} incidents into ChromaDB.")
def query_store(query: str, n_results: int = 5) -> list[dict]:
"""
Semantic search over incident logs.
Returns list of metadata dicts for top-n matching incidents.
"""
col = _get_collection()
results = col.query(query_texts=[query], n_results=n_results)
if not results["metadatas"] or not results["metadatas"][0]:
return []
return results["metadatas"][0]
|