128 lines
3.9 KiB
Python
128 lines
3.9 KiB
Python
|
"""
|
|||
|
|
用户会话记忆持久化：与 Cache 节点 user_memory_* 键对齐，写入 MySQL，Redis 仍作热缓存。
|
|||
|
|
"""
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import logging
|
|||
|
|
import uuid
|
|||
|
|
from typing import Any, Dict, Optional, Tuple
|
|||
|
|
|
|||
|
|
from sqlalchemy.orm import Session
|
|||
|
|
|
|||
|
|
from app.core.config import settings
|
|||
|
|
from app.models.persistent_user_memory import PersistentUserMemory
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
USER_MEMORY_PREFIX = "user_memory_"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def parse_memory_scope(workflow_id: str) -> Tuple[Optional[str], Optional[str]]:
|
|||
|
|
"""
|
|||
|
|
Celery 任务里 workflow_id 对 Agent 为 agent_{uuid}，对工作流为 UUID 字符串。
|
|||
|
|
返回 (scope_kind, scope_id) 或 (None, None) 表示不做 DB 持久化。
|
|||
|
|
"""
|
|||
|
|
if not workflow_id:
|
|||
|
|
return None, None
|
|||
|
|
if workflow_id.startswith("agent_") and len(workflow_id) > 6:
|
|||
|
|
return "agent", workflow_id[6:]
|
|||
|
|
# 标准 UUID
|
|||
|
|
s = workflow_id.strip()
|
|||
|
|
if len(s) == 36 and s.count("-") == 4:
|
|||
|
|
return "workflow", s
|
|||
|
|
return None, None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def is_user_memory_redis_key(key: str) -> bool:
|
|||
|
|
return isinstance(key, str) and key.startswith(USER_MEMORY_PREFIX)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def session_key_from_user_memory_key(key: str) -> Optional[str]:
|
|||
|
|
if not is_user_memory_redis_key(key):
|
|||
|
|
return None
|
|||
|
|
rest = key[len(USER_MEMORY_PREFIX) :]
|
|||
|
|
return rest if rest else None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def merge_memory_payloads(
|
|||
|
|
base: Optional[Dict[str, Any]], overlay: Optional[Dict[str, Any]]
|
|||
|
|
) -> Optional[Dict[str, Any]]:
|
|||
|
|
"""合并 DB 与 Redis：对话历史取更长的一条；画像/上下文做浅合并。"""
|
|||
|
|
if not base and not overlay:
|
|||
|
|
return None
|
|||
|
|
if not base:
|
|||
|
|
return dict(overlay) if isinstance(overlay, dict) else None
|
|||
|
|
if not overlay:
|
|||
|
|
return dict(base)
|
|||
|
|
out = dict(base)
|
|||
|
|
if not isinstance(overlay, dict):
|
|||
|
|
return out
|
|||
|
|
for k, v in overlay.items():
|
|||
|
|
if k == "conversation_history" and isinstance(v, list) and isinstance(out.get(k), list):
|
|||
|
|
if len(v) > len(out[k]):
|
|||
|
|
out[k] = v
|
|||
|
|
elif k == "user_profile" and isinstance(v, dict):
|
|||
|
|
out[k] = {**(out.get(k) or {}), **v}
|
|||
|
|
elif k == "context" and isinstance(v, dict):
|
|||
|
|
out[k] = {**(out.get(k) or {}), **v}
|
|||
|
|
else:
|
|||
|
|
out[k] = v
|
|||
|
|
return out
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_persistent_memory(
|
|||
|
|
db: Session, scope_kind: str, scope_id: str, session_key: str
|
|||
|
|
) -> Optional[Dict[str, Any]]:
|
|||
|
|
row = (
|
|||
|
|
db.query(PersistentUserMemory)
|
|||
|
|
.filter(
|
|||
|
|
PersistentUserMemory.scope_kind == scope_kind,
|
|||
|
|
PersistentUserMemory.scope_id == scope_id,
|
|||
|
|
PersistentUserMemory.session_key == session_key,
|
|||
|
|
)
|
|||
|
|
.first()
|
|||
|
|
)
|
|||
|
|
if not row or not isinstance(row.payload, dict):
|
|||
|
|
return None
|
|||
|
|
return dict(row.payload)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def save_persistent_memory(
|
|||
|
|
db: Session, scope_kind: str, scope_id: str, session_key: str, payload: Dict[str, Any]
|
|||
|
|
) -> None:
|
|||
|
|
row = (
|
|||
|
|
db.query(PersistentUserMemory)
|
|||
|
|
.filter(
|
|||
|
|
PersistentUserMemory.scope_kind == scope_kind,
|
|||
|
|
PersistentUserMemory.scope_id == scope_id,
|
|||
|
|
PersistentUserMemory.session_key == session_key,
|
|||
|
|
)
|
|||
|
|
.first()
|
|||
|
|
)
|
|||
|
|
if row:
|
|||
|
|
row.payload = payload
|
|||
|
|
else:
|
|||
|
|
row = PersistentUserMemory(
|
|||
|
|
id=str(uuid.uuid4()),
|
|||
|
|
scope_kind=scope_kind,
|
|||
|
|
scope_id=scope_id,
|
|||
|
|
session_key=session_key,
|
|||
|
|
payload=payload,
|
|||
|
|
)
|
|||
|
|
db.add(row)
|
|||
|
|
db.commit()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def delete_persistent_memory(db: Session, scope_kind: str, scope_id: str, session_key: str) -> None:
|
|||
|
|
q = db.query(PersistentUserMemory).filter(
|
|||
|
|
PersistentUserMemory.scope_kind == scope_kind,
|
|||
|
|
PersistentUserMemory.scope_id == scope_id,
|
|||
|
|
PersistentUserMemory.session_key == session_key,
|
|||
|
|
)
|
|||
|
|
q.delete()
|
|||
|
|
db.commit()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def persist_enabled() -> bool:
|
|||
|
|
return bool(getattr(settings, "MEMORY_PERSIST_DB_ENABLED", True))
|