jarvis-ai-assistant 0.7.8__py3-none-any.whl → 1.0.2__py3-none-any.whl

jarvis/jarvis_mcp/streamable_mcp_client.py

@@ -1,12 +1,17 @@
 # -*- coding: utf-8 -*-
 import json
 import threading
-from typing import Any, Callable, Dict, List
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import List
+from typing import cast
 from urllib.parse import urljoin
 
-import requests  # type: ignore[import-untyped]
+import requests
 
 from jarvis.jarvis_mcp import McpClient
+from jarvis.jarvis_utils.output import PrettyOutput
 
 
 class StreamableMcpClient(McpClient):
@@ -48,12 +53,12 @@ class StreamableMcpClient(McpClient):
         self.timeout = config.get("timeout", (10, 300))
 
         # 请求相关属性
-        self.pending_requests: Dict[str, threading.Event] = (
-            {}
-        )  # 存储等待响应的请求 {id: Event}
-        self.request_results: Dict[str, Dict[str, Any]] = (
-            {}
-        )  # 存储请求结果 {id: result}
+        self.pending_requests: Dict[
+            str, threading.Event
+        ] = {}  # 存储等待响应的请求 {id: Event}
+        self.request_results: Dict[
+            str, Dict[str, Any]
+        ] = {}  # 存储请求结果 {id: result}
         self.notification_handlers: Dict[str, List[Callable]] = {}
         self.event_lock = threading.Lock()
         self.request_id_counter = 0
@@ -85,7 +90,7 @@ class StreamableMcpClient(McpClient):
             self._send_notification("notifications/initialized", {})
 
         except Exception as e:
-            print(f"❌ MCP初始化失败: {str(e)}")
+            PrettyOutput.auto_print(f"❌ MCP初始化失败: {str(e)}")
             raise
 
     def register_notification_handler(self, method: str, handler: Callable) -> None:
@@ -176,30 +181,34 @@ class StreamableMcpClient(McpClient):
                             notify_method = data.get("method", "")
                             params = data.get("params", {})
                             if notify_method in self.notification_handlers:
-                                for handler in self.notification_handlers[notify_method]:
+                                for handler in self.notification_handlers[
+                                    notify_method
+                                ]:
                                     try:
                                         handler(params)
                                     except Exception as e:
-                                        error_lines.append(f"处理通知时出错 ({notify_method}): {e}")
+                                        error_lines.append(
+                                            f"处理通知时出错 ({notify_method}): {e}"
+                                        )
                     except Exception:
                         warning_lines.append(f"无法解析响应: {line}")
                         continue
 
             if warning_lines:
-                joined_warnings = '\n'.join(warning_lines)
-                print(f"⚠️ {joined_warnings}")
+                joined_warnings = "\n".join(warning_lines)
+                PrettyOutput.auto_print(f"⚠️ {joined_warnings}")
             if error_lines:
-                joined_errors = '\n'.join(error_lines)
-                print(f"❌ {joined_errors}")
+                joined_errors = "\n".join(error_lines)
+                PrettyOutput.auto_print(f"❌ {joined_errors}")
             # Ensure response is closed after streaming
             response.close()
             if result is None:
                 raise RuntimeError(f"未收到响应: {method}")
 
-            return result
+            return cast(Dict[str, Any], result)
 
         except Exception as e:
-            print(f"❌ 发送请求失败: {str(e)}")
+            PrettyOutput.auto_print(f"❌ 发送请求失败: {str(e)}")
             raise
         finally:
             # 清理请求状态
@@ -227,7 +236,7 @@ class StreamableMcpClient(McpClient):
             response.close()
 
         except Exception as e:
-            print(f"❌ 发送通知失败: {str(e)}")
+            PrettyOutput.auto_print(f"❌ 发送通知失败: {str(e)}")
             raise
 
     def get_tool_list(self) -> List[Dict[str, Any]]:
@@ -270,10 +279,10 @@ class StreamableMcpClient(McpClient):
                 else:
                     error_msg += ": 未知错误"
 
-                print(f"❌ {error_msg}")
+                PrettyOutput.auto_print(f"❌ {error_msg}")
                 return []
         except Exception as e:
