reme-ai 0.1.0__py3-none-any.whl

reme_ai/summary/task/comparative_extraction_op.py

@@ -0,0 +1,233 @@
+from typing import List, Tuple, Optional
+
+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
+
+
+@C.register_op()
+class ComparativeExtractionOp(BaseLLMOp):
+    file_path: str = __file__
+
+    def execute(self):
+        """Extract comparative task memories by comparing different scoring trajectories"""
+        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", [])
+
+        comparative_task_memories = []
+
+        # Soft comparison: highest score vs lowest score
+        if len(all_trajectories) >= 2 and self.op_params.get("enable_soft_comparison", True):
+            highest_traj, lowest_traj = self._find_highest_lowest_scoring_trajectories(all_trajectories)
+            if highest_traj and lowest_traj and highest_traj.score > lowest_traj.score:
+                logger.info(
+                    f"Extracting soft comparative task memories: highest ({highest_traj.score:.2f}) vs lowest ({lowest_traj.score:.2f})")
+                soft_task_memories = self._extract_soft_comparative_task_memory(highest_traj, lowest_traj)
+                comparative_task_memories.extend(soft_task_memories)
+
+        # Hard comparison: success vs failure (if similarity search is enabled)
+        if (success_trajectories and failure_trajectories and
+                self.op_params.get("enable_similarity_comparison", False)):
+
+            similar_pairs = self._find_similar_step_sequences(success_trajectories, failure_trajectories)
+            logger.info(f"Found {len(similar_pairs)} similar pairs for hard comparison")
+
+            for success_steps, failure_steps, similarity_score in similar_pairs:
+                hard_task_memories = self._extract_hard_comparative_task_memory(success_steps, failure_steps,
+                                                                                similarity_score)
+                comparative_task_memories.extend(hard_task_memories)
+
+        logger.info(f"Extracted {len(comparative_task_memories)} comparative task memories")
+
+        # Add task memories to context
+        self.context.comparative_task_memories = comparative_task_memories
+
+    @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
+
+        # Filter trajectories with valid scores
+        valid_trajectories = [traj for traj in trajectories if traj.score is not None]
+
+        if len(valid_trajectories) < 2:
+            logger.warning("Not enough trajectories with valid scores for comparison")
+            return None, None
+
+        # Sort by score
+        sorted_trajectories = sorted(valid_trajectories, key=lambda x: x.score, reverse=True)
+
+        highest_traj = sorted_trajectories[0]
+        lowest_traj = sorted_trajectories[-1]
+
+        return highest_traj, lowest_traj
+
+    @staticmethod
+    def _get_trajectory_score(trajectory: Trajectory) -> Optional[float]:
+        """Get trajectory score"""
+        return trajectory.score
+
+    def _extract_soft_comparative_task_memory(self, higher_traj: Trajectory, lower_traj: Trajectory) -> List[
+        BaseMemory]:
+        """Extract soft comparative task memory (high score vs low score)"""
+        higher_steps = self._get_trajectory_steps(higher_traj)
+        lower_steps = self._get_trajectory_steps(lower_traj)
+        higher_score = self._get_trajectory_score(higher_traj)
+        lower_score = self._get_trajectory_score(lower_traj)
+
+        prompt = self.prompt_format(
+            prompt_name="soft_comparative_step_task_memory_prompt",
+            higher_steps=merge_messages_content(higher_steps),
+            lower_steps=merge_messages_content(lower_steps),
+            higher_score=f"{higher_score:.2f}",
+            lower_score=f"{lower_score:.2f}"
+        )
+
+        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)
+
+    def _extract_hard_comparative_task_memory(self, success_steps: List[Message],
+                                              failure_steps: List[Message], similarity_score: float) -> List[
+        BaseMemory]:
+        """Extract hard comparative task memory (success vs failure)"""
+        prompt = self.prompt_format(
+            prompt_name="hard_comparative_step_task_memory_prompt",
+            success_steps=merge_messages_content(success_steps),
+            failure_steps=merge_messages_content(failure_steps),
+            similarity_score=similarity_score
+        )
+
+        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)
+
+    @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
+            all_steps = []
+            for segment in trajectory.segments:
+                all_steps.extend(segment)
+            return all_steps
+        else:
+            return trajectory.messages
+
+    def _find_similar_step_sequences(self, success_trajectories: List[Trajectory],
+                                     failure_trajectories: List[Trajectory]) -> List[
+        Tuple[List[Message], List[Message], float]]:
+        """Find similar step sequences for comparison"""
+        if not self.op_params.get("enable_similarity_comparison", False):
+            return []
+
+        try:
+            similar_pairs = []
+
+            # Get step sequences
+            success_step_sequences = []
+            for traj in success_trajectories:
+                if hasattr(traj.metadata, 'segments') and traj.metadata["segments"]:
+                    success_step_sequences.extend(traj.metadata["segments"])
+                else:
+                    success_step_sequences.append(traj.messages)
+
+            failure_step_sequences = []
+            for traj in failure_trajectories:
+                if hasattr(traj.metadata, 'segments') and traj.metadata["segments"]:
+                    failure_step_sequences.extend(traj.metadata["segments"])
+                else:
+                    failure_step_sequences.append(traj.messages)
+
+            # Limit comparison count to avoid computational overload
+            max_sequences = self.op_params.get("max_similarity_sequences", 5)
+            success_step_sequences = success_step_sequences[:max_sequences]
+            failure_step_sequences = failure_step_sequences[:max_sequences]
+
+            if not success_step_sequences or not failure_step_sequences:
+                return []
+
+            # Generate text representation for embedding
+            success_texts = [merge_messages_content(seq) for seq in success_step_sequences]
+            failure_texts = [merge_messages_content(seq) for seq in failure_step_sequences]
+
+            # Get embedding vectors
+            if hasattr(self.context, 'vector_store') and self.context.vector_store and hasattr(
+                    self.context.vector_store, 'embedding_model'):
+                success_embeddings = self.context.vector_store.embedding_model.get_embeddings(success_texts)
+                failure_embeddings = self.context.vector_store.embedding_model.get_embeddings(failure_texts)
+
+                # Calculate similarity and find most similar pairs
+                similarity_threshold = self.op_params.get("similarity_threshold", 0.3)
+
+                for i, s_emb in enumerate(success_embeddings):
+                    for j, f_emb in enumerate(failure_embeddings):
+                        similarity = self._calculate_cosine_similarity(s_emb, f_emb)
+
+                        if similarity > similarity_threshold:
+                            similar_pairs.append((
+                                success_step_sequences[i],
+                                failure_step_sequences[j],
+                                similarity
+                            ))
+
+                # Return top most similar pairs
+                max_pairs = self.op_params.get("max_similarity_pairs", 3)
+                return sorted(similar_pairs, key=lambda x: x[2], reverse=True)[:max_pairs]
+
+        except Exception as e:
+            logger.error(f"Error finding similar step sequences: {e}")
+
+        return []
+
+    @staticmethod
+    def _calculate_cosine_similarity(embedding1: List[float], embedding2: List[float]) -> float:
+        """Calculate cosine similarity"""
+        import numpy as np
+
+        vec1 = np.array(embedding1)
+        vec2 = np.array(embedding2)
+
+        # Calculate cosine similarity
+        dot_product = np.dot(vec1, vec2)
+        norm1 = np.linalg.norm(vec1)
+        norm2 = np.linalg.norm(vec2)
+
+        if norm1 == 0 or norm2 == 0:
+            return 0.0
+
+        return dot_product / (norm1 * norm2)

