From b9da8d2a7fc6d10e27826e5f4da97c3cdbd2dc80 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Thu, 16 Oct 2025 10:50:27 +0800 Subject: [PATCH] feat(memory): add tool call result handling and memory management - Add new ToolCallResult and ToolMemory schemas for tracking tool executions - Implement tool memory retrieval and summarization flows in default config - Register new parse_tool_call_result_op for processing tool call results - Extend memory conversion logic to support tool memory type - Add test cases for tool memory serialization and deserialization - Include token counting utilities for text processing tasks --- reme_ai/config/default.yaml | 32 ++- reme_ai/retrieve/__init__.py | 1 + reme_ai/retrieve/tool/__init__.py | 0 reme_ai/schema/memory.py | 197 +++++++++++++++++- reme_ai/summary/__init__.py | 1 + reme_ai/summary/tool/__init__.py | 1 + .../summary/tool/parse_tool_call_result_op.py | 13 ++ test/http_client_test.py | 31 ++- test/test5.py | 6 + test/test6.py | 13 ++ 10 files changed, 280 insertions(+), 15 deletions(-) create mode 100644 reme_ai/retrieve/tool/__init__.py create mode 100644 reme_ai/summary/tool/__init__.py create mode 100644 reme_ai/summary/tool/parse_tool_call_result_op.py create mode 100644 test/test5.py create mode 100644 test/test6.py diff --git a/reme_ai/config/default.yaml b/reme_ai/config/default.yaml index 7984694a..4cefb810 100644 --- a/reme_ai/config/default.yaml +++ b/reme_ai/config/default.yaml @@ -49,13 +49,31 @@ flow: description: "A list of conversation trajectory information, including message content and score. This field does not need to be filled in, the system will complete it automatically." required: false - retrieve_task_memory_simple: - flow_content: build_query_op >> recall_vector_store_op >> merge_memory_op + retrieve_tool_memory: + flow_content: parse_tool_call_result_op + description: "Retrieve tool related memories from the memory database based on queries to improve tool recall" + input_schema: + tool_names: + type: string + description: "tool_name1,tool_name2" + required: true + + add_tool_call_result: + flow_content: parse_tool_call_result_op description: "Retrieves the most relevant top-k memory experiences from historical data based on the current query with simplified processing" input_schema: - query: + tool_call_results: + type: array + description: "tool_call_result" + required: true + + summary_tool_call_result: + flow_content: parse_tool_call_result_op + description: "Retrieves the most relevant top-k memory experiences from historical data based on the current query with simplified processing" + input_schema: + tool_names: type: string - description: "current query" + description: "tool_name1,tool_name2" required: true summary_task_memory_simple: @@ -143,6 +161,12 @@ llm: params: temperature: 0.6 + wk1: + backend: openai_compatible + model_name: qwen3-30b-a3b-instruct-2507 + params: + temperature: 0.6 + qwen3_30b_instruct: backend: openai_compatible model_name: qwen3-30b-a3b-instruct-2507 diff --git a/reme_ai/retrieve/__init__.py b/reme_ai/retrieve/__init__.py index 8922c1ab..4f26e94d 100644 --- a/reme_ai/retrieve/__init__.py +++ b/reme_ai/retrieve/__init__.py @@ -1,2 +1,3 @@ from . import personal from . import task +from . import tool diff --git a/reme_ai/retrieve/tool/__init__.py b/reme_ai/retrieve/tool/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme_ai/schema/memory.py b/reme_ai/schema/memory.py index 8ed090cd..63a4c004 100644 --- a/reme_ai/schema/memory.py +++ b/reme_ai/schema/memory.py @@ -1,8 +1,10 @@ import datetime from abc import ABC +from typing import List from uuid import uuid4 from flowllm.schema.vector_node import VectorNode +from mcp.types import CallToolResult, TextContent from pydantic import BaseModel, Field @@ -105,6 +107,118 @@ class PersonalMemory(BaseMemory): metadata=metadata.pop("metadata", {})) +class ToolCallResult(BaseModel): + create_time: str = Field(default="", description="Time of tool invocation") + tool_name: str = Field(default=..., description="Name of the tool") + input: dict | str = Field(default="", description="Tool input") + output: str = Field(default="", description="Tool output") + token_cost: int = Field(default=-1, description="Token consumption of the tool") + success: bool = Field(default=True, description="Whether the tool invocation was successful") + time_cost: float = Field(default=0, description="Time consumed by the tool invocation, in seconds") + evaluation: str = Field(default="", description="Evaluation for the tool invocation") + score: float = Field(default=0, description="Score of the Evaluation") + # updated: bool = Field(default=False, description="Whether tool memory has been updated by `ToolCallMetadata`") + + metadata: dict = Field(default_factory=dict) + + def from_mcp_tool_result(self, tool_result: CallToolResult, max_char_len: int = None): + text_list = [] + for content in tool_result.content: + if isinstance(content, TextContent): + text_list.append(content.text) + + else: + raise NotImplementedError(f"content.type={type(content)} not supported") + content = "\n".join(text_list) + + if max_char_len: + content = content[:max_char_len] + self.output = content + + self.success = tool_result.is_error + self.metadata.update(tool_result.meta) + + +class ToolMemory(BaseMemory): + memory_type: str = Field(default="tool") + tool_call_results: List[ToolCallResult] = Field(default_factory=list) + + def to_vector_node(self) -> VectorNode: + return VectorNode(unique_id=self.memory_id, + workspace_id=self.workspace_id, + content=self.when_to_use, + metadata={ + "memory_type": self.memory_type, + "content": self.content, + "score": self.score, + "time_created": self.time_created, + "time_modified": self.time_modified, + "author": self.author, + "tool_call_results": [x.model_dump() for x in self.tool_call_results], + "metadata": self.metadata, + }) + + def statistic(self, recent_frequency: int = 20) -> dict: + """ + Calculate statistical information for the most recent 20 tool calls. + Returns avg token_cost, success rate, avg time_cost, and avg score. + """ + if not self.tool_call_results: + return { + "total_calls": 0, + "recent_calls_analyzed": 0, + "avg_token_cost": 0.0, + "success_rate": 0.0, + "avg_time_cost": 0.0, + "avg_score": 0.0 + } + + # Get the most recent 20 tool calls (or all if less than 20) + recent_calls = self.tool_call_results[-20:] + total_calls = len(self.tool_call_results) + recent_calls_count = len(recent_calls) + + # Calculate statistics + total_token_cost = sum(call.token_cost for call in recent_calls if call.token_cost >= 0) + valid_token_calls = [call for call in recent_calls if call.token_cost >= 0] + avg_token_cost = total_token_cost / len(valid_token_calls) if valid_token_calls else 0.0 + + successful_calls = sum(1 for call in recent_calls if call.success) + success_rate = successful_calls / recent_calls_count if recent_calls_count > 0 else 0.0 + + total_time_cost = sum(call.time_cost for call in recent_calls) + avg_time_cost = total_time_cost / recent_calls_count if recent_calls_count > 0 else 0.0 + + total_score = sum(call.score for call in recent_calls) + avg_score = total_score / recent_calls_count if recent_calls_count > 0 else 0.0 + + return { + "total_calls": total_calls, + "recent_calls_analyzed": recent_calls_count, + "avg_token_cost": round(avg_token_cost, 2), + "success_rate": round(success_rate, 4), + "avg_time_cost": round(avg_time_cost, 3), + "avg_score": round(avg_score, 3) + } + + @classmethod + def from_vector_node(cls, node: VectorNode) -> "ToolMemory": + metadata = node.metadata.copy() + tool_call_results = [ToolCallResult(**result) for result in metadata.pop("tool_call_results", [])] + return cls(workspace_id=node.workspace_id, + memory_id=node.unique_id, + when_to_use=node.content, + memory_type=metadata.pop("memory_type"), + content=metadata.pop("content"), + score=metadata.pop("score"), + time_created=metadata.pop("time_created"), + time_modified=metadata.pop("time_modified"), + author=metadata.pop("author"), + tool_call_results=tool_call_results, + metadata=metadata.pop("metadata", {})) + + + def vector_node_to_memory(node: VectorNode) -> BaseMemory: memory_type = node.metadata.get("memory_type") if memory_type == "task": @@ -113,6 +227,9 @@ def vector_node_to_memory(node: VectorNode) -> BaseMemory: elif memory_type == "personal": return PersonalMemory.from_vector_node(node) + elif memory_type == "tool": + return ToolMemory.from_vector_node(node) + else: raise RuntimeError(f"memory_type={memory_type} not supported!") @@ -125,11 +242,14 @@ def dict_to_memory(memory_dict: dict): elif memory_type == "personal": return PersonalMemory(**memory_dict) + elif memory_type == "tool": + return ToolMemory(**memory_dict) + else: raise RuntimeError(f"memory_type={memory_type} not supported!") -if __name__ == "__main__": +def task_main(): e1 = TaskMemory( workspace_id="w_1024", memory_id="123", @@ -142,3 +262,78 @@ if __name__ == "__main__": print(v1.model_dump_json(indent=2)) e2 = vector_node_to_memory(v1) print(e2.model_dump_json(indent=2)) + + +def personal_main(): + p1 = PersonalMemory( + workspace_id="w_2048", + memory_id="456", + when_to_use="personal memory test case", + content="personal test content", + target="user_preferences", + reflection_subject="learning_style", + score=0.85, + metadata={"category": "user_profile"}) + print("PersonalMemory test:") + print(p1.model_dump_json(indent=2)) + v1 = p1.to_vector_node() + print("VectorNode:") + print(v1.model_dump_json(indent=2)) + p2 = vector_node_to_memory(v1) + print("Reconstructed PersonalMemory:") + print(p2.model_dump_json(indent=2)) + + +def tool_main(): + # Create sample tool call results + tool_result1 = ToolCallResult( + create_time="2025-10-15 10:30:00", + tool_name="file_reader", + input={"file_path": "/test/file.txt"}, + output="File content successfully read", + token_cost=50, + success=True, + time_cost=0.5, + evaluation="Successfully executed", + score=0.95 + ) + + tool_result2 = ToolCallResult( + create_time="2025-10-15 10:31:00", + tool_name="data_processor", + input={"data": "sample_data", "format": "json"}, + output="Data processed successfully", + token_cost=75, + success=True, + time_cost=1.2, + evaluation="Good performance", + score=0.88 + ) + + t1 = ToolMemory( + workspace_id="w_4096", + memory_id="789", + memory_type="tool", + when_to_use="tool execution memory test", + content="tool execution test content", + score=0.92, + tool_call_results=[tool_result1, tool_result2], + metadata={"execution_context": "test_environment"}) + + print("ToolMemory test:") + print(t1.model_dump_json(indent=2)) + v1 = t1.to_vector_node() + print("VectorNode:") + print(v1.model_dump_json(indent=2)) + t2 = ToolMemory.from_vector_node(v1) + print("Reconstructed ToolMemory:") + print(t2.model_dump_json(indent=2)) + + +if __name__ == "__main__": + print("=== Task Memory Test ===") + # task_main() + print("\n=== Personal Memory Test ===") + # personal_main() + print("\n=== Tool Memory Test ===") + tool_main() diff --git a/reme_ai/summary/__init__.py b/reme_ai/summary/__init__.py index 8922c1ab..4f26e94d 100644 --- a/reme_ai/summary/__init__.py +++ b/reme_ai/summary/__init__.py @@ -1,2 +1,3 @@ from . import personal from . import task +from . import tool diff --git a/reme_ai/summary/tool/__init__.py b/reme_ai/summary/tool/__init__.py new file mode 100644 index 00000000..384b8389 --- /dev/null +++ b/reme_ai/summary/tool/__init__.py @@ -0,0 +1 @@ +from .parse_tool_call_result_op import ParseToolCallResultOp diff --git a/reme_ai/summary/tool/parse_tool_call_result_op.py b/reme_ai/summary/tool/parse_tool_call_result_op.py new file mode 100644 index 00000000..b5198279 --- /dev/null +++ b/reme_ai/summary/tool/parse_tool_call_result_op.py @@ -0,0 +1,13 @@ +from flowllm import C, BaseAsyncOp +from loguru import logger + +@C.register_op() +class ParseToolCallResultOp(BaseAsyncOp): + file_path: str = __file__ + + async def async_execute(self): + tool_call_results: list = self.context.get("tool_call_results", []) + logger.info(f"tool_call_results={tool_call_results}") + + + self.context.response.answer = "haha" \ No newline at end of file diff --git a/test/http_client_test.py b/test/http_client_test.py index 0e276ff3..064d9c9f 100644 --- a/test/http_client_test.py +++ b/test/http_client_test.py @@ -110,23 +110,34 @@ async def run2(session): result = await response.json() print(json.dumps(result, ensure_ascii=False)) + +async def run3(session): + workspace_id = "default2" + + async with session.post( + f"{base_url}/add_tool_call_result", + json={ + "tool_call_results": [ + {"a": 1}, + {"a": 2}, + ], + "workspace_id": workspace_id, + }, + headers={"Content-Type": "application/json"} + ) as response: + result = await response.json() + print(json.dumps(result, ensure_ascii=False)) + + async def main(): async with aiohttp.ClientSession() as session: # 获取工具列表 print("获取工具列表...") - async with session.get(f"{base_url}/list") as response: - if response.status == 200: - tools = await response.json() - print("可用工具:") - for tool in tools: - print(json.dumps(tool, ensure_ascii=False)) - else: - print(f"获取工具列表失败: {response.status}") - return # await run1(session) - await run2(session) + # await run2(session) + await run3(session) if __name__ == "__main__": asyncio.run(main()) diff --git a/test/test5.py b/test/test5.py new file mode 100644 index 00000000..d0845c55 --- /dev/null +++ b/test/test5.py @@ -0,0 +1,6 @@ +import tiktoken +enc = tiktoken.get_encoding("o200k_base") + +# r = enc.encode("我爱吃西瓜，你说啥") +r = enc.encode("hello world aaaaaaaaaaaa") +print(len(r)) \ No newline at end of file diff --git a/test/test6.py b/test/test6.py new file mode 100644 index 00000000..66ceae98 --- /dev/null +++ b/test/test6.py @@ -0,0 +1,13 @@ +import tiktoken + +def count_tokens(text: str) -> int: + """计算给定文本在指定模型下的 token 数量""" + encoding = tiktoken.get_encoding("o200k_base") + tokens = encoding.encode(text) + return len(tokens) + +# 示例使用 +text = "你好，世界！Hello, world!" +token_count = count_tokens(text) +print(f"Token 数量: {token_count}") +print(len(text) / 4)