Spaces:
Running
Running
File size: 1,170 Bytes
8b3905d | 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 | """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
|