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

jarvis/__init__.py

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

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,67 @@ 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())
+                    # 找到对应的原始描述
+                    desc = next((desc for p, _, desc in initial_results if p == file_path), "")
+                    scored_results.append((file_path, score/100.0, 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)
@@ -357,7 +422,7 @@ 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表示无效结果
@@ -369,13 +434,18 @@ class CodeBase:
                 
             file_path = self.file_paths[i]
             data = self.vector_cache[file_path]
-            ret.append((file_path, similarity, data["description"]))
-        return ret
+            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]]:
+    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 +454,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 +474,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
@@ -413,7 +486,7 @@ class CodeBase:
 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

@@ -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:
@@ -305,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})
@@ -438,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/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.67.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.67
+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"

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

@@ -1,17 +1,18 @@
-jarvis/__init__.py,sha256=YEUpLs5Xl5YrTW6S6tl7I42rz6TcpyCTJZRczXXMmi0,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=LdJGg4dUaK8naNkGoqF9i53OaYndKGs3jXfZQd4w-3M,18699
-jarvis/jarvis_coder/main.py,sha256=L_i1Zp3UJbc891WrpxKgoD4EzoqvLRnqznmLQP2ZN2U,21721
+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/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
@@ -19,14 +20,14 @@ 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.67.dist-info/LICENSE,sha256=AGgVgQmTqFvaztRtCAXsAMryUymB18gZif7_l2e1XOg,1063
-jarvis_ai_assistant-0.1.67.dist-info/METADATA,sha256=Vpf669oYa2bpkQDzfAsgSP0bv0geylGiyuPq6ezarJY,12373
-jarvis_ai_assistant-0.1.67.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-jarvis_ai_assistant-0.1.67.dist-info/entry_points.txt,sha256=QNUeqmUJd7nHufel2FO7cRttS1uKFfnbIyObv8eVyOY,140
-jarvis_ai_assistant-0.1.67.dist-info/top_level.txt,sha256=1BOxyWfzOP_ZXj8rVTDnNCJ92bBGB0rwq8N1PCpoMIs,7
-jarvis_ai_assistant-0.1.67.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,,