jarvis-ai-assistant 0.1.132__py3-none-any.whl → 0.1.138__py3-none-any.whl

jarvis/jarvis_utils/methodology.py

@@ -18,7 +18,6 @@ import faiss
 from typing import Dict, Any, List, Tuple, Optional
 from jarvis.jarvis_utils.output import PrettyOutput, OutputType
 from jarvis.jarvis_utils.embedding import load_embedding_model
-from jarvis.jarvis_utils.config import dont_use_local_model
 
 # 全局缓存，避免重复计算嵌入向量
 _methodology_embeddings_cache = {}
@@ -27,7 +26,7 @@ _methodology_index_cache: Optional[Tuple[faiss.IndexIDMap, List[Dict[str, str]],
 def _get_cache_directory() -> str:
     """
     获取缓存目录路径，如果不存在则创建
-    
+
     返回：
         str: 缓存目录的路径
     """
@@ -42,7 +41,7 @@ def _get_cache_directory() -> str:
 def _get_embeddings_cache_path() -> str:
     """
     获取嵌入向量缓存文件的路径
-    
+
     返回：
         str: 嵌入向量缓存文件的路径
     """
@@ -51,7 +50,7 @@ def _get_embeddings_cache_path() -> str:
 def _get_index_cache_path() -> str:
     """
     获取索引缓存文件的路径
-    
+
     返回：
         str: 索引缓存文件的路径
     """
@@ -60,7 +59,7 @@ def _get_index_cache_path() -> str:
 def _get_index_metadata_path() -> str:
     """
     获取索引元数据文件的路径
-    
+
     返回：
         str: 索引元数据文件的路径
     """
@@ -69,14 +68,14 @@ def _get_index_metadata_path() -> str:
 def _load_embeddings_cache() -> Dict[int, np.ndarray]:
     """
     从文件系统加载嵌入向量缓存
-    
+
     返回：
         Dict[int, np.ndarray]: 嵌入向量缓存字典
     """
     cache_path = _get_embeddings_cache_path()
     if not os.path.exists(cache_path):
         return {}
-    
+
     try:
         with open(cache_path, "rb") as f:
             embeddings_cache = pickle.load(f)
@@ -87,18 +86,18 @@ def _load_embeddings_cache() -> Dict[int, np.ndarray]:
 def _save_embeddings_cache(cache: Dict[int, np.ndarray]) -> bool:
     """
     将嵌入向量缓存保存到文件系统
-    
+
     参数：
         cache: 要保存的嵌入向量缓存字典
-        
+
     返回：
         bool: 保存是否成功
     """
     if not cache:
         return False
-    
+
     cache_path = _get_embeddings_cache_path()
-    
+
     try:
         with open(cache_path, "wb") as f:
             pickle.dump(cache, f)
@@ -110,62 +109,62 @@ def _save_embeddings_cache(cache: Dict[int, np.ndarray]) -> bool:
 def _load_index_cache() -> Optional[Tuple[faiss.IndexIDMap, List[Dict[str, str]], str]]:
     """
     从文件系统加载索引缓存
-    
+
     返回：
         Optional[Tuple[faiss.IndexIDMap, List[Dict[str, str]], str]]: 索引缓存元组
     """
     index_path = _get_index_cache_path()
     metadata_path = _get_index_metadata_path()
-    
+
     if not os.path.exists(index_path) or not os.path.exists(metadata_path):
         return None
-    
+
     try:
         # 加载索引
         index = faiss.read_index(index_path)
-        
+
         # 加载元数据
         with open(metadata_path, "rb") as f:
             metadata = pickle.load(f)
-            
+
         methodology_data = metadata.get("methodology_data", [])
         methodology_hash = metadata.get("methodology_hash", "")
-        
+
         if isinstance(index, faiss.IndexIDMap) and methodology_data and methodology_hash:
             return index, methodology_data, methodology_hash
     except Exception as e:
         PrettyOutput.print(f"加载索引缓存失败: {str(e)}", OutputType.WARNING)
-    
+
     return None
 
 def _save_index_cache(index: faiss.IndexIDMap, methodology_data: List[Dict[str, str]], methodology_hash: str) -> bool:
     """
     将索引缓存保存到文件系统
-    
+
     参数：
         index: FAISS索引
         methodology_data: 方法论数据列表
         methodology_hash: 方法论文件哈希值
-        
+
     返回：
         bool: 保存是否成功
     """
     index_path = _get_index_cache_path()
     metadata_path = _get_index_metadata_path()
-    
+
     try:
         # 保存索引
         faiss.write_index(index, index_path)
-        
+
         # 保存元数据
         metadata = {
             "methodology_data": methodology_data,
             "methodology_hash": methodology_hash
         }
-        
+
         with open(metadata_path, "wb") as f:
             pickle.dump(metadata, f)
-            
+
         return True
     except Exception as e:
         PrettyOutput.print(f"保存索引缓存失败: {str(e)}", OutputType.WARNING)
@@ -174,11 +173,11 @@ def _save_index_cache(index: faiss.IndexIDMap, methodology_data: List[Dict[str,
 def _create_methodology_embedding(embedding_model: Any, methodology_text: str) -> np.ndarray:
     """
     为方法论文本创建嵌入向量。
-    
+
     参数：
         embedding_model: 使用的嵌入模型
         methodology_text: 要创建嵌入的文本
-        
+
     返回：
         np.ndarray: 嵌入向量
     """
@@ -187,21 +186,21 @@ def _create_methodology_embedding(embedding_model: Any, methodology_text: str) -
         cache_key = hash(methodology_text)
         if cache_key in _methodology_embeddings_cache:
             return _methodology_embeddings_cache[cache_key]
-            
+
         # 截断长文本
         max_length = 512
         text = ' '.join(methodology_text.split()[:max_length])
-        
+
         # 使用sentence_transformers模型获取嵌入向量
-        embedding = embedding_model.encode([text], 
+        embedding = embedding_model.encode([text],
                                           convert_to_tensor=True,
                                           normalize_embeddings=True)
         vector = np.array(embedding.cpu().numpy(), dtype=np.float32)
         result = vector[0]  # 返回第一个向量，因为我们只编码了一个文本
-        
+
         # 缓存嵌入向量以便后续使用
         _methodology_embeddings_cache[cache_key] = result
-        
+
         return result
     except Exception as e:
         PrettyOutput.print(f"创建方法论嵌入向量失败: {str(e)}", OutputType.ERROR)
@@ -210,14 +209,14 @@ def _create_methodology_embedding(embedding_model: Any, methodology_text: str) -
 def make_methodology_prompt(data: Dict[str, str]) -> str:
     """
     从方法论数据生成格式化提示
-    
+
     参数：
         data: 方法论数据字典
-        
+
     返回：
         str: 格式化后的提示字符串
     """
-    ret = """这是处理以往问题的标准方法论，如果当前任务类似，可以参考使用，如果不相关，请忽略：\n""" 
+    ret = """这是处理以往问题的标准方法论，如果当前任务类似，可以参考使用，如果不相关，请忽略：\n"""
     for key, value in data.items():
         ret += f"问题: {key}\n方法论: {value}\n"
     return ret
@@ -225,7 +224,7 @@ def make_methodology_prompt(data: Dict[str, str]) -> str:
 def _get_methodology_directory() -> str:
     """
     获取方法论目录路径，如果不存在则创建
-    
+
     返回：
         str: 方法论目录的路径
     """
@@ -240,46 +239,46 @@ def _get_methodology_directory() -> str:
 def _get_methodology_files_hash() -> str:
     """
     计算所有方法论文件的组合哈希值，用于检测文件变化
-    
+
     返回：
         str: 所有方法论文件的组合哈希值
     """
     methodology_dir = _get_methodology_directory()
     if not os.path.exists(methodology_dir):
         return ""
-    
+
     # 获取所有方法论文件的路径和修改时间
     files_data = []
     for filepath in glob.glob(os.path.join(methodology_dir, "*.json")):
         mtime = os.path.getmtime(filepath)
         files_data.append((filepath, mtime))
-    
+
     # 按路径排序，保证哈希值的一致性
     files_data.sort(key=lambda x: x[0])
-    
+
     # 计算组合哈希值
     if not files_data:
         return ""
-    
+
     hasher = hashlib.md5()
     for filepath, mtime in files_data:
         hasher.update(f"{filepath}:{mtime}".encode("utf-8"))
-    
+
     return hasher.hexdigest()
 
 def _load_all_methodologies() -> Dict[str, str]:
     """
     加载所有方法论文件
-    
+
     返回：
         Dict[str, str]: 方法论字典，键为问题类型，值为方法论内容
     """
     methodology_dir = _get_methodology_directory()
     all_methodologies = {}
-    
+
     if not os.path.exists(methodology_dir):
         return all_methodologies
-    
+
     for filepath in glob.glob(os.path.join(methodology_dir, "*.json")):
         try:
             with open(filepath, "r", encoding="utf-8", errors="ignore") as f:
@@ -291,7 +290,7 @@ def _load_all_methodologies() -> Dict[str, str]:
         except Exception as e:
             filename = os.path.basename(filepath)
             PrettyOutput.print(f"加载方法论文件 {filename} 失败: {str(e)}", OutputType.WARNING)
-    
+
     return all_methodologies
 
 def _migrate_from_old_format():
@@ -301,34 +300,34 @@ def _migrate_from_old_format():
     old_format_file = os.path.expanduser("~/.jarvis/methodology")
     if not os.path.exists(old_format_file):
         return
-    
+
     try:
         # 加载旧格式文件
         with open(old_format_file, "r", encoding="utf-8", errors="ignore") as f:
             old_data = yaml.safe_load(f) or {}
-        
+
         if not old_data:
             return
-        
+
         # 创建新目录
         methodology_dir = _get_methodology_directory()
-        
+
         # 迁移每个方法论
         migrated_count = 0
         for problem_type, content in old_data.items():
             # 为每个方法论创建文件名（使用问题类型的MD5哈希作为文件名）
             safe_filename = hashlib.md5(problem_type.encode('utf-8')).hexdigest()
             file_path = os.path.join(methodology_dir, f"{safe_filename}.json")
-            
+
             # 保存为新格式
             with open(file_path, "w", encoding="utf-8", errors="ignore") as f:
                 json.dump({
                     "problem_type": problem_type,
                     "content": content
                 }, f, ensure_ascii=False, indent=2)
-            
+
             migrated_count += 1
-        
+
         if migrated_count > 0:
             # 备份旧文件
             backup_path = old_format_file + ".bak"
@@ -343,15 +342,15 @@ def _migrate_from_old_format():
 def load_methodology(user_input: str) -> str:
     """
     加载方法论并构建向量索引以进行相似性搜索。
-    
+
     参数：
         user_input: 要搜索方法论的输入文本
-        
+
     返回：
         str: 相关的方法论提示，如果未找到方法论则返回空字符串
     """
     from yaspin import yaspin
-    
+
     # 加载嵌入向量缓存
     global _methodology_embeddings_cache
     if not _methodology_embeddings_cache:
@@ -359,29 +358,22 @@ def load_methodology(user_input: str) -> str:
             _methodology_embeddings_cache = _load_embeddings_cache()
             spinner.text = f"加载嵌入向量缓存完成 ({len(_methodology_embeddings_cache)} 个向量)"
             spinner.ok("✅")
-    
+
     # 检查是否需要从旧格式迁移
     with yaspin(text="检查方法论格式...", color="yellow") as spinner:
         _migrate_from_old_format()
         spinner.text = "检查方法论格式完成"
         spinner.ok("✅")
-    
+
     # 获取方法论目录
     methodology_dir = _get_methodology_directory()
     if not os.path.exists(methodology_dir) or not glob.glob(os.path.join(methodology_dir, "*.json")):
         return ""
-    
+
     try:
         # 获取文件的修改时间戳组合哈希，用于检测文件是否被修改
         methodology_hash = _get_methodology_files_hash()
-        
-        with yaspin(text="加载方法论文件...", color="yellow") as spinner:
-            data = _load_all_methodologies()
-            if dont_use_local_model():
-                spinner.text = "加载方法论文件完成"
-                spinner.ok("✅")
-                return make_methodology_prompt(data)
-        
+
         # 检查缓存的索引是否可用且方法论文件未被修改
         global _methodology_index_cache
         if _methodology_index_cache is None:
@@ -394,7 +386,7 @@ def load_methodology(user_input: str) -> str:
                 else:
                     spinner.text = "没有可用的索引缓存"
                     spinner.fail("❌")
-        
+
         if _methodology_index_cache is not None:
             cached_index, cached_data, cache_hash = _methodology_index_cache
             if cache_hash == methodology_hash:
@@ -404,12 +396,12 @@ def load_methodology(user_input: str) -> str:
                     methodology_data = cached_data
                     spinner.text = "使用缓存的方法论索引完成"
                     spinner.ok("✅")
-                
+
                 with yaspin(text="加载嵌入模型...", color="yellow") as spinner:
                     embedding_model = load_embedding_model()
                     spinner.text = "加载嵌入模型完成"
                     spinner.ok("✅")
-                
+
                 with yaspin(text="执行搜索...", color="yellow") as spinner:
                     # 使用缓存构造输入文本的嵌入
                     query_embedding = _create_methodology_embedding(embedding_model, user_input)
@@ -419,7 +411,7 @@ def load_methodology(user_input: str) -> str:
                     ) # type: ignore
                     spinner.text = "执行搜索完成"
                     spinner.ok("✅")
-                
+
                 with yaspin(text="处理搜索结果...", color="yellow") as spinner:
                     relevant_methodologies = {}
                     for dist, idx in zip(distances[0], indices[0]):
@@ -430,11 +422,12 @@ def load_methodology(user_input: str) -> str:
                                 relevant_methodologies[methodology["key"]] = methodology["value"]
                     spinner.text = "处理搜索结果完成"
                     spinner.ok("✅")
-                
+
                 if relevant_methodologies:
+                    spinner.text = f"找到相关方法论: {', '.join(relevant_methodologies.keys())}"
+                    spinner.ok("✅")
                     return make_methodology_prompt(relevant_methodologies)
-                return make_methodology_prompt(data)
-        
+
         # 如果缓存无效，从头构建索引
         with yaspin(text="初始化数据结构...", color="yellow") as spinner:
             methodology_data: List[Dict[str, str]] = []
@@ -442,77 +435,85 @@ def load_methodology(user_input: str) -> str:
             ids: List[int] = []
             spinner.text = "初始化数据结构完成"
             spinner.ok("✅")
-            
-            with yaspin(text="加载嵌入模型...", color="yellow") as spinner:
-                embedding_model = load_embedding_model()
-                spinner.text = "加载嵌入模型完成"
+
+        with yaspin(text="加载嵌入模型...", color="yellow") as spinner:
+            embedding_model = load_embedding_model()
+            spinner.text = "加载嵌入模型完成"
+            spinner.ok("✅")
+
+        with yaspin(text="创建测试嵌入...", color="yellow") as spinner:
+            test_embedding = _create_methodology_embedding(embedding_model, "test")
+            embedding_dimension = len(test_embedding)
+            spinner.text = "创建测试嵌入完成"
+            spinner.ok("✅")
+
+        with yaspin(text="加载方法论文件...", color="yellow") as spinner:
+            data = _load_all_methodologies()
+            spinner.text = "加载方法论文件完成"
+            spinner.ok("✅")
+
+        with yaspin(text="处理方法论数据...", color="yellow") as spinner:
+            for i, (key, value) in enumerate(data.items()):
+                methodology_text = f"{key}\n{value}"
+                embedding = _create_methodology_embedding(embedding_model, methodology_text)
+                vectors.append(embedding)
+                ids.append(i)
+                methodology_data.append({"key": key, "value": value})
+            spinner.text = "处理方法论数据完成"
+            spinner.ok("✅")
+
+        if vectors:
+            with yaspin(text="构建索引...", color="yellow") as spinner:
+                vectors_array = np.vstack(vectors)
+                hnsw_index = faiss.IndexHNSWFlat(embedding_dimension, 16)
+                hnsw_index.hnsw.efConstruction = 40
+                hnsw_index.hnsw.efSearch = 16
+                methodology_index = faiss.IndexIDMap(hnsw_index)
+                methodology_index.add_with_ids(vectors_array, np.array(ids)) # type: ignore
+                # 缓存构建好的索引和数据以及时间戳哈希
+                _methodology_index_cache = (methodology_index, methodology_data, methodology_hash)
+
+                # 将索引和嵌入向量缓存保存到文件系统
+                _save_index_cache(methodology_index, methodology_data, methodology_hash)
+
+                spinner.text = "构建索引完成"
                 spinner.ok("✅")
-            
-            with yaspin(text="创建测试嵌入...", color="yellow") as spinner:
-                test_embedding = _create_methodology_embedding(embedding_model, "test")
-                embedding_dimension = len(test_embedding)
-                spinner.text = "创建测试嵌入完成"
+
+            with yaspin(text="执行搜索...", color="yellow") as spinner:
+                query_embedding = _create_methodology_embedding(embedding_model, user_input)
+                k = min(10, len(methodology_data))
+                distances, indices = methodology_index.search(
+                    query_embedding.reshape(1, -1), k
+                ) # type: ignore
+                spinner.text = "执行搜索完成"
                 spinner.ok("✅")
-            
-            with yaspin(text="处理方法论数据...", color="yellow") as spinner:
-                for i, (key, value) in enumerate(data.items()):
-                    methodology_text = f"{key}\n{value}"
-                    embedding = _create_methodology_embedding(embedding_model, methodology_text)
-                    vectors.append(embedding)
-                    ids.append(i)
-                    methodology_data.append({"key": key, "value": value})
-                spinner.text = "处理方法论数据完成"
+
+            with yaspin(text="处理搜索结果...", color="yellow") as spinner:
+                relevant_methodologies = {}
+                for dist, idx in zip(distances[0], indices[0]):
+                    if idx >= 0:
+                        similarity = 1.0 / (1.0 + float(dist))
+                        methodology = methodology_data[idx]
+                        if similarity >= 0.5:
+                            relevant_methodologies[methodology["key"]] = methodology["value"]
+                spinner.text = "处理搜索结果完成"
                 spinner.ok("✅")
-            
-            if vectors:
-                with yaspin(text="构建索引...", color="yellow") as spinner:
-                    vectors_array = np.vstack(vectors)
-                    hnsw_index = faiss.IndexHNSWFlat(embedding_dimension, 16)
-                    hnsw_index.hnsw.efConstruction = 40
-                    hnsw_index.hnsw.efSearch = 16
-                    methodology_index = faiss.IndexIDMap(hnsw_index)
-                    methodology_index.add_with_ids(vectors_array, np.array(ids)) # type: ignore
-                    # 缓存构建好的索引和数据以及时间戳哈希
-                    _methodology_index_cache = (methodology_index, methodology_data, methodology_hash)
-                    
-                    # 将索引和嵌入向量缓存保存到文件系统
-                    _save_index_cache(methodology_index, methodology_data, methodology_hash)
-                    
-                    spinner.text = "构建索引完成"
-                    spinner.ok("✅")
-                
-                with yaspin(text="执行搜索...", color="yellow") as spinner:
-                    query_embedding = _create_methodology_embedding(embedding_model, user_input)
-                    k = min(3, len(methodology_data))
-                    distances, indices = methodology_index.search(
-                        query_embedding.reshape(1, -1), k
-                    ) # type: ignore
-                    spinner.text = "执行搜索完成"
-                    spinner.ok("✅")
-                
-                with yaspin(text="处理搜索结果...", color="yellow") as spinner:
-                    relevant_methodologies = {}
-                    for dist, idx in zip(distances[0], indices[0]):
-                        if idx >= 0:
-                            similarity = 1.0 / (1.0 + float(dist))
-                            methodology = methodology_data[idx]
-                            if similarity >= 0.5:
-                                relevant_methodologies[methodology["key"]] = methodology["value"]
-                    spinner.text = "处理搜索结果完成"
+
+            # 保存嵌入向量缓存到文件系统
+            with yaspin(text="保存嵌入向量缓存...", color="yellow") as spinner:
+                if _save_embeddings_cache(_methodology_embeddings_cache):
+                    spinner.text = f"保存嵌入向量缓存完成 ({len(_methodology_embeddings_cache)} 个向量)"
                     spinner.ok("✅")
-                
-                # 保存嵌入向量缓存到文件系统
-                with yaspin(text="保存嵌入向量缓存...", color="yellow") as spinner:
-                    if _save_embeddings_cache(_methodology_embeddings_cache):
-                        spinner.text = f"保存嵌入向量缓存完成 ({len(_methodology_embeddings_cache)} 个向量)"
-                        spinner.ok("✅")
-                    else:
-                        spinner.text = "保存嵌入向量缓存失败"
-                        spinner.fail("❌")
-                
-                if relevant_methodologies:
-                    return make_methodology_prompt(relevant_methodologies)
-            return make_methodology_prompt(data)
+                else:
+                    spinner.text = "保存嵌入向量缓存失败"
+                    spinner.fail("❌")
+
+            if relevant_methodologies:
+                spinner.write(f"找到相关方法论: {', '.join(relevant_methodologies.keys())}")
+                return make_methodology_prompt(relevant_methodologies)
+        spinner.write(f"未找到相关的方法论")
+        spinner.fail("❌")
+        return ""
     except Exception as e:
         PrettyOutput.print(f"加载方法论失败: {str(e)}", OutputType.ERROR)
         return ""

jarvis/jarvis_utils/output.py

@@ -17,11 +17,11 @@ from rich.syntax import Syntax
 from rich.style import Style as RichStyle
 from pygments.lexers import guess_lexer
 from pygments.util import ClassNotFound
-from .globals import console, get_agent_list
+from jarvis.jarvis_utils.globals import console, get_agent_list
 class OutputType(Enum):
     """
     输出类型枚举，用于分类和样式化不同类型的消息。
-    
+
     属性：
         SYSTEM: AI助手消息
         CODE: 代码相关输出
@@ -51,7 +51,7 @@ class OutputType(Enum):
 class PrettyOutput:
     """
     使用rich库格式化和显示富文本输出的类。
-    
+
     提供以下方法：
     - 使用适当的样式格式化不同类型的输出
     - 代码块的语法高亮
@@ -110,11 +110,11 @@ class PrettyOutput:
     def _detect_language(text: str, default_lang: str = 'markdown') -> str:
         """
         检测给定文本的编程语言。
-        
+
         参数：
             text: 要分析的文本
             default_lang: 如果检测失败，默认返回的语言
-            
+
         返回：
             str: 检测到的语言名称
         """
@@ -128,11 +128,11 @@ class PrettyOutput:
     def _format(output_type: OutputType, timestamp: bool = True) -> Text:
         """
         使用时间戳和图标格式化输出头。
-        
+
         参数：
             output_type: 输出类型
             timestamp: 是否包含时间戳
-            
+
         返回：
             Text: 格式化后的rich Text对象
         """
@@ -149,7 +149,7 @@ class PrettyOutput:
     def print(text: str, output_type: OutputType, timestamp: bool = True, lang: Optional[str] = None, traceback: bool = False):
         """
         使用样式和语法高亮打印格式化输出。
-        
+
         参数：
             text: 要打印的文本内容
             output_type: 输出类型（影响样式）
@@ -186,13 +186,13 @@ class PrettyOutput:
         )
         console.print()
         console.print(panel)
-        if traceback:
+        if traceback or output_type == OutputType.ERROR:
             console.print_exception()
     @staticmethod
     def section(title: str, output_type: OutputType = OutputType.INFO):
         """
         在样式化面板中打印章节标题。
-        
+
         参数：
             title: 章节标题文本
             output_type: 输出类型（影响样式）
@@ -208,7 +208,7 @@ class PrettyOutput:
     def print_stream(text: str, is_thinking: bool = False):
         """
         打印流式输出，不带换行符。
-        
+
         参数：
             text: 要打印的文本
         """
@@ -228,10 +228,10 @@ class PrettyOutput:
     def _get_style(output_type: OutputType) -> RichStyle:
         """
         获取预定义的RichStyle用于输出类型。
-        
+
         参数：
             output_type: 要获取样式的输出类型
-            
+
         返回：
             RichStyle: 对应的样式
         """