jarvis-ai-assistant 0.1.91__py3-none-any.whl → 0.1.92__py3-none-any.whl

jarvis/jarvis_rag/main.py

@@ -1,13 +1,10 @@
 import os
-import hashlib
 import numpy as np
 import faiss
 from typing import List, Tuple, Optional, Dict
-from sentence_transformers import SentenceTransformer
 import pickle
-from jarvis.utils import OutputType, PrettyOutput, find_git_root, get_max_context_length, load_embedding_model, load_rerank_model
+from jarvis.utils import OutputType, PrettyOutput, get_file_md5, get_max_context_length, load_embedding_model, load_rerank_model
 from jarvis.utils import load_env_from_file
-import tiktoken
 from dataclasses import dataclass
 from tqdm import tqdm
 import fitz  # PyMuPDF for PDF files
@@ -16,12 +13,16 @@ from pathlib import Path
 from jarvis.models.registry import PlatformRegistry
 import shutil
 from datetime import datetime
+import lzma  # 添加 lzma 导入
+from concurrent.futures import ThreadPoolExecutor
+from threading import Lock
 
 @dataclass
 class Document:
     """文档类，用于存储文档内容和元数据"""
     content: str  # 文档内容
     metadata: Dict  # 元数据(文件路径、位置等)
+    md5: str = ""  # 文件MD5值，用于增量更新检测
 
 class FileProcessor:
     """文件处理器基类"""
@@ -163,7 +164,9 @@ class RAGTool:
         # 初始化缓存和索引
         self.cache_path = os.path.join(self.data_dir, "cache.pkl")
         self.documents: List[Document] = []
-        self.index = None
+        self.index = None  # 用于搜索的IVF索引
+        self.flat_index = None  # 用于存储原始向量
+        self.file_md5_cache = {}  # 用于存储文件的MD5值
         
         # 加载缓存
         self._load_cache()
@@ -175,17 +178,23 @@ class RAGTool:
             DocxProcessor()
         ]
 
+        # 添加线程相关配置
+        self.thread_count = int(os.environ.get("JARVIS_THREAD_COUNT", os.cpu_count() or 4))
+        self.vector_lock = Lock()  # 用于保护向量列表的并发访问
+
     def _load_cache(self):
         """加载缓存数据"""
         if os.path.exists(self.cache_path):
             try:
-                with open(self.cache_path, 'rb') as f:
+                with lzma.open(self.cache_path, 'rb') as f:
                     cache_data = pickle.load(f)
                     self.documents = cache_data["documents"]
                     vectors = cache_data["vectors"]
+                    self.file_md5_cache = cache_data.get("file_md5_cache", {})  # 加载MD5缓存
                     
                 # 重建索引
-                self._build_index(vectors)
+                if vectors is not None:
+                    self._build_index(vectors)
                 PrettyOutput.print(f"加载了 {len(self.documents)} 个文档片段", 
                                 output_type=OutputType.INFO)
             except Exception as e:
@@ -193,16 +202,18 @@ class RAGTool:
                                 output_type=OutputType.WARNING)
                 self.documents = []
                 self.index = None
+                self.flat_index = None
+                self.file_md5_cache = {}
 
     def _save_cache(self, vectors: np.ndarray):
         """优化缓存保存"""
         try:
-            # 添加版本号和时间戳
             cache_data = {
                 "version": "1.0",
                 "timestamp": datetime.now().isoformat(),
                 "documents": self.documents,
-                "vectors": vectors,
+                "vectors": vectors.copy() if vectors is not None else None,  # 创建数组的副本
+                "file_md5_cache": dict(self.file_md5_cache),  # 创建字典的副本
                 "metadata": {
                     "vector_dim": self.vector_dim,
                     "total_docs": len(self.documents),
@@ -210,9 +221,12 @@ class RAGTool:
                 }
             }
             
-            # 使用压缩存储
-            with open(self.cache_path, 'wb') as f:
-                pickle.dump(cache_data, f, protocol=pickle.HIGHEST_PROTOCOL)
+            # 先将数据序列化为字节流
+            data = pickle.dumps(cache_data, protocol=pickle.HIGHEST_PROTOCOL)
+            
+            # 然后使用 LZMA 压缩字节流
+            with lzma.open(self.cache_path, 'wb') as f:
+                f.write(data)
             
             # 创建备份
             backup_path = f"{self.cache_path}.backup"