reme_ai/summary/task/failure_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 FailureExtractionOp(BaseLLMOp):
+    file_path: str = __file__
+
+    def execute(self):
+        """Extract task memories from failed trajectories"""
+        failure_trajectories: List[Trajectory] = self.context.get("failure_trajectories", [])
+
+        if not failure_trajectories:
+            logger.info("No failure trajectories found for extraction")
+            return
+
+        logger.info(f"Extracting task memories from {len(failure_trajectories)} failed trajectories")
+
+        failure_task_memories = []
+
+        # Process trajectories
+        for trajectory in failure_trajectories:
+            if hasattr(trajectory, 'segments') and trajectory.segments:
+                # Process segmented step sequences
+                for segment in trajectory.segments:
+                    task_memories = self._extract_failure_task_memory_from_steps(segment, trajectory)
+                    failure_task_memories.extend(task_memories)
+            else:
+                # Process entire trajectory
+                task_memories = self._extract_failure_task_memory_from_steps(trajectory.messages, trajectory)
+                failure_task_memories.extend(task_memories)
+
+        logger.info(f"Extracted {len(failure_task_memories)} failure task memories")
+
+        # Add task memories to context
+        self.context.failure_task_memories = failure_task_memories
+
+    def _extract_failure_task_memory_from_steps(self, steps: List[Message], trajectory: Trajectory) -> List[BaseMemory]:
+        """Extract task memory from failed step sequences"""
+        step_content = merge_messages_content(steps)
+        context = get_trajectory_context(trajectory, steps)
+
+        prompt = self.prompt_format(
+            prompt_name="failure_step_task_memory_prompt",
+            query=trajectory.metadata.get('query', ''),
+            step_sequence=step_content,
+            context=context,
+            outcome="failed"
+        )
+
+        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/memory_deduplication_op.py

