jarvis-ai-assistant 0.1.85__py3-none-any.whl → 0.1.87__py3-none-any.whl

jarvis/__init__.py

@@ -1,3 +1,3 @@
 """Jarvis AI Assistant"""
 
-__version__ = "0.1.85" 
+__version__ = "0.1.87" 

jarvis/agent.py

@@ -8,7 +8,7 @@ import json
 
 from .models.registry import PlatformRegistry
 from .tools import ToolRegistry
-from .utils import PrettyOutput, OutputType, get_multiline_input, load_embedding_model, while_success
+from .utils import PrettyOutput, OutputType, get_max_context_length, get_multiline_input, load_embedding_model, while_success
 import os
 from datetime import datetime
 from prompt_toolkit import prompt
@@ -33,7 +33,7 @@ class Agent:
         
         # 从环境变量加载配置
         self.embedding_dimension = 1536  # Default for many embedding models
-        self.max_context_length = int(os.getenv('JARVIS_MAX_CONTEXT_LENGTH', '65536'))  # 默认30k
+        self.max_context_length = get_max_context_length()
         
         # 初始化嵌入模型
         try:
@@ -72,10 +72,10 @@ class Agent:
         
         # 逐行处理
         for line in lines:
-            if '<START_TOOL_CALL>' in line:
+            if '<TOOL_CALL>' in line:
                 in_tool_call = True
                 continue
-            elif '<END_TOOL_CALL>' in line:
+            elif '</TOOL_CALL>' in line:
                 if in_tool_call and tool_call_lines:
                     try:
                         # 直接解析YAML
@@ -344,7 +344,7 @@ class Agent:
 4. 生成解决方案：生成一个或者多个具备可操作性的解决方案
 5. 评估解决方案：从众多解决方案中选择一种最优的方案
 6. 制定行动计划：根据目前可以使用的工具制定行动计划，使用 PlantUML 格式输出明确的执行流程
-7. 执行行动计划：每步执行一个步骤，最多使用一个工具（工具执行完成后，等待工具结果再执行下一步）
+7. 执行行动计划：每步执行一个步骤，**最多使用一个工具**（工具执行完成后，等待工具结果再执行下一步）
 8. 监控与调整：如果执行结果与预期不符，则反思并调整行动计划，迭代之前的步骤
 9. 方法论：如果当前任务具有通用性且执行过程中遇到了值得记录的经验，使用方法论工具记录方法论，以提升后期处理类似问题的能力
 10. 任务结束：如果任务已经完成，使用任务结束指令结束任务
@@ -362,12 +362,12 @@ class Agent:
 
 工具使用格式：
 
-<START_TOOL_CALL>
+<TOOL_CALL>
 name: tool_name
 arguments:
     param1: value1
     param2: value2
-<END_TOOL_CALL>
+</TOOL_CALL>
 
 -------------------------------------------------------------
 
@@ -381,6 +381,7 @@ arguments:
 7. 在执行一些可能对系统或者用户代码库造成破坏的工具时，请先询问用户
 8. 在多次迭代却没有任何进展时，可请求用户指导
 9. 如果返回的yaml字符串中包含冒号，请将整个字符串用引号包裹，避免yaml解析错误
+10. yaml中包含多行字符串时，请使用 | 语法
 
 {methodology_prompt}
 

jarvis/jarvis_codebase/main.py

@@ -7,7 +7,7 @@ from jarvis.models.registry import PlatformRegistry
 import concurrent.futures
 from threading import Lock
 from concurrent.futures import ThreadPoolExecutor
-from jarvis.utils import OutputType, PrettyOutput, find_git_root, load_embedding_model
+from jarvis.utils import OutputType, PrettyOutput, find_git_root, get_max_context_length, get_thread_count, load_embedding_model, load_rerank_model
 from jarvis.utils import load_env_from_file
 import argparse
 from sentence_transformers import SentenceTransformer
