src/storage/vector_db.py - develop - dsboikov/aiBoardRoom

src/storage/vector_db.py
129 строк4 KB
Денис
Project updated
03 авг 2026, 12:39
03 авг 2026, 12:39d7da468
"""Local ChromaDB vector store for room context."""

from __future__ import annotations

import hashlib
from typing import Any

import chromadb
from chromadb.config import Settings

from config import CHROMA_PATH, ETL_TOP_K
from utils.logger import get_logger

logger = get_logger("vector_db")


class VectorStore:
    """Persistent local ChromaDB with one collection per room."""

    def __init__(self, path: str | None = None) -> None:
        self.path = str(path or CHROMA_PATH)
        self._client = chromadb.PersistentClient(
            path=self.path,
            settings=Settings(anonymized_telemetry=False, allow_reset=True),
        )

    def _collection_name(self, room_id: int) -> str:
        return f"room_{room_id}"

    def get_or_create_collection(self, room_id: int):
        return self._client.get_or_create_collection(
            name=self._collection_name(room_id),
            metadata={"room_id": str(room_id)},
        )

    @staticmethod
    def _doc_id(room_id: int, text: str, source: str) -> str:
        digest = hashlib.sha256(f"{room_id}:{source}:{text}".encode("utf-8")).hexdigest()[:24]
        return f"{room_id}_{digest}"

    def add_document(
        self,
        room_id: int,
        text: str,
        *,
        source: str = "message",
        metadata: dict[str, Any] | None = None,
    ) -> str:
        """Index a text chunk for a room. Returns document id."""
        if not text or not text.strip():
            return ""
        collection = self.get_or_create_collection(room_id)
        doc_id = self._doc_id(room_id, text, source)
        meta = {"source": source, "room_id": room_id, **(metadata or {})}
        # Chroma metadata values must be scalars
        clean_meta = {k: v for k, v in meta.items() if isinstance(v, (str, int, float, bool))}
        collection.upsert(ids=[doc_id], documents=[text.strip()], metadatas=[clean_meta])
        logger.debug("Indexed doc %s for room %s (%s chars)", doc_id, room_id, len(text))
        return doc_id

    def add_documents(
        self,
        room_id: int,
        texts: list[str],
        *,
        source: str = "ingest",
        metadatas: list[dict[str, Any]] | None = None,
    ) -> list[str]:
        ids: list[str] = []
        for i, text in enumerate(texts):
            meta = metadatas[i] if metadatas and i < len(metadatas) else None
            doc_id = self.add_document(room_id, text, source=source, metadata=meta)
            if doc_id:
                ids.append(doc_id)
        return ids

    def query(
        self,
        room_id: int,
        query_text: str,
        *,
        top_k: int | None = None,
    ) -> list[dict[str, Any]]:
        """Return relevant documents for a query."""
        collection = self.get_or_create_collection(room_id)
        n = collection.count()
        if n == 0 or not query_text.strip():
            return []
        k = min(top_k or ETL_TOP_K, n)
        result = collection.query(query_texts=[query_text], n_results=k)
        docs = result.get("documents", [[]])[0] or []
        metas = result.get("metadatas", [[]])[0] or []
        distances = result.get("distances", [[]])[0] or []
        ids = result.get("ids", [[]])[0] or []
        out: list[dict[str, Any]] = []
        for i, doc in enumerate(docs):
            out.append(
                {
                    "id": ids[i] if i < len(ids) else "",
                    "text": doc,
                    "metadata": metas[i] if i < len(metas) else {},
                    "distance": distances[i] if i < len(distances) else None,
                }
            )
        return out

    def clear_room(self, room_id: int) -> None:
        name = self._collection_name(room_id)
        try:
            self._client.delete_collection(name)
            logger.info("Cleared vector collection %s", name)
        except Exception:  # noqa: BLE001
            logger.debug("Collection %s already absent", name)

    def delete_room(self, room_id: int) -> None:
        self.clear_room(room_id)


_store: VectorStore | None = None


def get_vector_store(path: str | None = None) -> VectorStore:
    """Return process-wide VectorStore singleton."""
    global _store
    if path is not None:
        return VectorStore(path)
    if _store is None:
        _store = VectorStore()
    return _store