From ab4bf43a24a1568c154e68f4aa4eb5c99db2cb36 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Tue, 26 Aug 2025 23:32:57 +0800 Subject: [PATCH] refactor(summary): optimize task memory processing and validation - Rename and restructure task memory related operations - Implement static methods for utility functions - Improve type safety and error handling - Update memory validation and deduplication logic - Refactor comparative summary generation --- reme_ai/config/default.yaml | 21 ++++++++++++++++++- .../summary/task/comparative_extraction_op.py | 12 +++++++---- ...ation_op.py => memory_deduplication_op.py} | 9 ++++---- ...lidation_op.py => memory_validation_op.py} | 17 +++++++++------ ...mpt.yaml => memory_validation_prompt.yaml} | 0 .../task/simple_comparative_summary_op.py | 12 ++++++----- .../task/trajectory_segmentation_op.py | 6 ++++-- 7 files changed, 55 insertions(+), 22 deletions(-) rename reme_ai/summary/task/{task_memory_deduplication_op.py => memory_deduplication_op.py} (95%) rename reme_ai/summary/task/{experience_validation_op.py => memory_validation_op.py} (85%) rename reme_ai/summary/task/{experience_validation_prompt.yaml => memory_validation_prompt.yaml} (100%) diff --git a/reme_ai/config/default.yaml b/reme_ai/config/default.yaml index bd6c213c..9198bd99 100644 --- a/reme_ai/config/default.yaml +++ b/reme_ai/config/default.yaml @@ -44,7 +44,8 @@ flow: required: true summary_task_memory: - flow_content: trajectory_preprocess_op->[success_extraction_op|failure_extraction_op|comparative_extraction_op]->experience_validation_op->experience_deduplication_op->update_vector_store_op + # memory_deduplication_op + flow_content: trajectory_preprocess_op >> (success_extraction_op|failure_extraction_op|comparative_extraction_op) >> memory_validation_op >> update_vector_store_op description: "Summarize trajectories or messages into memories" input_schema: trajectories: @@ -52,6 +53,24 @@ flow: description: "A list of conversation trajectory information, including message content and score. This field does not need to be filled in, the system will complete it automatically." required: false + agent_task: + flow_content: react_op + description: "A React-capable agent that can utilize web search and code execution tools." + input_schema: + query: + type: "str" + description: "current query" + required: true + + vector_store: + flow_content: vector_store_action_op + description: "A React-capable agent that can utilize web search and code execution tools." + input_schema: + query: + type: "str" + description: "current query" + required: true + # vector_store: vector_store_action_op # agent: react_op diff --git a/reme_ai/summary/task/comparative_extraction_op.py b/reme_ai/summary/task/comparative_extraction_op.py index 14d52550..cdc1cd2d 100644 --- a/reme_ai/summary/task/comparative_extraction_op.py +++ b/reme_ai/summary/task/comparative_extraction_op.py @@ -46,7 +46,8 @@ class ComparativeExtractionOp(BaseLLMOp): # Add task memories to context self.context.comparative_task_memories = comparative_task_memories - def _find_highest_lowest_scoring_trajectories(self, trajectories: List[Trajectory]) -> Tuple[Optional[Trajectory], Optional[Trajectory]]: + @staticmethod + def _find_highest_lowest_scoring_trajectories(trajectories: List[Trajectory]) -> Tuple[Optional[Trajectory], Optional[Trajectory]]: """Find the highest and lowest scoring trajectories""" if len(trajectories) < 2: return None, None @@ -66,7 +67,8 @@ class ComparativeExtractionOp(BaseLLMOp): return highest_traj, lowest_traj - def _get_trajectory_score(self, trajectory: Trajectory) -> Optional[float]: + @staticmethod + def _get_trajectory_score(trajectory: Trajectory) -> Optional[float]: """Get trajectory score""" return trajectory.score @@ -133,7 +135,8 @@ class ComparativeExtractionOp(BaseLLMOp): return self.llm.chat(messages=[Message(content=prompt)], callback_fn=parse_task_memories) - def _get_trajectory_steps(self, trajectory: Trajectory) -> List[Message]: + @staticmethod + def _get_trajectory_steps(trajectory: Trajectory) -> List[Message]: """Get trajectory steps, prioritizing segmented steps""" if hasattr(trajectory, 'segments') and trajectory.segments: # If there are segments, merge all segments @@ -208,7 +211,8 @@ class ComparativeExtractionOp(BaseLLMOp): return [] - def _calculate_cosine_similarity(self, embedding1: List[float], embedding2: List[float]) -> float: + @staticmethod + def _calculate_cosine_similarity(embedding1: List[float], embedding2: List[float]) -> float: """Calculate cosine similarity""" import numpy as np diff --git a/reme_ai/summary/task/task_memory_deduplication_op.py b/reme_ai/summary/task/memory_deduplication_op.py similarity index 95% rename from reme_ai/summary/task/task_memory_deduplication_op.py rename to reme_ai/summary/task/memory_deduplication_op.py index c57af218..8982a576 100644 --- a/reme_ai/summary/task/task_memory_deduplication_op.py +++ b/reme_ai/summary/task/memory_deduplication_op.py @@ -12,7 +12,7 @@ class TaskMemoryDeduplicationOp(BaseOp): def execute(self): """Remove duplicate task memories""" # Get task memories to deduplicate - task_memories: List[BaseMemory] = self.context.get("task_memories", []) + task_memories: List[BaseMemory] = self.context.memory_list if not task_memories: logger.info("No task memories found for deduplication") @@ -26,7 +26,7 @@ class TaskMemoryDeduplicationOp(BaseOp): logger.info(f"Deduplication complete: {len(deduplicated_task_memories)} deduplicated task memories out of {len(task_memories)}") # Update context - self.context.deduplicated_task_memories = deduplicated_task_memories + self.context.memory_list = deduplicated_task_memories def _deduplicate_task_memories(self, task_memories: List[BaseMemory]) -> List[BaseMemory]: """Remove duplicate task memories""" @@ -91,7 +91,7 @@ class TaskMemoryDeduplicationOp(BaseOp): logger.warning(f"Failed to retrieve existing task memory embeddings: {e}") return [] - def _get_task_memory_embedding(self, task_memory: BaseMemory) -> List[float]: + def _get_task_memory_embedding(self, task_memory: BaseMemory) -> List[float] | None: """Generate embedding for task memory""" try: if not hasattr(self.context, 'vector_store') or not self.context.vector_store: @@ -137,7 +137,8 @@ class TaskMemoryDeduplicationOp(BaseOp): return True return False - def _calculate_cosine_similarity(self, embedding1: List[float], embedding2: List[float]) -> float: + @staticmethod + def _calculate_cosine_similarity(embedding1: List[float], embedding2: List[float]) -> float: """Calculate cosine similarity""" try: import numpy as np diff --git a/reme_ai/summary/task/experience_validation_op.py b/reme_ai/summary/task/memory_validation_op.py similarity index 85% rename from reme_ai/summary/task/experience_validation_op.py rename to reme_ai/summary/task/memory_validation_op.py index bf959a85..66cb2e9e 100644 --- a/reme_ai/summary/task/experience_validation_op.py +++ b/reme_ai/summary/task/memory_validation_op.py @@ -10,12 +10,17 @@ from reme_ai.schema.memory import BaseMemory @C.register_op() -class TaskMemoryValidationOp(BaseLLMOp): +class MemoryValidationOp(BaseLLMOp): file_path: str = __file__ def execute(self): """Validate quality of extracted task memories""" - task_memories: List[BaseMemory] = self.context.get("task_memories", []) + self.context.memory_list = [] + self.context.memory_list.append(self.context.success_task_memories) + self.context.memory_list.append(self.context.failure_task_memories) + self.context.memory_list.append(self.context.comparative_task_memories) + + task_memories: List[BaseMemory] = self.context.memory_list if not task_memories: logger.info("No task memories found for validation") @@ -37,7 +42,7 @@ class TaskMemoryValidationOp(BaseLLMOp): logger.info(f"Validated {len(validated_task_memories)} out of {len(task_memories)} task memories") # Update context - self.context.validated_task_memories = validated_task_memories + self.context.memory_list = validated_task_memories def _validate_single_task_memory(self, task_memory: BaseMemory) -> Dict[str, Any]: """Validate single task memory""" @@ -81,13 +86,13 @@ class TaskMemoryValidationOp(BaseLLMOp): "reason": "" if (is_valid and score >= validation_threshold) else f"Low validation score ({score:.2f}) or marked as invalid" } - except Exception as e: - logger.error(f"Error parsing validation response: {e}") + except Exception as e_inner: + logger.exception(f"Error parsing validation response: {e_inner}") return { "is_valid": False, "score": 0.0, "feedback": "", - "reason": f"Parse error: {str(e)}" + "reason": f"Parse error: {str(e_inner)}" } return self.llm.chat(messages=[Message(content=prompt)], callback_fn=parse_validation) diff --git a/reme_ai/summary/task/experience_validation_prompt.yaml b/reme_ai/summary/task/memory_validation_prompt.yaml similarity index 100% rename from reme_ai/summary/task/experience_validation_prompt.yaml rename to reme_ai/summary/task/memory_validation_prompt.yaml diff --git a/reme_ai/summary/task/simple_comparative_summary_op.py b/reme_ai/summary/task/simple_comparative_summary_op.py index a8569020..d9b32c24 100644 --- a/reme_ai/summary/task/simple_comparative_summary_op.py +++ b/reme_ai/summary/task/simple_comparative_summary_op.py @@ -45,7 +45,8 @@ class SimpleComparativeSummaryOp(BaseLLMOp): return self.llm.chat(messages=[Message(content=summary_prompt)], callback_fn=parse_content) def execute(self): - trajectories: List[Trajectory] = self.context.get("trajectories", []) + trajectories: list = self.context.get("trajectories", []) + trajectories: List[Trajectory] = [Trajectory(**x) if isinstance(x, dict) else x for x in trajectories] task_id_dict: Dict[str, List[Trajectory]] = {} for trajectory in trajectories: @@ -53,7 +54,7 @@ class SimpleComparativeSummaryOp(BaseLLMOp): task_id_dict[trajectory.task_id] = [] task_id_dict[trajectory.task_id].append(trajectory) - task_memory_list = [] + memory_list = [] for task_id, task_trajectories in task_id_dict.items(): task_trajectories: List[Trajectory] = sorted(task_trajectories, key=lambda x: x.score, reverse=True) if len(task_trajectories) < 2: @@ -62,8 +63,9 @@ class SimpleComparativeSummaryOp(BaseLLMOp): if task_trajectories[0].score > task_trajectories[-1].score: task_memories = self.compare_summary_trajectory(trajectory_a=task_trajectories[0], trajectory_b=task_trajectories[-1]) - task_memory_list.extend(task_memories) + memory_list.extend(task_memories) - self.context.comparative_summary_task_memories = task_memory_list - for tm in task_memory_list: + self.context.response.answer = json.dumps([x.model_dump() for x in memory_list]) + self.context.memory_list = memory_list + for tm in memory_list: logger.info(f"add task memory when_to_use={tm.when_to_use}\ncontent={tm.content}") diff --git a/reme_ai/summary/task/trajectory_segmentation_op.py b/reme_ai/summary/task/trajectory_segmentation_op.py index f81d7d98..0e5b4ea6 100644 --- a/reme_ai/summary/task/trajectory_segmentation_op.py +++ b/reme_ai/summary/task/trajectory_segmentation_op.py @@ -85,14 +85,16 @@ class TrajectorySegmentationOp(BaseLLMOp): return self.llm.chat(messages=[Message(content=prompt)], callback_fn=parse_segmentation, default_value=[trajectory.messages]) - def _format_trajectory_content(self, trajectory: Trajectory) -> str: + @staticmethod + def _format_trajectory_content(trajectory: Trajectory) -> str: """Format trajectory content for LLM processing""" content = "" for i, step in enumerate(trajectory.messages): content += f"Step {i + 1} ({step.role.value}):\n{step.content}\n\n" return content - def _parse_segmentation_response(self, response: str) -> List[int]: + @staticmethod + def _parse_segmentation_response(response: str) -> List[int]: """Parse segmentation response from LLM""" segment_points = []