jarvis-ai-assistant 0.1.66__py3-none-any.whl → 0.1.72__py3-none-any.whl

jarvis/__init__.py

@@ -1,3 +1,3 @@
 """Jarvis AI Assistant"""
 
-__version__ = "0.1.66" 
+__version__ = "0.1.72" 

jarvis/jarvis_codebase/main.py

@@ -1,7 +1,5 @@
 import hashlib
 import os
-import sqlite3
-import time
 import numpy as np
 import faiss
 from typing import List, Tuple, Optional
@@ -13,24 +11,29 @@ from jarvis.utils import OutputType, PrettyOutput, find_git_root
 from jarvis.utils import load_env_from_file
 import argparse
 from sentence_transformers import SentenceTransformer
+import pickle
 
 class CodeBase:
-    def __init__(self, root_dir: str, thread_count: int = 10):
+    def __init__(self, root_dir: str):
         load_env_from_file()
         self.root_dir = root_dir
         os.chdir(self.root_dir)
-        self.thread_count = thread_count
+        self.thread_count = int(os.environ.get("JARVIS_THREAD_COUNT") or 10)
         self.cheap_platform = os.environ.get("JARVIS_CHEAP_PLATFORM") or os.environ.get("JARVIS_PLATFORM") or "kimi"
         self.cheap_model = os.environ.get("JARVIS_CHEAP_MODEL") or os.environ.get("JARVIS_MODEL") or "kimi"
         self.normal_platform = os.environ.get("JARVIS_PLATFORM") or "kimi"
+        self.codegen_platform = os.environ.get("JARVIS_CODEGEN_PLATFORM") or os.environ.get("JARVIS_PLATFORM") or "kimi"
+        self.codegen_model = os.environ.get("JARVIS_CODEGEN_MODEL") or os.environ.get("JARVIS_MODEL") or "kimi"
         self.normal_model = os.environ.get("JARVIS_MODEL") or "kimi"
         self.embedding_model_name = os.environ.get("JARVIS_EMBEDDING_MODEL") or "BAAI/bge-large-zh-v1.5"
-        if not self.cheap_platform or not self.cheap_model or not self.embedding_model_name or not self.normal_platform or not self.normal_model:
-            raise ValueError("JARVIS_CHEAP_PLATFORM or JARVIS_CHEAP_MODEL or JARVIS_EMBEDDING_MODEL or JARVIS_PLATFORM or JARVIS_MODEL is not set")
+        if not self.cheap_platform or not self.cheap_model or not self.codegen_platform or not self.codegen_model or not self.embedding_model_name or not self.normal_platform or not self.normal_model:
+            raise ValueError("JARVIS_CHEAP_PLATFORM or JARVIS_CHEAP_MODEL or JARVIS_CODEGEN_PLATFORM or JARVIS_CODEGEN_MODEL or JARVIS_EMBEDDING_MODEL or JARVIS_PLATFORM or JARVIS_MODEL is not set")
         
         PrettyOutput.print(f"廉价模型使用平台: {self.cheap_platform} 模型: {self.cheap_model}", output_type=OutputType.INFO)
+        PrettyOutput.print(f"代码生成模型使用平台: {self.codegen_platform} 模型: {self.codegen_model}", output_type=OutputType.INFO)
         PrettyOutput.print(f"分析模型使用平台: {self.normal_platform} 模型: {self.normal_model}", output_type=OutputType.INFO)
         PrettyOutput.print(f"嵌入模型: {self.embedding_model_name}", output_type=OutputType.INFO)
+        PrettyOutput.print(f"索引建立线程数: {self.thread_count}", output_type=OutputType.INFO)
         PrettyOutput.print(f"检索算法：分层导航小世界算法", output_type=OutputType.INFO)
             
         # 初始化数据目录
@@ -40,6 +43,7 @@ class CodeBase:
             
         # 初始化嵌入模型，使用系统默认缓存目录
         try:
+            os.environ["TOKENIZERS_PARALLELISM"] = "false"
             PrettyOutput.print("正在加载/下载模型，请稍候...", output_type=OutputType.INFO)
             self.embedding_model = SentenceTransformer(self.embedding_model_name)
             
@@ -59,59 +63,37 @@ class CodeBase:
             
         self.vector_dim = self.embedding_model.get_sentence_embedding_dimension()
 
-            
-        self.db_path = os.path.join(self.data_dir, "codebase.db")
-        if not os.path.exists(self.db_path):
-            self.create_db()
         self.git_file_list = self.get_git_file_list()
         self.platform_registry = PlatformRegistry().get_global_platform_registry()
