rag_module/vector_store.py

"""
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]