iflow-mcp_hanw39_reasoning-bank-mcp 0.2.0__py3-none-any.whl

src/tools/extract_memory.py

@@ -0,0 +1,285 @@
+"""extract_memory 工具 - 提取记忆（支持异步）"""
+import json
+import uuid
+import asyncio
+from typing import Dict, List, Optional
+from datetime import datetime, timezone
+
+
+class ExtractMemoryTool:
+    """记忆提取工具"""
+
+    def __init__(self, config, storage_backend, llm_provider, embedding_provider, memory_manager=None):
+        """
+        初始化提取工具
+
+        Args:
+            config: 配置对象
+            storage_backend: 存储后端实例
+            llm_provider: LLM Provider 实例
+            embedding_provider: 嵌入 Provider 实例
+            memory_manager: 记忆管理器实例（可选，用于去重和合并）
+        """
+        self.config = config
+        self.storage = storage_backend
+        self.llm = llm_provider
+        self.embedding = embedding_provider
+        self.memory_manager = memory_manager
+
+        # 提取配置
+        extraction_config = config.get("extraction", default={})
+        self.max_memories = extraction_config.get("max_memories_per_trajectory", 3)
+        self.judge_temp = extraction_config.get("judge_temperature", 0.0)
+        self.extract_temp = extraction_config.get("extract_temperature", 1.0)
+        self.async_by_default = extraction_config.get("async_by_default", True)
+
+    async def execute(
+        self,
+        trajectory: List[Dict],
+        query: str,
+        success_signal: Optional[bool] = None,
+        async_mode: bool = None,
+        agent_id: str = None
+    ) -> Dict:
+        """
+        执行记忆提取
+
+        Args:
+            trajectory: 轨迹步骤列表
+            query: 任务查询
+            success_signal: 成功/失败标记，None 时自动判断
+            async_mode: 是否异步处理，None 时使用配置默认值
+            agent_id: Agent ID，用于多租户隔离
+
+        Returns:
+            提取结果字典
+        """
+        # 确定是否异步
+        if async_mode is None:
+            async_mode = self.async_by_default
+
+        # 生成任务 ID
+        task_id = f"extract_{uuid.uuid4().hex[:8]}"
+
+        if async_mode:
+            # 异步模式：立即返回，后台处理
+            asyncio.create_task(
+                self._extract_async(task_id, trajectory, query, success_signal, agent_id)
+            )
+            return {
+                "status": "processing",
+                "message": "记忆提取任务已提交，正在后台处理",
+                "task_id": task_id,
+                "async_mode": True
+            }
+        else:
+            # 同步模式：等待处理完成
+            result = await self._extract_sync(trajectory, query, success_signal, agent_id)
+            return {
+                **result,
+                "task_id": task_id,
+                "async_mode": False
+            }
+
+    async def _extract_sync(
+        self,
+        trajectory: List[Dict],
+        query: str,
+        success_signal: Optional[bool],
+        agent_id: str = None
+    ) -> Dict:
+        """同步提取记忆"""
+        try:
+            # 1. 判断成功/失败（如果未提供）
+            if success_signal is None:
+                success_signal = await self._judge_trajectory(trajectory, query)
+
+            # 2. 格式化轨迹
+            from ..prompts.formatters import format_trajectory
+            trajectory_text = format_trajectory(trajectory)
+
+            # 3. 提取记忆项
+            from ..prompts.templates import get_extract_prompt
+            extract_prompt = get_extract_prompt(query, trajectory_text, success_signal)
+
+            response = await self.llm.chat(
+                messages=[{"role": "user", "content": extract_prompt}],
+                temperature=self.extract_temp
+            )
+
+            # 4. 解析 LLM 响应
+            memories = self._parse_llm_response(response)
+
+            # 限制数量
+            memories = memories[:self.max_memories]
+
+            # 5. 构建记忆项和嵌入
+            new_memories = []
+            embeddings_dict = {}
+            current_time = datetime.now(timezone.utc).isoformat()
+
+            for mem_data in memories:
+                memory_id = f"mem_{uuid.uuid4().hex}"
+
+                # 构建完整记忆项
+                memory = {
+                    "memory_id": memory_id,
+                    "agent_id": agent_id,
+                    "timestamp": current_time,
+                    "success": success_signal,
+                    "title": mem_data["title"],
+                    "description": mem_data["description"],
+                    "content": mem_data["content"],
+                    "query": query,
+                    "retrieval_count": 0,
+                    "last_retrieved": None,
+                    "tags": self._extract_tags(mem_data, query),
+                    "metadata": {
+                        "extraction_model": self.llm.get_provider_name(),
+                        "embedding_model": self.embedding.get_provider_name()
+                    }
+                }
+
+                # 计算嵌入：对记忆内容进行向量化，而不是查询
+                memory_text = f"{mem_data['title']}  {mem_data['description']} {mem_data['content']}"
+                embedding = await self.embedding.embed(memory_text)
+
+                new_memories.append(memory)
+                embeddings_dict[memory_id] = embedding
+
+            # 6. 通过 MemoryManager 处理（去重和合并检测）
+            memories_to_save = new_memories
+            management_result = None
+
+            if self.memory_manager:
+                try:
+                    management_result = await self.memory_manager.on_memory_created(
+                        new_memories=new_memories,
+                        embeddings=embeddings_dict,
+                        agent_id=agent_id
+                    )
+
+                    # 使用去重后的记忆列表
+                    if management_result.success:
+                        memories_to_save = management_result.metadata.get("unique_memories", new_memories)
+
+                except Exception as e:
+                    # MemoryManager 失败不影响主流程
+                    import logging
+                    logging.warning(f"MemoryManager 处理失败: {e}", exc_info=True)
+
+            # 7. 保存记忆到存储
+            saved_memories = []
+            for memory in memories_to_save:
+                memory_id = memory["memory_id"]
+                embedding = embeddings_dict[memory_id]
+
+                await self.storage.add_memory(memory, embedding)
+
+                saved_memories.append({
+                    "memory_id": memory_id,
+                    "title": memory["title"],
+                    "description": memory["description"]
+                })
+
+            # 8. 构建返回结果
+            result = {
+                "status": "completed",
+                "success": success_signal,
+                "extracted_count": len(saved_memories),
+                "memories": saved_memories
+            }
+
+            # 添加管理信息
+            if management_result:
+                result["management"] = {
+                    "duplicates_skipped": management_result.duplicates_found,
+                    "merges_triggered": management_result.merged_count,
+                    "message": management_result.message
+                }
+
+            return result
+
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": f"提取失败: {str(e)}",
+                "success": None,
+                "extracted_count": 0,
+                "memories": []
+            }
+
+    async def _extract_async(
+        self,
+        task_id: str,
+        trajectory: List[Dict],
+        query: str,
+        success_signal: Optional[bool],
+        agent_id: str = None
+    ):
+        """异步提取记忆（后台任务）"""
+        # 直接调用同步提取逻辑
+        await self._extract_sync(trajectory, query, success_signal, agent_id)
+        # 注意：异步模式下，结果不返回给调用者，只记录到日志或存储
+
+    async def _judge_trajectory(self, trajectory: List[Dict], query: str) -> bool:
+        """判断轨迹是否成功"""
+        try:
+            from ..prompts.formatters import format_trajectory
+            from ..prompts.templates import get_judge_prompt
+
+            trajectory_text = format_trajectory(trajectory)
+            # todo 轨迹分段，A task may involve success and failure..
+            judge_prompt = get_judge_prompt(query, trajectory_text)
+
+            response = await self.llm.chat(
+                messages=[{"role": "user", "content": judge_prompt}],
+                temperature=self.judge_temp
+            )
+
+            # 解析判断结果
+            result = self._parse_json_response(response)
+            return result.get("result", "failure") == "success"
+
+        except Exception:
+            # 判断失败时，默认为失败轨迹
+            return False
+
+    def _parse_llm_response(self, response: str) -> List[Dict]:
+        """解析 LLM 返回的记忆项"""
+        try:
+            # 尝试提取 JSON
+            data = self._parse_json_response(response)
+
+            if "memories" in data:
+                return data["memories"]
+            else:
+                return []
+
+        except Exception:
+            return []
+
+    def _parse_json_response(self, response: str) -> Dict:
+        """从响应中提取 JSON"""
+        # 移除可能的 markdown 代码块标记
+        response = response.strip()
+        if response.startswith("```json"):
+            response = response[7:]
+        if response.startswith("```"):
+            response = response[3:]
+        if response.endswith("```"):
+            response = response[:-3]
+
+        response = response.strip()
+        return json.loads(response)
+
+    def _extract_tags(self, memory_data: Dict, query: str) -> List[str]:
+        """从记忆内容中提取标签"""
+        # 简单的标签提取逻辑
+        tags = []
+
+        # 基于成功/失败
+        # （在调用处已经有 success 信息）
+
+        # 可以添加更复杂的标签提取逻辑
+        return tags

