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,107 @@
+"""memory_unit 长期记忆仓储。"""
+
+from __future__ import annotations
+
+from datetime import datetime
+
+from sqlalchemy import select
+
+from model.memory_unit import MemoryUnit
+from repositories.base import BaseRepository
+
+
+ACTIVE_STATUSES = ("candidate", "confirmed")
+INACTIVE_STATUSES = ("expired", "rejected", "archived")
+
+
+class MemoryUnitRepo(BaseRepository):
+    """负责 MySQL 主体记忆的隔离、去重和同步状态管理。"""
+
+    model = MemoryUnit
+
+    async def add_memory(self, values: dict) -> MemoryUnit:
+        """先写入 MySQL 主体并返回数据库 ID。"""
+        obj = MemoryUnit(**values)
+        self.db.add(obj)
+        await self.db.commit()
+        await self.db.refresh(obj)
+        return obj
+
+    async def find_exact(
+        self, customer_id: int, memory_type: str, tag: str, content: str
+    ) -> MemoryUnit | None:
+        """按客户、类型、标签和内容精确查找去重对象。"""
+        return await self.db.scalar(
+            select(MemoryUnit).where(
+                MemoryUnit.customer_id == customer_id,
+                MemoryUnit.memory_type == memory_type,
+                MemoryUnit.tag == tag,
+                MemoryUnit.content == content,
+                MemoryUnit.status.not_in(INACTIVE_STATUSES),
+            )
+        )
+
+    async def list_for_customer(
+        self,
+        customer_id: int,
+        *,
+        memory_type: str | None = None,
+        tag: str | None = None,
+        now: datetime | None = None,
+        limit: int = 100,
+    ) -> list[MemoryUnit]:
+        """查询默认可召回的客户记忆，过滤拒绝、归档和过期记录。"""
+        now = now or datetime.now()
+        conditions = [
+            MemoryUnit.customer_id == customer_id,
+            MemoryUnit.status.in_(ACTIVE_STATUSES),
+            (MemoryUnit.valid_until.is_(None) | (MemoryUnit.valid_until > now)),
+        ]
+        if memory_type:
+            conditions.append(MemoryUnit.memory_type == memory_type)
+        if tag:
+            conditions.append(MemoryUnit.tag == tag)
+        statement = (
+            select(MemoryUnit)
+            .where(*conditions)
+            .order_by(MemoryUnit.update_time.desc(), MemoryUnit.id.desc())
+            .limit(limit)
+        )
+        return list((await self.db.scalars(statement)).all())
+
+    async def merge_evidence(
+        self, memory: MemoryUnit, *, evidence_count: int = 1, conflict_count: int = 0
+    ) -> MemoryUnit:
+        """合并证据和冲突计数，不改变客户隔离范围。"""
+        memory.evidence_count = (memory.evidence_count or 0) + evidence_count
+        memory.conflict_count = (memory.conflict_count or 0) + conflict_count
+        memory.update_time = datetime.now()
+        await self.db.commit()
+        await self.db.refresh(memory)
+        return memory
+
+    async def update_sync_status(self, memory_id: int, **values) -> None:
+        """更新向量/图谱索引 ID 和同步状态。"""
+        memory = await self.db.get(MemoryUnit, memory_id)
+        if memory is None:
+            return
+        for key, value in values.items():
+            setattr(memory, key, value)
+        memory.update_time = datetime.now()
+        await self.db.commit()
+
+    async def list_pending_sync(self, limit: int = 100) -> list[MemoryUnit]:
+        """返回存在失败或空索引的记忆，供重试入口使用。"""
+        statement = (
+            select(MemoryUnit)
+            .where(
+                (MemoryUnit.milvus_id.is_(None) | (MemoryUnit.milvus_sync_status == "failed"))
+                | (MemoryUnit.graph_node_id.is_(None) | (MemoryUnit.neo4j_sync_status == "failed"))
+            )
+            .order_by(MemoryUnit.id)
+            .limit(limit)
+        )
+        return list((await self.db.scalars(statement)).all())
+
+
+__all__ = ["MemoryUnitRepo"]