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

jarvis/jarvis_utils/methodology.py

@@ -9,12 +9,168 @@
 """
 import os
 import yaml
+import glob
+import json
+import hashlib
+import pickle
 import numpy as np
 import faiss
-from typing import Dict, Any, List
+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 = {}
+_methodology_index_cache: Optional[Tuple[faiss.IndexIDMap, List[Dict[str, str]], str]] = None
+
+def _get_cache_directory() -> str:
+    """
+    获取缓存目录路径，如果不存在则创建
+    
+    返回：
+        str: 缓存目录的路径
+    """
+    cache_dir = os.path.expanduser("~/.jarvis/cache")
+    if not os.path.exists(cache_dir):
+        try:
+            os.makedirs(cache_dir, exist_ok=True)
+        except Exception as e:
+            PrettyOutput.print(f"创建缓存目录失败: {str(e)}", OutputType.ERROR)
+    return cache_dir
+
+def _get_embeddings_cache_path() -> str:
+    """
+    获取嵌入向量缓存文件的路径
+    
+    返回：
+        str: 嵌入向量缓存文件的路径
+    """
+    return os.path.join(_get_cache_directory(), "methodology_embeddings.pkl")
+
+def _get_index_cache_path() -> str:
+    """
+    获取索引缓存文件的路径
+    
+    返回：
+        str: 索引缓存文件的路径
+    """
+    return os.path.join(_get_cache_directory(), "methodology_index.faiss")
+
+def _get_index_metadata_path() -> str:
+    """
+    获取索引元数据文件的路径
+    
+    返回：
+        str: 索引元数据文件的路径
+    """
+    return os.path.join(_get_cache_directory(), "methodology_index_metadata.pkl")
+
+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)
+        return embeddings_cache
+    except Exception as e:
+        return {}
+
+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)
+        return True
+    except Exception as e:
+        PrettyOutput.print(f"保存嵌入向量缓存失败: {str(e)}", OutputType.WARNING)
+        return False
+
+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)
+        return False
+
 def _create_methodology_embedding(embedding_model: Any, methodology_text: str) -> np.ndarray:
     """
     为方法论文本创建嵌入向量。
@@ -27,6 +183,11 @@ def _create_methodology_embedding(embedding_model: Any, methodology_text: str) -
         np.ndarray: 嵌入向量
     """
     try:
+        # 检查缓存中是否已有此文本的嵌入向量
+        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])
@@ -36,10 +197,16 @@ def _create_methodology_embedding(embedding_model: Any, methodology_text: str) -
                                           convert_to_tensor=True,
                                           normalize_embeddings=True)
         vector = np.array(embedding.cpu().numpy(), dtype=np.float32)
-        return vector[0]  # 返回第一个向量，因为我们只编码了一个文本
+        result = vector[0]  # 返回第一个向量，因为我们只编码了一个文本
+        
+        # 缓存嵌入向量以便后续使用
+        _methodology_embeddings_cache[cache_key] = result
+        
+        return result
     except Exception as e:
         PrettyOutput.print(f"创建方法论嵌入向量失败: {str(e)}", OutputType.ERROR)
         return np.zeros(1536, dtype=np.float32)
+
 def make_methodology_prompt(data: Dict[str, str]) -> str:
     """
     从方法论数据生成格式化提示
@@ -55,6 +222,124 @@ def make_methodology_prompt(data: Dict[str, str]) -> str:
         ret += f"问题: {key}\n方法论: {value}\n"
     return ret
 