-        self.index_path = os.path.join(self.data_dir, "vectors.index")
-        self.index = None
-        if  os.path.exists(self.index_path):
-            PrettyOutput.print("正在加载向量数据库", output_type=OutputType.INFO)
-            self.index = faiss.read_index(self.index_path)
+        
+        # 初始化缓存和索引
+        self.cache_path = os.path.join(self.data_dir, "cache.pkl")
+        self.vector_cache = {}
+        self.file_paths = []
+        
+        # 加载缓存
+        if os.path.exists(self.cache_path):
+            try:
+                with open(self.cache_path, 'rb') as f:
+                    cache_data = pickle.load(f)
+                    self.vector_cache = cache_data["vectors"]
+                    self.file_paths = cache_data["file_paths"]
+                PrettyOutput.print(f"加载了 {len(self.vector_cache)} 个向量缓存", 
+                                 output_type=OutputType.INFO)
+                # 从缓存重建索引
+                self.build_index()
+            except Exception as e:
+                PrettyOutput.print(f"加载缓存失败: {str(e)}", 
+                                 output_type=OutputType.WARNING)
+                self.vector_cache = {}
+                self.file_paths = []
+                self.index = None
 
     def get_git_file_list(self):
-        return os.popen("git ls-files").read().splitlines()
-
-    def get_db_connection(self):
-        """创建并返回一个新的数据库连接"""
-        return sqlite3.connect(self.db_path)
-
-    def clean_db(self) -> bool:
-        """清理数据库和向量索引中的过期记录"""
-        db = self.get_db_connection()
-        try:
-            # 获取所有数据库记录
-            all_records = db.execute("SELECT path FROM codebase").fetchall()
-            files_to_delete = []
-            
-            # 找出需要删除的文件
-            for row in all_records:
-                if row[0] not in self.git_file_list:
-                    files_to_delete.append(row[0])
-            
-            if not files_to_delete:
-                return False
-                
-            for file_path in files_to_delete:
-                db.execute("DELETE FROM codebase WHERE path = ?", (file_path,))
-            
-            db.commit()
-            
-            PrettyOutput.print(f"清理了 {len(files_to_delete)} 个文件的记录", 
-                             output_type=OutputType.INFO)
-            return True
-        finally:
-            db.close()
-        
-    def create_db(self):
-        db = self.get_db_connection()
-        try:
-            db.execute("CREATE TABLE IF NOT EXISTS codebase (path TEXT, md5 TEXT ,description TEXT)")
-            db.commit()
-        finally:
-            db.close()
+        """获取 git 仓库中的文件列表，排除 .jarvis-codebase 目录"""
+        files = os.popen("git ls-files").read().splitlines()
+        # 过滤掉 .jarvis-codebase 目录下的文件
+        return [f for f in files if not f.startswith(".jarvis-codebase/")]
 
     def is_text_file(self, file_path: str):
         with open(file_path, "r", encoding="utf-8") as f:
@@ -144,6 +126,74 @@ class CodeBase:
         response = model.chat(prompt)
         return response
 
+    def save_cache(self):
+        """保存缓存数据"""
+        try:
+            cache_data = {
+                "vectors": self.vector_cache,
+                "file_paths": self.file_paths
+            }
+            with open(self.cache_path, 'wb') as f:
+                pickle.dump(cache_data, f)
+            PrettyOutput.print(f"保存了 {len(self.vector_cache)} 个向量缓存", 
+                             output_type=OutputType.INFO)
+        except Exception as e:
+            PrettyOutput.print(f"保存缓存失败: {str(e)}", 
+                             output_type=OutputType.ERROR)
+
+    def get_cached_vector(self, file_path: str, description: str) -> Optional[np.ndarray]:
+        """从缓存获取文件的向量表示"""
+        if file_path not in self.vector_cache:
+            return None
+        
+        # 检查文件是否被修改
+        try:
+            with open(file_path, "rb") as f:
+                current_md5 = hashlib.md5(f.read()).hexdigest()
+        except Exception as e:
+            PrettyOutput.print(f"计算文件MD5失败 {file_path}: {str(e)}", 
+                              output_type=OutputType.ERROR)
+            return None
+        
+        cached_data = self.vector_cache[file_path]
+        if cached_data["md5"] != current_md5:
+            return None
+        
+        # 检查描述是否变化
+        if cached_data["description"] != description:
+            return None
+        
+        return cached_data["vector"]
+
+    def cache_vector(self, file_path: str, vector: np.ndarray, description: str):
+        """缓存文件的向量表示"""
+        try:
+            with open(file_path, "rb") as f:
+                file_md5 = hashlib.md5(f.read()).hexdigest()
+        except Exception as e:
+            PrettyOutput.print(f"计算文件MD5失败 {file_path}: {str(e)}", 
+                              output_type=OutputType.ERROR)
+            file_md5 = ""
+        
+        self.vector_cache[file_path] = {
+            "path": file_path,  # 保存文件路径
+            "md5": file_md5,    # 保存文件MD5
+            "description": description,  # 保存文件描述
+            "vector": vector    # 保存向量
+        }
+        
+        # 保存缓存到文件
+        try:
+            with open(self.cache_path, 'wb') as f:
+                cache_data = {
+                    "vectors": self.vector_cache,
+                    "file_paths": self.file_paths
+                }
+                pickle.dump(cache_data, f)
+        except Exception as e:
+            PrettyOutput.print(f"保存向量缓存失败: {str(e)}", 
+                              output_type=OutputType.ERROR)
+
     def get_embedding(self, text: str) -> np.ndarray:
         """使用 transformers 模型获取文本的向量表示"""
         # 对长文本进行截断
