jarvis-ai-assistant 0.1.134__py3-none-any.whl → 0.1.138__py3-none-any.whl

jarvis/jarvis_code_agent/code_agent.py

@@ -1,7 +1,12 @@
-import subprocess
+"""Jarvis代码代理模块。
+
+该模块提供CodeAgent类，用于处理代码修改任务。
+"""
+
 import os
+import sys
+import subprocess
 import argparse
-from token import OP
 from typing import Optional
 
 from yaspin import yaspin
@@ -13,41 +18,51 @@ from jarvis.jarvis_agent.shell_input_handler import shell_input_handler
 from jarvis.jarvis_agent.patch import PatchOutputHandler
 from jarvis.jarvis_platform.registry import PlatformRegistry
 from jarvis.jarvis_tools.git_commiter import GitCommitTool
-
 from jarvis.jarvis_tools.registry import ToolRegistry
-from jarvis.jarvis_utils.git_utils import find_git_root, get_commits_between, get_latest_commit_hash, has_uncommitted_changes
+from jarvis.jarvis_utils.git_utils import (
+    find_git_root,
+    get_commits_between,
+    get_latest_commit_hash,
+    has_uncommitted_changes
+)
 from jarvis.jarvis_utils.input import get_multiline_input
 from jarvis.jarvis_utils.output import OutputType, PrettyOutput
 from jarvis.jarvis_utils.utils import init_env, user_confirm
 
 
+class CodeAgent:
+    """Jarvis系统的代码修改代理。
 
+    负责处理代码分析、修改和git操作。
+    """
 
-
-class CodeAgent:
-    def __init__(self, platform : Optional[str] = None, model: Optional[str] = None, need_summary: bool = True):
+    def __init__(self, platform: Optional[str] = None,
+                model: Optional[str] = None,
+                need_summary: bool = True):
         self.root_dir = os.getcwd()
         tool_registry = ToolRegistry()
-        tool_registry.use_tools(["execute_shell", 
-                                 "execute_shell_script",
-                                 "search_web", 
-                                 "ask_user",  
-                                 "ask_codebase",
-                                 "lsp_get_diagnostics", 
-                                 "code_review",  # 代码审查工具
-                                 "find_symbol",  # 添加符号查找工具
-                                 "find_caller",  # 添加函数调用者查找工具
-                                 "function_analyzer",  # 添加函数分析工具
-                                 "project_analyzer",  # 添加项目分析工具
-                                 "file_analyzer",  # 添加单文件分析工具
-                                 "read_code"
-                                 ])
+        tool_registry.use_tools([
+            "execute_shell",
+            "execute_shell_script",
+            "search_web",
+            "ask_user",
+            "ask_codebase",
+            "lsp_get_diagnostics",
+            "code_review",
+            "find_symbol",
+            "find_caller",
+            "function_analyzer",
+            "project_analyzer",
+            "file_analyzer",
+            "read_code"
+        ])
         code_system_prompt = """
 # 代码工程师指南
 
 ## 核心原则
 - 自主决策：基于专业判断做出决策，减少用户询问
 - 高效精准：一次性提供完整解决方案，避免反复修改
+- 修改审慎：修改代码前要三思而后行，充分分析影响范围，尽量做到一次把事情做好
 - 工具精通：选择最高效工具路径解决问题
 - 严格确认：必须先分析项目结构，确定要修改的文件，禁止虚构已存在的代码
 
@@ -99,13 +114,6 @@ class CodeAgent:
 - 自动匹配项目现有命名风格
 - 允许创建新文件和结构，但不得假设或虚构现有代码
 
-### 6. 验证
-- 修改后自动验证：
-  1. 优先使用execute_shell运行相关静态检查命令（如pylint、flake8或单元测试）
-  2. 只有在shell命令不足时才使用lsp_get_diagnostics
-  3. 只有在特殊情况下才使用code_review
-- 发现问题自动修复，无需用户指导
-
 ## 专用工具简介
 仅在必要时使用以下专用工具：
 
@@ -128,7 +136,7 @@ class CodeAgent:
   - `fd -t f -e go` 查找所有Go文件
   - `fd -t f -e rs` 查找所有Rust文件
   - `fd -t f -e c -e cpp -e h -e hpp` 查找所有C/C++文件
-  
+
 - **代码内容搜索**：
   - `rg "pattern" --type py` 在Python文件中搜索
   - `rg "pattern" --type js` 在JavaScript文件中搜索
@@ -139,11 +147,7 @@ class CodeAgent:
   - `rg -w "word"` 精确匹配单词
 
 - **代码统计分析**：
-  - `loc <file_path>` 统计单个文件
-  - `loc --include="*.py"` 统计所有Python文件
-  - `loc --include="*.js" --include="*.ts"` 统计所有JavaScript/TypeScript文件
-  - `loc --exclude="test"` 排除测试文件
-  - `loc --sort=code` 按代码量排序
+  - `loc` 统计当前目录代码行数
 
 - **代码质量检查**：
   - Python: `pylint <file_path>`, `flake8 <file_path>`
@@ -174,27 +178,37 @@ class CodeAgent:
 - fd比find更快更易用，应优先使用
 - loc比wc -l提供更多代码统计信息，应优先使用
 - 针对不同编程语言选择对应的代码质量检查工具
+- 不要留下未实现的代码
 """
         # Dynamically add ask_codebase based on task complexity if really needed
         # 处理platform参数