@@ -18,8 +18,8 @@ class CodeBase:
         load_env_from_file()
         self.root_dir = root_dir
         os.chdir(self.root_dir)
-        self.thread_count = int(os.environ.get("JARVIS_THREAD_COUNT") or 1)
-        self.max_context_length = int(os.getenv("JARVIS_MAX_CONTEXT_LENGTH", 65536))
+        self.thread_count = get_thread_count()
+        self.max_context_length = get_max_context_length()
             
         # 初始化数据目录
         self.data_dir = os.path.join(self.root_dir, ".jarvis-codebase")
@@ -98,7 +98,7 @@ class CodeBase:
 5. 关键业务逻辑和处理流程
 6. 特殊功能点和亮点特性
 
-请用简洁专业的语言描述，突出代码的技术特征和功能特点，以便后续进行相似代码检索。
+请用简洁专业的语言描述，突出代码的技术特征和功能特点，以便后续进行关联代码检索。
 
 文件路径：{file_path}
 代码内容：
@@ -196,10 +196,15 @@ class CodeBase:
             if cached_vector is not None:
                 return cached_vector
                 
-            # 组合文件信息
+            # 读取文件内容并组合信息
+            with open(file_path, "r", encoding="utf-8") as f:
+                content = f.read()[:self.max_context_length]  # 限制文件内容长度
+            
+            # 组合文件信息，包含文件内容
             combined_text = f"""
 文件路径: {file_path}
 文件描述: {description}
+文件内容: {content}
 """
             vector = self.get_embedding(combined_text)
             
@@ -398,116 +403,126 @@ class CodeBase:
             PrettyOutput.print("没有检测到文件变更，无需重建索引", output_type=OutputType.INFO)
 
     def rerank_results(self, query: str, initial_results: List[Tuple[str, float, str]]) -> List[Tuple[str, float, str]]:
-        """使用大模型对搜索结果重新排序"""
+        """使用 BAAI/bge-reranker-v2-m3 对搜索结果重新排序"""
         if not initial_results:
             return []
-
-        model = PlatformRegistry.get_global_platform_registry().get_codegen_platform()
-        # model.set_suppress_output(True)
-
-        # 构建重排序的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)
-        
+            
+        try:
+            import torch
+            
+            # 加载模型和分词器
+            model, tokenizer = load_rerank_model()
+            
+            # 准备数据 - 加入文件内容进行更准确的重排序
+            pairs = []
+            for path, _, desc in initial_results:
+                try:
+                    with open(path, "r", encoding="utf-8") as f:
+                        content = f.read()[:512]  # 限制内容长度
+                    # 组合文件路径、描述和内容
+                    doc_content = f"文件: {path}\n描述: {desc}\n内容: {content}"
+                    pairs.append([query, doc_content])
+                except Exception as e:
+                    PrettyOutput.print(f"读取文件失败 {path}: {str(e)}", 
+                                    output_type=OutputType.ERROR)
+                    doc_content = f"文件: {path}\n描述: {desc}"
+                    pairs.append([query, doc_content])
+            
+            # 使用更大的batch size提高处理速度
+            batch_size = 16  # 根据GPU显存调整
+            
+            with torch.no_grad():
+                for i in range(0, len(pairs), batch_size):
+                    batch_pairs = pairs[i:i + batch_size]
+                    encoded = tokenizer(
+                        batch_pairs,
+                        padding=True,
+                        truncation=True,
+                        max_length=512,
+                        return_tensors='pt'
+                    )
+                    
+                    if torch.cuda.is_available():
+                        encoded = {k: v.cuda() for k, v in encoded.items()}
+                        
+                    outputs = model(**encoded)
+                    # 修改这里：直接使用 outputs.logits 作为分数
+                    batch_scores = outputs.logits.squeeze(-1).cpu().numpy()
+            
+            # 归一化分数到 0-1 范围
+            if batch_scores:
+                min_score = min(batch_scores)
+                max_score = max(batch_scores)
+                if max_score > min_score:
+                    batch_scores = [(s - min_score) / (max_score - min_score) for s in batch_scores]
+            
+            # 将分数与原始结果组合并排序
+            scored_results = []
+            for (path, _, desc), score in zip(initial_results, batch_scores):
+                if score >= 0.5:  # 只保留相关度大于 0.5 的结果
+                    scored_results.append((path, float(score), desc))
+                    
+            # 按分数降序排序
+            scored_results.sort(key=lambda x: x[1], reverse=True)
+            
+            return scored_results
+            
+        except Exception as e:
+            PrettyOutput.print(f"重排序失败，使用原始排序: {str(e)}", output_type=OutputType.WARNING)
+            return initial_results
 
     def search_similar(self, query: str, top_k: int = 30) -> List[Tuple[str, float, str]]:
