PikoAi 0.1.6__py3-none-any.whl → 0.1.9__py3-none-any.whl

Agents/Executor/executor.py

@@ -6,15 +6,17 @@ import sys
 import time
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../')))
 from Utils.ter_interface import TerminalInterface
-from Utils.executor_utils import parse_tool_call, parse_code, parse_shell_command
+from Utils.executor_utils import parse_tool_call
 from Agents.Executor.prompts import get_system_prompt, get_task_prompt # Import prompts
 
 from typing import Optional
 from mistralai.models.sdkerror import SDKError # This might be an issue if LiteLLM doesn't use SDKError
                                               # LiteLLM maps exceptions to OpenAI exceptions.
                                               # We'll keep it for now and see if errors arise during testing.
-from Env import python_executor
-from Env.shell import ShellExecutor # Import ShellExecutor
+# from Env import python_executor # Will be replaced by BaseEnv
+# from Env.shell import ShellExecutor # Will be replaced by BaseEnv
+from Src.Env.base_env import create_environment, BaseEnv # Added
+from Env import python_executor # Keep for type hint in the old execute method if needed, or remove if execute is fully removed
 from llm_interface.llm import LiteLLMInterface # Import LiteLLMInterface
 
 from Tools import tool_manager
@@ -38,8 +40,8 @@ class executor:
         self.max_iter = max_iter
         self.rate_limiter = RateLimiter(wait_time=5.0, max_retries=3)
         self.executor_prompt_init()  # Update system_prompt
-        self.python_executor = python_executor.PythonExecutor()  # Initialize PythonExecutor
-        self.shell_executor = ShellExecutor() # Initialize ShellExecutor
+        # self.python_executor = python_executor.PythonExecutor()  # Initialize PythonExecutor
+        # self.shell_executor = ShellExecutor() # Initialize ShellExecutor
         self.message = [
             {"role": "system", "content": self.system_prompt},
             {"role": "user", "content": self.task_prompt}
@@ -122,89 +124,29 @@ class executor:
             # Check for tool calls
             response = self.run_inference()
             tool_call = parse_tool_call(response)
+
             if tool_call:
-                print(f"\nCalling tool: {tool_call['tool_name']}")
+                tool_name = tool_call['tool_name']
+                tool_input = tool_call['input']
+                print(f"\nIdentified tool call: {tool_name} with input {tool_input}")
+
+                # Call the tool and append the result (no confirmation or special logic)
                 try:
-                    # Pass tool name and input as separate arguments
-                    tool_output = tool_manager.call_tool(tool_call["tool_name"], tool_call["input"])
-                    self.terminal.tool_output_log(tool_output, tool_call["tool_name"])
-                    self.message.append({"role": "user", "content": f"Tool Output: {tool_output}"})
+                    tool_output_result = tool_manager.call_tool(tool_name, tool_input)
+                    self.terminal.tool_output_log(tool_output_result, tool_name)
+                    print(tool_output_result)
+                    self.message.append({"role": "user", "content": tool_output_result})
                 except ValueError as e:
                     error_msg = str(e)
+                    print(f"Tool Error: {error_msg}")
                     self.message.append({"role": "user", "content": f"Tool Error: {error_msg}"})
-            
-            else: # Not a tool call, check for code or shell command
-                code = parse_code(response)
-                shell_command = parse_shell_command(response)
-
-                if code:
-                    # Ask user for confirmation before executing the code
-                    user_confirmation = input("Do you want to execute the Python code?")
-                    if user_confirmation.lower() == 'y':
-                        exec_result = self.python_executor.execute(code)
-                        if exec_result['output'] == "" and not exec_result['success']:
-                            error_msg = (
-                                f"Python execution failed.\n"
-                                f"Error: {exec_result.get('error', 'Unknown error')}"
-                            )
-                            print(f"there was an error in the python code execution {exec_result.get('error', 'Unknown error')}")
-                            self.message.append({"role": "user", "content": error_msg})
-
-                        elif exec_result['output'] == "":
-                            no_output_msg = (
-                                "Python execution completed but no output was shown. "
-                                "Please add print statements to show the results. This isn't a jupyter notebook environment. "
-                                "For example: print(your_variable) or print('Your message')"
-                            )
-                            self.message.append({"role": "user", "content": no_output_msg})
-                        
-                        #if there is an output (partial or full exeuction)
-                        else:
-                            # First, show the program output
-                            if exec_result['output'].strip():
-                                print(f"Program Output:\n{exec_result['output']}")
-                            
-                            # Then handle success/failure cases
-                            if exec_result['success']:
-                                self.message.append({"role": "user", "content": f"Program Output:\n{exec_result['output']}"})
-                            else:
-                                self.message.append({"role": "user", "content": f"Program Output:\n{exec_result['output']}\n{exec_result.get('error', 'Unknown error')}"})
-    
-                    else:
-                        self.message.append({"role":"user","content":"User chose not to execute the Python code."})
-                        print("Python code execution skipped by the user.")
-                
-                elif shell_command:
-                    user_confirmation = input(f"Do you want to execute the shell command: '{shell_command}'?\n ")
-                    if user_confirmation.lower() == 'y':
-                        shell_result = self.shell_executor.execute(shell_command)
-                        if shell_result['output'] == "" and not shell_result['success']:
-                            error_msg = (
-                                f"Shell command execution failed.\n"
-                                f"Error: {shell_result.get('error', 'Unknown error')}"
-                            )
-                            print(f"there was an error in the shell command execution {shell_result.get('error', 'Unknown error')}")
-                            self.message.append({"role": "user", "content": error_msg})
-
-                        elif shell_result['output'] == "":
-                            print("command executed")
-                            self.message.append({"role": "user", "content": "command executed"})
-                        
-                        #if there is an output (partial or full execution)
-                        else:
-                            # First, show the command output
-                            if shell_result['output'].strip():
-                                print(f"Command Output:\n{shell_result['output']}")
-                            
-                            # Then handle success/failure cases
-                            if shell_result['success']:
-                                self.message.append({"role": "user", "content": f"Command Output:\n{shell_result['output']}"})
-                            else:
-                                self.message.append({"role": "user", "content": f"Command Output:\n{shell_result['output']}\n{shell_result.get('error', 'Unknown error')}"})
-    
-                    else:
-                        self.message.append({"role":"user","content":"User chose not to execute the shell command."})
-                        print("Shell command execution skipped by the user.")
+
+            else: # Not a tool call, could be a direct response or requires clarification
+                # This part handles responses that are not formatted as tool calls.
+                # It might be a final answer, a question, or just conversational text.
+                # The existing logic for TASK_DONE or asking for next step handles this.
+                # No specific code/shell parsing here anymore as they are tools.
+                pass # Explicitly pass if no tool call and no old code/shell logic.
 
             # Check if task is done
             if "TASK_DONE" in response:
@@ -218,24 +160,26 @@ class executor:
         if not task_done:
             print(f"Task could not be completed within {self.max_iter} iterations.")
 
-    def execute(self, code: str, exec_env: python_executor.PythonExecutor):
-        """Executes the given Python code using the provided execution environment."""
-        result = exec_env.execute(code)
-        return result
+    # This method is superseded by the BaseEnv approach in run_task
+    # def execute(self, code: str, exec_env: python_executor.PythonExecutor):
+    #     """Executes the given Python code using the provided execution environment."""
+    #     result = exec_env.execute(code)
+    #     return result
 
 if __name__ == "__main__":
-    e1 = executor("")
-    user_prompt = input("Please enter your prompt: ")
-    e1.user_prompt = user_prompt
-    e1.executor_prompt_init()  # Update system_prompt
-    e1.message = [
-        {"role": "system", "content": e1.system_prompt},
-        {"role": "user", "content": e1.task_prompt}
-    ]  # Reset message list properly
-    e1.run()
-
-    while True:
-        user_prompt = input("Please enter your prompt: ")
-        e1.message.append({"role": "user", "content": user_prompt})
-        # e1.message.append({"role":"user","content":e1.system_prompt})
-        e1.run()
+    # e1 = executor("") # Commenting out example usage for now as it might need adjustment
+    # user_prompt = input("Please enter your prompt: ")
+    # e1.user_prompt = user_prompt
+    # e1.executor_prompt_init()  # Update system_prompt
+    # e1.message = [
+    #     {"role": "system", "content": e1.system_prompt},
+    #     {"role": "user", "content": e1.task_prompt}
+    # ]  # Reset message list properly
+    # e1.run()
+
+    # while True:
+    #     user_prompt = input("Please enter your prompt: ")
+    #     e1.message.append({"role": "user", "content": user_prompt})
+    #     # e1.message.append({"role":"user","content":e1.system_prompt})
+    #     e1.run()
+    pass # Placeholder if main execution is commented out

Agents/Executor/prompts.py

@@ -17,9 +17,7 @@ You have access to the following tools:
 
 Your primary objective is to accomplish the user's goal by performing step-by-step actions. These actions can include:
 1. Calling a tool
-2. Executing Python code
-3. Executing Shell commands
-4. Providing a direct response
+2. Providing a direct response
 
 You must break down the user's goal into smaller steps and perform one action at a time. After each action, carefully evaluate the output to determine the next step.
 
@@ -33,14 +31,9 @@ You must break down the user's goal into smaller steps and perform one action at
     }}
   }}
   <<END_TOOL_CALL>>
