85 lines
2.2 KiB
Python
85 lines
2.2 KiB
Python
"""向量存储接口与工厂。
|
|
|
|
通过 ``get_vector_store()`` 获取当前配置的后端实现：
|
|
- ``chromadb``（默认）：ChromaDB，内存/HTTP 双模式
|
|
- ``memory``：纯内存实现（测试用）
|
|
|
|
通过环境变量 ``CHROMADB_HOST`` / ``CHROMADB_PORT`` 控制 ChromaDB 连接模式。
|
|
通过 ``ENABLE_EMBEDDING=false`` 关闭 embedding（纯 SQL 检索）。
|
|
通过 ``VECTOR_DB_TYPE=chroma`` 选择向量库后端。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Protocol, runtime_checkable
|
|
|
|
|
|
@runtime_checkable
|
|
class VectorStoreProtocol(Protocol):
|
|
"""向量存储接口。所有后端实现必须遵循此协议。"""
|
|
|
|
def add_memory(
|
|
self,
|
|
character_id: str,
|
|
memory_id: str,
|
|
text: str,
|
|
embedding: list[float],
|
|
*,
|
|
metadata: dict[str, Any] | None = None,
|
|
) -> None: ...
|
|
|
|
def update_memory(
|
|
self,
|
|
character_id: str,
|
|
memory_id: str,
|
|
text: str,
|
|
embedding: list[float],
|
|
*,
|
|
metadata: dict[str, Any] | None = None,
|
|
) -> None: ...
|
|
|
|
def delete_memory(self, character_id: str, memory_id: str) -> None: ...
|
|
|
|
def search(
|
|
self,
|
|
character_id: str,
|
|
query: list[float],
|
|
*,
|
|
limit: int = 5,
|
|
) -> list[dict[str, Any]]: ...
|
|
|
|
def count(self, character_id: str) -> int: ...
|
|
|
|
|
|
_vector_store: VectorStoreProtocol | None = None
|
|
|
|
|
|
def get_vector_store() -> VectorStoreProtocol | None:
|
|
"""根据配置返回向量存储后端；embedding 关闭时返回 None。"""
|
|
from app.core.config import settings
|
|
|
|
if not settings.vector_db_enabled:
|
|
return None
|
|
|
|
global _vector_store # noqa: PLW0603
|
|
if _vector_store is not None:
|
|
return _vector_store
|
|
|
|
backend = settings.vector_db_type
|
|
if backend == "memory":
|
|
from app.integrations.chromadb.memory_store import MemoryVectorStore
|
|
|
|
_vector_store = MemoryVectorStore()
|
|
else:
|
|
# 默认 chromadb
|
|
from app.integrations.chromadb.client import ChromaDBVectorStore
|
|
|
|
_vector_store = ChromaDBVectorStore()
|
|
return _vector_store
|
|
|
|
|
|
def reset_vector_store() -> None:
|
|
"""重置单例（测试用）。"""
|
|
global _vector_store # noqa: PLW0603
|
|
_vector_store = None
|