jarvis-ai-assistant 0.1.98__py3-none-any.whl → 0.1.100__py3-none-any.whl

jarvis/jarvis_coder/patch_handler.py

@@ -1,40 +1,70 @@
 import re
 import os
 from typing import List, Tuple, Dict
+import yaml
+from pathlib import Path
 
+from jarvis.jarvis_coder.git_utils import generate_commit_message, init_git_repo, save_edit_record
 from jarvis.models.base import BasePlatform
 from jarvis.models.registry import PlatformRegistry
-from jarvis.utils import OutputType, PrettyOutput, get_multiline_input, while_success
+from jarvis.utils import OutputType, PrettyOutput, get_multiline_input, get_single_line_input, while_success
 
 class Patch:
-    def __init__(self, old_code: str, new_code: str):
-        self.old_code = old_code
-        self.new_code = new_code
+    def __init__(self, start: int, end: int, new_code: str):
+        self.start = start  # Line number where patch starts (inclusive)
+        self.end = end      # Line number where patch ends (exclusive) 
+        self.new_code = new_code  # New code to insert/replace
 
 class PatchHandler:
+    def __init__(self):
+        self.prompt_file = Path.home() / ".jarvis-coder-patch-prompt"
+        self.additional_info = self._load_additional_info()
+        self.root_dir = init_git_repo(os.getcwd())
+        self.record_dir = os.path.join(self.root_dir, ".jarvis-coder", "record")
+        if not os.path.exists(self.record_dir):
+            os.makedirs(self.record_dir)
+    def _load_additional_info(self) -> str:
+        """Load saved additional info from prompt file"""
+        if not self.prompt_file.exists():
+            return ""
+        try:
+            with open(self.prompt_file, 'r') as f:
+                data = yaml.safe_load(f)
+                return data.get('additional_info', '') if data else ''
+        except Exception as e:
+            PrettyOutput.print(f"Failed to load additional info: {e}", OutputType.WARNING)
+            return ""
 
+    def _save_additional_info(self, info: str):
+        """Save additional info to prompt file"""
+        try:
+            with open(self.prompt_file, 'w') as f:
+                yaml.dump({'additional_info': info}, f)
+        except Exception as e:
+            PrettyOutput.print(f"Failed to save additional info: {e}", OutputType.WARNING)
 
     def _extract_patches(self, response: str) -> List[Patch]:
-        """Extract patches from response
+        """Extract patches from response with hexadecimal line numbers
         
         Args:
             response: Model response content
             
         Returns:
