hunjing/backend/app/integrations/chromadb/client.py

"""ChromaDB 向量存储客户端。

实现 ``VectorStoreProtocol``。

内存模式（无外部依赖，本地开发默认）：``chromadb.Client()`` 自动用 SQLite 持久化到
本地文件，无需启动独立服务。生产环境配置 ``CHROMADB_HOST`` / ``CHROMADB_PORT`` 环境变量
切换为 HTTP 客户端模式。

字符级隔离：每个 character_id 对应一个 ChromaDB collection，互不干扰。
"""

from __future__ import annotations

from typing import Any

import chromadb
from chromadb.config import Settings as ChromaSettings

from app.core.config import settings
from app.core.logging import logger

# 默认持久化路径（如果未通过 env 配置则 fallback）
_DEFAULT_PERSIST_DIR = settings.chroma_persist_directory or str(
    settings.backend_dir / "../db/chroma"
)


class ChromaDBVectorStore:
    """ChromaDB 向量存储后端。

    用法::

        store = ChromaDBVectorStore()
        store.add_memory(character_id="char_xxx", memory_id="mem_xxx",
                         text="记忆内容", embedding=[0.1, 0.2, ...])
        results = store.search(character_id="char_xxx", query=[0.1, ...], limit=5)
    """

    def __init__(self) -> None:
        self._client: chromadb.Client | None = None
        self._collections: dict[str, chromadb.Collection] = {}

    # ── 生命周期 ──────────────────────────────────────────────────────────

    @property
    def client(self) -> chromadb.Client:
        if self._client is None:
            self._client = self._create_client()
        return self._client

    def _create_client(self) -> chromadb.Client:
        """创建 ChromaDB 客户端（内存模式优先）。"""
        host = settings.chromadb_host or ""
        if host:
            return chromadb.HttpClient(
                host=host,
                port=settings.chromadb_port,
            )
        # 内存模式：SQLite 持久化到本地文件
        return chromadb.Client(
            ChromaSettings(
                persist_directory=_DEFAULT_PERSIST_DIR,
                anonymized_telemetry=False,
            ),
        )

    # ── Collection ────────────────────────────────────────────────────────

    def _collection(self, character_id: str) -> chromadb.Collection:
        """按 character_id 取 collection（延迟创建）。"""
        if character_id not in self._collections:
            name = f"char_{character_id}"
            try:
                self._collections[character_id] = self.client.get_or_create_collection(
                    name=name,
                    metadata={"character_id": character_id, "hnsw:space": "cosine"},
                )
            except Exception as exc:
                logger.warning(f"ChromaDB collection 创建失败: {exc}")
                raise
        return self._collections[character_id]

    # ── 增 ────────────────────────────────────────────────────────────────

    def add_memory(
        self,
        character_id: str,
        memory_id: str,
        text: str,
        embedding: list[float],
        *,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        """添加一条记忆向量。"""
        try:
            coll = self._collection(character_id)
            coll.add(
                ids=[memory_id],
                embeddings=[embedding],
                metadatas=[{
                    "text": text[:1000],
                    **(metadata or {}),
                }],
            )
        except Exception as exc:
            logger.warning(
                "vector_store_add_failed",
                memory_id=memory_id,
                error=str(exc),
            )

    def update_memory(
        self,
        character_id: str,
        memory_id: str,
        text: str,
        embedding: list[float],
        *,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        """更新记忆向量（覆盖原有 embedding）。"""
        try:
            coll = self._collection(character_id)
            coll.update(
                ids=[memory_id],
                embeddings=[embedding],
                metadatas=[{
                    "text": text[:1000],
                    **(metadata or {}),
                }],
            )
        except Exception as exc:
            logger.warning(
                "vector_store_update_failed",
                memory_id=memory_id,
                error=str(exc),
            )

    def delete_memory(self, character_id: str, memory_id: str) -> None:
        """删除记忆向量。"""
        try:
            coll = self._collection(character_id)
            coll.delete(ids=[memory_id])
        except Exception as exc:
            logger.warning(
                "vector_store_delete_failed",
                memory_id=memory_id,
                error=str(exc),
            )

    # ── 查 ────────────────────────────────────────────────────────────────

    def search(
        self,
        character_id: str,
        query: list[float],
        *,
        limit: int = 5,
    ) -> list[dict[str, Any]]:
        """向量相似度搜索，返回 ``[{id, score, metadata}]``。"""
        try:
            coll = self._collection(character_id)
            results = coll.query(
                query_embeddings=[query],
                n_results=limit,
                include=["distances", "metadatas"],
            )
        except Exception as exc:
            logger.warning("vector_store_search_failed", error=str(exc))
            return []

        ids = results.get("ids", [[]])[0] or []
        distances = results.get("distances", [[]])[0] or []
        metadatas = results.get("metadatas", [[]])[0] or []

        return [
            {"id": mid, "score": 1.0 - dist, "metadata": meta or {}}
            for mid, dist, meta in zip(ids, distances, metadatas, strict=False)
        ]

    def count(self, character_id: str) -> int:
        """collection 中向量总数。"""
        try:
            coll = self._collection(character_id)
            return coll.count()
        except Exception:
            return 0