-- **Code Execution**: Write Python code when no tool is suitable or when custom logic is needed. Format:
-  <<CODE>>
-  your_python_code_here
-  <<CODE>>
-- **Shell Command Execution**: Execute shell commands when needed. Format:
-  <<SHELL_COMMAND>>
-  your_shell_command_here
-  <<END_SHELL_COMMAND>>
+  This includes executing Python code and shell commands:
+  `execute_python_code`: {{"code": "your_python_code_here"}}
+  `execute_shell_command`: {{"command": "your_shell_command_here"}}
 - **Direct Response**: Provide a direct answer if the task doesn't require tools or code.
 
 ### Important Notes:
@@ -68,16 +61,7 @@ Following are the things that you must read carefully and remember:
             }
         }
         <<END_TOOL_CALL>>
-
-        - For code execution, use:
-        <<CODE>>
-        your_python_code_here
-        <<CODE>>
-
-        - For shell command execution, use:
-        <<SHELL_COMMAND>>
-        your_shell_command_here
-        <<END_SHELL_COMMAND>>
+        Remember that executing Python code and shell commands is now done through specific tool calls (`execute_python_code` and `execute_shell_command`).
 
         After each action, always evaluate the output to decide your next step. Only include 'TASK_DONE'
         When the entire task is completed. Do not end the task immediately after a tool call or code execution without

Env/base_env.py

@@ -1,25 +1,30 @@
-from Env.js_executor import JavaScriptExecutor #class import
-from Env.python_executor import PythonExecutor #class import
-
 # to perform funciton ask whether to execute code
     
 
 class BaseEnv:
 
 
-    def __init__(self, language,code):
-        self.language = language
+    def __init__(self):
+        pass
+
 
+    def execute(self, code_or_command: str):
+        raise NotImplementedError("This method should be overridden by subclasses")
 
-    def execute(self):
+    def stop_execution(self):
         raise NotImplementedError("This method should be overridden by subclasses")
 
       
 def create_environment(language):
+    # Moved imports inside the function to avoid circular dependencies during testing
+    from Env.js_executor import JavaScriptExecutor #class import
+    from Env.python_executor import PythonExecutor #class import
+    from Env.shell import ShellExecutor #class import
+
     if language == "python":
         return PythonExecutor()
-    elif language == "javascript":
-        return JavaScriptExecutor()
+    elif language == "shell":
+        return ShellExecutor()
     else:
         raise ValueError(f"Unsupported language: {language}")
     