@@ -152,163 +202,291 @@ class CodeBase:
         
         # 获取嵌入向量
         embedding = self.embedding_model.encode(text, 
-                                             normalize_embeddings=True,  # L2归一化
-                                             show_progress_bar=False)
-        return np.array(embedding, dtype=np.float32)
+                                                 normalize_embeddings=True,  # L2归一化
+                                                 show_progress_bar=False)
+        vector = np.array(embedding, dtype=np.float32)
+        return vector
 
     def vectorize_file(self, file_path: str, description: str) -> np.ndarray:
         """将文件内容和描述向量化"""
-        try:            
+        try:
+            # 先尝试从缓存获取
+            cached_vector = self.get_cached_vector(file_path, description)
+            if cached_vector is not None:
+                return cached_vector
+                
             # 组合文件信息
             combined_text = f"""
 文件路径: {file_path}
 文件描述: {description}
 """
-            return self.get_embedding(combined_text)
+            vector = self.get_embedding(combined_text)
+            
+            # 保存到缓存
+            self.cache_vector(file_path, vector, description)
+            return vector
         except Exception as e:
             PrettyOutput.print(f"Error vectorizing file {file_path}: {str(e)}", 
                              output_type=OutputType.ERROR)
             return np.zeros(self.vector_dim, dtype=np.float32)
 
-    def process_file(self, file):
-        """处理单个文件的辅助方法"""
-        db = self.get_db_connection()
+    def clean_cache(self) -> bool:
+        """清理过期的缓存记录，返回是否有文件被删除"""
+        files_to_delete = []
+        for file_path in list(self.vector_cache.keys()):
+            if file_path not in self.git_file_list:
+                del self.vector_cache[file_path]
+                files_to_delete.append(file_path)
+        
+        if files_to_delete:
+            self.save_cache()
+            PrettyOutput.print(f"清理了 {len(files_to_delete)} 个文件的缓存", 
+                             output_type=OutputType.INFO)
+            return True
+        return False
+
+    def process_file(self, file_path: str):
+        """处理单个文件"""
         try:
-            if not self.is_text_file(file):
+            # 跳过不存在的文件
+            if not os.path.exists(file_path):
                 return None
-            md5 = hashlib.md5(open(file, "rb").read()).hexdigest()
-            if db.execute("SELECT path FROM codebase WHERE md5 = ?", (md5,)).fetchone():
+                
+            if not self.is_text_file(file_path):
                 return None
-            description = self.make_description(file)
-            return (file, md5, description)
-        finally:
-            db.close()
+                
+            md5 = hashlib.md5(open(file_path, "rb").read()).hexdigest()
+            
+            # 检查文件是否已经处理过且内容未变
+            if file_path in self.vector_cache:
+                if self.vector_cache[file_path].get("md5") == md5:
+                    return None
+                    
+            description = self.make_description(file_path)
+            vector = self.vectorize_file(file_path, description)
+            
+            # 保存到缓存，使用实际文件路径作为键
+            self.vector_cache[file_path] = {
+                "vector": vector,
+                "description": description,
+                "md5": md5
+            }
+            
+            return file_path
+            
+        except Exception as e:
+            PrettyOutput.print(f"处理文件失败 {file_path}: {str(e)}", 
+                             output_type=OutputType.ERROR,
+                             traceback=True)
+            return None
 
-    def gen_vector_db_from_sqlite(self):
-        self.index = faiss.IndexHNSWFlat(self.vector_dim, 16)
-        self.index.hnsw.efConstruction = 40
-        self.index.hnsw.efSearch = 16
-        db = self.get_db_connection()
-        try:
-            all_records = db.execute("SELECT path, description FROM codebase").fetchall()
-            for row in all_records:
-                file, description = row
-                PrettyOutput.print(f"正在向量化文件: {file}", output_type=OutputType.INFO)
-                vector = self.vectorize_file(file, description)
-                vector = vector.reshape(1, -1)  
-                self.index.add(vector)
-            faiss.write_index(self.index, self.index_path)
-        finally:
-            db.close()
+    def build_index(self):
+        """从向量缓存构建 faiss 索引"""
+        # 创建底层 HNSW 索引
+        hnsw_index = faiss.IndexHNSWFlat(self.vector_dim, 16)
+        hnsw_index.hnsw.efConstruction = 40
+        hnsw_index.hnsw.efSearch = 16
+        
+        # 用 IndexIDMap 包装 HNSW 索引
+        self.index = faiss.IndexIDMap(hnsw_index)
+        
+        vectors = []
+        ids = []
+        self.file_paths = []  # 重置文件路径列表
+        
+        for i, (file_path, data) in enumerate(self.vector_cache.items()):
+            vectors.append(data["vector"].reshape(1, -1))
+            ids.append(i)
+            self.file_paths.append(file_path)
+            
+        if vectors:
+            vectors = np.vstack(vectors)
+            self.index.add_with_ids(vectors, np.array(ids))
+        else:
+            self.index = None
+
+    def gen_vector_db_from_cache(self):
+        """从缓存生成向量数据库"""
+        self.build_index()
+        self.save_cache()
 
     def generate_codebase(self):