-            List[Tuple[str, str, str]]: Patch list, each patch is a tuple of (format, file path, patch content)
+            List[Patch]: List of patches, each containing the line range and new code
         """
-        # 修改后的正则表达式匹配三种补丁格式
-        fmt_pattern = r'<PATCH>\n>>>>>> SEARCH\n(.*?)\n?(={5,})\n(.*?)\n?<<<<<< REPLACE\n</PATCH>'
+        fmt_pattern = r'<PATCH>\n\[([0-9a-f]+),([0-9a-f]+)\)\n(.*?\n)</PATCH>'
         ret = []
         for m in re.finditer(fmt_pattern, response, re.DOTALL):
-            ret.append(Patch(m.group(1), m.group(3)))   
+            start = int(m.group(1), 16)  # Convert hex to decimal
+            end = int(m.group(2), 16)
+            new_code = m.group(3)
+            ret.append(Patch(start, end, new_code))
         return ret
 
-
     def _confirm_and_apply_changes(self, file_path: str) -> bool:
         """Confirm and apply changes"""
         os.system(f"git diff --cached {file_path}")
-        confirm = input(f"\nAccept {file_path} changes? (y/n) [y]: ").lower() or "y"
+        confirm = get_single_line_input(f"Accept {file_path} changes? (y/n) [y]").lower() or "y"
         if confirm == "y":
             return True
         else:
@@ -43,87 +73,208 @@ class PatchHandler:
             os.system(f"git checkout -- {file_path}")
             PrettyOutput.print(f"Changes to {file_path} have been rolled back", OutputType.WARNING)
             return False
-
+        
     
+
+    def _finalize_changes(self) -> None:
+        """Complete changes and commit"""
+        PrettyOutput.print("Modification confirmed, committing...", OutputType.INFO)
+
+        # Add only modified files under git control
+        os.system("git add -u")
+        
+        # Then get git diff
+        git_diff = os.popen("git diff --cached").read()
+        
+        # Automatically generate commit information, pass in feature
+        commit_message = generate_commit_message(git_diff)
+        
+        # Display and confirm commit information
+        PrettyOutput.print(f"Automatically generated commit information: {commit_message}", OutputType.INFO)
+        user_confirm = get_single_line_input("Use this commit information? (y/n) [y]").lower() or "y"
+        
+        if user_confirm.lower() != "y":
+            commit_message = get_single_line_input("Please enter a new commit information")
+        
+        # No need to git add again, it has already been added
+        os.system(f"git commit -m '{commit_message}'")
+        save_edit_record(self.record_dir, commit_message, git_diff)
+
+    def _revert_changes(self) -> None:
+        """Revert all changes"""
+        PrettyOutput.print("Modification cancelled, reverting changes", OutputType.INFO)
+        os.system(f"git reset --hard")
+        os.system(f"git clean -df")
+
+    def _check_patches_overlap(self, patches: List[Patch]) -> bool:
+        """Check if any patches overlap with each other
+        
+        Args:
+            patches: List of patches to check
+            
+        Returns:
+            bool: True if patches overlap, False otherwise
+        """
+        if not patches:
+            return False
+        
+        # Sort patches by start line
+        sorted_patches = sorted(patches, key=lambda x: x.start)
+        
+        # Check for overlaps
+        for i in range(len(sorted_patches) - 1):
+            current = sorted_patches[i]
+            next_patch = sorted_patches[i + 1]
+            
+            if current.end > next_patch.start:
+                PrettyOutput.print(
+                    f"Overlapping patches detected: [{current.start:04x},{current.end:04x}) and [{next_patch.start:04x},{next_patch.end:04x})",
+                    OutputType.WARNING
+                )
+                return True
+            
+        return False
+
     def apply_file_patch(self, file_path: str, patches: List[Patch]) -> bool:
-        """Apply file patch"""
+        """Apply file patches using line numbers"""
         if not os.path.exists(file_path):
             base_dir = os.path.dirname(file_path)
             os.makedirs(base_dir, exist_ok=True)
             open(file_path, "w", encoding="utf-8").close()
-        file_content = open(file_path, "r", encoding="utf-8").read()
+        
+        # Check for overlapping patches
+        if self._check_patches_overlap(patches):
+            PrettyOutput.print("Cannot apply overlapping patches", OutputType.ERROR)
+            os.system(f"git reset {file_path}")
+            os.system(f"git checkout -- {file_path}")
+            return False
+        
+        with open(file_path, "r", encoding="utf-8") as f:
+            lines = f.readlines()
+        
+        # Sort patches by start line in reverse order to apply from bottom to top
+        patches.sort(key=lambda x: x.start, reverse=True)
+        
         for i, patch in enumerate(patches):
-            if patch.old_code == "" and patch.new_code == "":
-                PrettyOutput.print(f"Apply patch {i+1}/{len(patches)}: Delete file {file_path}", OutputType.INFO)
-                file_content = ""
-                os.system(f"git rm {file_path}")
-                PrettyOutput.print(f"Apply patch {i+1}/{len(patches)} successfully", OutputType.SUCCESS)
-            elif patch.old_code == "":
-                PrettyOutput.print(f"Apply patch {i+1}/{len(patches)}: Replace file {file_path} content: \n{patch.new_code}", OutputType.INFO)
-                file_content = patch.new_code
-                open(file_path, "w", encoding="utf-8").write(patch.new_code)
-                os.system(f"git add {file_path}")
-                PrettyOutput.print(f"Apply patch {i+1}/{len(patches)} successfully", OutputType.SUCCESS)
+            PrettyOutput.print(f"Applying patch {i+1}/{len(patches)} at lines [{patch.start},{patch.end})", OutputType.INFO)
+            
+            if patch.start > len(lines):
+                PrettyOutput.print(f"Invalid patch: start line {patch.start} exceeds file length {len(lines)}", OutputType.WARNING)
+                os.system(f"git reset {file_path}")
+                os.system(f"git checkout -- {file_path}")
+                return False
+            
+            if patch.new_code:
+                new_lines = patch.new_code.splitlines(keepends=True)
+                lines[patch.start:patch.end] = new_lines
             else:
-                PrettyOutput.print(f"Apply patch {i+1}/{len(patches)}: File original content: \n{patch.old_code}\nReplace with: \n{patch.new_code}", OutputType.INFO)
-                if file_content.find(patch.old_code) == -1:
-                    PrettyOutput.print(f"File {file_path} does not contain {patch.old_code}", OutputType.WARNING)
-                    os.system(f"git reset {file_path}")
-                    os.system(f"git checkout -- {file_path}")
-                    return False
-                else:
-                    file_content = file_content.replace(patch.old_code, patch.new_code, 1)
-                    open(file_path, "w", encoding="utf-8").write(file_content)
-                    os.system(f"git add {file_path}")
-                    PrettyOutput.print(f"Apply patch {i+1}/{len(patches)} successfully", OutputType.SUCCESS)
+                del lines[patch.start:patch.end]
+                
+        # Write modified content back to file
+        with open(file_path, "w", encoding="utf-8") as f:
+            f.writelines(lines)
+        
+        os.system(f"git add {file_path}")
+        PrettyOutput.print(f"Successfully applied all patches to {file_path}", OutputType.SUCCESS)
         return True
             
     
-    def retry_comfirm(self) -> Tuple[str, str]:# Restore user selection logic
-        choice = input("\nPlease choose an action: (1) Retry (2) Skip (3) Completely stop [1]: ") or "1"
+    def retry_comfirm(self) -> Tuple[str, str]:
+        choice = get_single_line_input("\nPlease choose an action: (1) Retry (2) Skip (3) Completely stop [1]: ") or "1"
         if choice == "2":
             return "skip", ""
         if choice == "3":
             return "break", ""
-        return "continue", get_multiline_input("Please enter additional information and requirements:")
+            
+        feedback = get_multiline_input("Please enter additional information and requirements:")
+        if feedback:
+            save_prompt = get_single_line_input("Would you like to save this as general feedback for future patches? (y/n) [n]: ").lower() or "n"
+            if save_prompt == "y":
+                self._save_additional_info(feedback)
+                PrettyOutput.print("Feedback saved for future use", OutputType.SUCCESS)
+            
+        return "continue", feedback
 
-    def apply_patch(self, feature: str, raw_plan: str, structed_plan: Dict[str, str]) -> Tuple[bool, str]:
+    def apply_patch(self, feature: str, structed_plan: Dict[str, str]) -> Tuple[bool, str]:
         """Apply patch (main entry)"""
+        feedback = ""
         for file_path, current_plan in structed_plan.items():
-            additional_info = ""            
+            additional_info = self.additional_info  # Initialize with saved info
             while True:
-                
                 if os.path.exists(file_path):
-                    content = open(file_path, "r", encoding="utf-8").read()
+                    # Read file and add line numbers
+                    lines = []
+                    with open(file_path, "r", encoding="utf-8") as f:
+                        for i, line in enumerate(f):
+                            lines.append(f"{i:04x}{line}")  # Changed from i+1 to i for 0-based indexing
+                    content = "".join(lines)
                 else:
                     content = "<File does not exist, need to create>"
+                    
                 prompt = """You are a senior software development expert who can generate code patches based on the complete modification plan, current original code file path, code content, and current file's modification plan. The output format should be as follows:
