130 lines
4.1 KiB
Python
130 lines
4.1 KiB
Python
"""Beaver 内置 memory tool。
|
||
|
||
这个文件的职责很单纯：把 `MemoryStore` 暴露成一个 agent runtime 可以调用的统一工具。
|
||
|
||
设计边界：
|
||
1. `store.py` 负责底层数据与并发安全
|
||
2. 本文件负责工具接口、参数校验分发、JSON 响应
|
||
3. 更高层的 engine / loader 之后再决定如何把这个工具注册进 runtime
|
||
|
||
换句话说，本文件是“memory 能力的工具化外壳”，不是记忆实现本身。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
from dataclasses import dataclass, field
|
||
from typing import Any
|
||
|
||
from beaver.memory.curated.store import MemoryStore
|
||
|
||
MEMORY_TOOL_DESCRIPTION = (
|
||
"Save durable information to persistent memory that survives across sessions. "
|
||
"Use this proactively for user corrections, preferences, environment facts, "
|
||
"project conventions, and stable tool quirks. Do not store temporary task "
|
||
"progress or raw session logs here; use session search for historical detail."
|
||
)
|
||
|
||
MEMORY_TOOL_PARAMETERS: dict[str, Any] = {
|
||
"type": "object",
|
||
"properties": {
|
||
"action": {
|
||
"type": "string",
|
||
"enum": ["add", "replace", "remove"],
|
||
"description": "The memory operation to perform.",
|
||
},
|
||
"target": {
|
||
"type": "string",
|
||
"enum": ["memory", "user"],
|
||
"description": "Which curated store to update.",
|
||
},
|
||
"content": {
|
||
"type": "string",
|
||
"description": "The new entry content. Required for add and replace.",
|
||
},
|
||
"old_text": {
|
||
"type": "string",
|
||
"description": "A short unique substring identifying the entry to replace or remove.",
|
||
},
|
||
},
|
||
"required": ["action", "target"],
|
||
}
|
||
|
||
|
||
def memory_tool(
|
||
*,
|
||
action: str,
|
||
target: str = "memory",
|
||
content: str | None = None,
|
||
old_text: str | None = None,
|
||
store: MemoryStore | None = None,
|
||
) -> str:
|
||
"""分发 Hermes 风格的 CRUD memory API，并返回 JSON 字符串。
|
||
|
||
这里统一采用 JSON 返回，是为了兼容常见 tool-calling 场景：
|
||
- LLM 更容易消费结构化结果
|
||
- Web/API/日志层也更容易透传和记录
|
||
"""
|
||
|
||
if store is None:
|
||
return json.dumps(
|
||
{
|
||
"success": False,
|
||
"error": "Memory store is not available for this runtime.",
|
||
},
|
||
ensure_ascii=False,
|
||
)
|
||
|
||
if target not in {"memory", "user"}:
|
||
return json.dumps(
|
||
{
|
||
"success": False,
|
||
"error": f"Invalid target '{target}'. Use 'memory' or 'user'.",
|
||
},
|
||
ensure_ascii=False,
|
||
)
|
||
|
||
if action == "add":
|
||
if not content:
|
||
result = {"success": False, "error": "content is required for add."}
|
||
else:
|
||
result = store.add(target, content)
|
||
elif action == "replace":
|
||
if not old_text:
|
||
result = {"success": False, "error": "old_text is required for replace."}
|
||
elif not content:
|
||
result = {"success": False, "error": "content is required for replace."}
|
||
else:
|
||
result = store.replace(target, old_text, content)
|
||
elif action == "remove":
|
||
if not old_text:
|
||
result = {"success": False, "error": "old_text is required for remove."}
|
||
else:
|
||
result = store.remove(target, old_text)
|
||
else:
|
||
result = {
|
||
"success": False,
|
||
"error": f"Unknown action '{action}'. Use add, replace, or remove.",
|
||
}
|
||
|
||
return json.dumps(result, ensure_ascii=False)
|
||
|
||
|
||
@dataclass(slots=True)
|
||
class MemoryTool:
|
||
"""面向 runtime 的轻量工具封装。
|
||
|
||
这里故意保持很薄：
|
||
1. 不重复实现业务逻辑
|
||
2. 不重复维护 schema
|
||
3. 只做 `execute()` 到 `memory_tool()` 的桥接
|
||
"""
|
||
|
||
store: MemoryStore
|
||
name: str = "memory"
|
||
description: str = MEMORY_TOOL_DESCRIPTION
|
||
parameters: dict[str, Any] = field(default_factory=lambda: dict(MEMORY_TOOL_PARAMETERS))
|
||
|
||
async def execute(self, **kwargs: Any) -> str:
|
||
return memory_tool(store=self.store, **kwargs)
|