-        updated =self.clean_db()
-        db_lock = Lock()
-        processed_files = []  # 用于跟踪已处理的文件
-        
-        def process_and_save(file):
-            result = self.process_file(file)
-            if result:
-                file, md5, description = result
-                db = self.get_db_connection()
-                try:
-                    with db_lock:
-                        db.execute("DELETE FROM codebase WHERE path = ?", (file,))
-                        db.execute("INSERT INTO codebase (path, md5, description) VALUES (?, ?, ?)", 
-                                 (file, md5, description))
-                        db.commit()
-                        PrettyOutput.print(f"索引文件: {file}", output_type=OutputType.INFO)
-                        processed_files.append(file)
-                finally:
-                    db.close()
+        """生成代码库索引"""
+        files_deleted = self.clean_cache()  # 清理过期缓存
+        processed_files = []
         
-        # 使用 ThreadPoolExecutor 并等待所有任务完成
+        # 使用线程池处理文件
         with ThreadPoolExecutor(max_workers=self.thread_count) as executor:
-            futures = [executor.submit(process_and_save, file) for file in self.git_file_list]
-            # 等待所有任务完成
-            concurrent.futures.wait(futures)
+            futures = [executor.submit(self.process_file, file) for file in self.git_file_list]
+            for future in concurrent.futures.as_completed(futures):
+                result = future.result()
+                if result:
+                    processed_files.append(result)
+                    PrettyOutput.print(f"索引文件: {result}", output_type=OutputType.INFO)
 
-        if updated or len(processed_files) > 0:
-            PrettyOutput.print("有新的文件被删除或添加，正在重新生成向量数据库", output_type=OutputType.INFO)
-            self.gen_vector_db_from_sqlite()
+        if files_deleted or processed_files:
+            PrettyOutput.print("重新生成向量数据库", output_type=OutputType.INFO)
+            self.gen_vector_db_from_cache()
         else:
-            PrettyOutput.print("没有新的文件被删除或添加，跳过向量数据库生成", output_type=OutputType.INFO)
+            PrettyOutput.print("没有新的文件变更，跳过向量数据库生成", output_type=OutputType.INFO)
             
-        PrettyOutput.print(f"成功索引 {len(processed_files)} 个文件", output_type=OutputType.INFO)
+        PrettyOutput.print(f"成功为 {len(processed_files)} 个文件生成索引", output_type=OutputType.INFO)
 
-    def search_similar(self, query: str, top_k: int = 5) -> List[Tuple[str, float, str]]:
-        """搜索与查询最相似的文件
-        
-        Args:
-            query: 查询文本
-            top_k: 返回结果数量
+    def rerank_results(self, query: str, initial_results: List[Tuple[str, float, str]]) -> List[Tuple[str, float, str]]:
+        """使用大模型对搜索结果重新排序"""
+        if not initial_results:
+            return []
+
+        model = self.platform_registry.create_platform(self.normal_platform)
+        model.set_model_name(self.normal_model)
+        model.set_suppress_output(True)
+
+        try:
+            # 构建重排序的prompt
+            prompt = f"""请根据用户的查询，对以下代码文件进行相关性排序。对每个文件给出0-100的相关性分数，分数越高表示越相关。
+只需要输出每个文件的分数，格式为：
+<RERANK_START>
+文件路径: 分数
+文件路径: 分数
+<RERANK_END>
+
+用户查询: {query}
+
+待评估文件:
+"""
+            for path, _, desc in initial_results:
+                prompt += f"""
+文件: {path}
+描述: {desc}
+---
+"""
+            
+            response = model.chat(prompt)
+            
+            # 提取<RERANK_START>和<RERANK_END>之间的内容
+            start_tag = "<RERANK_START>"
+            end_tag = "<RERANK_END>"
+            if start_tag in response and end_tag in response:
+                response = response[response.find(start_tag) + len(start_tag):response.find(end_tag)]
+            
+            # 解析响应，提取文件路径和分数
+            scored_results = []
+            for line in response.split('\n'):
+                if ':' not in line:
+                    continue
+                try:
+                    file_path, score_str = line.split(':', 1)
+                    file_path = file_path.strip()
+                    score = float(score_str.strip())
+                    # 找到对应的原始描述
+                    desc = next((desc for p, _, desc in initial_results if p == file_path), "")
+                    scored_results.append((file_path, score/100.0, desc))
+                except:
+                    continue
             
-        Returns:
-            List of (file_path, similarity_score, description) tuples
-        """
-        # 获取查询文本的向量表示
-        query_vector = self.get_embedding(query)
-        query_vector = query_vector.reshape(1, -1)
-
-        # 搜索最相似的向量
-        distances, indices = self.index.search(query_vector, top_k)
+            # 按分数降序排序
+            return sorted(scored_results, key=lambda x: x[1], reverse=True)
+            
+        finally:
+            model.delete_chat()
+        
+        return initial_results
+
+    def search_similar(self, query: str, top_k: int = 20) -> List[Tuple[str, float, str]]:
+        """搜索相似文件"""
+        model = self.platform_registry.create_platform(self.normal_platform)
+        model.set_model_name(self.normal_model)
+        model.set_suppress_output(True)
         