-                        <PATCH>
-                        >>>>>> SEARCH
-                        old_code
-                        ======
-                        new_code
-                        <<<<<< REPLACE
-                        </PATCH>
-                        Rules:
-                        1. When old_code is empty, it means replace everything from start to end
-                        2. When new_code is empty, it means delete old_code
-                        3. When both old_code and new_code are empty, it means delete the file
-                        Notes:
-                        1. Multiple patches can be generated
-                        2. old_code will be replaced with new_code, pay attention to context continuity
-                        3. Avoid breaking existing code logic when generating patches, e.g., don't insert function definitions inside existing function bodies
-                        4. Include sufficient context to avoid ambiguity
-                        5. Patches will be merged using file_content.replace(patch.old_code, patch.new_code, 1), so old_code and new_code need to match exactly, including empty lines, line breaks, whitespace, tabs, and comments
-                        6. Ensure generated code has correct format (syntax, indentation, line breaks)
-                        7. Ensure new_code's indentation and format matches old_code
-                        8. Ensure code is inserted in appropriate locations, e.g., code using variables should be after declarations/definitions
-                        9. Provide at least 3 lines of context before and after modified code for location
-
-
-                        """
+
+                <PATCH>
+                [start,end)
+                new_code
+                </PATCH>
+
+                Example:
+                <PATCH>
+                [0004,0004)
+                def new_function():
+    pass
+                </PATCH>
+
+                means:
+                Insert code BEFORE line 4:
+                ```
+                def new_function():
+    pass
+                ```
+                
+                Example 2:
+                <PATCH>
+                [0004,000b)
+                aa
+                bb
+                cc
+                </PATCH>
+
+                means:
+                Replace lines [4,11) with:
+                ```
+                aa
+                bb
+                cc
+                ```
+
+                Rules:
+                1. start and end are hexadecimal line numbers (e.g., 000a)
+                2. The patch will replace lines [start,end) with new_code (including start, excluding end)
+                3. If start equals end, new_code will be inserted BEFORE that line
+                4. If new_code is empty, lines [start,end) will be deleted
+                5. Multiple patches can be generated
+                6. Each line in the input file starts with its 4-digit hexadecimal line number (0-based)
+                7. Your new_code should NOT include line numbers
+                8. CRITICAL: Patches MUST NOT overlap - ensure each line is modified by at most one patch
+                9. Generate patches from bottom to top of the file
+                10. Ensure new_code maintains correct indentation and formatting
+                11. Each patch should modify no more than 20 lines
+                12. Include sufficient context in new_code to maintain code consistency
+                13. `[` and `)` must be included in the line range
+                14. Line numbers start from 0
+                15. Example of INVALID overlapping patches:
+                    <PATCH>
+                    [0001,0005)
+                    code1
+                    </PATCH>
+                    <PATCH>
+                    [0003,0007)  # This overlaps with the previous patch
+                    code2
+                    </PATCH>
+                """
+                
                 prompt += f"""# Original requirement: {feature}
-                    # Complete modification plan: {raw_plan}
                     # Current file path: {file_path}
                     # Current file content:
                     <CONTENT>
@@ -146,21 +297,25 @@ class PatchHandler:
                     act, msg = self.retry_comfirm()
                     if act == "break":
                         PrettyOutput.print("Terminate patch application", OutputType.WARNING)
-                        return False, msg
+                        additional_info = get_multiline_input("Please enter your additional information or suggestions (press Enter to cancel):")
+                        return False, additional_info
                     if act == "skip":
                         PrettyOutput.print(f"Skip file {file_path}", OutputType.WARNING)
+                        feedback += f"Skip file {file_path}\n"
+                        feedback += "Reason: " + get_multiline_input("Please enter your reason:") + "\n"
                         break
                     else:
                         additional_info += msg + "\n"
                         continue
                 else:
+                    self._finalize_changes()
                     break
         
-        return True, ""
+        return True, feedback
 
 
 
-    def handle_patch_application(self, feature: str, raw_plan: str, structed_plan: Dict[str,str]) -> bool:
+    def handle_patch_application(self, feature: str, structed_plan: Dict[str,str]) -> Tuple[bool, str]:
         """Process patch application process
         
         Args:
@@ -176,17 +331,10 @@ class PatchHandler:
             PrettyOutput.print(f"\nFile: {file_path}", OutputType.INFO)
             PrettyOutput.print(f"Modification plan: \n{patches_code}", OutputType.INFO)
         # 3. Apply patches
-        success, error_msg = self.apply_patch(feature, raw_plan, structed_plan)
+        success, additional_info = self.apply_patch(feature, structed_plan)
         if not success:
             os.system("git reset --hard")
-            return False
+            return False, additional_info
         # 6. Apply successfully, let user confirm changes
         PrettyOutput.print("\nPatches applied, please check the modification effect.", OutputType.SUCCESS)
-        confirm = input("\nKeep these changes? (y/n) [y]: ").lower() or "y"
-        if confirm != "y":
-            PrettyOutput.print("User cancelled changes, rolling back", OutputType.WARNING)
-            os.system("git reset --hard")  # Rollback all changes
-            return False
-        else:
-            return True
-            
+        return True, "Modification applied successfully"

jarvis/jarvis_coder/plan_generator.py

@@ -1,7 +1,7 @@
 import re
 from typing import Dict, List, Tuple
 from jarvis.models.registry import PlatformRegistry
-from jarvis.utils import PrettyOutput, OutputType, get_multiline_input
+from jarvis.utils import PrettyOutput, OutputType, get_multiline_input, get_single_line_input, is_long_context
 
 class PlanGenerator:
     """Modification plan generator"""   