-        """搜索相似文件"""
-        model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
-        model.set_suppress_output(True)
-        
-
-        prompt = f"""请根据以下查询，生成意思完全相同的另一个表述。这个表述将用于代码搜索，所以要保持专业性和准确性。
+        """搜索关联文件"""
+        try:
+            # 生成多个查询变体以提高召回率
+            model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
+            model.set_suppress_output(True)
+            
+            prompt = f"""请根据以下查询，生成3个不同的表述，每个表述都要完整表达原始查询的意思。这些表述将用于代码搜索，要保持专业性和准确性。
 原始查询: {query}
 
-请直接输出新表述，不要有编号或其他标记。
+请直接输出3个表述，用换行分隔，不要有编号或其他标记。
 """
-        
-        query = model.chat(prompt)
+            query_variants = model.chat(prompt).strip().split('\n')
+            query_variants.append(query)  # 添加原始查询
             
-
-
-        PrettyOutput.print(f"查询:  {query}", output_type=OutputType.INFO)
-        
-        # 为每个查询获取相似文件
-        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
+            # 对每个查询变体进行搜索
+            all_results = {}
+            for q in query_variants:
+                q_vector = self.get_embedding(q)
+                q_vector = q_vector.reshape(1, -1)
                 
-            similarity = 1.0 / (1.0 + float(distance))
-            # 只保留相似度大于等于0.5的结果
-            if similarity >= 0.5:
-                PrettyOutput.print(f"  {self.file_paths[i]} : 距离 {distance:.3f}, 相似度 {similarity:.3f}", 
-                                 output_type=OutputType.INFO)
+                distances, indices = self.index.search(q_vector, top_k)
                 
-                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)
+                for i, distance in zip(indices[0], distances[0]):
+                    if i == -1:
+                        continue
+                        
+                    similarity = 1.0 / (1.0 + float(distance))
+                    if similarity >= 0.5:
+                        file_path = self.file_paths[i]
+                        # 使用最高的相似度分数
+                        if file_path not in all_results or similarity > all_results[file_path][1]:
+                            data = self.vector_cache[file_path]
+                            all_results[file_path] = (file_path, similarity, data["description"])
+            
+            # 转换为列表并排序
+            results = list(all_results.values())
+            results.sort(key=lambda x: x[1], reverse=True)
+            
+            return results[:top_k]
+            
+        except Exception as e:
+            PrettyOutput.print(f"搜索失败: {str(e)}", output_type=OutputType.ERROR)
             return []
 
-        # 使用大模型重新排序
-        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)
+        if not results:
+            PrettyOutput.print("没有找到关联的文件", output_type=OutputType.WARNING)
+            return ""
+        
         PrettyOutput.print(f"找到的关联文件: ", output_type=OutputType.SUCCESS)
         for path, score, _ in results:
             PrettyOutput.print(f"文件: {path} 关联度: {score:.3f}", 

jarvis/jarvis_coder/main.py

@@ -6,7 +6,7 @@ from typing import Dict, Any, List, Tuple
 
 import yaml
 from jarvis.models.base import BasePlatform
-from jarvis.utils import OutputType, PrettyOutput, find_git_root, get_multiline_input, load_env_from_file
+from jarvis.utils import OutputType, PrettyOutput, find_git_root, get_max_context_length, get_multiline_input, load_env_from_file
 from jarvis.models.registry import PlatformRegistry
 from jarvis.jarvis_codebase.main import CodeBase
 from prompt_toolkit import PromptSession
@@ -22,7 +22,7 @@ class JarvisCoder:
     def __init__(self, root_dir: str, language: str):
         """初始化代码修改工具"""
 
-        self.max_context_length = int(os.getenv("JARVIS_MAX_CONTEXT_LENGTH", 65536))
+        self.max_context_length = get_max_context_length()
 
 
         self.root_dir = find_git_root(root_dir)
@@ -113,8 +113,8 @@ class JarvisCoder:
             return []
             
         try:
-            patches = re.findall(r'<PATCH_START>.*?<PATCH_END>', response, re.DOTALL)
-            return [patch.replace('<PATCH_START>', '').replace('<PATCH_END>', '').strip() 
+            patches = re.findall(r'<PATCH>.*?</PATCH>', response, re.DOTALL)
+            return [patch.replace('<PATCH>', '').replace('</PATCH>', '').strip() 
                    for patch in patches if patch.strip()]
         except Exception as e:
             PrettyOutput.print(f"解析patch失败: {str(e)}", OutputType.WARNING)
@@ -126,29 +126,57 @@ class JarvisCoder:
 
 修改格式说明：
 1. 每个修改块格式如下：
-<PATCH_START>
->>>>>> path/to/file
-要替换的内容
+<PATCH>
+> path/to/file
+def old_function():
+    print("old code")
+    return False
 =======
-新的内容
->>>>>>
-<PATCH_END>
+def old_function():
+    print("new code")
+    return True
+</PATCH>
 
 2. 如果是新文件或者替换整个文件内容，格式如下：
-<PATCH_START>
->>>>>> path/to/new/file
+<PATCH>
+> src/new_module.py
 =======
-新文件的完整内容
->>>>>>
-<PATCH_END>
+from typing import List
+
+def new_function():
+    return "This is a new file"
+</PATCH>
 
 3. 如果要删除文件中的某一段，格式如下：
-<PATCH_START>
->>>>>> path/to/file
-要删除的内容
+<PATCH>
+> path/to/file
+    # 这是要删除的注释
+    deprecated_code = True
+    if deprecated_code:
+        print("old feature")
 =======
->>>>>>
-<PATCH_END>
+</PATCH>
+
+4. 如果要修改导入语句，格式如下：
+<PATCH>
+> src/main.py
+from old_module import old_class
+=======
+from new_module import new_class
+</PATCH>
+
+5. 如果要修改类定义，格式如下：
+<PATCH>
+> src/models.py
+class OldModel:
+    def __init__(self):
+        self.value = 0
+=======
+class OldModel:
+    def __init__(self):
+        self.value = 1
+        self.name = "new"
+</PATCH>
 
 文件列表如下：
 """