src/tools/retrieve_memory.py

@@ -0,0 +1,139 @@
+"""retrieve_memory 工具 - 检索相关记忆"""
+from typing import Dict, List
+from datetime import datetime, timezone
+import numpy as np
+
+
+class RetrieveMemoryTool:
+    """检索记忆工具"""
+
+    def __init__(self, config, storage_backend, embedding_provider, retrieval_strategy):
+        """
+        初始化检索工具
+
+        Args:
+            config: 配置对象
+            storage_backend: 存储后端实例
+            embedding_provider: 嵌入 Provider 实例
+            retrieval_strategy: 检索策略实例
+        """
+        self.config = config
+        self.storage = storage_backend
+        self.embedding = embedding_provider
+        self.retrieval = retrieval_strategy
+
+    async def execute(self, query: str, top_k: int = None, agent_id: str = None) -> Dict:
+        """
+        执行记忆检索
+
+        Args:
+            query: 任务查询
+            top_k: 检索数量，默认使用配置中的值
+            agent_id: Agent ID，用于多租户隔离
+
+        Returns:
+            检索结果字典
+        """
+        # 使用配置的默认值
+        if top_k is None:
+            top_k = self.config.get("retrieval", "default_top_k", default=1)
+
+        # 限制最大值
+        max_top_k = self.config.get("retrieval", "max_top_k", default=10)
+        top_k = min(top_k, max_top_k)
+
+        # 获取最小分数阈值
+        min_score_threshold = self.config.get("retrieval", "min_score_threshold", default=0.85)
+
+        try:
+            # 1. 对查询进行嵌入
+            query_embedding = await self.embedding.embed(query)
+            query_vec = np.array(query_embedding)
+
+            # 2. 使用策略检索
+            top_k_results = await self.retrieval.retrieve(
+                query=query,
+                query_embedding=query_vec,
+                storage_backend=self.storage,
+                top_k=top_k,
+                agent_id=agent_id
+            )
+
+            if not top_k_results:
+                return {
+                    "status": "no_memories",
+                    "message": "记忆库为空或没有找到相关记忆",
+                    "query": query,
+                    "memories": [],
+                    "formatted_prompt": ""
+                }
+
+            # 3. 获取完整记忆内容并更新统计（过滤低分记忆）
+            retrieved_memories = []
+            current_time = datetime.now(timezone.utc).isoformat()
+            filtered_count = 0  # 记录被过滤的数量
+
+            for memory_id, score in top_k_results:
+                # 过滤低于阈值的记忆
+                if score < min_score_threshold:
+                    filtered_count += 1
+                    continue
+
+                memory = await self.storage.get_memory_by_id(memory_id)
+                if memory:
+                    # 更新检索统计
+                    await self.storage.update_retrieval_stats(memory_id, current_time)
+
+                    # 添加到结果
+                    retrieved_memories.append({
+                        "memory_id": memory_id,
+                        "score": float(score),
+                        "title": memory["title"],
+                        "content": memory["content"],
+                        "success": memory.get("success", True),
+                        "tags": memory.get("tags", []),
+                        "description": memory.get("description", "")
+                    })
+
+            # 如果所有记忆都被过滤了
+            if not retrieved_memories:
+                return {
+                    "status": "no_memories",
+                    "message": f"没有找到相关度高于 {min_score_threshold} 的记忆（过滤了 {filtered_count} 条低相关度记忆）",
+                    "query": query,
+                    "min_score_threshold": min_score_threshold,
+                    "filtered_count": filtered_count,
+                    "memories": [],
+                    "formatted_prompt": ""
+                }
+
+            # 4. 格式化为 LLM 提示
+            formatted_prompt = self._format_for_prompt(retrieved_memories)
+
+            return {
+                "status": "success",
+                "query": query,
+                "retrieval_strategy": self.retrieval.get_name(),
+                "top_k": top_k,
+                "min_score_threshold": min_score_threshold,
+                "filtered_count": filtered_count,
+                "memories": retrieved_memories,
+                "formatted_prompt": formatted_prompt
+            }
+
+        except Exception as e:
+            return {
+                "status": "error",
+                "message": f"检索失败: {str(e)}",
+                "query": query,
+                "memories": [],
+                "formatted_prompt": ""
+            }
+
+    def _format_for_prompt(self, memories: List[Dict]) -> str:
+        """格式化为可直接用于 LLM 提示的文本"""
+        if not memories:
+            return ""
+
+        from ..prompts.formatters import format_memory_for_prompt
+        return format_memory_for_prompt(memories)