-            print(f"❌ 获取工具列表失败: {str(e)}")
+            PrettyOutput.auto_print(f"❌ 获取工具列表失败: {str(e)}")
             return []
 
     def execute(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
@@ -312,7 +321,7 @@ class StreamableMcpClient(McpClient):
                     "stderr": response.get("error", "Unknown error"),
                 }
         except Exception as e:
-            print(f"❌ 执行工具失败: {str(e)}")
+            PrettyOutput.auto_print(f"❌ 执行工具失败: {str(e)}")
             return {"success": False, "stdout": "", "stderr": str(e)}
 
     def get_resource_list(self) -> List[Dict[str, Any]]:
@@ -328,17 +337,17 @@ class StreamableMcpClient(McpClient):
         try:
             response = self._send_request("resources/list", {})
             if "result" in response and "resources" in response["result"]:
-                return response["result"]["resources"]
+                return cast(List[Dict[str, Any]], response["result"]["resources"])
             else:
                 error_msg = "获取资源列表失败"
                 if "error" in response:
                     error_msg += f": {response['error']}"
                 else:
                     error_msg += ": 未知错误"
-                print(f"❌ {error_msg}")
+                PrettyOutput.auto_print(f"❌ {error_msg}")
                 return []
         except Exception as e:
-            print(f"❌ 获取资源列表失败: {str(e)}")
+            PrettyOutput.auto_print(f"❌ 获取资源列表失败: {str(e)}")
             return []
 
     def get_resource(self, uri: str) -> Dict[str, Any]:
@@ -379,11 +388,11 @@ class StreamableMcpClient(McpClient):
                     error_msg += f": {response['error']}"
                 else:
                     error_msg += ": 未知错误"
-                print(f"❌ {error_msg}")
+                PrettyOutput.auto_print(f"❌ {error_msg}")
                 return {"success": False, "stdout": "", "stderr": error_msg}
         except Exception as e:
             error_msg = f"获取资源内容失败: {str(e)}"
-            print(f"❌ {error_msg}")
+            PrettyOutput.auto_print(f"❌ {error_msg}")
             return {"success": False, "stdout": "", "stderr": error_msg}
 
     def __del__(self):

jarvis/jarvis_memory_organizer/memory_organizer.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
 """
 记忆整理工具 - 用于合并具有相似标签的记忆
 
@@ -9,16 +7,23 @@
 import json
 from collections import defaultdict
 from pathlib import Path
-from typing import Dict, List, Set, Any, Optional
+
+from jarvis.jarvis_utils.output import PrettyOutput
+
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Set
 
 import typer
 
-from jarvis.jarvis_utils.config import (
-    get_data_dir,
-    get_normal_platform_name,
-    get_normal_model_name,
-)
 from jarvis.jarvis_platform.registry import PlatformRegistry
+from jarvis.jarvis_utils.config import get_data_dir
+from jarvis.jarvis_utils.config import get_normal_model_name
+from jarvis.jarvis_utils.config import get_normal_platform_name
 from jarvis.jarvis_utils.utils import init_env
 
 
@@ -74,8 +79,8 @@ class MemoryOrganizer:
                 error_lines.append(f"读取记忆文件 {memory_file} 失败: {str(e)}")
 
         if error_lines:
-            joined_errors = '\n'.join(error_lines)
-            print(f"⚠️ {joined_errors}")
+            joined_errors = "\n".join(error_lines)
+            PrettyOutput.auto_print(f"⚠️ {joined_errors}")
 
         return memories
 
@@ -161,11 +166,11 @@ class MemoryOrganizer:
         prompt = f"""请将以下{len(memories)}个相关记忆合并成一个综合性的记忆。
 
 原始记忆（按时间从新到旧排序）：
-{"="*50}
+{"=" * 50}
 {memory_contents_str}
-{"="*50}
+{"=" * 50}
 
-原始标签集合：{', '.join(sorted(all_tags))}
+原始标签集合：{", ".join(sorted(all_tags))}
 
 请完成以下任务：
 1. 分析这些记忆的共同主题和关键信息