@@ -48,7 +48,42 @@ class PlanGenerator:
         return prompt
     
     
-    def generate_plan(self, feature: str, related_files: List[Dict]) -> Tuple[str, Dict[str,str]]:
+    @staticmethod
+    def get_key_code(files: List[str], feature: str)->List[Dict[str, List[str]]]:
+        """Extract relevant key code snippets from files"""
+        ret = []
+        for file in files:
+            PrettyOutput.print(f"Analyzing file: {file}", OutputType.INFO)
+            model = PlatformRegistry.get_global_platform_registry().get_codegen_platform()
+            model.set_suppress_output(True)
+            file_path = file
+            content = open(file_path, "r", encoding="utf-8").read()
+
+            try:
+                prompt = f"""You are a code analysis expert who can extract relevant snippets from code.
+Please return in the following format:
+<PART>
+content
+</PART>
+
+Multiple snippets can be returned. If the file content is not relevant to the requirement, return empty.
+
+Requirement: {feature}
+File path: {file_path}
+Code content:
+{content}
+"""
+
+                # 调用大模型进行分析
+                response = model.chat_until_success(prompt)
+
+                parts = re.findall(r'<PART>\n(.*?)\n</PART>', response, re.DOTALL)
+                ret.append({"file_path": file, "parts": parts})
+            except Exception as e:
+                PrettyOutput.print(f"Failed to analyze file: {str(e)}", OutputType.ERROR)
+        return ret
+    
+    def generate_plan(self, feature: str, related_files: List[str]) -> Dict[str,str]:
         """Generate modification plan
         
         Args:
@@ -59,8 +94,15 @@ class PlanGenerator:
             Tuple[str, Dict[str,str]]: Modification plan, return None if user cancels
         """
         additional_info = ""
+        file_info = []
+        if is_long_context(related_files):
+            file_info =  PlanGenerator.get_key_code(related_files, feature)
+        else:
+            for file in related_files:
+                file_info.append({"file_path": file, "parts": [open(file, "r", encoding="utf-8").read()]})
+
         while True:
-            prompt = self._build_prompt(feature, related_files, additional_info)
+            prompt = self._build_prompt(feature, file_info, additional_info)
             # Build prompt
             PrettyOutput.print("Start generating modification plan...", OutputType.PROGRESS)
             
@@ -71,17 +113,17 @@ class PlanGenerator:
                 PrettyOutput.print("Modification plan generation failed, please try again", OutputType.ERROR)
                 tmp = get_multiline_input("Please enter your additional information or suggestions (press Enter to cancel):")
                 if tmp == "__interrupt__" or prompt == "":