src/utils/__init__.py

@@ -0,0 +1,7 @@
+"""工具函数包"""
+from .similarity import cosine_similarity, find_top_k_similar
+
+__all__ = [
+    "cosine_similarity",
+    "find_top_k_similar",
+]

src/utils/similarity.py

@@ -0,0 +1,54 @@
+"""工具函数 - 相似度计算"""
+import numpy as np
+from typing import List, Tuple, Dict
+
+
+def cosine_similarity(vec1: np.ndarray, vec2: np.ndarray) -> float:
+    """
+    计算两个向量的余弦相似度
+
+    Args:
+        vec1: 向量1
+        vec2: 向量2
+
+    Returns:
+        余弦相似度值 [-1, 1]
+    """
+    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 float(dot_product / (norm1 * norm2))
+
+
+def find_top_k_similar(
+    query_embedding: List[float],
+    memory_embeddings: Dict[str, np.ndarray],
+    top_k: int = 1
+) -> List[Tuple[str, float]]:
+    """
+    找到与查询最相似的 Top-K 记忆
+
+    Args:
+        query_embedding: 查询的嵌入向量
+        memory_embeddings: 记忆嵌入字典 {memory_id: embedding}
+        top_k: 返回的数量
+
+    Returns:
+        [(memory_id, similarity_score), ...] 按相似度降序排列
+    """
+    query_vec = np.array(query_embedding)
+
+    similarities = []
+    for memory_id, memory_vec in memory_embeddings.items():
+        sim = cosine_similarity(query_vec, memory_vec)
+        similarities.append((memory_id, sim))
+
+    # 按相似度降序排序
+    similarities.sort(key=lambda x: x[1], reverse=True)
+
+    # 返回 Top-K
+    return similarities[:top_k]