@@ -223,22 +237,29 @@ class RAGTool:
         except Exception as e:
             PrettyOutput.print(f"保存缓存失败: {str(e)}", 
                             output_type=OutputType.ERROR)
+            raise
 
     def _build_index(self, vectors: np.ndarray):
         """构建FAISS索引"""
-        # 添加IVF索引以提高大规模检索性能
+        if vectors.shape[0] == 0:
+            self.index = None
+            self.flat_index = None
+            return
+            
+        # 创建扁平索引存储原始向量，用于重建
+        self.flat_index = faiss.IndexFlatIP(self.vector_dim)
+        self.flat_index.add(vectors)
+        
+        # 创建IVF索引用于快速搜索
         nlist = max(4, int(vectors.shape[0] / 1000))  # 每1000个向量一个聚类中心
         quantizer = faiss.IndexFlatIP(self.vector_dim)
         self.index = faiss.IndexIVFFlat(quantizer, self.vector_dim, nlist, faiss.METRIC_INNER_PRODUCT)
         
-        if vectors.shape[0] > 0:
-            # 训练IVF索引
-            self.index.train(vectors)
-            self.index.add(vectors)
-            # 设置搜索时探测的聚类数
-            self.index.nprobe = min(nlist, 10)
-        else:
-            self.index = None
+        # 训练并添加向量
+        self.index.train(vectors)
+        self.index.add(vectors)
+        # 设置搜索时探测的聚类数
+        self.index.nprobe = min(nlist, 10)
 
     def _split_text(self, text: str) -> List[str]:
         """使用更智能的分块策略"""
@@ -302,16 +323,58 @@ class RAGTool:
                                             show_progress_bar=False)
         return np.array(embedding, dtype=np.float32)
 
-    def _process_file(self, file_path: str) -> List[Document]:
-        """处理单个文件
+    def _get_embedding_batch(self, texts: List[str]) -> np.ndarray:
+        """批量获取文本的向量表示
         
         Args:
-            file_path: 文件路径
+            texts: 文本列表
             
         Returns:
