jarvis-ai-assistant 0.1.67__tar.gz → 0.1.73__tar.gz

{jarvis_ai_assistant-0.1.67/src/jarvis_ai_assistant.egg-info → jarvis_ai_assistant-0.1.73}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: jarvis-ai-assistant
-Version: 0.1.67
+Version: 0.1.73
 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"

{jarvis_ai_assistant-0.1.67 → jarvis_ai_assistant-0.1.73}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "jarvis-ai-assistant"
-version = "0.1.67"
+version = "0.1.73"
 description = "Jarvis: An AI assistant that uses tools to interact with the system"
 readme = "README.md"
 authors = [{ name = "Your Name", email = "your.email@example.com" }]
@@ -26,9 +26,10 @@ dependencies = [
     "prompt_toolkit>=3.0.0",
     "openai>=1.20.0",
     "playwright>=1.41.1",
-    "numpy>=1.26.0",
-    "faiss-cpu>=1.8.1",
+    "numpy>=1.24.0",
+    "faiss-cpu>=1.8.0",
     "sentence-transformers>=2.2.2",
+    "bs4>=0.0.1",
 ]
 requires-python = ">=3.8"
 

{jarvis_ai_assistant-0.1.67 → jarvis_ai_assistant-0.1.73}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="jarvis-ai-assistant",
-    version="0.1.67",
+    version="0.1.73",
     author="skyfire",
     author_email="skyfireitdiy@hotmail.com",
     description="An AI assistant that uses various tools to interact with the system",
@@ -19,9 +19,10 @@ setup(
         "prompt_toolkit>=3.0.0",
         "openai>=1.20.0",
         "playwright>=1.41.1",
-        "numpy>=1.26.0",
-        "faiss-cpu>=1.8.1",
+        "numpy>=1.24.0",
+        "faiss-cpu>=1.8.0",
         "sentence-transformers>=2.2.2",
+        "bs4>=0.0.1",
     ],
     entry_points={
         "console_scripts": [

{jarvis_ai_assistant-0.1.67 → jarvis_ai_assistant-0.1.73}/src/jarvis/__init__.py

@@ -1,3 +1,3 @@
 """Jarvis AI Assistant"""
 
-__version__ = "0.1.67" 
+__version__ = "0.1.73" 

{jarvis_ai_assistant-0.1.67 → jarvis_ai_assistant-0.1.73}/src/jarvis/jarvis_codebase/main.py

@@ -18,18 +18,22 @@ class CodeBase:
         load_env_from_file()
         self.root_dir = root_dir
         os.chdir(self.root_dir)
-        self.thread_count = os.environ.get("JARVIS_THREAD_COUNT") or 10
+        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)
             
         # 初始化数据目录
@@ -39,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)
             
@@ -327,7 +332,69 @@ class CodeBase:
             
         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]]:
+    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()) / 100.0  # 转换为0-1范围
+                    # 只保留相关度大于等于0.7的结果
+                    if score >= 0.7:
+                        # 找到对应的原始描述
+                        desc = next((desc for p, _, desc in initial_results if p == file_path), "")
+                        scored_results.append((file_path, score, desc))
+                except:
+                    continue
+            
+            # 按分数降序排序
+            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)
@@ -348,8 +415,6 @@ class CodeBase:
         PrettyOutput.print(f"查询:  {query}", output_type=OutputType.INFO)
         
         # 为每个查询获取相似文件
-        all_results = {}  # 文件路径 -> (总分数, 出现次数, 描述)
-        
         q_vector = self.get_embedding(query)
         q_vector = q_vector.reshape(1, -1)
             
@@ -357,25 +422,36 @@ class CodeBase:
 
         PrettyOutput.print(f"查询 {query} 的结果: ", output_type=OutputType.INFO)
 
-        ret = []
+        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)
+            # 只保留相似度大于等于0.5的结果
+            if similarity >= 0.5:
+                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]
-            ret.append((file_path, similarity, data["description"]))
-        return ret
+                file_path = self.file_paths[i]
+                data = self.vector_cache[file_path]
+                initial_results.append((file_path, similarity, data["description"]))
+
+        if not initial_results:
+            PrettyOutput.print("没有找到相似度大于0.5的文件", output_type=OutputType.WARNING)
+            return []
 
-    def ask_codebase(self, query: str, top_k: int = 5) -> List[Tuple[str, float, str]]:
+        # 使用大模型重新排序
+        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=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)
@@ -384,9 +460,12 @@ class CodeBase:
 """
         for path, _, _ in results:
             try:
+                if len(prompt) > 30 * 1024:
+                    PrettyOutput.print(f"避免上下文超限，丢弃低相关度文件：{path}", OutputType.WARNING)
+                    continue
                 content = open(path, "r", encoding="utf-8").read()
                 prompt += f"""
-文件路径: {path}
+文件路径: {path}prompt
 文件内容:
 {content}
 ========================================