@@ -193,14 +198,18 @@ class MemoryOrganizer:
 """
 
         try:
+            if self.platform is None:
+                raise ValueError("Platform is not initialized")
+
             # 调用大模型 - 收集完整响应
             response_parts = []
-            for chunk in self.platform.chat(prompt):  # type: ignore
+            for chunk in self.platform.chat(prompt):
                 response_parts.append(chunk)
             response = "".join(response_parts)
 
             # 解析响应
             import re
+
             from jarvis.jarvis_utils.jsonnet_compat import loads as json5_loads
 
             # 提取 <merged_memory> 标签内的内容
@@ -226,7 +235,7 @@ class MemoryOrganizer:
                 raise ValueError("无法从模型响应中提取 <merged_memory> 标签内容")
 
         except Exception as e:
-            print(f"⚠️ 调用大模型合并记忆失败: {str(e)}")
+            PrettyOutput.auto_print(f"⚠️ 调用大模型合并记忆失败: {str(e)}")
             # 返回 None 表示合并失败，跳过这组记忆
             return None
 
@@ -247,15 +256,17 @@ class MemoryOrganizer:
         返回：
             整理结果统计
         """
-        print(f"ℹ️ 开始整理{memory_type}类型的记忆，最小重叠标签数: {min_overlap}")
+        PrettyOutput.auto_print(
+            f"ℹ️ 开始整理{memory_type}类型的记忆，最小重叠标签数: {min_overlap}"
+        )
 
         # 加载记忆
         memories = self._load_memories(memory_type)
         if not memories:
-            print("ℹ️ 没有找到需要整理的记忆")
+            PrettyOutput.auto_print("ℹ️ 没有找到需要整理的记忆")
             return {"processed": 0, "merged": 0}
 
-        print(f"ℹ️ 加载了 {len(memories)} 个记忆")
+        PrettyOutput.auto_print(f"ℹ️ 加载了 {len(memories)} 个记忆")
 
         # 统计信息
         stats = {
@@ -288,7 +299,9 @@ class MemoryOrganizer:
 
             if overlap_count in overlap_groups:
                 groups = overlap_groups[overlap_count]
-                print(f"\nℹ️ 发现 {len(groups)} 个具有 {overlap_count} 个重叠标签的记忆组")
+                PrettyOutput.auto_print(
+                    f"\nℹ️ 发现 {len(groups)} 个具有 {overlap_count} 个重叠标签的记忆组"
+                )
 
                 for group in groups:
                     # 将活跃索引转换回原始索引
@@ -306,8 +319,8 @@ class MemoryOrganizer:
                             f"  - ID: {mem.get('id', '未知')}, "
                             f"标签: {', '.join(mem.get('tags', []))[:50]}..."
                         )
-                    joined_lines = '\n'.join(lines)
-                    print(f"ℹ️ {joined_lines}")
+                    joined_lines = "\n".join(lines)
+                    PrettyOutput.auto_print(f"ℹ️ {joined_lines}")
 
                     if not dry_run:
                         # 合并记忆
@@ -315,7 +328,7 @@ class MemoryOrganizer:
 
                         # 如果合并失败，跳过这组
                         if merged_memory is None:
-                            print("  ⚠️ 跳过这组记忆的合并")
+                            PrettyOutput.auto_print("  ⚠️ 跳过这组记忆的合并")
                             continue
 
                         # 保存新记忆
@@ -332,14 +345,14 @@ class MemoryOrganizer:
                         # 标记这些记忆已被删除
                         deleted_indices.update(original_indices)
                     else:
-                        print("  ℹ️ [模拟运行] 跳过实际合并")
+                        PrettyOutput.auto_print("  ℹ️ [模拟运行] 跳过实际合并")
 
         # 显示统计信息
-        print("\n✅ 整理完成！")
-        print(f"ℹ️ 总记忆数: {stats['total_memories']}")
-        print(f"ℹ️ 处理的组数: {stats['processed_groups']}")
-        print(f"ℹ️ 合并的记忆数: {stats['merged_memories']}")
-        print(f"ℹ️ 创建的新记忆数: {stats['created_memories']}")
+        PrettyOutput.auto_print("\n✅ 整理完成！")
+        PrettyOutput.auto_print(f"ℹ️ 总记忆数: {stats['total_memories']}")
+        PrettyOutput.auto_print(f"ℹ️ 处理的组数: {stats['processed_groups']}")
+        PrettyOutput.auto_print(f"ℹ️ 合并的记忆数: {stats['merged_memories']}")
+        PrettyOutput.auto_print(f"ℹ️ 创建的新记忆数: {stats['created_memories']}")
 
         return stats
 