-            文档对象列表
+            np.ndarray: 向量表示数组
         """
         try:
+            embeddings = self.embedding_model.encode(texts, 
+                                                normalize_embeddings=True,
+                                                show_progress_bar=False,
+                                                batch_size=32)  # 使用批处理提高效率
+            return np.array(embeddings, dtype=np.float32)
+        except Exception as e:
+            PrettyOutput.print(f"获取向量表示失败: {str(e)}", 
+                            output_type=OutputType.ERROR)
+            return np.zeros((len(texts), self.vector_dim), dtype=np.float32)
+
+    def _process_document_batch(self, documents: List[Document]) -> List[np.ndarray]:
+        """处理一批文档的向量化
+        
+        Args:
+            documents: 文档列表
+            
+        Returns:
+            List[np.ndarray]: 向量列表
+        """
+        texts = []
+        for doc in documents:
+            # 组合文档信息
+            combined_text = f"""
+文件: {doc.metadata['file_path']}
+内容: {doc.content}
+"""
+            texts.append(combined_text)
+            
+        return self._get_embedding_batch(texts)
+
+    def _process_file(self, file_path: str) -> List[Document]:
+        """处理单个文件"""
+        try:
+            # 计算文件MD5
+            current_md5 = get_file_md5(file_path)
+            if not current_md5:
+                return []
+
+            # 检查文件是否需要重新处理
+            if file_path in self.file_md5_cache and self.file_md5_cache[file_path] == current_md5:
+                return []
+
             # 查找合适的处理器
             processor = None
             for p in self.file_processors:
@@ -320,18 +383,14 @@ class RAGTool:
                     break
                     
             if not processor:
-                PrettyOutput.print(f"跳过不支持的文件: {file_path}", 
-                                output_type=OutputType.WARNING)
+                # 如果找不到合适的处理器，则返回一个空的文档
                 return []
             
             # 提取文本内容
             content = processor.extract_text(file_path)
             if not content.strip():
-                PrettyOutput.print(f"文件内容为空: {file_path}", 
-                                output_type=OutputType.WARNING)
                 return []
             
-                
             # 分割文本
             chunks = self._split_text(content)
             
@@ -345,10 +404,13 @@ class RAGTool:
                         "file_type": Path(file_path).suffix.lower(),
                         "chunk_index": i,
                         "total_chunks": len(chunks)
-                    }
+                    },
+                    md5=current_md5
                 )
                 documents.append(doc)
-                
+            
+            # 更新MD5缓存
+            self.file_md5_cache[file_path] = current_md5
             return documents
             
         except Exception as e:
@@ -361,43 +423,117 @@ class RAGTool:
         # 获取所有文件
         all_files = []
         for root, _, files in os.walk(dir):
-            # 忽略特定目录
             if any(ignored in root for ignored in ['.git', '__pycache__', 'node_modules']) or \
                any(part.startswith('.jarvis-') for part in root.split(os.sep)):
                 continue
             for file in files:
-                # 忽略 .jarvis- 开头的文件
                 if file.startswith('.jarvis-'):
                     continue
                     
                 file_path = os.path.join(root, file)
-                # 跳过大文件
                 if os.path.getsize(file_path) > 100 * 1024 * 1024:  # 100MB
                     PrettyOutput.print(f"跳过大文件: {file_path}", 
                                     output_type=OutputType.WARNING)
                     continue
                 all_files.append(file_path)
 
-        # 处理所有文件
-        self.documents = []
-        for file_path in tqdm(all_files, desc="处理文件"):
-            docs = self._process_file(file_path)
-            self.documents.extend(docs)
+        # 清理已删除文件的缓存
+        deleted_files = set(self.file_md5_cache.keys()) - set(all_files)
+        for file_path in deleted_files:
+            del self.file_md5_cache[file_path]
+            # 移除相关的文档
+            self.documents = [doc for doc in self.documents if doc.metadata['file_path'] != file_path]
 
-        # 获取所有文档的向量表示
-        vectors = []
-        for doc in tqdm(self.documents, desc="生成向量"):
-            vector = self._get_embedding(doc.content)
-            vectors.append(vector)
+        # 检查文件变化
+        files_to_process = []
+        unchanged_files = []
+        
+        with tqdm(total=len(all_files), desc="检查文件状态") as pbar:
+            for file_path in all_files:
+                current_md5 = get_file_md5(file_path)
+                if current_md5:  # 只处理能成功计算MD5的文件
+                    if file_path in self.file_md5_cache and self.file_md5_cache[file_path] == current_md5:
+                        # 文件未变化，记录但不重新处理
+                        unchanged_files.append(file_path)
+                    else:
+                        # 新文件或已修改的文件
+                        files_to_process.append(file_path)
+                pbar.update(1)
+
+        # 保留未变化文件的文档
+        unchanged_documents = [doc for doc in self.documents 
+                            if doc.metadata['file_path'] in unchanged_files]
+
+        # 处理新文件和修改的文件
+        new_documents = []
+        if files_to_process:
+            with tqdm(total=len(files_to_process), desc="处理文件") as pbar:
+                for file_path in files_to_process:
+                    try:
+                        docs = self._process_file(file_path)
+                        if len(docs) > 0:
+                            new_documents.extend(docs)
+                    except Exception as e:
+                        PrettyOutput.print(f"处理文件失败 {file_path}: {str(e)}", 
+                                        output_type=OutputType.ERROR)
+                    pbar.update(1)
+
+        # 更新文档列表
+        self.documents = unchanged_documents + new_documents
+
+        if not self.documents:
+            PrettyOutput.print("没有需要处理的文档", output_type=OutputType.WARNING)
+            return
+
+        # 只对新文档进行向量化
+        if new_documents:
+            PrettyOutput.print(f"开始处理 {len(new_documents)} 个新文档", 
+                            output_type=OutputType.INFO)
+            
+            # 使用线程池并发处理向量化
+            batch_size = 32
+            new_vectors = []
+            
+            with tqdm(total=len(new_documents), desc="生成向量") as pbar:
+                with ThreadPoolExecutor(max_workers=self.thread_count) as executor:
+                    for i in range(0, len(new_documents), batch_size):
+                        batch = new_documents[i:i + batch_size]
+                        future = executor.submit(self._process_document_batch, batch)
+                        batch_vectors = future.result()
+                        
+                        with self.vector_lock:
+                            new_vectors.extend(batch_vectors)
+                        
+                        pbar.update(len(batch))
+
+            # 合并新旧向量
+            if self.flat_index is not None:
+                # 获取未变化文档的向量
+                unchanged_vectors = []
+                for doc in unchanged_documents:
+                    # 从现有索引中提取向量
+                    doc_idx = next((i for i, d in enumerate(self.documents) 
+                                if d.metadata['file_path'] == doc.metadata['file_path']), None)
+                    if doc_idx is not None:
+                        # 从扁平索引中重建向量
+                        vector = np.zeros((1, self.vector_dim), dtype=np.float32)
+                        self.flat_index.reconstruct(doc_idx, vector.ravel())
+                        unchanged_vectors.append(vector)
+                
+                if unchanged_vectors:
+                    unchanged_vectors = np.vstack(unchanged_vectors)
+                    vectors = np.vstack([unchanged_vectors, np.vstack(new_vectors)])
+                else:
+                    vectors = np.vstack(new_vectors)
+            else:
+                vectors = np.vstack(new_vectors)
 
-        if vectors:
-            vectors = np.vstack(vectors)
             # 构建索引
             self._build_index(vectors)
             # 保存缓存
             self._save_cache(vectors)
-            
-        PrettyOutput.print(f"成功索引了 {len(self.documents)} 个文档片段", 
+        
+        PrettyOutput.print(f"成功索引了 {len(self.documents)} 个文档片段 (新增/修改: {len(new_documents)}, 未变化: {len(unchanged_documents)})", 
                         output_type=OutputType.SUCCESS)
 
     def search(self, query: str, top_k: int = 30) -> List[Tuple[Document, float]]:

jarvis/jarvis_smart_shell/main.py

@@ -4,6 +4,8 @@ import os
 import sys
 import readline
 from typing import Optional
+from yaspin import yaspin
+from yaspin.spinners import Spinners
 
 from jarvis.models.registry import PlatformRegistry
 from jarvis.utils import PrettyOutput, OutputType, load_env_from_file
@@ -11,7 +13,7 @@ from jarvis.utils import PrettyOutput, OutputType, load_env_from_file
 def execute_command(command: str) -> None:
     """显示命令并允许用户编辑，回车执行，Ctrl+C取消"""
     try:
-        print("生成的命令 (可以编辑，回车执行，Ctrl+C取消):")
+        print("\n生成的命令 (可以编辑，回车执行，Ctrl+C取消):")
         # 预填充输入行
         readline.set_startup_hook(lambda: readline.insert_text(command))
         try:
@@ -68,14 +70,19 @@ find . -name "*.py"
         prefix = f"当前路径: {current_path}\n"
         prefix += f"当前shell: {shell}\n"
         
-        # 处理请求
-        result = model.chat(prefix + request)
-        
-        # 提取命令 - 简化处理逻辑，因为现在应该只返回纯命令
-        if result and isinstance(result, str):
-            return result.strip()
-        
-        return None
+        # 使用yaspin显示Thinking状态
+        with yaspin(Spinners.dots, text="Thinking", color="yellow") as spinner:
+            # 处理请求
+            result = model.chat(prefix + request)
+            
+            # 提取命令
+            if result and isinstance(result, str):
+                command = result.strip()
+                spinner.ok("✓")
+                return command
+            
+            spinner.fail("✗")
+            return None
         
     except Exception as e:
         PrettyOutput.print(f"处理请求时发生错误: {str(e)}", OutputType.ERROR)

jarvis/main.py

@@ -54,6 +54,15 @@ def load_tasks() -> dict:
                 PrettyOutput.print("Warning: .jarvis file should contain a dictionary of task_name: task_description", OutputType.ERROR)
         except Exception as e:
             PrettyOutput.print(f"Error loading .jarvis file: {str(e)}", OutputType.ERROR)
+
+    # 读取方法论
+    method_path = os.path.expanduser("~/.jarvis_methodology")
+    if os.path.exists(method_path):
+        with open(method_path, "r", encoding="utf-8") as f:
+            methodology = yaml.safe_load(f)
+        if isinstance(methodology, dict):
+            for name, desc in methodology.items():
+                tasks[f"执行方法论：{str(name)}"] = str(desc)
     
     return tasks
 

jarvis/models/ollama.py

@@ -29,15 +29,15 @@ class OllamaPlatform(BasePlatform):
                 PrettyOutput.print("1. 安装 Ollama: https://ollama.ai", OutputType.INFO)
                 PrettyOutput.print("2. 下载模型:", OutputType.INFO)
                 PrettyOutput.print(f"   ollama pull {self.model_name}", OutputType.INFO)
-                raise Exception("No available models found")
+                PrettyOutput.print("Ollama没有可用的模型", OutputType.WARNING)
                 
         except requests.exceptions.ConnectionError:
-            PrettyOutput.print("\nOllama 服务未启动或无法连接", OutputType.ERROR)
+            PrettyOutput.print("\nOllama 服务未启动或无法连接", OutputType.WARNING)
             PrettyOutput.print("请确保已经：", OutputType.INFO)
             PrettyOutput.print("1. 安装了 Ollama: https://ollama.ai", OutputType.INFO)
             PrettyOutput.print("2. 启动了 Ollama 服务", OutputType.INFO)
             PrettyOutput.print("3. 服务地址配置正确 (默认: http://localhost:11434)", OutputType.INFO)
-            raise Exception("Ollama service is not available")
+            
             
         self.messages = []
         self.system_message = ""

jarvis/tools/thinker.py

@@ -17,11 +17,10 @@ class ThinkerTool:
                 "description": "问题相关的上下文信息或背景知识",
                 "default": ""
             },
-            "approach": {
+            "goal": {
                 "type": "string",
-                "enum": ["chain_of_thought", "tree_of_thought", "step_by_step"],
-                "description": "思考方式：chain_of_thought(思维链)、tree_of_thought(思维树)、step_by_step(步骤分解)",
-                "default": "chain_of_thought"
+                "description": "期望达成的具体目标或结果",
+                "default": ""
             }
         },
         "required": ["question"]
@@ -31,89 +30,36 @@ class ThinkerTool:
         """初始化思考工具"""
         self.model = PlatformRegistry.get_global_platform_registry().get_thinking_platform()
 
-    def _generate_prompt(self, question: str, context: str, approach: str) -> str:
+    def _generate_prompt(self, question: str, context: str, goal: str) -> str:
         """生成提示词
         
         Args:
             question: 问题
             context: 上下文