-        # 获取对应的文件信息
-        db = self.get_db_connection()
         try:
-            results = []
-            for i, distance in zip(indices[0], distances[0]):
-                if i == -1:  # faiss返回-1表示无效结果
-                    continue
-                    
-                # 将numpy.int64转换为Python int
-                offset = int(i)
-                # 获取文件路径和描述
-                cursor = db.execute("SELECT path, description FROM codebase LIMIT 1 OFFSET ?", (offset,))
-                row = cursor.fetchone()
-                if row:
-                    path, description = row
-                    # 将distance转换为相似度分数（0-1之间）
-                    similarity = 1.0 / (1.0 + float(distance))  # 确保使用Python float
-                    results.append((path, similarity, description))
+            prompt = f"""请根据以下查询，生成意思完全相同的另一个表述。这个表述将用于代码搜索，所以要保持专业性和准确性。
+原始查询: {query}
+
+请直接输出新表述，不要有编号或其他标记。
+"""
+            
+            query = model.chat(prompt)
             
-            return results
         finally:
-            db.close()
+            model.delete_chat()
+
+        PrettyOutput.print(f"查询:  {query}", output_type=OutputType.INFO)
+        
+        # 为每个查询获取相似文件
+        all_results = {}  # 文件路径 -> (总分数, 出现次数, 描述)
+        
+        q_vector = self.get_embedding(query)
+        q_vector = q_vector.reshape(1, -1)
+            
+        distances, indices = self.index.search(q_vector, top_k)
+
+        PrettyOutput.print(f"查询 {query} 的结果: ", output_type=OutputType.INFO)
+
+        initial_results = []
+        
+        for i, distance in zip(indices[0], distances[0]):
+            if i == -1:  # faiss返回-1表示无效结果
+                continue
+                
+            similarity = 1.0 / (1.0 + float(distance))
+            PrettyOutput.print(f"  {self.file_paths[i]} : 距离 {distance:.3f}, 相似度 {similarity:.3f}", 
+                                output_type=OutputType.INFO)
+                
+            file_path = self.file_paths[i]
+            data = self.vector_cache[file_path]
+            initial_results.append((file_path, similarity, data["description"]))
+
+        # 使用大模型重新排序
+        PrettyOutput.print("使用大模型重新排序...", output_type=OutputType.INFO)
+        reranked_results = self.rerank_results(query, initial_results)
+        
+        return reranked_results
 
-    def ask_codebase(self, query: str, top_k: int = 5) -> List[Tuple[str, float, str]]:
-        """Ask a question about the codebase"""
-        # 使用搜索函数获取相似文件
+    def ask_codebase(self, query: str, top_k: int=20) -> str:
+        """查询代码库"""
         results = self.search_similar(query, top_k)
-        PrettyOutput.print(f"找到的关联文件: ", output_type=OutputType.INFO)
+        PrettyOutput.print(f"找到的关联文件: ", output_type=OutputType.SUCCESS)
         for path, score, _ in results:
-            PrettyOutput.print(f"文件: {path} 关联度: {score:.3f}", output_type=OutputType.INFO)
+            PrettyOutput.print(f"文件: {path} 关联度: {score:.3f}", 
+                             output_type=OutputType.INFO)
         
         prompt = f"""你是一个代码专家，请根据以下文件信息回答用户的问题：
 """
         for path, _, _ in results:
