feat:客户agent以及记忆模块功能开发

2026-09-11 22:38:15 +08:00
parent 0197ac2ef2
commit 4da4775e8c
39 changed files with 2508 additions and 14 deletions
@@ -0,0 +1,114 @@
+"""客户长期记忆的 Milvus 向量镜像。"""
+
+from __future__ import annotations
+
+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()
+        return 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"],
+        )
+
+    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",
+]
+