-            approach: 思考方式
+            goal: 期望目标
             
         Returns:
             str: 完整的提示词
         """
         # 基础提示词
-        base_prompt = "你是一个擅长深度思考和逻辑推理的助手。"
-        
-        # 根据不同的思考方式添加具体指导
-        approach_prompts = {
-            "chain_of_thought": """请使用思维链方式分析问题：
-1. 仔细阅读问题和上下文
-2. 逐步推理，每一步都要说明推理依据
-3. 考虑多个可能的角度
-4. 得出最终结论
-
-请按以下格式输出：
-思考过程：
-1. [第一步推理]
-2. [第二步推理]
-...
-
-结论：
-[最终结论]""",
-
-            "tree_of_thought": """请使用思维树方式分析问题：
-1. 将问题分解为多个子问题
-2. 对每个子问题进行分支探索
-3. 评估每个分支的可行性
-4. 整合最优路径
-
-请按以下格式输出：
-问题分解：
-- 子问题1
-  - 分支1.1
-  - 分支1.2
-- 子问题2
-  - 分支2.1
-  - 分支2.2
-
-分析过程：
-[详细分析每个分支]
-
-最优路径：
-[说明选择原因]
-
-结论：
-[最终结论]""",
+        prompt = f"""你是一个擅长深度思考和逻辑推理的助手。请帮助分析问题并给出解决方案。
 
