128 lines
4.8 KiB
Python
128 lines
4.8 KiB
Python
|
"""记忆模块：对话历史管理"""
|
|||
|
|
"""
|
|||
|
|
memory/memory_store.py
|
|||
|
|
Agent 记忆模块：管理对话历史（短期记忆）与关键信息摘要（长期记忆）
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from collections import deque
|
|||
|
|
from dataclasses import dataclass, field
|
|||
|
|
from datetime import datetime
|
|||
|
|
from typing import Literal
|
|||
|
|
|
|||
|
|
from utils.logger import get_logger
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ── 消息数据结构 ───────────────────────────────────────────────
|
|||
|
|
@dataclass
|
|||
|
|
class Message:
|
|||
|
|
"""单条对话消息"""
|
|||
|
|
role: Literal["user", "assistant", "tool"]
|
|||
|
|
content: str
|
|||
|
|
timestamp: str = field(default_factory=lambda: datetime.now().strftime("%H:%M:%S"))
|
|||
|
|
metadata: dict = field(default_factory=dict)
|
|||
|
|
|
|||
|
|
def to_dict(self) -> dict:
|
|||
|
|
return {
|
|||
|
|
"role": self.role,
|
|||
|
|
"content": self.content,
|
|||
|
|
"timestamp": self.timestamp,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ── 记忆存储 ───────────────────────────────────────────────────
|
|||
|
|
class MemoryStore:
|
|||
|
|
"""
|
|||
|
|
对话记忆存储
|
|||
|
|
|
|||
|
|
短期记忆: 使用 deque 保存最近 N 轮对话，自动滚动淘汰旧消息
|
|||
|
|
长期记忆: 保存关键事实摘要（生产环境可替换为向量数据库）
|
|||
|
|
|
|||
|
|
使用示例:
|
|||
|
|
memory = MemoryStore(max_history=10)
|
|||
|
|
memory.add_user_message("你好")
|
|||
|
|
memory.add_assistant_message("你好！有什么可以帮你？")
|
|||
|
|
history = memory.get_history()
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, max_history: int = 20):
|
|||
|
|
"""
|
|||
|
|
Args:
|
|||
|
|
max_history: 短期记忆保留的最大消息条数
|
|||
|
|
"""
|
|||
|
|
self.logger = get_logger("MEMORY")
|
|||
|
|
self.max_history = max_history
|
|||
|
|
self._history: deque[Message] = deque(maxlen=max_history)
|
|||
|
|
self._facts: list[str] = [] # 长期记忆：关键事实
|
|||
|
|
|
|||
|
|
self.logger.info(f"💾 Memory 初始化，最大历史: {max_history} 条")
|
|||
|
|
|
|||
|
|
# ── 写入接口 ────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def add_user_message(self, content: str) -> None:
|
|||
|
|
"""记录用户消息"""
|
|||
|
|
self._add(Message(role="user", content=content))
|
|||
|
|
|
|||
|
|
def add_assistant_message(self, content: str) -> None:
|
|||
|
|
"""记录 Agent 回复"""
|
|||
|
|
self._add(Message(role="assistant", content=content))
|
|||
|
|
|
|||
|
|
def add_tool_result(self, tool_name: str, result: str) -> None:
|
|||
|
|
"""记录工具调用结果"""
|
|||
|
|
self._add(Message(
|
|||
|
|
role="tool",
|
|||
|
|
content=result,
|
|||
|
|
metadata={"tool": tool_name},
|
|||
|
|
))
|
|||
|
|
|
|||
|
|
def add_fact(self, fact: str) -> None:
|
|||
|
|
"""向长期记忆中添加关键事实"""
|
|||
|
|
self._facts.append(fact)
|
|||
|
|
self.logger.debug(f"📌 长期记忆新增: {fact}")
|
|||
|
|
|
|||
|
|
# ── 读取接口 ────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def get_history(self, last_n: int | None = None) -> list[dict]:
|
|||
|
|
"""
|
|||
|
|
获取对话历史（LLM 上下文格式）
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
last_n: 仅返回最近 N 条，None 表示全部
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
消息字典列表，格式: [{"role": ..., "content": ...}, ...]
|
|||
|
|
"""
|
|||
|
|
messages = list(self._history)
|
|||
|
|
if last_n:
|
|||
|
|
messages = messages[-last_n:]
|
|||
|
|
return [m.to_dict() for m in messages]
|
|||
|
|
|
|||
|
|
def get_facts(self) -> list[str]:
|
|||
|
|
"""获取所有长期记忆事实"""
|
|||
|
|
return list(self._facts)
|
|||
|
|
|
|||
|
|
def get_context_summary(self) -> str:
|
|||
|
|
"""生成上下文摘要字符串，供 LLM Prompt 使用"""
|
|||
|
|
history = self.get_history(last_n=6)
|
|||
|
|
lines = [f"[{m['role'].upper()}] {m['content'][:80]}" for m in history]
|
|||
|
|
return "\n".join(lines) if lines else "（暂无对话历史）"
|
|||
|
|
|
|||
|
|
# ── 管理接口 ────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def clear_history(self) -> None:
|
|||
|
|
"""清空短期对话历史"""
|
|||
|
|
self._history.clear()
|
|||
|
|
self.logger.info("🗑 对话历史已清空")
|
|||
|
|
|
|||
|
|
def stats(self) -> dict:
|
|||
|
|
"""返回记忆统计信息"""
|
|||
|
|
return {
|
|||
|
|
"history_count": len(self._history),
|
|||
|
|
"facts_count": len(self._facts),
|
|||
|
|
"max_history": self.max_history,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# ── 私有方法 ────────────────────────────────────────────────
|
|||
|
|
|
|||
|
|
def _add(self, message: Message) -> None:
|
|||
|
|
self._history.append(message)
|
|||
|
|
self.logger.debug(f"💬 [{message.role.upper()}] {message.content[:60]}...")
|