jarvis-ai-assistant 0.1.102__py3-none-any.whl → 0.1.104__py3-none-any.whl

jarvis/{jarvis_coder → jarvis_code_agent}/file_select.py

@@ -1,10 +1,9 @@
-
 import os
 import re
 from typing import Dict, List
 from prompt_toolkit import PromptSession
 from prompt_toolkit.completion import WordCompleter, Completer, Completion
-from jarvis.utils import OutputType, PrettyOutput, get_single_line_input
+from jarvis.utils import OutputType, PrettyOutput, get_single_line_input, user_confirm
 
 
 def _parse_file_selection(input_str: str, max_index: int) -> List[int]:
@@ -64,25 +63,23 @@ def _get_file_completer(root_dir: str) -> Completer:
             self.root_dir = root_dir
             
         def get_completions(self, document, complete_event):
-            # Get the text of the current input
             text = document.text_before_cursor
             
-            # If the input is empty, return all files in the root directory
             if not text:
                 for path in self._list_files(""):
                     yield Completion(path, start_position=0)
                 return
                 
-            # Get the current directory and partial file name
-            current_dir = os.path.dirname(text)
-            file_prefix = os.path.basename(text)
-            
-            # List matching files
-            search_dir = os.path.join(self.root_dir, current_dir) if current_dir else self.root_dir
-            if os.path.isdir(search_dir):
-                for path in self._list_files(current_dir):
-                    if path.startswith(text):
-                        yield Completion(path, start_position=-len(text))
+            # 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)"""
@@ -93,7 +90,6 @@ def _get_file_completer(root_dir: str) -> Completer:
                 for filename in filenames:
                     full_path = os.path.join(root, filename)
                     rel_path = os.path.relpath(full_path, self.root_dir)
-                    # Ignore .git directory and other hidden files
                     if not any(part.startswith('.') for part in rel_path.split(os.sep)):
                         files.append(rel_path)
             
@@ -133,14 +129,16 @@ def select_files(related_files: List[str], root_dir: str) -> List[str]:
     """Let the user select and supplement related files"""
     PrettyOutput.section("Related files", OutputType.INFO)
     
+    output = ""
     # Display found files
     selected_files = list(related_files)  # Default select all
     for i, file in enumerate(related_files, 1):
-        PrettyOutput.print(f"[{i}] {file}", OutputType.INFO)
+        output += f"[{i}] {file}\n"
+
+    PrettyOutput.print(output, OutputType.INFO, lang="markdown")
     
     # Ask the user if they need to adjust the file list
-    user_input = get_single_line_input("Do you need to adjust the file list? (y/n) [n]").strip().lower() or 'n'
-    if user_input == 'y':
+    if user_confirm("Do you need to adjust the file list?", False):
         # Let the user select files
         numbers = get_single_line_input("Please enter the file numbers to include (support: 1,3-6 format, press Enter to keep the current selection)").strip()
         if numbers:
@@ -151,8 +149,7 @@ def select_files(related_files: List[str], root_dir: str) -> List[str]:
                 PrettyOutput.print("No valid files selected, keep the current selection", OutputType.WARNING)
     
     # Ask if they need to supplement files
-    user_input = get_single_line_input("Do you need to supplement other files? (y/n) [n]").strip().lower() or 'n'
-    if user_input == 'y':
+    if user_confirm("Do you need to supplement other files?", False):
         # Create file completion session
         session = PromptSession(
             completer=_get_file_completer(root_dir),
@@ -160,7 +157,7 @@ def select_files(related_files: List[str], root_dir: str) -> List[str]:
         )
         
         while True:
-            PrettyOutput.print("\nPlease enter the file path to supplement (support Tab completion and *? wildcard, input empty line to end):", OutputType.INFO)
+            PrettyOutput.print("Please enter the file path to supplement (support Tab completion and *? wildcard, input empty line to end):", OutputType.INFO)
             try:
                 file_path = session.prompt(">>> ").strip()
             except KeyboardInterrupt:
@@ -177,7 +174,7 @@ def select_files(related_files: List[str], root_dir: str) -> List[str]:
                     continue
                     
                 # Display matching files