+def _get_methodology_directory() -> str:
+    """
+    获取方法论目录路径，如果不存在则创建
+    
+    返回：
+        str: 方法论目录的路径
+    """
+    methodology_dir = os.path.expanduser("~/.jarvis/methodologies")
+    if not os.path.exists(methodology_dir):
+        try:
+            os.makedirs(methodology_dir, exist_ok=True)
+        except Exception as e:
+            PrettyOutput.print(f"创建方法论目录失败: {str(e)}", OutputType.ERROR)
+    return methodology_dir
+
+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:
+                methodology = json.load(f)
+                problem_type = methodology.get("problem_type", "")
+                content = methodology.get("content", "")
+                if problem_type and content:
+                    all_methodologies[problem_type] = content
+        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():
+    """
+    从旧的单文件格式迁移到新的多文件格式
+    """
+    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"
+            try:
+                os.rename(old_format_file, backup_path)
+                PrettyOutput.print(f"已成功迁移 {migrated_count} 个方法论，旧文件已备份为 {backup_path}", OutputType.INFO)
+            except Exception as e:
+                PrettyOutput.print(f"备份旧文件失败，但已完成迁移: {str(e)}", OutputType.WARNING)
+    except Exception as e:
+        PrettyOutput.print(f"迁移方法论失败: {str(e)}", OutputType.ERROR)
+
 def load_methodology(user_input: str) -> str:
     """
     加载方法论并构建向量索引以进行相似性搜索。
@@ -66,19 +351,91 @@ def load_methodology(user_input: str) -> str:
         str: 相关的方法论提示，如果未找到方法论则返回空字符串
     """
     from yaspin import yaspin
-    user_jarvis_methodology = os.path.expanduser("~/.jarvis/methodology")
-    if not os.path.exists(user_jarvis_methodology):
+    
+    # 加载嵌入向量缓存
+    global _methodology_embeddings_cache
+    if not _methodology_embeddings_cache:
+        with yaspin(text="加载嵌入向量缓存...", color="yellow") as spinner:
+            _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:
-            with open(user_jarvis_methodology, "r", encoding="utf-8", errors="ignore") as f:
-                data = yaml.safe_load(f)
+            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:
+            # 尝试从文件系统加载索引缓存
+            with yaspin(text="加载索引缓存...", color="yellow") as spinner:
+                _methodology_index_cache = _load_index_cache()
+                if _methodology_index_cache:
+                    spinner.text = "加载索引缓存完成"
+                    spinner.ok("✅")
+                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:
+                # 直接使用缓存的索引和数据
+                with yaspin(text="使用缓存的方法论索引...", color="yellow") as spinner:
+                    methodology_index = cached_index
+                    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)
+                    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 = "处理搜索结果完成"
+                    spinner.ok("✅")
+                
+                if relevant_methodologies:
+                    return make_methodology_prompt(relevant_methodologies)
+                return make_methodology_prompt(data)
+        
+        # 如果缓存无效，从头构建索引
         with yaspin(text="初始化数据结构...", color="yellow") as spinner:
             methodology_data: List[Dict[str, str]] = []
             vectors: List[np.ndarray] = []
@@ -115,6 +472,12 @@ def load_methodology(user_input: str) -> str:
                     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("✅")
                 
@@ -137,10 +500,19 @@ def load_methodology(user_input: str) -> str:
                                 relevant_methodologies[methodology["key"]] = methodology["value"]
                     spinner.text = "处理搜索结果完成"
                     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)
     except Exception as e:
+        PrettyOutput.print(f"加载方法论失败: {str(e)}", OutputType.ERROR)
         return ""

jarvis/jarvis_utils/output.py

@@ -186,7 +186,7 @@ class PrettyOutput:
         )
         console.print()
         console.print(panel)
-        if traceback or output_type == OutputType.ERROR:
+        if traceback:
             console.print_exception()
     @staticmethod
     def section(title: str, output_type: OutputType = OutputType.INFO):
@@ -205,14 +205,16 @@ class PrettyOutput:
         console.print(panel)
         console.print()
     @staticmethod
-    def print_stream(text: str):
+    def print_stream(text: str, is_thinking: bool = False):
         """
         打印流式输出，不带换行符。
         
         参数：
             text: 要打印的文本
         """
-        style = RichStyle(color="bright_cyan", italic=True, frame=True, meta={"icon": "🤖"})
+        style = RichStyle(color="bright_cyan", bold=True, frame=True, meta={"icon": "🤖"})
+        if is_thinking:
+            style = RichStyle(color="grey58", italic=True, frame=True, meta={"icon": "🤖"})
         console.print(text, style=style, end="")
     @staticmethod
     def print_stream_end():

jarvis/jarvis_utils/utils.py