@@ -158,27 +186,27 @@ class JarvisCoder:
                 continue
             prompt += f"""{i}. {file["file_path"]}\n"""
             prompt += f"""文件内容:\n"""
-            prompt += f"<FILE_CONTENT_START>\n"
+            prompt += f"<FILE_CONTENT>\n"
             prompt += f'{file["file_content"]}\n'
-            prompt += f"<FILE_CONTENT_END>\n"
+            prompt += f"</FILE_CONTENT>\n"
         
         prompt += f"\n需求描述: {feature}\n"
         prompt += """
 注意事项：
-1、仅输出补丁内容，不要输出任何其他内容，每个补丁必须用<PATCH_START>和<PATCH_END>标记
+1、仅输出补丁内容，不要输出任何其他内容，每个补丁必须用<PATCH>和</PATCH>标记
 2、如果在大段代码中有零星修改，生成多个补丁
 3、要替换的内容，一定要与文件内容完全一致，不要有任何多余或者缺失的内容
 4、每个patch不超过20行，超出20行，请生成多个patch
 """
-        
+
         success, response = self._call_model_with_retry(self.main_model, prompt)
         if not success:
             return []
             
         try:
             # 使用正则表达式匹配每个patch块
-            patches = re.findall(r'<PATCH_START>.*?<PATCH_END>', response, re.DOTALL)
-            return [patch.replace('<PATCH_START>', '').replace('<PATCH_END>', '').strip() 
+            patches = re.findall(r'<PATCH>.*?</PATCH>', response, re.DOTALL)
+            return [patch.replace('<PATCH>', '').replace('</PATCH>', '').strip() 
                    for patch in patches if patch.strip()]
         except Exception as e:
             PrettyOutput.print(f"解析patch失败: {str(e)}", OutputType.WARNING)
@@ -204,7 +232,7 @@ class JarvisCoder:
                     continue
                     
                 # 获取文件路径
-                file_path_match = re.search(r'>>>>>> (.*)', lines[0])
+                file_path_match = re.search(r'> (.*)', lines[0])
                 if not file_path_match:
                     error_info.append(f"无法解析文件路径: {lines[0]}")
                     return False, "\n".join(error_info)
@@ -220,7 +248,7 @@ class JarvisCoder:
                     return False, "\n".join(error_info)
                 
                 old_content = parts[0]
-                new_content = parts[1].split(">>>>>>")[0]
+                new_content = parts[1].split("</PATCH>")[0]
                 
                 # 处理新文件
                 if not old_content:
@@ -348,19 +376,59 @@ class JarvisCoder:
                     }
                 else:
                     self._revert_changes()
