diff --git a/reme/core/utils/llm_utils.py b/reme/core/utils/llm_utils.py index 28935c73..805ac890 100644 --- a/reme/core/utils/llm_utils.py +++ b/reme/core/utils/llm_utils.py @@ -2,9 +2,10 @@ import json import re +from loguru import logger from ..enumeration import Role -from ..schema import Message, MemoryNode +from ..schema import Message, Trajectory, MemoryNode def format_messages(messages: list[Message | dict], enable_system: bool = False) -> str: @@ -29,6 +30,140 @@ def format_messages(messages: list[Message | dict], enable_system: bool = False) return "\n".join(formatted_lines) +def merge_messages_content(messages: list[Message | dict]) -> str: + """Merge messages content into a formatted string representation. + + This function processes a list of messages (either Message objects or dicts) + and formats them into a structured string. Different message roles are + formatted differently: + - ASSISTANT: Includes reasoning content, main content, and tool calls + - USER: Includes the user content + - TOOL: Includes tool call results + + Each message is prefixed with a step number (starting from 0) to indicate + its position in the conversation sequence. + + Args: + messages: List of Message objects or dictionaries to merge. If a dict + is provided, it will be converted to a Message object. + + Returns: + Formatted string representation of all messages with step numbers. + Each message is separated by newlines and includes role information. + + Example: + ```python + messages = [ + Message(role=Role.USER, content="What's the weather?"), + Message(role=Role.ASSISTANT, content="Let me check", + tool_calls=[ToolCall(name="get_weather", arguments={})]) + ] + result = merge_messages_content(messages) + # Returns formatted string with step numbers and role information + ``` + """ + content_collector = [] + for i, message in enumerate(messages): + if isinstance(message, dict): + message = Message(**message) + + if message.role is Role.ASSISTANT: + line = ( + f"### step.{i} role={message.role.value} content=\n{message.reasoning_content}\n\n{message.content}\n" + ) + if message.tool_calls: + for tool_call in message.tool_calls: + line += f" - tool call={tool_call.name}\n params={tool_call.arguments}\n" + content_collector.append(line) + + elif message.role is Role.USER: + line = f"### step.{i} role={message.role.value} content=\n{message.content}\n" + content_collector.append(line) + + elif message.role is Role.TOOL: + line = f"### step.{i} role={message.role.value} tool call result=\n{message.content}\n" + content_collector.append(line) + + return "\n".join(content_collector) + + +def parse_json_experience_response(response: str) -> list[dict]: + """Parse JSON formatted experience response""" + try: + # Extract JSON blocks + json_pattern = r"```json\s*([\s\S]*?)\s*```" + json_blocks = re.findall(json_pattern, response) + + if json_blocks: + parsed = json.loads(json_blocks[0]) + + # Handle array format + if isinstance(parsed, list): + experiences = [] + for exp_data in parsed: + if isinstance(exp_data, dict) and ( + ("when_to_use" in exp_data and "experience" in exp_data) + or ("condition" in exp_data and "experience" in exp_data) + ): + experiences.append(exp_data) + + return experiences + + # Handle single object + elif isinstance(parsed, dict) and ( + ("when_to_use" in parsed and "experience" in parsed) + or ("condition" in parsed and "experience" in parsed) + ): + return [parsed] + + # Fallback: try to parse entire response + parsed = json.loads(response) + if isinstance(parsed, list): + return parsed + elif isinstance(parsed, dict): + return [parsed] + + except json.JSONDecodeError as e: + logger.warning(f"Failed to parse JSON experience response: {e}") + + return [] + + +def get_trajectory_context(trajectory: Trajectory, step_sequence: list[Message]) -> str: + """Get context of step sequence within trajectory""" + try: + # Find position of step sequence in trajectory + start_idx = 0 + for i, step in enumerate(trajectory.messages): + if step == step_sequence[0]: + start_idx = i + break + + # Extract before and after context + context_before = trajectory.messages[max(0, start_idx - 2) : start_idx] + context_after = trajectory.messages[start_idx + len(step_sequence) : start_idx + len(step_sequence) + 2] + + context = f"Query: {trajectory.metadata.get('query', 'N/A')}\n" + + if context_before: + context += ( + "Previous steps:\n" + + "\n".join( + [f"- {step.content[:100]}..." for step in context_before], + ) + + "\n" + ) + + if context_after: + context += "Following steps:\n" + "\n".join([f"- {step.content[:100]}..." for step in context_after]) + + return context + + except Exception as e: + logger.error(f"Error getting trajectory context: {e}") + return f"Query: {trajectory.metadata.get('query', 'N/A')}" + + def extract_content(text: str, language_tag: str = "json", greedy: bool = False): """Extracts content from Markdown code blocks and parses it if the tag is JSON.""" quantifier = ".*" if greedy else ".*?" diff --git a/reme/workflow/procedural_memory/__init__.py b/reme/workflow/procedural_memory/__init__.py index a3388dbb..91f1fcda 100644 --- a/reme/workflow/procedural_memory/__init__.py +++ b/reme/workflow/procedural_memory/__init__.py @@ -43,4 +43,4 @@ # examples # bench 里的llm ，辛苦改成 app = ReMeApp() app.default_llm -# clear && pre-commit run --all-files \ No newline at end of file +# clear && pre-commit run --all-files diff --git a/reme/workflow/procedural_memory/summarizer/__init__.py b/reme/workflow/procedural_memory/summarizer/__init__.py new file mode 100644 index 00000000..fefa976d --- /dev/null +++ b/reme/workflow/procedural_memory/summarizer/__init__.py @@ -0,0 +1,15 @@ +"""Summarizer operators for procedural memory workflow. + +This package exposes and registers summarization-related operators such as +`TrajectoryPreprocess` and `SuccessExtraction` to the global operator registry. +""" + +from ....core import R +from .trajectory_preprocess import TrajectoryPreprocess +from .success_extraction import SuccessExtraction + +__all__ = ["TrajectoryPreprocess", "SuccessExtraction"] + +for name in __all__: + tool_class = globals()[name] + R.op.register()(tool_class) diff --git a/reme/workflow/procedural_memory/summarizer/success_extraction.py b/reme/workflow/procedural_memory/summarizer/success_extraction.py new file mode 100644 index 00000000..f4b56c90 --- /dev/null +++ b/reme/workflow/procedural_memory/summarizer/success_extraction.py @@ -0,0 +1,95 @@ +"""Success extraction operation for task memory generation. + +This module provides operations to extract task memories from successful +trajectories, identifying patterns and strategies that lead to success. +""" + +from typing import List + +from loguru import logger + +from ....core.enumeration import MemoryType, Role +from ....core.op import BaseOp +from ....core.schema.memory_node import MemoryNode +from ....core.schema.message import Message, Trajectory +from ....core.utils.llm_utils import ( + get_trajectory_context, + merge_messages_content, + parse_json_experience_response, +) + + +class SuccessExtraction(BaseOp): + """Extract task memories from successful trajectories. + + This operation analyzes successful trajectories (or their segments) to + extract reusable patterns, strategies, and best practices that can be + applied to similar future tasks. + """ + + async def execute(self): + """Extract task memories from successful trajectories""" + success_trajectories: List[Trajectory] = self.context.success_trajectories + + if not success_trajectories: + logger.info("No success trajectories found for extraction") + return + + logger.info(f"Extracting task memories from {len(success_trajectories)} successful trajectories") + + success_task_memories = [] + + # Process trajectories + for trajectory in success_trajectories: + if "segments" in trajectory.metadata: + # Process segmented step sequences + for segment in trajectory.metadata["segments"]: + task_memories = await self._extract_success_task_memory_from_steps(segment, trajectory) + success_task_memories.extend(task_memories) + else: + # Process entire trajectory + task_memories = await self._extract_success_task_memory_from_steps(trajectory.messages, trajectory) + success_task_memories.extend(task_memories) + + logger.info(f"Extracted {len(success_task_memories)} success task memories") + + # Add task memories to context + self.context.success_task_memories = success_task_memories + + async def _extract_success_task_memory_from_steps( + self, + steps: List[Message], + trajectory: Trajectory, + ) -> List[MemoryNode]: + """Extract task memory from successful step sequences""" + step_content = merge_messages_content(steps) + context = get_trajectory_context(trajectory, steps) + + prompt = self.prompt_format( + prompt_name="success_step_task_memory_prompt", + query=trajectory.metadata.get("query", ""), + step_sequence=step_content, + context=context, + outcome="successful", + ) + + def parse_task_memories(message: Message) -> list[MemoryNode]: + task_memories_data = parse_json_experience_response(message.content) # extract content + task_memories = [] + + for tm_data in task_memories_data: + task_memory = MemoryNode( + memory_type=MemoryType.PROCEDURAL, + when_to_use=tm_data.get("when_to_use", tm_data.get("condition", "")), + content=tm_data.get("experience", ""), + author=getattr(self.llm, "model_name", "system"), + metadata=tm_data, + ) + task_memories.append(task_memory) + + return task_memories + + return await self.llm.chat( + messages=[Message(role=Role.USER, content=prompt)], + callback_fn=parse_task_memories, + ) diff --git a/reme/workflow/procedural_memory/summarizer/success_extraction.yaml b/reme/workflow/procedural_memory/summarizer/success_extraction.yaml new file mode 100644 index 00000000..3ceb7ce6 --- /dev/null +++ b/reme/workflow/procedural_memory/summarizer/success_extraction.yaml @@ -0,0 +1,42 @@ +success_step_task_memory_prompt: | + You are an expert AI analyst reviewing successful step sequences from an AI agent execution. + + Your task is to extract reusable, actionable step-level task memories that can guide future agent executions. + Focus on identifying specific patterns, techniques, and decision points that contributed to success. + + ANALYSIS FRAMEWORK: + ● STEP PATTERN ANALYSIS: Identify the specific sequence of actions that led to success + ● DECISION POINTS: Highlight critical decisions made during these steps + ● TECHNIQUE EFFECTIVENESS: Analyze why specific approaches worked well + ● REUSABILITY: Extract patterns that can be applied to similar scenarios + + EXTRACTION PRINCIPLES: + ● Focus on TRANSFERABLE TECHNIQUES and decision frameworks + ● Frame insights as actionable guidelines and best practices + + # Original Query + {query} + + # Step Sequence Analysis + {step_sequence} + + # Context Information + {context} + + # Outcome + This step sequence was part of a {outcome} trajectory. + + OUTPUT FORMAT: + Generate 1-3 step-level success insights as JSON objects: + ```json + [ + {{ + "when_to_use": "Specific conditions when this step pattern should be applied", + "experience": "Detailed description of the successful step pattern and why it works", + "tags": ["relevant", "keywords", "for", "categorization"], + "confidence": 0.8, + "step_type": "reasoning|action|observation|decision", + "tools_used": ["list", "of", "tools"] + }} + ] + ``` \ No newline at end of file diff --git a/reme/workflow/procedural_memory/summarizer/trajectory_preprocess.py b/reme/workflow/procedural_memory/summarizer/trajectory_preprocess.py new file mode 100644 index 00000000..e099f483 --- /dev/null +++ b/reme/workflow/procedural_memory/summarizer/trajectory_preprocess.py @@ -0,0 +1,59 @@ +"""Trajectory preprocessing operation for task memory generation. + +This module provides operations to preprocess and classify trajectories +into success and failure categories based on score thresholds. +""" + +from typing import Dict, List + +from loguru import logger + +from ....core.op import BaseOp +from ....core.schema.message import Trajectory + + +class TrajectoryPreprocess(BaseOp): + """Preprocess trajectories: validate and classify by success/failure. + + This operation classifies trajectories into success and failure categories + based on score thresholds, preparing them for downstream memory extraction + operations. + """ + + async def execute(self): + """Preprocess trajectories: validate and classify""" + trajectories: list = self.context.get("trajectories", []) + trajectories: List[Trajectory] = [Trajectory(**x) if isinstance(x, dict) else x for x in trajectories] + + # Classify trajectories + classified = self._classify_trajectories(trajectories) + logger.info( + f"Classified trajectories - Success: {len(classified['success'])}, " + f"Failure: {len(classified['failure'])}, All: {len(classified['all'])}", + ) + + # Set context for downstream operators + self.context.success_trajectories = classified["success"] + self.context.failure_trajectories = classified["failure"] + self.context.all_trajectories = classified["all"] + + def _classify_trajectories(self, trajectories: List[Trajectory]) -> Dict[str, List[Trajectory]]: + """Classify trajectories based on score threshold""" + success_trajectories = [] + failure_trajectories = [] + + success_threshold = self.context.get("success_threshold", 1.0) + + for traj in trajectories: + is_success = traj.score >= success_threshold + + if is_success: + success_trajectories.append(traj) + else: + failure_trajectories.append(traj) + + return { + "success": success_trajectories, + "failure": failure_trajectories, + "all": trajectories, + }