-                    return "", {}
+                    return {}
                 additional_info += tmp + "\n"
                 continue
-            user_input = input("\nDo you agree with this modification plan? (y/n) [y]: ").strip().lower() or 'y'
+            user_input = get_single_line_input("Do you agree with this modification plan? (y/n) [y]").strip().lower() or 'y'
             if user_input == 'y' or user_input == '':
-                return raw_plan, structed_plan
+                return structed_plan
             elif user_input == 'n':
                 # Get user feedback
                 tmp = get_multiline_input("Please enter your additional information or suggestions (press Enter to cancel):")
                 if prompt == "__interrupt__" or prompt == "":
-                    return "", {}
+                    return {}
                 additional_info += tmp + "\n"
                 continue 
 

jarvis/jarvis_platform/main.py

@@ -1,5 +1,5 @@
 from jarvis.models.registry import PlatformRegistry
-from jarvis.utils import PrettyOutput, OutputType, load_env_from_file, get_multiline_input
+from jarvis.utils import PrettyOutput, OutputType, init_env, get_multiline_input
 
 def list_platforms():
     """List all supported platforms and models"""
@@ -106,7 +106,7 @@ def main():
     """Main function"""
     import argparse
 
-    load_env_from_file()
+    init_env()
     
     parser = argparse.ArgumentParser(description='Jarvis AI Platform')
     subparsers = parser.add_subparsers(dest='command', help='Available subcommands')

jarvis/jarvis_rag/main.py

@@ -4,7 +4,7 @@ import faiss
 from typing import List, Tuple, Optional, Dict
 import pickle
 from jarvis.utils import OutputType, PrettyOutput, get_file_md5, get_max_context_length, load_embedding_model, load_rerank_model
-from jarvis.utils import load_env_from_file
+from jarvis.utils import init_env
 from dataclasses import dataclass
 from tqdm import tqdm
 import fitz  # PyMuPDF for PDF files
@@ -90,7 +90,7 @@ class TextFileProcessor(FileProcessor):
                     continue
                     
             if not detected_encoding:
-                raise UnicodeDecodeError(f"Failed to decode file with supported encodings: {file_path}")
+                raise UnicodeDecodeError(f"Failed to decode file with supported encodings: {file_path}") # type: ignore
                 
             # Use the detected encoding to read the file
             with open(file_path, 'r', encoding=detected_encoding, errors='replace') as f:
@@ -114,7 +114,7 @@ class PDFProcessor(FileProcessor):
     @staticmethod
     def extract_text(file_path: str) -> str:
         text_parts = []
-        with fitz.open(file_path) as doc:
+        with fitz.open(file_path) as doc: # type: ignore
             for page in doc:
                 text_parts.append(page.get_text())
         return "\n".join(text_parts)
@@ -137,7 +137,7 @@ class RAGTool:
         Args:
             root_dir: Project root directory
         """
-        load_env_from_file()
+        init_env()
         self.root_dir = root_dir
         os.chdir(self.root_dir)
         
@@ -248,7 +248,7 @@ class RAGTool:
             
         # Create a flat index to store original vectors, for reconstruction
         self.flat_index = faiss.IndexFlatIP(self.vector_dim)
-        self.flat_index.add(vectors)
+        self.flat_index.add(vectors) # type: ignore
         
         # Create an IVF index for fast search
         nlist = max(4, int(vectors.shape[0] / 1000))  # 每1000个向量一个聚类中心
@@ -256,8 +256,8 @@ class RAGTool:
         self.index = faiss.IndexIVFFlat(quantizer, self.vector_dim, nlist, faiss.METRIC_INNER_PRODUCT)
         
         # Train and add vectors
-        self.index.train(vectors)
-        self.index.add(vectors)
+        self.index.train(vectors) # type: ignore
+        self.index.add(vectors) # type: ignore
         # Set the number of clusters to probe during search
         self.index.nprobe = min(nlist, 10)
 
@@ -341,7 +341,7 @@ class RAGTool:
         except Exception as e:
             PrettyOutput.print(f"Failed to get vector representation: {str(e)}", 
                             output_type=OutputType.ERROR)
-            return np.zeros((len(texts), self.vector_dim), dtype=np.float32)
+            return np.zeros((len(texts), self.vector_dim), dtype=np.float32) # type: ignore
 
     def _process_document_batch(self, documents: List[Document]) -> List[np.ndarray]:
         """Process a batch of documents vectorization
