diff --git a/benchmark/halumem/eval_reme.py b/benchmark/halumem/eval_reme.py index a68ac3b2..eb6b18ef 100644 --- a/benchmark/halumem/eval_reme.py +++ b/benchmark/halumem/eval_reme.py @@ -42,6 +42,7 @@ class EvalConfig: batch_size: int = 20 output_dir: str = "bench_results/reme" eval_model_name: str = "qwen3-max" + algo_version: str = "v1" # ==================== Utilities ==================== @@ -217,7 +218,7 @@ async def answer_question_with_memories( result = await reme.llm.simple_request_for_json( prompt=prompt, - model_name=model_name + model_name="qwen3-30b-a3b-instruct-2507" ) return result @@ -269,9 +270,10 @@ async def evaluation_for_question( class MemoryProcessor: """Handles ReMe memory operations.""" - def __init__(self, reme: ReMe, eval_model_name: str = "qwen3-max"): + def __init__(self, reme: ReMe, eval_model_name: str = "qwen3-max", algo_version: str = "v1"): self.reme = reme self.eval_model_name = eval_model_name + self.algo_version = algo_version async def add_memories( self, @@ -297,7 +299,7 @@ class MemoryProcessor: result = await self.reme.summary_memory( messages=batch, user_name=user_id, - version="default", + version=self.algo_version, return_dict=True, ) @@ -305,7 +307,7 @@ class MemoryProcessor: total_duration_ms += duration_ms extracted_memories.extend([m.model_dump(exclude_none=True) for m in result["answer"]]) - summary_messages.extend([m.simple_dump() for m in result["messages"]]) + summary_messages.extend([m.simple_dump(enable_argument_dict=True) for m in result["messages"]]) return extracted_memories, summary_messages, total_duration_ms @@ -329,13 +331,13 @@ class MemoryProcessor: query=query, retrieve_top_k=top_k, user_name=user_id, - version="default", + version=self.algo_version, return_dict=True, ) # Extract memories from response memories = result["answer"] - agent_messages = [x.model_dump(exclude_none=True) for x in result["messages"]] + agent_messages = [x.simple_dump(enable_argument_dict=True) for x in result["messages"]] retrieved_nodes = [x.model_dump(exclude_none=True) for x in result["retrieved_nodes"]] # Use LLM to generate structured answer from memories @@ -521,7 +523,11 @@ class HaluMemEvaluator: self.reme.prompt_handler.load_prompt_by_file(prompts_yaml_path) self.file_manager = FileManager(config.output_dir) - self.memory_processor = MemoryProcessor(self.reme, config.eval_model_name) + self.memory_processor = MemoryProcessor( + self.reme, + config.eval_model_name, + config.algo_version + ) self.qa_evaluator = QuestionAnsweringEvaluator( self.memory_processor, self.reme, @@ -763,7 +769,8 @@ async def main_async( top_k: int, user_num: int, max_concurrency: int, - eval_model_name: str = "qwen3-max" + eval_model_name: str = "qwen3-max", + algo_version: str = "v1" ): """Main async entry point for ReMe evaluation with proper resource cleanup.""" config = EvalConfig( @@ -771,7 +778,8 @@ async def main_async( top_k=top_k, user_num=user_num, max_concurrency=max_concurrency, - eval_model_name=eval_model_name + eval_model_name=eval_model_name, + algo_version=algo_version ) # Use async context manager for automatic cleanup @@ -784,7 +792,8 @@ def main( top_k: int, user_num: int, max_concurrency: int, - eval_model_name: str = "qwen3-max" + eval_model_name: str = "qwen3-max", + algo_version: str = "v1" ): """Main entry point for ReMe evaluation.""" asyncio.run(main_async( @@ -792,7 +801,8 @@ def main( top_k=top_k, user_num=user_num, max_concurrency=max_concurrency, - eval_model_name=eval_model_name + eval_model_name=eval_model_name, + algo_version=algo_version )) @@ -833,6 +843,12 @@ if __name__ == "__main__": # default="qwen3-235b-a22b-instruct-2507", help="Model name for evaluation (default: qwen3-max)" ) + parser.add_argument( + "--algo_version", + type=str, + default="v1", + help="Algorithm version for summary and retrieval (default: v1)" + ) args = parser.parse_args() @@ -841,5 +857,6 @@ if __name__ == "__main__": top_k=args.top_k, user_num=args.user_num, max_concurrency=args.max_concurrency, - eval_model_name=args.eval_model_name + eval_model_name=args.eval_model_name, + algo_version=args.algo_version ) diff --git a/benchmark/halumem/scripts.sh b/benchmark/halumem/scripts.sh new file mode 100755 index 00000000..0f589175 --- /dev/null +++ b/benchmark/halumem/scripts.sh @@ -0,0 +1 @@ +cat bench_results/reme/Martin\ Mark/session* | grep '"result_type": "' | awk -F'"' '{total++; if($4=="Correct") count++} END {printf "Correct Rate: %.2f%% (%d/%d)\n", (count/total)*100, count, total}' \ No newline at end of file diff --git a/reme/agent/memory/__init__.py b/reme/agent/memory/__init__.py index c1bac883..024b68c4 100644 --- a/reme/agent/memory/__init__.py +++ b/reme/agent/memory/__init__.py @@ -1,9 +1,32 @@ """memory agent""" -from . import default from .base_memory_agent import BaseMemoryAgent +from .personal.personal_retriever import PersonalRetriever +from .personal.personal_summarizer import PersonalSummarizer +from .personal.personal_v1_retriever import PersonalV1Retriever +from .personal.personal_v1_summarizer import PersonalV1Summarizer +from .procedural.procedural_retriever import ProceduralRetriever +from .procedural.procedural_summarizer import ProceduralSummarizer +from .reme_retriever import ReMeRetriever +from .reme_summarizer import ReMeSummarizer +from .tool.tool_retriever import ToolRetriever +from .tool.tool_summarizer import ToolSummarizer +from ...core import R __all__ = [ - "default", "BaseMemoryAgent", + "PersonalRetriever", + "PersonalSummarizer", + "PersonalV1Retriever", + "PersonalV1Summarizer", + "ProceduralRetriever", + "ProceduralSummarizer", + "ReMeRetriever", + "ReMeSummarizer", + "ToolRetriever", + "ToolSummarizer", ] + +for name in __all__: + tool_class = globals()[name] + R.op.register()(tool_class) diff --git a/reme/agent/memory/default/__init__.py b/reme/agent/memory/default/__init__.py deleted file mode 100644 index 65eb3463..00000000 --- a/reme/agent/memory/default/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -"""Default memory agents for personal, procedural, tool and ReMe memory operations.""" - -from .personal_retriever import PersonalRetriever -from .personal_summarizer import PersonalSummarizer -from .procedural_retriever import ProceduralRetriever -from .procedural_summarizer import ProceduralSummarizer -from .reme_retriever import ReMeRetriever -from .reme_summarizer import ReMeSummarizer -from .tool_retriever import ToolRetriever -from .tool_summarizer import ToolSummarizer -from ....core import R - -__all__ = [ - "PersonalRetriever", - "PersonalSummarizer", - "ProceduralRetriever", - "ProceduralSummarizer", - "ReMeRetriever", - "ReMeSummarizer", - "ToolRetriever", - "ToolSummarizer", -] - -for name in __all__: - tool_class = globals()[name] - R.op.register()(tool_class) diff --git a/reme/agent/memory/default/tool_retriever.py b/reme/agent/memory/default/tool_retriever.py deleted file mode 100644 index 6a2d206b..00000000 --- a/reme/agent/memory/default/tool_retriever.py +++ /dev/null @@ -1,10 +0,0 @@ -"""Tool memory retriever agent implementation.""" - -from ..base_memory_agent import BaseMemoryAgent -from ....core.enumeration import MemoryType - - -class ToolRetriever(BaseMemoryAgent): - """Agent responsible for retrieving tool-related memories.""" - - memory_type: MemoryType = MemoryType.TOOL diff --git a/reme/agent/memory/default/tool_summarizer.py b/reme/agent/memory/default/tool_summarizer.py deleted file mode 100644 index 85222d4f..00000000 --- a/reme/agent/memory/default/tool_summarizer.py +++ /dev/null @@ -1,10 +0,0 @@ -"""Tool memory summarizer agent implementation.""" - -from ..base_memory_agent import BaseMemoryAgent -from ....core.enumeration import MemoryType - - -class ToolSummarizer(BaseMemoryAgent): - """Agent responsible for summarizing tool-related memories.""" - - memory_type: MemoryType = MemoryType.TOOL diff --git a/reme/agent/memory/personal/__init__.py b/reme/agent/memory/personal/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme/agent/memory/default/personal_retriever.py b/reme/agent/memory/personal/personal_retriever.py similarity index 100% rename from reme/agent/memory/default/personal_retriever.py rename to reme/agent/memory/personal/personal_retriever.py diff --git a/reme/agent/memory/default/personal_retriever.yaml b/reme/agent/memory/personal/personal_retriever.yaml similarity index 100% rename from reme/agent/memory/default/personal_retriever.yaml rename to reme/agent/memory/personal/personal_retriever.yaml diff --git a/reme/agent/memory/default/personal_summarizer.py b/reme/agent/memory/personal/personal_summarizer.py similarity index 100% rename from reme/agent/memory/default/personal_summarizer.py rename to reme/agent/memory/personal/personal_summarizer.py diff --git a/reme/agent/memory/default/personal_summarizer.yaml b/reme/agent/memory/personal/personal_summarizer.yaml similarity index 100% rename from reme/agent/memory/default/personal_summarizer.yaml rename to reme/agent/memory/personal/personal_summarizer.yaml diff --git a/reme/agent/memory/personal/personal_v1_retriever.py b/reme/agent/memory/personal/personal_v1_retriever.py new file mode 100644 index 00000000..34b53e09 --- /dev/null +++ b/reme/agent/memory/personal/personal_v1_retriever.py @@ -0,0 +1,86 @@ +"""Personal memory retriever agent for retrieving personal memories through vector search.""" + +from ..base_memory_agent import BaseMemoryAgent +from ....core.enumeration import Role, MemoryType +from ....core.op import BaseTool +from ....core.schema import Message +from ....core.utils import format_messages + + +class PersonalV1Retriever(BaseMemoryAgent): + """Retrieve personal memories through vector search and history reading.""" + + memory_type: MemoryType = MemoryType.PERSONAL + + async def build_messages(self) -> list[Message]: + if self.context.get("query"): + context = self.context.query + elif self.context.get("messages"): + context = self.description + "\n" + format_messages(self.context.messages) + else: + raise ValueError("input must have either `query` or `messages`") + + read_all_profiles_tool: BaseTool | None = self.pop_tool("read_all_profiles") + if read_all_profiles_tool is not None: + all_profiles = await read_all_profiles_tool.call( + memory_target=self.memory_target, + service_context=self.service_context, + ) + else: + all_profiles = "" + + return [ + Message( + role=Role.SYSTEM, + content=self.prompt_format( + prompt_name="system_prompt", + memory_type=self.memory_type.value, + memory_target=self.memory_target, + user_profile=all_profiles, + context=context.strip(), + ), + ), + Message( + role=Role.USER, + content=self.get_prompt("user_message"), + ), + ] + + async def _acting_step( + self, + assistant_message: Message, + tools: list[BaseTool], + step: int, + stage: str = "", + **kwargs, + ) -> tuple[list[BaseTool], list[Message]]: + """Execute tool calls with memory context.""" + return await super()._acting_step( + assistant_message, + tools, + step, + memory_type=self.memory_type.value, + memory_target=self.memory_target, + retrieved_nodes=self.retrieved_nodes, + **kwargs, + ) + + async def execute(self): + result = await super().execute() + answer = result["answer"] + if "MEMORY_NOT_FOUND" in answer: + result["answer"] = "\n".join( + [ + n.format( + include_memory_id=False, + include_when_to_use=False, + include_content=True, + include_message_time=False, + ref_memory_id_key="", + ) + for n in self.retrieved_nodes + ], + ) + + result["retrieved_nodes"] = self.retrieved_nodes + return result diff --git a/reme/agent/memory/personal/personal_v1_summarizer.py b/reme/agent/memory/personal/personal_v1_summarizer.py new file mode 100644 index 00000000..864cbd3b --- /dev/null +++ b/reme/agent/memory/personal/personal_v1_summarizer.py @@ -0,0 +1,117 @@ +"""Personal memory summarizer agent for two-phase personal memory processing.""" + +from loguru import logger + +from ..base_memory_agent import BaseMemoryAgent +from ....core.enumeration import Role, MemoryType +from ....core.op import BaseTool +from ....core.schema import Message + + +class PersonalV1Summarizer(BaseMemoryAgent): + """Two-phase personal memory processor: retrieve/add memories then update profile.""" + + memory_type: MemoryType = MemoryType.PERSONAL + + async def _build_s1_messages(self) -> list[Message]: + return [ + Message( + role=Role.SYSTEM, + content=self.prompt_format( + prompt_name="system_prompt_s1", + context=self.context.history_node.content, + memory_type=self.memory_type.value, + memory_target=self.memory_target, + ), + ), + Message( + role=Role.USER, + content=self.get_prompt("user_message_s1"), + ), + ] + + async def _build_s2_messages(self) -> list[Message]: + return [ + Message( + role=Role.SYSTEM, + content=self.prompt_format( + prompt_name="system_prompt_s2", + context=self.context.history_node.content, + memory_type=self.memory_type.value, + memory_target=self.memory_target, + ), + ), + Message( + role=Role.USER, + content=self.get_prompt("user_message_s2"), + ), + ] + + async def _acting_step( + self, + assistant_message: Message, + tools: list[BaseTool], + step: int, + stage: str = "", + **kwargs, + ) -> tuple[list[BaseTool], list[Message]]: + """Execute tool calls with memory context.""" + return await super()._acting_step( + assistant_message, + tools, + step, + stage=stage, + memory_type=self.memory_type.value, + memory_target=self.memory_target, + history_node=self.history_node, + author=self.author, + retrieved_nodes=self.retrieved_nodes, + **kwargs, + ) + + async def execute(self): + memory_tools = [] + profile_tools = [] + for i, tool in enumerate(self.tools): + tool_name = tool.tool_call.name + if "_memory" in tool_name: + memory_tools.append(tool) + elif "_profile" in tool_name: + profile_tools.append(tool) + else: + raise ValueError(f"[{self.__class__.__name__}] unknown tool_name={tool_name}") + logger.info(f"[{self.__class__.__name__}] tool_call[{i}]={tool.tool_call.simple_input_dump(as_dict=False)}") + + stage = "s1-memory" + messages_s1 = await self._build_s1_messages() + for i, message in enumerate(messages_s1): + role = message.name or message.role + logger.info(f"[{self.__class__.__name__} {stage}] role={role} {message.simple_dump(as_dict=False)}") + tools_s1, messages_s1, success_s1 = await self.react(messages_s1, memory_tools, stage=stage) + + if profile_tools: + stage = "s2-profile" + messages_s2 = await self._build_s2_messages() + for i, message in enumerate(messages_s2): + role = message.name or message.role + logger.info(f"[{self.__class__.__name__} {stage}] role={role} {message.simple_dump(as_dict=False)}") + tools_s2, messages_s2, success_s2 = await self.react(messages_s2, profile_tools, stage=stage) + else: + tools_s2, messages_s2, success_s2 = [], [], True + + answer = (messages_s1[-1].content if success_s1 else "") + (messages_s2[-1].content if success_s2 else "") + success = success_s1 and success_s2 + messages = messages_s1 + messages_s2 + tools = tools_s1 + tools_s2 + memory_nodes = [] + for tool in tools: + if tool.memory_nodes: + memory_nodes.extend(tool.memory_nodes) + + return { + "answer": answer, + "success": success, + "messages": messages, + "tools": tools, + "memory_nodes": memory_nodes, + } diff --git a/reme/agent/memory/personal/personal_v1_summarizer.yaml b/reme/agent/memory/personal/personal_v1_summarizer.yaml new file mode 100644 index 00000000..f30dabf6 --- /dev/null +++ b/reme/agent/memory/personal/personal_v1_summarizer.yaml @@ -0,0 +1,65 @@ +system_prompt_s1: | + You are a Memory Agent responsible for managing {memory_type} memories about {memory_target}. + + ## Latest Conversation + Format: round [] : + {context} + + ## Task + ### Step 1: Create Memory Draft + Create a memory draft in `add_draft_and_retrieve_similar_memory` based on the latest conversation. + - Use actual names from the conversation (e.g., "Bob likes apples") instead of generic references (e.g., "user likes apples") + - Always record memories with real names + - The tool will retrieve similar historical memories via vector search to help you consolidate in Step 2 + + ### Step 2: Update Memory Store + Update the vector store using `update_memory` to keep it well-organized and consolidated: + + **What to Delete** (via `memory_ids_to_delete`): + - Duplicate memories with identical or highly similar content + - Memories that should be merged into a single consolidated entry + + **What to Add** (via `memories_to_add` with message_time and memory_content): + - For each topic with changes: add ONE consolidated memory that merges related information + - New distinct memories that don't overlap with existing ones + - Updated memories that capture the latest state while preserving temporal evolution + + ## Requirements + - Extract only what's explicitly stated—no inferences, assumptions, or fabrications + - Preserve temporal evolution: capture how things change over time within the same topic + - Maintain organization: group related memories by topic and eliminate all redundancy + +user_message_s1: | + Complete the task by following Step 1 and Step 2 in order + +system_prompt_s2: | + You are a Profile Agent responsible for managing profiles about {memory_target}. + + ## Latest Conversation + Format: round [] : + {context} + + ## Task + ### Step 1: Create Profile Draft + Create a profile draft in `add_draft_and_read_all_profiles` based on the latest conversation. + - The tool will return all existing profiles to help you maintain the profile store in Step 2 + + ### Step 2: Update Profile Store + Update the profile store using `update_profile` to keep it well-organized and consolidated: + + **What to Delete** (via `profile_ids_to_delete`): + - Duplicate profiles with identical keys or values + - Conflicting profiles that contradict the new information + - Profiles that should be merged into a single consolidated entry + + **What to Add** (via `profiles_to_add` with message_time, profile_key, and profile_value): + - For each profile key with changes: add ONE consolidated profile that merges related information + - New distinct profiles that don't overlap with existing ones + - Updated profiles that capture the latest state + + ## Requirements + - Extract only what's explicitly stated—no inferences, assumptions, or fabrications + - Maintain organization: group related profiles by key and eliminate all redundancy + +user_message_s2: | + Complete the task by following Step 1 and Step 2 in order diff --git a/reme/agent/memory/default/procedural_retriever.py b/reme/agent/memory/procedural/procedural_retriever.py similarity index 100% rename from reme/agent/memory/default/procedural_retriever.py rename to reme/agent/memory/procedural/procedural_retriever.py diff --git a/reme/agent/memory/default/procedural_summarizer.py b/reme/agent/memory/procedural/procedural_summarizer.py similarity index 100% rename from reme/agent/memory/default/procedural_summarizer.py rename to reme/agent/memory/procedural/procedural_summarizer.py diff --git a/reme/agent/memory/default/reme_retriever.py b/reme/agent/memory/reme_retriever.py similarity index 93% rename from reme/agent/memory/default/reme_retriever.py rename to reme/agent/memory/reme_retriever.py index 2f3f78bb..9de51aa2 100644 --- a/reme/agent/memory/default/reme_retriever.py +++ b/reme/agent/memory/reme_retriever.py @@ -1,10 +1,10 @@ """ReMe retriever agent that orchestrates multiple memory agents to retrieve information.""" -from ..base_memory_agent import BaseMemoryAgent -from ....core.enumeration import Role -from ....core.op import BaseTool -from ....core.schema import Message -from ....core.utils import format_messages +from .base_memory_agent import BaseMemoryAgent +from ...core.enumeration import Role +from ...core.op import BaseTool +from ...core.schema import Message +from ...core.utils import format_messages class ReMeRetriever(BaseMemoryAgent): diff --git a/reme/agent/memory/default/reme_retriever.yaml b/reme/agent/memory/reme_retriever.yaml similarity index 93% rename from reme/agent/memory/default/reme_retriever.yaml rename to reme/agent/memory/reme_retriever.yaml index 84fb6b4e..9b5ea8af 100644 --- a/reme/agent/memory/default/reme_retriever.yaml +++ b/reme/agent/memory/reme_retriever.yaml @@ -23,7 +23,5 @@ system_prompt: | - Each memory_target should be assigned **only once** - do not duplicate assignments 3. Use the `delegate_task` tool **once** with all relevant memory_target(s) to enable parallel processing by specialized agents - Note: If the context contains no memorable information (e.g., simple greetings), return ``. - user_message: | Please analyze the context and delegate retrieval tasks to the appropriate specialized agents. \ No newline at end of file diff --git a/reme/agent/memory/default/reme_summarizer.py b/reme/agent/memory/reme_summarizer.py similarity index 94% rename from reme/agent/memory/default/reme_summarizer.py rename to reme/agent/memory/reme_summarizer.py index 934863bb..8b263f64 100644 --- a/reme/agent/memory/default/reme_summarizer.py +++ b/reme/agent/memory/reme_summarizer.py @@ -1,10 +1,10 @@ """ReMe summarizer agent that orchestrates multiple memory agents to summarize information.""" -from ..base_memory_agent import BaseMemoryAgent -from ....core.enumeration import Role -from ....core.op import BaseTool -from ....core.schema import Message -from ....core.utils import format_messages +from .base_memory_agent import BaseMemoryAgent +from ...core.enumeration import Role +from ...core.op import BaseTool +from ...core.schema import Message +from ...core.utils import format_messages class ReMeSummarizer(BaseMemoryAgent): diff --git a/reme/agent/memory/default/reme_summarizer.yaml b/reme/agent/memory/reme_summarizer.yaml similarity index 100% rename from reme/agent/memory/default/reme_summarizer.yaml rename to reme/agent/memory/reme_summarizer.yaml diff --git a/reme/core/schema/message.py b/reme/core/schema/message.py index 54acde63..d9201979 100644 --- a/reme/core/schema/message.py +++ b/reme/core/schema/message.py @@ -84,6 +84,7 @@ class Message(BaseModel): add_reasoning: bool = True, add_time_created: bool = False, add_metadata: bool = False, + enable_argument_dict: bool = False, as_dict: bool = True, ) -> dict | str: """Transforms the message into a simplified dictionary for standard APIs.""" @@ -98,7 +99,13 @@ class Message(BaseModel): result["reasoning_content"] = self.reasoning_content if self.tool_calls: - result["tool_calls"] = [tc.simple_output_dump() for tc in self.tool_calls] + result["tool_calls"] = [ + tc.simple_output_dump( + as_dict=True, + enable_argument_dict=enable_argument_dict, + ) + for tc in self.tool_calls + ] if self.tool_call_id: result["tool_call_id"] = self.tool_call_id diff --git a/reme/core/schema/tool_call.py b/reme/core/schema/tool_call.py index 9cd58c7f..19d7e106 100644 --- a/reme/core/schema/tool_call.py +++ b/reme/core/schema/tool_call.py @@ -145,17 +145,13 @@ class ToolCall(BaseModel): } return result if as_dict else json.dumps(result, ensure_ascii=False) - def simple_output_dump(self, as_dict: bool = True) -> dict | str: - """Convert ToolCall to output format dictionary or JSON string for API responses. - - Args: - as_dict: If True, returns dict; if False, returns JSON string. - """ + def simple_output_dump(self, as_dict: bool = True, enable_argument_dict: bool = False) -> dict | str: + """Convert ToolCall to output format dictionary or JSON string for API responses.""" result = { "index": self.index, "id": self.id, self.type: { - "arguments": self.arguments, + "arguments": self.argument_dict if enable_argument_dict else self.arguments, "name": self.name, }, "type": self.type, diff --git a/reme/reme.py b/reme/reme.py index 326cf7a6..b7c1ea5e 100644 --- a/reme/reme.py +++ b/reme/reme.py @@ -3,15 +3,17 @@ import sys from pathlib import Path -from reme.agent.memory import BaseMemoryAgent -from .agent.memory.default import ( +from .agent.memory import ( + BaseMemoryAgent, ReMeSummarizer, + ReMeRetriever, + PersonalV1Summarizer, + PersonalV1Retriever, PersonalSummarizer, PersonalRetriever, - ReMeRetriever, ProceduralSummarizer, - ToolSummarizer, ProceduralRetriever, + ToolSummarizer, ToolRetriever, ) from .config import ReMeConfigParser @@ -116,7 +118,7 @@ class ReMe(Application): format_messages.append(message) personal_summarizer: BaseMemoryAgent - if version: + if version == "default": personal_summarizer = PersonalSummarizer( tools=[ AddDraftAndRetrieveSimilarMemory( @@ -134,17 +136,33 @@ class ReMe(Application): ), ], ) + + elif version == "v1": + personal_summarizer = PersonalV1Summarizer( + tools=[ + AddDraftAndRetrieveSimilarMemory( + enable_thinking_params=enable_thinking_params, + top_k=retrieve_top_k, + ), + UpdateMemoryV2(enable_thinking_params=enable_thinking_params), + AddDraftAndReadAllProfiles( + enable_thinking_params=enable_thinking_params, + profile_dir=self.profile_dir, + ), + ], + ) + else: raise NotImplementedError procedural_summarizer: BaseMemoryAgent - if version == "default": + if version in ["default", "v1"]: procedural_summarizer = ProceduralSummarizer(tools=[]) else: raise NotImplementedError tool_summarizer: BaseMemoryAgent - if version == "default": + if version in ["default", "v1"]: tool_summarizer = ToolSummarizer(tools=[]) else: raise NotImplementedError @@ -186,7 +204,7 @@ class ReMe(Application): memory_agents = [personal_summarizer, procedural_summarizer, tool_summarizer] reme_summarizer: BaseMemoryAgent - if version == "default": + if version in ["default", "v1"]: reme_summarizer = ReMeSummarizer(tools=[AddHistory(), DelegateTask(memory_agents=memory_agents)]) else: raise NotImplementedError @@ -221,7 +239,7 @@ class ReMe(Application): """Retrieve relevant personal, procedural and tool memories for a query.""" personal_retriever: BaseMemoryAgent - if version: + if version == "default": personal_retriever = PersonalRetriever( tools=[ ReadAllProfiles( @@ -236,17 +254,33 @@ class ReMe(Application): ReadHistory(enable_thinking_params=enable_thinking_params), ], ) + + elif version == "v1": + personal_retriever = PersonalV1Retriever( + tools=[ + ReadAllProfiles( + enable_thinking_params=enable_thinking_params, + profile_dir=self.profile_dir, + ), + RetrieveMemory( + enable_thinking_params=enable_thinking_params, + top_k=retrieve_top_k, + enable_time_filter=enable_time_filter, + ), + ], + ) + else: raise NotImplementedError procedural_retriever: BaseMemoryAgent - if version == "default": + if version in ["default", "v1"]: procedural_retriever = ProceduralRetriever(tools=[]) else: raise NotImplementedError tool_retriever: BaseMemoryAgent - if version == "default": + if version in ["default", "v1"]: tool_retriever = ToolRetriever(tools=[]) else: raise NotImplementedError @@ -286,7 +320,7 @@ class ReMe(Application): memory_agents = [personal_retriever, procedural_retriever, tool_retriever] reme_retriever: BaseMemoryAgent - if version == "default": + if version in ["default", "v1"]: reme_retriever = ReMeRetriever(tools=[DelegateTask(memory_agents=memory_agents)]) else: raise NotImplementedError diff --git a/reme/tool/memory/__init__.py b/reme/tool/memory/__init__.py index cab968dd..ff58d44f 100644 --- a/reme/tool/memory/__init__.py +++ b/reme/tool/memory/__init__.py @@ -1,42 +1,48 @@ """memory tools""" -from .add_draft_and_read_all_profiles import AddDraftAndReadAllProfiles -from .add_draft_and_retrieve_similar_memory import AddDraftAndRetrieveSimilarMemory -from .add_history import AddHistory -from .add_memory import AddMemory from .base_memory_tool import BaseMemoryTool from .delegate_task import DelegateTask -from .delete_memory import DeleteMemory -from .memory_handler import MemoryHandler -from .profile_handler import ProfileHandler -from .read_all_profiles import ReadAllProfiles -from .read_history import ReadHistory -from .retrieve_memory import RetrieveMemory -from .retrieve_recent_memory import RetrieveRecentMemory -from .update_memory import UpdateMemory -from .update_memory_v2 import UpdateMemoryV2 -from .update_profile import UpdateProfile +from .history.add_history import AddHistory +from .history.read_history import ReadHistory +from .profiles.add_draft_and_read_all_profiles import AddDraftAndReadAllProfiles +from .profiles.profile_handler import ProfileHandler +from .profiles.read_all_profiles import ReadAllProfiles +from .profiles.update_profile import UpdateProfile +from .vector.add_draft_and_retrieve_similar_memory import AddDraftAndRetrieveSimilarMemory +from .vector.add_memory import AddMemory +from .vector.delete_memory import DeleteMemory +from .vector.memory_handler import MemoryHandler +from .vector.retrieve_memory import RetrieveMemory +from .vector.retrieve_recent_memory import RetrieveRecentMemory +from .vector.update_memory import UpdateMemory +from .vector.update_memory_v2 import UpdateMemoryV2 from ...core import R __all__ = [ - "AddDraftAndReadAllProfiles", - "AddDraftAndRetrieveSimilarMemory", - "AddHistory", - "AddMemory", + # Base "BaseMemoryTool", "DelegateTask", - "DeleteMemory", - "MemoryHandler", + # History + "AddHistory", + "ReadHistory", + # Profiles + "AddDraftAndReadAllProfiles", "ProfileHandler", "ReadAllProfiles", - "ReadHistory", + "UpdateProfile", + # Vector + "AddDraftAndRetrieveSimilarMemory", + "AddMemory", + "DeleteMemory", + "MemoryHandler", "RetrieveMemory", "RetrieveRecentMemory", "UpdateMemory", "UpdateMemoryV2", - "UpdateProfile", ] for name in __all__: tool_class = globals()[name] - R.op.register()(tool_class) + # Only register classes that inherit from BaseMemoryTool + if isinstance(tool_class, type) and issubclass(tool_class, BaseMemoryTool) and tool_class is not BaseMemoryTool: + R.op.register()(tool_class) diff --git a/reme/tool/memory/history/__init__.py b/reme/tool/memory/history/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme/tool/memory/add_history.py b/reme/tool/memory/history/add_history.py similarity index 89% rename from reme/tool/memory/add_history.py rename to reme/tool/memory/history/add_history.py index 57a0097e..629f4532 100644 --- a/reme/tool/memory/add_history.py +++ b/reme/tool/memory/history/add_history.py @@ -2,10 +2,10 @@ from loguru import logger -from .base_memory_tool import BaseMemoryTool -from ...core.enumeration import MemoryType -from ...core.schema import ToolCall, MemoryNode, Message -from ...core.utils import format_messages +from ..base_memory_tool import BaseMemoryTool +from ....core.enumeration import MemoryType +from ....core.schema import ToolCall, MemoryNode, Message +from ....core.utils import format_messages class AddHistory(BaseMemoryTool): diff --git a/reme/tool/memory/read_history.py b/reme/tool/memory/history/read_history.py similarity index 93% rename from reme/tool/memory/read_history.py rename to reme/tool/memory/history/read_history.py index 500d4bd9..1f328eb2 100644 --- a/reme/tool/memory/read_history.py +++ b/reme/tool/memory/history/read_history.py @@ -2,8 +2,8 @@ from loguru import logger -from .base_memory_tool import BaseMemoryTool -from ...core.schema import MemoryNode, ToolCall +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import MemoryNode, ToolCall class ReadHistory(BaseMemoryTool): diff --git a/reme/tool/memory/profiles/__init__.py b/reme/tool/memory/profiles/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme/tool/memory/add_draft_and_read_all_profiles.py b/reme/tool/memory/profiles/add_draft_and_read_all_profiles.py similarity index 97% rename from reme/tool/memory/add_draft_and_read_all_profiles.py rename to reme/tool/memory/profiles/add_draft_and_read_all_profiles.py index 0f9ba241..716c4e5e 100644 --- a/reme/tool/memory/add_draft_and_read_all_profiles.py +++ b/reme/tool/memory/profiles/add_draft_and_read_all_profiles.py @@ -2,9 +2,9 @@ from loguru import logger -from .base_memory_tool import BaseMemoryTool from .profile_handler import ProfileHandler -from ...core.schema import ToolCall +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import ToolCall class AddDraftAndReadAllProfiles(BaseMemoryTool): diff --git a/reme/tool/memory/profile_handler.py b/reme/tool/memory/profiles/profile_handler.py similarity index 98% rename from reme/tool/memory/profile_handler.py rename to reme/tool/memory/profiles/profile_handler.py index 0f51d014..6f646ac3 100644 --- a/reme/tool/memory/profile_handler.py +++ b/reme/tool/memory/profiles/profile_handler.py @@ -4,9 +4,9 @@ from pathlib import Path from loguru import logger -from ...core.enumeration import MemoryType -from ...core.schema import MemoryNode -from ...core.utils import CacheHandler, deduplicate_memories +from ....core.enumeration import MemoryType +from ....core.schema import MemoryNode +from ....core.utils import CacheHandler, deduplicate_memories class ProfileHandler: diff --git a/reme/tool/memory/read_all_profiles.py b/reme/tool/memory/profiles/read_all_profiles.py similarity index 92% rename from reme/tool/memory/read_all_profiles.py rename to reme/tool/memory/profiles/read_all_profiles.py index d18be220..892bbb58 100644 --- a/reme/tool/memory/read_all_profiles.py +++ b/reme/tool/memory/profiles/read_all_profiles.py @@ -2,9 +2,9 @@ from loguru import logger -from .base_memory_tool import BaseMemoryTool from .profile_handler import ProfileHandler -from ...core.schema import ToolCall +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import ToolCall class ReadAllProfiles(BaseMemoryTool): diff --git a/reme/tool/memory/update_profile.py b/reme/tool/memory/profiles/update_profile.py similarity index 97% rename from reme/tool/memory/update_profile.py rename to reme/tool/memory/profiles/update_profile.py index 5ab7c962..5edd990b 100644 --- a/reme/tool/memory/update_profile.py +++ b/reme/tool/memory/profiles/update_profile.py @@ -2,9 +2,9 @@ from loguru import logger -from .base_memory_tool import BaseMemoryTool from .profile_handler import ProfileHandler -from ...core.schema import ToolCall +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import ToolCall class UpdateProfile(BaseMemoryTool): diff --git a/reme/tool/memory/vector/__init__.py b/reme/tool/memory/vector/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme/tool/memory/add_draft_and_retrieve_similar_memory.py b/reme/tool/memory/vector/add_draft_and_retrieve_similar_memory.py similarity index 96% rename from reme/tool/memory/add_draft_and_retrieve_similar_memory.py rename to reme/tool/memory/vector/add_draft_and_retrieve_similar_memory.py index fd3d10c8..040c475a 100644 --- a/reme/tool/memory/add_draft_and_retrieve_similar_memory.py +++ b/reme/tool/memory/vector/add_draft_and_retrieve_similar_memory.py @@ -2,10 +2,10 @@ from loguru import logger -from .base_memory_tool import BaseMemoryTool from .memory_handler import MemoryHandler -from ...core.schema import ToolCall, MemoryNode -from ...core.utils import deduplicate_memories +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import ToolCall, MemoryNode +from ....core.utils import deduplicate_memories class AddDraftAndRetrieveSimilarMemory(BaseMemoryTool): diff --git a/reme/tool/memory/add_memory.py b/reme/tool/memory/vector/add_memory.py similarity index 98% rename from reme/tool/memory/add_memory.py rename to reme/tool/memory/vector/add_memory.py index 833c2b54..3a049fc9 100644 --- a/reme/tool/memory/add_memory.py +++ b/reme/tool/memory/vector/add_memory.py @@ -2,9 +2,9 @@ from loguru import logger -from .base_memory_tool import BaseMemoryTool from .memory_handler import MemoryHandler -from ...core.schema import ToolCall +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import ToolCall class AddMemory(BaseMemoryTool): diff --git a/reme/tool/memory/delete_memory.py b/reme/tool/memory/vector/delete_memory.py similarity index 95% rename from reme/tool/memory/delete_memory.py rename to reme/tool/memory/vector/delete_memory.py index e95cfa40..da8094fe 100644 --- a/reme/tool/memory/delete_memory.py +++ b/reme/tool/memory/vector/delete_memory.py @@ -2,9 +2,9 @@ from loguru import logger -from .base_memory_tool import BaseMemoryTool from .memory_handler import MemoryHandler -from ...core.schema import ToolCall +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import ToolCall class DeleteMemory(BaseMemoryTool): diff --git a/reme/tool/memory/memory_handler.py b/reme/tool/memory/vector/memory_handler.py similarity index 97% rename from reme/tool/memory/memory_handler.py rename to reme/tool/memory/vector/memory_handler.py index 5b896449..8b6f8d78 100644 --- a/reme/tool/memory/memory_handler.py +++ b/reme/tool/memory/vector/memory_handler.py @@ -1,9 +1,9 @@ """Memory handler""" -from ...core.context import ServiceContext -from ...core.enumeration import MemoryType -from ...core.schema import MemoryNode -from ...core.vector_store import BaseVectorStore +from ....core.context import ServiceContext +from ....core.enumeration import MemoryType +from ....core.schema import MemoryNode +from ....core.vector_store import BaseVectorStore class MemoryHandler: diff --git a/reme/tool/memory/retrieve_memory.py b/reme/tool/memory/vector/retrieve_memory.py similarity index 96% rename from reme/tool/memory/retrieve_memory.py rename to reme/tool/memory/vector/retrieve_memory.py index db542643..af19bc1e 100644 --- a/reme/tool/memory/retrieve_memory.py +++ b/reme/tool/memory/vector/retrieve_memory.py @@ -2,10 +2,10 @@ from loguru import logger -from .base_memory_tool import BaseMemoryTool from .memory_handler import MemoryHandler -from ...core.schema import ToolCall, MemoryNode -from ...core.utils import deduplicate_memories +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import ToolCall, MemoryNode +from ....core.utils import deduplicate_memories class RetrieveMemory(BaseMemoryTool): diff --git a/reme/tool/memory/retrieve_recent_memory.py b/reme/tool/memory/vector/retrieve_recent_memory.py similarity index 92% rename from reme/tool/memory/retrieve_recent_memory.py rename to reme/tool/memory/vector/retrieve_recent_memory.py index 05b1a777..fc94d23d 100644 --- a/reme/tool/memory/retrieve_recent_memory.py +++ b/reme/tool/memory/vector/retrieve_recent_memory.py @@ -2,10 +2,10 @@ from loguru import logger -from .base_memory_tool import BaseMemoryTool from .memory_handler import MemoryHandler -from ...core.schema import ToolCall, MemoryNode -from ...core.utils import deduplicate_memories +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import ToolCall, MemoryNode +from ....core.utils import deduplicate_memories class RetrieveRecentMemory(BaseMemoryTool): diff --git a/reme/tool/memory/update_memory.py b/reme/tool/memory/vector/update_memory.py similarity index 98% rename from reme/tool/memory/update_memory.py rename to reme/tool/memory/vector/update_memory.py index ba4f8d5b..dbdf9ce6 100644 --- a/reme/tool/memory/update_memory.py +++ b/reme/tool/memory/vector/update_memory.py @@ -2,9 +2,9 @@ from loguru import logger -from .base_memory_tool import BaseMemoryTool from .memory_handler import MemoryHandler -from ...core.schema import ToolCall +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import ToolCall class UpdateMemory(BaseMemoryTool): diff --git a/reme/tool/memory/update_memory_v2.py b/reme/tool/memory/vector/update_memory_v2.py similarity index 98% rename from reme/tool/memory/update_memory_v2.py rename to reme/tool/memory/vector/update_memory_v2.py index b855c15b..5b8f0950 100644 --- a/reme/tool/memory/update_memory_v2.py +++ b/reme/tool/memory/vector/update_memory_v2.py @@ -2,9 +2,9 @@ from loguru import logger -from .base_memory_tool import BaseMemoryTool from .memory_handler import MemoryHandler -from ...core.schema import ToolCall +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import ToolCall class UpdateMemoryV2(BaseMemoryTool): diff --git a/tests/test_reme.py b/tests/test_reme.py index 52b11a64..7ac30e3c 100644 --- a/tests/test_reme.py +++ b/tests/test_reme.py @@ -17,34 +17,42 @@ async def test_reme(): { "role": "user", "content": "你好，我是张伟，今年28岁，是一名软件工程师。", + "time_created": "2026-01-29 10:00:00", }, { "role": "assistant", "content": "你好张伟！很高兴认识你。作为一名软件工程师，你主要从事什么方向的开发工作呢？", + "time_created": "2026-01-29 10:00:05", }, { "role": "user", "content": "我主要做后端开发，擅长Python和Go语言。最近在研究AI Agent相关的技术。", + "time_created": "2026-01-29 10:00:30", }, { "role": "assistant", "content": "很棒！Python和Go都是非常实用的语言。AI Agent是当前很热门的方向，你在这方面有什么具体的研究重点吗？", + "time_created": "2026-01-29 10:00:35", }, { "role": "user", "content": "我特别关注记忆系统的设计，希望能让AI Agent具有长期记忆能力。我的工作地点在北京，平时喜欢看技术博客和参加技术分享会。", + "time_created": "2026-01-29 10:01:00", }, { "role": "assistant", "content": "记忆系统确实是AI Agent的核心能力之一。北京有很多优秀的技术社区和活动，相信你能找到很多志同道合的朋友。", + "time_created": "2026-01-29 10:01:05", }, { "role": "user", "content": "是的，我每周末都会去参加一些技术沙龙。对了，我的邮箱是zhangwei@example.com，如果有好的技术资料可以发给我。", + "time_created": "2026-01-29 10:01:30", }, { "role": "assistant", "content": "好的，我记下了。保持学习的热情很重要，祝你在AI Agent领域的研究越来越深入！", + "time_created": "2026-01-29 10:01:35", }, ]