-                PrettyOutput.print("\nFound the following matching files:", OutputType.INFO)
+                PrettyOutput.print("Found the following matching files:", OutputType.INFO)
                 for i, path in enumerate(matches, 1):
                     PrettyOutput.print(f"[{i}] {path}", OutputType.INFO)
                     

jarvis/jarvis_code_agent/patch.py

@@ -0,0 +1,120 @@
+import re
+from typing import Dict, Any, List, Tuple
+import os
+from jarvis.tools.git_commiter import GitCommitTool
+from jarvis.tools.read_code import ReadCodeTool
+from jarvis.utils import OutputType, PrettyOutput, has_uncommitted_changes, make_choice_input, user_confirm
+
+
+def _parse_patch(patch_str: str) -> Dict[str, List[Dict[str, Any]]]:
+    """Parse patches from string with format:
+    <PATCH>
+    > /path/to/file start_line,end_line
+    content_line1
+    content_line2
+    ...
+    </PATCH>
+    """
+    result = {}
+    patches = re.findall(r"<PATCH>(.*?)</PATCH>", patch_str, re.DOTALL)
+    
+    for patch in patches:
+        lines = patch.strip().split('\n')
+        if not lines:
+            continue
+            
+        # Parse file path and line range
+        file_info = lines[0].strip()
+        if not file_info.startswith('>'):
+            continue
+            
+        # Extract file path and line range
+        match = re.match(r'>\s*([^\s]+)\s+(\d+),(\d+)', file_info)
+        if not match:
+            continue
+            
+        filepath = match.group(1).strip()
+        start_line = int(match.group(2))
+        end_line = int(match.group(3))
+        
+        # Get content lines (skip the first line with file info)
+        content = '\n'.join(lines[1:])
+
+        if filepath not in result:
+            result[filepath] = []
+        
+        # Store in result dictionary
+        result[filepath].append({
+            'start_line': start_line,
+            'end_line': end_line,
+            'content': content
+        })
+    
+    # Sort patches by start line in reverse order to apply from bottom to top
+    for filepath in result:
+        result[filepath].sort(key=lambda x: x['start_line'], reverse=True)
+    
+    return result
+
+
+def apply_patch(output_str: str)->str:
+    """Apply patches to files"""
+    patches = _parse_patch(output_str)
+    if not patches:
+        return ""
+        
+    for filepath, patch_info in patches.items():
+        try:
+            # Check if file exists
+            if not os.path.exists(filepath):
+                PrettyOutput.print(f"File not found: {filepath}", OutputType.WARNING)
+                continue
+                
+            # Read original file content
+            lines = open(filepath, 'r', encoding='utf-8').readlines()
+            
+            # Apply patch
+            for patch in patch_info:
+                start_line = patch['start_line']
+                end_line = patch['end_line']
+                new_content = patch['content'].splitlines(keepends=True)
+
+                if new_content and new_content[-1] and new_content[-1][-1] != '\n':
+                    new_content[-1] += '\n'
+                # Validate line numbers
+                if start_line < 0 or end_line > len(lines) + 1 or start_line > end_line:
+                    PrettyOutput.print(f"Invalid line range [{start_line}, {end_line}) for file: {filepath}", OutputType.WARNING)
+                    continue
+                    
+                # Create new content
+                lines[start_line:end_line] = new_content
+                
+                # Write back to file
+                open(filepath, 'w', encoding='utf-8').writelines(lines)
+
+                PrettyOutput.print(f"Applied patch to {filepath} successfully\n", OutputType.SUCCESS)
+            
+        except Exception as e:
+            PrettyOutput.print(f"Error applying patch to {filepath}: {str(e)}", OutputType.ERROR)
+            continue
+    
+    if has_uncommitted_changes():
+        handle_commit_workflow()
+    return ""
+    
+def handle_commit_workflow()->bool:
+    """Handle the git commit workflow and return the commit details.
+    
+    Returns:
+        tuple[bool, str, str]: (continue_execution, commit_id, commit_message)
+    """
+    diff = os.popen("git diff HEAD").read()
+    PrettyOutput.print(diff, OutputType.CODE, lang="diff")
+    if not user_confirm("Do you want to commit the code?", default=True):
+        os.system("git reset HEAD")
+        os.system("git checkout -- .")
+        return False
+
+    git_commiter = GitCommitTool()
+    commit_result = git_commiter.execute({})
+    return commit_result["success"]

