108 lines
3.7 KiB
Python
108 lines
3.7 KiB
Python
"""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"]
|