services/chroma_memory.py

"""ChromaDB vector memory for incident similarity (optional)."""

from __future__ import annotations

import os
from typing import Any, Optional

_client: Any = None


def get_collection():
    global _client
    if _client is not None:
        return _client
    try:
        import chromadb  # type: ignore
    except ImportError:
        return None
    path = os.getenv("CHROMA_PERSIST_DIR", "./chroma_data")
    client = chromadb.PersistentClient(path=path)
    _client = client.get_or_create_collection("sentinel_incidents")
    return _client


def remember_incident(text: str, metadata: dict[str, Any]) -> None:
    col = get_collection()
    if col is None:
        return
    import uuid

    col.add(ids=[str(uuid.uuid4())], documents=[text], metadatas=[metadata])


def similar_incidents(text: str, k: int = 3) -> list[dict[str, Any]]:
    col = get_collection()
    if col is None:
        return []
    res = col.query(query_texts=[text], n_results=k)
    out: list[dict[str, Any]] = []
    for md, dist in zip(res.get("metadatas", [[]])[0] or [], res.get("distances", [[]])[0] or []):
        out.append({"metadata": md, "distance": dist})
    return out