reme-ai 0.1.0__py3-none-any.whl

reme_ai/summary/task/simple_comparative_summary_op.py

@@ -0,0 +1,71 @@
+import json
+from typing import List, Dict
+
+from flowllm import C, BaseLLMOp
+from loguru import logger
+
+from reme_ai.schema import Message, Trajectory
+from reme_ai.schema.memory import BaseMemory, TaskMemory
+from reme_ai.utils.op_utils import merge_messages_content
+
+
+@C.register_op()
+class SimpleComparativeSummaryOp(BaseLLMOp):
+    file_path: str = __file__
+
+    def compare_summary_trajectory(self, trajectory_a: Trajectory, trajectory_b: Trajectory) -> List[BaseMemory]:
+        summary_prompt = self.prompt_format(prompt_name="summary_prompt",
+                                            execution_process_a=merge_messages_content(trajectory_a.messages),
+                                            execution_process_b=merge_messages_content(trajectory_b.messages),
+                                            summary_example=self.get_prompt("summary_example"))
+
+        def parse_content(message: Message):
+            content = message.content
+            task_memory_list = []
+            try:
+                content = content.split("```")[1].strip()
+                if content.startswith("json"):
+                    content = content.strip("json")
+
+                for tm_dict in json.loads(content):
+                    when_to_use = tm_dict.get("when_to_use", "").strip()
+                    task_memory_content = tm_dict.get("experience", "").strip()
+                    if when_to_use and task_memory_content:
+                        task_memory_list.append(TaskMemory(workspace_id=self.context.get("workspace_id", ""),
+                                                           when_to_use=when_to_use,
+                                                           content=task_memory_content,
+                                                           author=getattr(self.llm, 'model_name', 'system')))
+
+                return task_memory_list
+
+            except Exception as e:
+                logger.exception(f"parse content failed!\n{content}")
+                raise e
+
+        return self.llm.chat(messages=[Message(content=summary_prompt)], callback_fn=parse_content)
+
+    def execute(self):
+        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:
+            if trajectory.task_id not in task_id_dict:
+                task_id_dict[trajectory.task_id] = []
+            task_id_dict[trajectory.task_id].append(trajectory)
+
+        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:
+                continue
+
+            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])
+                memory_list.extend(task_memories)
+
+        self.context.response.answer = json.dumps([x.model_dump() for x in memory_list])
+        self.context.response.metadata["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}")

reme_ai/summary/task/simple_summary_op.py

@@ -0,0 +1,67 @@
+import json
+from typing import List
+
+from flowllm import C, BaseLLMOp
+from loguru import logger
+
+from reme_ai.schema import Message, Trajectory
+from reme_ai.schema.memory import BaseMemory, TaskMemory
+from reme_ai.utils.op_utils import merge_messages_content
+
+
+@C.register_op()
+class SimpleSummaryOp(BaseLLMOp):
+    file_path: str = __file__
+
+    def summary_trajectory(self, trajectory: Trajectory) -> List[BaseMemory]:
+        execution_process = merge_messages_content(trajectory.messages)
+        success_score_threshold: float = self.op_params.get("success_score_threshold", 0.9)
+        logger.info(f"success_score_threshold={success_score_threshold}")
+
+        execution_result = "success" if trajectory.score >= success_score_threshold else "fail"
+        summary_prompt = self.prompt_format(prompt_name="summary_prompt",
+                                            execution_process=execution_process,
+                                            execution_result=execution_result,
+                                            summary_example=self.get_prompt("summary_example"))
+
+        def parse_content(message: Message):
+            content = message.content
+            memory_list = []
+            try:
+                if "```" in content:
+                    content = content.split("```")[1].strip()
+
+                if content.startswith("json"):
+                    content = content.strip("json")
+
+                for exp_dict in json.loads(content):
+                    when_to_use = exp_dict.get("when_to_use", "").strip()
+                    memory = exp_dict.get("memory", "").strip()
+                    if when_to_use and memory:
+                        memory_list.append(TaskMemory(workspace_id=self.context.get("workspace_id", ""),
+                                                      when_to_use=when_to_use,
+                                                      content=memory,
+                                                      author=getattr(self.llm, "model_name", "system")))
+
+                return memory_list
+
+            except Exception as e:
+                logger.exception(f"parse content failed!\n{content}")
+                raise e
+
+        return self.llm.chat(messages=[Message(content=summary_prompt)], callback_fn=parse_content)
+
+    def execute(self):
+        trajectories: list = self.context.trajectories
+        trajectories: List[Trajectory] = [Trajectory(**x) if isinstance(x, dict) else x for x in trajectories]
+
+        memory_list: List[BaseMemory] = []
+        for trajectory in trajectories:
+            memories = self.summary_trajectory(trajectory)
+            if memories:
+                memory_list.extend(memories)
+
+        self.context.response.answer = json.dumps([x.model_dump() for x in memory_list])
+        self.context.response.metadata["memory_list"] = memory_list
+        for memory in memory_list:
+            logger.info(f"add memory: when_to_use={memory.when_to_use}\ncontent={memory.content}")