@@ -0,0 +1,163 @@
+from typing import List
+
+from flowllm import C, BaseOp
+from loguru import logger
+
+from reme_ai.schema.memory import BaseMemory
+
+
+@C.register_op()
+class MemoryDeduplicationOp(BaseOp):
+    file_path: str = __file__
+
+    def execute(self):
+        """Remove duplicate task memories"""
+        # Get task memories to deduplicate
+        task_memories: List[BaseMemory] = self.context.memory_list
+
+        if not task_memories:
+            logger.info("No task memories found for deduplication")
+            return
+
+        logger.info(f"Starting deduplication for {len(task_memories)} task memories")
+
+        # Perform deduplication
+        deduplicated_task_memories = self._deduplicate_task_memories(task_memories)
+
+        logger.info(
+            f"Deduplication complete: {len(deduplicated_task_memories)} deduplicated task memories out of {len(task_memories)}")
+
+        # Update context
+        self.context.memory_list = deduplicated_task_memories
+
+    def _deduplicate_task_memories(self, task_memories: List[BaseMemory]) -> List[BaseMemory]:
+        """Remove duplicate task memories"""
+        if not task_memories:
+            return task_memories
+
+        similarity_threshold = self.op_params.get("similarity_threshold", 0.5)
+        workspace_id = self.context.get("workspace_id")
+
+        unique_task_memories = []
+
+        # Get existing task memory embeddings
+        existing_embeddings = self._get_existing_task_memory_embeddings(workspace_id)
+
+        for task_memory in task_memories:
+            # Generate embedding for current task memory
+            current_embedding = self._get_task_memory_embedding(task_memory)
+
+            if current_embedding is None:
+                logger.warning(f"Failed to generate embedding for task memory: {str(task_memory.when_to_use)[:50]}...")
+                continue
+
+            # Check similarity with existing task memories
+            if self._is_similar_to_existing_task_memories(current_embedding, existing_embeddings, similarity_threshold):
+                logger.debug(f"Skipping similar task memory: {str(task_memory.when_to_use)[:50]}...")
+                continue
+
+            # Check similarity with current batch task memories
+            if self._is_similar_to_current_task_memories(current_embedding, unique_task_memories, similarity_threshold):
+                logger.debug(f"Skipping duplicate in current batch: {str(task_memory.when_to_use)[:50]}...")
+                continue
+
+            # Add to unique task memories list
+            unique_task_memories.append(task_memory)
+            logger.debug(f"Added unique task memory: {str(task_memory.when_to_use)[:50]}...")
+
+        return unique_task_memories
+
+    def _get_existing_task_memory_embeddings(self, workspace_id: str) -> List[List[float]]:
+        """Get embeddings of existing task memories"""
+        try:
+            if not hasattr(self.context, 'vector_store') or not self.context.vector_store or not workspace_id:
+                return []
+
+            # Query existing task memory nodes
+            existing_nodes = self.context.vector_store.search(
+                query="...",  # Empty query to get all
+                workspace_id=workspace_id,
+                top_k=self.op_params.get("max_existing_task_memories", 1000)
+            )
+
+            # Extract embeddings
+            existing_embeddings = []
+            for node in existing_nodes:
+                if hasattr(node, 'embedding') and node.embedding:
+                    existing_embeddings.append(node.embedding)
+
+            logger.debug(
+                f"Retrieved {len(existing_embeddings)} existing task memory embeddings from workspace {workspace_id}")
+            return existing_embeddings
+
+        except Exception as e:
+            logger.warning(f"Failed to retrieve existing task memory embeddings: {e}")
+            return []
+
+    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:
+                return None
+
+            # Combine task memory description and content for embedding
+            text_for_embedding = f"{task_memory.when_to_use} {task_memory.content}"
+            embeddings = self.context.vector_store.embedding_model.get_embeddings([text_for_embedding])
+
+            if embeddings and len(embeddings) > 0:
+                return embeddings[0]
+            else:
+                logger.warning("Empty embedding generated for task memory")
+                return None
+
+        except Exception as e:
+            logger.error(f"Error generating embedding for task memory: {e}")
+            return None
+
+    def _is_similar_to_existing_task_memories(self, current_embedding: List[float],
+                                              existing_embeddings: List[List[float]],
+                                              threshold: float) -> bool:
+        """Check if current embedding is similar to existing embeddings"""
+        for existing_embedding in existing_embeddings:
+            similarity = self._calculate_cosine_similarity(current_embedding, existing_embedding)
+            if similarity > threshold:
+                logger.debug(f"Found similar existing task memory with similarity: {similarity:.3f}")
+                return True
+        return False
+
+    def _is_similar_to_current_task_memories(self, current_embedding: List[float],
+                                             current_task_memories: List[BaseMemory],
+                                             threshold: float) -> bool:
+        for existing_task_memory in current_task_memories:
+            existing_embedding = self._get_task_memory_embedding(existing_task_memory)
+            if existing_embedding is None:
+                continue
+
+            similarity = self._calculate_cosine_similarity(current_embedding, existing_embedding)
+            if similarity > threshold:
+                logger.debug(f"Found similar task memory in current batch with similarity: {similarity:.3f}")
+                return True
+        return False
+
+    @staticmethod
+    def _calculate_cosine_similarity(embedding1: List[float], embedding2: List[float]) -> float:
+        """Calculate cosine similarity"""
+        try:
+            import numpy as np
+
+            vec1 = np.array(embedding1)
+            vec2 = np.array(embedding2)
+
+            # Calculate cosine similarity
+            dot_product = np.dot(vec1, vec2)
+            norm1 = np.linalg.norm(vec1)
+            norm2 = np.linalg.norm(vec2)
+
+            if norm1 == 0 or norm2 == 0:
+                return 0.0
+
+            return dot_product / (norm1 * norm2)
+
+        except Exception as e:
+            logger.error(f"Error calculating cosine similarity: {e}")
+            return 0.0

