diff --git a/reme/config/default.yaml b/reme/config/default.yaml index 3a7e7073..43e2ba2d 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -59,3 +59,145 @@ token_counters: model_name: Qwen/Qwen3-Coder-30B-A3B-Instruct use_mirror: true +flow: + retrieve_task_memory: + flow_content: BuildQuery() >> MemoryRetrieval() >> RerankMemory() >> RewriteMemory() + description: "Retrieves the most relevant top-k memory experiences from historical data based on the current query to enhance task-solving capabilities" + parameters: + type: object + properties: + query: + type: string + description: "The search query string for retrieving relevant memories. Either query or messages must be provided." + messages: + type: array + description: "A list of conversation messages to build the query from. Either query or messages must be provided." + enable_llm_build: + type: boolean + description: "Whether to use LLM to build query from messages (default: true)." + top_k: + type: integer + description: "Number of top results to retrieve (default: 5)." + threshold_score: + type: number + description: "Optional minimum score threshold for filtering retrieved memories." + enable_llm_rerank: + type: boolean + description: "Whether to enable LLM-based reranking (default: false)." + enable_score_filter: + type: boolean + description: "Whether to enable score-based filtering (default: false)." + min_score_threshold: + type: number + description: "Minimum combined score threshold for filtering memories (default: 0.3)." + enable_llm_rewrite: + type: boolean + description: "Whether to use LLM to rewrite context messages (default: false)." + required: [] + + summary_task_memory: + flow_content: TrajectoryPreprocess() >> (SuccessExtraction()|FailureExtraction()|ComparativeExtraction()) >> MemoryValidation() >> MemoryDeduplication() + description: "Summarizes conversation trajectories or messages into structured memory representations for long-term storage" + parameters: + type: object + properties: + trajectories: + type: array + description: "A list of conversation trajectory information, including message content and score." + success_threshold: + type: number + description: "Score threshold for classifying trajectories as successful (default: 1.0)." + enable_soft_comparison: + type: boolean + description: "Whether to enable soft comparison between highest and lowest scoring trajectories (default: true)." + enable_similarity_comparison: + type: boolean + description: "Whether to enable similarity-based comparison between success and failure trajectories (default: true)." + max_similarity_sequences: + type: integer + description: "Maximum number of sequences to compare for similarity (default: 5)." + similarity_threshold: + type: number + description: "Similarity threshold for comparing trajectories (default: 0.5)." + max_similarity_pairs: + type: integer + description: "Maximum number of similar pairs to extract from comparison (default: 3)." + validation_threshold: + type: number + description: "Minimum validation score threshold for accepting task memories (default: 0.5)." + max_existing_task_memories: + type: integer + description: "Maximum number of existing task memories to check for deduplication (default: 1000)." + required: + - trajectories + + add_task_memory: + flow_content: MemoryAddition() + description: "Add task memories to the vector store" + parameters: + type: object + properties: + memory_list: + type: array + description: "A list of task memory to add to the vector store." + required: + - memory_list + + delete_task_memory: + flow_content: MemoryDeletion() + description: "Delete task memories when utility/freq < utility_threshold and freq >= freq_threshold" + parameters: + type: object + properties: + freq_threshold: + type: integer + description: "The retrieved frequency threshold for deleting task memory." + utility_threshold: + type: number + description: "The utility/freq threshold for deleting task memory." + required: + - freq_threshold + - utility_threshold + + record_task_memory: + flow_content: UpdateMemoryMetadata() + description: "Update the freq & utility attributes of retrieved task memories" + parameters: + type: object + properties: + memory_list: + type: array + description: "A list of retrieved task memory corresponding to the current task." + update_utility: + type: boolean + description: "Whether to update the utility attribute of the retrieved task memory." + required: + - memory_list + - update_utility + + load_memory: + flow_content: LoadMemory() + description: "Load memories from disk into the vector store" + parameters: + type: object + properties: + load_file_path: + type: string + description: "The path to the memories file." + clear_existing: + type: boolean + description: "If True, clears existing memories before loading (default: False)." + required: + - load_file_path + + dump_memory: + flow_content: DumpMemory() + description: "Dump the vector store memories to disk" + parameters: + type: object + properties: + dump_file_path: + type: string + description: "The path to the memories file." + required: + - dump_file_path diff --git a/reme/reme.py b/reme/reme.py index cab3a2c4..5e5b9caf 100644 --- a/reme/reme.py +++ b/reme/reme.py @@ -287,8 +287,20 @@ class ReMe(Application): raise NotImplementedError procedural_summarizer: BaseMemoryAgent - if version in ["default", "v1", "v2", "halumem"]: - procedural_summarizer = ProceduralSummarizer(tools=[]) + if version in ["default", "v1"]: + procedural_summarizer = ProceduralSummarizer( + tools=[ + VectorRetrieveMemory( + add_memory_type_target=False, + enable_thinking_params=True, + metadata_desc=None, + top_k=5, + ), + AddMemory(enable_thinking_params=True,add_when_to_use=True, metadata_desc=metadata_summary), + DeleteMemory(enable_thinking_params=True), + UpdateMemory(enable_thinking_params=True,add_when_to_use=True, metadata_desc=metadata_summary), + ], + ) else: raise NotImplementedError @@ -466,8 +478,18 @@ class ReMe(Application): raise NotImplementedError procedural_retriever: BaseMemoryAgent - if version in ["default", "v1", "v2", "halumem"]: - procedural_retriever = ProceduralRetriever(tools=[]) + if version in ["default", "v1"]: + procedural_retriever = ProceduralRetriever( + tools=[ + VectorRetrieveMemory( + add_memory_type_target=True, + enable_thinking_params=True, + metadata_desc=metadata_retrieve, + top_k=top_k, + ), + ReadHistory(enable_thinking_params=True), + ], + ) else: raise NotImplementedError diff --git a/reme/workflow/procedural_memory/__init__.py b/reme/workflow/procedural_memory/__init__.py index 91f1fcda..5f901ce0 100644 --- a/reme/workflow/procedural_memory/__init__.py +++ b/reme/workflow/procedural_memory/__init__.py @@ -1,46 +1,45 @@ -# summary retriever vector store 的 op +from ...core import R -# 1. BaseAsyncOp -> reme.core.op.BaseOp -# 2. @C.register_op() -> R.op.register()(MergeMemoryOp) -# for name in __all__: -# tool_class = globals()[name] -# R.op.register()(tool_class) -# 3. class name 不一定叫 op -# 4. async def async_execute(self): —》async def execute(self): -# 5. self.llm.chat -# 6. file_path: str = __file__ 不需要 -# 7. self.op_params.get("enable_llm_rerank", True) 都改成 self.context.get("enable_llm_rerank", True) -# -# -# # 跑起来 reme = "reme_ai.main:main" -> reme = "reme.reme_app:main" -# -# app = ReMeApp() -# -# -# def test_search(): -# """Test search tool operations. -# -# Tests DashscopeSearch, MockSearch, and TavilySearch operations -# with a sample query to verify they work correctly. -# """ -# from reme.tool.search import DashscopeSearch, MockSearch, TavilySearch -# -# query = "今天杭州的天气如何？" -# -# for op in [ -# DashscopeSearch(), -# MockSearch(), -# TavilySearch(), -# ]: -# print("\n" + "=" * 60) -# print(f"Testing {op.__class__.__name__}") -# print("=" * 60) -# print(f"Query: {query}") -# output = asyncio.run(op.call(query=query, service_context=app.service_context)) -# -# self.context.query -# app.service_context 保证了 self.llm emb vectorstore +from .dump_memory import DumpMemory +from .load_memory import LoadMemory -# examples -# bench 里的llm ，辛苦改成 app = ReMeApp() app.default_llm -# clear && pre-commit run --all-files +from .summary.trajectory_preprocess import TrajectoryPreprocess +from .summary.trajectory_segmentation import TrajectorySegmentation +from .summary.success_extraction import SuccessExtraction +from .summary.failure_extraction import FailureExtraction +from .summary.comparative_extraction import ComparativeExtraction +from .summary.memory_validation import MemoryValidation +from .summary.memory_deduplication import MemoryDeduplication +from .summary.memory_addition import MemoryAddition + +from .retrieve.build_query import BuildQuery +from .retrieve.memory_deletion import MemoryDeletion +from .retrieve.memory_retrieval import MemoryRetrieval +from .retrieve.merge_memory import MergeMemory +from .retrieve.rerank_memory import RerankMemory +from .retrieve.rewrite_memory import RewriteMemory +from .retrieve.update_memory_metadata import UpdateMemoryMetadata + +__all__ = [ + "DumpMemory", + "LoadMemory", + "TrajectoryPreprocess", + "TrajectorySegmentation", + "SuccessExtraction", + "FailureExtraction", + "ComparativeExtraction", + "MemoryValidation", + "MemoryDeduplication", + "MemoryAddition", + "BuildQuery", + "MemoryDeletion", + "MemoryRetrieval", + "MergeMemory", + "RerankMemory", + "RewriteMemory", + "UpdateMemoryMetadata", +] + +for name in __all__: + tool_class = globals()[name] + R.op.register()(tool_class) diff --git a/reme/workflow/procedural_memory/dump_memory.py b/reme/workflow/procedural_memory/dump_memory.py new file mode 100644 index 00000000..8ca993d0 --- /dev/null +++ b/reme/workflow/procedural_memory/dump_memory.py @@ -0,0 +1,63 @@ +"""Operation for dumping memories from vector store to JSONL file.""" + +import json +from pathlib import Path +from typing import List + +from loguru import logger + +from ...core.op import BaseOp +from ...core.schema.memory_node import MemoryNode +from ...core.schema.vector_node import VectorNode + + +class DumpMemory(BaseOp): + """Operation that dumps memories from vector store to a JSONL file. + + This operation retrieves all memories from the vector store, converts them + to MemoryNode objects, and writes them to a JSONL file (one JSON object + per line) for backup or export purposes. + """ + + async def execute(self): + """Execute the memory dump operation. + + Dumps all memories from the vector store to a JSONL file: + 1. Retrieves all VectorNodes from the vector store + 2. Converts them to MemoryNode objects + 3. Writes each MemoryNode as a JSON line to the output file + + Expected context attributes: + dump_file_path: Path to the output JSONL file. + + Sets context attributes: + dumped_count: Number of memories dumped to the file. + """ + # Support both dump_file_path and path for backward compatibility + dump_file_path: str = self.context.dump_file_path + if not dump_file_path: + logger.error("dump_file_path is required in context") + return + + file_path = Path(dump_file_path) + file_path.parent.mkdir(parents=True, exist_ok=True) + + # Retrieve all nodes from vector store + vector_nodes: List[VectorNode] = await self.vector_store.list() + logger.info(f"Retrieved {len(vector_nodes)} nodes from vector store") + + # Convert to MemoryNodes and write to JSONL file + dumped_count = 0 + with open(file_path, "w", encoding="utf-8") as f: + for node in vector_nodes: + try: + memory = MemoryNode.from_vector_node(node) + # Write as JSON line (one JSON object per line) + json_line = json.dumps(memory.model_dump(exclude_none=True), ensure_ascii=False) + f.write(json_line + "\n") + dumped_count += 1 + except Exception as e: + logger.warning(f"Failed to convert and dump node {node.vector_id}: {e}") + continue + + logger.info(f"Dumped {dumped_count} memories to {dump_file_path}") diff --git a/reme/workflow/procedural_memory/load_memory.py b/reme/workflow/procedural_memory/load_memory.py new file mode 100644 index 00000000..e0236656 --- /dev/null +++ b/reme/workflow/procedural_memory/load_memory.py @@ -0,0 +1,82 @@ +"""Operation for loading memories from JSONL file to vector store.""" + +import json +import asyncio +from pathlib import Path +from typing import List + +from loguru import logger + +from ...core.op import BaseOp +from ...core.schema.memory_node import MemoryNode +from ...core.schema.vector_node import VectorNode + + +class LoadMemory(BaseOp): + """Operation that loads memories from a JSONL file to vector store. + + This operation reads MemoryNode objects from a JSONL file (one JSON object + per line), converts them to VectorNode objects, and inserts them into the + vector store. + """ + + async def execute(self): + """Execute the memory load operation. + + Loads memories from a JSONL file to the vector store: + 1. Reads each line from the JSONL file + 2. Parses JSON and creates MemoryNode objects + 3. Converts MemoryNodes to VectorNodes + 4. Inserts them into the vector store + + Expected context attributes: + load_file_path: Path to the input JSONL file. + clear_existing: Optional. If True, clears existing memories before loading (default: False). + + Sets context attributes: + loaded_count: Number of memories loaded from the file. + """ + load_file_path: str = self.context.load_file_path + if not load_file_path: + logger.error("load_file_path is required in context") + return + + file_path = Path(load_file_path) + if not file_path.exists(): + logger.error(f"File not found: {load_file_path}") + return + + try: + # Attempt to retrieve the event loop associated with the current thread + loop = asyncio.get_running_loop() + print(f"Running event loop found: {loop}") + except RuntimeError: + # Start a new event loop to run the coroutine to completion + print("No running event loop found, starting a new one") + + clear_existing: bool = self.context.get("clear_existing", False) + if clear_existing: + await self.vector_store.delete_all() + logger.info("Cleared existing memories from vector store") + + # Read and parse JSONL file + memory_nodes: List[MemoryNode] = [] + with open(file_path, "r", encoding="utf-8") as f: + for line_num, line in enumerate(f, 1): + line = line.strip() + if not line: + continue + try: + data = json.loads(line) + memory = MemoryNode.model_validate(data) + memory_nodes.append(memory) + except Exception as e: + logger.warning(f"Failed to parse line {line_num} in {load_file_path}: {e}") + continue + logger.info(f"Parsed {len(memory_nodes)} memories from {load_file_path}") + + # Convert to VectorNodes and insert into vector store + if memory_nodes: + vector_nodes: List[VectorNode] = [memory.to_vector_node() for memory in memory_nodes] + await self.vector_store.insert(nodes=vector_nodes) + logger.info(f"Loaded {len(memory_nodes)} memories into vector store") diff --git a/reme/workflow/procedural_memory/retriever/__init__.py b/reme/workflow/procedural_memory/retrieve/__init__.py similarity index 100% rename from reme/workflow/procedural_memory/retriever/__init__.py rename to reme/workflow/procedural_memory/retrieve/__init__.py diff --git a/reme/workflow/procedural_memory/retrieve/build_query.py b/reme/workflow/procedural_memory/retrieve/build_query.py new file mode 100644 index 00000000..dc121508 --- /dev/null +++ b/reme/workflow/procedural_memory/retrieve/build_query.py @@ -0,0 +1,58 @@ +"""Query building operation module. + +This module provides functionality to build retrieval queries from either +explicit query strings or conversation messages, optionally using LLM to +generate optimized queries. +""" + +from loguru import logger + +from ....core.enumeration import Role +from ....core.op import BaseOp +from ....core.schema.message import Message +from ..utils import merge_messages_content + + +class BuildQuery(BaseOp): + """Build retrieval query from context or messages. + + This operation constructs a query string for memory retrieval. It can use + an explicit query from context, or generate one from conversation messages + using either LLM-based generation or simple message concatenation. + """ + + async def execute(self): + """Execute the query building operation. + + Builds a query string from either: + 1. An explicit query in the context + 2. Conversation messages (using LLM or simple concatenation) + + Stores the built query in context.query. + """ + if "query" in self.context: + query = self.context.query + + elif "messages" in self.context: + if self.context.get("enable_llm_build", True): + execution_process = merge_messages_content(self.context.messages) + prompt = self.prompt_format(prompt_name="query_build", execution_process=execution_process) + message = await self.llm.chat(messages=[Message(role=Role.USER, content=prompt)]) + query = message.content + + else: + context_parts = [] + message_summaries = [] + for message in self.context.messages[-3:]: # Last 3 messages + content = message.content[:200] + "..." if len(message.content) > 200 else message.content + message_summaries.append(f"- {message.role.value}: {content}") + if message_summaries: + context_parts.append("Recent messages:\n" + "\n".join(message_summaries)) + + query = "\n\n".join(context_parts) + + else: + raise RuntimeError("query or messages is required!") + + logger.info(f"build.query={query}") + self.context.query = query diff --git a/reme/workflow/procedural_memory/retrieve/build_query.yaml b/reme/workflow/procedural_memory/retrieve/build_query.yaml new file mode 100644 index 00000000..7017cec7 --- /dev/null +++ b/reme/workflow/procedural_memory/retrieve/build_query.yaml @@ -0,0 +1,6 @@ +query_build: | + # Execution Process + {execution_process} + + Read through the entire execution process to understand which part is currently being executed. + Generate a `query` that reflects the current state, which will later be used to search for similar problems in the database and help resolve the issue at hand. diff --git a/reme/workflow/procedural_memory/retrieve/memory_deletion.py b/reme/workflow/procedural_memory/retrieve/memory_deletion.py new file mode 100644 index 00000000..90315f74 --- /dev/null +++ b/reme/workflow/procedural_memory/retrieve/memory_deletion.py @@ -0,0 +1,54 @@ +"""Operation for deleting memories from the vector store.""" + +import json +from typing import List + +from loguru import logger + +from ....core.op import BaseOp +from ....core.schema.vector_node import VectorNode + + +class MemoryDeletion(BaseOp): + """Operation that deletes memories from the vector store. + + This operation identifies memories to delete based on frequency and utility + thresholds, then deletes them. Memories with frequency >= freq_threshold + and utility/frequency ratio < utility_threshold are deleted. + """ + + async def execute(self): + """Execute the memory deletion operation. + + Identifies and deletes memories from the vector store: + 1. Lists all nodes from the vector store + 2. Identifies memories that meet deletion criteria based on thresholds + 3. Deletes identified memories from the vector store + 4. Stores deletion count in response.metadata["result"] + + The deletion criteria: + - Memory frequency must be >= freq_threshold + - Memory utility/frequency ratio must be < utility_threshold + + Expected context attributes: + freq_threshold: Minimum frequency threshold for consideration. + utility_threshold: Maximum utility/frequency ratio threshold. + """ + + # Step 1: Identify memories to delete based on thresholds + freq_threshold: int = self.context.freq_threshold + utility_threshold: float = self.context.utility_threshold + nodes: List[VectorNode] = await self.vector_store.list() + + deleted_memory_ids = [] + for node in nodes: + freq = node.metadata.get("freq", 0) + utility = node.metadata.get("utility", 0) + if freq >= freq_threshold: + if freq > 0 and utility * 1.0 / freq < utility_threshold: + deleted_memory_ids.append(node.vector_id) + + # Step 2: Execute deletion if there are any IDs to delete + if deleted_memory_ids: + await self.vector_store.delete(vector_ids=deleted_memory_ids) + logger.info(f"Deleted {len(deleted_memory_ids)} memories: {json.dumps(deleted_memory_ids, indent=2)}") diff --git a/reme/workflow/procedural_memory/retrieve/memory_retrieval.py b/reme/workflow/procedural_memory/retrieve/memory_retrieval.py new file mode 100644 index 00000000..c5c9b5b8 --- /dev/null +++ b/reme/workflow/procedural_memory/retrieve/memory_retrieval.py @@ -0,0 +1,68 @@ +"""Operation for recalling memories from the vector store based on a query.""" + +from typing import List + +from loguru import logger + +from ....core.op import BaseOp +from ....core.schema.memory_node import MemoryNode +from ....core.schema.vector_node import VectorNode + + +class MemoryRetrieval(BaseOp): + """Operation that retrieves relevant memories from the vector store. + + This operation performs a semantic search on the vector store to find + memories relevant to a given query. It supports optional score filtering + and deduplication based on memory content. + """ + + async def execute(self): + """Execute the memory recall operation. + + Performs a semantic search in the vector store using the provided query, + retrieves the top-k most relevant memories, and optionally filters them + by a score threshold. Duplicate memories (based on content) are removed. + + Expected context attributes: + query: The search query string. + top_k: Number of top results to retrieve (default: 3). + + Expected context attributes (optional): + threshold_score: Optional minimum score threshold for filtering. + + Sets response.metadata: + memory_list: List of retrieved MemoryNode objects. + """ + top_k: int = self.context.get("top_k", 5) + + query: str = self.context.get("query", "") + assert query, "query should be not empty!" + + # Perform semantic search + nodes: List[VectorNode] = await self.vector_store.search( + query=query, + limit=top_k, + filters=None, + ) + + # Convert VectorNodes to MemoryNodes and deduplicate by content + memory_list: List[MemoryNode] = [] + memory_content_set: set[str] = set() # for deduplication + for node in nodes: + try: + memory = MemoryNode.from_vector_node(node) + if memory.content not in memory_content_set: + memory_list.append(memory) + memory_content_set.add(memory.content) + except Exception as e: + logger.warning(f"Failed to convert VectorNode to MemoryNode: {e}") + continue + logger.info(f"Retrieved memory.size={len(memory_list)}") + + threshold_score: float | None = self.context.get("threshold_score", None) + if threshold_score is not None: + memory_list = [mem for mem in memory_list if mem.score >= threshold_score] + logger.info(f"After threshold filter: {len(memory_list)} memories retained") + + self.context.response.metadata["memory_list"] = memory_list diff --git a/reme/workflow/procedural_memory/retrieve/merge_memory.py b/reme/workflow/procedural_memory/retrieve/merge_memory.py new file mode 100644 index 00000000..c54f5806 --- /dev/null +++ b/reme/workflow/procedural_memory/retrieve/merge_memory.py @@ -0,0 +1,45 @@ +"""Memory merging operation module. + +This module provides functionality to merge multiple retrieved memories +into a single formatted context string for use in LLM responses. +""" + +from typing import List + +from loguru import logger + +from ....core.op import BaseOp +from ....core.schema.memory_node import MemoryNode + + +class MergeMemory(BaseOp): + """Merge multiple memories into a single formatted context. + + This operation takes a list of retrieved memories and formats them into + a single context string that can be used to guide LLM responses. It includes + instructions for the LLM to consider the helpful parts from these memories. + """ + + async def execute(self): + """Execute the memory merging operation. + + Merges memories from context metadata into a formatted string with + instructions for the LLM. Stores the merged result in response.answer. + """ + memory_list: List[MemoryNode] = self.context.response.metadata["memory_list"] + + if not memory_list: + return + + content_collector = ["Previous Memory"] + for memory in memory_list: + if not memory.content: + continue + + content_collector.append(f"- {memory.when_to_use} {memory.content}\n") + content_collector.append( + "Please consider the helpful parts from these in answering the question, " + "to make the response more comprehensive and substantial.", + ) + self.context.response.answer = "\n".join(content_collector) + logger.info(f"response.answer={self.context.response.answer}") diff --git a/reme/workflow/procedural_memory/retrieve/rerank_memory.py b/reme/workflow/procedural_memory/retrieve/rerank_memory.py new file mode 100644 index 00000000..12d73815 --- /dev/null +++ b/reme/workflow/procedural_memory/retrieve/rerank_memory.py @@ -0,0 +1,186 @@ +"""Memory reranking operation module. + +This module provides functionality to rerank and filter retrieved memories +using LLM-based reranking and score-based filtering to select the most relevant +memories for the current task. +""" + +import json +import re +from typing import List + +from loguru import logger + +from ....core.enumeration import Role +from ....core.op import BaseOp +from ....core.schema.message import Message +from ....core.schema.memory_node import MemoryNode + + +class RerankMemory(BaseOp): + """Rerank and filter recalled experiences using LLM and score-based filtering. + + This operation takes recalled memories and applies multiple filtering and + ranking strategies to select the most relevant memories for the current task. + It supports LLM-based reranking and score-based filtering. + """ + + async def execute(self): + """Execute the memory reranking operation. + + Applies LLM-based reranking (optional) and score-based filtering (optional) + to rerank retrieved memories. Stores the reranked results + in the context response metadata. + """ + memory_list: List[MemoryNode] = self.context.response.metadata["memory_list"] + retrieval_query: str = self.context.query + enable_llm_rerank = self.context.get("enable_llm_rerank", False) + enable_score_filter = self.context.get("enable_score_filter", False) + min_score_threshold = self.context.get("min_score_threshold", 0.3) + + if not memory_list: + logger.info("No recalled memory_list to rerank") + return + + logger.info(f"Reranking {len(memory_list)} memories") + + # Step 1: LLM reranking (optional) + if enable_llm_rerank: + memory_list = await self._llm_rerank(retrieval_query, memory_list) + logger.info(f"After LLM reranking: {len(memory_list)} memories") + + # Step 2: Score-based filtering (optional) + if enable_score_filter: + memory_list = self._score_based_filter(memory_list, min_score_threshold) + logger.info(f"After score filtering: {len(memory_list)} memories") + + # Store results in context + self.context.response.metadata["memory_list"] = memory_list + + async def _llm_rerank(self, query: str, candidates: List[MemoryNode]) -> List[MemoryNode]: + """LLM-based reranking of candidate experiences. + + Args: + query: The retrieval query used to rank candidates. + candidates: List of memory candidates to rerank. + + Returns: + List of memories reranked by relevance to the query. + """ + if not candidates: + return candidates + + # Format candidates for LLM evaluation + candidates_text = self._format_candidates_for_rerank(candidates) + + prompt = self.prompt_format( + prompt_name="memory_rerank_prompt", + query=query, + candidates=candidates_text, + num_candidates=len(candidates), + ) + + response = await self.llm.chat(messages=[Message(role=Role.USER, content=prompt)]) + + # Parse reranking results + reranked_indices = self._parse_rerank_response(response.content) + + # Reorder candidates based on LLM ranking + if reranked_indices: + reranked_candidates = [] + for idx in reranked_indices: + if 0 <= idx < len(candidates): + reranked_candidates.append(candidates[idx]) + + # Add any remaining candidates that weren't explicitly ranked + ranked_indices_set = set(reranked_indices) + for i, candidate in enumerate(candidates): + if i not in ranked_indices_set: + reranked_candidates.append(candidate) + + return reranked_candidates + + return candidates + + @staticmethod + def _score_based_filter(memories: List[MemoryNode], min_score: float) -> List[MemoryNode]: + """Filter memories based on quality scores. + + Args: + memories: List of memories to filter. + min_score: Minimum combined score threshold for filtering. + + Returns: + List of memories that meet the minimum score threshold. + """ + filtered_memories = [] + + for memory in memories: + # Get confidence score from metadata + confidence = memory.metadata.get("confidence", 0.5) + validation_score = memory.score or 0.5 + + # Calculate combined score + combined_score = (confidence + validation_score) / 2 + + if combined_score >= min_score: + filtered_memories.append(memory) + else: + logger.debug(f"Filtered out memory with score {combined_score:.2f}") + + logger.info(f"Score filtering: {len(filtered_memories)}/{len(memories)} memories retained") + return filtered_memories + + @staticmethod + def _format_candidates_for_rerank(candidates: List[MemoryNode]) -> str: + """Format candidates for LLM reranking. + + Args: + candidates: List of memory candidates to format. + + Returns: + Formatted string representation of candidates for LLM evaluation. + """ + formatted_candidates = [] + + for i, candidate in enumerate(candidates): + condition = candidate.when_to_use + content = candidate.content + + candidate_text = f"Candidate {i}:\n" + candidate_text += f"Condition: {condition}\n" + candidate_text += f"Experience: {content}\n" + + formatted_candidates.append(candidate_text) + + return "\n---\n".join(formatted_candidates) + + @staticmethod + def _parse_rerank_response(response: str) -> List[int]: + """Parse LLM reranking response to extract ranked indices. + + Args: + response: The LLM response containing ranked indices. + + Returns: + List of indices representing the reranked order. + """ + try: + # Try to extract JSON format + json_pattern = r"```json\s*([\s\S]*?)\s*```" + json_blocks = re.findall(json_pattern, response) + + if json_blocks: + parsed = json.loads(json_blocks[0]) + if isinstance(parsed, dict) and "ranked_indices" in parsed: + return parsed["ranked_indices"] + elif isinstance(parsed, list): + return parsed + + # Try to extract numbers from text + numbers = re.findall(r"\b\d+\b", response) + return [int(num) for num in numbers if int(num) < 100] # Reasonable upper bound + + except Exception as e: + logger.error(f"Error parsing rerank response: {e}") + return [] diff --git a/reme/workflow/procedural_memory/retrieve/rerank_memory.yaml b/reme/workflow/procedural_memory/retrieve/rerank_memory.yaml new file mode 100644 index 00000000..e5e53337 --- /dev/null +++ b/reme/workflow/procedural_memory/retrieve/rerank_memory.yaml @@ -0,0 +1,25 @@ +memory_rerank_prompt: | + You are an expert AI analyst tasked with reranking retrieved experiences based on their relevance to a specific query. + + Your task is to analyze the candidates and rank them by relevance, considering: + ● DIRECT RELEVANCE: How directly applicable the experience is to the current query + ● SITUATION SIMILARITY: How similar the experience context is to the current situation + ● ACTIONABILITY: How actionable and specific the experience is + ● QUALITY: The overall quality and clarity of the experience + + # Current Query + {query} + + # Candidate Experiences (Total: {num_candidates}) + {candidates} + + OUTPUT FORMAT: + Provide a ranked list of candidate indices (0-based) from most relevant to least relevant: + ```json + {{ + "ranked_indices": [2, 0, 4, 1, 3], + "reasoning": "Brief explanation of ranking rationale" + }} + ``` + + Note: Include ALL candidate indices in the ranking, even if some are less relevant. \ No newline at end of file diff --git a/reme/workflow/procedural_memory/retrieve/rewrite_memory.py b/reme/workflow/procedural_memory/retrieve/rewrite_memory.py new file mode 100644 index 00000000..0190b32b --- /dev/null +++ b/reme/workflow/procedural_memory/retrieve/rewrite_memory.py @@ -0,0 +1,201 @@ +"""Memory rewriting operation module. + +This module provides functionality to rewrite and format retrieved memories +into context messages that can be used by LLMs for task completion. +""" + +import json +import re +from typing import List + +from loguru import logger + +from ....core.enumeration import Role +from ....core.op import BaseOp +from ....core.schema.message import Message +from ....core.schema.memory_node import MemoryNode + + +class RewriteMemory(BaseOp): + """Generate and rewrite context messages from reranked experiences. + + This operation takes reranked memories and formats them into context messages + that can be used by LLMs. It optionally uses LLM-based rewriting to make + the context more relevant and actionable for the current task. + """ + + async def execute(self): + """Execute the memory rewrite operation. + + Retrieves memories from context metadata, formats them, and optionally + rewrites them using LLM to make them more relevant for the current query. + Stores the rewritten context in the response answer field. + """ + memory_list: List[MemoryNode] = self.context.response.metadata["memory_list"] + query: str = self.context.query + messages: List[Message] = [Message(**x) if isinstance(x, dict) else x for x in self.context.get("messages", [])] + + if not memory_list: + logger.info("No reranked memories to rewrite") + self.context.response.answer = "" + return + + logger.info(f"Generating context from {len(memory_list)} memories") + + # Generate initial context message + rewritten_memory = await self._generate_context_message(query, messages, memory_list) + + # Store results in context + self.context.response.answer = rewritten_memory + self.context.response.metadata["memory_list"] = [memory.model_dump() for memory in memory_list] + + async def _generate_context_message(self, query: str, messages: List[Message], memories: List[MemoryNode]) -> str: + """Generate context message from retrieved memories. + + Args: + query: The current query string. + messages: List of conversation messages for context. + memories: List of retrieved memories to format. + + Returns: + Formatted context string, optionally rewritten by LLM. + """ + if not memories: + return "" + + try: + logger.info("memories") + # Format retrieved memories + formatted_memories = self._format_memories_for_context(memories) + + if self.context.get("enable_llm_rewrite", False): + context_content = await self._rewrite_context(query, formatted_memories, messages) + else: + context_content = formatted_memories + + return context_content + + except Exception as e: + logger.error(f"Error generating context message: {e}") + return self._format_memories_for_context(memories) + + async def _rewrite_context(self, query: str, context_content: str, messages: List[Message]) -> str: + """LLM-based context rewriting to make experiences more relevant and actionable. + + Args: + query: The current query string. + context_content: The formatted context content to rewrite. + messages: List of conversation messages for additional context. + + Returns: + Rewritten context string optimized for the current task. + """ + if not context_content: + return context_content + + try: + # Extract current context + current_context = self._extract_context(messages) + + prompt = self.prompt_format( + prompt_name="memory_rewrite_prompt", + current_query=query, + current_context=current_context, + original_context=context_content, + ) + + response = await self.llm.chat(messages=[Message(role=Role.USER, content=prompt)]) + + # Extract rewritten context + rewritten_context = self._parse_json_response(response.content, "rewritten_context") + + if rewritten_context and rewritten_context.strip(): + logger.info("Context successfully rewritten for current task") + return rewritten_context.strip() + + return context_content + + except Exception as e: + logger.error(f"Error in context rewriting: {e}") + return context_content + + @staticmethod + def _format_memories_for_context(memories: List[MemoryNode]) -> str: + """Format memories for context generation. + + Args: + memories: List of memories to format. + + Returns: + Formatted string containing all memories with their conditions and content. + """ + formatted_memories = [] + + for i, memory in enumerate(memories, 1): + condition = memory.when_to_use + memory_content = memory.content + memory_text = f"Experience {i} :\n When to use: {condition}\n Content: {memory_content}\n" + + formatted_memories.append(memory_text) + + return "\n".join(formatted_memories) + + @staticmethod + def _extract_context(messages: List[Message]) -> str: + """Extract relevant context from messages. + + Args: + messages: List of conversation messages. + + Returns: + Formatted string containing recent conversation context. + """ + if not messages: + return "" + + context_parts = [] + + # Add recent messages if available + recent_messages = messages[-3:] # Last 3 messages + message_summaries = [] + for message in recent_messages: + content = message.content[:300] + "..." if len(message.content) > 300 else message.content + message_summaries.append(f"- {message.role.value}: {content}") + + if message_summaries: + context_parts.append("Recent conversation:\n" + "\n".join(message_summaries)) + + return "\n\n".join(context_parts) + + @staticmethod + def _parse_json_response(response: str, key: str) -> str: + """Parse JSON response to extract specific key. + + Args: + response: The response string that may contain JSON. + key: The key to extract from the JSON object. + + Returns: + The value associated with the key, or the response string if parsing fails. + """ + try: + # Try to 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]) + if isinstance(parsed, dict) and key in parsed: + return parsed[key] + + # Fallback: try to parse the entire response as JSON + parsed = json.loads(response) + if isinstance(parsed, dict) and key in parsed: + return parsed[key] + + except json.JSONDecodeError: + logger.warning(f"Failed to parse JSON response for key '{key}', using raw response") + # If JSON parsing fails, return the response as-is for fallback + return response.strip() + + return "" diff --git a/reme/workflow/procedural_memory/retrieve/rewrite_memory.yaml b/reme/workflow/procedural_memory/retrieve/rewrite_memory.yaml new file mode 100644 index 00000000..93e899b1 --- /dev/null +++ b/reme/workflow/procedural_memory/retrieve/rewrite_memory.yaml @@ -0,0 +1,34 @@ +memory_rewrite_prompt: | + You are an expert AI assistant tasked with rewriting and reorganizing context content to make it more relevant and actionable for the current task. + + Your task is to take the original context (containing multiple experiences) and rewrite it as a cohesive, task-specific guidance that directly addresses the current situation. + + REWRITING GUIDELINES: + ● RELEVANCE FOCUS: Emphasize the most relevant aspects of each experience. Prioritize the most relevant experiences. Use clear, direct language. + ● ACTIONABLE INSIGHTS: Extract specific, actionable guidance. Make the context immediately actionable + ● COHERENT NARRATIVE: Create a flowing narrative rather than disconnected tips + ● SITUATIONAL AWARENESS: Adapt the guidance to the current situation + + # Current Task/Query + {current_query} + + # Current Trajectory + {current_context} + + # Original Context Content (Multiple Experiences) + {original_context} + + OUTPUT FORMAT: + Provide the rewritten context: + ```json + {{ + "rewritten_context": "A cohesive, task-specific context message that reorganizes and adapts the original experiences for the current task. This should be written as a unified guidance rather than separate experience items.", + }} + ``` + + Guidelines: + - Rewrite as a unified, flowing guidance + - Adapt terminology and examples to match the current task domain + - Consolidate overlapping insights into coherent recommendations + - Prioritize experiences most relevant to the current situation + - Make the guidance feel custom-written for this specific task diff --git a/reme/workflow/procedural_memory/retrieve/update_memory_metadata.py b/reme/workflow/procedural_memory/retrieve/update_memory_metadata.py new file mode 100644 index 00000000..14d7e57a --- /dev/null +++ b/reme/workflow/procedural_memory/retrieve/update_memory_metadata.py @@ -0,0 +1,51 @@ +"""Operation for updating memory metadata (frequency and utility). + +This module provides a unified operation to update frequency counters and +optionally utility scores for recalled memories, directly updating the +vector store. +""" + +from typing import List + +from loguru import logger + +from ....core.op import BaseOp +from ....core.schema.memory_node import MemoryNode +from ....core.schema.vector_node import VectorNode + + +class UpdateMemoryMetadata(BaseOp): + """Update memory metadata: frequency and optionally utility. + + This operation (1) increments each memory's frequency counter; + (2) optionally increments utility when update_utility is True; + (3) directly updates the VectorNode in the vector store using the update method. + + Expected context attributes: + memory_list: List of MemoryNode objects to update (already loaded from + previous operations like rerank_memory). + update_utility: Boolean flag. If True, also increment utility for each memory. + """ + + async def execute(self): + """Run frequency update, optional utility update, and directly update vector store.""" + memory_list: List[MemoryNode] = [MemoryNode(**node) for node in self.context.memory_list] + update_utility = self.context.update_utility + + if not memory_list: + logger.info("No memories to update metadata") + return + + updated_nodes: List[VectorNode] = [] + for memory in memory_list: + meta = memory.metadata + meta["freq"] = meta.get("freq", 0) + 1 + if update_utility: + meta["utility"] = meta.get("utility", 0) + 1 + memory.metadata = meta + vector_node = memory.to_vector_node() + updated_nodes.append(vector_node) + + if updated_nodes: + await self.vector_store.update(nodes=updated_nodes) + logger.info(f"Updated metadata for {len(updated_nodes)} memories in vector store") diff --git a/reme/workflow/procedural_memory/summarizer/__init__.py b/reme/workflow/procedural_memory/summarizer/__init__.py deleted file mode 100644 index 34644cd0..00000000 --- a/reme/workflow/procedural_memory/summarizer/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -"""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.ops.register(tool_class) diff --git a/reme/workflow/procedural_memory/summary/comparative_extraction.py b/reme/workflow/procedural_memory/summary/comparative_extraction.py new file mode 100644 index 00000000..9ca78784 --- /dev/null +++ b/reme/workflow/procedural_memory/summary/comparative_extraction.py @@ -0,0 +1,274 @@ +"""Comparative extraction operation for task memory generation. + +This module provides operations to extract comparative task memories by comparing +different trajectories with varying scores or success/failure outcomes. +""" + +from typing import List, Tuple, Optional + +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 ..utils import ( + merge_messages_content, + parse_json_experience_response, +) + + +class ComparativeExtraction(BaseOp): + """Extract comparative task memories by comparing different scoring trajectories. + + This operation performs two types of comparisons: + 1. Soft comparison: Compares highest vs lowest scoring trajectories + 2. Hard comparison: Compares similar success vs failure step sequences + + The extracted memories help identify what makes some trajectories more successful + than others. + """ + + async 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 self.context.get("enable_soft_comparison", True) and len(all_trajectories) >= 2: + 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: " + f"highest ({highest_traj.score:.2f}) vs lowest ({lowest_traj.score:.2f})", + ) + soft_task_memories = await 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 self.context.get("enable_similarity_comparison", True) and success_trajectories and failure_trajectories: + 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 = await 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 + + async def _extract_soft_comparative_task_memory( + self, + higher_traj: Trajectory, + lower_traj: Trajectory, + ) -> List[MemoryNode]: + """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[MemoryNode]: + task_memories_data = parse_json_experience_response(message.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, + ) + + async def _extract_hard_comparative_task_memory( + self, + success_steps: List[Message], + failure_steps: List[Message], + similarity_score: float, + ) -> List[MemoryNode]: + """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[MemoryNode]: + task_memories_data = parse_json_experience_response(message.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, + ) + + @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""" + 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.context.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, "vector_store") + and self.vector_store + and hasattr( + self.vector_store, + "embedding_model", + ) + ): + success_embeddings = self.vector_store.embedding_model.get_embeddings(success_texts) + failure_embeddings = self.vector_store.embedding_model.get_embeddings(failure_texts) + + # Calculate similarity and find most similar pairs + similarity_threshold = self.context.get("similarity_threshold", 0.5) + + 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.context.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) diff --git a/reme/workflow/procedural_memory/summary/comparative_extraction.yaml b/reme/workflow/procedural_memory/summary/comparative_extraction.yaml new file mode 100644 index 00000000..626cc925 --- /dev/null +++ b/reme/workflow/procedural_memory/summary/comparative_extraction.yaml @@ -0,0 +1,79 @@ +soft_comparative_step_task_memory_prompt: | + You are an expert AI analyst comparing higher-scoring and lower-scoring step sequences to extract performance insights. + + Your task is to identify the key differences between higher and lower performing approaches at the step level. + Focus on what made the higher-scoring approach more effective, even when both approaches may have had partial success. + + SOFT COMPARATIVE ANALYSIS FRAMEWORK: + ● PERFORMANCE FACTORS: Identify what specifically contributed to the higher score + ● APPROACH DIFFERENCES: Compare methodologies and execution strategies + ● EFFICIENCY ANALYSIS: Analyze why one approach was more efficient or effective + ● OPTIMIZATION INSIGHTS: Extract lessons for improving performance + + EXTRACTION PRINCIPLES: + ● Focus on INCREMENTAL IMPROVEMENTS and performance optimization + ● Extract QUALITY INDICATORS that differentiate better vs good approaches + ● Identify REFINEMENT STRATEGIES that lead to higher scores + ● Frame insights as PERFORMANCE ENHANCEMENT guidelines + + # Higher-Scoring Step Sequence (Score: {higher_score}) + {higher_steps} + + # Lower-Scoring Step Sequence (Score: {lower_score}) + {lower_steps} + + + OUTPUT FORMAT: + Generate 1-2 performance improvement insights as JSON objects: + ```json + [ + {{ + "when_to_use": "Specific scenarios where this performance insight applies", + "experience": "Detailed analysis of what made the higher-scoring approach more effective", + "tags": ["performance_optimization", "score_improvement", "relevant_keywords"], + "confidence": 0.7, + "step_type": "reasoning|action|observation|decision", + "tools_used": ["list", "of", "tools"] + }} + ] + ``` + +hard_comparative_step_task_memory_prompt: | + You are an expert AI analyst comparing successful and failed step sequences to extract differential insights. + + Your task is to identify the key differences between success and failure patterns at the step level. + Focus on critical decision points, technique variations, and approach differences. + + COMPARATIVE ANALYSIS FRAMEWORK: + ● DECISION CONTRAST: Compare critical decisions made in success vs failure cases + ● TECHNIQUE VARIATIONS: Identify different approaches and their outcomes + ● TIMING DIFFERENCES: Analyze when certain actions were taken and their impact + ● SUCCESS FACTORS: Extract what specifically made the difference + + EXTRACTION PRINCIPLES: + ● Frame comparisons as PRINCIPLES as well as case-specific SOLUTIONS + ● Identify PATTERNS that differentiate effective vs ineffective approaches + ● Extract RULES that can guide future similar situations + ● Focus on UNDERLYING MECHANISMS rather than surface-level differences + + # Successful Step Sequence + {success_steps} + + # Failed Step Sequence + {failure_steps} + + # Similarity Score: {similarity_score} + + OUTPUT FORMAT: + Generate 1-2 comparative insights as JSON objects: + ```json + [ + {{ + "when_to_use": "Specific scenarios where this comparative insight applies", + "experience": "Detailed comparison highlighting why success approach works better", + "tags": ["comparative_analysis", "success_factors", "relevant_keywords"], + "confidence": 0.8, + "step_type": "reasoning|action|observation|decision" + }} + ] + ``` \ No newline at end of file diff --git a/reme/workflow/procedural_memory/summary/failure_extraction.py b/reme/workflow/procedural_memory/summary/failure_extraction.py new file mode 100644 index 00000000..42e682d8 --- /dev/null +++ b/reme/workflow/procedural_memory/summary/failure_extraction.py @@ -0,0 +1,95 @@ +"""Failure extraction operation for task memory generation. + +This module provides operations to extract task memories from failed trajectories, +identifying mistakes, pitfalls, and lessons learned from failures. +""" + +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 ..utils import ( + get_trajectory_context, + merge_messages_content, + parse_json_experience_response, +) + + +class FailureExtraction(BaseOp): + """Extract task memories from failed trajectories. + + This operation analyzes failed trajectories (or their segments) to extract + lessons learned, common mistakes, and anti-patterns that should be avoided + in similar future tasks. + """ + + async 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 "segments" in trajectory.metadata: + # Process segmented step sequences + for segment in trajectory.metadata["segments"]: + task_memories = await self._extract_failure_task_memory_from_steps(segment, trajectory) + failure_task_memories.extend(task_memories) + else: + # Process entire trajectory + task_memories = await 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 + + async def _extract_failure_task_memory_from_steps( + self, + steps: List[Message], + trajectory: Trajectory, + ) -> List[MemoryNode]: + """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[MemoryNode]: + task_memories_data = parse_json_experience_response(message.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/summary/failure_extraction.yaml b/reme/workflow/procedural_memory/summary/failure_extraction.yaml new file mode 100644 index 00000000..be660ccf --- /dev/null +++ b/reme/workflow/procedural_memory/summary/failure_extraction.yaml @@ -0,0 +1,42 @@ +failure_step_task_memory_prompt: | + You are an expert AI analyst reviewing failed step sequences from an AI agent execution. + + Your task is to extract learning task memories from failures to prevent similar mistakes in future executions. + Focus on identifying error patterns, missed opportunities, and alternative approaches. + + ANALYSIS FRAMEWORK: + ● FAILURE POINT IDENTIFICATION: Pinpoint where and why the steps went wrong + ● ERROR PATTERN ANALYSIS: Identify recurring mistakes or problematic approaches + ● ALTERNATIVE APPROACHES: Suggest what could have been done differently + ● PREVENTION STRATEGIES: Extract actionable insights to avoid similar failures + + EXTRACTION PRINCIPLES: + ● Extract GENERAL PRINCIPLES as well as SPECIFIC INSTRUCTIONS + ● Focus on PATTERNS and RULES as well as particular instances + + # 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 failure prevention insights as JSON objects: + ```json + [ + {{ + "when_to_use": "Specific situations where this lesson should be remembered", + "experience": "Universal principle or rule extracted from the failure pattern ", + "tags": ["error_prevention", "failure_analysis", "relevant_keywords"], + "confidence": 0.7, + "step_type": "reasoning|action|observation|decision", + "tools_used": ["list", "of", "tools"] + }} + ] + ``` \ No newline at end of file diff --git a/reme/workflow/procedural_memory/summary/memory_addition.py b/reme/workflow/procedural_memory/summary/memory_addition.py new file mode 100644 index 00000000..17ece4f1 --- /dev/null +++ b/reme/workflow/procedural_memory/summary/memory_addition.py @@ -0,0 +1,30 @@ +"""Operation for adding memories to the vector store.""" + +from typing import List +from loguru import logger +from ....core.op import BaseOp +from ....core.schema.vector_node import VectorNode +from ....core.schema.memory_node import MemoryNode + + +class MemoryAddition(BaseOp): + """Operation that adds new or updated memories to the vector store. + + This operation inserts memories into the vector store. It reads the list + of memories to insert from response.metadata and performs the actual + database insertion operations. + """ + + async def execute(self): + """Execute the memory insertion operation. + + Inserts new or updated memories into the vector store: + 1. Reads memory_list from response.metadata + 2. Converts MemoryNode objects to VectorNode objects + 3. Inserts them into the vector store + """ + insert_memory_list: List[MemoryNode] = self.context.memory_list + if insert_memory_list: + insert_nodes: List[VectorNode] = [x.to_vector_node() for x in insert_memory_list] + await self.vector_store.insert(nodes=insert_nodes) + logger.info(f"insert insert_node.size={len(insert_nodes)}") diff --git a/reme/workflow/procedural_memory/summary/memory_deduplication.py b/reme/workflow/procedural_memory/summary/memory_deduplication.py new file mode 100644 index 00000000..66222461 --- /dev/null +++ b/reme/workflow/procedural_memory/summary/memory_deduplication.py @@ -0,0 +1,179 @@ +"""Memory deduplication operation for task memory management. + +This module provides operations to remove duplicate or highly similar task +memories by comparing embeddings and calculating similarity scores. +""" + +from typing import List + +from loguru import logger + +from ....core.op import BaseOp +from ....core.schema.memory_node import MemoryNode + + +class MemoryDeduplication(BaseOp): + """Remove duplicate task memories using embedding similarity. + + This operation identifies and removes duplicate or highly similar task + memories by comparing their embeddings against both existing memories + in the vector store and other memories in the current batch. + """ + + async def execute(self): + """Remove duplicate task memories""" + # Get task memories to deduplicate + task_memories: List[MemoryNode] = self.context.response.metadata.get("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 = await self._deduplicate_task_memories(task_memories) + + logger.info( + f"Deduplication complete: {len(deduplicated_task_memories)} deduplicated " + f"task memories out of {len(task_memories)}", + ) + + # Update context + self.context.response.metadata["memory_list"] = deduplicated_task_memories + + async def _deduplicate_task_memories(self, task_memories: List[MemoryNode]) -> List[MemoryNode]: + """Remove duplicate task memories""" + if not task_memories: + return task_memories + + similarity_threshold = self.context.get("similarity_threshold", 0.5) + + unique_task_memories = [] + + # Get existing task memory embeddings + existing_embeddings = await self._get_existing_task_memory_embeddings() + + 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 + + async def _get_existing_task_memory_embeddings(self) -> List[List[float]]: + """Get embeddings of existing task memories""" + try: + if not hasattr(self, "vector_store") or not self.vector_store: + return [] + + # List all existing task memory nodes + existing_nodes = await self.vector_store.list( + filters=None, # No filters to get all nodes + limit=self.context.get("max_existing_task_memories", 1000), + ) + + # Extract embeddings + existing_embeddings = [] + for node in existing_nodes: + if node.vector: + existing_embeddings.append(node.vector) + + logger.debug( + f"Retrieved {len(existing_embeddings)} existing task memory embeddings", + ) + 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: MemoryNode) -> List[float] | None: + """Generate embedding for task memory""" + try: + + # Combine task memory description and content for embedding + text_for_embedding = f"{task_memory.when_to_use} {task_memory.content}" + embeddings = self.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[MemoryNode], + threshold: float, + ) -> bool: + """Check if current embedding is similar to other memories in current batch.""" + 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 diff --git a/reme/workflow/procedural_memory/summary/memory_validation.py b/reme/workflow/procedural_memory/summary/memory_validation.py new file mode 100644 index 00000000..07def492 --- /dev/null +++ b/reme/workflow/procedural_memory/summary/memory_validation.py @@ -0,0 +1,127 @@ +"""Memory validation operation for task memory quality control. + +This module provides operations to validate the quality of extracted task +memories using LLM-based evaluation, ensuring only high-quality memories +are stored. +""" + +import json +import re +from typing import List, Dict, Any + +from loguru import logger + +from ....core.enumeration import Role +from ....core.op import BaseOp +from ....core.schema.message import Message +from ....core.schema.memory_node import MemoryNode + + +class MemoryValidation(BaseOp): + """Validate quality of extracted task memories. + + This operation uses LLM-based evaluation to assess the quality of extracted + task memories, filtering out low-quality or invalid memories based on + validation scores and criteria. + """ + + async def execute(self): + """Validate quality of extracted task memories""" + + task_memories: List[MemoryNode] = [] + 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 = await 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 + + async def _validate_single_task_memory(self, task_memory: MemoryNode) -> Dict[str, Any]: + """Validate single task memory""" + validation_info = await self._llm_validate_task_memory(task_memory) + logger.info(f"Validating: {validation_info}") + return validation_info + + async def _llm_validate_task_memory(self, task_memory: MemoryNode) -> 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.context.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 await self.llm.chat( + messages=[Message(role=Role.USER, 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)}", + } diff --git a/reme/workflow/procedural_memory/summary/memory_validation.yaml b/reme/workflow/procedural_memory/summary/memory_validation.yaml new file mode 100644 index 00000000..44287e28 --- /dev/null +++ b/reme/workflow/procedural_memory/summary/memory_validation.yaml @@ -0,0 +1,29 @@ +task_memory_validation_prompt: | + You are an expert AI analyst tasked with validating the quality and usefulness of extracted step-level task memories. + + Your task is to access whether the extracted task memory is actionable, accurate, and valuable for future agent executions. + + VALIDATION CRITERIA: + ● ACTIONABILITY: Is the task memory specific enough to guide future actions? + ● ACCURACY: Does the task memory correctly reflect the patterns observed? + ● RELEVANCE: Is the task memory applicable to similar future scenarios? + ● CLARITY: Is the task memory clearly articulated and understandable? + ● UNIQUENESS: Does the task memory provide novel insights or common knowledge? + + # Task Memory to Validate + Condition: {condition} + Task Memory Content: {task_memory_content} + + OUTPUT FORMAT: + Provide validation assessment: + ```json + {{ + "is_valid": true/false, + "score": 0.8, + "feedback": "Detailed explanation of validation decision", + "recommendations": "Suggestions for improvement if applicable" + }} + ``` + + Score should be between 0.0 (poor quality) and 1.0 (excellent quality). + Mark as invalid if score is below 0.3 or if there are fundamental issues with the task memory. diff --git a/reme/workflow/procedural_memory/summarizer/success_extraction.py b/reme/workflow/procedural_memory/summary/success_extraction.py similarity index 98% rename from reme/workflow/procedural_memory/summarizer/success_extraction.py rename to reme/workflow/procedural_memory/summary/success_extraction.py index f4b56c90..4e45459f 100644 --- a/reme/workflow/procedural_memory/summarizer/success_extraction.py +++ b/reme/workflow/procedural_memory/summary/success_extraction.py @@ -12,7 +12,7 @@ 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 ( +from ..utils import ( get_trajectory_context, merge_messages_content, parse_json_experience_response, diff --git a/reme/workflow/procedural_memory/summarizer/success_extraction.yaml b/reme/workflow/procedural_memory/summary/success_extraction.yaml similarity index 100% rename from reme/workflow/procedural_memory/summarizer/success_extraction.yaml rename to reme/workflow/procedural_memory/summary/success_extraction.yaml diff --git a/reme/workflow/procedural_memory/summarizer/trajectory_preprocess.py b/reme/workflow/procedural_memory/summary/trajectory_preprocess.py similarity index 100% rename from reme/workflow/procedural_memory/summarizer/trajectory_preprocess.py rename to reme/workflow/procedural_memory/summary/trajectory_preprocess.py diff --git a/reme/workflow/procedural_memory/summary/trajectory_segmentation.py b/reme/workflow/procedural_memory/summary/trajectory_segmentation.py new file mode 100644 index 00000000..9673ab76 --- /dev/null +++ b/reme/workflow/procedural_memory/summary/trajectory_segmentation.py @@ -0,0 +1,139 @@ +"""Trajectory segmentation operation for task memory generation. + +This module provides operations to segment trajectories into meaningful step +sequences that can be used for more granular memory extraction. +""" + +import json +import re +from typing import List + +from loguru import logger + +from ....core.enumeration import Role +from ....core.op import BaseOp +from ....core.schema.message import Message, Trajectory + + +class TrajectorySegmentation(BaseOp): + """Segment trajectories into meaningful step sequences. + + This operation uses LLM to identify natural breakpoints in trajectories, + allowing for more granular analysis and memory extraction from specific + segments rather than entire trajectories. + """ + + async 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 = await 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.context.get("segment_target", "all") + + if segment_target == "success": + return success_trajectories + elif segment_target == "failure": + return failure_trajectories + else: + return all_trajectories + + async 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 await self.llm.chat( + messages=[Message(role=Role.USER, 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))) diff --git a/reme/workflow/procedural_memory/summary/trajectory_segmentation.yaml b/reme/workflow/procedural_memory/summary/trajectory_segmentation.yaml new file mode 100644 index 00000000..8a22f3fe --- /dev/null +++ b/reme/workflow/procedural_memory/summary/trajectory_segmentation.yaml @@ -0,0 +1,31 @@ +step_segmentation_prompt: | + You are an expert AI analyst tasked with segmenting a trajectory into meaningful step sequences. + + Your task is to identify natural breakpoints in the execution where one logical unit of work ends and another begins. + Consider factors like: task completion, context switches, tool changes, reasoning phases, and logical groupings. + + SEGMENTATION CRITERIA: + ● LOGICAL COMPLETION: Steps that complete a specific sub-task or reasoning phase + ● CONTEXT SWITCHES: Points where the agent shifts focus or approach + ● TOOL BOUNDARIES: Natural breaks around tool usage patterns + ● REASONING PHASES: Distinct phases of analysis, planning, or execution + + # Original Query + {query} + + # Full Trajectory (Total steps: {total_steps}) + {trajectory_content} + + OUTPUT FORMAT: + Provide segmentation points as a JSON array of step indices where splits should occur: + ```json + {{ + "segment_points": [3, 7, 12, 18], + "reasoning": "Brief explanation of segmentation logic" + }} + ``` + + Note: Segment points indicate the END of each segment. For example, [3, 7] means: + - Segment 1: steps 0-3 + - Segment 2: steps 4-7 + - Segment 3: steps 8-end \ No newline at end of file diff --git a/reme/workflow/procedural_memory/utils.py b/reme/workflow/procedural_memory/utils.py new file mode 100644 index 00000000..11c99fa0 --- /dev/null +++ b/reme/workflow/procedural_memory/utils.py @@ -0,0 +1,142 @@ +"""Utility functions for processing and formatting LLM-related message data.""" + +import json +import re +from loguru import logger + +from ...core.enumeration import Role +from ...core.schema.message import Message, Trajectory + +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')}" +