jarvis/jarvis_code_agent/relevant_files.py

@@ -0,0 +1,97 @@
+import os
+import re
+from typing import List
+
+import yaml
+from jarvis.agent import Agent
+from jarvis.jarvis_code_agent.file_select import select_files
+from jarvis.jarvis_codebase.main import CodeBase
+from jarvis.models.registry import PlatformRegistry
+from jarvis.tools.registry import ToolRegistry
+from jarvis.utils import OutputType, PrettyOutput, is_disable_codebase
+
+
+def find_relevant_files(user_input: str, root_dir: str) -> List[str]:
+    try:
+        files_from_codebase = []
+        if not is_disable_codebase():
+            PrettyOutput.print("Find files from codebase...", OutputType.INFO)
+            codebase = CodeBase(root_dir)
+            files_from_codebase = codebase.search_similar(user_input)
+
+        PrettyOutput.print("Find files by agent...", OutputType.INFO)
+        find_file_tool_registry = ToolRegistry()
+        find_file_tool_registry.use_tools(["read_code", "execute_shell"])
+        find_file_agent = Agent(
+            system_prompt="""You are a file agent, you are responsible for finding files related to the user's requirement.
+
+        SEARCH STRATEGY:
+        1. First Pass - Quick Search:
+           - Use `execute_shell` with git grep/find to locate potential files
+           - Search for key terms, function names, and relevant patterns
+           - Example: execute_shell("git grep -l 'search_term'")
+
+        2. Content Analysis:
+           - For each potential file, analyze its content
+           - Follow the file reading guidelines for large files
+           - Look for:
+             * Direct matches to requirement terms
+             * Related functionality
+             * Imported/referenced files
+             * Test files for modified code
+
+        FILE READING GUIDELINES:
+        1. For Large Files (>200 lines):
+           - Do NOT read the entire file at once
+           - First use grep/ctags to locate relevant sections
+           - Then read specific sections with context
+           - Example:
+             * execute_shell("grep -n 'function_name' path/to/file")
+             * read_code("path/to/file", start_line=found_line-10, end_line=found_line+20)
+
+        2. For Small Files:
+           - Can read entire file directly
+
+        IMPORTANT RULES:
+        - Only return files that are DIRECTLY related to the requirement
+        - Exclude false positives and loosely related files
+        - If a file only contains imports/references, don't include it
+        - Include both implementation and test files when relevant
+        - If unsure about a file, use grep/read_code to verify relevance
+        - Return empty list if no truly relevant files are found
+        - Do NOT modify any code, only find files
+
+        OUTPUT FORMAT:
+        - Only provide file paths in the specified YAML format
+        - No additional explanations or comments
+        """, 
+            name="FindFileAgent", 
+            is_sub_agent=True,
+            tool_registry=find_file_tool_registry, 
+            platform=PlatformRegistry().get_normal_platform(),
+            auto_complete=True,
+            summary_prompt="""Please provide the file paths as YAML list:
+            <FILE_PATH>
+            - file_path1
+            - file_path2
+            </FILE_PATH>
+            """)
+        prompt = f"Find files related about '{user_input}'\n"
+        if files_from_codebase:
+            prompt += f"\n\nFiles maybe related: {files_from_codebase}\n\n Please read above files first"
+        output = find_file_agent.run(prompt)
+
+        rsp_from_agent = re.findall(r'<FILE_PATH>(.*?)</FILE_PATH>', output, re.DOTALL)
+        files_from_agent = []
+        if rsp_from_agent:
+            try:
+                files_from_agent = yaml.safe_load(rsp_from_agent[0])
+            except Exception as e:
+                files_from_agent = []
+        else:
+            files_from_agent = []
+
+        selected_files = select_files(files_from_agent, os.getcwd())
+        return selected_files
+    except Exception as e:
+        return []