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

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_rag/main.py

@@ -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)
@@ -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,24 +4,23 @@ 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, get_single_line_input, load_env_from_file
 
 def execute_command(command: str) -> None:
     """Show command and allow user to edit, then execute, Ctrl+C to cancel"""
     try:
-        print("\nGenerated command (can be edited, press Enter to execute, Ctrl+C to cancel):")
         # Pre-fill input line
         readline.set_startup_hook(lambda: readline.insert_text(command))
         try:
-            edited_command = input("> ")
+            edited_command = get_single_line_input("Generated command (can be edited, press Enter to execute, Ctrl+C to cancel)")
             if edited_command.strip():  # Ensure command is not empty
                 os.system(edited_command)
         except KeyboardInterrupt:
-            print("\nExecution cancelled")
+            PrettyOutput.print("Execution cancelled", OutputType.INFO)
         finally:
             readline.set_startup_hook()  # Clear pre-filled
     except Exception as e:
@@ -38,7 +37,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)
 

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]:

jarvis/models/ollama.py

@@ -1,7 +1,7 @@
 import requests
 from typing import List, Dict, Tuple
 from jarvis.models.base import BasePlatform
-from jarvis.utils import OutputType, PrettyOutput
+from jarvis.utils import OutputType, PrettyOutput, get_single_line_input
 import os
 import json
 
@@ -136,7 +136,7 @@ if __name__ == "__main__":
         ollama = OllamaPlatform()
         while True:
             try:
-                message = input("\nInput question (Ctrl+C to exit): ")
+                message = get_single_line_input("\nInput question (Ctrl+C to exit)")
                 ollama.chat_until_success(message)
             except KeyboardInterrupt:
                 print("\nGoodbye!")

jarvis/models/registry.py

@@ -3,8 +3,8 @@ import inspect
 import os
 import sys
 from typing import Dict, Type, Optional, List
-from .base import BasePlatform
-from ..utils import PrettyOutput, OutputType
+from jarvis.models.base import BasePlatform
+from jarvis.utils import PrettyOutput, OutputType
 
 REQUIRED_METHODS = [
     ('chat', ['message']),           # 方法名和参数列表
@@ -23,7 +23,6 @@ class PlatformRegistry:
 
     global_platform_name = "kimi"
     global_platform_registry = None
-    suppress_output = False
 
     @staticmethod
     def get_platform_dir() -> str:
@@ -117,7 +116,7 @@ class PlatformRegistry:
                         module = importlib.import_module(module_name)
                     
                     # 遍历模块中的所有类
-                    for name, obj in inspect.getmembers(module):
+                    for _, obj in inspect.getmembers(module):
                         # 检查是否是BasePlatform的子类，但不是BasePlatform本身
                         if (inspect.isclass(obj) and 
                             issubclass(obj, BasePlatform) and 
@@ -126,8 +125,6 @@ class PlatformRegistry:
                             # 检查平台实现
                             if not PlatformRegistry.check_platform_implementation(obj):
                                 continue
-                            if not PlatformRegistry.suppress_output:
-                                PrettyOutput.print(f"Load platform from {os.path.join(directory, filename)}: {obj.platform_name}", OutputType.SUCCESS) # type: ignore
                             platforms[obj.platform_name] = obj # type: ignore
                             break
                 except Exception as e:

jarvis/tools/ask_user.py

@@ -18,22 +18,22 @@ class AskUserTool:
         
 
     def execute(self, args: Dict[str, Any]) -> Dict[str, Any]:
-        """执行询问用户操作
+        """Execute the operation of asking the user
         
         Args:
-            args: 包含问题的字典
+            args: A dictionary containing the question
             
         Returns:
-            Dict: 包含用户响应的字典
+            Dict: A dictionary containing the user's response
         """
         try:
             question = args["question"]
             
-            # 显示问题
-            PrettyOutput.print("\n问题:", OutputType.SYSTEM)
+            # Display the question
+            PrettyOutput.print("\nQuestion:", OutputType.SYSTEM)
             PrettyOutput.print(question, OutputType.SYSTEM)
             
-            # 获取用户输入
+            # Get user input
             user_response = get_multiline_input("Please enter your answer (input empty line to end)")
             
             if user_response == "__interrupt__":

jarvis/tools/codebase_qa.py

@@ -4,7 +4,7 @@ from jarvis.jarvis_codebase.main import CodeBase
 from jarvis.utils import find_git_root, PrettyOutput, OutputType
 
 class CodebaseQATool:
-    """代码库问答工具，用于回答关于代码库的问题"""
+    """Codebase QA Tool"""
     
     name = "codebase_qa"
     description = "Answer questions about the codebase, can query and understand code functionality, structure, and implementation details"
@@ -29,7 +29,7 @@ class CodebaseQATool:
     }
     
     def execute(self, params: Dict[str, Any]) -> Dict[str, Any]:
-        """执行代码问答"""
+        """Execute codebase QA"""
         try:
             dir = params.get("dir")
             question = params["question"]
@@ -42,7 +42,7 @@ class CodebaseQATool:
                 return {
                     "success": False,
                     "stdout": "",
-                    "stderr": "错误：当前目录不在Git仓库中",
+                    "stderr": "Error: Current directory is not in a Git repository",
                     "error": "NotInGitRepository"
                 }
 
@@ -62,13 +62,11 @@ class CodebaseQATool:
             }
             
         except Exception as e:
-            PrettyOutput.print(f"代码问答出错: {str(e)}", output_type=OutputType.ERROR)
+            PrettyOutput.print(f"Codebase QA error: {str(e)}", output_type=OutputType.ERROR)
             return {
                 "success": False,
                 "stdout": "",
-                "stderr": f"执行代码问答时发生错误: {str(e)}",
+                "stderr": f"Error executing codebase QA: {str(e)}",
                 "error": str(type(e).__name__)
             }
 
-def register():
-    return CodebaseQATool() 

jarvis/tools/create_code_sub_agent.py

@@ -0,0 +1,55 @@
+from typing import Dict, Any
+
+from jarvis.agent import Agent
+from jarvis.utils import OutputType, PrettyOutput
+from jarvis.jarvis_code_agent.main import system_prompt
+
+
+class CodeSubAgentTool:
+    name = "create_code_sub_agent"
+    description = "Create a sub-agent to handle specific code development subtasks"
+    parameters = {
+        "type": "object",
+        "properties": {
+            "name": {
+                "type": "string",
+                "description": "The name of the sub-agent"
+            },
+            "subtask": {
+                "type": "string",
+                "description": "The specific code development subtask to complete"
+            },
+        },
+        "required": ["subtask", "name"]
+    }
+
+    def execute(self, args: Dict) -> Dict[str, Any]:
+        """Execute code development subtask"""
+        try:
+            subtask = args["subtask"]
+            name = args["name"]
+
+            PrettyOutput.print(f"Creating code sub-agent {name} for subtask: {subtask}", OutputType.INFO)
+
+            # Create sub-agent
+            sub_agent = Agent(
+                system_prompt=system_prompt,
+                name=name,
+                is_sub_agent=True
+            )
+
+            # Execute subtask
+            result = sub_agent.run(subtask)
+
+            return {
+                "success": True,
+                "stdout": f"Code Development Subtask Results:\n\n{result}",
+                "stderr": ""
+            }
+
+        except Exception as e:
+            PrettyOutput.print(str(e), OutputType.ERROR)
+            return {
+                "success": False,
+                "error": f"Failed to execute code development subtask: {str(e)}"
+            }

jarvis/tools/{sub_agent.py → create_sub_agent.py}

@@ -1,7 +1,7 @@
 from typing import Dict, Any
 
 
-from jarvis.agent import Agent
+from jarvis.agent import Agent, origin_agent_system_prompt
 from jarvis.utils import OutputType, PrettyOutput
 
 
@@ -58,8 +58,11 @@ class SubAgentTool:
             if goal:
                 task_description += f"\n\nCompletion goal:\n{goal}"
 
+            
+
             # Create sub-agent
             sub_agent = Agent(
+                system_prompt=origin_agent_system_prompt,
                 name=agent_name,
                 is_sub_agent=True
             )

jarvis/tools/execute_code_modification.py

@@ -0,0 +1,72 @@
+from typing import Dict, Any
+
+from jarvis.agent import Agent
+from jarvis.utils import OutputType, PrettyOutput
+from jarvis.jarvis_coder.patch_handler import PatchHandler
+
+
+class CodeModifyTool:
+    name = "execute_code_modification"
+    description = "Execute code modifications according to the provided plan"
+    parameters = {
+        "type": "object",
+        "properties": {
+            "task": {
+                "type": "string",
+                "description": "The code modification task description"
+            },
+            "structured_plan": {
+                "type": "object",
+                "description": "Dictionary mapping file paths to their modification plans. Example: {'path/to/file.py': 'Add function foo() to handle...'}", 
+                "additionalProperties": {
+                    "type": "string",
+                    "description": "Modification plan for a specific file"
+                },
+                "examples": [{
+                    "src/file1.py": "Add error handling to process_data()",
+                    "src/file2.py": "Update API endpoint URL in get_data()"
+                }]
+            }
+        },
+        "required": ["task", "raw_plan", "structured_plan"]
+    }
+
+    def execute(self, args: Dict) -> Dict[str, Any]:
+        """Execute code modifications using PatchHandler"""
+        try:
+            task = args["task"]
+            structured_plan = args["structured_plan"]
+
+            PrettyOutput.print("Executing code modifications...", OutputType.INFO)
+
+            # Create patch handler instance
+            patch_handler = PatchHandler()
+
+            # Apply patches and handle the process
+            success = patch_handler.handle_patch_application(
+                feature=task,
+                structed_plan=structured_plan
+            )
+
+            if not success:
+                return {
+                    "success": False,
+                    "error": "Code modification was cancelled or failed",
+                    "stdout": "Changes have been rolled back",
+                    "stderr": ""
+                }
+
+            return {
+                "success": True,
+                "stdout": "Code modifications have been successfully applied and committed",
+                "stderr": ""
+            }
+
+        except Exception as e:
+            PrettyOutput.print(str(e), OutputType.ERROR)
+            return {
+                "success": False,
+                "error": f"Failed to execute code modifications: {str(e)}",
+                "stdout": "",
+                "stderr": str(e)
+            }

jarvis/tools/{file_ops.py → file_operation.py}

@@ -1,6 +1,5 @@
-from typing import Dict, Any, Protocol
+from typing import Dict, Any
 import os
-from enum import Enum
 
 from jarvis.utils import OutputType, PrettyOutput
 
@@ -36,15 +35,15 @@ class FileOperationTool:
 
 
     def execute(self, args: Dict) -> Dict[str, Any]:
-        """执行文件操作"""
+        """Execute file operations"""
         try:
             operation = args["operation"]
             filepath = args["filepath"]
             encoding = args.get("encoding", "utf-8")
             
-            # 记录操作和完整路径
+            # Record the operation and the full path
             abs_path = os.path.abspath(filepath)
-            PrettyOutput.print(f"文件操作: {operation} - {abs_path}", OutputType.INFO)
+            PrettyOutput.print(f"File operation: {operation} - {abs_path}", OutputType.INFO)
             
             if operation == "exists":
                 exists = os.path.exists(filepath)
@@ -61,15 +60,15 @@ class FileOperationTool:
                         "error": f"文件不存在: {filepath}"
                     }
                     
-                # 检查文件大小
+                # Check file size
                 if os.path.getsize(filepath) > 10 * 1024 * 1024:  # 10MB
                     return {
                         "success": False,
-                        "error": "文件过大 (>10MB)"
+                        "error": "File too large (>10MB)"
                     }
                     
-                with open(filepath, 'r', encoding=encoding) as f:
-                    content = f.read()
+                content = open(filepath, 'r', encoding=encoding).read()
+                PrettyOutput.print(content, OutputType.INFO)
                 return {
                     "success": True,
                     "stdout": content,
@@ -80,10 +79,10 @@ class FileOperationTool:
                 if not args.get("content"):
                     return {
                         "success": False,
-                        "error": "写入/追加操作需要提供content参数"
+                        "error": "Write/append operation requires providing the content parameter"
                     }
                 
-                # 创建目录（如果不存在）
+                # Create directory (if it doesn't exist)
                 os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
                     
                 mode = 'a' if operation == "append" else 'w'
@@ -92,19 +91,19 @@ class FileOperationTool:
                     
                 return {
                     "success": True,
-                    "stdout": f"成功{operation}内容到 {filepath}",
+                    "stdout": f"Successfully {operation} content to {filepath}",
                     "stderr": ""
                 }
                 
             else:
                 return {
                     "success": False,
-                    "error": f"未知操作: {operation}"
+                    "error": f"Unknown operation: {operation}"
                 }
                 
         except Exception as e:
             PrettyOutput.print(str(e), OutputType.ERROR)
             return {
                 "success": False,
-                "error": f"文件操作失败: {str(e)}"
+                "error": f"File operation failed: {str(e)}"
             } 

jarvis/tools/find_related_files.py

@@ -0,0 +1,86 @@
+import os
+from typing import Any, Dict
+from jarvis.jarvis_codebase.main import CodeBase
+from jarvis.utils import find_git_root, PrettyOutput, OutputType
+
+class CodebaseSearchTool:
+    """Codebase Search Tool"""
+    
+    name = "find_related_files"
+    description = "Find code files related to the given requirement by searching the codebase using NLP-based semantic search. IMPORTANT: Requires complete sentence descriptions - keywords or phrases alone will not yield comprehensive results."
+    parameters = {
+        "type": "object",
+        "properties": {
+            "dir": {
+                "type": "string",
+                "description": "Project root directory"
+            },
+            "query": {
+                "type": "string",
+                "description": """Requirement description to find related code files. Must be a complete sentence that clearly describes what you're looking for.
+
+Good examples:
+- 'Need to modify the user authentication process to add password validation'
+- 'Want to update the database connection configuration to support multiple databases'
+- 'Looking for the code that handles file upload functionality in the system'
+
+Bad examples (will not work well):
+- 'user auth'
+- 'database config'
+- 'upload'""",
+                "examples": [
+                    "Need to modify the error handling in the authentication process",
+                    "Want to update how the system processes user uploads",
+                    "Looking for the code that manages database connections"
+                ]
+            },
+            "top_k": {
+                "type": "integer",
+                "description": "Number of most relevant files to return",
+                "default": 5
+            }
+        },
+        "required": ["query"]
+    }
+    
+    def execute(self, params: Dict[str, Any]) -> Dict[str, Any]:
+        """Execute codebase search"""
+        try:
+            dir = params.get("dir")
+            query = params["query"]
+            top_k = params.get("top_k", 5)
+            
+            # Find the root directory of the codebase
+            current_dir = os.getcwd()
+            root_dir = find_git_root(dir or current_dir)
+            if not root_dir:
+                return {
+                    "success": False,
+                    "stdout": "",
+                    "stderr": "Error: Current directory is not in a Git repository",
+                    "error": "NotInGitRepository"
+                }
+
+            os.chdir(root_dir)
+            codebase = CodeBase(root_dir)
+            # Generate index
+
+            codebase.generate_codebase()
+            # Execute search
+            response = codebase.search_similar(query, top_k)
+            os.chdir(current_dir)
+            return {
+                "success": True,
+                "stdout": str(response),
+                "stderr": "",
+                "error": None
+            }
+            
+        except Exception as e:
+            PrettyOutput.print(f"Codebase QA error: {str(e)}", output_type=OutputType.ERROR)
+            return {
+                "success": False,
+                "stdout": "",
+                "stderr": f"Error executing codebase QA: {str(e)}",
+                "error": str(type(e).__name__)
+            }