-        platform_instance = (PlatformRegistry().create_platform(platform) 
-                            if platform 
-                            else PlatformRegistry().get_normal_platform())
+        platform_instance = (PlatformRegistry().create_platform(platform)
+                             if platform
+                             else PlatformRegistry().get_normal_platform())
         if model:
-            platform_instance.set_model_name(model) # type: ignore
-        
+            platform_instance.set_model_name(model)  # type: ignore
+
         self.agent = Agent(system_prompt=code_system_prompt,
                            name="CodeAgent",
-                           auto_complete=False, 
-                           output_handler=[tool_registry, PatchOutputHandler()], 
+                           auto_complete=False,
+                           output_handler=[tool_registry,
+                                           PatchOutputHandler()],
                            platform=platform_instance,
-                           input_handler=[shell_input_handler, file_input_handler, builtin_input_handler],
+                           input_handler=[
+                               shell_input_handler, file_input_handler, builtin_input_handler],
                            need_summary=need_summary)
+        self.agent.set_addon_prompt("请使用工具充分理解用户需求，然后根据需求一步步执行代码修改/开发")
+
+    def get_root_dir(self) -> str:
+        """获取项目根目录
 
-    
+        返回:
+            str: 项目根目录路径
+        """
+        return self.root_dir
 
     def _init_env(self):
-        with yaspin(text="正在初始化环境...", color="cyan") as spinner: 
+        with yaspin(text="正在初始化环境...", color="cyan") as spinner:
             curr_dir = os.getcwd()
             git_dir = find_git_root(curr_dir)
             self.root_dir = git_dir
@@ -205,62 +219,86 @@ class CodeAgent:
             spinner.text = "环境初始化完成"
             spinner.ok("✅")
 