-            "step_by_step": """请使用步骤分解方式分析问题：
-1. 将问题分解为具体步骤
-2. 详细说明每个步骤的执行方法
-3. 考虑每个步骤可能的问题
-4. 提供完整的解决方案
-
-请按以下格式输出：
-步骤分解：
-步骤1: [具体内容]
-步骤2: [具体内容]
-...
-
-执行分析：
-[详细分析每个步骤]
-
-解决方案：
-[完整方案]"""
-        }
-
-        # 构建完整提示词
-        prompt = f"""{base_prompt}
-
-{approach_prompts[approach]}
+请按以下方式思考：
+1. 仔细理解问题和目标
+2. 进行系统性分析和推理
+3. 考虑多个可能的解决方案
+4. 给出最佳建议和具体行动步骤
 
 问题：
 {question}
-
 """
+        # 如果有目标，添加到提示词中
+        if goal:
+            prompt += f"""
+期望目标：
+{goal}
+"""
+
         # 如果有上下文，添加到提示词中
         if context:
             prompt += f"""
@@ -131,7 +77,7 @@ class ThinkerTool:
             args: 包含参数的字典
                 - question: 问题
                 - context: 上下文（可选）
-                - approach: 思考方式（可选）
+                - goal: 期望目标（可选）
                 
         Returns:
             Dict[str, Any]: 执行结果
@@ -140,16 +86,17 @@ class ThinkerTool:
             # 获取参数
             question = args["question"]
             context = args.get("context", "")
-            approach = args.get("approach", "chain_of_thought")
+            goal = args.get("goal", "")
             
             # 生成提示词
-            prompt = self._generate_prompt(question, context, approach)
+            prompt = self._generate_prompt(question, context, goal)
             
             # 记录开始分析
             PrettyOutput.print(f"开始分析问题: {question}", OutputType.INFO)
             if context:
                 PrettyOutput.print("包含上下文信息", OutputType.INFO)
-            PrettyOutput.print(f"使用{approach}方式思考", OutputType.INFO)
+            if goal:
+                PrettyOutput.print(f"目标: {goal}", OutputType.INFO)
             
             # 调用模型进行分析
             response = self.model.chat(prompt)
@@ -182,15 +129,14 @@ def main():
     parser = argparse.ArgumentParser(description='深度思考分析工具')
     parser.add_argument('--question', required=True, help='需要分析的问题')
     parser.add_argument('--context', help='问题相关的上下文信息')
-    parser.add_argument('--approach', choices=['chain_of_thought', 'tree_of_thought', 'step_by_step'],
-                      default='chain_of_thought', help='思考方式')
+    parser.add_argument('--goal', help='期望达成的具体目标或结果')
     args = parser.parse_args()
     
     tool = ThinkerTool()
     result = tool.execute({
         "question": args.question,
         "context": args.context,
-        "approach": args.approach
+        "goal": args.goal
     })
     
     if result["success"]:

jarvis/utils.py

@@ -1,3 +1,4 @@
+import hashlib
 from pathlib import Path
 import sys
 import time
@@ -276,4 +277,7 @@ def get_max_context_length():
     return int(os.getenv('JARVIS_MAX_CONTEXT_LENGTH', '131072'))  # 默认128k
 
 def get_thread_count():
-    return int(os.getenv('JARVIS_THREAD_COUNT', '1'))  
+    return int(os.getenv('JARVIS_THREAD_COUNT', '1'))  
+
+def get_file_md5(filepath: str)->str:    
+    return hashlib.md5(open(filepath, "rb").read(100*1024*1024)).hexdigest()

{jarvis_ai_assistant-0.1.91.dist-info → jarvis_ai_assistant-0.1.92.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: jarvis-ai-assistant
-Version: 0.1.91
+Version: 0.1.92
 Summary: Jarvis: An AI assistant that uses tools to interact with the system
 Home-page: https://github.com/skyfireitdiy/Jarvis
 Author: skyfire
@@ -53,6 +53,8 @@ Requires-Dist: PyMuPDF>=1.21.0
 Requires-Dist: python-docx>=0.8.11
 Requires-Dist: tiktoken>=0.3.0
 Requires-Dist: tqdm>=4.65.0
+Requires-Dist: docx>=0.2.4
+Requires-Dist: yaspin>=2.5.0
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: black; extra == "dev"