@@ -361,7 +361,7 @@ Content: {doc.content}
 """
             texts.append(combined_text)
             
-        return self._get_embedding_batch(texts)
+        return self._get_embedding_batch(texts) # type: ignore
 
     def _process_file(self, file_path: str) -> List[Document]:
         """Process a single file"""
@@ -516,7 +516,7 @@ Content: {doc.content}
                                 if d.metadata['file_path'] == doc.metadata['file_path']), None)
                     if doc_idx is not None:
                         # Reconstruct vectors from flat index
-                        vector = np.zeros((1, self.vector_dim), dtype=np.float32)
+                        vector = np.zeros((1, self.vector_dim), dtype=np.float32) # type: ignore
                         self.flat_index.reconstruct(doc_idx, vector.ravel())
                         unchanged_vectors.append(vector)
                 
@@ -585,7 +585,7 @@ Content: {doc.content}
         
         # Initial search more results for MMR
         initial_k = min(top_k * 2, len(self.documents))
-        distances, indices = self.index.search(query_vector, initial_k)
+        distances, indices = self.index.search(query_vector, initial_k) # type: ignore
         
         # Get valid results
         valid_indices = indices[0][indices[0] != -1]

jarvis/jarvis_smart_shell/main.py

@@ -4,11 +4,11 @@ import os
 import sys
 import readline
 from typing import Optional
-from yaspin import yaspin
-from yaspin.spinners import Spinners
+from yaspin import yaspin # type: ignore
+from yaspin.spinners import Spinners # type: ignore
 
 from jarvis.models.registry import PlatformRegistry
-from jarvis.utils import PrettyOutput, OutputType, load_env_from_file
+from jarvis.utils import PrettyOutput, OutputType, init_env
 
 def execute_command(command: str) -> None:
     """Show command and allow user to edit, then execute, Ctrl+C to cancel"""
@@ -27,6 +27,7 @@ def execute_command(command: str) -> None:
     except Exception as e:
         PrettyOutput.print(f"Failed to execute command: {str(e)}", OutputType.ERROR)
 
+
 def process_request(request: str) -> Optional[str]:
     """Process user request and return corresponding shell command
     
@@ -38,7 +39,6 @@ def process_request(request: str) -> Optional[str]:
     """
     try:
         # Get language model instance
-        PlatformRegistry.suppress_output = True
         model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
         model.set_suppress_output(True)
 
@@ -90,7 +90,7 @@ Remember: Only return the command itself, without any additional content.
 
 def main():
     # 创建参数解析器
-    load_env_from_file()
+    init_env()
     parser = argparse.ArgumentParser(
         description="Convert natural language requirements to shell commands",
         formatter_class=argparse.RawDescriptionHelpFormatter,

jarvis/models/base.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+import re
 from typing import Dict, List, Tuple
 
 from jarvis.utils import OutputType, PrettyOutput, while_success, while_true
@@ -26,7 +27,11 @@ class BasePlatform(ABC):
         raise NotImplementedError("chat is not implemented")
 
     def chat_until_success(self, message: str) -> str:
-        return while_true(lambda: while_success(lambda: self.chat(message), 5), 5)
+        def _chat():
+            response = self.chat(message)
+            response = re.sub(r'<think>(.*?)</think>', '', response, flags=re.DOTALL)
+            return response
+        return while_true(lambda: while_success(lambda: _chat(), 5), 5)
 
     @abstractmethod
     def upload_files(self, file_list: List[str]) -> List[Dict]: