jarvis-ai-assistant 0.1.101__py3-none-any.whl → 0.1.103__py3-none-any.whl

jarvis/__init__.py

@@ -1,3 +1,3 @@
 """Jarvis AI Assistant"""
 
-__version__ = "0.1.101" 
+__version__ = "0.1.103" 

jarvis/agent.py

@@ -1,22 +1,38 @@
 import argparse
 import time
-from typing import Dict, List, Optional
+from typing import Callable, Dict, List, Optional
 
 from prompt_toolkit import prompt
 import yaml
 
+from jarvis.models.base import BasePlatform
 from jarvis.models.registry import PlatformRegistry
-from jarvis.tools import ToolRegistry
-from jarvis.tools.registry import load_tools
-from jarvis.utils import PrettyOutput, OutputType, get_single_line_input, load_methodology, add_agent, delete_current_agent, get_max_context_length, get_multiline_input, load_embedding_model, init_env
+from jarvis.tools.registry import ToolRegistry, tool_call_help
+from jarvis.utils import PrettyOutput, OutputType, is_auto_complete, load_methodology, add_agent, delete_current_agent, get_max_context_length, get_multiline_input, init_env
 import os
 
 class Agent:
 
     def __del__(self):
         delete_current_agent()
+
+    def set_summary_prompt(self, summary_prompt: str):
+        self.summary_prompt = summary_prompt
+
+    def set_output_filter(self, output_filter: List[Callable]):
+        self.output_filter = output_filter
         
-    def __init__(self, system_prompt: str, name: str = "Jarvis", is_sub_agent: bool = False, tool_registry: Optional[ToolRegistry] = None):
+    def __init__(self, 
+                 system_prompt: str, 
+                 name: str = "Jarvis", 
+                 is_sub_agent: bool = False, 
+                 tool_registry: Optional[ToolRegistry] = None, 
+                 platform: Optional[BasePlatform] = None, 
+                 summary_prompt: Optional[str] = None, 
+                 auto_complete: bool = False, 
+                 record_methodology: bool = True,
+                 output_filter: Optional[List[Callable]] = None,
+                 need_summary: bool = True):
         """Initialize Agent with a model, optional tool registry and name
         
         Args:
@@ -24,45 +40,70 @@ class Agent:
             name: Agent name, default is "Jarvis"
             is_sub_agent: Whether it is a sub-agent, default is False
             tool_registry: Tool registry instance
+            platform: Optional platform instance, default uses normal platform
         """
         add_agent(name)
         PrettyOutput.print(f"Welcome to Jarvis, your AI assistant, Initiating...", OutputType.SYSTEM)
-        self.model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
-        self.tool_registry = tool_registry if tool_registry else ToolRegistry.get_global_tool_registry()
+        if platform is not None:
+            self.model = platform
+        else:
+            self.model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
+        self.tool_registry = tool_registry if tool_registry else ToolRegistry()
+        self.record_methodology = record_methodology
         self.name = name
         self.is_sub_agent = is_sub_agent
         self.prompt = ""
         self.conversation_length = 0  # Use length counter instead
         self.system_prompt = system_prompt
+        self.need_summary = need_summary
         # Load configuration from environment variables
-        self.embedding_dimension = 1536  # Default for many embedding models
+        self.output_filter = output_filter if output_filter else []
+
+        self.summary_prompt = summary_prompt if summary_prompt else f"""Please generate a concise summary report of the task execution, including:
+
+1. Task Objective: Task restatement
+2. Execution Result: Success/Failure
+3. Key Information: Important information extracted during execution
+4. Important Findings: Any noteworthy discoveries
+5. Follow-up Suggestions: If any
+
+Please describe in concise bullet points, highlighting important information.
+"""
+        
         self.max_context_length = get_max_context_length()
+
+        self.auto_complete = auto_complete
+
+
         