-    
-
-    def run(self, user_input: str) :
-        """Run the code agent with the given user input.
-        
-        Args:
-            user_input: The user's requirement/request
-            
-        Returns:
-            str: Output describing the execution result
+    def _handle_uncommitted_changes(self):
+        """处理未提交的修改"""
+        if has_uncommitted_changes():
+            PrettyOutput.print("检测到未提交的修改，是否要提交？", OutputType.WARNING)
+            if user_confirm("是否要提交？", True):
+                git_commiter = GitCommitTool()
+                git_commiter.execute({})
+
+    def _show_commit_history(self, start_commit, end_commit):
+        """显示提交历史"""
+        if start_commit and end_commit:
+            commits = get_commits_between(start_commit, end_commit)
+        else:
+            commits = []
+
+        if commits:
+            commit_messages = "检测到以下提交记录:\n" + \
+                "\n".join(
+                    [f"- {commit_hash[:7]}: {message}" for commit_hash, message in commits])
+            PrettyOutput.print(commit_messages, OutputType.INFO)
+        return commits
+
+    def _handle_commit_confirmation(self, commits, start_commit):
+        """处理提交确认和可能的重置"""
+        if commits and user_confirm("是否接受以上提交记录？", True):
+            if len(commits) > 1 and user_confirm(
+                "是否要合并为一个更清晰的提交记录？", True
+            ):
+                # Reset to start commit
+                subprocess.run(
+                    ["git", "reset", "--mixed", start_commit],
+                    stdout=subprocess.DEVNULL,
+                    stderr=subprocess.DEVNULL,
+                    check=True
+                )
+                # Create new commit
+                git_commiter = GitCommitTool()
+                git_commiter.execute({})
+        elif start_commit:
+            os.system(f"git reset --hard {start_commit}")
+            PrettyOutput.print("已重置到初始提交", OutputType.INFO)
+
+    def run(self, user_input: str) -> Optional[str]:
+        """使用给定的用户输入运行代码代理。
+
+        参数:
+            user_input: 用户的需求/请求
+
+        返回:
+            str: 描述执行结果的输出，成功时返回None
         """
         try:
             self._init_env()
-            
             start_commit = get_latest_commit_hash()
-            
+
             try:
                 self.agent.run(user_input)
-            except Exception as e:
+            except RuntimeError as e:
                 PrettyOutput.print(f"执行失败: {str(e)}", OutputType.WARNING)
-            
+                return str(e)
+
+            self._handle_uncommitted_changes()
             end_commit = get_latest_commit_hash()
-            # Print commit history between start and end commits
-            if start_commit and end_commit:
-                commits = get_commits_between(start_commit, end_commit)
-            else:
-                commits = []
-            
-            if commits:
-                commit_messages = "检测到以下提交记录:\n" + "\n".join([f"- {commit_hash[:7]}: {message}" for commit_hash, message in commits])
-                PrettyOutput.print(commit_messages, OutputType.INFO)
-
-            if commits and user_confirm("是否接受以上提交记录？", True):
-                if len(commits) > 1 and user_confirm("是否要合并为一个更清晰的提交记录？", True):
-                    # Reset to start commit
-                    subprocess.run(["git", "reset", "--mixed", start_commit], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
-                    # Create new commit
-                    git_commiter = GitCommitTool()
-                    git_commiter.execute({})
-            elif start_commit:
-                os.system(f"git reset --hard {start_commit}")
-                PrettyOutput.print("已重置到初始提交", OutputType.INFO)
-                
-        except Exception as e:
+            commits = self._show_commit_history(start_commit, end_commit)
+            self._handle_commit_confirmation(commits, start_commit)
+            return None
+
+        except RuntimeError as e:
             return f"Error during execution: {str(e)}"
-        
 
 
 def main():
-    """Jarvis main entry point"""
-    # Add argument parser
+    """Jarvis主入口点。"""
     init_env()
 
     parser = argparse.ArgumentParser(description='Jarvis Code Agent')
-    parser.add_argument('-p', '--platform', type=str, help='Target platform name', default=None)
-    parser.add_argument('-m', '--model', type=str, help='Model name to use', default=None)
+    parser.add_argument('-p', '--platform', type=str,
+                      help='Target platform name', default=None)
+    parser.add_argument('-m', '--model', type=str,
+                      help='Model name to use', default=None)
     args = parser.parse_args()
 
     curr_dir = os.getcwd()
@@ -268,21 +306,18 @@ def main():
     PrettyOutput.print(f"当前目录: {git_dir}", OutputType.INFO)
 
     try:
-        try:
-            user_input = get_multiline_input("请输入你的需求（输入空行退出）:")
-            if not user_input:
-                return 0
-            agent = CodeAgent(platform=args.platform, model=args.model, need_summary=False)
-            agent.run(user_input)
-            
-        except Exception as e:
-            PrettyOutput.print(f"错误: {str(e)}", OutputType.ERROR)
+        user_input = get_multiline_input("请输入你的需求（输入空行退出）:")
+        if not user_input:
+            sys.exit(0)
+        agent = CodeAgent(platform=args.platform,
+                        model=args.model,
+                        need_summary=False)
+        agent.run(user_input)
 
-    except Exception as e:
-        PrettyOutput.print(f"初始化错误: {str(e)}", OutputType.ERROR)
-        return 1
+    except RuntimeError as e:
+        PrettyOutput.print(f"错误: {str(e)}", OutputType.ERROR)
+        sys.exit(1)
 
-    return 0
 
 if __name__ == "__main__":
-    exit(main())
+    main()

jarvis/jarvis_code_agent/file_select.py

@@ -11,15 +11,15 @@ from jarvis.jarvis_utils.utils import user_confirm
 
 def _parse_file_selection(input_str: str, max_index: int) -> List[int]:
     selected = set()
-    
+
     # Remove all whitespace characters
     input_str = "".join(input_str.split())
-    
+
     # Process comma-separated parts
     for part in input_str.split(","):
         if not part:
             continue
-        
+
         # Process range (e.g.: 3-6)
         if "-" in part:
             try:
@@ -41,7 +41,7 @@ def _parse_file_selection(input_str: str, max_index: int) -> List[int]:
                     PrettyOutput.print(f"忽略超出范围的索引: {part}", OutputType.WARNING)
             except ValueError:
                 PrettyOutput.print(f"忽略无效的数字: {part}", OutputType.WARNING)
-    
+
     return sorted(list(selected))
 
 def _get_file_completer(root_dir: str) -> Completer:
@@ -49,58 +49,58 @@ def _get_file_completer(root_dir: str) -> Completer:
     class FileCompleter(Completer):
         def __init__(self, root_dir: str):
             self.root_dir = root_dir
-            
+
         def get_completions(self, document, complete_event):
             text = document.text_before_cursor
-            
+
             if not text:
                 for path in self._list_files(""):
                     yield Completion(path, start_position=0)
                 return
-                
+
             # Generate fuzzy matching pattern
             pattern = '.*'.join(map(re.escape, text))
             try:
                 regex = re.compile(pattern, re.IGNORECASE)
             except re.error:
                 return
-                
+
             for path in self._list_files(""):
                 if regex.search(path):
                     yield Completion(path, start_position=-len(text))
-        
+
         def _list_files(self, current_dir: str) -> List[str]:
             """List all files in the specified directory (recursively)"""
             files = []
             search_dir = os.path.join(self.root_dir, current_dir)
-            
+
             for root, _, filenames in os.walk(search_dir):
                 for filename in filenames:
                     full_path = os.path.join(root, filename)
                     rel_path = os.path.relpath(full_path, self.root_dir)
                     if not any(part.startswith('.') for part in rel_path.split(os.sep)):
                         files.append(rel_path)
-            
+
             return sorted(files)
 
     return FileCompleter(root_dir)
 
 def _fuzzy_match_files(root_dir: str, pattern: str) -> List[str]:
     """Fuzzy match file path
