jarvis-ai-assistant 0.1.91__py3-none-any.whl → 0.1.93__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)}\n {str(desc)}" ] = str(desc)
     
     return tasks
 

jarvis/models/ai8.py

@@ -20,7 +20,7 @@ class AI8Model(BasePlatform):
         """Initialize model"""
         super().__init__()
         self.system_message = ""
-        self.conversation = None
+        self.conversation = {}
         self.files = []
         self.models = {}  # 存储模型信息
 
@@ -112,6 +112,7 @@ class AI8Model(BasePlatform):
                 "name": name,
                 "data": f"data:image/png;base64,{base64_data}"
             })
+        return self.files
     
     def set_system_message(self, message: str):
         """Set system message"""
@@ -138,7 +139,7 @@ class AI8Model(BasePlatform):
             
             payload = {
                 "text": message,
-                "sessionId": self.conversation['id'],
+                "sessionId": self.conversation['id'] if self.conversation else None,
                 "files": []
             }
             
@@ -307,6 +308,6 @@ class AI8Model(BasePlatform):
             return list(self.models.keys())
             
         except Exception as e:
-            PrettyOutput.print(f"获取模型列表异常: {str(e)}", OutputType.ERROR)
+            PrettyOutput.print(f"获取模型列表异常: {str(e)}", OutputType.WARNING)
             return []
         

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/models/openai.py

@@ -69,9 +69,9 @@ class OpenAIModel(BasePlatform):
             
             response = self.client.chat.completions.create(
                 model=self.model_name,  # 使用配置的模型名称
-                messages=self.messages,
+                messages=self.messages, # type: ignore
                 stream=True
-            )
+            ) # type: ignore
             
             full_response = ""
             

jarvis/models/oyi.py

@@ -23,7 +23,7 @@ class OyiModel(BasePlatform):
         self.messages = []
         self.system_message = ""
         self.conversation = None
-        self.upload_files = [] 
+        self.files = []
         self.first_chat = True
         
         self.token = os.getenv("OYI_API_KEY")
@@ -122,7 +122,7 @@ class OyiModel(BasePlatform):
             }
             
             payload = {
-                "topicId": self.conversation['result']['id'],
+                "topicId": self.conversation['result']['id'] if self.conversation else None,
                 "messages": self.messages,
                 "content": message,
                 "contentFiles": []
@@ -130,8 +130,8 @@ class OyiModel(BasePlatform):
             
             # 如果有上传的文件，添加到请求中
             if self.first_chat:
-                if self.upload_files:
-                    for file_data in self.upload_files:
+                if self.files:
+                    for file_data in self.files:
                         file_info = {
                             "contentType": 1,  # 1 表示图片
                             "fileUrl": file_data['result']['url'],
@@ -140,7 +140,7 @@ class OyiModel(BasePlatform):
                         }
                         payload["contentFiles"].append(file_info)
                     # 清空已使用的文件列表
-                    self.upload_files = []
+                    self.files = []
                 message = self.system_message + "\n" + message
                 payload["content"] = message
                 self.first_chat = False
@@ -195,7 +195,7 @@ class OyiModel(BasePlatform):
         """Reset model state"""
         self.messages = []
         self.conversation = None
-        self.upload_files = []
+        self.files = []
         self.first_chat = True
             
     def delete_chat(self) -> bool:
@@ -251,7 +251,7 @@ class OyiModel(BasePlatform):
             model_info = self.models.get(self.model_name)
             if not model_info or not model_info.get('uploadFile', False):
                 PrettyOutput.print(f"当前模型 {self.model_name} 不支持文件上传", OutputType.WARNING)
-                return None
+                return []
             
             headers = {
                 'Authorization': f'Bearer {self.token}',
@@ -283,18 +283,18 @@ class OyiModel(BasePlatform):
                     if response.status_code == 200:
                         data = response.json()
                         if data.get('code') == 200:
-                            self.upload_files.append(data)
-                            return data
+                            self.files.append(data)
                         else:
                             PrettyOutput.print(f"文件上传失败: {data.get('message')}", OutputType.ERROR)
-                            return None
+                            return []
                     else:
                         PrettyOutput.print(f"文件上传失败: {response.status_code}", OutputType.ERROR)
-                        return None
+                        return []
                 
+            return self.files
         except Exception as e:
             PrettyOutput.print(f"文件上传异常: {str(e)}", OutputType.ERROR)
-            return None
+            return []
 
     def get_available_models(self) -> List[str]:
         """获取可用的模型列表
@@ -364,5 +364,5 @@ class OyiModel(BasePlatform):
             return sorted(models)
             
         except Exception as e:
-            PrettyOutput.print(f"获取模型列表异常: {str(e)}", OutputType.ERROR)
+            PrettyOutput.print(f"获取模型列表异常: {str(e)}", OutputType.WARNING)
             return []

jarvis/tools/ask_user.py

@@ -34,8 +34,7 @@ class AskUserTool:
             PrettyOutput.print(question, OutputType.SYSTEM)
             
             # 获取用户输入
-            PrettyOutput.print("\n请输入您的回答（输入空行结束）:", OutputType.INPUT)
-            user_response = get_multiline_input()
+            user_response = get_multiline_input("请输入您的回答（输入空行结束）")
             
             if user_response == "__interrupt__":
                 return {

jarvis/tools/coder.py

@@ -0,0 +1,69 @@
+import os
+from typing import Dict, Any, Optional
+from jarvis.jarvis_coder.main import JarvisCoder
+from jarvis.utils import PrettyOutput, OutputType
+
+class CoderTool:
+    """代码修改工具"""
+    
+    name = "coder"
+    description = "分析并修改现有代码，用于实现新功能、修复bug、重构代码等。能理解代码上下文并进行精确的代码编辑。"
+    parameters = {
+        "feature": {
+            "type": "string",
+            "description": "要实现的功能描述或需要修改的内容，例如：'添加日志功能'、'修复内存泄漏'、'优化性能'等",
+            "required": True
+        },
+        "dir": {
+            "type": "string", 
+            "description": "项目根目录，默认为当前目录",
+            "required": False
+        },
+        "language": {
+            "type": "string",
+            "description": "项目的主要编程语言，默认为python",
+            "required": False
+        }
+    }
+
+    def __init__(self):
+        self._coder = None
+
+
+    def _init_coder(self, dir: Optional[str] = None, language: Optional[str] = "python") -> None:
+        """初始化JarvisCoder实例"""
+        if not self._coder:
+            import os
+            work_dir = dir or os.getcwd()
+            self._coder = JarvisCoder(work_dir, language)
+
+    def execute(self, args: Dict) -> Dict[str, Any]:
+        """执行代码修改
+        
+        Args:
+            feature: 要实现的功能描述
+            dir: 可选，项目根目录
+            language: 可选，编程语言
+            
+        Returns:
+            Dict[str, Any]: 执行结果
+        """
+        feature = args.get("feature")
+        dir = args.get("dir")
+        language = args.get("language", "python")
+        
+        try:
+            self.current_dir = os.getcwd()
+            self._init_coder(dir, language)
+            result = self._coder.execute(str(feature)) # type: ignore
+            return result
+        except Exception as e:
+            PrettyOutput.print(f"代码修改失败: {str(e)}", OutputType.ERROR)
+            return {
+                "success": False,
+                "stdout": "",
+                "stderr": f"执行失败: {str(e)}",
+                "error": e
+            } 
+        finally:
+            os.chdir(self.current_dir)