@@ -401,8 +480,8 @@ class CodeBase:
 
 请用专业的语言回答用户的问题，如果给出的文件内容不足以回答用户的问题，请告诉用户，绝对不要胡编乱造。
 """
-        model = self.platform_registry.create_platform(self.normal_platform)
-        model.set_model_name(self.normal_model)
+        model = self.platform_registry.create_platform(self.codegen_platform)
+        model.set_model_name(self.codegen_model)
         try:
             response = model.chat(prompt)
             return response
@@ -410,21 +489,25 @@ class CodeBase:
             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')
+    parser.add_argument('--generate', action='store_true', help='Generate codebase index')
     args = parser.parse_args()
     
     current_dir = find_git_root()
     codebase = CodeBase(current_dir)
 
-    try:
-        codebase.generate_codebase()
-        PrettyOutput.print("\nCodebase generation completed", output_type=OutputType.SUCCESS)
-    except Exception as e:
-        PrettyOutput.print(f"Error during codebase generation: {str(e)}", output_type=OutputType.ERROR)
+
+    if args.generate:
+        try:
+            codebase.generate_codebase()
+            PrettyOutput.print("\nCodebase generation completed", output_type=OutputType.SUCCESS)
+        except Exception as e:
+            PrettyOutput.print(f"Error during codebase generation: {str(e)}", output_type=OutputType.ERROR)
     
     if args.search:
         results = codebase.search_similar(args.search, args.top_k)

{jarvis_ai_assistant-0.1.67 → jarvis_ai_assistant-0.1.73}/src/jarvis/jarvis_coder/main.py

@@ -124,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]:
@@ -138,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"
@@ -177,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]:
@@ -216,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:
@@ -297,17 +308,37 @@ class JarvisCoder:
 
 
 
+
     def _prepare_execution(self) -> None:
         """准备执行环境"""
         self.main_model = self._new_model()
-        self._codebase.generate_codebase()
+        
+
+        # 询问用户是否更新索引数据库
+        while True:
+            answer = get_multiline_input("是否要更新代码库索引数据库？(y/n)").strip().lower()
+            if answer == 'y':
+                PrettyOutput.print("正在更新代码库索引数据库...", OutputType.PROGRESS)
+                self._codebase.generate_codebase()
+                break
+            elif answer == 'n' or not answer:
+                PrettyOutput.print("跳过代码库索引数据库更新", OutputType.INFO)
+                break
+            else:
+                PrettyOutput.print("请输入 y 或 n", OutputType.WARNING)
+
 
     def _load_related_files(self, feature: str) -> List[Dict]:
         """加载相关文件内容"""
         ret = []
-        related_files = self._codebase.search_similar(feature, top_k=5)
+        # 确保索引数据库已生成
+        if not self._codebase.is_index_generated():
+            PrettyOutput.print("检测到索引数据库未生成，正在生成...", OutputType.WARNING)
+            self._codebase.generate_codebase()
+            
+        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})
@@ -438,7 +469,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_ai_assistant-0.1.67 → jarvis_ai_assistant-0.1.73}/src/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_ai_assistant-0.1.67 → jarvis_ai_assistant-0.1.73}/src/jarvis/tools/codebase_qa.py

@@ -45,8 +45,11 @@ class CodebaseQATool:
                     "stderr": "错误：当前目录不在Git仓库中",
                     "error": "NotInGitRepository"
                 }
+
             os.chdir(root_dir)
             codebase = CodeBase(root_dir)
+            # 生成索引
+            codebase.generate_codebase()
             # 执行问答
             response = codebase.ask_codebase(question, top_k)
             os.chdir(current_dir)

{jarvis_ai_assistant-0.1.67 → jarvis_ai_assistant-0.1.73}/src/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_ai_assistant-0.1.67 → jarvis_ai_assistant-0.1.73}/src/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.67 → jarvis_ai_assistant-0.1.73/src/jarvis_ai_assistant.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: jarvis-ai-assistant
-Version: 0.1.67
+Version: 0.1.73
 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"

{jarvis_ai_assistant-0.1.67 → jarvis_ai_assistant-0.1.73}/src/jarvis_ai_assistant.egg-info/SOURCES.txt

@@ -9,6 +9,7 @@ src/jarvis/main.py
 src/jarvis/utils.py
 src/jarvis/jarvis_codebase/__init__.py
 src/jarvis/jarvis_codebase/main.py
+src/jarvis/jarvis_coder/__init__.py
 src/jarvis/jarvis_coder/main.py
 src/jarvis/models/__init__.py
 src/jarvis/models/ai8.py

{jarvis_ai_assistant-0.1.67 → jarvis_ai_assistant-0.1.73}/src/jarvis_ai_assistant.egg-info/requires.txt

@@ -4,9 +4,10 @@ colorama>=0.4.6
 prompt_toolkit>=3.0.0
 openai>=1.20.0
 playwright>=1.41.1
-numpy>=1.26.0
-faiss-cpu>=1.8.1
+numpy>=1.24.0
+faiss-cpu>=1.8.0
 sentence-transformers>=2.2.2
+bs4>=0.0.1
 
 [dev]
 pytest