145 lines
5.3 KiB
Python
145 lines
5.3 KiB
Python
"""客户长期记忆的 Milvus 向量镜像。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from pymilvus import AsyncMilvusClient, DataType
|
|
|
|
from config.database.milvus import client as configured_client
|
|
from rag.embedding import EMBEDDING_DIMENSION
|
|
|
|
|
|
CUSTOMER_MEMORY_COLLECTION = "customer_memory"
|
|
|
|
|
|
def build_memory_schema():
|
|
"""构造客户记忆向量集合结构，维度来自 LLM_EMBED_DIMENSIONS。"""
|
|
schema = AsyncMilvusClient.create_schema(auto_id=False, enable_dynamic_field=False)
|
|
schema.add_field("memory_id", DataType.VARCHAR, is_primary=True, max_length=128)
|
|
schema.add_field("customer_id", DataType.INT64)
|
|
schema.add_field("memory_type", DataType.VARCHAR, max_length=32)
|
|
schema.add_field("tag", DataType.VARCHAR, max_length=64)
|
|
schema.add_field("content", DataType.VARCHAR, max_length=2048)
|
|
schema.add_field("status", DataType.VARCHAR, max_length=16)
|
|
schema.add_field("vector", DataType.FLOAT_VECTOR, dim=EMBEDDING_DIMENSION)
|
|
return schema
|
|
|
|
|
|
def build_memory_index_params():
|
|
"""构造客户记忆向量索引。"""
|
|
params = AsyncMilvusClient.prepare_index_params()
|
|
params.add_index(
|
|
field_name="vector",
|
|
index_type="HNSW",
|
|
metric_type="COSINE",
|
|
params={"M": 16, "efConstruction": 200},
|
|
)
|
|
return params
|
|
|
|
|
|
class MilvusMemoryStore:
|
|
"""封装客户记忆向量写入、查询和删除。"""
|
|
|
|
def __init__(self, client=None, *, collection_name=CUSTOMER_MEMORY_COLLECTION):
|
|
self.client = client or configured_client()
|
|
self.collection_name = collection_name
|
|
|
|
async def ensure_collection(self) -> None:
|
|
"""创建集合或校验已有集合的向量维度。"""
|
|
if not await self.client.has_collection(self.collection_name):
|
|
await self.client.create_collection(
|
|
collection_name=self.collection_name,
|
|
schema=build_memory_schema(),
|
|
index_params=build_memory_index_params(),
|
|
)
|
|
return
|
|
desc = await self.client.describe_collection(self.collection_name)
|
|
for field in desc.get("fields", []):
|
|
if field.get("name") == "vector":
|
|
dim = field.get("params", {}).get("dim")
|
|
if dim is not None and int(dim) != EMBEDDING_DIMENSION:
|
|
raise RuntimeError(
|
|
f"Milvus collection {self.collection_name!r} vector dim={dim}, "
|
|
f"expected {EMBEDDING_DIMENSION}"
|
|
)
|
|
|
|
async def upsert(self, memory, vector: list[float]) -> str:
|
|
"""写入一条客户记忆向量并返回 Milvus 主键。"""
|
|
await self.ensure_collection()
|
|
memory_id = str(memory.id)
|
|
await self.client.delete(
|
|
collection_name=self.collection_name,
|
|
filter=f'memory_id == "{memory_id}"',
|
|
)
|
|
await self.client.insert(
|
|
collection_name=self.collection_name,
|
|
data=[
|
|
{
|
|
"memory_id": memory_id,
|
|
"customer_id": memory.customer_id,
|
|
"memory_type": memory.memory_type,
|
|
"tag": memory.tag,
|
|
"content": memory.content,
|
|
"status": memory.status,
|
|
"vector": vector,
|
|
}
|
|
],
|
|
)
|
|
return memory_id
|
|
|
|
async def search(self, vector: list[float], customer_id: int, *, limit: int = 10) -> list[dict]:
|
|
"""按客户 ID 过滤向量查询，返回归一化的命中列表。"""
|
|
await self.ensure_collection()
|
|
raw = await self.client.search(
|
|
collection_name=self.collection_name,
|
|
data=[vector],
|
|
limit=limit,
|
|
filter=f"customer_id == {int(customer_id)}",
|
|
output_fields=["memory_id", "customer_id", "memory_type", "tag", "content", "status"],
|
|
)
|
|
return self._normalize_hits(raw)
|
|
|
|
@staticmethod
|
|
def _normalize_hits(raw: Any) -> list[dict]:
|
|
"""将 pymilvus 返回结构收敛为 memory_id + distance 的扁平列表。"""
|
|
hits: list[dict] = []
|
|
for batch in raw or []:
|
|
for hit in batch or []:
|
|
if not isinstance(hit, dict):
|
|
continue
|
|
entity = hit.get("entity") or {}
|
|
memory_id = entity.get("memory_id") or hit.get("id")
|
|
if memory_id is None:
|
|
continue
|
|
try:
|
|
distance = float(hit.get("distance", 0.0))
|
|
except (TypeError, ValueError):
|
|
distance = 0.0
|
|
hits.append(
|
|
{
|
|
"memory_id": str(memory_id),
|
|
"distance": max(0.0, min(1.0, distance)),
|
|
"tag": entity.get("tag"),
|
|
"content": entity.get("content"),
|
|
"status": entity.get("status"),
|
|
}
|
|
)
|
|
return hits
|
|
|
|
async def delete(self, memory_id: int | str) -> None:
|
|
"""删除一条客户记忆向量。"""
|
|
await self.client.delete(
|
|
collection_name=self.collection_name,
|
|
filter=f'memory_id == "{memory_id}"',
|
|
)
|
|
|
|
|
|
__all__ = [
|
|
"CUSTOMER_MEMORY_COLLECTION",
|
|
"MilvusMemoryStore",
|
|
"build_memory_index_params",
|
|
"build_memory_schema",
|
|
]
|
|
|