-            content = open(path, "r", encoding="utf-8").read()
-            prompt += f"""
-文件路径: {path}
+            try:
+                if len(prompt) > 30 * 1024:
+                    PrettyOutput.print(f"避免上下文超限，丢弃低相关度文件：{path}", OutputType.WARNING)
+                    continue
+                content = open(path, "r", encoding="utf-8").read()
+                prompt += f"""
+文件路径: {path}prompt
 文件内容:
 {content}
 ========================================
 """
+            except Exception as e:
+                PrettyOutput.print(f"读取文件失败 {path}: {str(e)}", 
+                                 output_type=OutputType.ERROR)
+                continue
+                
         prompt += f"""
 用户问题: {query}
 
 请用专业的语言回答用户的问题，如果给出的文件内容不足以回答用户的问题，请告诉用户，绝对不要胡编乱造。
 """
-        model = self.platform_registry.create_platform(self.normal_platform)
-        model.set_model_name(self.normal_model)
-        response = model.chat(prompt)
-        return response
+        model = self.platform_registry.create_platform(self.codegen_platform)
+        model.set_model_name(self.codegen_model)
+        try:
+            response = model.chat(prompt)
+            return response
+        finally:
+            model.delete_chat()
 
 
 def main():
     parser = argparse.ArgumentParser(description='Codebase management and search tool')
     parser.add_argument('--search', type=str, help='Search query to find similar code files')
-    parser.add_argument('--top-k', type=int, default=5, help='Number of results to return (default: 5)')
+    parser.add_argument('--top-k', type=int, default=20, help='Number of results to return (default: 20)')
     parser.add_argument('--ask', type=str, help='Ask a question about the codebase')
     args = parser.parse_args()
     

jarvis/jarvis_coder/main.py

@@ -1,11 +1,8 @@
-from concurrent.futures import ThreadPoolExecutor, as_completed
-import hashlib
 import os
 import re
-import sqlite3
 import threading
 import time
-from typing import Dict, Any, List, Optional, Tuple
+from typing import Dict, Any, List, Tuple
 
 import yaml
 from jarvis.models.base import BasePlatform
@@ -127,7 +124,7 @@ class JarvisCoder:
             return [patch.replace('<PATCH_START>', '').replace('<PATCH_END>', '').strip() 
                    for patch in patches if patch.strip()]
         except Exception as e:
-            PrettyOutput.print(f"解析patch失败: {str(e)}", OutputType.ERROR)
+            PrettyOutput.print(f"解析patch失败: {str(e)}", OutputType.WARNING)
             return []
         
     def _make_patch(self, related_files: List[Dict], feature: str) -> List[str]:
@@ -141,20 +138,31 @@ class JarvisCoder:
 要替换的内容
 =======
 新的内容
-<<<<<<
+>>>>>>
 <PATCH_END>
 
-2. 如果是新文件，格式如下：
+2. 如果是新文件或者替换整个文件内容，格式如下：
 <PATCH_START>
 >>>>>> path/to/new/file
 =======
 新文件的完整内容
-<<<<<<
+>>>>>>
+<PATCH_END>
+
+3. 如果要删除文件中的某一段，格式如下：
+<PATCH_START>
+>>>>>> path/to/file
+要删除的内容
+=======
+>>>>>>
 <PATCH_END>
 
 文件列表如下：
 """
         for i, file in enumerate(related_files):
+            if len(prompt) > 30 * 1024:
+                PrettyOutput.print(f'避免上下文超限，丢弃低相关度文件：{file["file_path"]}', OutputType.WARNING)
+                continue
             prompt += f"""{i}. {file["file_path"]}\n"""
             prompt += f"""文件内容:\n"""
             prompt += f"<FILE_CONTENT_START>\n"
@@ -180,7 +188,7 @@ class JarvisCoder:
             return [patch.replace('<PATCH_START>', '').replace('<PATCH_END>', '').strip() 
                    for patch in patches if patch.strip()]
         except Exception as e:
-            PrettyOutput.print(f"解析patch失败: {str(e)}", OutputType.ERROR)
+            PrettyOutput.print(f"解析patch失败: {str(e)}", OutputType.WARNING)
             return []
 
     def _apply_patch(self, related_files: List[Dict], patches: List[str]) -> Tuple[bool, str]:
@@ -219,7 +227,7 @@ class JarvisCoder:
                     return False, "\n".join(error_info)
                 
                 old_content = parts[0]
-                new_content = parts[1].split("<<<<<<")[0]
+                new_content = parts[1].split(">>>>>>")[0]
                 
                 # 处理新文件
                 if not old_content:
@@ -308,9 +316,9 @@ class JarvisCoder:
     def _load_related_files(self, feature: str) -> List[Dict]:
         """加载相关文件内容"""
         ret = []
-        related_files = self._codebase.search_similar(feature, top_k=5)
+        related_files = self._codebase.search_similar(feature)
         for file, score, _ in related_files:
-            PrettyOutput.print(f"相关文件: {file} 相关度: {score:.3f}", OutputType.INFO)
+            PrettyOutput.print(f"相关文件: {file} 相关度: {score:.3f}", OutputType.SUCCESS)
             with open(file, "r", encoding="utf-8") as f:
                 content = f.read()
             ret.append({"file_path": file, "file_content": content})
@@ -441,7 +449,7 @@ def main():
                 PrettyOutput.print(result["stdout"], OutputType.SUCCESS)
             else:
                 if result["stderr"]:
-                    PrettyOutput.print(result["stderr"], OutputType.ERROR)
+                    PrettyOutput.print(result["stderr"], OutputType.WARNING)
                 if result["error"]:
                     PrettyOutput.print(f"错误类型: {type(result['error']).__name__}", OutputType.WARNING)
                 

jarvis/models/registry.py

@@ -121,7 +121,7 @@ class PlatformRegistry:
                             # 检查平台实现
                             if not PlatformRegistry.check_platform_implementation(obj):
                                 continue
-
+                            PrettyOutput.print(f"从 {os.path.join(directory, filename)} 加载平台：{obj.platform_name}", OutputType.SUCCESS)
                             platforms[obj.platform_name] = obj
                             break
                 except Exception as e:

jarvis/tools/__init__.py

@@ -1,9 +1,6 @@
 from .registry import ToolRegistry
-from jarvis.tools.codebase_qa import CodebaseQATool
 
 __all__ = [
     'ToolRegistry',
 ] 
 
-def register_tools():
-    register_tool(CodebaseQATool()) 

jarvis/tools/registry.py

@@ -99,7 +99,7 @@ class ToolRegistry:
                         parameters=tool_instance.parameters,
                         func=tool_instance.execute
                     )