reme_ai/summary/task/memory_validation_op.py

@@ -0,0 +1,108 @@
+import json
+import re
+from typing import List, Dict, Any
+
+from flowllm import C, BaseLLMOp
+from loguru import logger
+
+from reme_ai.schema import Message
+from reme_ai.schema.memory import BaseMemory
+
+
+@C.register_op()
+class MemoryValidationOp(BaseLLMOp):
+    file_path: str = __file__
+
+    def execute(self):
+        """Validate quality of extracted task memories"""
+
+        task_memories: List[BaseMemory] = []
+        task_memories.extend(self.context.get("success_task_memories", []))
+        task_memories.extend(self.context.get("failure_task_memories", []))
+        task_memories.extend(self.context.get("comparative_task_memories", []))
+
+        if not task_memories:
+            logger.info("No task memories found for validation")
+            return
+
+        logger.info(f"Validating {len(task_memories)} extracted task memories")
+
+        # Validate task memories
+        validated_task_memories = []
+
+        for task_memory in task_memories:
+            validation_result = self._validate_single_task_memory(task_memory)
+            if validation_result and validation_result.get("is_valid", False):
+                task_memory.score = validation_result.get("score", 0.0)
+                validated_task_memories.append(task_memory)
+            else:
+                reason = validation_result.get("reason", "Unknown reason") if validation_result else "Validation failed"
+                logger.warning(f"Task memory validation failed: {reason}")
+
+        logger.info(f"Validated {len(validated_task_memories)} out of {len(task_memories)} task memories")
+
+        # Update context
+        self.context.response.answer = json.dumps([x.model_dump() for x in validated_task_memories])
+        self.context.response.metadata["memory_list"] = validated_task_memories
+
+    def _validate_single_task_memory(self, task_memory: BaseMemory) -> Dict[str, Any]:
+        """Validate single task memory"""
+        validation_info = self._llm_validate_task_memory(task_memory)
+        logger.info(f"Validating: {validation_info}")
+        return validation_info
+
+    def _llm_validate_task_memory(self, task_memory: BaseMemory) -> Dict[str, Any]:
+        """Validate task memory using LLM"""
+        try:
+            prompt = self.prompt_format(
+                prompt_name="task_memory_validation_prompt",
+                condition=task_memory.when_to_use,
+                task_memory_content=task_memory.content)
+
+            def parse_validation(message: Message) -> Dict[str, Any]:
+                try:
+                    response_content = message.content
+
+                    # Parse validation result
+                    # Extract JSON blocks
+                    json_pattern = r'```json\s*([\s\S]*?)\s*```'
+                    json_blocks = re.findall(json_pattern, response_content)
+
+                    if json_blocks:
+                        parsed = json.loads(json_blocks[0])
+                    else:
+                        parsed = {}
+
+                    is_valid = parsed.get("is_valid", True)
+                    score = parsed.get("score", 0.5)
+
+                    # Set validation threshold
+                    validation_threshold = self.op_params.get("validation_threshold", 0.5)
+
+                    return {
+                        "is_valid": is_valid and score >= validation_threshold,
+                        "score": score,
+                        "feedback": response_content,
+                        "reason": "" if (
+                                    is_valid and score >= validation_threshold) else f"Low validation score ({score:.2f}) or marked as invalid"
+                    }
+
+                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_inner)}"
+                    }
+
+            return self.llm.chat(messages=[Message(content=prompt)], callback_fn=parse_validation)
+
+        except Exception as e:
+            logger.error(f"LLM validation failed: {e}")
+            return {
+                "is_valid": False,
+                "score": 0.0,
+                "feedback": "",
+                "reason": f"LLM validation error: {str(e)}"
+            }