Env/python_executor.py

@@ -21,9 +21,13 @@ import os
 from typing import Dict
 import textwrap
 import sys
+from Src.Env.base_env import BaseEnv
+import time
 
-class PythonExecutor:
+class PythonExecutor(BaseEnv):
     def __init__(self):
+        super().__init__()
+        self.process = None
         self.forbidden_terms = [
             'import os', 'import sys', 'import subprocess',
             'open(', 'exec(', 'eval(',
@@ -34,52 +38,80 @@ class PythonExecutor:
         code_lower = code.lower()
         return not any(term.lower() in code_lower for term in self.forbidden_terms)
 
-    def execute(self, code: str) -> Dict[str, str]:
+    def execute(self, code_or_command: str) -> Dict[str, str]:
         """Executes Python code in a separate process and returns the result"""
         
         # Basic safety check
-        if not self.basic_code_check(code):
+        if not self.basic_code_check(code_or_command):
             return {
                 'success': False,
                 'output': 'Error: Code contains potentially unsafe operations. You can try and use tools to achieve same functionality.',
                 'error': 'Security check failed'
             }
 
-        # Create a temporary file to store the code
-        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
-            # Properly indent the code to fit inside the try block
-            indented_code = textwrap.indent(code, '    ')
-            # Wrap the indented code to capture output
-            wrapped_code = f"""
+        # Properly indent the code to fit inside the try block
+        indented_code = textwrap.indent(code_or_command, '    ')
+        # Wrap the indented code to capture output
+        wrapped_code = f"""
 try:
 {indented_code}
 except Exception as e:
     print(f"Error: {{str(e)}}")
 """
-            f.write(wrapped_code)
-            temp_file = f.name
+           
 
         try:
             # Execute the code in a subprocess
-            result = subprocess.run(
-                [sys.executable, temp_file],
-                capture_output=True,
+            self.process = subprocess.Popen(
+                [sys.executable, "-u", "-c", wrapped_code],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
                 text=True,
-                timeout=30  # 30 second timeout
+                bufsize=1,  # Line buffered
+                universal_newlines=True
             )
-
+            
+            stdout_data = []
+            stderr_data = []
+            start_time = time.time()
+            
+            # First read all stdout
+            for line in self.process.stdout:
+                # Check for timeout
+                if time.time() - start_time > 30:
+                    self.process.kill()
+                    return {
+                        'success': False,
+                        'output': 'Execution timed out after 30 seconds',
+                        'error': 'Timeout error'
+                    }
+                
+                stdout_data.append(line)
+                print(line, end='', flush=True)  # Print in real-time
+            
+            # Then read all stderr
+            for line in self.process.stderr:
+                # Check for timeout
+                if time.time() - start_time > 30:
+                    self.process.kill()
+                    return {
+                        'success': False,
+                        'output': 'Execution timed out after 30 seconds',
+                        'error': 'Timeout error'
+                    }
+                
+                stderr_data.append(line)
+                print(line, end='', file=sys.stderr, flush=True)  # Print in real-time
+            
+            # Wait for process to complete
+            returncode = self.process.wait()
+            
             return {
-                'success': result.returncode == 0,
-                'output': result.stdout if result.returncode == 0 else result.stderr,
-                'error': result.stderr if result.returncode != 0 else ''
+                'success': returncode == 0,
+                'output': ''.join(stdout_data) if returncode == 0 else ''.join(stderr_data),
+                'error': ''.join(stderr_data) if returncode != 0 else ''
             }
 
-        except subprocess.TimeoutExpired:
-            return {
-                'success': False,
-                'output': 'Execution timed out after 30 seconds',
-                'error': 'Timeout error'
-            }
         except Exception as e:
             return {
                 'success': False,
@@ -87,10 +119,18 @@ except Exception as e:
                 'error': str(e)
             }
         finally:
-            # Clean up the temporary file
+            self.process = None # Reset process
+
+
+    def stop_execution(self):
+        if self.process and hasattr(self.process, 'pid') and self.process.pid is not None:
             try:
-                os.unlink(temp_file)
-            except:
-                pass  # Ignore cleanup errors
-    
+                self.process.terminate()
+                print(f"Attempted to terminate Python process with PID: {self.process.pid}")
+            except Exception as e:
+                print(f"Error terminating Python process with PID {self.process.pid}: {e}")
+            finally:
+                self.process = None
+        else:
+            print("No active Python process to stop.")
 

Env/shell.py

@@ -1,12 +1,21 @@
 import subprocess
+import time
+import sys
+from Src.Env.base_env import BaseEnv
+import re
 
-class ShellExecutor:
-    def execute(self, command: str) -> dict:
+class ShellExecutor(BaseEnv):
+    def __init__(self):
+        super().__init__()
+        self.process = None
+
+    def execute(self, code_or_command: str) -> dict:
         """
-        Executes a shell command and captures its output, error, and success status.
+        Executes a shell command and streams its output in real-time.
+        Strictly prevents execution of harmful commands and access to sensitive directories.
 
         Args:
-            command: The shell command to execute.
+            code_or_command: The shell command to execute.
 
         Returns:
             A dictionary with the following keys:
@@ -15,24 +24,142 @@ class ShellExecutor:
             - 'success': A boolean indicating whether the command executed successfully.
         """
         try:
-            process = subprocess.run(
-                command,
+            # Strict security check for harmful commands and sensitive directories
+            forbidden_patterns = [
+                r'rm\s+-rf\s+/?(\s|$)',
+                r'rm\s+-rf\s+--no-preserve-root',
+                r'rm\s+-rf\s+/\\?',
+                r'shutdown(\s|$)',
+                r'reboot(\s|$)',
+                r'halt(\s|$)',
+                r':(){:|:&};:',  # fork bomb
+                r'chmod\s+777\s+/(\s|$)',
+                r'chown\s+root',
+                r'\bmkfs\b',
+                r'\bdd\b.*\bif=\/dev\/zero\b',
+                r'\bdd\b.*\bof=\/dev\/sda',
+                r'\bpoweroff\b',
+                r'\binit\s+0\b',
+                r'\bsudo\s+rm\s+-rf\s+/?',
+                r'\bsudo\s+shutdown',
+                r'\bsudo\s+reboot',
+                r'\bsudo\s+halt',
+                r'\bsudo\s+mkfs',
+                r'\bsudo\s+dd',
+                r'\bsudo\s+init\s+0',
+                r'\bsudo\s+poweroff',
+                r'\bsudo\s+chmod\s+777\s+/',
+                r'\bsudo\s+chown\s+root',
+                r'\bdel\b.*\/s.*\/q.*\/f.*C:\\',  # Windows
+                r'format\s+C:',
+                r'rd\s+/?s\s+/?q\s+C:\\',
+                r'\bshutdown\b.*\/s',
+                r'\bshutdown\b.*\/r',
+                r'\bshutdown\b.*\/f',
+                r'\bshutdown\b.*\/p',
+                r'\bshutdown\b.*\/t',
+                r'\bshutdown\b.*\/a',
+                r'\bnet\s+user\s+.*\s+/delete',
+                r'\bnet\s+user\s+administrator\s+/active:no',
+                r'\bnet\s+user\s+administrator\s+/active:yes',
+                r'\bnet\s+localgroup\s+administrators\s+.*\s+/delete',
+                r'\bnet\s+localgroup\s+administrators\s+.*\s+/add',
+            ]
+            sensitive_dirs = [
+                '/', '/etc', '/bin', '/usr', '/var', '/root', '/boot', '/dev', '/proc', '/sys', '/lib', '/lib64',
+                'C:\\', 'C:/', 'C:\\Windows', 'C:/Windows', 'C:\\System32', 'C:/System32',
+                'D:\\', 'D:/', 'E:\\', 'E:/'
+            ]
+            # Check for forbidden patterns
+            for pattern in forbidden_patterns:
+                if re.search(pattern, code_or_command, re.IGNORECASE):
+                    return {
+                        'success': False,
+                        'output': 'Blocked potentially harmful command.',
+                        'error': f'Command matches forbidden pattern: {pattern}'
+                    }
+            # Check for sensitive directory access
+            for sensitive_dir in sensitive_dirs:
+                # Only block if the command is trying to directly access or operate on the sensitive dir
+                if re.search(rf'\b{sensitive_dir}\b', code_or_command, re.IGNORECASE):
+                    return {
+                        'success': False,
+                        'output': f'Blocked access to sensitive directory: {sensitive_dir}',
+                        'error': f'Attempted access to sensitive directory: {sensitive_dir}'
+                    }
+            
+            # Execute the command in a subprocess
+            self.process = subprocess.Popen(
+                code_or_command,
                 shell=True,
-                capture_output=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
                 text=True,
-                check=False  # Don't raise an exception on non-zero exit codes
+                bufsize=1,  # Line buffered
+                universal_newlines=True
             )
+            
+            stdout_data = []
+            stderr_data = []
+            start_time = time.time()
+            
+            # First read all stdout
+            for line in self.process.stdout:
+                # Check for timeout
+                if time.time() - start_time > 30:
+                    self.process.kill()
+                    return {
+                        'success': False,
+                        'output': 'Execution timed out after 30 seconds',
+                        'error': 'Timeout error'
+                    }
+                
+                stdout_data.append(line)
+                print(line, end='', flush=True)  # Print in real-time
+            
+            # Then read all stderr
+            for line in self.process.stderr:
+                # Check for timeout
+                if time.time() - start_time > 30:
+                    self.process.kill()
+                    return {
+                        'success': False,
+                        'output': 'Execution timed out after 30 seconds',
+                        'error': 'Timeout error'
+                    }
+                
+                stderr_data.append(line)
+                print(line, end='', file=sys.stderr, flush=True)  # Print in real-time
+            
+            # Wait for process to complete
+            returncode = self.process.wait()
+            
             return {
-                "output": process.stdout,
-                "error": process.stderr,
-                "success": process.returncode == 0,
+                'success': returncode == 0,
+                'output': ''.join(stdout_data) if returncode == 0 else ''.join(stderr_data),
+                'error': ''.join(stderr_data) if returncode != 0 else ''
             }
+
         except Exception as e:
             return {
-                "output": "",
-                "error": str(e),
-                "success": False,
+                'success': False,
+                'output': f'Error: {str(e)}',
+                'error': str(e)
             }
+        finally:
+            self.process = None  # Reset process
+
+    def stop_execution(self):
+        if self.process and hasattr(self.process, 'pid') and self.process.pid is not None:
+            try:
+                self.process.terminate()
+                print(f"Attempted to terminate shell process with PID: {self.process.pid}")
+            except Exception as e:
+                print(f"Error terminating shell process with PID {self.process.pid}: {e}")
+            finally:
+                self.process = None
+        else:
+            print("No active shell process to stop.")
 
 if __name__ == '__main__':
     # Example usage (optional, for testing)

Env/tests/test_python_executor.py

@@ -0,0 +1,70 @@
+import unittest
+from Src.Env.python_executor import PythonExecutor
+import time
+import threading
+
+class TestPythonExecutor(unittest.TestCase):
+
+    def test_execute_simple_script(self):
+        executor = PythonExecutor()
+        result = executor.execute("print('hello python')")
+        self.assertTrue(result["success"])
+        self.assertEqual(result["output"].strip(), "hello python")
+        self.assertEqual(result["error"], "")
+
+    def test_execute_script_with_error(self):
+        executor = PythonExecutor()
+        result = executor.execute("1/0")
+        self.assertFalse(result["success"])
+        self.assertTrue("ZeroDivisionError" in result["output"]) # Error message goes to stdout for python_executor
+        self.assertEqual(result["error"], "") # stderr should be empty as Popen merges stderr to stdout in this case
+
+    def test_stop_execution_long_script(self):
+        executor = PythonExecutor()
+        long_script = "import time; time.sleep(5); print('should not print')"
+
+        execute_result = {}
+        def target():
+            res = executor.execute(long_script)
+            execute_result.update(res)
+
+        thread = threading.Thread(target=target)
+        thread.start()
+
+        time.sleep(1) # Give the script time to start
+        executor.stop_execution()
+
+        thread.join(timeout=2) # Wait for the thread to finish (should be quick after stop)
+
+        self.assertFalse(execute_result.get("success", True), "Script execution should have failed or been stopped.")
+        # Depending on timing, the process might be killed before it produces output,
+        # or it might produce a timeout error, or a specific error from being terminated.
+        # We check if the output indicates it didn't complete normally.
+        output = execute_result.get("output", "")
+        error = execute_result.get("error", "")
+
+        # Check if 'should not print' is NOT in the output
+        self.assertNotIn("should not print", output, "Script should have been terminated before completion.")
+
+        # Check for signs of termination or timeout
+        # This part is a bit tricky as the exact message can vary.
+        # If basic_code_check fails, output is "Error: Code contains potentially unsafe operations..."
+        # If timeout in communicate(), output is "Execution timed out..."
+        # If process is terminated, output might be empty or contain partial error.
+        # For now, we'll accept that if "should not print" is not there, it's a good sign.
+        # A more robust check might involve looking for specific error messages related to termination if available.
+
+        # A simple check that process is no longer listed in executor
+        self.assertIsNone(executor.process, "Executor process should be None after stopping.")
+
+
+    def test_stop_execution_no_script(self):
+        executor = PythonExecutor()
+        try:
+            executor.stop_execution()
+        except Exception as e:
+            self.fail(f"stop_execution with no script raised an exception: {e}")
+        self.assertIsNone(executor.process, "Executor process should be None.")
+
+if __name__ == '__main__':
+    unittest.main()

Env/tests/test_shell_executor.py

@@ -0,0 +1,29 @@
+import unittest
+from Src.Env.shell import ShellExecutor
+
+class TestShellExecutor(unittest.TestCase):
+
+    def test_execute_success(self):
+        executor = ShellExecutor()
+        result = executor.execute("echo hello")
+        self.assertTrue(result["success"])
+        self.assertEqual(result["output"].strip(), "hello")
+        self.assertEqual(result["error"], "")
+
+    def test_execute_error(self):
+        executor = ShellExecutor()
+        result = executor.execute("ls non_existent_directory_for_test")
+        self.assertFalse(result["success"])
+        self.assertEqual(result["output"], "")
+        self.assertTrue("non_existent_directory_for_test" in result["error"])
+
+    def test_stop_execution(self):
+        executor = ShellExecutor()
+        # This test mainly ensures that calling stop_execution doesn't raise an error
+        try:
+            executor.stop_execution()
+        except Exception as e:
+            self.fail(f"stop_execution raised an exception: {e}")
+
+if __name__ == '__main__':
+    unittest.main()

OpenCopilot.py

@@ -6,9 +6,40 @@ from prompt_toolkit import PromptSession, HTML
 from prompt_toolkit.completion import Completer, Completion
 from prompt_toolkit.shortcuts import print_formatted_text
 from prompt_toolkit.formatted_text import FormattedText
+from prompt_toolkit.key_binding import KeyBindings
+from prompt_toolkit.application.current import get_app
 
+from Tools.file_task import file_reader
 from Agents.Executor.executor import executor
 
+# Custom key bindings for handling Enter key
+kb = KeyBindings()
+
+@kb.add('enter')
+def handle_enter(event):
+    """Handle Enter key press based on completion state"""
+    buff = event.current_buffer
+    app = get_app()
+    
+    # If completion menu is open, select the current completion
+    if app.current_buffer.complete_state:
+        # Get the current completion
+        current_completion = app.current_buffer.complete_state.current_completion
+        if current_completion:
+            # Get the text before cursor
+            text_before_cursor = buff.text[:buff.cursor_position]
+            # Find the last @ symbol
+            last_at_pos = text_before_cursor.rindex('@')
+            # Delete text from @ to cursor position
+            buff.delete_before_cursor(count=buff.cursor_position - last_at_pos - 1)
+            # Insert the completion
+            buff.insert_text(current_completion.text)
+        # Close the completion menu
+        buff.complete_state = None
+    else:
+        # If no completion menu, submit the command
+        buff.validate_and_handle()
+
 class FilePathCompleter(Completer):
     def get_completions(self, document, complete_event):
         text_before_cursor = document.text_before_cursor
@@ -74,7 +105,12 @@ class FilePathCompleter(Completer):
 class OpenCopilot:
     def __init__(self):
         self.e1 = None  # Initialize as None, will be set in run()
-        self.session = PromptSession(completer=FilePathCompleter())
+        # Initialize session with custom key bindings
+        self.session = PromptSession(
+            completer=FilePathCompleter(),
+            key_bindings=kb,  # Add custom key bindings
+            complete_while_typing=True  # Enable completion while typing
+        )
 
     def extract_files_and_process_prompt(self, user_input):
         """Extract file paths from @ commands and process the prompt."""
@@ -93,10 +129,11 @@ class OpenCopilot:
             
             if os.path.exists(expanded_path):
                 if os.path.isfile(expanded_path):
-                    try:
-                        with open(expanded_path, 'r', encoding='utf-8') as f:
-                            content = f.read()
-                        
+                    # Call file_reader to get file content
+                    file_read_result = file_reader(file_path=expanded_path)
+                    
+                    if file_read_result["success"]:
+                        content = file_read_result["output"]
                         # Add file content with clear formatting
                         file_contents.append(f"=== Content of file: {expanded_path} ===\n{content}\n=== End of file: {expanded_path} ===\n")
                         # Remove the @file pattern from the processed prompt
@@ -105,10 +142,10 @@ class OpenCopilot:
                         print_formatted_text(FormattedText([
                             ('class:success', f"✓ Loaded file: {expanded_path}")
                         ]))
-                        
-                    except Exception as e:
+                    else:
+                        error_message = file_read_result["output"]
                         print_formatted_text(FormattedText([
-                            ('class:error', f"✗ Error reading file {expanded_path}: {str(e)}")
+                            ('class:error', f"✗ Error reading file {expanded_path}: {error_message}")
                         ]))
                 else:
                     # For directories, just append the path to the processed prompt

Tools/tool_dir.json

@@ -70,5 +70,21 @@
             "prompt": "The message to display to the user",
             "input_type": "Type of input to validate (text, number, boolean, optional, defaults to text)"
         }
+    },
+    {
+        "name": "execute_python_code",
+        "summary": "Executes a given snippet of Python code. The code should be self-contained. Input should be a JSON object with a 'code' key containing the Python code string.",
+        "arguments": {
+            "code": "The Python code to execute."
+        },
+        "input_schema": {"type": "object", "properties": {"code": {"type": "string"}}, "required": ["code"]}
+    },
+    {
+        "name": "execute_shell_command",
+        "summary": "Executes a given shell command. Input should be a JSON object with a 'command' key containing the shell command string.",
+        "arguments": {
+            "command": "The shell command to execute."
+        },
+        "input_schema": {"type": "object", "properties": {"command": {"type": "string"}}, "required": ["command"]}
     }
 ]

Tools/tool_manager.py

@@ -7,26 +7,64 @@ from Tools.web_search import web_search
 from Tools.file_task import file_reader, file_maker, file_writer, directory_maker
 from Tools.system_details import get_os_details, get_datetime, get_memory_usage, get_cpu_info
 from Tools.userinp import get_user_input
+from Src.Env.python_executor import PythonExecutor
+from Src.Env.shell import ShellExecutor
 
 #need to transform it into map of dictionary
 #name : [function : xyz,description : blah bah]
 
-tools_function_map = {
-    "web_loader": load_data,
-    "web_search": web_search,
-    "file_maker": file_maker,
-    "file_reader":file_reader,
-    "directory_maker":directory_maker,
-    "file_writer":file_writer,
-    "get_os_details": get_os_details,
-    "get_datetime": get_datetime,
-    "get_memory_usage": get_memory_usage,
-    "get_cpu_info": get_cpu_info,
-    "get_user_input": get_user_input
-}
 
 
 
+def execute_python_code_tool(code: str) -> str:
+    """ 
+    Prompts for confirmation, then executes the given Python code and returns a formatted result string.
+    """
+    user_confirmation = input(f"Do you want to execute this Python code snippet?\n```python\n{code}\n```\n(y/n): ")
+    if user_confirmation.lower() != 'y':
+        print("Python code execution skipped by the user.")
+        return "User chose not to execute the Python code."
+    executor = PythonExecutor()
+    result = executor.execute(code)
+    if result['output'] == "" and not result['success']:
+        error_msg = (
+            f"Python execution failed.\n"
+            f"Error: {result.get('error', 'Unknown error')}"
+        )
+        return error_msg
+    elif result['output'] == "":
+        no_output_msg = (
+            "Python execution completed but no output was produced. "
+            "Ensure your code includes print() statements to show results."
+        )
+        return no_output_msg
+    else:
+        if result['success']:
+            return f"Program Output:\n{result['output']}"
+        else:
+            return f"Program Output:\n{result['output']}\nError: {result.get('error', 'Unknown error')}"
+
+def execute_shell_command_tool(command: str) -> str:
+    """
+    Prompts for confirmation, then executes the given shell command and returns a formatted result string.
+    """
+    user_confirmation = input(f"Do you want to execute the shell command: '{command}'? (y/n): ")
+    if user_confirmation.lower() != 'y':
+        print("Shell command execution skipped by the user.")
+        return "User chose not to execute the shell command."
+    executor = ShellExecutor()
+    result = executor.execute(command)
+    if result['output'] == "":
+        if result['success']:
+            return "Shell command executed successfully with no output."
+        else:
+            return f"Shell command executed with no output, but an error occurred: {result.get('error', 'Unknown error')}"
+    else:
+        if result['success']:
+            return f"Command Output:\n{result['output']}"
+        else:
+            return f"Command Output:\n{result['output']}\nError: {result.get('error', 'Unknown error')}"
+
 def call_tool(tool_name, tool_input):
     """
     Calls the appropriate tool function with the given input.
@@ -41,6 +79,21 @@ def call_tool(tool_name, tool_input):
     else:
         raise ValueError(f"Tool '{tool_name}' not found. Check the tools available in the tool directory")
 
+tools_function_map = {
+    "web_loader": load_data,
+    "web_search": web_search,
+    "file_maker": file_maker,
+    "file_reader":file_reader,
+    "directory_maker":directory_maker,
+    "file_writer":file_writer,
+    "get_os_details": get_os_details,
+    "get_datetime": get_datetime,
+    "get_memory_usage": get_memory_usage,
+    "get_cpu_info": get_cpu_info,
+    "get_user_input": get_user_input,
+    "execute_python_code": execute_python_code_tool,
+    "execute_shell_command": execute_shell_command_tool,
+}
 
 # print(call_tool("web_loader","https://www.toastmasters.org"))
 # print(call_tool("web_search","manus ai"))

Utils/executor_utils.py

@@ -13,21 +13,3 @@ def parse_tool_call(response: str) -> Optional[dict]:
         except json.JSONDecodeError:
             return None
     return None
-
-def parse_code(response: str) -> Optional[str]:
-    """
-    Parses code from the response.
-    """
-    if "<<CODE>>" in response and "<<CODE>>" in response: # There was a typo in the original file, it checked for <<CODE>> twice
-        code = response.split("<<CODE>>")[1].split("<<CODE>>")[0].strip()
-        return code
-    return None
-
-def parse_shell_command(response: str) -> Optional[str]:
-    """
-    Parses a shell command from the response.
-    """
-    if "<<SHELL_COMMAND>>" in response and "<<END_SHELL_COMMAND>>" in response:
-        shell_command = response.split("<<SHELL_COMMAND>>")[1].split("<<END_SHELL_COMMAND>>")[0].strip()
-        return shell_command
-    return None

cli.py

@@ -296,5 +296,64 @@ def list_models():
         for model_name_full in model_list: # model_name_full is "openai/gpt-4o", etc.
             click.echo(f"  - {model_name_full}")
 
+@cli.command('set-api-key')
+@click.option('--provider', '-p', type=click.Choice(list(AVAILABLE_MODELS.keys())), 
+              help='The LLM provider to set API key for')
+@click.option('--key', '-k', help='The API key to set (if not provided, will prompt for it)')
+def set_api_key(provider, key):
+    """Set or update API key for a specific LLM provider"""
+    if not provider:
+        # If no provider specified, ask user to choose
+        questions = [
+            inquirer.List('provider_key',
+                message="Select LLM Provider to update API key",
+                choices=list(AVAILABLE_MODELS.keys())
+            )
+        ]
+        provider = inquirer.prompt(questions)['provider_key']
+    
+    # Get the environment variable name for this provider
+    env_var = API_KEYS.get(provider)
+    if not env_var:
+        raise ValueError(f"Unknown provider: {provider}")
+    
+    # Get the API key (either from command line or prompt)
+    if not key:
+        questions = [
+            inquirer.Text('api_key',
+                message=f"Enter your {provider.upper()} API key",
+                validate=lambda _, x: len(x.strip()) > 0
+            )
+        ]
+        key = inquirer.prompt(questions)['api_key']
+    
+    # Get the path to .env file
+    env_path = os.path.join(os.path.dirname(__file__), '../.env')
+    
+    # Read existing .env file if it exists
+    lines = []
+    if os.path.exists(env_path):
+        with open(env_path, 'r') as f:
+            lines = f.readlines()
+    
+    # Update or add the API key
+    key_line = f"{env_var}={key}\n"
+    key_exists = False
+    
+    for i, line in enumerate(lines):
+        if line.strip().startswith(f"{env_var}=") or line.strip().startswith(f"#{env_var}="):
+            lines[i] = key_line
+            key_exists = True
+            break
+    
+    if not key_exists:
+        lines.append(key_line)
+    
+    # Write back to .env file
+    with open(env_path, 'w') as f:
+        f.writelines(lines)
+    
+    click.echo(f"API key for {provider.upper()} has been updated successfully in {env_path}")
+
 if __name__ == '__main__':
     cli() 

{pikoai-0.1.6.dist-info → pikoai-0.1.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: PikoAi
-Version: 0.1.6
+Version: 0.1.9
 Summary: An AI-powered task automation tool
 Home-page: https://github.com/nihaaaar22/OS-Assistant
 Author: Nihar S

pikoai-0.1.9.dist-info/RECORD

@@ -0,0 +1,35 @@
+OpenCopilot.py,sha256=vBAuU6MNhYdnyIIrDWOBSlxa0iyjADJLq5eCxUhlffw,12240
+cli.py,sha256=hY6KUxvKvJOFThZT--r6m2nzOEMIOVtiVptewsi9Z9w,13868
+Agents/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+Agents/Executor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+Agents/Executor/executor.py,sha256=lbtJejQBt7O4kBf-_JG1QM9QG7RHCO8VVqlLn00S7FE,8812
+Agents/Executor/prompts.py,sha256=bcVxBBCEda7JZtnA1KoMM8KayKAqFeQCP4lY6EbodzQ,3607
+Env/__init__.py,sha256=KLe7UcNV5L395SxhMwbYGyu7KPrSNaoV_9QJo3mLop0,196
+Env/base_env.py,sha256=K4PoWwPXn3pKeu7_-JOlUuyNbyYQ9itMhQybFOm-3K4,1563
+Env/base_executor.py,sha256=awTwJ44CKWV4JO2KUHfHDX0p1Ujw55hlaL5oNYTEW9M,893
+Env/env.py,sha256=I3lOoyBJR5QNq3tWk_HVH6ncRx1vnilRYmj7b7h4jic,114
+Env/js_executor.py,sha256=tEAg5ho8Pa8LzxUbS1Idau8XuJWZZqPiNlvFPDwGkgc,2690
+Env/python_executor.py,sha256=i1SelJ2GdQpGVjtAVQmp__0bA95OurnU4DUyu_Yfb9g,4818
+Env/shell.py,sha256=7jbSMPh-1Cx6xCzvS78MQbneYfBSNlItl-EHjDR1Zw4,7294
+Env/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+Env/tests/test_python_executor.py,sha256=5kfs0cOf-RgWTOers_1wB0yvYSF-HrHPsraJ-PxgR50,3156
+Env/tests/test_shell_executor.py,sha256=-RcCdSUMoRAXHISIh0IR5MCutco80fX2S3sQBcinc_Q,1034
+Tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+Tools/file_task.py,sha256=VUhWq_G-SWvGahQo8PG7TOpElHUW3BGLUabrTdJS89o,12151
+Tools/system_details.py,sha256=7-mTm3CG4NoatHcvcosalEgEcpWlNsCsZ7kuS3y_EmY,2262
+Tools/tool_dir.json,sha256=pEqI-1rKud-VEWHYLF2Pbv6xwQZWy62HUF_1hjcdhCY,3128
+Tools/tool_manager.py,sha256=Q42jnN9eoBm98cgyEpAnINPZYKegdKO-T6zQY2Ns2jE,4129
+Tools/userinp.py,sha256=vUhEj3y1W1_ZFHqo2xQwvqDyeOg3VsisSKTI0EurUH8,1205
+Tools/web_loader.py,sha256=PyZk2g7WngZT0tCLs9Danx20dYspnaZwy4rlVE9Sx_4,5054
+Tools/web_search.py,sha256=4EGq1VZqfDgG-_yXTd4_Ha1iEUcR-szdlgRV7oFPru4,1259
+Utils/__init__.py,sha256=oukU0ufroPRd8_N8d2xiFes9CTxSaw4NA6p2nS1kkSg,16
+Utils/executor_utils.py,sha256=WwK3TKgw_hG_crg7ijRaqfidYnnNXYbbs37vKZRYK-0,491
+Utils/ter_interface.py,sha256=Zay9mwyAyKYTNQAKOWXHAa3upo9TWprSf26STiHXk0g,6255
+llm_interface/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+llm_interface/llm.py,sha256=tI_KDOW14QLWowA7bB3GPe2qjlk0sjS5fBavs9XD1fo,5185
+pikoai-0.1.9.dist-info/licenses/LICENSE,sha256=cELUVOboOAderKFp8bdtcM5VyJi61YH1oDbRhOuoQZw,1067
+pikoai-0.1.9.dist-info/METADATA,sha256=4S0C2M5G7pjDx37Ydv40ks8jTL94-r1Nah00AJcOj6Q,2961
+pikoai-0.1.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+pikoai-0.1.9.dist-info/entry_points.txt,sha256=xjZnheDymNDnQ0o84R0jZKEITrhNbzQWN-AhqfA_d6s,50
+pikoai-0.1.9.dist-info/top_level.txt,sha256=hWzBNE7UQsuNcENIOksGcJED08k3ZGRRn2X5jnStICU,53
+pikoai-0.1.9.dist-info/RECORD,,

pikoai-0.1.6.dist-info/RECORD

@@ -1,32 +0,0 @@
-OpenCopilot.py,sha256=6210OONj0yV-LPzkCdZ1RUiHC5NZvLsD-y93Y1zurX0,10687
-cli.py,sha256=o2V_DETYXEiYm6617Xg8rDcAgb0ofmojPvdlOpNHwMY,11809
-Agents/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-Agents/Executor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-Agents/Executor/executor.py,sha256=zaFO4IQNE1PscvGPV8zoRsQAJAYVTz2JsD_zUBLyYwg,12360
-Agents/Executor/prompts.py,sha256=wLS3lPAYWjeKF02LzJ8vP5bZ2VQrMJUd4A7rBfl6qSQ,3846
-Env/__init__.py,sha256=KLe7UcNV5L395SxhMwbYGyu7KPrSNaoV_9QJo3mLop0,196
-Env/base_env.py,sha256=ORM6U5qwj7cTuSHFtSmCSsE0cl6pZ28D97CEyyFnucI,1323
-Env/base_executor.py,sha256=awTwJ44CKWV4JO2KUHfHDX0p1Ujw55hlaL5oNYTEW9M,893
-Env/env.py,sha256=I3lOoyBJR5QNq3tWk_HVH6ncRx1vnilRYmj7b7h4jic,114
-Env/js_executor.py,sha256=tEAg5ho8Pa8LzxUbS1Idau8XuJWZZqPiNlvFPDwGkgc,2690
-Env/python_executor.py,sha256=KiQxyxLgdEph1iRBLteqYy6KJGE8ag_jxiYJ_-BEEB4,3221
-Env/shell.py,sha256=gr6czmeuSWtB3xSA9TZN7wnK2BENOuA9zjNttwbxztU,1877
-Tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-Tools/file_task.py,sha256=VUhWq_G-SWvGahQo8PG7TOpElHUW3BGLUabrTdJS89o,12151
-Tools/system_details.py,sha256=7-mTm3CG4NoatHcvcosalEgEcpWlNsCsZ7kuS3y_EmY,2262
-Tools/tool_dir.json,sha256=1zF0Z8ATZn5GT5lGq0cGRlBQoh0F655RP_htkZwnJHI,2302
-Tools/tool_manager.py,sha256=0i3bd_VxhbpWKLzyfSeYyv_33Z6HmvQDBUxPUxjLYlU,1736
-Tools/userinp.py,sha256=vUhEj3y1W1_ZFHqo2xQwvqDyeOg3VsisSKTI0EurUH8,1205
-Tools/web_loader.py,sha256=PyZk2g7WngZT0tCLs9Danx20dYspnaZwy4rlVE9Sx_4,5054
-Tools/web_search.py,sha256=4EGq1VZqfDgG-_yXTd4_Ha1iEUcR-szdlgRV7oFPru4,1259
-Utils/__init__.py,sha256=oukU0ufroPRd8_N8d2xiFes9CTxSaw4NA6p2nS1kkSg,16
-Utils/executor_utils.py,sha256=eVVxZDcAGxENRV5DhCkgh1AKiJEgfV_1q-iSQNWR5kI,1180
-Utils/ter_interface.py,sha256=Zay9mwyAyKYTNQAKOWXHAa3upo9TWprSf26STiHXk0g,6255
-llm_interface/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-llm_interface/llm.py,sha256=tI_KDOW14QLWowA7bB3GPe2qjlk0sjS5fBavs9XD1fo,5185
-pikoai-0.1.6.dist-info/licenses/LICENSE,sha256=cELUVOboOAderKFp8bdtcM5VyJi61YH1oDbRhOuoQZw,1067
-pikoai-0.1.6.dist-info/METADATA,sha256=3EMhRNO8Abcy9dy0QUJEwYYJvajnvbqCmqqhccpD8es,2961
-pikoai-0.1.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-pikoai-0.1.6.dist-info/entry_points.txt,sha256=xjZnheDymNDnQ0o84R0jZKEITrhNbzQWN-AhqfA_d6s,50
-pikoai-0.1.6.dist-info/top_level.txt,sha256=hWzBNE7UQsuNcENIOksGcJED08k3ZGRRn2X5jnStICU,53
-pikoai-0.1.6.dist-info/RECORD,,