-                    return {
-                        "success": False,
-                        "stdout": "",
-                        "stderr": "修改被用户取消，文件未发生任何变化",
-                        "error": UserWarning("用户取消修改")
-                    }
+                    
+                    # 让用户输入调整意见
+                    user_feedback = get_multiline_input("""
+请提供修改建议，帮助生成更好的补丁：
+1. 修改的位置是否正确？
+2. 修改的内容是否合适？
+3. 是否有遗漏的修改？
+4. 其他调整建议？
+
+请输入调整意见(直接回车跳过):""")
+                    
+                    if not user_feedback:
+                        return {
+                            "success": False,
+                            "stdout": "",
+                            "stderr": "修改被用户取消，文件未发生任何变化",
+                            "error": UserWarning("用户取消修改")
+                        }
+                    
+                    retry_prompt = f"""补丁被用户拒绝，请根据用户意见重新生成补丁：
+
+用户意见：
+{user_feedback}
+
+请重新生成补丁，确保：
+1. 按照用户意见调整修改内容
+2. 准确定位要修改的代码位置
+3. 正确处理代码缩进
+4. 考虑代码上下文
+"""
+                    patches = self._remake_patch(retry_prompt)
+                    continue
             else:
-                PrettyOutput.print(f"补丁应用失败，请求重新生成: {error_info}", OutputType.WARNING)
-                retry_prompt = f"""补丁应用失败，请根据以下错误信息重新生成补丁：
+                PrettyOutput.print(f"补丁应用失败: {error_info}", OutputType.WARNING)
+                
+                # 让用户输入补充信息
+                user_info = get_multiline_input("""
+补丁应用失败。请提供更多信息来帮助修复问题：
+1. 是否需要调整代码位置？
+2. 是否有特殊的格式要求？
+3. 是否需要考虑其他文件的依赖？
+4. 其他补充说明？
+
+请输入补充信息(直接回车跳过):""")
+                
+                retry_prompt = f"""补丁应用失败，请根据以下信息重新生成补丁：
 
 错误信息：
 {error_info}
 
+用户补充信息：
+{user_info if user_info else "用户未提供补充信息"}
+
 请确保：
 1. 准确定位要修改的代码位置
 2. 正确处理代码缩进