初始化提交
This commit is contained in:
chenxin committed
commit
887c1f449e
963 files changed
+166869
No files matched your search
@@ -0,0 +1,187 @@
|
||||
"""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
|
||||
Reference in new issue
Block a user