-        # Initialize embedding model
-        try:
-            self.embedding_model = load_embedding_model()
-            
-            # Warm up model and get correct dimension
-            test_text = "This is a test text to ensure the model is fully loaded."
-            test_embedding = self.embedding_model.encode(
-                test_text, 
-                convert_to_tensor=True,
-                normalize_embeddings=True
-            )
-            self.embedding_dimension = len(test_embedding)
-            PrettyOutput.print("Successfully loaded embedding model", OutputType.SUCCESS)
-            
-            # Initialize HNSW index (use correct dimension)
-            
-            
-        except Exception as e:
-            PrettyOutput.print(f"Failed to load embedding model: {str(e)}", OutputType.ERROR)
-            raise
             
         # Initialize methodology related attributes
         self.methodology_data = []
 
         PrettyOutput.section(f"Jarvis initialized - With {self.model.name()}", OutputType.SYSTEM)
+        tools = self.tool_registry.get_all_tools()
+        if tools:
+            PrettyOutput.section(f"Available tools: {', '.join([tool['name'] for tool in tools])}", OutputType.SYSTEM)
+
+
+        # Load methodology
+        
+        tools_prompt = self.tool_registry.load_tools()
+        complete_prompt = """"""
+        if self.auto_complete:
+            complete_prompt = """
+            When the task is completed, you should print the following message:
+            <!!!COMPLETE!!!>
+            """
+
+        self.model.set_system_message(f"""
+{self.system_prompt}
+
+{tools_prompt}
+
+{complete_prompt}
+""")
+        self.first = True
 
     @staticmethod
     def extract_tool_calls(content: str) -> List[Dict]:
@@ -72,14 +113,6 @@ class Agent:
         tool_call_lines = []
         in_tool_call = False
 
-        tool_call_help = """Tool Usage Format:
-
-<TOOL_CALL>
-name: tool_name
-arguments:
-    param1: value1
-    param2: value2
-</TOOL_CALL>"""
 
         # Process line by line
         for line in lines:
@@ -88,27 +121,19 @@ arguments:
                 continue
             elif '</TOOL_CALL>' in line:
                 if in_tool_call and tool_call_lines:
-                    try:
-                        # Parse YAML directly
-                        tool_call_text = '\n'.join(tool_call_lines)
-                        tool_call_data = yaml.safe_load(tool_call_text)
-                        
-                        # Validate necessary fields
-                        if "name" in tool_call_data and "arguments" in tool_call_data:
-                            # Return content before tool call and tool call
-                            return [{
-                                "name": tool_call_data["name"],
-                                "arguments": tool_call_data["arguments"]
-                            }]
-                        else:
-                            PrettyOutput.print("Tool call missing necessary fields", OutputType.ERROR)
-                            raise Exception("Tool call missing necessary fields, " + tool_call_help)
-                    except yaml.YAMLError as e:
-                        PrettyOutput.print(f"YAML parsing error: {str(e)}", OutputType.ERROR)
-                        raise Exception(f"YAML parsing error: {str(e)}")
-                    except Exception as e:
-                        PrettyOutput.print(f"Error processing tool call: {str(e)}", OutputType.ERROR)
-                        raise Exception(f"Error processing tool call: {str(e)}")
+                    # Parse YAML directly
+                    tool_call_text = '\n'.join(tool_call_lines)
+                    tool_call_data = yaml.safe_load(tool_call_text)
+                    
+                    # Validate necessary fields
+                    if "name" in tool_call_data and "arguments" in tool_call_data:
+                        # Return content before tool call and tool call
+                        return [{
+                            "name": tool_call_data["name"],
+                            "arguments": tool_call_data["arguments"]
+                        }]
+                    else:
+                        raise Exception("Tool call missing necessary fields")
                 in_tool_call = False
                 continue
             
