118 lines
4.6 KiB
Python
118 lines
4.6 KiB
Python
"""MemoryTool — Agent 可在对话中读写记忆的工具.
|
|
|
|
操作:
|
|
- add: 追加内容到指定 section
|
|
- replace: 替换 section 内的文本
|
|
- remove: 删除整个 section
|
|
- read: 读取文件内容
|
|
|
|
file 参数: soul | user | memory | daily
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from agentkit.memory.profile import MemoryStore
|
|
from agentkit.tools.base import Tool
|
|
|
|
|
|
VALID_FILES = {"soul", "user", "memory", "daily"}
|
|
VALID_ACTIONS = {"add", "replace", "remove", "read"}
|
|
|
|
|
|
class MemoryTool(Tool):
|
|
"""Agent 可调用的记忆操作工具.
|
|
|
|
让 Agent 在对话中读写 SOUL/USER/MEMORY/DAILY 记忆文件。
|
|
"""
|
|
|
|
def __init__(self, memory_store: MemoryStore):
|
|
super().__init__(
|
|
name="memory",
|
|
description="Read and write persistent memory files. Use to remember user preferences, project info, and notes across sessions.",
|
|
input_schema={
|
|
"type": "object",
|
|
"properties": {
|
|
"action": {
|
|
"type": "string",
|
|
"enum": list(VALID_ACTIONS),
|
|
"description": "Operation: add, replace, remove, read",
|
|
},
|
|
"file": {
|
|
"type": "string",
|
|
"enum": list(VALID_FILES),
|
|
"description": "Memory file: soul (agent identity), user (user profile), memory (work notes), daily (today's log)",
|
|
},
|
|
"section": {
|
|
"type": "string",
|
|
"description": "Section name within the file (e.g. '项目信息', '偏好')",
|
|
},
|
|
"content": {
|
|
"type": "string",
|
|
"description": "Content to add or new text for replace",
|
|
},
|
|
"old_text": {
|
|
"type": "string",
|
|
"description": "Text to find for replace action",
|
|
},
|
|
"new_text": {
|
|
"type": "string",
|
|
"description": "Replacement text for replace action",
|
|
},
|
|
},
|
|
"required": ["action", "file"],
|
|
},
|
|
)
|
|
self._store = memory_store
|
|
|
|
async def execute(self, **kwargs) -> dict[str, Any]:
|
|
action = kwargs.get("action", "")
|
|
file_key = kwargs.get("file", "")
|
|
|
|
# Validate
|
|
if file_key not in VALID_FILES:
|
|
return {"success": False, "error": f"Invalid file: {file_key}. Must be one of {VALID_FILES}"}
|
|
if action not in VALID_ACTIONS:
|
|
return {"success": False, "error": f"Unknown action: {action}. Must be one of {VALID_ACTIONS}"}
|
|
|
|
try:
|
|
mf = self._store.get_file(file_key)
|
|
|
|
if action == "read":
|
|
content = mf.read()
|
|
return {"success": True, "content": content}
|
|
|
|
elif action == "add":
|
|
section = kwargs.get("section", "")
|
|
content = kwargs.get("content", "")
|
|
if not section:
|
|
return {"success": False, "error": "section is required for add action"}
|
|
mf.add_section(section, content)
|
|
return {"success": True, "message": f"Added to {file_key}/{section}"}
|
|
|
|
elif action == "replace":
|
|
section = kwargs.get("section", "")
|
|
old_text = kwargs.get("old_text", "")
|
|
new_text = kwargs.get("new_text", "")
|
|
if not section:
|
|
return {"success": False, "error": "section is required for replace action"}
|
|
if not old_text:
|
|
return {"success": False, "error": "old_text is required for replace action"}
|
|
success = mf.replace_section(section, old_text, new_text)
|
|
if not success:
|
|
return {"success": False, "error": f"old_text not found in {file_key}/{section}"}
|
|
return {"success": True, "message": f"Replaced in {file_key}/{section}"}
|
|
|
|
elif action == "remove":
|
|
section = kwargs.get("section", "")
|
|
if not section:
|
|
return {"success": False, "error": "section is required for remove action"}
|
|
mf.remove_section(section)
|
|
return {"success": True, "message": f"Removed {file_key}/{section}"}
|
|
|
|
return {"success": False, "error": f"Unhandled action: {action}"}
|
|
|
|
except Exception as e:
|
|
return {"success": False, "error": str(e)}
|