reme_ai/summary/task/success_extraction_op.py

@@ -0,0 +1,73 @@
+from typing import List
+
+from flowllm import C, BaseLLMOp
+from loguru import logger
+
+from reme_ai.schema import Message, Trajectory
+from reme_ai.schema.memory import BaseMemory, TaskMemory
+from reme_ai.utils.op_utils import merge_messages_content, parse_json_experience_response, get_trajectory_context
+
+
+@C.register_op()
+class SuccessExtractionOp(BaseLLMOp):
+    file_path: str = __file__
+
+    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 = self._extract_success_task_memory_from_steps(segment, trajectory)
+                    success_task_memories.extend(task_memories)
+            else:
+                # Process entire trajectory
+                task_memories = 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
+
+    def _extract_success_task_memory_from_steps(self, steps: List[Message], trajectory: Trajectory) -> List[BaseMemory]:
+        """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[BaseMemory]:
+            task_memories_data = parse_json_experience_response(message.content)
+            task_memories = []
+
+            for tm_data in task_memories_data:
+                task_memory = TaskMemory(
+                    workspace_id=self.context.get("workspace_id", ""),
+                    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 self.llm.chat(messages=[Message(content=prompt)], callback_fn=parse_task_memories)

reme_ai/summary/task/trajectory_preprocess_op.py

@@ -0,0 +1,76 @@
+import json
+from typing import List, Dict
+
+from flowllm import C, BaseOp
+from loguru import logger
+
+from reme_ai.schema import Trajectory
+
+
+@C.register_op()
+class TrajectoryPreprocessOp(BaseOp):
+    file_path: str = __file__
+
+    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]
+        new_trajectories: List[Trajectory] = []
+        for x in trajectories:
+            if isinstance(x, dict):
+                x["messages"] = self._modify_tool_calls(x["messages"])
+                new_trajectories.append(Trajectory(**x))
+            else:
+                new_trajectories.append(x)
+        trajectories = new_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.op_params.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
+        }
+
+    def _modify_tool_calls(self, messages: List[Dict]) -> List[Dict]:
+        new_messages = []
+
+        for msg in messages:
+            if 'tool_calls' in msg:
+                processed_tool_calls = []
+                for tool_call in msg['tool_calls']:
+                    tool_type = tool_call.get("type", "function")
+                    nested_data = tool_call.get(tool_type, {})
+                    tool_call.update({
+                        "arguments": json.loads(nested_data.get("arguments", "")),
+                        "name": nested_data.get("name", "")
+                    })
+                    tool_call.pop(tool_type)
+                    processed_tool_calls.append(tool_call)
+                msg['tool_calls'] = processed_tool_calls
+            new_messages.append(msg)
+
+        return new_messages

reme_ai/summary/task/trajectory_segmentation_op.py

@@ -0,0 +1,118 @@
+import json
+import re
+from typing import List
+
+from flowllm import C, BaseLLMOp
+from loguru import logger
+
+from reme_ai.schema import Message, Trajectory
+
+
+@C.register_op()
+class TrajectorySegmentationOp(BaseLLMOp):
+    file_path: str = __file__
+
+    def execute(self):
+        """Segment trajectories into meaningful steps"""
+        # Get trajectories from context
+        all_trajectories: List[Trajectory] = self.context.get("all_trajectories", [])
+        success_trajectories: List[Trajectory] = self.context.get("success_trajectories", [])
+        failure_trajectories: List[Trajectory] = self.context.get("failure_trajectories", [])
+
+        if not all_trajectories:
+            logger.warning("No trajectories found in context")
+            return
+
+        # Determine which trajectories to segment
+        target_trajectories = self._get_target_trajectories(all_trajectories, success_trajectories,
+                                                            failure_trajectories)
+
+        # Add segmentation info to trajectories
+        segmented_count = 0
+        for trajectory in target_trajectories:
+            segments = self._llm_segment_trajectory(trajectory)
+            trajectory.metadata["segments"] = segments
+            segmented_count += 1
+
+        logger.info(f"Segmented {segmented_count} trajectories")
+
+        # Update context with segmented trajectories
+
+    def _get_target_trajectories(self, all_trajectories: List[Trajectory],
+                                 success_trajectories: List[Trajectory],
+                                 failure_trajectories: List[Trajectory]) -> List[Trajectory]:
+        """Determine which trajectories to segment based on configuration"""
+        segment_target = self.op_params.get("segment_target", "all")
+
+        if segment_target == "success":
+            return success_trajectories
+        elif segment_target == "failure":
+            return failure_trajectories
+        else:
+            return all_trajectories
+
+    def _llm_segment_trajectory(self, trajectory: Trajectory) -> List[List[Message]]:
+        """Use LLM for trajectory segmentation"""
+        trajectory_content = self._format_trajectory_content(trajectory)
+
+        prompt = self.prompt_format(
+            prompt_name="step_segmentation_prompt",
+            query=trajectory.metadata.get('query', ''),
+            trajectory_content=trajectory_content,
+            total_steps=len(trajectory.messages))
+
+        def parse_segmentation(message: Message) -> List[List[Message]]:
+            content = message.content
+            segment_points = self._parse_segmentation_response(content)
+
+            # Segment trajectory based on segmentation points
+            segments = []
+            start_idx = 0
+
+            for end_idx in segment_points:
+                if start_idx < end_idx <= len(trajectory.messages):
+                    segments.append(trajectory.messages[start_idx:end_idx])
+                    start_idx = end_idx
+
+            # Add remaining steps
+            if start_idx < len(trajectory.messages):
+                segments.append(trajectory.messages[start_idx:])
+
+            return segments if segments else [trajectory.messages]
+
+        return self.llm.chat(messages=[Message(content=prompt)], callback_fn=parse_segmentation,
+                             default_value=[trajectory.messages])
+
+    @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
+
+    @staticmethod
+    def _parse_segmentation_response(response: str) -> List[int]:
+        """Parse segmentation response from LLM"""
+        segment_points = []
+
+        # Try to extract JSON format
+        json_pattern = r'```json\s*([\s\S]*?)\s*```'
+        json_blocks = re.findall(json_pattern, response)
+
+        if json_blocks:
+            try:
+                parsed = json.loads(json_blocks[0])
+                if isinstance(parsed, dict) and "segment_points" in parsed:
+                    segment_points = parsed["segment_points"]
+                elif isinstance(parsed, list):
+                    segment_points = parsed
+            except json.JSONDecodeError:
+                pass
+
+        # Fallback: extract numbers
+        if not segment_points:
+            numbers = re.findall(r'\b\d+\b', response)
+            segment_points = [int(num) for num in numbers if int(num) > 0]
+
+        return sorted(list(set(segment_points)))