@@ -185,53 +210,39 @@ Please continue the task based on the above information.
         """
         PrettyOutput.section("Task completed", OutputType.SUCCESS)
         
-        # 询问是否生成方法论，带输入验证
-        while True:
-            user_input = get_single_line_input("Generate methodology for this task? (y/n)").strip().lower()
-            if user_input in ['y', 'n', '']:
-                break
-            PrettyOutput.print("Invalid input, please enter y or n", OutputType.WARNING)
-        
-        if user_input == 'y':
-            try:
-                # 让模型判断是否需要生成方法论
-                analysis_prompt = """The current task has ended, please analyze whether a methodology needs to be generated.
-If you think a methodology should be generated, first determine whether to create a new methodology or update an existing one. If updating an existing methodology, use 'update', otherwise use 'add'.
-If you think a methodology is not needed, please explain why.
-The methodology should be applicable to general scenarios, do not include task-specific information such as code commit messages.
-The methodology should include: problem restatement, optimal solution, notes (as needed), and nothing else.
-Only output the methodology tool call instruction, or the explanation for not generating a methodology. Do not output anything else.
-"""
-                self.prompt = analysis_prompt
-                response = self._call_model(self.prompt)
-                
-                # 检查是否包含工具调用
+        if not self.is_sub_agent:
+            if self.record_methodology:
+
                 try:
-                    tool_calls = Agent.extract_tool_calls(response)
-                    if tool_calls:
-                        self.tool_registry.handle_tool_calls(tool_calls)
+                    # 让模型判断是否需要生成方法论
+                    analysis_prompt = """The current task has ended, please analyze whether a methodology needs to be generated.
+    If you think a methodology should be generated, first determine whether to create a new methodology or update an existing one. If updating an existing methodology, use 'update', otherwise use 'add'.
+    If you think a methodology is not needed, please explain why.
+    The methodology should be applicable to general scenarios, do not include task-specific information such as code commit messages.
+    The methodology should include: problem restatement, optimal solution, notes (as needed), and nothing else.
+    Only output the methodology tool call instruction, or the explanation for not generating a methodology. Do not output anything else.
+    """
+                    self.prompt = analysis_prompt
+                    response = self._call_model(self.prompt)
+                    
+                    # 检查是否包含工具调用
+                    try:
+                        tool_calls = Agent.extract_tool_calls(response)
+                        if tool_calls:
+                            self.tool_registry.handle_tool_calls(tool_calls)
+                    except Exception as e:
+                        PrettyOutput.print(f"Failed to handle methodology generation: {str(e)}", OutputType.ERROR)
+                    
                 except Exception as e:
-                    PrettyOutput.print(f"Failed to handle methodology generation: {str(e)}", OutputType.ERROR)
-                
-            except Exception as e:
-                PrettyOutput.print(f"Error generating methodology: {str(e)}", OutputType.ERROR)
-        
-        if not self.is_sub_agent:
+                    PrettyOutput.print(f"Error generating methodology: {str(e)}", OutputType.ERROR)
+            
             return "Task completed"
         
-        # 生成任务总结
-        summary_prompt = f"""Please generate a concise summary report of the task execution, including:
-
-1. Task Objective: Task restatement
-2. Execution Result: Success/Failure
-3. Key Information: Important information extracted during execution
-4. Important Findings: Any noteworthy discoveries
-5. Follow-up Suggestions: If any
-
-Please describe in concise bullet points, highlighting important information.
-"""
-        self.prompt = summary_prompt
-        return self._call_model(self.prompt)
+        if self.need_summary:
+            self.prompt = self.summary_prompt
+            return self._call_model(self.prompt)
+        
+        return "Task completed"
 
 
     def run(self, user_input: str, file_list: Optional[List[str]] = None) -> str:
@@ -244,28 +255,22 @@ Please describe in concise bullet points, highlighting important information.
         Returns:
             str: Task summary report
         """
+
+        
+
         try:
             PrettyOutput.section("Preparing environment", OutputType.PLANNING)
             if file_list:
                 self.model.upload_files(file_list)
 
-            # Load methodology
-            methodology_prompt = load_methodology(user_input)
-            tools_prompt = load_tools()
-
             # 显示任务开始
             PrettyOutput.section(f"Starting new task: {self.name}", OutputType.PLANNING)
 
-            self.clear_history()  
-
-            self.model.set_system_message(f"""
-{self.system_prompt}
-
-{tools_prompt}
-
-{methodology_prompt}
-""")
-            self.prompt = f"{user_input}"
+            if self.first:
+                self.prompt = f"{user_input}\n\n{load_methodology(user_input)}"
+                self.first = False
+            else:
+                self.prompt = f"{user_input}"
 
             while True:
                 try:
@@ -281,25 +286,32 @@ Please describe in concise bullet points, highlighting important information.
                         continue
                     else:
                         current_response = self._call_model(self.prompt)
-                        self.conversation_length += len(current_response)  # Add response length
+                        self.prompt = ""
+                        self.conversation_length += len(current_response) 
+
+                    for filter in self.output_filter:
+                        self.prompt += filter(current_response)
+
                     try:
                         result = Agent.extract_tool_calls(current_response)
                     except Exception as e:
                         PrettyOutput.print(f"Tool call error: {str(e)}", OutputType.ERROR)
-                        self.prompt = f"Tool call error: {str(e)}"
+                        self.prompt += f"Tool call error: {str(e)}"
                         continue
                     
                     if len(result) > 0:
                         PrettyOutput.print("Executing tool call...", OutputType.PROGRESS)
                         tool_result = self.tool_registry.handle_tool_calls(result)
-                        self.prompt = tool_result
+                        self.prompt += tool_result
+                    
+                    if self.prompt:
                         continue
+
+                    if self.auto_complete and "<!!!COMPLETE!!!>" in current_response:
+                        return self._complete_task()
                     
                     # 获取用户输入
                     user_input = get_multiline_input(f"{self.name}: You can continue to input, or enter an empty line to end the current task")
-                    if user_input == "__interrupt__":
-                        PrettyOutput.print("Task cancelled by user", OutputType.WARNING)
-                        return "Task cancelled by user"
 
                     if user_input:
                         self.prompt = user_input
@@ -382,10 +394,11 @@ def select_task(tasks: dict) -> str:
     # Convert tasks to list for ordered display
     task_names = list(tasks.keys())
     
-    PrettyOutput.print("\nAvailable tasks:", OutputType.INFO)
+    task_list = ["Available tasks:"]
     for i, name in enumerate(task_names, 1):
-        PrettyOutput.print(f"[{i}] {name}", OutputType.INFO)
-    PrettyOutput.print("[0] Skip predefined tasks", OutputType.INFO)
+        task_list.append(f"[{i}] {name}")
+    task_list.append("[0] Skip predefined tasks")
+    PrettyOutput.print("\n".join(task_list), OutputType.INFO)
     
     
     while True:
@@ -426,26 +439,13 @@ When users need to execute tasks, you will strictly follow these steps to handle
 7. Execute Action Plan: Execute one step at a time, **use at most one tool** (wait for tool execution results before proceeding)
 8. Monitor and Adjust: If execution results don't match expectations, reflect and adjust the action plan, iterate previous steps
 9. Methodology: If the current task has general applicability and valuable experience is gained, use methodology tools to record it for future similar problems
-10. Task Completion: End the task using task completion command when finished
+10. Auto check the task goal completion status: If the task goal is completed, use the task completion command to end the task
+11. Task Completion: End the task using task completion command when finished
 
 Methodology Template:
 1. Problem Restatement
 2. Optimal Solution
 3. Optimal Solution Steps (exclude failed actions)
-                                          
-Strict Rules:
-- Execute only one tool at a time
-- Tool execution must strictly follow the tool usage format
-- Wait for user to provide execution results
-- Don't assume or imagine results
-- Don't create fake dialogues
-- If current information is insufficient, you may ask the user
-- Not all problem-solving steps are mandatory, skip as appropriate
-- Ask user before executing tools that might damage system or user's codebase
-- Request user guidance when multiple iterations show no progress
-- If yaml string contains colons, wrap the entire string in quotes to avoid yaml parsing errors
-- Use | syntax for multi-line strings in yaml
-- If you can start executing the task, please start directly without asking the user if you can begin.
 
 -------------------------------------------------------------"""
 
@@ -459,7 +459,7 @@ def main():
 
     try:
         # 获取全局模型实例
-        agent = Agent(system_prompt=origin_agent_system_prompt)
+        agent = Agent(system_prompt=origin_agent_system_prompt, tool_registry=ToolRegistry())
 
         # 加载预定义任务
         tasks = load_tasks()
@@ -474,7 +474,7 @@ def main():
         while True:
             try:
                 user_input = get_multiline_input("Please enter your task (input empty line to exit):")
-                if not user_input or user_input == "__interrupt__":
+                if not user_input:
                     break
                 agent.run(user_input, args.files)
             except Exception as e:

jarvis/jarvis_code_agent/code_agent.py

@@ -0,0 +1,234 @@
+from enum import auto
+import os
+import re
+from typing import List
+
+import yaml
+from jarvis.agent import Agent
+from jarvis.jarvis_code_agent.patch import apply_patch
+from jarvis.jarvis_code_agent.file_select import select_files
+from jarvis.jarvis_code_agent.relevant_files import find_relevant_files
+from jarvis.models.registry import PlatformRegistry
+from jarvis.tools.git_commiter import GitCommitTool
+from jarvis.tools.registry import ToolRegistry
+from jarvis.utils import OutputType, PrettyOutput, get_file_line_count, get_multiline_input, get_single_line_input, has_uncommitted_changes, init_env, find_git_root, is_disable_codebase, make_choice_input, user_confirm
+
+
+
+
+
+class CodeAgent:
+    def __init__(self):
+        self.root_dir = os.getcwd()
+        tool_registry = ToolRegistry()
+        tool_registry.use_tools(["read_code", "execute_shell", "search", "ask_user", "ask_codebase"])
+        code_system_prompt = """
+You are a code agent, you are responsible for modifying the code.
+
+You should read the code and analyze the code, and then provide a plan for the code modification.
+
+## Workflow Steps
+
+1. ANALYSIS
+- Understand the requirement thoroughly
+- Identify which files need to be modified
+- Review the current implementation
+- Consider potential impacts
+
+2. PLANNING
+- Break down the changes into logical steps
+- Consider dependencies between changes
+- Plan the implementation sequence
+- Think about potential risks
+
+3. IMPLEMENTATION
+For each file that needs changes:
+a. Read and understand the current code
+b. Plan the specific modifications
+c. Write the patch in the required format
+d. Review the patch for correctness
+
+## File Reading Guidelines
+
+1. For Large Files (>200 lines):
+- Do NOT read the entire file at once using 'read_code'
+- First use 'execute_shell' with grep/find to locate relevant sections
+- Then use 'read_code' with specific line ranges to read only necessary portions
+- Example: 
+  * Use: execute_shell("grep -n 'function_name' path/to/file")
+  * Then: read_code("path/to/file", start_line=found_line-10, end_line=found_line+20)
+
+2. For Small Files:
+- Can read entire file using 'read_code' directly
+
+## Patch Format and Guidelines
+
+1. Basic Format:
+<PATCH>
+> /path/to/file start_line,end_line
+new_content_line1
+new_content_line2
+</PATCH>
+
+2. Rules:
+- Each <PATCH> block MUST contain exactly ONE patch for ONE location
+- Multiple changes to different locations require separate <PATCH> blocks
+- Line Numbers Behavior:
+  * start_line (first number): This line WILL be replaced
+  * end_line (second number): This line will NOT be replaced
+  * The patch replaces content from start_line (inclusive) to end_line (exclusive)
+- Use absolute paths relative to the project root
+- Maintain consistent indentation
+- Include enough context for precise location
+
+3. Multiple Changes Example:
+Before:
+```
+Line 0: first line
+Line 1: second line
+Line 2: third line
+Line 3: fourth line
+```
+
+For multiple changes, use separate patches:
+```
+<PATCH>
+> /path/to/file 0,1
+new first line
+</PATCH>
+
+<PATCH>
+> /path/to/file 2,3
+new third line
+</PATCH>
+```
+
+After:
+```
+new first line
+Line 1: second line
+new third line
+Line 3: fourth line
+```
+
+Note: In this example:
+- Each change is in its own <PATCH> block
+- Changes are applied sequentially
+- Line numbers are based on the original file
+
+## Implementation Guidelines
+
+1. Code Quality:
+- Keep changes minimal and focused
+- Maintain consistent style
+- Add clear comments for complex logic
+- Follow project patterns
+- Ensure proper error handling
+
+2. Tools Available:
+- Use 'read_code/ask_codebase' to examine file contents
+- Use 'execute_shell' for grep/find/ctags operations
+- Use 'search' to search on web
+- Use 'ask_user' when clarification is needed
+
+Please proceed with the analysis and implementation following this workflow.
+Start by examining the files and planning your changes.
+Then provide the necessary patches in the specified format.
+"""
+        self.agent = Agent(system_prompt=code_system_prompt, 
+                           name="CodeAgent", 
+                           auto_complete=False,
+                           is_sub_agent=False, 
+                           tool_registry=tool_registry, 
+                           platform=PlatformRegistry().get_codegen_platform(), 
+                           record_methodology=False,
+                           output_filter=[apply_patch],
+                           need_summary=False)
+
+    
+
+    def _init_env(self):
+        curr_dir = os.getcwd()
+        git_dir = find_git_root(curr_dir)
+        self.root_dir = git_dir
+        if has_uncommitted_changes():
+            git_commiter = GitCommitTool()
+            git_commiter.execute({})
+
+    
+    def make_files_prompt(self, files: List[str]) -> str:
+        """Make the files prompt.
+        
+        Args:
+            files: The files to be modified
+            
+        """
+        return "\n".join(
+            f"- {file} ({get_file_line_count(file)} lines)"
+            for file in files
+        )
+
+    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
+        """
+        try:
+            self._init_env()
+            files = find_relevant_files(user_input, self.root_dir)
+            self.agent.run(self._build_first_edit_prompt(user_input, self.make_files_prompt(files)))
+            
+        except Exception as e:
+            return f"Error during execution: {str(e)}"
+        
+
+
+    def _build_first_edit_prompt(self, user_input: str, files_prompt: str) -> str:
+        """Build the initial prompt for the agent.
+        
+        Args:
+            user_input: The user's requirement
+            files_prompt: The formatted list of relevant files
+            
+        Returns:
+            str: The formatted prompt
+        """
+        return f"""# Code Modification Task
+
+## User Requirement
+{user_input}
+
+## Available Files
+{files_prompt}
+"""
+def main():
+    """Jarvis main entry point"""
+    # Add argument parser
+    init_env()
+
+
+    try:
+        # Interactive mode
+        while True:
+            try:
+                user_input = get_multiline_input("Please enter your requirement (input empty line to exit):")
+                if not user_input:
+                    break
+                agent = CodeAgent()
+                agent.run(user_input)
+                
+            except Exception as e:
+                PrettyOutput.print(f"Error: {str(e)}", OutputType.ERROR)
+
+    except Exception as e:
+        PrettyOutput.print(f"Initialization error: {str(e)}", OutputType.ERROR)
+        return 1
+
+    return 0
+
+if __name__ == "__main__":
+    exit(main())