diff --git a/.gitignore b/.gitignore index efc9bea1..7ee1a323 100644 --- a/.gitignore +++ b/.gitignore @@ -34,7 +34,7 @@ test_compact_storage/* test_working_memory/* *.code-workspace local_vector_store/* -reme_local_memory/* +reme_profile/* chroma_vector_store/* bench_results/* meta_memory/* diff --git a/bench_old/halumem/eval_reme_simple_v4.py b/benchmark/halumem/eval_reme.py similarity index 70% rename from bench_old/halumem/eval_reme_simple_v4.py rename to benchmark/halumem/eval_reme.py index c68ee4d1..eb6b18ef 100644 --- a/bench_old/halumem/eval_reme_simple_v4.py +++ b/benchmark/halumem/eval_reme.py @@ -8,7 +8,7 @@ A modular evaluation pipeline that: 4. Generates comprehensive metrics Usage: - python bench/halumem/eval_reme_simple_v4.py \ + python benchmark/halumem/eval_reme.py \ --data_path /Users/yuli/workspace/HaluMem/data/HaluMem-Medium.jsonl \ --top_k 20 --user_num 100 --max_concurrency 20 """ @@ -26,14 +26,12 @@ from typing import Any from loguru import logger -from eval_tools import evaluation_for_question2, answer_question_with_memories -from reme_ai.core.enumeration import MemoryType -from reme_ai.core.schema import MemoryNode -from reme_ai.reme import ReMe +from reme.reme import ReMe # ==================== Configuration ==================== + @dataclass class EvalConfig: """Evaluation configuration parameters.""" @@ -42,7 +40,9 @@ class EvalConfig: user_num: int = 1 max_concurrency: int = 2 batch_size: int = 20 - output_dir: str = "bench_results/reme_simple_v4" + output_dir: str = "bench_results/reme" + eval_model_name: str = "qwen3-max" + algo_version: str = "v1" # ==================== Utilities ==================== @@ -106,7 +106,7 @@ class FileManager: def __init__(self, base_dir: str): self.base_dir = Path(base_dir) - self.tmp_dir = self.base_dir / "tmp" + self.tmp_dir = self.base_dir self.tmp_dir.mkdir(parents=True, exist_ok=True) def get_user_dir(self, user_name: str) -> Path: @@ -138,7 +138,7 @@ class FileManager: """Check if user has cached results.""" user_dir = self.get_user_dir(user_name) return any(f.name.startswith("session_") and f.suffix == ".json" - for f in user_dir.iterdir()) + for f in user_dir.iterdir()) def combine_results(self, output_file: str): """Combine all user session files into a single JSONL file.""" @@ -177,68 +177,145 @@ class FileManager: f_out.write(json.dumps(user_data, ensure_ascii=False) + "\n") +# ==================== Evaluation Functions ==================== + +async def answer_question_with_memories( + reme: ReMe, + question: str, + memories: str, + user_id: str = None, + model_name: str = "qwen3-30b-a3b-instruct-2507" +): + """ + Answer a question using retrieved memories with PROMPT_MEMZERO_JSON template. + + Args: + reme: ReMe instance with default_llm and prompt_handler + question: The question to answer + memories: The retrieved memories (formatted as context) + user_id: Optional user ID for context formatting + model_name: Model name to use for LLM request + + Returns: + dict with 'reasoning' and 'answer' fields + """ + # Format context with memories + if user_id: + context = reme.prompt_handler.prompt_format( + "TEMPLATE_MEMOS", + user_id=user_id, + memories=memories + ) + else: + context = f"Memories:\n{memories}" + + # Use PROMPT_MEMZERO_JSON template for structured JSON response + prompt = reme.prompt_handler.prompt_format( + "PROMPT_MEMZERO_JSON", + context=context, + question=question + ) + + result = await reme.llm.simple_request_for_json( + prompt=prompt, + model_name="qwen3-30b-a3b-instruct-2507" + ) + + return result + + +async def evaluation_for_question( + reme: ReMe, + question: str, + reference_answer: str, + key_memory_points: str, + response: str, + dialogue: str = None, + model_name: str = "qwen3-max" +): + """ + Question-Answering Evaluation with optional Dialogue Context. + + Args: + reme: ReMe instance with default_llm and prompt_handler + question: The question string to be evaluated. + reference_answer: The reference (gold-standard) answer. + key_memory_points: The memory points used to derive the reference answer. + response: The answer produced by the memory system. + dialogue: Optional formatted dialogue history (role, content, time_created). + model_name: Model name to use for LLM request + + Returns: + dict with 'reasoning' and 'evaluation_result' fields + """ + prompt = reme.prompt_handler.prompt_format( + "EVALUATION_PROMPT_FOR_QUESTION2", + question=question, + reference_answer=reference_answer, + key_memory_points=key_memory_points, + response=response, + dialogue=dialogue if dialogue else "" + ) + + result = await reme.llm.simple_request_for_json( + prompt=prompt, + model_name=model_name + ) + + return result + + # ==================== Memory Operations ==================== class MemoryProcessor: """Handles ReMe memory operations.""" - def __init__(self, reme: ReMe): + 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, - user_id: str, - messages: list[dict], - batch_size: int = 10000 - ) -> tuple[list[str], list[list[dict]], float]: + self, + user_id: str, + messages: list[dict], + batch_size: int = 10000 + ) -> tuple[list[str], list, float]: """ Add memories in batches using ReMe and return extracted memory contents. Returns: tuple: (extracted_memories, agent_messages, total_duration_ms) """ - added_memories: list[MemoryNode] = [] - deleted_memories: list[str] = [] - all_agent_messages: list = [] + extracted_memories = [] + summary_messages = [] total_duration_ms = 0 for i in range(0, len(messages), batch_size): batch = messages[i:i + batch_size] start = time.time() - memory_nodes, agent_messages, success = await self.reme.summary_v4( + # Use new summary API + result = await self.reme.summary_memory( messages=batch, - user_id=user_id + user_name=user_id, + version=self.algo_version, + return_dict=True, ) duration_ms = (time.time() - start) * 1000 total_duration_ms += duration_ms - # Save agent messages for this batch - if agent_messages: - all_agent_messages.extend(agent_messages) + extracted_memories.extend([m.model_dump(exclude_none=True) for m in result["answer"]]) + summary_messages.extend([m.simple_dump(enable_argument_dict=True) for m in result["messages"]]) - if memory_nodes: - for node in memory_nodes: - if isinstance(node, MemoryNode) and node.memory_type == MemoryType.HISTORY: - continue - - if isinstance(node, MemoryNode): - added_memories.append(node) - - if isinstance(node, str): - deleted_memories.append(node) - - extracted_memories = deleted_memories - extracted_memories += ["[delete]" + n.format_memory() for n in added_memories if n.memory_id in deleted_memories] - extracted_memories += ["[add]" + n.format_memory() for n in added_memories if n.memory_id not in deleted_memories] - return extracted_memories, all_agent_messages, total_duration_ms + return extracted_memories, summary_messages, total_duration_ms async def search_memory( - self, - query: str, - user_id: str, - top_k: int = 20 + self, + query: str, + user_id: str, + top_k: int = 20 ) -> tuple[dict, list, float]: """ Search memory using ReMe and return structured answer with reasoning. @@ -249,22 +326,32 @@ class MemoryProcessor: """ start = time.time() - # Retrieve memories from ReMe - memories_response, agent_messages, success = await self.reme.retrieve_v4( + # Retrieve memories from ReMe using new API + result = await self.reme.retrieve_memory( query=query, - user_id=user_id, - top_k=top_k + retrieve_top_k=top_k, + user_name=user_id, + version=self.algo_version, + return_dict=True, ) + # Extract memories from response + memories = result["answer"] + 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 answer_result = await answer_question_with_memories( + reme=self.reme, question=query, - memories=memories_response, - user_id=user_id + memories=memories, + user_id=user_id, + model_name=self.eval_model_name ) # Add original memories to the result - answer_result["memories"] = memories_response + answer_result["memories"] = memories + answer_result["retrieved_nodes"] = retrieved_nodes duration_ms = (time.time() - start) * 1000 return answer_result, agent_messages, duration_ms @@ -275,17 +362,19 @@ class MemoryProcessor: class QuestionAnsweringEvaluator: """Evaluates question answering performance.""" - def __init__(self, memory_processor: MemoryProcessor, top_k: int): + def __init__(self, memory_processor: MemoryProcessor, reme: ReMe, top_k: int, eval_model_name: str = "qwen3-max"): self.memory_processor = memory_processor + self.reme = reme self.top_k = top_k + self.eval_model_name = eval_model_name async def evaluate_questions( - self, - questions: list[dict], - user_name: str, - uuid: str, - session_id: int, - formatted_dialogue: str + self, + questions: list[dict], + user_name: str, + uuid: str, + session_id: int, + formatted_dialogue: str ) -> list[dict]: """Evaluate all questions for a session.""" results = [] @@ -301,15 +390,18 @@ class QuestionAnsweringEvaluator: system_answer = answer_dict.get("answer", "") system_reasoning = answer_dict.get("reasoning", "") retrieved_memories = answer_dict.get("memories", "") + retrieved_nodes = answer_dict.get("retrieved_nodes", "") # Evaluate response evidence_text = "\n".join([e["memory_content"] for e in qa["evidence"]]) - eval_result = await evaluation_for_question2( - qa["question"], - qa["answer"], - evidence_text, - system_answer, - formatted_dialogue + eval_result = await evaluation_for_question( + reme=self.reme, + question=qa["question"], + reference_answer=qa["answer"], + key_memory_points=evidence_text, + response=system_answer, + dialogue=formatted_dialogue, + model_name=self.eval_model_name ) # Build result record @@ -320,7 +412,8 @@ class QuestionAnsweringEvaluator: "system_response": system_answer, "system_reasoning": system_reasoning, "retrieved_memories": retrieved_memories, - "retrieve_messages": [m.model_dump() for m in agent_messages], + "retrieved_nodes": retrieved_nodes, + "retrieve_messages": agent_messages, "search_duration_ms": duration_ms, "result_type": eval_result.get("evaluation_result"), "question_answering_reasoning": eval_result.get("reasoning", "") @@ -418,25 +511,55 @@ class MetricsAggregator: # ==================== Main Pipeline ==================== -class HaluMemEvaluatorV4: +class HaluMemEvaluator: + """HaluMem evaluator with proper resource management.""" def __init__(self, config: EvalConfig): self.config = config self.reme = ReMe() + + # Load evaluation prompts into ReMe's prompt handler + prompts_yaml_path = Path(__file__).parent / "eval_reme.yaml" + self.reme.prompt_handler.load_prompt_by_file(prompts_yaml_path) + self.file_manager = FileManager(config.output_dir) - self.memory_processor = MemoryProcessor(self.reme) + self.memory_processor = MemoryProcessor( + self.reme, + config.eval_model_name, + config.algo_version + ) self.qa_evaluator = QuestionAnsweringEvaluator( self.memory_processor, - config.top_k + self.reme, + config.top_k, + config.eval_model_name ) self.data_loader = DataLoader() + async def __aenter__(self): + """Async context manager entry.""" + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + """Async context manager exit with cleanup.""" + await self.reme.close() + return False + + def __enter__(self): + """Sync context manager entry.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Sync context manager exit with cleanup.""" + self.reme.close_sync() + return False + async def process_session( - self, - session: dict, - session_id: int, - user_name: str, - uuid: str + self, + session: dict, + session_id: int, + user_name: str, + uuid: str ) -> dict: """Process a single session using ReMe.""" session_data = { @@ -463,7 +586,7 @@ class HaluMemEvaluatorV4: session_data.update({ "dialogue": dialogue, "extracted_memories": extracted_memories, - "summary_messages": [m.model_dump() for m in agent_messages], + "summary_messages": agent_messages, "add_dialogue_duration_ms": duration_ms }) @@ -509,6 +632,20 @@ class HaluMemEvaluatorV4: """Run the complete evaluation pipeline using ReMe.""" start_time = time.time() + # Load user data first to get user names + all_users = self.data_loader.load_jsonl(self.config.data_path) + users_to_process = all_users[:self.config.user_num] + + # Extract all user names and delete all profiles + all_user_names = [ + self.data_loader.extract_user_name(user_data["persona_info"]) + for user_data in all_users + ] + if all_user_names: + for user_name in all_user_names: + self.reme.get_profile_handler(user_name).delete_all() + logger.info(f"Deleted all profiles for {len(all_user_names)} users") + # Clear existing data await self.reme.vector_store.delete_all() @@ -519,10 +656,6 @@ class HaluMemEvaluatorV4: logger.info(f"Cleared meta_memory directory: {meta_memory_path}") meta_memory_path.mkdir(parents=True, exist_ok=True) - # Load user data - all_users = self.data_loader.load_jsonl(self.config.data_path) - users_to_process = all_users[:self.config.user_num] - print("\n" + "=" * 80) print("HALUMEM EVALUATION - REME - QUESTION ANSWERING") print(f"Users: {len(users_to_process)} | Concurrency: {self.config.max_concurrency}") @@ -631,22 +764,46 @@ class HaluMemEvaluatorV4: # ==================== Entry Point ==================== -def main( - data_path: str, - top_k: int = 20, - user_num: int = 1, - max_concurrency: int = 2 +async def main_async( + data_path: str, + top_k: int, + user_num: int, + max_concurrency: int, + eval_model_name: str = "qwen3-max", + algo_version: str = "v1" ): - """Main entry point for ReMe evaluation.""" + """Main async entry point for ReMe evaluation with proper resource cleanup.""" config = EvalConfig( data_path=data_path, top_k=top_k, user_num=user_num, - max_concurrency=max_concurrency + max_concurrency=max_concurrency, + eval_model_name=eval_model_name, + algo_version=algo_version ) - evaluator = HaluMemEvaluatorV4(config) - asyncio.run(evaluator.run_evaluation()) + # Use async context manager for automatic cleanup + async with HaluMemEvaluator(config) as evaluator: + await evaluator.run_evaluation() + + +def main( + data_path: str, + top_k: int, + user_num: int, + max_concurrency: int, + eval_model_name: str = "qwen3-max", + algo_version: str = "v1" +): + """Main entry point for ReMe evaluation.""" + asyncio.run(main_async( + data_path=data_path, + top_k=top_k, + user_num=user_num, + max_concurrency=max_concurrency, + eval_model_name=eval_model_name, + algo_version=algo_version + )) if __name__ == "__main__": @@ -676,8 +833,21 @@ if __name__ == "__main__": parser.add_argument( "--max_concurrency", type=int, - default=2, - help="Maximum concurrent user processing (default: 2)" + default=100, + help="Maximum concurrent user processing (default: 100)" + ) + parser.add_argument( + "--eval_model_name", + type=str, + default="qwen3-max", + # 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() @@ -686,5 +856,7 @@ if __name__ == "__main__": data_path=args.data_path, top_k=args.top_k, user_num=args.user_num, - max_concurrency=args.max_concurrency + max_concurrency=args.max_concurrency, + eval_model_name=args.eval_model_name, + algo_version=args.algo_version ) diff --git a/benchmark/halumem/eval_reme.yaml b/benchmark/halumem/eval_reme.yaml new file mode 100644 index 00000000..56dde2db --- /dev/null +++ b/benchmark/halumem/eval_reme.yaml @@ -0,0 +1,548 @@ +TEMPLATE_MEMOS: | + Memories for user {user_id}: + {memories} + +PROMPT_MEMZERO_JSON: | + # CONTEXT: + {context} + + # CONTEXT PRIORITY: + When the context contains information from multiple sources, follow this strict priority order: + 1. **Historical Dialogue** (highest priority) - Direct conversation content + 2. **Extracted Memories** (medium priority) - Summarized memory points + 3. **User Profile** (lowest priority) - General user information + + # Question: + {question} + + # OUTPUT FORMAT: + Do not hallucinate; strictly answer the user's question based on the content of the CONTEXT. + Please provide your response in the following JSON format: + + ```json + {{ + "reasoning": "reasoning content", + "answer": "Provide a detailed answer" + }} + ``` + +PROMPT_MEMZERO_JSON2: | + # CONTEXT: + {context} + + # CONTEXT PRIORITY: + When the context contains information from multiple sources, follow this strict priority order: + 1. **Historical Dialogue** (highest priority) - Direct conversation content + 2. **Extracted Memories** (medium priority) - Summarized memory points + 3. **User Profile** (lowest priority) - General user information + + # Question: + {question} + + # OUTPUT FORMAT: + Do not hallucinate; strictly answer the user's question based on the content of the CONTEXT. + Please provide your response in the following JSON format: + + ```json + {{ + "reasoning": "reasoning content", + "answer": "Provide a detailed answer" + }} + ``` + +PROMPT_MEMZERO: | + You are an intelligent memory assistant tasked with retrieving accurate information from conversation memories. + + # CONTEXT: + You have access to memories from two speakers in a conversation. These memories contain + timestamped information that may be relevant to answering the question. + + # INSTRUCTIONS: + 1. Carefully analyze all provided memories from both speakers + 2. Pay special attention to the timestamps to determine the answer + 3. If the question asks about a specific event or fact, look for direct evidence in the memories + 4. If the memories contain contradictory information, prioritize the most recent memory + 5. If there is a question about time references (like "last year", "two months ago", etc.), + calculate the actual date based on the memory timestamp. For example, if a memory from + 4 May 2022 mentions "went to India last year," then the trip occurred in 2021. + 6. Always convert relative time references to specific dates, months, or years. For example, + convert "last year" to "2022" or "two months ago" to "March 2023" based on the memory + timestamp. Ignore the reference while answering the question. + 7. Focus only on the content of the memories from both speakers. Do not confuse character + names mentioned in memories with the actual users who created those memories. + 8. The answer should be less than 5-6 words. + + # APPROACH (Think step by step): + 1. First, examine all memories that contain information related to the question + 2. Examine the timestamps and content of these memories carefully + 3. Look for explicit mentions of dates, times, locations, or events that answer the question + 4. If the answer requires calculation (e.g., converting relative time references), show your work + 5. Formulate a precise, concise answer based solely on the evidence in the memories + 6. Double-check that your answer directly addresses the question asked + 7. Ensure your final answer is specific and avoids vague time references + + {context} + + Question: {question} + + Answer: + +PROMPT_ZEP: | + You are an intelligent memory assistant tasked with retrieving accurate information from conversation memories. + + # CONTEXT: + You have access to memories from a conversation. These memories contain + timestamped information that may be relevant to answering the question. + + # INSTRUCTIONS: + 1. Carefully analyze all provided memories + 2. Pay special attention to the timestamps to determine the answer + 3. If the question asks about a specific event or fact, look for direct evidence in the memories + 4. If the memories contain contradictory information, prioritize the most recent memory + 5. If there is a question about time references (like "last year", "two months ago", etc.), + calculate the actual date based on the memory timestamp. For example, if a memory from + 4 May 2022 mentions "went to India last year," then the trip occurred in 2021. + 6. Always convert relative time references to specific dates, months, or years. For example, + convert "last year" to "2022" or "two months ago" to "March 2023" based on the memory + timestamp. Ignore the reference while answering the question. + 7. Focus only on the content of the memories. Do not confuse character + names mentioned in memories with the actual users who created those memories. + 8. The answer should be less than 5-6 words. + + # APPROACH (Think step by step): + 1. First, examine all memories that contain information related to the question + 2. Examine the timestamps and content of these memories carefully + 3. Look for explicit mentions of dates, times, locations, or events that answer the question + 4. If the answer requires calculation (e.g., converting relative time references), show your work + 5. Formulate a precise, concise answer based solely on the evidence in the memories + 6. Double-check that your answer directly addresses the question asked + 7. Ensure your final answer is specific and avoids vague time references + + Context: + + {context} + + Question: {question} + Answer: + +PROMPT_MEMOS: | + You are a knowledgeable and helpful AI assistant. + + # CONTEXT: + You have access to memories from two speakers in a conversation. These memories contain + timestamped information that may be relevant to answering the question. + + # INSTRUCTIONS: + 1. Carefully analyze all provided memories. Synthesize information across different entries if needed to form a complete answer. + 2. Pay close attention to the timestamps to determine the answer. If memories contain contradictory information, the **most recent memory** is the source of truth. + 3. If the question asks about a specific event or fact, look for direct evidence in the memories. + 4. Your answer must be grounded in the memories. However, you may use general world knowledge to interpret or complete information found within a memory (e.g., identifying a landmark mentioned by description). + 5. If the question involves time references (like "last year", "two months ago", etc.), you **must** calculate the actual date based on the memory's timestamp. For example, if a memory from 4 May 2022 mentions "went to India last year," then the trip occurred in 2021. + 6. Always convert relative time references to specific dates, months, or years in your final answer. + 7. Do not confuse character names mentioned in memories with the actual users who created them. + 8. The answer must be brief (under 5-6 words) and direct, with no extra description. + + # APPROACH (Think step by step): + 1. First, examine all memories that contain information related to the question. + 2. Synthesize findings from multiple memories if a single entry is insufficient. + 3. Examine timestamps and content carefully, looking for explicit dates, times, locations, or events. + 4. If the answer requires calculation (e.g., converting relative time references), perform the calculation. + 5. Formulate a precise, concise answer based on the evidence from the memories (and allowed world knowledge). + 6. Double-check that your answer directly addresses the question asked and adheres to all instructions. + 7. Ensure your final answer is specific and avoids vague time references. + + {context} + + Question: {question} + + Answer: + +PROMPT_MEMOBASE: | + You are an intelligent memory assistant tasked with retrieving accurate information from conversation memories. + + # CONTEXT: + You have access to memories from two speakers in a conversation. These memories contain + timestamped information that may be relevant to answering the question. + + # INSTRUCTIONS: + 1. Carefully analyze all provided memories from both speakers + 2. Pay special attention to the timestamps to determine the answer + 3. If the question asks about a specific event or fact, look for direct evidence in the memories + 4. If the memories contain contradictory information, prioritize the most recent memory + 5. If there is a question about time references (like "last year", "two months ago", etc.), calculate the actual date based on the memory timestamp. For example, if a memory from 4 May 2022 mentions "went to India last year," then the trip occurred in 2021. + 6. Always convert relative time references to specific dates, months, or years. For example, convert "last year" to "2022" or "two months ago" to "March 2023" based on the memory timestamp. Ignore the reference while answering the question. + 7. Focus only on the content of the memories from both speakers. Do not confuse character names mentioned in memories with the actual users who created those memories. + 8. The answer should be less than 5-6 words. + + # APPROACH (Think step by step): + 1. First, examine all memories that contain information related to the question + 2. Examine the timestamps and content of these memories carefully + 3. Look for explicit mentions of dates, times, locations, or events that answer the question + 4. If the answer requires calculation (e.g., converting relative time references), show your work + 5. Formulate a precise, concise answer based solely on the evidence in the memories + 6. Double-check that your answer directly addresses the question asked + 7. Ensure your final answer is specific and avoids vague time references + + {context} + + Question: {question} + + Answer: + + +EVALUATION_PROMPT_FOR_MEMORY_INTEGRITY: | + You are a strict **"Memory Integrity" evaluator**. + Your core task is to assess whether an AI memory system has **missed any key memory points** after processing a conversation. This evaluation measures the system’s **memory integrity**, i.e., its ability to resist **amnesia** or **omission**. + + # Evaluation Context & Data: + + 1. **Extracted Memories:** + These are all the memory items actually extracted by the memory system. + {memories} + + 2. **Expected Memory Point:** + The key memory point that *should* have been extracted. + {expected_memory_point} + + # Evaluation Instructions: + + 1. For each **Expected Memory Point**, search within the **Extracted Memories** list for corresponding or related information. Ignore unrelated items. + 2. Based on the following scoring rubric, rate how well the memory system captured the **Expected Memory Point** and provide a detailed explanation. + + # Scoring Rubric: + + * **2:** Fully covered or implied. + One or more items in “Extracted Memories” fully cover or logically imply all information in the “Expected Memory Point.” + + * **1:** Partially covered or mentioned. + Some information in “Extracted Memories” mentions part of the “Expected Memory Point,” but key information is missing, inaccurate, or slightly incorrect. + + * **0:** Not mentioned or incorrect. + “Extracted Memories” contains no mention of the “Expected Memory Point,” or the corresponding information is entirely wrong. + + # Scoring Notes: + + * For **compound Expected Memory Points** (with multiple elements such as person/event/time/location/preference, etc.): + + * All elements correct → **2 points** + * Some elements correct / uncertain → **1 point** + * Key elements missing or wrong → **0 points** + + * Semantic matching is acceptable; exact wording is **not** required. + + * If “Extracted Memories” contains **conflicting information**, assign the **best possible coverage score** and mention the conflict in your reasoning. + + * Extra or stylistically different memories do **not** reduce the score; only the coverage of the **Expected Memory Point** matters. + + * For uncertain wording (“might,” “probably,” “tends to,” etc.): + + * If the Expected Memory Point is a definite statement, usually assign **1 point**. + + * If critical fields (e.g., time, entity name, relationship) are partly wrong but others match → **1 point**. + + * If all key fields are wrong or missing → **0 points**. + + # Output Format: + + Please output your result in the following JSON format: + + ```json + {{ + "reasoning": "Provide a concise justification for the score", + "score": "2|1|0" + }} + ``` + +EVALUATION_PROMPT_FOR_MEMORY_ACCURACY: | + You are a **Dialogue Memory Accuracy Evaluator.** Your task is to evaluate the **accuracy** of a memory extracted by an AI memory system, based on three given inputs: the dialogue content, the *target (gold)* memory points (the correct annotated memories), and the *candidate* memory to be evaluated. The goal is to output a **structured evaluation result**. + + # Input Content + + * **Dialogue:** + {dialogue} + + * **Golden Memories (Target Memory Points):** + The correct memory points pre-annotated for this dialogue in the evaluation dataset. + {golden_memories} + + * **Candidate Memory:** + The memory extracted by the system to be evaluated. + {candidate_memory} + + # Evaluation Principles and Definitions + + ### 1) Support / Entailment + + * An **information point** (atomic fact) in the candidate memory is considered *supported* if it can be directly stated or semantically entailed (via synonym, paraphrase, or equivalent expression) by the *Dialogue* or *Golden Memories*. + * Only the given dialogue and golden memories can be used for judgment — **no external knowledge** or assumptions are allowed. + Any information not appearing in or inferable from these two sources is considered *unsupported*. + * Pay careful attention to **negation**, **quantities**, **time**, and **subjects**. + If the candidate statement contradicts the dialogue or golden memories, it is considered a **conflict**. + + ### 2) Memory Accuracy Score (integer: 0 / 1 / 2) + + * **2 points:** Every information point in the candidate memory is supported by the dialogue or golden memories, with **no contradictions or hallucinations**. + * **1 point:** The candidate memory is *partially correct* (at least one supported information point) but also includes *unsupported* or *contradictory* content. + * **0 points:** The candidate memory is **entirely unsupported or contradictory** to the sources (i.e., a “hallucinated memory”). + + > Note: + > + > * If a candidate memory contains multiple information points, **any unsupported or contradictory element** prevents a full score (2). + > * If both supported and unsupported/conflicting content appear, assign a score of **1**. + + ### 3) Inclusion in Golden Memories (Boolean field-level judgment) + + **Definition:** + + * **Atomic information point:** the smallest factual unit in the candidate memory (e.g., *name = Li Si*, *age = 25*, *location = Beijing*, *preference = coffee*, *budget ≤ 2000*, *meeting_time = Wednesday 10:00*, *tool = Zoom*, etc.). + * **Field / Slot:** the semantic dimension of an information point (e.g., *name*, *age*, *residence*, *food preference*, *budget*, *meeting time*, *meeting tool*, etc.). + + **Judgment Rules (independent of correctness):** + + * **true:** + Every atomic information point in the candidate memory has a corresponding **field** in the golden memories (allowing for synonyms, paraphrases, or equivalent expressions; ignore value, polarity, or quantity differences). + + * Note: A single field in the gold list may match multiple candidate points (e.g., multiple “drink preference” facts can be covered by one “drink preference” field in gold). + * **false:** + If **any** atomic information point’s field in the candidate memory cannot be found in the golden memories, mark as *false*. + + **Important Notes:** + + * Field matching is restricted to fields that are **explicitly present or semantically recognizable** in the golden memories — no external knowledge may be used to expand the field set. + * Differences in **values** (e.g., “Zhang San” vs. “Li Si”), **polarity** (like/dislike), or **exact number/time** do **not** affect this Boolean judgment. + + # Evaluation Procedure + + For each candidate memory: + + 1. **Decompose** it into atomic information points (e.g., name, number, location, preference). + 2. For each information point, **search** the dialogue and golden memories for supporting or contradictory evidence. + 3. Assign the **accuracy_score** (0 / 1 / 2) according to the rules above. + 4. Determine **is_included_in_golden_memories (true/false)**: + + * Identify each information point’s field; + * If *all* fields exist in the golden memories, mark as *true*; otherwise, *false*. + 5. Provide a **concise Chinese explanation** in `"reason"`, citing key evidence (short excerpts allowed), and clearly state any unsupported or contradictory parts if applicable. + + # Output Format (strictly required) + + Output **only one JSON object**, with the following three fields: + + * `"accuracy_score"`: `"0"` or `"1"` or `"2"` + * `"is_included_in_golden_memories"`: `"true"` or `"false"` + * `"reason"`: `"brief explanation in Chinese"` + + Do **not** include any other text, explanation, or fields. + Do **not** include the candidate memory text inside the JSON. + + Please output **only** the following JSON (in a code block): + + ```json + {{ + "accuracy_score": "2 | 1 | 0", + "is_included_in_golden_memories": "true | false", + "reason": "Brief explanation in Chinese" + }} + ``` + +EVALUATION_PROMPT_FOR_UPDATE_MEMORY: | + Your task is to **evaluate the update accuracy** of an AI memory system. + Based on the information provided below, determine whether the system-generated **“Generated Memories”** correctly **includes** the **Target Memory for Update**. + + # Background Information + + The following information is provided for evaluation: + + 1. **Generated Memories:** + This is the list of memory points generated by the system after the current dialogue. + {memories} + + 2. **Target Memory for Update:** + This is the correct, updated version of the memory point that should have been produced — the one we focus on in this evaluation. + {updated_memory} + + 3. **Original Memory Content:** + This is the original version of the target memory before the update. + {original_memory} + + # Evaluation Criteria + + Please make your judgment **strictly based on the content update of the “Target Memory for Update.”** + Use the following categories: + + ### Correct Update + + * **Generated Memories** **contains all information points** from the “Target Memory for Update,” accurately and completely reflecting the intended update. + * **Key fields** (e.g., date, time, values, proper nouns, etc.) must match exactly. + * The **original memory** is effectively replaced or marked as outdated. + * Synonymous or slightly rephrased expressions are acceptable. + + ### Hallucinated Update + + * **Factual error:** The **Generated Memories** includes a new memory related to the “Target Memory for Update,” but its content contains factual mistakes or contradictions compared to the correct update. + + ### Omitted Update + + * **Completely omitted:** The **Generated Memories** contains no new memory related to the “Target Memory for Update.” + * **Partially omitted:** A related new memory was generated in **Generated Memories**, but it **misses key information** that should have been included. + + ### Other + + Used for update failures that do **not clearly fall** into the above categories of “Hallucination” or “Omission.” + + # Output Requirements + + Please return your evaluation strictly in the following JSON format and provide a concise explanation. + + ```json + {{ + "reason": "Briefly explain your reasoning here and why it fits this category.", + "evaluation_result": "Correct | Hallucination | Omission | Other" + }} + ``` + +EVALUATION_PROMPT_FOR_QUESTION: | + You are an **evaluation expert for AI memory system question answering**. + Based **only** on the provided **“Question”**, **“Reference Answer”**, and **“Key Memory Points”** (the essential facts needed to derive the reference answer), strictly evaluate the **accuracy** of the **“Memory System Response.”** Classify it as one of **“Correct”**, **“Hallucination”**, or **“Omission.”** Do **not** use any external knowledge or subjective inference. Finally, output your judgment **strictly** in the specified JSON format. + + # Evaluation Criteria + + ## Answer Type Classification + + ### 1. Correct + + * The “Memory System Response” accurately answers the “Question,” and its content is **semantically equivalent** to the “Reference Answer.” + * It contains **no contradictions** with the “Key Memory Points” or “Reference Answer.” + * It introduces **no unsupported details** beyond the “Key Memory Points” that could alter the conclusion. + * Synonyms, paraphrasing, and reasonable summarization are acceptable. + + ### 2. Hallucination + + * The “Memory System Response” includes information or facts that **contradict or are inconsistent** with the “Reference Answer” or the “Key Memory Points.” + * When the “Reference Answer” is labeled as *unknown/uncertain*, yet the response provides a specific verifiable fact or conclusion. + * Extra irrelevant information that does **not change** the conclusion is **not** considered hallucination by itself; however, if it **changes or misleads** the conclusion, or **contradicts** the “Key Memory Points,” it should be judged as a **Hallucination**. + + ### 3. Omission + + * The response is **incomplete** compared to the “Reference Answer.” + * It explicitly states “don’t know,” “can’t remember,” or “no related memory,” even though relevant information exists in the “Key Memory Points.” + * For multi-element questions, **all elements must be correct and present**; omission of **any** element is considered an **Omission**. + + ## Priority Rules (Conflict Handling) + + * If the response contains **both missing necessary information** and **fabricated/contradictory information**, classify it as **Hallucination**. + * If there is **no fabrication/contradiction** but some necessary information is missing, classify it as **Omission**. + * Only when the meaning is **fully equivalent** to the reference answer should it be classified as **Correct**. + + ## Detailed Guidelines and Tolerance + + * Equivalent expressions of numbers, times, and units are acceptable, but the **numerical values themselves must not differ**. + * For multi-element questions, **all elements must be complete and accurate**; missing any element counts as **Omission**. + * If the reference answer is *“unknown / cannot be determined”* and the system provides a definite fact, that is a **Hallucination**. + If the system also answers *“unknown”* (without guessing), it may be **Correct**. + * The evaluation must rely **only** on the *Reference Answer*, *Key Memory Points*, and *System Response* — no external context, world knowledge, or speculative reasoning is allowed. + + # Information for Evaluation + + * **Question:** + {question} + + * **Reference Answer:** + {reference_answer} + + * **Key Memory Points:** + {key_memory_points} + + * **Memory System Response:** + {response} + + # Output Requirements + + Please provide your evaluation result **strictly** in the JSON format below. + Do **not** add any extra explanation or comments outside the JSON block. + + ```json + {{ + "reasoning": "Provide a concise and traceable evaluation rationale: first compare the system’s response with the Key Memory Points (which were correctly used, which were missing, and whether there was any fabrication/contradiction), then assess its consistency with the Reference Answer, and finally state the classification basis.", + "evaluation_result": "Correct | Hallucination | Omission" + }} + ``` + + +EVALUATION_PROMPT_FOR_QUESTION2: | + You are an **evaluation expert for AI memory system question answering**. + + Based **only** on the provided **"Question"**, **"Reference Answer"**, and **"Key Memory Points"** (the essential facts needed to derive the reference answer), strictly evaluate the **accuracy** of the **"Memory System Response."** Classify it as one of **"Correct"**, **"Hallucination"**, or **"Omission."** Do **not** use any external knowledge or subjective inference. Finally, output your judgment **strictly** in the specified JSON format. + + # Evaluation Criteria + + ## Answer Type Classification + + ### 1. Correct + + * The "Memory System Response" accurately answers the "Question," and its content is **semantically equivalent** to the "Reference Answer." + * It contains **no contradictions** with the "Key Memory Points" or "Reference Answer." + * **Extra details not present in the Key Memory Points are allowed and should not be penalized**, as long as they: + - Do not contradict the Key Memory Points or Reference Answer + - Do not change or mislead the core conclusion + - Are reasonable additional context that the memory system may have retained from the conversation + * The memory system may have stored additional information beyond the Key Memory Points. Such extra information should be treated as **supplementary context** rather than hallucination, provided it does not conflict with the core answer. + * Synonyms, paraphrasing, and reasonable summarization are acceptable. + + ### 2. Hallucination + + * The "Memory System Response" includes information or facts that **contradict or are inconsistent** with the "Reference Answer" or the "Key Memory Points." + * The response provides information that **directly contradicts** known facts from the Key Memory Points. + * When the "Reference Answer" is labeled as *unknown/uncertain*, yet the response provides a specific verifiable fact or conclusion. + * **Important:** Extra information that is NOT in Key Memory Points is **NOT automatically a hallucination**. Only classify as hallucination if the extra information: + - Directly contradicts the Key Memory Points or Reference Answer + - Changes or misleads the core conclusion in a way that makes the answer incorrect + - Provides a definitive answer when the Reference Answer indicates uncertainty + + ### 3. Omission + + * The response is **incomplete** compared to the "Reference Answer." + * It explicitly states "don't know," "can't remember," or "no related memory," even though relevant information exists in the "Key Memory Points." + * For multi-element questions, **all elements must be correct and present**; omission of **any** element is considered an **Omission**. + + ## Priority Rules (Conflict Handling) + + * If the response contains **both missing necessary information** and **fabricated/contradictory information**, classify it as **Hallucination**. + * If there is **no fabrication/contradiction** but some necessary information is missing, classify it as **Omission**. + * If the core answer is correct and complete, classify as **Correct** even if there are extra details not in Key Memory Points (as long as they don't contradict or mislead). + + ## Detailed Guidelines and Tolerance + + * Equivalent expressions of numbers, times, and units are acceptable, but the **numerical values themselves must not differ**. + * For multi-element questions, **all elements must be complete and accurate**; missing any element counts as **Omission**. + * If the reference answer is *"unknown / cannot be determined"* and the system provides a definite fact, that is a **Hallucination**. + If the system also answers *"unknown"* (without guessing), it may be **Correct**. + * **Focus on evaluating whether the core answer to the question is correct**, not whether the response is limited to only the Key Memory Points. + * Extra contextual information (e.g., additional preferences, related details) should be viewed as enrichment, not as errors, unless they contradict or mislead. + + # Information for Evaluation + + * **Question:** + {question} + + * **Reference Answer:** + {reference_answer} + + * **Key Memory Points:** + {key_memory_points} + + * **Memory System Response:** + {response} + + # Output Requirements + + Please provide your evaluation result **strictly** in the JSON format below. + Do **not** add any extra explanation or comments outside the JSON block. + + ```json + {{ + "reasoning": "Provide a concise and traceable evaluation rationale: first verify that the system's response correctly includes all required elements from the Reference Answer, then check if any information contradicts the Key Memory Points or Reference Answer. Extra details not in Key Memory Points should be noted but not penalized unless they contradict or mislead. Finally state the classification basis.", + "evaluation_result": "Correct | Hallucination | Omission" + }} + ``` + """ \ No newline at end of file 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/__init__.py b/reme/agent/__init__.py index 45fed6e8..a2a0aa17 100644 --- a/reme/agent/__init__.py +++ b/reme/agent/__init__.py @@ -1,7 +1,9 @@ """A simple chatbot.""" from . import chat +from . import memory __all__ = [ "chat", + "memory", ] 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/base_memory_agent.py b/reme/agent/memory/base_memory_agent.py index 361a6676..6b97a1c2 100644 --- a/reme/agent/memory/base_memory_agent.py +++ b/reme/agent/memory/base_memory_agent.py @@ -1,9 +1,7 @@ """Base memory agent for handling memory operations with tool-based reasoning.""" +import json from abc import ABCMeta -from typing import Literal - -from loguru import logger from ...core.enumeration import MemoryType from ...core.op import BaseReact @@ -15,40 +13,6 @@ class BaseMemoryAgent(BaseReact, metaclass=ABCMeta): memory_type: MemoryType | None = None - @staticmethod - async def read_meta_memories(meta_memories: list[dict]) -> str: - """Read and format meta memory information from the provided metadata list.""" - from ...tool.memory import ReadMetaMemory - - meta_memory_info = ReadMetaMemory().format_memory_metadata(meta_memories) - logger.info(f"meta_memory_info={meta_memory_info}") - return meta_memory_info - - async def read_user_profile(self, show_id: Literal["profile", "history"] = "profile") -> str: - """Read current user profile.""" - from ...tool.memory import ReadUserProfile - - read_tool = ReadUserProfile(show_id=show_id) - await read_tool.call(memory_target=self.memory_target, service_context=self.service_context) - return str(read_tool.response.answer) - - async def add_history_node(self) -> MemoryNode: - """Add history node""" - from ...tool.memory import AddHistory - - add_history_tool = AddHistory() - await add_history_tool.call( - messages=self.messages, - description=self.description, - service_context=self.service_context, - ) - return add_history_tool.context.history_node - - @property - def memory_target(self) -> str: - """memory_target""" - return self.context.get("memory_target", "") - @property def query(self) -> str: """query""" @@ -64,6 +28,11 @@ class BaseMemoryAgent(BaseReact, metaclass=ABCMeta): """description""" return self.context.get("description", "") + @property + def memory_target(self) -> str: + """memory_target""" + return self.context.memory_target + @property def history_node(self) -> MemoryNode: """Returns the history node.""" @@ -73,3 +42,27 @@ class BaseMemoryAgent(BaseReact, metaclass=ABCMeta): def author(self) -> str: """Returns the LLM model name as the author identifier.""" return self.llm.model_name + + @property + def retrieved_nodes(self) -> list[MemoryNode]: + """Returns the retrieved nodes.""" + if "retrieved_nodes" not in self.context: + self.context.retrieved_nodes = [] + return self.context.retrieved_nodes + + @property + def memory_target_type_mapping(self) -> dict[str, MemoryType]: + """Get the memory target type mapping from context.""" + return self.context.service_context.memory_target_type_mapping + + @property + def meta_memory_info(self) -> str: + """Get the meta memory info from context.""" + lines = [] + for memory_target, memory_type in self.memory_target_type_mapping.items(): + line = { + "agent": f"Agent managing {memory_type} memories for {memory_target}", + "memory_target": memory_target, + } + lines.append(json.dumps(line, ensure_ascii=False)) + return "\n".join(lines) diff --git a/reme/agent/memory/default/__init__.py b/reme/agent/memory/default/__init__.py deleted file mode 100644 index 516a81c1..00000000 --- a/reme/agent/memory/default/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -"""Default memory agents for personal and ReMe memory operations.""" - -from .personal_retriever import PersonalRetriever -from .personal_summarizer import PersonalSummarizer -from .reme_retriever import ReMeRetriever -from .reme_summarizer import ReMeSummarizer - -__all__ = [ - "PersonalRetriever", - "PersonalSummarizer", - "ReMeRetriever", - "ReMeSummarizer", -] diff --git a/reme/agent/memory/default/personal_retriever.yaml b/reme/agent/memory/default/personal_retriever.yaml deleted file mode 100644 index 7b81e4cc..00000000 --- a/reme/agent/memory/default/personal_retriever.yaml +++ /dev/null @@ -1,48 +0,0 @@ -system_prompt: | - You are a memory agent managing **{memory_type}** memories about **{memory_target}**. - - ## User Profile - {user_profile} - - ## Question - {context} - - ## Retrieval Strategy - - **Tool 1: Vector Search (`retrieve_memory`)** - - Purpose: Search for relevant memories using semantic similarity - - Try at least 3-5 different queries before moving to next tool: - * Direct question reformulation - * Different phrasings and perspectives - * Entity-focused queries (names, places, events) - * Various keyword combinations - - Time range filtering (optional): - * Format: single date '20200101' or range '20200101,20200102' - * Example: '20200101,20200102' for 20200101 <= time <= 20200102 - * Single-sided: '0,20200102' (before date) or '20200101,99999999' (after date) - - If no results: retry with different time ranges or remove time constraints - - **Tool 2: Read History (`read_history`) - ONLY AFTER Tool 1** - - Purpose: Read full original conversation context - - Use this ONLY after completing multiple retrieve_memory attempts - - Extract history_id from retrieved memory results - - Prioritize most relevant or recent history entries - - Read multiple histories if needed for complete understanding - - ## Response Requirements - - Answer ONLY based on retrieved memories and user profile - NO hallucination or inference - - Always cite the source: reference specific memories with their timestamps - - If information conflicts, present all versions with their respective times - - Try multiple search angles before concluding no information exists - - ## Output Format - When answering, structure your response as follows: - - - [timestamp][Relevant memory content from search results] - - [timestamp][Relevant user profile information] - - If no relevant information found after thorough search (5+ queries), state: - "No relevant information found after thorough search using multiple query strategies." - -user_message: | - Answer the question following the retrieval strategy and response requirements above. \ No newline at end of file diff --git a/reme/agent/memory/default/personal_summarizer.py b/reme/agent/memory/default/personal_summarizer.py deleted file mode 100644 index ef14a1d0..00000000 --- a/reme/agent/memory/default/personal_summarizer.py +++ /dev/null @@ -1,103 +0,0 @@ -"""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, MemoryNode - - -class PersonalSummarizer(BaseMemoryAgent): - """Two-phase personal memory processor: retrieve/add memories then update profile.""" - - memory_type: MemoryType = MemoryType.PERSONAL - - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.retrieved_nodes: list[MemoryNode] = [] - - async def _build_phase1_messages(self) -> list[Message]: - """Build messages for phase 1: retrieve and add memory.""" - return [ - Message( - role=Role.SYSTEM, - content=self.prompt_format( - prompt_name="system_prompt_phase1", - 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_phase1"), - ), - ] - - async def _build_phase2_messages(self) -> list[Message]: - """Build messages for phase 2: update user profile.""" - return [ - Message( - role=Role.SYSTEM, - content=self.prompt_format( - prompt_name="system_prompt_phase2", - context=self.context.history_node.content, - memory_type=self.memory_type.value, - memory_target=self.memory_target, - user_profile=await self.read_user_profile(show_id="profile"), - ), - ), - Message( - role=Role.USER, - content=self.get_prompt("user_message_phase2"), - ), - ] - - 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): - """Execute two-phase memory processing: retrieve/add -> update profile.""" - tools = self.tools - for i, tool in enumerate(tools): - logger.info(f"[{self.__class__.__name__}] tool_call[{i}]={tool.tool_call.simple_input_dump(as_dict=False)}") - - messages_phase1 = await self._build_phase1_messages() - for i, message in enumerate(messages_phase1): - role = message.name or message.role - logger.info(f"[{self.__class__.__name__}-S1] role={role} {message.simple_dump(as_dict=False)}") - tools_phase1, messages_phase1, success_phase1 = await self.react(messages_phase1, tools[:-1], stage="S1") - - messages_phase2 = await self._build_phase2_messages() - for i, message in enumerate(messages_phase2): - role = message.name or message.role - logger.info(f"[{self.__class__.__name__}-S2] role={role} {message.simple_dump(as_dict=False)}") - tools_phase2, messages_phase2, success_phase2 = await self.react(messages_phase2, tools[-1:], stage="S2") - - return { - "answer": (messages_phase1[-1].content if success_phase1 else "") - + (messages_phase2[-1].content if success_phase2 else ""), - "success": success_phase1 and success_phase2, - "messages": messages_phase1 + messages_phase2, - "tools": tools_phase1 + tools_phase2, - } diff --git a/reme/agent/memory/default/personal_summarizer.yaml b/reme/agent/memory/default/personal_summarizer.yaml deleted file mode 100644 index 77da7843..00000000 --- a/reme/agent/memory/default/personal_summarizer.yaml +++ /dev/null @@ -1,45 +0,0 @@ -system_prompt_phase1: | - You are a memory agent managing **{memory_type}** memories about **{memory_target}**. - - ## Latest Conversation: - Message format: `round [] : ` (timestamp: YYYY-MM-DD HH:MM:SS). - {context} - - ## Task: Retrieve Similar Memories and Add New Memories - **CRITICAL**: Extract ONLY explicitly stated information. DO NOT infer, assume, or fabricate. - - ### Step 1: Retrieve Similar Memories - Use `retrieve_memory` to search for existing similar memories about **{memory_target}**. - - Use appropriate queries to find relevant existing memories - - Check if new information already exists in the memory store - - ### Step 2: Add New Memories - Use `add_memory` to add new memories: - - Extract and summarize important information about **{memory_target}** - - Set `conversation_time` (format: 2020-01-01 00:00:00; use 0000-00-00 00:00:00 if unavailable) - - If the information is completely identical to existing memory, skip adding - -user_message_phase1: | - First retrieve similar memories, then extract and add new personal memories from the conversation. - -system_prompt_phase2: | - You are a memory agent managing **{memory_type}** memories about **{memory_target}**. - - ## Latest Conversation: - Message format: `round [] : ` (timestamp: YYYY-MM-DD HH:MM:SS). - {context} - - ## Current User Profile: - UserProfile format: `profile_id= conversation_time= `. - {user_profile} - - ## Task: Update Profile with `UpdateUserProfile` - **CRITICAL**: Extract ONLY explicitly stated information. DO NOT infer, assume, or fabricate. - - Synchronize profile/memories with new information from the conversation, including **{memory_target}**' current status: - - `profile_ids_to_delete`: Remove outdated, conflicting, or redundant entries. - - `profiles_to_add`: Add new profiles/memories with `conversation_time`, e.g. `YYYY-MM-DD HH:MM:SS`, {memory_target} did something. - - Maintain profiles that are concise, mutually exclusive, and collectively comprehensive with no information loss. - -user_message_phase2: | - Update user profile using `UpdateUserProfile` based on the conversation and current profile. \ No newline at end of file diff --git a/reme/agent/memory/default/reme_retriever.yaml b/reme/agent/memory/default/reme_retriever.yaml deleted file mode 100644 index 849f562e..00000000 --- a/reme/agent/memory/default/reme_retriever.yaml +++ /dev/null @@ -1,23 +0,0 @@ -system_prompt: | - You are a Memory Orchestrator responsible for routing memory retrieval tasks to specialized agents based on the user query. - - # User Query - {context} - - ## Available Memory Agents - Each line indicates a specialized Memory Agent dedicated to storing and retrieving memories within a specific dimension (). - Format: "- (): " - {meta_memory_info} - - ## Your Task - Use the `hands_off` tool to retrieve information from specialized agents: - 1. Analyze the user query and identify which memory dimensions are relevant - 2. Specify `memory_type` and `memory_target` for each retrieval task - - The `memory_type` and `memory_target` must **exactly match** existing entries in the "Available Memory Agents" listed above - - Do NOT query agents that don't exist above - 3. Multiple tasks can be specified to enable parallel retrieval from specialized agents - - Note: If the retrieved information is insufficient to answer the query, respond: "nothing found after thorough search." - -user_message: | - Please analyze the user query and retrieve relevant information from the appropriate existing agents. \ No newline at end of file diff --git a/reme/agent/memory/default/reme_summarizer.py b/reme/agent/memory/default/reme_summarizer.py deleted file mode 100644 index 9cae5cbd..00000000 --- a/reme/agent/memory/default/reme_summarizer.py +++ /dev/null @@ -1,76 +0,0 @@ -"""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 - - -class ReMeSummarizer(BaseMemoryAgent): - """Orchestrates multiple memory agents to summarize and store information across different memory types.""" - - def __init__(self, meta_memories: list[dict] | None = None, **kwargs): - super().__init__(**kwargs) - self.meta_memories: list[dict] = meta_memories or [] - - async def build_messages(self) -> list[Message]: - self.context.history_node = await self.add_history_node() - - messages = [ - Message( - role=Role.SYSTEM, - content=self.prompt_format( - prompt_name="system_prompt", - meta_memory_info=await self.read_meta_memories(self.meta_memories), - context=self.context.history_node.content, - ), - ), - Message( - role=Role.USER, - content=self.get_prompt("user_message"), - ), - ] - - return messages - - async def _acting_step( - self, - assistant_message: Message, - tools: list[BaseTool], - step: int, - stage: str = "", - **kwargs, - ) -> tuple[list[BaseTool], list[Message]]: - return await super()._acting_step( - assistant_message, - tools, - step, - description=self.description, - messages=self.messages, - history_node=self.history_node, - author=self.author, - **kwargs, - ) - - async def execute(self): - result = await super().execute() - tools: list[BaseTool] = result["tools"] - hands_off_tool = tools[0] - agents: list[BaseMemoryAgent] = hands_off_tool.response.metadata["agents"] - - answer = "" - success = True - messages = [] - tools = [] - for agent in agents: - answer += "\n" + agent.response.answer - success = success and agent.response.success - messages += agent.response.metadata["messages"] - tools += agent.response.metadata["tools"] - - return { - "answer": answer.strip(), - "success": True, - "messages": self.messages, - "tools": tools, - } diff --git a/reme/agent/memory/default/reme_summarizer.yaml b/reme/agent/memory/default/reme_summarizer.yaml deleted file mode 100644 index d31e0213..00000000 --- a/reme/agent/memory/default/reme_summarizer.yaml +++ /dev/null @@ -1,23 +0,0 @@ -system_prompt: | - You are a Memory Orchestrator responsible for routing memory tasks to specialized agents based on the context. - - # Context - {context} - - ## Available Memory Agents - Each line indicates a specialized Memory Agent dedicated to deep summarization and updating of memories within a specific dimension (). - Format: "- (): " - {meta_memory_info} - - ## Your Task - Use the `hands_off` tool to distribute memory tasks to specialized agents: - 1. Analyze the context and identify which memory dimensions require updates - 2. Specify `memory_type` and `memory_target` for each task - - The `memory_type` and `memory_target` must **exactly match** existing entries in the "Available Memory Agents" listed above - - Do NOT create new agents or use () combinations that don't exist above - 3. Multiple tasks can be specified 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 route memory tasks to the appropriate existing agents. diff --git a/reme/tool/memory/hands_off/__init__.py b/reme/agent/memory/personal/__init__.py similarity index 100% rename from reme/tool/memory/hands_off/__init__.py rename to reme/agent/memory/personal/__init__.py diff --git a/reme/agent/memory/default/personal_retriever.py b/reme/agent/memory/personal/personal_retriever.py similarity index 62% rename from reme/agent/memory/default/personal_retriever.py rename to reme/agent/memory/personal/personal_retriever.py index 13197304..5b2fa88c 100644 --- a/reme/agent/memory/default/personal_retriever.py +++ b/reme/agent/memory/personal/personal_retriever.py @@ -3,7 +3,7 @@ from ..base_memory_agent import BaseMemoryAgent from ....core.enumeration import Role, MemoryType from ....core.op import BaseTool -from ....core.schema import Message, MemoryNode +from ....core.schema import Message from ....core.utils import format_messages @@ -12,10 +12,6 @@ class PersonalRetriever(BaseMemoryAgent): memory_type: MemoryType = MemoryType.PERSONAL - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.retrieved_nodes: list[MemoryNode] = [] - async def build_messages(self) -> list[Message]: if self.context.get("query"): context = self.context.query @@ -24,6 +20,15 @@ class PersonalRetriever(BaseMemoryAgent): 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, @@ -31,7 +36,7 @@ class PersonalRetriever(BaseMemoryAgent): prompt_name="system_prompt", memory_type=self.memory_type.value, memory_target=self.memory_target, - user_profile=await self.read_user_profile(show_id="history"), + user_profile=all_profiles, context=context.strip(), ), ), @@ -59,3 +64,23 @@ class PersonalRetriever(BaseMemoryAgent): 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_retriever.yaml b/reme/agent/memory/personal/personal_retriever.yaml new file mode 100644 index 00000000..32b85e45 --- /dev/null +++ b/reme/agent/memory/personal/personal_retriever.yaml @@ -0,0 +1,66 @@ +system_prompt: | + You are a Memory Retrieval Agent specialized in retrieving {memory_type} memories about {memory_target}. + + ## User Profile + {user_profile} + + ## User Question + {context} + + ## Multi-Phase Retrieval Strategy + Follow these phases sequentially to gather comprehensive information: + + ### Phase 1: Semantic Search (No Time Filter) + **Tool**: `retrieve_memory` (without time constraints) + **Objective**: Cast a wide net to find potentially relevant memories + **Approach**: + - Execute 3-5 diverse search queries using different formulations: + * Original question verbatim + * Rephrased variations (different wording, synonyms) + * Entity-focused queries (extract and search specific names, places, events) + * Keyword-based searches (core concepts, topics) + * Related context queries (broader themes) + - Review all results before proceeding to next phase + + ### Phase 2: Temporal Search (Optional) + **Tool**: `retrieve_memory` (with time filter) + **When to use**: Only if the user question contains temporal references + **Time Filter Format**: + - Single date: `20200101` + - Date range: `20200101,20200102` (inclusive: 20200101 ≤ time ≤ 20200102) + - Before date: `0,20200102` (up to and including 20200102) + - After date: `20200101,99999999` (from 20200101 onwards) + **Approach**: + - Identify temporal constraints from the user question + - Refine Phase 1 queries with appropriate time filters + - Try multiple time ranges if initial searches yield no results + + ### Phase 3: Deep Dive into History + **Tool**: `read_history` + **When to use**: After exhausting retrieval attempts OR when specific conversation context is needed + **Approach**: + - Extract `history_id` from retrieved memory references + - Prioritize histories that are most relevant or recent + - Read multiple histories if necessary for complete context + - Use this to understand the full conversation surrounding a memory + + ## Response Guidelines + **Critical Rules**: + - Base your answer EXCLUSIVELY on retrieved memories, user profile, and history data + - Never infer, assume, or hallucinate information + - Always cite sources with timestamps: `[timestamp] Memory content` + - Present conflicting information transparently with respective timestamps + - Exhaust all search strategies before concluding information doesn't exist + + **Output Format**: + - When information is found: + + [timestamp] All relevant memory/profile/history content + + - When no information is found after thorough search (5+ queries across phases): + + No relevant information found after exhaustive search using multiple query strategies and retrieval phases. + + +user_message: | + Retrieve relevant memories following the multi-phase strategy outlined above. \ No newline at end of file diff --git a/reme/agent/memory/personal/personal_summarizer.py b/reme/agent/memory/personal/personal_summarizer.py new file mode 100644 index 00000000..624cbcce --- /dev/null +++ b/reme/agent/memory/personal/personal_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 PersonalSummarizer(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_summarizer.yaml b/reme/agent/memory/personal/personal_summarizer.yaml new file mode 100644 index 00000000..f30dabf6 --- /dev/null +++ b/reme/agent/memory/personal/personal_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/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_retriever.yaml b/reme/agent/memory/personal/personal_v1_retriever.yaml new file mode 100644 index 00000000..32b85e45 --- /dev/null +++ b/reme/agent/memory/personal/personal_v1_retriever.yaml @@ -0,0 +1,66 @@ +system_prompt: | + You are a Memory Retrieval Agent specialized in retrieving {memory_type} memories about {memory_target}. + + ## User Profile + {user_profile} + + ## User Question + {context} + + ## Multi-Phase Retrieval Strategy + Follow these phases sequentially to gather comprehensive information: + + ### Phase 1: Semantic Search (No Time Filter) + **Tool**: `retrieve_memory` (without time constraints) + **Objective**: Cast a wide net to find potentially relevant memories + **Approach**: + - Execute 3-5 diverse search queries using different formulations: + * Original question verbatim + * Rephrased variations (different wording, synonyms) + * Entity-focused queries (extract and search specific names, places, events) + * Keyword-based searches (core concepts, topics) + * Related context queries (broader themes) + - Review all results before proceeding to next phase + + ### Phase 2: Temporal Search (Optional) + **Tool**: `retrieve_memory` (with time filter) + **When to use**: Only if the user question contains temporal references + **Time Filter Format**: + - Single date: `20200101` + - Date range: `20200101,20200102` (inclusive: 20200101 ≤ time ≤ 20200102) + - Before date: `0,20200102` (up to and including 20200102) + - After date: `20200101,99999999` (from 20200101 onwards) + **Approach**: + - Identify temporal constraints from the user question + - Refine Phase 1 queries with appropriate time filters + - Try multiple time ranges if initial searches yield no results + + ### Phase 3: Deep Dive into History + **Tool**: `read_history` + **When to use**: After exhausting retrieval attempts OR when specific conversation context is needed + **Approach**: + - Extract `history_id` from retrieved memory references + - Prioritize histories that are most relevant or recent + - Read multiple histories if necessary for complete context + - Use this to understand the full conversation surrounding a memory + + ## Response Guidelines + **Critical Rules**: + - Base your answer EXCLUSIVELY on retrieved memories, user profile, and history data + - Never infer, assume, or hallucinate information + - Always cite sources with timestamps: `[timestamp] Memory content` + - Present conflicting information transparently with respective timestamps + - Exhaust all search strategies before concluding information doesn't exist + + **Output Format**: + - When information is found: + + [timestamp] All relevant memory/profile/history content + + - When no information is found after thorough search (5+ queries across phases): + + No relevant information found after exhaustive search using multiple query strategies and retrieval phases. + + +user_message: | + Retrieve relevant memories following the multi-phase strategy outlined above. \ No newline at end of file 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/tool/memory/identity/__init__.py b/reme/agent/memory/procedural/__init__.py similarity index 100% rename from reme/tool/memory/identity/__init__.py rename to reme/agent/memory/procedural/__init__.py diff --git a/reme/agent/memory/procedural/procedural_retriever.py b/reme/agent/memory/procedural/procedural_retriever.py new file mode 100644 index 00000000..090d66f0 --- /dev/null +++ b/reme/agent/memory/procedural/procedural_retriever.py @@ -0,0 +1,10 @@ +"""Procedural memory retriever agent implementation.""" + +from ..base_memory_agent import BaseMemoryAgent +from ....core.enumeration import MemoryType + + +class ProceduralRetriever(BaseMemoryAgent): + """Agent responsible for retrieving procedural memories.""" + + memory_type: MemoryType = MemoryType.PROCEDURAL diff --git a/reme/agent/memory/procedural/procedural_summarizer.py b/reme/agent/memory/procedural/procedural_summarizer.py new file mode 100644 index 00000000..72876568 --- /dev/null +++ b/reme/agent/memory/procedural/procedural_summarizer.py @@ -0,0 +1,10 @@ +"""Procedural memory summarizer agent implementation.""" + +from ..base_memory_agent import BaseMemoryAgent +from ....core.enumeration import MemoryType + + +class ProceduralSummarizer(BaseMemoryAgent): + """Agent responsible for summarizing procedural memories.""" + + memory_type: MemoryType = MemoryType.PROCEDURAL diff --git a/reme/agent/memory/default/reme_retriever.py b/reme/agent/memory/reme_retriever.py similarity index 56% rename from reme/agent/memory/default/reme_retriever.py rename to reme/agent/memory/reme_retriever.py index 46657344..9de51aa2 100644 --- a/reme/agent/memory/default/reme_retriever.py +++ b/reme/agent/memory/reme_retriever.py @@ -1,19 +1,15 @@ """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): """Orchestrate multiple memory agents to retrieve information.""" - def __init__(self, meta_memories: list[dict] | None = None, **kwargs): - super().__init__(**kwargs) - self.meta_memories: list[dict] = meta_memories or [] - async def build_messages(self) -> list[Message]: if self.context.get("query"): context = self.context.query @@ -27,7 +23,7 @@ class ReMeRetriever(BaseMemoryAgent): role=Role.SYSTEM, content=self.prompt_format( prompt_name="system_prompt", - meta_memory_info=await self.read_meta_memories(self.meta_memories), + meta_memory_info=self.meta_memory_info, context=context.strip(), ), ), @@ -56,25 +52,41 @@ class ReMeRetriever(BaseMemoryAgent): **kwargs, ) + async def react(self, messages: list[Message], tools: list["BaseTool"], stage: str = ""): + """Run single ReAct step - only one tool call iteration.""" + used_tools: list[BaseTool] = [] + assistant_message, should_act = await self._reasoning_step(messages, tools, step=0, stage=stage) + success = True + + if should_act: + t_tools, tool_messages = await self._acting_step(assistant_message, tools, step=0, stage=stage) + used_tools.extend(t_tools) + messages.extend(tool_messages) + + return used_tools, messages, success + async def execute(self): result = await super().execute() tools: list[BaseTool] = result["tools"] - hands_off_tool = tools[0] - agents: list[BaseMemoryAgent] = hands_off_tool.response.metadata["agents"] + delegate_task_tool = tools[0] + agents: list[BaseMemoryAgent] = delegate_task_tool.response.metadata["agents"] - answer = "" + answer = [] success = True messages = [] tools = [] + retrieved_nodes = [] for agent in agents: - answer += "\n" + agent.response.answer + answer.append(agent.response.answer) success = success and agent.response.success - messages += agent.response.metadata["messages"] - tools += agent.response.metadata["tools"] + messages.extend(agent.response.metadata["messages"]) + tools.extend(agent.response.metadata["tools"]) + retrieved_nodes.extend(agent.response.metadata["retrieved_nodes"]) return { - "answer": answer.strip(), + "answer": "\n".join(answer), "success": True, - "messages": self.messages, + "messages": messages, "tools": tools, + "retrieved_nodes": retrieved_nodes, } diff --git a/reme/agent/memory/reme_retriever.yaml b/reme/agent/memory/reme_retriever.yaml new file mode 100644 index 00000000..9b5ea8af --- /dev/null +++ b/reme/agent/memory/reme_retriever.yaml @@ -0,0 +1,27 @@ +system_prompt: | + You are a Memory Orchestrator responsible for routing memory retrieval tasks to specialized agents based on the context. + + # Context + {context} + + ## Available Memory Agents + Each line below is a JSON object representing a specialized Memory Agent: + - `agent`: Description of what this agent specializes in + - `memory_target`: The unique identifier for this agent (THIS IS WHAT YOU MUST USE) + + {meta_memory_info} + + ## Your Task + Analyze the context and delegate retrieval tasks to appropriate specialized agents: + 1. Examine the context content and identify which memory_target(s) from the "Available Memory Agents" list above should be queried for relevant information + 2. For each relevant memory_target, delegate the retrieval task to its corresponding specialized agent + - **CRITICAL**: The memory_target must be EXACTLY one of the `memory_target` field values from the JSON objects listed in "Available Memory Agents" above + - **DO NOT** extract or create new memory_target names from the context content + - **DO NOT** use topic names, entity names, descriptions, or any other identifiers from the context as memory_targets + - **DO NOT** use the agent description text as memory_target + - **ONLY** use the exact string values from the `memory_target` fields in the JSON objects above + - 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 + +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/reme_summarizer.py b/reme/agent/memory/reme_summarizer.py new file mode 100644 index 00000000..8b263f64 --- /dev/null +++ b/reme/agent/memory/reme_summarizer.py @@ -0,0 +1,95 @@ +"""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 + + +class ReMeSummarizer(BaseMemoryAgent): + """Orchestrates multiple memory agents to summarize and store information across different memory types.""" + + async def build_messages(self) -> list[Message]: + add_history_tool: BaseTool | None = self.pop_tool("add_history") + if add_history_tool is not None: + await add_history_tool.call( + messages=self.messages, + description=self.description, + author=self.author, + service_context=self.service_context, + ) + self.context.history_node = add_history_tool.context.history_node + + context = self.context.description + "\n" + format_messages(self.context.messages) + messages = [ + Message( + role=Role.SYSTEM, + content=self.prompt_format( + prompt_name="system_prompt", + meta_memory_info=self.meta_memory_info, + context=context.strip(), + ), + ), + Message( + role=Role.USER, + content=self.get_prompt("user_message"), + ), + ] + + return messages + + async def _acting_step( + self, + assistant_message: Message, + tools: list[BaseTool], + step: int, + stage: str = "", + **kwargs, + ) -> tuple[list[BaseTool], list[Message]]: + return await super()._acting_step( + assistant_message, + tools, + step, + description=self.description, + messages=self.messages, + history_node=self.history_node, + author=self.author, + **kwargs, + ) + + async def react(self, messages: list[Message], tools: list["BaseTool"], stage: str = ""): + """Run single ReAct step - only one tool call iteration.""" + used_tools: list[BaseTool] = [] + assistant_message, should_act = await self._reasoning_step(messages, tools, step=0, stage=stage) + success = True + + if should_act: + t_tools, tool_messages = await self._acting_step(assistant_message, tools, step=0, stage=stage) + used_tools.extend(t_tools) + messages.extend(tool_messages) + + return used_tools, messages, success + + async def execute(self): + result = await super().execute() + tools: list[BaseTool] = result["tools"] + delegate_task_tool = tools[0] + agents: list[BaseMemoryAgent] = delegate_task_tool.response.metadata["agents"] + + success = True + messages = [] + tools = [] + memory_nodes = [] + for agent in agents: + success = success and agent.response.success + messages.extend(agent.response.metadata["messages"]) + tools.extend(agent.response.metadata["tools"]) + memory_nodes.extend(agent.response.metadata["memory_nodes"]) + + return { + "answer": memory_nodes, + "success": True, + "messages": messages, + "tools": tools, + } diff --git a/reme/agent/memory/reme_summarizer.yaml b/reme/agent/memory/reme_summarizer.yaml new file mode 100644 index 00000000..7d235060 --- /dev/null +++ b/reme/agent/memory/reme_summarizer.yaml @@ -0,0 +1,29 @@ +system_prompt: | + You are a Memory Orchestrator responsible for routing memory summarization tasks to specialized agents based on the context. + + # Context + {context} + + ## Available Memory Agents + Each line below is a JSON object representing a specialized Memory Agent: + - `agent`: Description of what this agent specializes in + - `memory_target`: The unique identifier for this agent (THIS IS WHAT YOU MUST USE) + + {meta_memory_info} + + ## Your Task + Analyze the context and delegate summarization tasks to appropriate specialized agents: + 1. Examine the context content and identify which memory_target(s) from the "Available Memory Agents" list above should receive this information + 2. For each relevant memory_target, delegate the summarization task to its corresponding specialized agent + - **CRITICAL**: The memory_target must be EXACTLY one of the `memory_target` field values from the JSON objects listed in "Available Memory Agents" above + - **DO NOT** extract or create new memory_target names from the context content + - **DO NOT** use topic names, preference categories, descriptions, or any other identifiers from the context as memory_targets + - **DO NOT** use the agent description text as memory_target + - **ONLY** use the exact string values from the `memory_target` fields in the JSON objects above + - 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 summarization tasks to the appropriate specialized agents. diff --git a/reme/tool/memory/meta/__init__.py b/reme/agent/memory/tool/__init__.py similarity index 100% rename from reme/tool/memory/meta/__init__.py rename to reme/agent/memory/tool/__init__.py diff --git a/reme/agent/memory/tool/tool_retriever.py b/reme/agent/memory/tool/tool_retriever.py new file mode 100644 index 00000000..6a2d206b --- /dev/null +++ b/reme/agent/memory/tool/tool_retriever.py @@ -0,0 +1,10 @@ +"""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/tool/tool_summarizer.py b/reme/agent/memory/tool/tool_summarizer.py new file mode 100644 index 00000000..85222d4f --- /dev/null +++ b/reme/agent/memory/tool/tool_summarizer.py @@ -0,0 +1,10 @@ +"""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/core/__init__.py b/reme/core/__init__.py index 251c2433..2b56fb08 100644 --- a/reme/core/__init__.py +++ b/reme/core/__init__.py @@ -11,6 +11,7 @@ from . import service from . import token_counter from . import utils from . import vector_store +from .application import Application from .context import R __all__ = [ @@ -25,5 +26,6 @@ __all__ = [ "token_counter", "utils", "vector_store", + "Application", "R", ] diff --git a/reme/core/application.py b/reme/core/application.py new file mode 100644 index 00000000..f96b1b03 --- /dev/null +++ b/reme/core/application.py @@ -0,0 +1,109 @@ +"""High-level entry point for configuring and running ReMe services and flows.""" + +import asyncio + +from .context import PromptHandler, ServiceContext +from .embedding import BaseEmbeddingModel +from .flow import BaseFlow +from .llm import BaseLLM +from .schema import Response +from .token_counter import BaseTokenCounter +from .utils import execute_stream_task, PydanticConfigParser +from .vector_store import BaseVectorStore + + +class Application: + """Application wrapper that wires together service context, flows, and runtimes.""" + + def __init__( + self, + *args, + llm_api_key: str | None = None, + llm_api_base: str | None = None, + embedding_api_key: str | None = None, + embedding_api_base: str | None = None, + enable_logo: bool = True, + parser: type[PydanticConfigParser] | None = None, + llm: dict | None = None, + embedding_model: dict | None = None, + vector_store: dict | None = None, + token_counter: dict | None = None, + **kwargs, + ): + # ServiceContext + self.service_context = ServiceContext( + *args, + llm_api_key=llm_api_key, + llm_api_base=llm_api_base, + embedding_api_key=embedding_api_key, + embedding_api_base=embedding_api_base, + service_config=None, + parser=parser, + config_path=None, + enable_logo=enable_logo, + llm=llm, + embedding_model=embedding_model, + vector_store=vector_store, + token_counter=token_counter, + **kwargs, + ) + + # PromptHandler + self.prompt_handler = PromptHandler(language=self.service_context.language) + + # LLM & EmbeddingModel & VectorStore & TokenCounter + self.llm: BaseLLM | None = self.service_context.llms.get("default", None) + self.embedding_model: BaseEmbeddingModel | None = self.service_context.embedding_models.get("default", None) + self.vector_store: BaseVectorStore | None = self.service_context.vector_stores.get("default", None) + self.token_counter: BaseTokenCounter | None = self.service_context.token_counters.get("default", None) + + async def __aenter__(self): + """Async context manager entry.""" + return self + + def __enter__(self): + """Context manager entry.""" + return self + + async def close(self): + """Close""" + return await self.service_context.close() + + def close_sync(self): + """Close synchronously""" + self.service_context.close_sync() + + async def __aexit__(self, exc_type=None, exc_val=None, exc_tb=None): + """Async context manager exit.""" + await self.close() + return False + + def __exit__(self, exc_type=None, exc_val=None, exc_tb=None): + """Context manager exit.""" + self.close_sync() + return False + + async def execute_flow(self, name: str, **kwargs) -> Response: + """Execute a flow with the given name and parameters.""" + assert name in self.service_context.flows, f"Flow {name} not found" + flow: BaseFlow = self.service_context.flows[name] + return await flow.call(**kwargs) + + async def execute_stream_flow(self, name: str, **kwargs): + """Execute a stream flow with the given name and parameters.""" + assert name in self.service_context.flows, f"Flow {name} not found" + flow: BaseFlow = self.service_context.flows[name] + assert flow.stream is True, "non-stream flow is not supported in execute_stream_flow!" + stream_queue = asyncio.Queue() + task = asyncio.create_task(flow.call(stream_queue=stream_queue, **kwargs)) + async for chunk in execute_stream_task( + stream_queue=stream_queue, + task=task, + task_name=name, + as_bytes=False, + ): + yield chunk + + def run_service(self): + """Run the configured service (HTTP, MCP, or CMD).""" + self.service_context.service.run() diff --git a/reme/core/llm/base_llm.py b/reme/core/llm/base_llm.py index be7b5e0d..df5e3c0d 100644 --- a/reme/core/llm/base_llm.py +++ b/reme/core/llm/base_llm.py @@ -12,6 +12,7 @@ from ..enumeration import ChunkEnum, Role from ..schema import Message from ..schema import StreamChunk from ..schema import ToolCall +from ..utils import extract_content class BaseLLM(ABC): @@ -458,6 +459,43 @@ class BaseLLM(ABC): time.sleep(1 + i) return default_value + async def simple_request( + self, + prompt: str, + model_name: str, + callback_fn: Callable[[Message], Any] | None = None, + default_value: Any = None, + **kwargs, + ) -> str: + """Make a simple request using the LLM.""" + assistant_message = await self.chat( + messages=[Message(role=Role.USER, content=prompt)], + model_name=model_name, + callback_fn=callback_fn, + default_value=default_value, + **kwargs, + ) + return assistant_message.content + + async def simple_request_for_json( + self, + prompt: str, + model_name: str, + **kwargs, + ) -> dict: + """Make a simple request using the LLM and extract JSON.""" + + def extract_fn(message: Message) -> dict: + return extract_content(message.content) + + return await self.chat( + messages=[Message(role=Role.USER, content=prompt)], + model_name=model_name, + callback_fn=extract_fn, + default_value={}, + **kwargs, + ) + async def close(self): """Release async resources.""" diff --git a/reme/core/op/base_op.py b/reme/core/op/base_op.py index 9c2dff71..6d32ddc0 100644 --- a/reme/core/op/base_op.py +++ b/reme/core/op/base_op.py @@ -115,6 +115,7 @@ class BaseOp(metaclass=ABCMeta): @property def service_context(self) -> ServiceContext: """Access the service context.""" + assert self.context, "Service context is not initialized!" return self.context.service_context @property diff --git a/reme/core/op/base_react.py b/reme/core/op/base_react.py index 95302138..38f3a42c 100644 --- a/reme/core/op/base_react.py +++ b/reme/core/op/base_react.py @@ -38,6 +38,13 @@ class BaseReact(BaseOp): """Return available tools for the agent.""" return self.sub_ops + def pop_tool(self, name: str) -> "BaseTool | None": + """Remove and return a tool from self.tools by name.""" + for i, tool in enumerate(self.sub_ops): + if tool.tool_call.name == name: + return self.sub_ops.pop(i) + return None + async def build_messages(self) -> list[Message]: """Build initial message list from context query or messages.""" if self.context.get("query"): @@ -93,7 +100,7 @@ class BaseReact(BaseOp): logger.warning(f"{prefix} unknown tool_call={tool_call.name}") continue - logger.info(f"{prefix} submit tool_calls={tool_call.simple_output_dump(as_dict=False)}") + logger.info(f"{prefix} submit tool_call[{tool_call.name}] arguments={tool_call.arguments}") # Create independent tool copy with unique ID tool_copy: BaseTool = tool_dict[tool_call.name].copy() diff --git a/reme/core/schema/memory_node.py b/reme/core/schema/memory_node.py index 67ed7c43..e6fde012 100644 --- a/reme/core/schema/memory_node.py +++ b/reme/core/schema/memory_node.py @@ -36,6 +36,7 @@ class MemoryNode(BaseModel): memory_target: Target or topic this memory relates to. when_to_use: Condition description for vector retrieval. content: Actual memory content. + message_time: Time of the message that generated this memory. ref_memory_id: Reference to related raw history memory. time_created: Creation timestamp. time_modified: Last modification timestamp. @@ -49,6 +50,7 @@ class MemoryNode(BaseModel): memory_target: str = Field(default="", description="Target or topic of the memory") when_to_use: str = Field(default="", description="Condition description for vector retrieval") content: str = Field(default="", description="Actual memory content") + message_time: str = Field(default="", description="Time of the message that generated this memory") ref_memory_id: str = Field(default="", description="Reference to related raw history memory ID") time_created: str = Field(default_factory=get_now_time, description="Creation timestamp") @@ -123,6 +125,7 @@ class MemoryNode(BaseModel): metadata: dict[str, Any] = { "memory_type": self.memory_type.value, "memory_target": self.memory_target, + "message_time": self.message_time, "ref_memory_id": self.ref_memory_id, "time_created": self.time_created, "time_modified": self.time_modified, @@ -145,6 +148,34 @@ class MemoryNode(BaseModel): metadata=metadata, ) + def format( + self, + include_memory_id: bool = True, + include_when_to_use: bool = True, + include_content: bool = True, + include_message_time: bool = True, + ref_memory_id_key: str = "", + ) -> str: + """Format memory node as string with configurable fields.""" + line = "" + + if include_memory_id and self.memory_id: + line += f"memory_id={self.memory_id} " + + if include_message_time and self.message_time: + line += f"[{self.message_time}] " + + if include_when_to_use and self.when_to_use: + line += f"{self.when_to_use} " + + if include_content and self.content: + line += self.content.strip() + + if ref_memory_id_key and self.ref_memory_id: + line += f" {ref_memory_id_key}={self.ref_memory_id}" + + return line.strip() + @classmethod def from_vector_node(cls, node: VectorNode) -> "MemoryNode": """Reconstruct MemoryNode from VectorNode. @@ -189,6 +220,7 @@ class MemoryNode(BaseModel): memory_target=metadata.pop("memory_target", ""), when_to_use=when_to_use, content=content, + message_time=metadata.pop("message_time", ""), ref_memory_id=metadata.pop("ref_memory_id", ""), time_created=metadata.pop("time_created", ""), time_modified=metadata.pop("time_modified", ""), 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/core/utils/llm_utils.py b/reme/core/utils/llm_utils.py index 805ac890..515c07b3 100644 --- a/reme/core/utils/llm_utils.py +++ b/reme/core/utils/llm_utils.py @@ -2,6 +2,7 @@ import json import re + from loguru import logger from ..enumeration import Role @@ -170,18 +171,18 @@ def extract_content(text: str, language_tag: str = "json", greedy: bool = False) pattern = rf"```\s*{re.escape(language_tag)}\s*({quantifier})\s*```" match = re.search(pattern, text, re.DOTALL) - if match: - result = match.group(1).strip() - else: - result = text + if not match: + return None + + content = match.group(1).strip() if language_tag == "json": try: - result = json.loads(result) + return json.loads(content) except json.JSONDecodeError: - result = None - - return result + return None + else: + return content def deduplicate_memories(memories: list[MemoryNode]) -> list[MemoryNode]: diff --git a/reme/reme.py b/reme/reme.py index 255c6a1d..5f587ea5 100644 --- a/reme/reme.py +++ b/reme/reme.py @@ -1,23 +1,42 @@ -"""ReMe application classes for simplified configuration and execution.""" +"""ReMe classes for simplified configuration and execution.""" -import asyncio import sys +from pathlib import Path -from .agent.memory.default import ReMeSummarizer, PersonalSummarizer, PersonalRetriever, ReMeRetriever +from .agent.memory import ( + BaseMemoryAgent, + ReMeSummarizer, + ReMeRetriever, + PersonalV1Summarizer, + PersonalV1Retriever, + PersonalSummarizer, + PersonalRetriever, + ProceduralSummarizer, + ProceduralRetriever, + ToolSummarizer, + ToolRetriever, +) from .config import ReMeConfigParser -from .core.context import ServiceContext -from .core.embedding import BaseEmbeddingModel -from .core.flow import BaseFlow -from .core.llm import BaseLLM -from .core.schema import Response, Message -from .core.token_counter import BaseTokenCounter -from .core.utils import execute_stream_task -from .core.vector_store import BaseVectorStore -from .tool.memory import UpdateUserProfile, RetrieveMemory, AddMemory, HandsOff, ReadHistory +from .core import Application +from .core.enumeration import MemoryType +from .core.schema import Message +from .tool.memory import ( + RetrieveMemory, + DelegateTask, + ReadHistory, + ProfileHandler, + MemoryHandler, + AddDraftAndRetrieveSimilarMemory, + UpdateMemoryV2, + AddDraftAndReadAllProfiles, + UpdateProfile, + AddHistory, + ReadAllProfiles, +) -class ReMe: - """ReMe application with config file support and flow execution methods.""" +class ReMe(Application): + """ReMe with config file support and flow execution methods.""" def __init__( self, @@ -31,224 +50,317 @@ class ReMe: embedding_model: dict | None = None, vector_store: dict | None = None, token_counter: dict | None = None, + personal_memory_target: list[str] | None = None, + procedural_memory_target: list[str] | None = None, + tool_memory_target: list[str] | None = None, + profile_dir: str = "reme_profile", **kwargs, ): - self.service_context = ServiceContext( + super().__init__( *args, llm_api_key=llm_api_key, llm_api_base=llm_api_base, embedding_api_key=embedding_api_key, embedding_api_base=embedding_api_base, - service_config=None, - parser=ReMeConfigParser, - config_path=None, enable_logo=enable_logo, + parser=ReMeConfigParser, llm=llm, embedding_model=embedding_model, vector_store=vector_store, token_counter=token_counter, **kwargs, ) + memory_target_type_mapping: dict[str, MemoryType] = {} + if personal_memory_target: + for name in personal_memory_target: + assert name not in memory_target_type_mapping, f"Memory target name {name} is already used." + memory_target_type_mapping[name] = MemoryType.PERSONAL - async def __aenter__(self): - """Async context manager entry.""" - return self + if procedural_memory_target: + for name in procedural_memory_target: + assert name not in memory_target_type_mapping, f"Memory target name {name} is already used." + memory_target_type_mapping[name] = MemoryType.PROCEDURAL - def __enter__(self): - """Context manager entry.""" - return self + if tool_memory_target: + for name in tool_memory_target: + assert name not in memory_target_type_mapping, f"Memory target name {name} is already used." + memory_target_type_mapping[name] = MemoryType.TOOL - async def close(self): - """Close the application.""" - return await self.service_context.close() + self.service_context.memory_target_type_mapping = memory_target_type_mapping + self.profile_dir: str = profile_dir - def close_sync(self): - """Close the application synchronously.""" - self.service_context.close_sync() + def add_meta_memory(self, memory_type: str | MemoryType, memory_target: str): + """Register or validate a memory target with the given memory type.""" + if memory_target in self.service_context.memory_target_type_mapping: + assert self.service_context.memory_target_type_mapping[memory_target] is memory_type + else: + self.service_context.memory_target_type_mapping[memory_target] = MemoryType(memory_type) - async def __aexit__(self, exc_type=None, exc_val=None, exc_tb=None): - """Async context manager exit.""" - await self.close() - return False - - def __exit__(self, exc_type=None, exc_val=None, exc_tb=None): - """Context manager exit.""" - self.close_sync() - return False - - @property - def default_llm(self) -> BaseLLM: - """Return the default LLM instance from the service context.""" - return self.service_context.llms["default"] - - @property - def default_embedding_model(self) -> BaseEmbeddingModel: - """Return the default embedding model instance from the service context.""" - return self.service_context.embedding_models["default"] - - @property - def default_vector_store(self) -> BaseVectorStore: - """Return the default vector store instance from the service context.""" - return self.service_context.vector_stores["default"] - - @property - def default_token_counter(self) -> BaseTokenCounter: - """Return the default token counter instance from the service context.""" - return self.service_context.token_counters["default"] - - async def summary( + async def summary_memory( self, messages: list[Message | dict], description: str = "", user_name: str | list[str] = "", - enable_thinking_params: bool = False, - meta_memories: list[dict] = None, + task_name: str | list[str] = "", + tool_name: str | list[str] = "", + enable_thinking_params: bool = True, version: str = "default", + retrieve_top_k: int = 20, + return_dict: bool = False, **kwargs, - ): - """Summarize messages and store them in memory for the specified user(s).""" - if user_name: + ) -> str | dict: + """Summarize personal, procedural and tool memories for the given context.""" + format_messages: list[Message] = [] + for message in messages: + if isinstance(message, dict): + assert message.get("time_created"), "message must have time_created field." + message = Message(**message) + format_messages.append(message) - if isinstance(user_name, str): - for message in messages: - if isinstance(message, dict) and not message.get("name"): - message["name"] = user_name - elif isinstance(message, Message) and not message.name: - message.name = user_name - user_name = [user_name] + personal_summarizer: BaseMemoryAgent + if version == "default": + personal_summarizer = PersonalSummarizer( + 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, + ), + UpdateProfile( + enable_thinking_params=enable_thinking_params, + profile_dir=self.profile_dir, + ), + ], + ) - if not meta_memories: - meta_memories = [ - { - "memory_type": "personal", - "memory_target": name, - } - for name in user_name - ] - - if version == "default": - reme_summarizer = ReMeSummarizer( - meta_memories=meta_memories, - tools=[ - HandsOff( - memory_agents=[ - PersonalSummarizer( - tools=[ - RetrieveMemory(enable_thinking_params=enable_thinking_params), - AddMemory(enable_thinking_params=enable_thinking_params), - UpdateUserProfile(enable_thinking_params=enable_thinking_params), - ], - ), - ], - ), - ], - ) - - else: - raise NotImplementedError - - return await reme_summarizer.call( - messages=messages, - description=description, - service_context=self.service_context, - **kwargs, + 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 - async def retrieve( + procedural_summarizer: BaseMemoryAgent + if version in ["default", "v1"]: + procedural_summarizer = ProceduralSummarizer(tools=[]) + else: + raise NotImplementedError + + tool_summarizer: BaseMemoryAgent + if version in ["default", "v1"]: + tool_summarizer = ToolSummarizer(tools=[]) + else: + raise NotImplementedError + + memory_agents = [] + if user_name: + if isinstance(user_name, str): + for message in format_messages: + message.name = user_name + self.add_meta_memory(MemoryType.PERSONAL, user_name) + elif isinstance(user_name, list): + for name in user_name: + self.add_meta_memory(MemoryType.PERSONAL, name) + else: + raise RuntimeError("user_name must be str or list[str]") + memory_agents.append(personal_summarizer) + + if task_name: + if isinstance(task_name, str): + self.add_meta_memory(MemoryType.PROCEDURAL, task_name) + elif isinstance(task_name, list): + for name in task_name: + self.add_meta_memory(MemoryType.PROCEDURAL, name) + else: + raise RuntimeError("task_name must be str or list[str]") + memory_agents.append(procedural_summarizer) + + if tool_name: + if isinstance(tool_name, str): + self.add_meta_memory(MemoryType.TOOL, tool_name) + elif isinstance(tool_name, list): + for name in tool_name: + self.add_meta_memory(MemoryType.TOOL, name) + else: + raise RuntimeError("tool_name must be str or list[str]") + memory_agents.append(tool_summarizer) + + if not memory_agents: + memory_agents = [personal_summarizer, procedural_summarizer, tool_summarizer] + + reme_summarizer: BaseMemoryAgent + if version in ["default", "v1"]: + reme_summarizer = ReMeSummarizer(tools=[AddHistory(), DelegateTask(memory_agents=memory_agents)]) + else: + raise NotImplementedError + + result = await reme_summarizer.call( + messages=format_messages, + description=description, + service_context=self.service_context, + **kwargs, + ) + + if return_dict: + return result + else: + return result["answer"] + + async def retrieve_memory( self, query: str = "", - top_k: int = 20, description: str = "", messages: list[dict] | None = None, user_name: str | list[str] = "", - enable_thinking_params: bool = False, - meta_memories: list[dict] = None, + task_name: str | list[str] = "", + tool_name: str | list[str] = "", + enable_thinking_params: bool = True, version: str = "default", + retrieve_top_k: int = 20, + enable_time_filter: bool = True, + return_dict: bool = False, **kwargs, - ): - """Retrieve relevant memories for the specified user(s) based on query or messages.""" - if user_name: - if isinstance(user_name, str): - if messages: - for message in messages: - if isinstance(message, dict) and not message.get("name"): - message["name"] = user_name - elif isinstance(message, Message) and not message.name: - message.name = user_name - user_name = [user_name] + ) -> str | dict: + """Retrieve relevant personal, procedural and tool memories for a query.""" - if not meta_memories: - meta_memories = [ - { - "memory_type": "personal", - "memory_target": name, - } - for name in user_name - ] + personal_retriever: BaseMemoryAgent + if version == "default": + personal_retriever = PersonalRetriever( + 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, + ), + ReadHistory(enable_thinking_params=enable_thinking_params), + ], + ) - if version == "default": - - reme_retriever = ReMeRetriever( - meta_memories=meta_memories, - tools=[ - HandsOff( - memory_agents=[ - PersonalRetriever( - tools=[ - RetrieveMemory(enable_thinking_params=enable_thinking_params, top_k=top_k), - ReadHistory(enable_thinking_params=enable_thinking_params), - ], - ), - ], - ), - ], - ) - - else: - raise NotImplementedError - - return await reme_retriever.call( - query=query, - messages=messages, - description=description, - service_context=self.service_context, - **kwargs, + 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, + ), + ReadHistory(enable_thinking_params=enable_thinking_params), + ], ) else: raise NotImplementedError - async def execute_flow(self, name: str, **kwargs) -> Response: - """Execute a flow with the given name and parameters.""" - assert name in self.service_context.flows, f"Flow {name} not found" - flow: BaseFlow = self.service_context.flows[name] - return await flow.call(**kwargs) + procedural_retriever: BaseMemoryAgent + if version in ["default", "v1"]: + procedural_retriever = ProceduralRetriever(tools=[]) + else: + raise NotImplementedError - async def execute_stream_flow(self, name: str, **kwargs): - """Execute a stream flow with the given name and parameters.""" - assert name in self.service_context.flows, f"Flow {name} not found" - flow: BaseFlow = self.service_context.flows[name] - assert flow.stream is True, "non-stream flow is not supported in execute_stream_flow!" - stream_queue = asyncio.Queue() - task = asyncio.create_task(flow.call(stream_queue=stream_queue, **kwargs)) - async for chunk in execute_stream_task( - stream_queue=stream_queue, - task=task, - task_name=name, - as_bytes=False, - ): - yield chunk + tool_retriever: BaseMemoryAgent + if version in ["default", "v1"]: + tool_retriever = ToolRetriever(tools=[]) + else: + raise NotImplementedError - def run_service(self): - """Run the configured service (HTTP, MCP, or CMD).""" - self.service_context.service.run() + memory_agents = [] + if user_name: + if isinstance(user_name, str): + self.add_meta_memory(MemoryType.PERSONAL, user_name) + elif isinstance(user_name, list): + for name in user_name: + self.add_meta_memory(MemoryType.PERSONAL, name) + else: + raise RuntimeError("user_name must be str or list[str]") + memory_agents.append(personal_retriever) + + if task_name: + if isinstance(task_name, str): + self.add_meta_memory(MemoryType.PROCEDURAL, task_name) + elif isinstance(task_name, list): + for name in task_name: + self.add_meta_memory(MemoryType.PROCEDURAL, name) + else: + raise RuntimeError("task_name must be str or list[str]") + memory_agents.append(procedural_retriever) + + if tool_name: + if isinstance(tool_name, str): + self.add_meta_memory(MemoryType.TOOL, tool_name) + elif isinstance(tool_name, list): + for name in tool_name: + self.add_meta_memory(MemoryType.TOOL, name) + else: + raise RuntimeError("tool_name must be str or list[str]") + memory_agents.append(tool_retriever) + + if not memory_agents: + memory_agents = [personal_retriever, procedural_retriever, tool_retriever] + + reme_retriever: BaseMemoryAgent + if version in ["default", "v1"]: + reme_retriever = ReMeRetriever(tools=[DelegateTask(memory_agents=memory_agents)]) + else: + raise NotImplementedError + + result = await reme_retriever.call( + query=query, + messages=messages, + description=description, + service_context=self.service_context, + **kwargs, + ) + + if return_dict: + return result + else: + return result["answer"] + + @property + def profile_path(self) -> Path: + """Get the path to the profile directory.""" + return Path(self.profile_dir) / self.vector_store.collection_name + + def get_memory_handler(self, memory_target: str) -> MemoryHandler: + """Get the memory handler for the specified memory target.""" + return MemoryHandler(memory_target=memory_target, service_context=self.service_context) + + def get_profile_handler(self, user_name: str) -> ProfileHandler: + """Get the profile handler for the specified user.""" + return ProfileHandler(memory_target=user_name, profile_path=self.profile_path) + + async def context_offload(self): + """working memory summary""" + + async def context_reload(self): + """working memory retrieve""" def main(): - """Main entry point for running ReMe application from command line.""" + """Main entry point for running ReMe from command line.""" with ReMe(*sys.argv[1:]) as app: app.run_service() diff --git a/reme/tool/memory/__init__.py b/reme/tool/memory/__init__.py index 9c8bbd11..ff58d44f 100644 --- a/reme/tool/memory/__init__.py +++ b/reme/tool/memory/__init__.py @@ -1,40 +1,48 @@ """memory tools""" from .base_memory_tool import BaseMemoryTool -from .hands_off.hands_off import HandsOff +from .delegate_task import DelegateTask from .history.add_history import AddHistory from .history.read_history import ReadHistory -from .identity.add_identity import AddIdentity -from .identity.read_identity import ReadIdentity -from .meta.add_meta_memory import AddMetaMemory -from .meta.read_meta_memory import ReadMetaMemory -from .user_profile.read_user_profile import ReadUserProfile -from .user_profile.update_user_profile import UpdateUserProfile +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__ = [ + # Base "BaseMemoryTool", - "HandsOff", + "DelegateTask", + # History "AddHistory", "ReadHistory", - "AddIdentity", - "ReadIdentity", - "AddMetaMemory", - "ReadMetaMemory", - "ReadUserProfile", - "UpdateUserProfile", + # Profiles + "AddDraftAndReadAllProfiles", + "ProfileHandler", + "ReadAllProfiles", + "UpdateProfile", + # Vector + "AddDraftAndRetrieveSimilarMemory", "AddMemory", "DeleteMemory", + "MemoryHandler", "RetrieveMemory", "RetrieveRecentMemory", "UpdateMemory", + "UpdateMemoryV2", ] 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/base_memory_tool.py b/reme/tool/memory/base_memory_tool.py index adc2f6b2..3431b164 100644 --- a/reme/tool/memory/base_memory_tool.py +++ b/reme/tool/memory/base_memory_tool.py @@ -6,7 +6,6 @@ from pathlib import Path from ...core.enumeration import MemoryType from ...core.op import BaseTool from ...core.schema import ToolCall, MemoryNode, ToolAttr -from ...core.utils import CacheHandler class BaseMemoryTool(BaseTool, metaclass=ABCMeta): @@ -16,14 +15,13 @@ class BaseMemoryTool(BaseTool, metaclass=ABCMeta): self, enable_multiple: bool = True, enable_thinking_params: bool = False, - local_memory_path: str = "./reme_local_memory", + profile_dir: str = "", **kwargs, ): super().__init__(**kwargs) self.enable_multiple: bool = enable_multiple self.enable_thinking_params: bool = enable_thinking_params - self.local_memory_path: str = local_memory_path - self.memory_nodes: list[MemoryNode | str] = [] + self.profile_dir: str = profile_dir def _build_tool_call(self) -> ToolCall: """Build and return the tool call schema""" @@ -60,37 +58,51 @@ class BaseMemoryTool(BaseTool, metaclass=ABCMeta): parameters.required = ["thinking"] return self._tool_call - @property - def local_memory(self) -> CacheHandler: - """Create the meta memory cache handler.""" - return CacheHandler(Path(self.local_memory_path) / self.vector_store.collection_name) - @property def memory_type(self) -> MemoryType: """Get the memory type from context.""" - return MemoryType(self.context.get("memory_type")) + return self.memory_target_type_mapping[self.memory_target] @property def memory_target(self) -> str: """Get the memory target from context.""" - return self.context.get("memory_target", "") + if "memory_target" in self.context: + return self.context.memory_target + elif len(self.memory_target_type_mapping) == 1: + return list(self.memory_target_type_mapping.keys())[0] + else: + raise ValueError("memory_target is not specified in context or memory_target_type_mapping!") @property - def memory_cache_key(self) -> str: - """Get the memory cache key from context.""" - return f"{self.memory_type.value}_{self.memory_target}".replace(" ", "_").lower() - - @property - def history_node(self) -> MemoryNode: + def history_id(self) -> str: """Get the history node from context.""" - return self.context.get("history_node") + if "history_node" in self.context: + return self.context.history_node.memory_id + return "" @property def retrieved_nodes(self) -> list[MemoryNode]: """Get the retrieved nodes from context.""" - return self.context.get("retrieved_nodes") + return self.context.retrieved_nodes @property def author(self) -> str: """Get the author from context.""" - return self.context.get("author", "") + return self.context.author + + @property + def memory_nodes(self) -> list[MemoryNode | str]: + """Get the memory nodes from context.""" + if "memory_nodes" not in self.context: + self.context.memory_nodes = [] + return self.context.memory_nodes + + @property + def memory_target_type_mapping(self) -> dict[str, MemoryType]: + """Get the memory target type mapping from context.""" + return self.context.service_context.memory_target_type_mapping + + @property + def profile_path(self) -> Path: + """Get the path to the profile directory for the current collection.""" + return Path(self.profile_dir) / self.vector_store.collection_name diff --git a/reme/tool/memory/delegate_task.py b/reme/tool/memory/delegate_task.py new file mode 100644 index 00000000..d27f4cd4 --- /dev/null +++ b/reme/tool/memory/delegate_task.py @@ -0,0 +1,85 @@ +"""Hands-off tool to delegate memory tasks to specific agents""" + +from loguru import logger + +from .base_memory_tool import BaseMemoryTool +from ...agent.memory import BaseMemoryAgent +from ...core.enumeration import MemoryType +from ...core.schema import ToolCall + + +class DelegateTask(BaseMemoryTool): + """Tool to delegate memory tasks to appropriate memory agents""" + + def __init__(self, memory_agents: list[BaseMemoryAgent] = None, **kwargs): + kwargs["enable_multiple"] = True + kwargs["sub_ops"] = memory_agents or [] + super().__init__(**kwargs) + self.sub_ops: list[BaseMemoryAgent] = [a for a in self.sub_ops if isinstance(a, BaseMemoryAgent)] + assert all(a.memory_type is not None for a in self.sub_ops) + + @property + def memory_agent_dict(self) -> dict[MemoryType, BaseMemoryAgent]: + """Map memory types to their corresponding agents""" + return {a.memory_type: a for a in self.sub_ops} + + def _build_multiple_tool_call(self) -> ToolCall: + """Build and return the multiple tool call schema""" + return ToolCall( + **{ + "description": "Delegate tasks to appropriate agents.", + "parameters": { + "type": "object", + "properties": { + "tasks": { + "type": "array", + "description": "List of tasks to delegate to specific memory agents", + "items": { + "type": "object", + "description": "A task item", + "properties": { + "memory_target": { + "type": "string", + "description": "The memory_target identifier to " + "delegate to the corresponding agent", + }, + }, + "required": ["memory_target"], + }, + }, + }, + "required": ["tasks"], + }, + }, + ) + + async def execute(self): + # Deduplicate and validate tasks + tasks = self.context.get("tasks", []) + memory_target_tasks = sorted(set(task["memory_target"] for task in tasks)) + + # Submit memory_target_tasks to agents + agent_list: list[BaseMemoryAgent] = [] + for i, memory_target in enumerate(memory_target_tasks): + memory_type = self.memory_target_type_mapping[memory_target] + agent = self.memory_agent_dict[memory_type].copy() + agent_list.append(agent) + + logger.info(f"Task {i}: {memory_type.value} agent for {memory_target}") + task_kwargs = {"memory_target": memory_target} + for k in ["query", "messages", "description", "history_node"]: + if k in self.context: + task_kwargs[k] = self.context[k] + self.submit_async_task(agent.call, service_context=self.service_context, **task_kwargs) + await self.join_async_tasks() + + # Collect results + results = [] + for agent in agent_list: + results.append(f"Task: {agent.memory_target}\n{agent.response.answer}") + + logger.info(f"Completed {len(results)} memory_target(s)") + return { + "answer": "\n\n".join(results), + "agents": agent_list, + } diff --git a/reme/tool/memory/hands_off/hands_off.py b/reme/tool/memory/hands_off/hands_off.py deleted file mode 100644 index 45822381..00000000 --- a/reme/tool/memory/hands_off/hands_off.py +++ /dev/null @@ -1,108 +0,0 @@ -"""Hands-off tool to delegate memory tasks to specific agents""" - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....agent.memory import BaseMemoryAgent -from ....core.enumeration import MemoryType -from ....core.schema import ToolCall - - -class HandsOff(BaseMemoryTool): - """Tool to delegate memory tasks to appropriate memory agents""" - - def __init__(self, memory_agents: list[BaseMemoryAgent] = None, **kwargs): - kwargs["enable_multiple"] = True - kwargs["sub_ops"] = memory_agents or [] - super().__init__(**kwargs) - self.sub_ops: list[BaseMemoryAgent] = [ - a for a in self.sub_ops if isinstance(a, BaseMemoryAgent) and a.memory_type is not None - ] - - @property - def memory_agent_dict(self) -> dict[MemoryType, BaseMemoryAgent]: - """Map memory types to their corresponding agents""" - return {a.memory_type: a for a in self.sub_ops} - - def _build_multiple_tool_call(self) -> ToolCall: - """Build and return the multiple tool call schema""" - return ToolCall( - **{ - "description": "Delegate memory tasks to appropriate agents.", - "parameters": { - "type": "object", - "properties": { - "memory_tasks": { - "type": "array", - "description": "Memory tasks to delegate to specific agents", - "items": { - "type": "object", - "properties": { - "memory_type": { - "type": "string", - "description": "Memory type to handle", - "enum": [k.value for k in self.memory_agent_dict if k], - }, - "memory_target": { - "type": "string", - "description": "Target or context for the memory operation", - }, - }, - "required": ["memory_type", "memory_target"], - }, - }, - }, - "required": ["memory_tasks"], - }, - }, - ) - - async def execute(self): - # Deduplicate and validate tasks - tasks = [] - seen = set() - for task in self.context.get("memory_tasks", []): - memory_type = MemoryType(task.get("memory_type", "")) - memory_target = task.get("memory_target", "") - - task_key = (memory_type, memory_target) - if task_key in seen: - logger.info(f"Skip duplicate: {memory_type.value} - {memory_target}") - continue - seen.add(task_key) - - tasks.append({"memory_type": memory_type, "memory_target": memory_target}) - - if not tasks: - return "No valid memory tasks to execute." - - # Submit tasks to agents - agent_list: list[BaseMemoryAgent] = [] - for i, task in enumerate(tasks): - memory_type: MemoryType = task["memory_type"] - memory_target: str = task["memory_target"] - - agent = self.memory_agent_dict[memory_type].copy() - agent_list.append(agent) - - logger.info(f"Task {i}: {memory_type.value} agent for {memory_target}") - task_kwargs = {"memory_type": memory_type, "memory_target": memory_target} - for k in ["query", "messages", "description", "history_node"]: - if k in self.context: - task_kwargs[k] = self.context[k] - self.submit_async_task(agent.call, service_context=self.service_context, **task_kwargs) - - await self.join_async_tasks() - - # Collect results - results = [] - for agent in agent_list: - memory_type = agent.memory_type - memory_target = agent.memory_target - results.append(f"{memory_type.value}({memory_target}): {agent.response.answer}") - - logger.info(f"Completed {len(results)} task(s)") - return { - "answer": "\n".join(results), - "agents": agent_list, - } diff --git a/reme/tool/memory/history/add_history.py b/reme/tool/memory/history/add_history.py index d9527ff8..629f4532 100644 --- a/reme/tool/memory/history/add_history.py +++ b/reme/tool/memory/history/add_history.py @@ -31,10 +31,11 @@ class AddHistory(BaseMemoryTool): async def execute(self): """Execute the add history operation""" self.context.messages = [Message(**m) if isinstance(m, dict) else m for m in self.context.messages] - history_content: str = (self.context.description + "\n" + format_messages(self.context.messages)).strip() + history_content: str = self.context.description + "\n" + format_messages(self.context.messages) + history_content = history_content.strip() history_node = MemoryNode( memory_type=MemoryType.HISTORY, - when_to_use=history_content[:100], + when_to_use=history_content[:1024], content=history_content, author=self.author, ) diff --git a/reme/tool/memory/history/read_history.py b/reme/tool/memory/history/read_history.py index 089a3309..1f328eb2 100644 --- a/reme/tool/memory/history/read_history.py +++ b/reme/tool/memory/history/read_history.py @@ -36,11 +36,12 @@ class ReadHistory(BaseMemoryTool): nodes = await self.vector_store.get(vector_ids=[history_id]) if not nodes: - output = f"No history: {history_id}" + output = f"No history_id={history_id} data." logger.warning(output) return output - memory = MemoryNode.from_vector_node(nodes[0]) - output = f"Historical Dialogue[{history_id}]\n{memory.content}" - logger.info(f"Successfully read history memory: {history_id}") + memory_node: MemoryNode = MemoryNode.from_vector_node(nodes[0]) + self.retrieved_nodes.append(memory_node) + output = f"Historical Dialogue[{history_id}]\n{memory_node.content}" + logger.info(f"Successfully read history memory_node: {history_id}") return output diff --git a/reme/tool/memory/identity/add_identity.py b/reme/tool/memory/identity/add_identity.py deleted file mode 100644 index 5f6ab984..00000000 --- a/reme/tool/memory/identity/add_identity.py +++ /dev/null @@ -1,42 +0,0 @@ -"""Add identity memory tool""" - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....core.schema import ToolCall - - -class AddIdentity(BaseMemoryTool): - """Tool to add or update agent identity memory""" - - def __init__(self, **kwargs): - kwargs["enable_multiple"] = False - super().__init__(**kwargs) - - def _build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "add or update agent identity memory.", - "parameters": { - "type": "object", - "properties": { - "identity_memory": { - "type": "string", - "description": "Agent identity content, such as role, personality, or current state.", - }, - }, - "required": ["identity_memory"], - }, - }, - ) - - async def execute(self): - identity_memory = self.context.get("identity_memory", "") - - if not identity_memory: - logger.warning("No valid identity memory provided") - return "No valid identity memory provided for update." - - self.local_memory.save("identity_memory", identity_memory) - logger.info(f"Successfully updated identity memory: {identity_memory}") - return "Successfully updated identity memory." diff --git a/reme/tool/memory/identity/read_identity.py b/reme/tool/memory/identity/read_identity.py deleted file mode 100644 index 856aea55..00000000 --- a/reme/tool/memory/identity/read_identity.py +++ /dev/null @@ -1,36 +0,0 @@ -"""Read identity memory tool""" - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....core.schema import ToolCall - - -class ReadIdentity(BaseMemoryTool): - """Tool to read agent identity memory""" - - def __init__(self, **kwargs): - kwargs["enable_multiple"] = False - super().__init__(**kwargs) - - def _build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "read agent identity memory.", - "parameters": { - "type": "object", - "properties": {}, - "required": [], - }, - }, - ) - - async def execute(self): - identity_memory = self.local_memory.load("identity_memory") - - if not identity_memory: - logger.info("No identity memory found") - return "No identity memory found." - - logger.info(f"Read identity memory: {identity_memory}") - return f"Identity\n{identity_memory}" diff --git a/reme/tool/memory/meta/add_meta_memory.py b/reme/tool/memory/meta/add_meta_memory.py deleted file mode 100644 index 43e9f7db..00000000 --- a/reme/tool/memory/meta/add_meta_memory.py +++ /dev/null @@ -1,89 +0,0 @@ -"""Add meta memory tool""" - -import json - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....core.enumeration import MemoryType -from ....core.schema import ToolCall - - -class AddMetaMemory(BaseMemoryTool): - """Tool to add memory metadata entries to meta storage""" - - def __init__(self, **kwargs): - kwargs["enable_multiple"] = True - super().__init__(**kwargs) - - def _build_multiple_tool_call(self) -> ToolCall: - """Build and return the multiple tool call schema""" - return ToolCall( - **{ - "description": "add memory metadata entries to register memory types and targets. " - "Before using, verify Main Agent's Meta Memory doesn't already contain the " - "same memory_type(memory_target) combinations.", - "parameters": { - "type": "object", - "properties": { - "meta_memories": { - "type": "array", - "description": "List of memory metadata entries to add", - "items": { - "type": "object", - "properties": { - "memory_type": { - "type": "string", - "description": "Type of memory: 'personal' for person-specific preferences, " - "'procedural' for how-to knowledge", - "enum": [MemoryType.PERSONAL.value, MemoryType.PROCEDURAL.value], - }, - "memory_target": { - "type": "string", - "description": "Target identifier, " - "e.g., person's name ('John') or domain ('deployment')", - }, - }, - "required": ["memory_type", "memory_target"], - }, - }, - }, - "required": ["meta_memories"], - }, - }, - ) - - async def execute(self): - existing_memories: list[dict] = self.local_memory.load("meta_memories") or [] - existing_set = {(m["memory_type"], m["memory_target"]) for m in existing_memories} - - # Filter and build new memories to add - new_memories: list[dict] = [] - meta_memories: list[dict] = self.context.get("meta_memories", []) - - for mem in meta_memories: - memory_type = mem.get("memory_type", "") - memory_target = mem.get("memory_target", "") - - # Check if valid and not duplicate - if ( - memory_type in [MemoryType.PERSONAL.value, MemoryType.PROCEDURAL.value] - and memory_target - and (memory_type, memory_target) not in existing_set - ): - new_memories.append({"memory_type": memory_type, "memory_target": memory_target}) - existing_set.add((memory_type, memory_target)) - - if not new_memories: - output = "No new meta memories to add (all entries already exist or invalid)." - logger.info(output) - return output - - # Merge, sort and save - all_memories = sorted(existing_memories + new_memories, key=lambda m: (m["memory_type"], m["memory_target"])) - self.local_memory.save("meta_memories", all_memories) - - # Format output - output = f"Successfully update meta memory entries: {json.dumps(new_memories, ensure_ascii=False)}" - logger.info(output) - return output diff --git a/reme/tool/memory/meta/read_meta_memory.py b/reme/tool/memory/meta/read_meta_memory.py deleted file mode 100644 index f1936394..00000000 --- a/reme/tool/memory/meta/read_meta_memory.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Read meta memory tool""" - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....core.enumeration import MemoryType -from ....core.schema import ToolCall - - -class ReadMetaMemory(BaseMemoryTool): - """Tool to read memory metadata from meta storage""" - - TYPE_DESC_DICT = { - MemoryType.IDENTITY.value: "self-cognition memory storing agent's identity and state", - MemoryType.PERSONAL.value: "person-specific memory storing preferences and context", - MemoryType.PROCEDURAL.value: "procedural memory storing how-to knowledge and processes", - } - - def __init__(self, enable_identity_memory: bool = False, **kwargs): - kwargs["enable_multiple"] = False - super().__init__(**kwargs) - self.enable_identity_memory = enable_identity_memory - - def _build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "read memory metadata registry to see what types of memories are being tracked.", - "parameters": { - "type": "object", - "properties": {}, - "required": [], - }, - }, - ) - - def format_memory_metadata(self, memories: list[dict[str, str]]) -> str: - """Format memory metadata into a readable string.""" - if not memories: - return "" - - lines = [] - for memory in memories: - memory_type = memory["memory_type"] - memory_target = memory["memory_target"] - description = self.TYPE_DESC_DICT[memory_type] - lines.append(f"- {memory_type}({memory_target}): {description}") - - return "\n".join(lines) - - async def execute(self): - # Load and filter meta memories - result = self.local_memory.load("meta_memories") - all_memories = result if result is not None else [] - - memories = [ - m for m in all_memories if m.get("memory_type") in [MemoryType.PERSONAL.value, MemoryType.PROCEDURAL.value] - ] - - if self.enable_identity_memory: - memories.append( - { - "memory_type": MemoryType.IDENTITY.value, - "memory_target": "self", - }, - ) - - # Format output - output = self.format_memory_metadata(memories) - if output: - logger.info(f"Retrieved {len(memories)} meta memory entries") - else: - output = "No memory metadata found." - logger.info(output) - - return output diff --git a/reme/tool/memory/user_profile/__init__.py b/reme/tool/memory/profiles/__init__.py similarity index 100% rename from reme/tool/memory/user_profile/__init__.py rename to reme/tool/memory/profiles/__init__.py diff --git a/reme/tool/memory/profiles/add_draft_and_read_all_profiles.py b/reme/tool/memory/profiles/add_draft_and_read_all_profiles.py new file mode 100644 index 00000000..716c4e5e --- /dev/null +++ b/reme/tool/memory/profiles/add_draft_and_read_all_profiles.py @@ -0,0 +1,108 @@ +"""Add draft profile and read all profiles from local storage""" + +from loguru import logger + +from .profile_handler import ProfileHandler +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import ToolCall + + +class AddDraftAndReadAllProfiles(BaseMemoryTool): + """Tool to add draft profile and read all profiles""" + + def __init__(self, enable_memory_target: bool = False, **kwargs): + super().__init__(**kwargs) + self.enable_memory_target: bool = enable_memory_target + + def _build_query_parameters(self) -> dict: + """Build the query parameters schema""" + properties = { + "message_time": { + "type": "string", + "description": "Message time, e.g. '2020-01-01 00:00:00'", + }, + "profile_key": { + "type": "string", + "description": "Profile key or category, e.g. 'name'", + }, + "profile_value": { + "type": "string", + "description": "Profile value or content, e.g. 'John Smith'", + }, + } + required = ["message_time", "profile_key", "profile_value"] + + if self.enable_memory_target: + properties["memory_target"] = { + "type": "string", + "description": "memory_target", + } + required.append("memory_target") + + return { + "type": "object", + "properties": properties, + "required": required, + } + + def _build_tool_call(self) -> ToolCall: + return ToolCall( + **{ + "description": "Add draft profile and read all profiles from local storage.", + "parameters": self._build_query_parameters(), + }, + ) + + def _build_multiple_tool_call(self) -> ToolCall: + return ToolCall( + **{ + "description": "Add draft profile and read all profiles from local storage.", + "parameters": { + "type": "object", + "properties": { + "draft_items": { + "type": "array", + "description": "List of draft profile items.", + "items": self._build_query_parameters(), + }, + }, + "required": ["draft_items"], + }, + }, + ) + + async def execute(self): + if self.enable_multiple: + draft_items = self.context.get("draft_items", []) + else: + draft_items = [self.context] + + # Collect all profiles from all targets + all_profiles = [] + targets_processed = set() + + for item in draft_items: + if self.enable_memory_target: + target = item["memory_target"] + else: + target = self.memory_target + + # Skip if already processed this target + if target in targets_processed: + continue + targets_processed.add(target) + + profile_handler = ProfileHandler(profile_path=self.profile_path, memory_target=target) + + profiles_str = profile_handler.read_all(add_profile_id=True) + if profiles_str: + all_profiles.append(f"## Profiles for {target}:\n{profiles_str}") + + if not all_profiles: + output = "No profiles found." + logger.info(output) + return output + + output = "\n\n".join(all_profiles) + logger.info(f"Successfully read profiles for {len(targets_processed)} target(s)") + return output diff --git a/reme/tool/memory/profiles/profile_handler.py b/reme/tool/memory/profiles/profile_handler.py new file mode 100644 index 00000000..6f646ac3 --- /dev/null +++ b/reme/tool/memory/profiles/profile_handler.py @@ -0,0 +1,199 @@ +"""Profile Handler for managing user profiles in local memory""" + +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 + + +class ProfileHandler: + """User profile CRUD handler""" + + def __init__(self, profile_path: str | Path, memory_target: str, max_capacity: int = 100): + """init""" + self.memory_target: str = memory_target + self.cache_key: str = self.memory_target.replace(" ", "_").lower() + self.cache_handler: CacheHandler = CacheHandler(profile_path) + self.max_capacity: int = max_capacity + + def _load_nodes(self) -> list[MemoryNode]: + """Load profile nodes""" + cached_data = self.cache_handler.load(self.cache_key, auto_clean=False) + if not cached_data: + return [] + return [MemoryNode(**data) for data in cached_data] + + def _save_nodes(self, nodes: list[MemoryNode], apply_limits: bool = True): + """Save nodes with optional deduplication and capacity enforcement""" + if apply_limits: + nodes = deduplicate_memories(nodes) + + # Enforce capacity limit by removing the oldest profiles + if len(nodes) > self.max_capacity: + sorted_nodes = sorted(nodes, key=lambda n: n.message_time) + removed_count = len(sorted_nodes) - self.max_capacity + nodes = sorted_nodes[removed_count:] + logger.info( + f"Capacity limit reached: removed {removed_count} oldest profiles " + f"(kept {len(nodes)}/{self.max_capacity})", + ) + + nodes_data = [node.model_dump(exclude_none=True) for node in nodes] + self.cache_handler.save(self.cache_key, nodes_data) + logger.info(f"Saved {len(nodes)} profiles to {self.cache_key}") + + def delete(self, profile_id: str | list[str]) -> bool | int: + """Delete profile by ID(s), returns True/False for single ID or count for batch delete""" + nodes = self._load_nodes() + original_count = len(nodes) + + # Batch delete mode + if isinstance(profile_id, list): + profile_ids_set = set(profile_id) + nodes = [n for n in nodes if n.memory_id not in profile_ids_set] + deleted_count = original_count - len(nodes) + + if deleted_count == 0: + logger.warning(f"No profiles found to delete from {len(profile_id)} IDs") + return 0 + + self._save_nodes(nodes, apply_limits=False) + logger.info(f"Batch deleted {deleted_count} profiles") + return deleted_count + + # Single delete mode + nodes = [n for n in nodes if n.memory_id != profile_id] + + if len(nodes) == original_count: + logger.warning(f"Profile {profile_id} not found") + return False + + self._save_nodes(nodes, apply_limits=False) + logger.info(f"Deleted profile {profile_id}") + return True + + def delete_all(self) -> int: + """Delete all profiles, returns count deleted""" + nodes = self._load_nodes() + count = len(nodes) + self._save_nodes([], apply_limits=False) + logger.info(f"Deleted all {count} profiles") + return count + + def add(self, message_time: str, profile_key: str, profile_value: str, ref_memory_id: str = "") -> MemoryNode: + """Add new profile, returns created MemoryNode""" + nodes = self._load_nodes() + + new_node = MemoryNode( + memory_type=MemoryType.PERSONAL, + memory_target=self.memory_target, + when_to_use=profile_key, + content=profile_value, + message_time=message_time, + ref_memory_id=ref_memory_id, + ) + + nodes.append(new_node) + self._save_nodes(nodes) + logger.info(f"Added profile: {profile_key}={profile_value}") + return new_node + + def add_batch(self, profiles: list[dict], ref_memory_id: str = "") -> list[MemoryNode]: + """Add multiple profiles in batch, returns list of created MemoryNodes""" + if not profiles: + return [] + + nodes = self._load_nodes() + + new_nodes = [ + MemoryNode( + memory_type=MemoryType.PERSONAL, + memory_target=self.memory_target, + when_to_use=p.get("profile_key", ""), + content=p.get("profile_value", ""), + message_time=p.get("message_time", ""), + ref_memory_id=ref_memory_id, + ) + for p in profiles + ] + + nodes.extend(new_nodes) + self._save_nodes(nodes) + logger.info(f"Batch added {len(new_nodes)} profiles") + return new_nodes + + def update(self, profile_id: str, message_time: str, profile_key: str, profile_value: str) -> MemoryNode | None: + """Update profile by ID, returns updated node or None if not found""" + nodes = self._load_nodes() + + target_node = None + for node in nodes: + if node.memory_id == profile_id: + node.when_to_use = profile_key + node.content = profile_value + node.message_time = message_time + target_node = node + break + + if target_node is None: + logger.warning(f"Profile {profile_id} not found") + return None + + self._save_nodes(nodes, apply_limits=False) + logger.info(f"Updated profile {profile_id}: {profile_key}={profile_value}") + return target_node + + def get_by(self, *, profile_id: str | None = None, profile_key: str | None = None) -> MemoryNode | None: + """Get profile by ID or key""" + if not profile_id and not profile_key: + raise ValueError("Must provide either profile_id or profile_key") + + nodes = self._load_nodes() + for node in nodes: + if profile_id and node.memory_id == profile_id: + return node + if profile_key and node.when_to_use == profile_key: + return node + return None + + def get_by_id(self, profile_id: str) -> MemoryNode | None: + """Get profile by ID (convenience method)""" + return self.get_by(profile_id=profile_id) + + def get_by_key(self, profile_key: str) -> MemoryNode | None: + """Get profile by key (convenience method)""" + return self.get_by(profile_key=profile_key) + + def get_all(self) -> list[MemoryNode]: + """Get all profiles, sorted by message_time""" + nodes = self._load_nodes() + nodes.sort(key=lambda n: n.message_time) + return nodes + + @staticmethod + def _format_node(node: MemoryNode, add_profile_id: bool = False, add_history_id: bool = False) -> str: + """Format a single node to string""" + parts = [] + + if add_profile_id: + parts.append(f"profile_id={node.memory_id}") + + if node.message_time: + parts.append(f"[{node.message_time}]") + + parts.append(f"{node.when_to_use}: {node.content}") + + if add_history_id: + parts.append(f"history_id={node.ref_memory_id}") + + return " ".join(parts) + + def read_all(self, add_profile_id: bool = False, add_history_id: bool = False) -> str: + """Read all profiles and return formatted string""" + nodes = self.get_all() + formatted_profiles = [self._format_node(node, add_profile_id, add_history_id) for node in nodes] + logger.info(f"Read {len(formatted_profiles)} profiles from {self.cache_key}") + return "\n".join(formatted_profiles).strip() diff --git a/reme/tool/memory/profiles/read_all_profiles.py b/reme/tool/memory/profiles/read_all_profiles.py new file mode 100644 index 00000000..892bbb58 --- /dev/null +++ b/reme/tool/memory/profiles/read_all_profiles.py @@ -0,0 +1,39 @@ +"""Read user profile tool""" + +from loguru import logger + +from .profile_handler import ProfileHandler +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import ToolCall + + +class ReadAllProfiles(BaseMemoryTool): + """Tool to read all user profiles""" + + def __init__(self, **kwargs): + kwargs["enable_multiple"] = False + super().__init__(**kwargs) + + def _build_tool_call(self) -> ToolCall: + """Build and return the tool call schema""" + return ToolCall( + **{ + "description": "Read all user profiles.", + "parameters": { + "type": "object", + "properties": {}, + "required": [], + }, + }, + ) + + async def execute(self): + profile_handler = ProfileHandler(profile_path=self.profile_path, memory_target=self.memory_target) + profiles_str = profile_handler.read_all(add_profile_id=True) + if not profiles_str: + output = "No profiles found." + logger.info(output) + return output + + logger.info("Successfully read profiles") + return profiles_str diff --git a/reme/tool/memory/user_profile/update_user_profile.py b/reme/tool/memory/profiles/update_profile.py similarity index 52% rename from reme/tool/memory/user_profile/update_user_profile.py rename to reme/tool/memory/profiles/update_profile.py index da561d8e..5edd990b 100644 --- a/reme/tool/memory/user_profile/update_user_profile.py +++ b/reme/tool/memory/profiles/update_profile.py @@ -2,13 +2,12 @@ from loguru import logger +from .profile_handler import ProfileHandler from ..base_memory_tool import BaseMemoryTool from ....core.schema import ToolCall -from ....core.schema.memory_node import MemoryNode -from ....core.utils import deduplicate_memories -class UpdateUserProfile(BaseMemoryTool): +class UpdateProfile(BaseMemoryTool): """Tool to update user profile by adding or removing profile entries""" def __init__(self, **kwargs): @@ -19,14 +18,16 @@ class UpdateUserProfile(BaseMemoryTool): """Build and return the multiple tool call schema""" return ToolCall( **{ - "description": "update user profile by adding or removing profile entries.", + "description": "update user profile by removing and adding profile entries.", "parameters": { "type": "object", "properties": { "profile_ids_to_delete": { "type": "array", "description": "List of profile IDs to delete", - "items": {"type": "string"}, + "items": { + "type": "string", + }, }, "profiles_to_add": { "type": "array", @@ -34,9 +35,9 @@ class UpdateUserProfile(BaseMemoryTool): "items": { "type": "object", "properties": { - "conversation_time": { + "message_time": { "type": "string", - "description": "Conversation time, e.g. '2020-01-01 00:00:00'", + "description": "Message time, e.g. '2020-01-01 00:00:00'", }, "profile_key": { "type": "string", @@ -47,7 +48,7 @@ class UpdateUserProfile(BaseMemoryTool): "description": "Profile value or content, e.g. 'John Smith'", }, }, - "required": ["conversation_time", "profile_key", "profile_value"], + "required": ["message_time", "profile_key", "profile_value"], }, }, }, @@ -57,53 +58,34 @@ class UpdateUserProfile(BaseMemoryTool): ) async def execute(self): - # Get and deduplicate profile IDs to delete + profile_handler = ProfileHandler(profile_path=self.profile_path, memory_target=self.memory_target) + + # Get parameters profile_ids_to_delete = self.context.get("profile_ids_to_delete", []) - profile_ids_to_delete = list(dict.fromkeys([pid for pid in profile_ids_to_delete if pid])) + profile_ids_to_delete = sorted({pid for pid in profile_ids_to_delete if pid}) profiles_to_add = self.context.get("profiles_to_add", []) if not profile_ids_to_delete and not profiles_to_add: - return "No profiles to remove or add. Operation completed." + return "No profiles to remove or add, operation completed." - # Load existing profiles from local memory - cached_data = self.local_memory.load(self.memory_cache_key, auto_clean=False) - existing_nodes = [MemoryNode(**data) for data in cached_data] if cached_data else [] - - # Remove profiles + # Delete profiles using ProfileHandler (batch mode) removed_count = 0 if profile_ids_to_delete: - original_count = len(existing_nodes) - existing_nodes = [n for n in existing_nodes if n.memory_id not in profile_ids_to_delete] - removed_count = original_count - len(existing_nodes) - logger.info(f"Removed {removed_count} profiles.") + removed_count = profile_handler.delete(profile_ids_to_delete) - # Add new profiles - new_nodes = [] + # Add new profiles using ProfileHandler (batch mode) + added_count = 0 if profiles_to_add: - for profile in profiles_to_add: - node = MemoryNode( - memory_type=self.memory_type, - memory_target=self.memory_target, - when_to_use=profile.get("profile_key", ""), - content=profile.get("profile_value", ""), - ref_memory_id=self.history_node.memory_id, - author=self.author, - metadata={"conversation_time": profile.get("conversation_time", "")}, - ) - new_nodes.append(node) - logger.info(f"Added {len(new_nodes)} new profiles.") - - # Deduplicate and save updated profiles - updated_nodes = deduplicate_memories(existing_nodes + new_nodes) - nodes_data = [node.model_dump(exclude_none=True) for node in updated_nodes] - self.local_memory.save(self.memory_cache_key, nodes_data) + new_nodes = profile_handler.add_batch(profiles=profiles_to_add, ref_memory_id=self.history_id) + self.memory_nodes.extend(new_nodes) + added_count = len(new_nodes) # Build output message operations = [] if removed_count > 0: operations.append(f"removed {removed_count} old profiles.") - if len(new_nodes) > 0: - operations.append(f"added {len(new_nodes)} new profiles.") + if added_count > 0: + operations.append(f"added {added_count} new profiles.") operations.append("Operation completed.") logger.info("\n".join(operations)) return "\n".join(operations) diff --git a/reme/tool/memory/user_profile/read_user_profile.py b/reme/tool/memory/user_profile/read_user_profile.py deleted file mode 100644 index 05346037..00000000 --- a/reme/tool/memory/user_profile/read_user_profile.py +++ /dev/null @@ -1,61 +0,0 @@ -"""Read user profile tool""" - -from typing import Literal - -from loguru import logger - -from ..base_memory_tool import BaseMemoryTool -from ....core.enumeration import MemoryType -from ....core.schema import ToolCall, MemoryNode - - -class ReadUserProfile(BaseMemoryTool): - """Tool to read user profile from local memory""" - - def __init__(self, show_id: Literal["profile", "history"] = "profile", **kwargs): - kwargs["enable_multiple"] = False - super().__init__(**kwargs) - self.show_id = show_id - - def _build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "read user profile.", - "parameters": { - "type": "object", - "properties": {}, - "required": [], - }, - }, - ) - - async def execute(self): - self.context.memory_type = MemoryType.PERSONAL - cached_data = self.local_memory.load(self.memory_cache_key, auto_clean=False) - - if not cached_data: - logger.info(f"No cached data found for {self.memory_cache_key}") - return "" - - nodes = [MemoryNode(**data) for data in cached_data] - nodes.sort(key=lambda n: n.metadata.get("conversation_time", "")) - - formatted_profiles = [] - for node in nodes: - parts = [] - if self.show_id == "profile": - parts.append(f"profile_id={node.memory_id}") - - if conv_time := node.metadata.get("conversation_time"): - parts.append(f"conversation_time={conv_time}") - - parts.append(f"{node.when_to_use}: {node.content}") - - if self.show_id == "history": - parts.append(f"history_id={node.ref_memory_id}") - - formatted_profiles.append(" ".join(parts)) - - logger.info(f"Read {len(formatted_profiles)} profiles from cache key: {self.memory_cache_key}") - - return "### User Profile\n" + "\n".join(formatted_profiles).strip() diff --git a/reme/tool/memory/vector/add_draft_and_retrieve_similar_memory.py b/reme/tool/memory/vector/add_draft_and_retrieve_similar_memory.py new file mode 100644 index 00000000..040c475a --- /dev/null +++ b/reme/tool/memory/vector/add_draft_and_retrieve_similar_memory.py @@ -0,0 +1,127 @@ +"""Add draft memory and retrieve similar memories from vector store""" + +from loguru import logger + +from .memory_handler import MemoryHandler +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import ToolCall, MemoryNode +from ....core.utils import deduplicate_memories + + +class AddDraftAndRetrieveSimilarMemory(BaseMemoryTool): + """Tool to add draft memory and retrieve similar memories""" + + def __init__( + self, + top_k: int = 20, + enable_memory_target: bool = False, + enable_when_to_use: bool = False, + **kwargs, + ): + super().__init__(**kwargs) + self.top_k: int = top_k + self.enable_memory_target: bool = enable_memory_target + self.enable_when_to_use: bool = enable_when_to_use + + def _build_query_parameters(self) -> dict: + """Build the query parameters schema""" + properties = { + "message_time": { + "type": "string", + "description": "message time, e.g. '2020-01-01 00:00:00'", + }, + "memory_content": { + "type": "string", + "description": "content of the memory.", + }, + } + required = ["message_time", "memory_content"] + + if self.enable_when_to_use: + properties["when_to_use"] = { + "type": "string", + "description": "description of when to use this memory.", + } + required.append("when_to_use") + + if self.enable_memory_target: + properties["memory_target"] = { + "type": "string", + "description": "target memory type for this memory.", + } + required.append("memory_target") + + return { + "type": "object", + "properties": properties, + "required": required, + } + + def _build_tool_call(self) -> ToolCall: + return ToolCall( + **{ + "description": "Add draft memory and retrieve similar memories from the vector store.", + "parameters": self._build_query_parameters(), + }, + ) + + def _build_multiple_tool_call(self) -> ToolCall: + return ToolCall( + **{ + "description": "Add draft memory and retrieve similar memories from the vector store.", + "parameters": { + "type": "object", + "properties": { + "draft_items": { + "type": "array", + "description": "draft_items", + "items": self._build_query_parameters(), + }, + }, + "required": ["draft_items"], + }, + }, + ) + + async def execute(self): + if self.enable_multiple: + draft_items = self.context.get("draft_items", []) + else: + draft_items = [self.context] + + queries_by_target: dict[str, list[dict]] = {} + for item in draft_items: + if self.enable_memory_target: + target = item["memory_target"] + else: + target = self.memory_target + if target not in queries_by_target: + queries_by_target[target] = [] + + queries_by_target[target].append( + { + "query": item["memory_content"], + "limit": self.top_k, + "filters": {}, + }, + ) + + # Execute batch searches for each target + memory_nodes: list[MemoryNode] = [] + for target, searches in queries_by_target.items(): + handler = MemoryHandler(target, self.service_context) + nodes = await handler.batch_search(searches) + memory_nodes.extend(nodes) + + memory_nodes = deduplicate_memories(memory_nodes) + retrieved_ids = {n.memory_id for n in self.retrieved_nodes if n.memory_id} + new_nodes = [n for n in memory_nodes if n.memory_id not in retrieved_ids] + self.retrieved_nodes.extend(new_nodes) + + if not new_nodes: + output = "No similar memories found." + else: + output = "\n".join([n.format(ref_memory_id_key="history_id") for n in new_nodes]) + + logger.info(f"Retrieved {len(memory_nodes)} similar memories, {len(new_nodes)} new after deduplication") + return output diff --git a/reme/tool/memory/vector/add_memory.py b/reme/tool/memory/vector/add_memory.py index 3371c790..3a049fc9 100644 --- a/reme/tool/memory/vector/add_memory.py +++ b/reme/tool/memory/vector/add_memory.py @@ -2,37 +2,67 @@ from loguru import logger +from .memory_handler import MemoryHandler from ..base_memory_tool import BaseMemoryTool -from ....core.schema import ToolCall, MemoryNode +from ....core.schema import ToolCall class AddMemory(BaseMemoryTool): """Tool to add memories to vector store""" + def __init__( + self, + enable_memory_target: bool = False, + enable_when_to_use: bool = False, + **kwargs, + ): + super().__init__(**kwargs) + self.enable_memory_target: bool = enable_memory_target + self.enable_when_to_use: bool = enable_when_to_use + + def _build_memory_parameters(self) -> dict: + """Build the memory parameters schema based on enabled features.""" + properties = { + "message_time": { + "type": "string", + "description": "message time, e.g. '2020-01-01 00:00:00'", + }, + "memory_content": { + "type": "string", + "description": "content of the memory.", + }, + } + required = ["message_time", "memory_content"] + + if self.enable_when_to_use: + properties["when_to_use"] = { + "type": "string", + "description": "description of when to use this memory.", + } + required.append("when_to_use") + + if self.enable_memory_target: + properties["memory_target"] = { + "type": "string", + "description": "target memory type for this memory.", + } + required.append("memory_target") + + return { + "type": "object", + "properties": properties, + "required": required, + } + def _build_tool_call(self) -> ToolCall: - """Build and return the single tool call schema""" return ToolCall( **{ "description": "add a memory to vector store for future retrieval.", - "parameters": { - "type": "object", - "properties": { - "conversation_time": { - "type": "string", - "description": "conversation time, e.g. '2020-01-01 00:00:00'", - }, - "memory_content": { - "type": "string", - "description": "content of the memory.", - }, - }, - "required": ["conversation_time", "memory_content"], - }, + "parameters": self._build_memory_parameters(), }, ) def _build_multiple_tool_call(self) -> ToolCall: - """Build and return the multiple tool call schema""" return ToolCall( **{ "description": "add multiple memories to vector store for future retrieval.", @@ -42,20 +72,7 @@ class AddMemory(BaseMemoryTool): "memories": { "type": "array", "description": "list of memories to store.", - "items": { - "type": "object", - "properties": { - "conversation_time": { - "type": "string", - "description": "conversation time, e.g. '2020-01-01 00:00:00'", - }, - "memory_content": { - "type": "string", - "description": "content of the memory.", - }, - }, - "required": ["conversation_time", "memory_content"], - }, + "items": self._build_memory_parameters(), }, }, "required": ["memories"], @@ -63,47 +80,58 @@ class AddMemory(BaseMemoryTool): }, ) - def _create_memory_node(self, data: dict) -> MemoryNode: - """Create a MemoryNode from a dictionary.""" - memory_content = data.get("memory_content", "") - conversation_time = data.get("conversation_time", "") - metadata: dict = {"conversation_time": conversation_time} - - try: - metadata["time_int"] = int(conversation_time.split(" ")[0].replace("-", "")) - except Exception: - logger.warning(f"Invalid conversation time format. {conversation_time}") - - return MemoryNode( - memory_type=self.memory_type, - memory_target=self.memory_target, - content=memory_content, - author=self.author, - metadata=metadata, - ) - async def execute(self): - memory_nodes: list[MemoryNode] = [] - memories: list[dict] = self.context.get("memories", []) - - if not memories: - memory_nodes.append(self._create_memory_node(self.context)) + if self.enable_multiple: + memories = self.context.get("memories", []) else: + memories = [self.context] + + # Group memories by memory_target if enabled + if self.enable_memory_target: + memories_by_target = {} for mem in memories: - memory_nodes.append(self._create_memory_node(mem)) + target = mem["memory_target"] + if target not in memories_by_target: + memories_by_target[target] = [] + memories_by_target[target].append(mem) + else: + memories_by_target = {self.memory_target: memories} - if not memory_nodes: - output = "No valid memories provided for addition." - logger.info(output) - return output + # Process each memory_target group + all_memory_nodes = [] + for target, target_memories in memories_by_target.items(): + # Parse and prepare memory data + memory_dicts = [] + for mem in target_memories: + memory_content = mem.get("memory_content", "") + message_time = mem.get("message_time", "") + when_to_use = mem.get("when_to_use", "") if self.enable_when_to_use else "" + metadata = {} + try: + metadata["time_int"] = int(message_time.split(" ")[0].replace("-", "")) + except Exception: + logger.warning(f"Invalid message time format: {message_time}") - vector_nodes = [node.to_vector_node() for node in memory_nodes] - vector_ids: list[str] = [node.vector_id for node in vector_nodes] + memory_dicts.append( + { + "content": memory_content, + "when_to_use": when_to_use, + "message_time": message_time, + "ref_memory_id": self.history_id, + "author": self.author, + "metadata": metadata, + }, + ) - await self.vector_store.delete(vector_ids=vector_ids) - await self.vector_store.insert(nodes=vector_nodes) - self.memory_nodes = memory_nodes + if memory_dicts: + handler = MemoryHandler(target, self.service_context) + memory_nodes = await handler.add_batch(memory_dicts) + all_memory_nodes.extend(memory_nodes) - output = f"Successfully added {len(memory_nodes)} memories to vector_store." + if not all_memory_nodes: + return "No valid memories provided." + + self.memory_nodes.extend(all_memory_nodes) + output = f"Successfully added {len(all_memory_nodes)} memories." logger.info(output) return output diff --git a/reme/tool/memory/vector/delete_memory.py b/reme/tool/memory/vector/delete_memory.py index 329cedf6..da8094fe 100644 --- a/reme/tool/memory/vector/delete_memory.py +++ b/reme/tool/memory/vector/delete_memory.py @@ -2,6 +2,7 @@ from loguru import logger +from .memory_handler import MemoryHandler from ..base_memory_tool import BaseMemoryTool from ....core.schema import ToolCall @@ -10,7 +11,6 @@ class DeleteMemory(BaseMemoryTool): """Tool to delete memories from vector store""" def _build_tool_call(self) -> ToolCall: - """Build and return the single tool call schema""" return ToolCall( **{ "description": "delete a memory from vector store using its unique ID.", @@ -19,7 +19,7 @@ class DeleteMemory(BaseMemoryTool): "properties": { "memory_id": { "type": "string", - "description": "unique identifier (memory_id) of the memory to delete.", + "description": "memory_id of the memory to delete.", }, }, "required": ["memory_id"], @@ -28,7 +28,6 @@ class DeleteMemory(BaseMemoryTool): ) def _build_multiple_tool_call(self) -> ToolCall: - """Build and return the multiple tool call schema""" return ToolCall( **{ "description": "delete multiple memories from vector store using their unique IDs.", @@ -37,7 +36,7 @@ class DeleteMemory(BaseMemoryTool): "properties": { "memory_ids": { "type": "array", - "description": "list of unique identifiers (memory_ids) of memories to delete.", + "description": "memory_ids of memories to delete.", "items": {"type": "string"}, }, }, @@ -47,25 +46,14 @@ class DeleteMemory(BaseMemoryTool): ) async def execute(self): - memory_ids: list[str] = [] - - # Handle multiple memories (array format) - ids_from_array = self.context.get("memory_ids", []) - if ids_from_array: - memory_ids = [m for m in ids_from_array if m] - else: - memory_id = self.context.get("memory_id", "") - if memory_id: - memory_ids = [memory_id] - + memory_ids = self.context.get("memory_ids") or [] if not memory_ids: - output = "No valid memory IDs provided for deletion." - logger.info(output) - return output + memory_ids = [self.context.get("memory_id", "")] - await self.vector_store.delete(vector_ids=memory_ids) - self.memory_nodes = memory_ids + handler = MemoryHandler(self.memory_target, self.service_context) + await handler.delete(memory_ids) + self.memory_nodes.extend(memory_ids) - output = f"Successfully deleted {len(memory_ids)} memories from vector_store." + output = f"Successfully deleted {len(memory_ids)} memories." logger.info(output) return output diff --git a/reme/tool/memory/vector/memory_handler.py b/reme/tool/memory/vector/memory_handler.py new file mode 100644 index 00000000..8b6f8d78 --- /dev/null +++ b/reme/tool/memory/vector/memory_handler.py @@ -0,0 +1,212 @@ +"""Memory handler""" + +from ....core.context import ServiceContext +from ....core.enumeration import MemoryType +from ....core.schema import MemoryNode +from ....core.vector_store import BaseVectorStore + + +class MemoryHandler: + """Handler for managing memory nodes in the vector store.""" + + def __init__(self, memory_target: str, service_context: ServiceContext): + self.memory_target: str = memory_target + self.memory_type: MemoryType = service_context.memory_target_type_mapping[memory_target] + self.vector_store: BaseVectorStore = service_context.vector_stores["default"] + + async def add_batch(self, memories: list[dict]) -> list[MemoryNode]: + """Add multiple memory nodes and return their memory_ids.""" + # First, delete existing memory nodes if memory_ids are provided + memory_ids_to_delete = [mem.get("memory_id") for mem in memories if mem.get("memory_id")] + if memory_ids_to_delete: + await self.vector_store.delete(memory_ids_to_delete) + + # Create MemoryNode objects + memory_nodes = [ + MemoryNode( + memory_type=self.memory_type, + memory_target=self.memory_target, + content=mem.get("content", ""), + when_to_use=mem.get("when_to_use", ""), + message_time=mem.get("message_time", ""), + ref_memory_id=mem.get("ref_memory_id", ""), + author=mem.get("author", ""), + score=mem.get("score", 0.0), + metadata=mem.get("metadata", {}), + ) + for mem in memories + ] + + # Deduplicate memory_nodes by content (keep last occurrence) + memory_dict = {node.content: node for node in memory_nodes} + memory_nodes = list(memory_dict.values()) + + # Convert to VectorNodes and insert + vector_nodes = [node.to_vector_node() for node in memory_nodes] + await self.vector_store.insert(vector_nodes) + return memory_nodes + + async def add( + self, + content: str, + when_to_use: str = "", + message_time: str = "", + ref_memory_id: str = "", + author: str = "", + score: float = 0.0, + **kwargs, + ) -> MemoryNode: + """Add a single memory node and return its memory_id.""" + memory_dict = { + "content": content, + "when_to_use": when_to_use, + "message_time": message_time, + "ref_memory_id": ref_memory_id, + "author": author, + "score": score, + "metadata": kwargs, + } + memory_nodes = await self.add_batch([memory_dict]) + return memory_nodes[0] + + async def delete(self, memory_ids: str | list[str]): + """Delete multiple memory nodes by their memory_ids.""" + # Deduplicate if input is a list + if isinstance(memory_ids, list): + memory_ids = list(dict.fromkeys(memory_ids)) + await self.vector_store.delete(memory_ids) + + async def delete_all(self): + """Delete all memory nodes.""" + await self.vector_store.delete_all() + + async def update_batch(self, updates: list[dict]) -> list[MemoryNode]: + """Update multiple memory nodes with their memory_ids and new values using delete + add.""" + # Deduplicate updates by memory_id (keep last occurrence) + updates_dict = {upd["memory_id"]: upd for upd in updates} + updates = list(updates_dict.values()) + memory_ids = list(updates_dict.keys()) + + # Get existing nodes + vector_nodes = await self.vector_store.get(memory_ids) + if not isinstance(vector_nodes, list): + vector_nodes = [vector_nodes] + + # Update and convert back + updated_nodes: list[MemoryNode] = [] + for vector_node, update in zip(vector_nodes, updates): + memory_node = MemoryNode.from_vector_node(vector_node) + memory_node.memory_target = self.memory_target + memory_node.memory_type = self.memory_type + + if "content" in update: + memory_node.content = update["content"] + if "when_to_use" in update: + memory_node.when_to_use = update["when_to_use"] + if "message_time" in update: + memory_node.message_time = update["message_time"] + if "ref_memory_id" in update: + memory_node.ref_memory_id = update["ref_memory_id"] + if "author" in update: + memory_node.author = update["author"] + if "score" in update: + memory_node.score = update["score"] + if "metadata" in update: + memory_node.metadata.update(update["metadata"]) + updated_nodes.append(memory_node) + + # Delete old nodes first + await self.vector_store.delete(memory_ids) + + # Then add updated nodes + vector_nodes = [node.to_vector_node() for node in updated_nodes] + await self.vector_store.insert(vector_nodes) + + return updated_nodes + + async def update( + self, + memory_id: str, + content: str | None = None, + when_to_use: str | None = None, + message_time: str | None = None, + ref_memory_id: str | None = None, + author: str | None = None, + score: float | None = None, + **kwargs, + ) -> MemoryNode: + """Update a memory node's content, when_to_use, or other fields.""" + update_dict: dict = {"memory_id": memory_id} + if content is not None: + update_dict["content"] = content + if when_to_use is not None: + update_dict["when_to_use"] = when_to_use + if message_time is not None: + update_dict["message_time"] = message_time + if ref_memory_id is not None: + update_dict["ref_memory_id"] = ref_memory_id + if author is not None: + update_dict["author"] = author + if score is not None: + update_dict["score"] = score + if kwargs is not None: + update_dict["metadata"] = kwargs + + memory_nodes = await self.update_batch([update_dict]) + return memory_nodes[0] + + async def search( + self, + query: str | list[str], + limit: int = 5, + filters: dict | None = None, + **kwargs, + ) -> list[MemoryNode]: + """Search for similar memory nodes based on query text.""" + filters = filters or {} + filters["memory_type"] = self.memory_type.value + filters["memory_target"] = self.memory_target + + # Handle single query + if isinstance(query, str): + vector_nodes = await self.vector_store.search(query, limit=limit, filters=filters, **kwargs) + return [MemoryNode.from_vector_node(node) for node in vector_nodes] + + # Handle multiple queries: search each query with the same limit + seen_ids: dict[str, MemoryNode] = {} + + for q in query: + vector_nodes = await self.vector_store.search(q, limit=limit, filters=filters, **kwargs) + for vector_node in vector_nodes: + memory_node = MemoryNode.from_vector_node(vector_node) + if memory_node.memory_id not in seen_ids: + seen_ids[memory_node.memory_id] = memory_node + + return list(seen_ids.values()) + + async def batch_search(self, searches: list[dict]) -> list[MemoryNode]: + """Execute multiple search queries in batch and return deduplicated results.""" + seen_ids: dict[str, MemoryNode] = {} + + for search_params in searches: + search_result = await self.search(**search_params) + for memory_node in search_result: + if memory_node.memory_id not in seen_ids: + seen_ids[memory_node.memory_id] = memory_node + + return list(seen_ids.values()) + + async def list( + self, + filters: dict | None = None, + limit: int | None = None, + sort_key: str | None = None, + reverse: bool = True, + ) -> list[MemoryNode]: + """List memory nodes with optional filtering and sorting.""" + filters = filters or {} + filters["memory_type"] = self.memory_type.value + filters["memory_target"] = self.memory_target + + vector_nodes = await self.vector_store.list(filters=filters, limit=limit, sort_key=sort_key, reverse=reverse) + return [MemoryNode.from_vector_node(node) for node in vector_nodes] diff --git a/reme/tool/memory/vector/retrieve_memory.py b/reme/tool/memory/vector/retrieve_memory.py index a83a1dae..af19bc1e 100644 --- a/reme/tool/memory/vector/retrieve_memory.py +++ b/reme/tool/memory/vector/retrieve_memory.py @@ -2,57 +2,70 @@ from loguru import logger +from .memory_handler import MemoryHandler from ..base_memory_tool import BaseMemoryTool -from ....core.schema import ToolCall, MemoryNode, VectorNode +from ....core.schema import ToolCall, MemoryNode from ....core.utils import deduplicate_memories class RetrieveMemory(BaseMemoryTool): - """Tool to retrieve memories from vector store using similarity search""" + """Tool to retrieve memories using similarity search""" - def __init__(self, top_k: int = 20, **kwargs): + def __init__(self, top_k: int = 20, enable_memory_target: bool = False, enable_time_filter: bool = False, **kwargs): super().__init__(**kwargs) self.top_k: int = top_k + self.enable_memory_target: bool = enable_memory_target + self.enable_time_filter: bool = enable_time_filter + + def _build_query_parameters(self) -> dict: + """Build the query parameters schema based on enabled features.""" + properties = { + "query": { + "type": "string", + "description": "query", + }, + } + required = ["query"] + + if self.enable_time_filter: + properties["time_filter"] = { + "type": "string", + "description": "Optional time filter to narrow down search results by date. " + "Format: single date '20200101' for exact date match, " + "or date range '20200101,20200102' for inclusive range filtering.", + } + + if self.enable_memory_target: + properties["memory_target"] = { + "type": "string", + "description": "memory_target", + } + required.append("memory_target") - @staticmethod - def _build_query_parameters() -> dict: - """Build query parameters schema for retrieval""" return { "type": "object", - "properties": { - "query": { - "type": "string", - "description": "query text for vector similarity search.", - }, - "time_range": { - "type": "string", - "description": "optional time range filter. " - "Format: single date '20200101' or range '20200101,20200102'", - }, - }, - "required": ["query"], + "properties": properties, + "required": required, } def _build_tool_call(self) -> ToolCall: - """Build and return the tool call schema""" return ToolCall( **{ - "description": "retrieve memories using vector similarity search.", + "description": "Retrieve relevant memories from the vector store using semantic similarity search.", "parameters": self._build_query_parameters(), }, ) def _build_multiple_tool_call(self) -> ToolCall: - """Build and return the multiple tool call schema""" return ToolCall( **{ - "description": "retrieve memories using multiple queries with vector similarity search.", + "description": "Retrieve relevant memories from the vector store using semantic similarity search.", "parameters": { "type": "object", "properties": { "query_items": { "type": "array", - "description": "list of query items for vector similarity search.", + "description": "List of query items.", "items": self._build_query_parameters(), }, }, @@ -61,77 +74,55 @@ class RetrieveMemory(BaseMemoryTool): }, ) - async def _retrieve_by_query( - self, - memory_type: str, - memory_target: str, - query: str, - time_range: str | None = None, - ) -> list[MemoryNode]: - """Retrieve memories by query with filters""" - filter_dict: dict = { - "memory_type": memory_type, - "memory_target": memory_target, - } - - if time_range: - time_range = time_range.strip() - if "," in time_range: - parts = time_range.split(",") - start_time = int(parts[0].strip()) - end_time = int(parts[1].strip()) - filter_dict["time_int"] = [start_time, end_time] - else: - single_time = int(time_range) - filter_dict["time_int"] = [single_time, single_time] - - nodes: list[VectorNode] = await self.vector_store.search(query=query, limit=self.top_k, filters=filter_dict) - return [MemoryNode.from_vector_node(n) for n in nodes] - async def execute(self): - memory_type: str = self.memory_type.value - memory_target: str = self.memory_target - if self.enable_multiple: - query_items: list[dict] = self.context.get("query_items", []) + query_items = self.context.get("query_items", []) else: - query_items: list[dict] = [ - { - "query": self.context.get("query", ""), - "time_range": self.context.get("time_range", ""), - }, - ] + query_items = [self.context] - query_items = [item for item in query_items if item.get("query")] - memory_nodes: list[MemoryNode] = [] + queries_by_target: dict[str, list[dict]] = {} for item in query_items: - retrieved = await self._retrieve_by_query( - memory_type=memory_type, - memory_target=memory_target, - query=item["query"], - time_range=item.get("time_range", ""), + if self.enable_memory_target: + target = item["memory_target"] + else: + target = self.memory_target + if target not in queries_by_target: + queries_by_target[target] = [] + + filters = {} + time_filter = item.get("time_filter") + if time_filter: + time_filter = time_filter.strip() + if "," in time_filter: + start, end = time_filter.split(",") + filters = {"time_int": [int(start.strip()), int(end.strip())]} + else: + filters = {"time_int": [int(time_filter), int(time_filter)]} + + queries_by_target[target].append( + { + "query": item["query"], + "limit": self.top_k, + "filters": filters, + }, ) - memory_nodes.extend(retrieved) + + # Execute batch searches for each target + memory_nodes: list[MemoryNode] = [] + for target, searches in queries_by_target.items(): + handler = MemoryHandler(target, self.service_context) + nodes = await handler.batch_search(searches) + memory_nodes.extend(nodes) memory_nodes = deduplicate_memories(memory_nodes) - retrieved_memory_ids = {node.memory_id for node in self.retrieved_nodes if node.memory_id} - new_memory_nodes = [node for node in memory_nodes if node.memory_id not in retrieved_memory_ids] - self.retrieved_nodes.extend(new_memory_nodes) - self.memory_nodes = new_memory_nodes + retrieved_ids = {n.memory_id for n in self.retrieved_nodes if n.memory_id} + new_nodes = [n for n in memory_nodes if n.memory_id not in retrieved_ids] + self.retrieved_nodes.extend(new_nodes) - if not new_memory_nodes: - output = "No new memory_nodes found matching the query (duplicates removed)." + if not new_nodes: + output = "No new memories found." else: - outputs = [] - for node in new_memory_nodes: - line = "" - if "conversation_time" in node.metadata and node.metadata["conversation_time"]: - line += f"conversation_time={node.metadata['conversation_time']} " - line += node.content.strip() + " " - if node.ref_memory_id: - line += f"history_id={node.ref_memory_id} " - outputs.append(line.strip()) - output = "\n".join(outputs) + output = "\n".join([n.format(ref_memory_id_key="history_id") for n in new_nodes]) - logger.info(f"Retrieved {len(memory_nodes)} memory_nodes, {len(new_memory_nodes)} new after deduplication") + logger.info(f"Retrieved {len(memory_nodes)} memories, {len(new_nodes)} new after deduplication") return output diff --git a/reme/tool/memory/vector/retrieve_recent_memory.py b/reme/tool/memory/vector/retrieve_recent_memory.py index 391fc72c..fc94d23d 100644 --- a/reme/tool/memory/vector/retrieve_recent_memory.py +++ b/reme/tool/memory/vector/retrieve_recent_memory.py @@ -2,13 +2,14 @@ from loguru import logger +from .memory_handler import MemoryHandler from ..base_memory_tool import BaseMemoryTool -from ....core.schema import ToolCall, MemoryNode, VectorNode +from ....core.schema import ToolCall, MemoryNode from ....core.utils import deduplicate_memories class RetrieveRecentMemory(BaseMemoryTool): - """Tool to retrieve most recent memories sorted by conversation time""" + """Tool to retrieve most recent memories sorted by time""" def __init__(self, top_k: int = 20, **kwargs): kwargs["enable_multiple"] = False @@ -16,10 +17,9 @@ class RetrieveRecentMemory(BaseMemoryTool): self.top_k: int = top_k def _build_tool_call(self) -> ToolCall: - """Build and return the tool call schema""" return ToolCall( **{ - "description": "retrieve the most recent memories sorted by conversation time (newest first).", + "description": "retrieve the most recent memories sorted by message time (newest first).", "parameters": { "type": "object", "properties": {}, @@ -28,35 +28,25 @@ class RetrieveRecentMemory(BaseMemoryTool): }, ) - async def _retrieve_recent(self) -> list[MemoryNode]: - """Retrieve recent memories sorted by conversation_time descending""" - filter_dict = { - "memory_type": self.memory_type.value, - "memory_target": self.memory_target, - } + async def execute(self): + handler = MemoryHandler(self.memory_target, self.service_context) - nodes: list[VectorNode] = await self.vector_store.list( - filters=filter_dict, + memory_nodes: list[MemoryNode] = await handler.list( limit=self.top_k, - sort_key="conversation_time", + sort_key="message_time", reverse=True, ) - - return [MemoryNode.from_vector_node(n) for n in nodes] - - async def execute(self): - memory_nodes: list[MemoryNode] = await self._retrieve_recent() memory_nodes = deduplicate_memories(memory_nodes) - retrieved_memory_ids = {node.memory_id for node in self.retrieved_nodes if node.memory_id} - new_memory_nodes = [node for node in memory_nodes if node.memory_id not in retrieved_memory_ids] - self.retrieved_nodes.extend(new_memory_nodes) - self.memory_nodes = new_memory_nodes + retrieved_ids = {n.memory_id for n in self.retrieved_nodes if n.memory_id} + new_nodes = [n for n in memory_nodes if n.memory_id not in retrieved_ids] + self.retrieved_nodes.extend(new_nodes) + self.memory_nodes.extend(new_nodes) - if not new_memory_nodes: - output = "No new memory_nodes found (duplicates removed)." + if not new_nodes: + output = "No new memories found." else: - output = "\n".join([m.format_memory() for m in new_memory_nodes]) + output = "\n".join([n.format(ref_memory_id_key="history_id") for n in new_nodes]) - logger.info(f"Retrieved {len(memory_nodes)} memory_nodes, {len(new_memory_nodes)} new after deduplication") + logger.info(f"Retrieved {len(memory_nodes)} memories, {len(new_nodes)} new after deduplication") return output diff --git a/reme/tool/memory/vector/update_memory.py b/reme/tool/memory/vector/update_memory.py index 3bb8f478..dbdf9ce6 100644 --- a/reme/tool/memory/vector/update_memory.py +++ b/reme/tool/memory/vector/update_memory.py @@ -2,41 +2,71 @@ from loguru import logger +from .memory_handler import MemoryHandler from ..base_memory_tool import BaseMemoryTool -from ....core.schema import ToolCall, MemoryNode +from ....core.schema import ToolCall class UpdateMemory(BaseMemoryTool): """Tool to update memories in vector store""" + def __init__( + self, + enable_memory_target: bool = False, + enable_when_to_use: bool = False, + **kwargs, + ): + super().__init__(**kwargs) + self.enable_memory_target: bool = enable_memory_target + self.enable_when_to_use: bool = enable_when_to_use + + def _build_update_parameters(self) -> dict: + """Build the update parameters schema based on enabled features.""" + properties = { + "memory_id": { + "type": "string", + "description": "unique identifier of memory to update.", + }, + "message_time": { + "type": "string", + "description": "message time, e.g. '2020-01-01 00:00:00'", + }, + "memory_content": { + "type": "string", + "description": "new content of the memory.", + }, + } + required = ["memory_id", "message_time", "memory_content"] + + if self.enable_when_to_use: + properties["when_to_use"] = { + "type": "string", + "description": "description of when to use this memory.", + } + required.append("when_to_use") + + if self.enable_memory_target: + properties["memory_target"] = { + "type": "string", + "description": "target memory type for this memory.", + } + required.append("memory_target") + + return { + "type": "object", + "properties": properties, + "required": required, + } + def _build_tool_call(self) -> ToolCall: - """Build and return the single tool call schema""" return ToolCall( **{ "description": "update a memory in vector store by replacing old memory with new content.", - "parameters": { - "type": "object", - "properties": { - "memory_id": { - "type": "string", - "description": "unique identifier of memory to update.", - }, - "conversation_time": { - "type": "string", - "description": "conversation time, e.g. '2020-01-01 00:00:00'", - }, - "memory_content": { - "type": "string", - "description": "new content of the memory.", - }, - }, - "required": ["memory_id", "conversation_time", "memory_content"], - }, + "parameters": self._build_update_parameters(), }, ) def _build_multiple_tool_call(self) -> ToolCall: - """Build and return the multiple tool call schema""" return ToolCall( **{ "description": "update multiple memories in vector store by replacing old memories with new content.", @@ -46,24 +76,7 @@ class UpdateMemory(BaseMemoryTool): "memories": { "type": "array", "description": "list of memory update objects.", - "items": { - "type": "object", - "properties": { - "memory_id": { - "type": "string", - "description": "unique identifier of memory to update.", - }, - "conversation_time": { - "type": "string", - "description": "conversation time, e.g. '2020-01-01 00:00:00'", - }, - "memory_content": { - "type": "string", - "description": "new content of the memory.", - }, - }, - "required": ["memory_id", "conversation_time", "memory_content"], - }, + "items": self._build_update_parameters(), }, }, "required": ["memories"], @@ -71,56 +84,58 @@ class UpdateMemory(BaseMemoryTool): }, ) - def _create_memory_node(self, data: dict) -> tuple[str, MemoryNode]: - """Create a MemoryNode from a dictionary.""" - memory_id = data.get("memory_id", "") - memory_content = data.get("memory_content", "") - conversation_time = data.get("conversation_time", "") - metadata: dict = {"conversation_time": conversation_time} - - try: - metadata["time_int"] = int(conversation_time.split(" ")[0].replace("-", "")) - except Exception: - logger.warning(f"Invalid conversation time format. {conversation_time}") - - memory_node = MemoryNode( - memory_type=self.memory_type, - memory_target=self.memory_target, - content=memory_content, - author=self.author, - metadata=metadata, - ) - - return memory_id, memory_node - async def execute(self): - old_memory_ids: list[str] = [] - memory_nodes: list[MemoryNode] = [] - memories: list[dict] = self.context.get("memories", []) - - if not memories: - old_id, node = self._create_memory_node(self.context) - old_memory_ids.append(old_id) - memory_nodes.append(node) + if self.enable_multiple: + memories = self.context.get("memories", []) else: + memories = [self.context] + + # Group memories by memory_target if enabled + if self.enable_memory_target: + memories_by_target = {} for mem in memories: - old_id, node = self._create_memory_node(mem) - old_memory_ids.append(old_id) - memory_nodes.append(node) + target = mem["memory_target"] + if target not in memories_by_target: + memories_by_target[target] = [] + memories_by_target[target].append(mem) + else: + memories_by_target = {self.memory_target: memories} - if not memory_nodes: - output = "No valid memories provided for update." - logger.info(output) - return output + # Process each memory_target group + all_memory_nodes = [] + for target, target_memories in memories_by_target.items(): + # Parse and prepare update data + update_dicts = [] + for mem in target_memories: + memory_content = mem.get("memory_content", "") + message_time = mem.get("message_time", "") + when_to_use = mem.get("when_to_use", "") if self.enable_when_to_use else "" + metadata = {} + try: + metadata["time_int"] = int(message_time.split(" ")[0].replace("-", "")) + except Exception: + logger.warning(f"Invalid message time format: {message_time}") - vector_nodes = [node.to_vector_node() for node in memory_nodes] - new_vector_ids: list[str] = [node.vector_id for node in vector_nodes] + update_dicts.append( + { + "memory_id": mem.get("memory_id", ""), + "content": memory_content, + "when_to_use": when_to_use, + "message_time": message_time, + "author": self.author, + "metadata": metadata, + }, + ) - all_ids_to_delete = list(set(old_memory_ids + new_vector_ids)) - await self.vector_store.delete(vector_ids=all_ids_to_delete) - await self.vector_store.insert(nodes=vector_nodes) - self.memory_nodes = memory_nodes + if update_dicts: + handler = MemoryHandler(target, self.service_context) + memory_nodes = await handler.update_batch(update_dicts) + all_memory_nodes.extend(memory_nodes) - output = f"Successfully updated {len(memory_nodes)} memories in vector_store." + if not all_memory_nodes: + return "No valid memories provided." + + self.memory_nodes.extend(all_memory_nodes) + output = f"Successfully updated {len(all_memory_nodes)} memories." logger.info(output) return output diff --git a/reme/tool/memory/vector/update_memory_v2.py b/reme/tool/memory/vector/update_memory_v2.py new file mode 100644 index 00000000..5b8f0950 --- /dev/null +++ b/reme/tool/memory/vector/update_memory_v2.py @@ -0,0 +1,157 @@ +"""Update memory in vector store""" + +from loguru import logger + +from .memory_handler import MemoryHandler +from ..base_memory_tool import BaseMemoryTool +from ....core.schema import ToolCall + + +class UpdateMemoryV2(BaseMemoryTool): + """Tool to update memories in vector store by deleting and adding memory entries""" + + def __init__( + self, + name="update_memory", + enable_memory_target: bool = False, + enable_when_to_use: bool = False, + **kwargs, + ): + kwargs["enable_multiple"] = True + super().__init__(name=name, **kwargs) + self.enable_memory_target: bool = enable_memory_target + self.enable_when_to_use: bool = enable_when_to_use + + def _build_add_memory_parameters(self) -> dict: + """Build the add memory parameters schema based on enabled features.""" + properties = { + "message_time": { + "type": "string", + "description": "message time, e.g. '2020-01-01 00:00:00'", + }, + "memory_content": { + "type": "string", + "description": "content of the memory.", + }, + } + required = ["message_time", "memory_content"] + + if self.enable_when_to_use: + properties["when_to_use"] = { + "type": "string", + "description": "description of when to use this memory.", + } + required.append("when_to_use") + + if self.enable_memory_target: + properties["memory_target"] = { + "type": "string", + "description": "target memory type for this memory.", + } + required.append("memory_target") + + return { + "type": "object", + "properties": properties, + "required": required, + } + + def _build_multiple_tool_call(self) -> ToolCall: + """Build and return the multiple tool call schema""" + return ToolCall( + **{ + "description": "update memories by removing and adding memory entries.", + "parameters": { + "type": "object", + "properties": { + "memory_ids_to_delete": { + "type": "array", + "description": "List of memory IDs to delete", + "items": { + "type": "string", + }, + }, + "memories_to_add": { + "type": "array", + "description": "List of memories to add", + "items": self._build_add_memory_parameters(), + }, + }, + "required": ["memory_ids_to_delete", "memories_to_add"], + }, + }, + ) + + async def execute(self): + # Get parameters + memory_ids_to_delete = self.context.get("memory_ids_to_delete", []) + memory_ids_to_delete = sorted({mid for mid in memory_ids_to_delete if mid}) + memories_to_add = self.context.get("memories_to_add", []) + + if not memory_ids_to_delete and not memories_to_add: + return "No memories to remove or add, operation completed." + + # Group memories by memory_target if enabled + if self.enable_memory_target: + memories_by_target = {} + for mem in memories_to_add: + target = mem.get("memory_target", self.memory_target) + if target not in memories_by_target: + memories_by_target[target] = [] + memories_by_target[target].append(mem) + else: + memories_by_target = {self.memory_target: memories_to_add} + + # Delete memories (all at once, regardless of target) + removed_count = 0 + if memory_ids_to_delete: + # Use the default memory_target handler for deletion + handler = MemoryHandler(self.memory_target, self.service_context) + await handler.delete(memory_ids_to_delete) + removed_count = len(memory_ids_to_delete) + + # Add new memories by target + added_count = 0 + all_memory_nodes = [] + for target, target_memories in memories_by_target.items(): + # Parse and prepare add data + add_dicts = [] + for mem in target_memories: + memory_content = mem.get("memory_content", "") + message_time = mem.get("message_time", "") + when_to_use = mem.get("when_to_use", "") if self.enable_when_to_use else "" + metadata = {} + try: + metadata["time_int"] = int(message_time.split(" ")[0].replace("-", "")) + except Exception: + logger.warning(f"Invalid message time format: {message_time}") + + add_dicts.append( + { + "content": memory_content, + "when_to_use": when_to_use, + "message_time": message_time, + "ref_memory_id": self.history_id, + "author": self.author, + "metadata": metadata, + }, + ) + + if add_dicts: + handler = MemoryHandler(target, self.service_context) + memory_nodes = await handler.add_batch(add_dicts) + all_memory_nodes.extend(memory_nodes) + added_count += len(memory_nodes) + + # Extend memory_nodes for tracking + self.memory_nodes.extend(all_memory_nodes) + + # Build output message + operations = [] + if removed_count > 0: + operations.append(f"removed {removed_count} old memories.") + if added_count > 0: + operations.append(f"added {added_count} new memories.") + operations.append("Operation completed.") + logger.info("\n".join(operations)) + return "\n".join(operations) diff --git a/tests/test_reme.py b/tests/test_reme.py index a4f383cf..7ac30e3c 100644 --- a/tests/test_reme.py +++ b/tests/test_reme.py @@ -11,40 +11,48 @@ reme = ReMe(vector_store={"collection_name": "reme"}) async def test_reme(): """Tests ReMe memory system with personal information storage and retrieval.""" # 构建一段包含个人信息的对话 - await reme.default_vector_store.delete_all() + await reme.vector_store.delete_all() messages = [ { "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", }, ] @@ -53,7 +61,7 @@ async def test_reme(): print("=" * 60) # 对对话进行总结，生成记忆 - await reme.summary( + await reme.summary_memory( messages=messages, user_name="zhangwei", description="用户自我介绍和技术兴趣分享", @@ -66,7 +74,7 @@ async def test_reme(): print("=" * 60) # 列出所有存储的记忆节点 - nodes: list[VectorNode] = await reme.default_vector_store.list() + nodes: list[VectorNode] = await reme.vector_store.list() for i, node in enumerate(nodes, 1): memory_node = MemoryNode.from_vector_node(node) print(f"{i} {memory_node.model_dump_json()}") @@ -78,25 +86,25 @@ async def test_reme(): # 测试问题1: 检索用户姓名 query1 = "用户叫什么名字？" print(f"\n问题1: {query1}") - result1 = await reme.retrieve(query=query1, user_name="zhangwei") + result1 = await reme.retrieve_memory(query=query1, user_name="zhangwei") print(f"检索结果:\n{result1}") # 测试问题2: 检索技术背景 query2 = "用户擅长什么编程语言和技术方向？" print(f"\n问题2: {query2}") - result2 = await reme.retrieve(query=query2, user_name="zhangwei") + result2 = await reme.retrieve_memory(query=query2, user_name="zhangwei") print(f"检索结果:\n{result2}") # 测试问题3: 检索个人信息 query3 = "用户的工作地点和联系方式是什么？" print(f"\n问题3: {query3}") - result3 = await reme.retrieve(query=query3, user_name="zhangwei") + result3 = await reme.retrieve_memory(query=query3, user_name="zhangwei") print(f"检索结果:\n{result3}") # 测试问题4: 检索兴趣爱好 query4 = "用户平时有什么爱好或活动？" print(f"\n问题4: {query4}") - result4 = await reme.retrieve(query=query4, user_name="zhangwei") + result4 = await reme.retrieve_memory(query=query4, user_name="zhangwei") print(f"检索结果:\n{result4}") print("\n" + "=" * 60)