import datetime import hashlib import json 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 class BaseMemory(BaseModel, ABC): workspace_id: str = Field(default="") memory_id: str = Field(default_factory=lambda: uuid4().hex) memory_type: str = Field(default=...) when_to_use: str = Field(default="") content: str | bytes = Field(default="") score: float = Field(default=0) time_created: str = Field(default_factory=lambda: datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")) time_modified: str = Field(default_factory=lambda: datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")) author: str = Field(default="") metadata: dict = Field(default_factory=dict) def update_modified_time(self): self.time_modified = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") def update_metadata(self, new_metadata): self.metadata = new_metadata def to_vector_node(self) -> VectorNode: raise NotImplementedError @classmethod def from_vector_node(cls, node: VectorNode): raise NotImplementedError class TaskMemory(BaseMemory): memory_type: str = Field(default="task") 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, "metadata": self.metadata, }) @classmethod def from_vector_node(cls, node: VectorNode) -> "TaskMemory": metadata = node.metadata.copy() return cls(workspace_id=node.workspace_id, memory_id=node.unique_id, memory_type=metadata.pop("memory_type"), when_to_use=node.content, content=metadata.pop("content"), score=metadata.pop("score"), time_created=metadata.pop("time_created"), time_modified=metadata.pop("time_modified"), author=metadata.pop("author"), metadata=metadata.pop("metadata", {})) class PersonalMemory(BaseMemory): memory_type: str = Field(default="personal") target: str = Field(default="") reflection_subject: str = Field(default="") # For storing reflection subject attributes 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, "target": self.target, "reflection_subject": self.reflection_subject, "score": self.score, "time_created": self.time_created, "time_modified": self.time_modified, "author": self.author, "metadata": self.metadata, }) @classmethod def from_vector_node(cls, node: VectorNode) -> "PersonalMemory": metadata = node.metadata.copy() return cls(workspace_id=node.workspace_id, memory_id=node.unique_id, memory_type=metadata.pop("memory_type"), when_to_use=node.content, content=metadata.pop("content"), target=metadata.pop("target", ""), reflection_subject=metadata.pop("reflection_subject", ""), score=metadata.pop("score"), time_created=metadata.pop("time_created"), time_modified=metadata.pop("time_modified"), author=metadata.pop("author"), 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") summary: str = Field(default="", description="Brief summary of the tool call result") evaluation: str = Field(default="", description="Detailed evaluation for the tool invocation") score: float = Field(default=0, description="Score of the Evaluation (0.0 for failure, 1.0 for complete success)") is_summarized: bool = Field(default=False, description="Whether this tool call has been included in a summary") call_hash: str = Field(default="", description="Hash value of input and output combined for deduplication") metadata: dict = Field(default_factory=dict) def generate_hash(self) -> str: """Generate hash value from tool input and output for deduplication""" # Convert input to string if it's a dict input_str = json.dumps(self.input, sort_keys=True) if isinstance(self.input, dict) else str(self.input) # Combine input and output combined = f"{input_str}|{self.output}" # Generate MD5 hash hash_value = hashlib.md5(combined.encode('utf-8')).hexdigest() return hash_value def ensure_hash(self): """Ensure call_hash is set, generate if empty""" if not self.call_hash: self.call_hash = self.generate_hash() 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 = not 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 N 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 N tool calls (or all if less than N) recent_calls = self.tool_call_results[-recent_frequency:] 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 { "avg_token_cost": round(avg_token_cost, 2), "avg_time_cost": round(avg_time_cost, 3), "success_rate": round(success_rate, 4), "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): memory_type = node.metadata.get("memory_type") if memory_type == "task": return TaskMemory.from_vector_node(node) 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!") def dict_to_memory(memory_dict: dict): memory_type = memory_dict.get("memory_type", "task") if memory_type == "task": return TaskMemory(**memory_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!") def task_main(): e1 = TaskMemory( workspace_id="w_1024", memory_id="123", when_to_use="test case use", content="test content", score=0.99, metadata={}) print(e1.model_dump_json(indent=2)) v1 = e1.to_vector_node() 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()