-    
+
     Args:
         pattern: Matching pattern
-        
+
     Returns:
         List[str]: List of matching file paths
     """
     matches = []
-    
+
     # 将模式转换为正则表达式
     pattern = pattern.replace('.', r'\.').replace('*', '.*').replace('?', '.')
     pattern = f".*{pattern}.*"  # 允许部分匹配
     regex = re.compile(pattern, re.IGNORECASE)
-    
+
     # 遍历所有文件
     for root, _, files in os.walk(root_dir):
         for file in files:
@@ -110,7 +110,7 @@ def _fuzzy_match_files(root_dir: str, pattern: str) -> List[str]:
             if not any(part.startswith('.') for part in rel_path.split(os.sep)):
                 if regex.match(rel_path):
                     matches.append(rel_path)
-    
+
     return sorted(matches)
 
 def select_files(related_files: List[Dict[str, str]], root_dir: str) -> List[Dict[str, str]]:
@@ -127,7 +127,7 @@ def select_files(related_files: List[Dict[str, str]], root_dir: str) -> List[Dic
     if output:
         PrettyOutput.section("相关文件", OutputType.INFO)
         PrettyOutput.print(output, OutputType.INFO, lang="markdown")
-    
+
     if len(related_files) > 0:
         # Ask the user if they need to adjust the file list
         if user_confirm("是否需要调整文件列表？", False):
@@ -139,7 +139,7 @@ def select_files(related_files: List[Dict[str, str]], root_dir: str) -> List[Dic
                     selected_files = [related_files[i] for i in selected_indices]
                 else:
                     PrettyOutput.print("没有有效的文件被选择, 保持当前选择", OutputType.WARNING)
-    
+
     tips = ""
     # Ask if they need to supplement files
     if user_confirm("是否需要补充其他文件？", False):
@@ -154,23 +154,23 @@ def select_files(related_files: List[Dict[str, str]], root_dir: str) -> List[Dic
                 file_path = session.prompt(">>> ").strip()
             except KeyboardInterrupt:
                 break
-                
+
             if not file_path:
                 break
-                
+
             # Process wildcard matching
             if '*' in file_path or '?' in file_path:
                 matches = _fuzzy_match_files(root_dir, file_path)
                 if not matches:
                     PrettyOutput.print("没有找到匹配的文件", OutputType.WARNING)
                     continue
-                    
+
                 # Display matching files
                 tips = "找到以下匹配的文件:"
                 for i, path in enumerate(matches, 1):
                     tips += f"\n[{i}] {path}"
                 PrettyOutput.print(tips, OutputType.INFO)
-                    
+
                 # Let the user select
                 numbers = get_single_line_input("请选择要添加的文件编号（支持: 1,3-6格式, 按回车选择所有）").strip()
                 if numbers:
@@ -182,7 +182,7 @@ def select_files(related_files: List[Dict[str, str]], root_dir: str) -> List[Dic
                     paths_to_add = matches
             else:
                 paths_to_add = [file_path]
-            
+
             # Add selected files
             tips = "添加以下文件:"
             for path in paths_to_add:
@@ -190,7 +190,7 @@ def select_files(related_files: List[Dict[str, str]], root_dir: str) -> List[Dic
                 if not os.path.isfile(full_path):
                     tips += f"\n文件不存在: {path}"
                     continue
-                
+
                 try:
                     selected_files.append({"file": path, "reason": "I Added"})
                     tips += f"\n文件已添加: {path}"