@@ -371,7 +384,9 @@ class MemoryOrganizer:
         with open(new_file, "w", encoding="utf-8") as f:
             json.dump(memory, f, ensure_ascii=False, indent=2)
 
-        print(f"✅ 创建新记忆: {memory['id']} (标签: {', '.join(memory['tags'][:3])}...)")
+        PrettyOutput.auto_print(
+            f"✅ 创建新记忆: {memory['id']} (标签: {', '.join(memory['tags'][:3])}...)"
+        )
 
         # 删除原始记忆文件（先汇总日志，最后统一打印）
         info_lines: List[str] = []
@@ -382,7 +397,9 @@ class MemoryOrganizer:
                     file_path = Path(orig_memory["file_path"])
                     if file_path.exists():
                         file_path.unlink()
-                        info_lines.append(f"删除原始记忆: {orig_memory.get('id', '未知')}")
+                        info_lines.append(
+                            f"删除原始记忆: {orig_memory.get('id', '未知')}"
+                        )
                     else:
                         info_lines.append(
                             f"原始记忆文件已不存在，跳过删除: {orig_memory.get('id', '未知')}"
@@ -392,11 +409,11 @@ class MemoryOrganizer:
                         f"删除记忆文件失败 {orig_memory.get('file_path', '')}: {str(e)}"
                     )
         if info_lines:
-            joined_info = '\n'.join(info_lines)
-            print(f"ℹ️ {joined_info}")
+            joined_info = "\n".join(info_lines)
+            PrettyOutput.auto_print(f"ℹ️ {joined_info}")
         if warn_lines:
-            joined_warn = '\n'.join(warn_lines)
-            print(f"⚠️ {joined_warn}")
+            joined_warn = "\n".join(warn_lines)
+            PrettyOutput.auto_print(f"⚠️ {joined_warn}")
 
     def export_memories(
         self,
@@ -441,15 +458,17 @@ class MemoryOrganizer:
 
         # 统一展示导出进度日志
         if progress_lines:
-            joined_progress = '\n'.join(progress_lines)
-            print(f"ℹ️ {joined_progress}")
+            joined_progress = "\n".join(progress_lines)
+            PrettyOutput.auto_print(f"ℹ️ {joined_progress}")
 
         # 保存到文件
         output_file.parent.mkdir(parents=True, exist_ok=True)
         with open(output_file, "w", encoding="utf-8") as f:
             json.dump(all_memories, f, ensure_ascii=False, indent=2)
 
-        print(f"✅ 成功导出 {len(all_memories)} 个记忆到 {output_file}")
+        PrettyOutput.auto_print(
+            f"✅ 成功导出 {len(all_memories)} 个记忆到 {output_file}"
+        )
 
         return len(all_memories)
 
@@ -478,7 +497,7 @@ class MemoryOrganizer:
         if not isinstance(memories, list):
             raise ValueError("导入文件格式错误，应为记忆列表")
 
-        print(f"ℹ️ 准备导入 {len(memories)} 个记忆")
+        PrettyOutput.auto_print(f"ℹ️ 准备导入 {len(memories)} 个记忆")
 
         # 统计导入结果
         import_stats: Dict[str, int] = defaultdict(int)
@@ -496,7 +515,7 @@ class MemoryOrganizer:
             elif memory_type == "global_long_term":
                 memory_dir = self.global_memory_dir / memory_type
             else:
-                print(f"⚠️ 跳过不支持的记忆类型: {memory_type}")
+                PrettyOutput.auto_print(f"⚠️ 跳过不支持的记忆类型: {memory_type}")
                 skipped_count += 1
                 continue
 
@@ -513,7 +532,7 @@ class MemoryOrganizer:
             memory_file = memory_dir / f"{memory_id}.json"
 
             if memory_file.exists() and not overwrite:
-                print(f"ℹ️ 跳过已存在的记忆: {memory_id}")
+                PrettyOutput.auto_print(f"ℹ️ 跳过已存在的记忆: {memory_id}")
                 skipped_count += 1
                 continue
 
@@ -535,14 +554,17 @@ class MemoryOrganizer:
             import_stats[memory_type] += 1
 
         # 显示导入结果
-        print("\n✅ 导入完成！")
+        PrettyOutput.auto_print("\n✅ 导入完成！")
         if import_stats:
-            lines = [f"{memory_type}: 导入了 {count} 个记忆" for memory_type, count in import_stats.items()]
-            joined_lines = '\n'.join(lines)
-            print(f"ℹ️ {joined_lines}")
+            lines = [
+                f"{memory_type}: 导入了 {count} 个记忆"
+                for memory_type, count in import_stats.items()
+            ]
+            joined_lines = "\n".join(lines)
+            PrettyOutput.auto_print(f"ℹ️ {joined_lines}")
 
         if skipped_count > 0:
-            print(f"⚠️ 跳过了 {skipped_count} 个记忆")
+            PrettyOutput.auto_print(f"⚠️ 跳过了 {skipped_count} 个记忆")
 
         return dict(import_stats)
 
@@ -570,7 +592,6 @@ def organize(
     llm_group: Optional[str] = typer.Option(
         None, "-g", "--llm-group", help="使用的模型组，覆盖配置文件中的设置"
     ),
-
 ):
     """
     整理和合并具有相似标签的记忆。
@@ -588,11 +609,13 @@ def organize(
     """
     # 验证参数
     if memory_type not in ["project_long_term", "global_long_term"]:
-        print(f"❌ 错误：不支持的记忆类型 '{memory_type}'，请选择 'project_long_term' 或 'global_long_term'")
+        PrettyOutput.auto_print(
+            f"❌ 错误：不支持的记忆类型 '{memory_type}'，请选择 'project_long_term' 或 'global_long_term'"
+        )
         raise typer.Exit(1)
 
     if min_overlap < 2:
-        print("❌ 错误：最小重叠数必须大于等于2")
+        PrettyOutput.auto_print("❌ 错误：最小重叠数必须大于等于2")
         raise typer.Exit(1)
 
     # 创建整理器并执行
@@ -612,7 +635,7 @@ def organize(
         # typer.Exit 是正常的退出方式，直接传播
         raise
     except Exception as e:
-        print(f"❌ 记忆整理失败: {str(e)}")
+        PrettyOutput.auto_print(f"❌ 记忆整理失败: {str(e)}")
         raise typer.Exit(1)
 
 
@@ -656,7 +679,7 @@ def export(
         invalid_types = [mt for mt in memory_types if mt not in valid_types]
         if invalid_types:
             invalid_str = ", ".join(f"'{mt}'" for mt in invalid_types)
-            print(f"❌ 错误：不支持的记忆类型: {invalid_str}")
+            PrettyOutput.auto_print(f"❌ 错误：不支持的记忆类型: {invalid_str}")
             raise typer.Exit(1)
 
         count = organizer.export_memories(
@@ -668,11 +691,11 @@ def export(
         if count > 0:
             raise typer.Exit(0)
         else:
-            print("⚠️ 没有找到要导出的记忆")
+            PrettyOutput.auto_print("⚠️ 没有找到要导出的记忆")
             raise typer.Exit(0)
 
     except Exception as e:
-        print(f"❌ 导出失败: {str(e)}")
+        PrettyOutput.auto_print(f"❌ 导出失败: {str(e)}")
         raise typer.Exit(1)
 
 
@@ -712,14 +735,14 @@ def import_memories(
         if total_imported > 0:
             raise typer.Exit(0)
         else:
-            print("⚠️ 没有导入任何记忆")
+            PrettyOutput.auto_print("⚠️ 没有导入任何记忆")
             raise typer.Exit(0)
 
     except FileNotFoundError as e:
-        print(f"❌ {str(e)}")
+        PrettyOutput.auto_print(f"❌ {str(e)}")
         raise typer.Exit(1)
     except Exception as e:
-        print(f"❌ 导入失败: {str(e)}")
+        PrettyOutput.auto_print(f"❌ 导入失败: {str(e)}")
         raise typer.Exit(1)