reme_ai/summary/task/pdf_preprocess_op_wrapper.py

@@ -0,0 +1,50 @@
+from flowllm import C, BaseOp
+from loguru import logger
+
+from reme_ai.utils.miner_u_pdf_processor import MinerUPDFProcessor, chunk_pdf_content
+
+
+@C.register_op()
+class PDFPreprocessOp(BaseOp):
+    file_path: str = __file__
+
+    def execute(self):
+        """Process PDF files using MinerU and chunk content"""
+        pdf_path = self.context.get("pdf_path")
+        output_dir = self.context.get("output_dir")
+
+        if not pdf_path:
+            logger.error("No PDF path provided in context")
+            return
+
+        # Process PDF
+        processor = MinerUPDFProcessor(log_level="INFO")
+
+        try:
+            content_list, markdown_content = processor.process_pdf(
+                pdf_path=pdf_path,
+                output_dir=output_dir,
+                method=self.op_params.get("method", "auto"),
+                lang=self.op_params.get("lang"),
+                backend=self.op_params.get("backend", "pipeline")
+            )
+
+            # Create chunks if requested
+            chunks = []
+            if self.op_params.get("create_chunks", True):
+                max_length = self.op_params.get("max_chunk_length", 4000)
+                chunks = chunk_pdf_content(content_list, max_length=max_length)
+
+            # Store results in context
+            self.context.pdf_content_list = content_list
+            self.context.pdf_markdown_content = markdown_content
+            self.context.pdf_chunks = chunks
+
+            logger.info(f"PDF processing completed: {len(content_list)} content blocks, "
+                        f"{len(chunks)} chunks, {len(markdown_content)} characters of markdown")
+
+        except Exception as e:
+            logger.error(f"PDF processing failed: {e}")
+            self.context.pdf_content_list = []
+            self.context.pdf_markdown_content = ""
+            self.context.pdf_chunks = []