-                    PrettyOutput.print(f"从 {file_path} 加载工具: {tool_instance.name}: {tool_instance.description}", OutputType.INFO)
+                    PrettyOutput.print(f"从 {file_path} 加载工具: {tool_instance.name}: {tool_instance.description}", OutputType.SUCCESS)
                     tool_found = True
                     break
                     

jarvis/utils.py

@@ -158,7 +158,7 @@ def get_multiline_input(tip: str) -> str:
             lines.append(line)
             
     except KeyboardInterrupt:
-        PrettyOutput.print("\n输入已取消", OutputType.ERROR)
+        PrettyOutput.print("\n输入已取消", OutputType.INFO)
         return "__interrupt__"
     
     return "\n".join(lines)

{jarvis_ai_assistant-0.1.66.dist-info → jarvis_ai_assistant-0.1.72.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: jarvis-ai-assistant
-Version: 0.1.66
+Version: 0.1.72
 Summary: Jarvis: An AI assistant that uses tools to interact with the system
 Home-page: https://github.com/skyfireitdiy/Jarvis
 Author: skyfire
@@ -44,9 +44,10 @@ Requires-Dist: colorama>=0.4.6
 Requires-Dist: prompt_toolkit>=3.0.0
 Requires-Dist: openai>=1.20.0
 Requires-Dist: playwright>=1.41.1
-Requires-Dist: numpy>=1.26.0
-Requires-Dist: faiss-cpu>=1.8.1
+Requires-Dist: numpy>=1.24.0
+Requires-Dist: faiss-cpu>=1.8.0
 Requires-Dist: sentence-transformers>=2.2.2
+Requires-Dist: bs4>=0.0.1
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: black; extra == "dev"
@@ -123,7 +124,6 @@ Jarvis supports configuration through environment variables that can be set in t
 |---------|------|--------|------|
 | JARVIS_PLATFORM | AI platform to use, supports kimi/openai/ai8 etc | kimi | Yes |
 | JARVIS_MODEL | Model name to use | - | No |
-
 | JARVIS_CODEGEN_PLATFORM | AI platform for code generation | Same as JARVIS_PLATFORM | No |
 | JARVIS_CODEGEN_MODEL | Model name for code generation | Same as JARVIS_MODEL | No |
 | JARVIS_CHEAP_PLATFORM | AI platform for cheap operations | Same as JARVIS_PLATFORM | No |
@@ -155,17 +155,17 @@ jarvis -p openai  # Use OpenAI platform
 
 ### Code Modification
 ```bash
-jarvis coder --feature "Add new feature"  # Modify code to add new feature
+jarvis-coder --feature "Add new feature"  # Modify code to add new feature
 ```
 
 ### Codebase Search
 ```bash
-jarvis codebase --search "database connection"  # Search codebase
+jarvis-codebase --search "database connection"  # Search codebase
 ```
 
 ### Codebase Question
 ```bash
-jarvis codebase --ask "How to use the database?"  # Ask about codebase
+jarvis-codebase --ask "How to use the database?"  # Ask about codebase
 ```
 
 ### Keep Chat History

{jarvis_ai_assistant-0.1.66.dist-info → jarvis_ai_assistant-0.1.72.dist-info}/RECORD

@@ -1,32 +1,33 @@
-jarvis/__init__.py,sha256=XwQusgz9_jvC5jEMwbV21qnY-dFIO3yH5W9FEENNxsU,50
+jarvis/__init__.py,sha256=MyFCehMdftT6QiIVeyWoT3cnA4reDk6ObCHGRDEr7JY,50
 jarvis/agent.py,sha256=kl6pwNrluzb-9eZKgwmsk5Jh4CpWi4F8B3RvEQNvc5U,14921
 jarvis/main.py,sha256=7EcSlxa5JFFXBujzKDWdNtwX6axLhFFdJMc2GxTjfdk,6295
-jarvis/utils.py,sha256=bjC0PAR58RvcXHgabIFmNmYL1L_GhhiEwMFytWurcN4,7499
+jarvis/utils.py,sha256=vZV8sHj0ggZy4Rb8RxIujQhRWgeNEomhqVl4WXmpq7c,7498
 jarvis/jarvis_codebase/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jarvis/jarvis_codebase/main.py,sha256=bncfOXKFthcr21cV9tONXzcchSWaaYXmo0o11R5hPzc,14799
-jarvis/jarvis_coder/main.py,sha256=mK68MJyOerVnY7Fr9ibQ1swQLVIWyBAFmmVj08SHKyk,21825
+jarvis/jarvis_codebase/main.py,sha256=A-KY9WzAIsCF0QPndm4b8-rAKEz4nLlm4NkfUMpnQUM,21965
+jarvis/jarvis_coder/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+jarvis/jarvis_coder/main.py,sha256=uQ7aQ3hmnxecacjLSBOV0rSLkacExGLgQRc9zD1ar1E,22085
 jarvis/models/__init__.py,sha256=mrOt67nselz_H1gX9wdAO4y2DY5WPXzABqJbr5Des8k,63
 jarvis/models/ai8.py,sha256=vgy-r_3HHxGMAalZrA65VWHC1PuwBTYgtprSgHkCbrk,12557
 jarvis/models/base.py,sha256=ShV1H8Unee4RMaiFO4idROQA0Hc6wu4dyeRPX5fcszk,1433
 jarvis/models/kimi.py,sha256=1iTB0Z_WOmCML3Ufsge6jmeKOYvccr7I5lS3JUXymU4,17611
 jarvis/models/openai.py,sha256=ayaBWAN5VexMcKVrjEPDNB-Q9wx0sCV9Z4BCrvwYJ9w,4315
 jarvis/models/oyi.py,sha256=X2c5SWDIuQDCCFBcEKbzIWEz3I34eOAi0d1XAFgxlpw,15001
-jarvis/models/registry.py,sha256=hJyaROiOF_TkbtIXsjOD8-ArOvAvtxviawyqBFfLV6s,7617
-jarvis/tools/__init__.py,sha256=xmROdzJTZz6JDLLuAbwVLjUD4xfUUYb6D1Ssu_desaE,183
+jarvis/models/registry.py,sha256=YpooKSpk5pSWfb5cBDz5wRfPK-abb9uuUZr4WBejqwI,7762
+jarvis/tools/__init__.py,sha256=7Rqyj5hBAv5cWDVr5T9ZTZASO7ssBHeQNm2_4ZARdkA,72
 jarvis/tools/base.py,sha256=EGRGbdfbLXDLwtyoWdvp9rlxNX7bzc20t0Vc2VkwIEY,652
 jarvis/tools/codebase_qa.py,sha256=AEpusYxyWtALVVwPk1DMUH9cVI73mE1e3WFHJXDpXto,2333
 jarvis/tools/coder.py,sha256=ZJfPInKms4Hj3-eQlBwamVsvZ-2nlZ-4jsqJ-tJc6mg,2040
 jarvis/tools/file_ops.py,sha256=h8g0eT9UvlJf4kt0DLXvdSsjcPj7x19lxWdDApeDfpg,3842
 jarvis/tools/generator.py,sha256=vVP3eN5cCDpRXf_fn0skETkPXAW1XZFWx9pt2_ahK48,5999
 jarvis/tools/methodology.py,sha256=UG6s5VYRcd9wrKX4cg6f7zJhet5AIcthFGMOAdevBiw,5175
-jarvis/tools/registry.py,sha256=mlOAmUq3yzRz-7yvwrrCwbe5Lmw8eh1v8-_Fa5sezwI,7209
+jarvis/tools/registry.py,sha256=MeTYNdZNRdhlgABviVxzbDPSgLpwDp2Nx2dGzedRu8U,7212
 jarvis/tools/search.py,sha256=1EqOVvLhg2Csh-i03-XeCrusbyfmH69FZ8khwZt8Tow,6131
 jarvis/tools/shell.py,sha256=UPKshPyOaUwTngresUw-ot1jHjQIb4wCY5nkJqa38lU,2520
 jarvis/tools/sub_agent.py,sha256=rEtAmSVY2ZjFOZEKr5m5wpACOQIiM9Zr_3dT92FhXYU,2621
 jarvis/tools/webpage.py,sha256=d3w3Jcjcu1ESciezTkz3n3Zf-rp_l91PrVoDEZnckOo,2391
-jarvis_ai_assistant-0.1.66.dist-info/LICENSE,sha256=AGgVgQmTqFvaztRtCAXsAMryUymB18gZif7_l2e1XOg,1063
-jarvis_ai_assistant-0.1.66.dist-info/METADATA,sha256=O3UdZLfScFyhR3ARPYaKu1d5hemg45HFN6bZKclTfYk,12374
-jarvis_ai_assistant-0.1.66.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-jarvis_ai_assistant-0.1.66.dist-info/entry_points.txt,sha256=QNUeqmUJd7nHufel2FO7cRttS1uKFfnbIyObv8eVyOY,140
-jarvis_ai_assistant-0.1.66.dist-info/top_level.txt,sha256=1BOxyWfzOP_ZXj8rVTDnNCJ92bBGB0rwq8N1PCpoMIs,7
-jarvis_ai_assistant-0.1.66.dist-info/RECORD,,
+jarvis_ai_assistant-0.1.72.dist-info/LICENSE,sha256=AGgVgQmTqFvaztRtCAXsAMryUymB18gZif7_l2e1XOg,1063
+jarvis_ai_assistant-0.1.72.dist-info/METADATA,sha256=4YOsbaLCJw40e43wAL34myfveTRLOmC-MIpBrkrt2zs,12399
+jarvis_ai_assistant-0.1.72.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+jarvis_ai_assistant-0.1.72.dist-info/entry_points.txt,sha256=QNUeqmUJd7nHufel2FO7cRttS1uKFfnbIyObv8eVyOY,140
+jarvis_ai_assistant-0.1.72.dist-info/top_level.txt,sha256=1BOxyWfzOP_ZXj8rVTDnNCJ92bBGB0rwq8N1PCpoMIs,7
+jarvis_ai_assistant-0.1.72.dist-info/RECORD,,