@@ -9,10 +9,12 @@ from jarvis.jarvis_utils.embedding import get_context_token_count
 from jarvis.jarvis_utils.input import get_single_line_input
 from jarvis.jarvis_utils.output import PrettyOutput, OutputType
 def init_env():
-    """Initialize environment variables from ~/.jarvis/env file.
+    """初始化环境变量从~/.jarvis/env文件
     
-    Creates the .jarvis directory if it doesn't exist and loads environment variables
-    from the env file. Handles file reading errors gracefully.
+    功能：
+    1. 创建不存在的.jarvis目录
+    2. 加载环境变量到os.environ
+    3. 处理文件读取异常
     """
     jarvis_dir = Path.home() / ".jarvis"
     env_file = jarvis_dir / "env"
@@ -34,6 +36,15 @@ def init_env():
         except Exception as e:
             PrettyOutput.print(f"警告: 读取 {env_file} 失败: {e}", OutputType.WARNING)
 def while_success(func, sleep_time: float = 0.1):
+    """循环执行函数直到成功
+    
+    参数：
+    func -- 要执行的函数
+    sleep_time -- 每次失败后的等待时间（秒）
+    
+    返回：
+    函数执行结果
+    """
     while True:
         try:
             return func()
@@ -87,10 +98,16 @@ def get_file_line_count(filename: str) -> int:
     except Exception as e:
         return 0
 def init_gpu_config() -> Dict:
-    """Initialize GPU configuration based on available hardware.
+    """初始化GPU配置
     
-    Returns:
-        Dict: GPU configuration including memory sizes and availability
+    功能：
+    1. 检测CUDA可用性
+    2. 计算设备内存和共享内存
+    3. 设置CUDA内存分配策略
+    4. 处理异常情况
+    
+    返回：
+    包含GPU配置信息的字典
     """
     config = {
         "has_gpu": False,
@@ -98,6 +115,8 @@ def init_gpu_config() -> Dict:
         "device_memory": 0,
         "memory_fraction": 0.8  # 默认使用80%的可用内存
     }
+
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
     
     try:
         import torch
@@ -130,7 +149,17 @@ def init_gpu_config() -> Dict:
 
 
 def is_long_context(files: list) -> bool:
-    """检查文件列表是否属于长上下文（总字符数超过最大上下文长度的80%）"""
+    """检查文件列表是否属于长上下文
+    
+    判断标准：
+    当总token数超过最大上下文长度的80%时视为长上下文
+    
+    参数：
+    files -- 要检查的文件路径列表
+    
+    返回：
+    布尔值表示是否属于长上下文
+    """
     max_token_count = get_max_token_count()
     threshold = max_token_count * 0.8
     total_tokens = 0
@@ -148,3 +177,26 @@ def is_long_context(files: list) -> bool:
             continue
             
     return total_tokens > threshold
+
+
+def ot(tag_name: str) -> str:
+    """生成HTML标签开始标记
+    
+    参数：
+    tag_name -- HTML标签名称
+    
+    返回：
+    格式化的开始标签字符串
+    """
+    return f"<{tag_name}>"
+
+def ct(tag_name: str) -> str:
+    """生成HTML标签结束标记
+    
+    参数：
+    tag_name -- HTML标签名称
+    
+    返回：
+    格式化的结束标签字符串
+    """
+    return f"</{tag_name}>"

{jarvis_ai_assistant-0.1.131.dist-info → jarvis_ai_assistant-0.1.132.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jarvis-ai-assistant
-Version: 0.1.131
+Version: 0.1.132
 Summary: Jarvis: An AI assistant that uses tools to interact with the system
 Home-page: https://github.com/skyfireitdiy/Jarvis
 Author: skyfire
@@ -64,7 +64,8 @@ Requires-Dist: jedi>=0.17.2
 Requires-Dist: psutil>=7.0.0
 Requires-Dist: fastapi>=0.115.4
 Requires-Dist: uvicorn>=0.33.0
-Requires-Dist: tree-sitter>=0.20.0
+Requires-Dist: python-pptx>=1.0.0
+Requires-Dist: pandas>=2.0.0
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: black; extra == "dev"