PikoAi 0.1.0__py3-none-any.whl

Agents/Executor/executor.py

@@ -0,0 +1,242 @@
+# the change in this executor is the the tasks will not be iterated in a for loop and execution will not be done one by one
+# instead it would be asked what is the next course of action
+
+import os
+import sys
+import time
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../')))
+from Src.Utils.ter_interface import TerminalInterface
+from Src.Utils.executor_utils import parse_tool_call, parse_code, parse_shell_command
+from Src.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 Src.Env import python_executor
+from Src.Env.shell import ShellExecutor # Import ShellExecutor
+from Src.llm_interface.llm import LiteLLMInterface # Import LiteLLMInterface
+
+from Src.Tools import tool_manager
+
+class RateLimiter:
+    def __init__(self, wait_time: float = 5.0, max_retries: int = 3):
+        self.wait_time = wait_time
+        self.max_retries = max_retries
+        self.last_call_time = None
+
+    def wait_if_needed(self):
+        if self.last_call_time is not None:
+            elapsed = time.time() - self.last_call_time
+            if elapsed < 1.0:
+                time.sleep(1.0 - elapsed)
+        self.last_call_time = time.time()
+
+class executor:
+    def __init__(self, user_prompt, max_iter=10):
+        self.user_prompt = user_prompt
+        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.message = [
+            {"role": "system", "content": self.system_prompt},
+            {"role": "user", "content": self.task_prompt}
+        ]
+        self.terminal = TerminalInterface()
+        self.initialize_llm()
+
+    def initialize_llm(self):
+        # Directly instantiate LiteLLMInterface. 
+        # It handles its own configuration loading (including model_name from config.json).
+        self.llm = LiteLLMInterface()
+
+    def get_tool_dir(self):
+        # Get the absolute path to the project root directory
+        project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../'))
+        tool_dir_path = os.path.join(project_root, 'Src', 'Tools', 'tool_dir.json')
+        with open(tool_dir_path, "r") as file:
+            return file.read()
+
+    def executor_prompt_init(self):
+        # Load tools details when initializing prompt
+        tools_details = self.get_tool_dir()
+
+        # Read working_directory from config.json
+        # This import needs to be here, or moved to the top if json is used elsewhere
+        import json 
+        with open(os.path.join(os.path.dirname(__file__), '../../../config.json'), "r") as config_file:
+            config = json.load(config_file)
+            working_dir = config.get("working_directory", "")
+
+        self.system_prompt = get_system_prompt(self.user_prompt, working_dir, tools_details)
+        self.task_prompt = get_task_prompt()
+
+    def run_inference(self):
+        retries = 0
+        while retries <= self.rate_limiter.max_retries:
+            try:
+                self.rate_limiter.wait_if_needed()
+
+                response = self.llm.chat(self.message) # LiteLLMInterface.chat() returns the full response string
+
+                # Streaming is handled within LiteLLMInterface.chat()
+                # and TerminalInterface.process_markdown_chunk()
+                self.message.append({"role": "assistant", "content": response})
+                return response
+
+            except Exception as e: # Catching generic Exception as LiteLLM maps to OpenAI exceptions
+                # Check if the error message contains "429" for rate limiting
+                if "429" in str(e) and retries < self.rate_limiter.max_retries:
+                    retries += 1
+                    print(f"\nRate limit error detected. Waiting {self.rate_limiter.wait_time} seconds before retry {retries}/{self.rate_limiter.max_retries}")
+                    time.sleep(self.rate_limiter.wait_time)
+                # Check if the error is an SDKError (though less likely with LiteLLM directly)
+                # or if it's any other exception that we should retry or raise.
+                elif isinstance(e, SDKError) and "429" in str(e) and retries < self.rate_limiter.max_retries: # Added SDKError check just in case
+                    retries += 1
+                    print(f"\nRate limit exceeded (SDKError). Waiting {self.rate_limiter.wait_time} seconds before retry {retries}/{self.rate_limiter.max_retries}")
+                    time.sleep(self.rate_limiter.wait_time)
+                else:
+                    print(f"\nError occurred during inference: {str(e)}")
+                    # You might want to log the full traceback here for debugging
+                    # import traceback
+                    # print(traceback.format_exc())
+                    raise
+        raise Exception("Failed to complete inference after maximum retries")
+
+    def run(self):
+
+        self.run_task()
+
+    def run_task(self):
+        # Remove tools_details parameter since it's in the prompt
+        task_message = self.task_prompt
+
+        self.message.append({"role": "user", "content": task_message})
+
+        iteration = 0
+        task_done = False
+
+        while iteration < self.max_iter and not task_done:
+            # 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']}")
+                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}"})
+                except ValueError as e:
+                    error_msg = str(e)
+                    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.")
+
+            # Check if task is done
+            if "TASK_DONE" in response:
+                
+                task_done = True
+
+            else:
+                self.message.append({"role": "user", "content": "If the task i mentioned is complete then output TASK_DONE .If not then run another iteration."})
+                iteration += 1
+
+        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
+
+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()

Agents/Executor/prompts.py

@@ -0,0 +1,99 @@
+# This file contains the prompts used by the Executor agent.
+
+import platform
+
+def get_system_prompt(user_prompt: str, working_dir: str, tools_details: str) -> str:
+    """
+    Returns the system prompt for the Executor agent.
+    """
+    os_name = platform.system()
+    return f"""You are a terminal-based operating system assistant designed to help users achieve their goals by executing tasks provided in text format. The current user goal is: {user_prompt}.
+
+Working Directory: {working_dir}
+Operating System: {os_name}
+
+You have access to the following tools:
+{tools_details}
+
+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
+
+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.
+
+### Action Guidelines:
+- **Tool Call**: Use when a specific tool can help with the current step. Format:
+  <<TOOL_CALL>>
+  {{
+    "tool_name": "name_of_tool",
+    "input": {{
+      "key": "value"   //Replace 'key' with the actual parameter name for the tool
+    }}
+  }}
+  <<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>>
+- **Direct Response**: Provide a direct answer if the task doesn't require tools or code.
+
+### Important Notes:
+- Perform only one action per step.
+- Always evaluate the output of each action before deciding the next step.
+- Continue performing actions until the user's goal is fully achieved. Only then, include 'TASK_DONE' in your response.
+- Do not end the task immediately after a tool call or code execution without evaluating its output.
+
+Now, carefully plan your approach and start with the first step to achieve the user's goal.
+"""
+
+def get_task_prompt() -> str:
+    """
+    Returns the task prompt for the Executor agent.
+    """
+    return """
+Following are the things that you must read carefully and remember:
+
+        - For tool calls, use:
+        <<TOOL_CALL>>
+        {
+            "tool_name": "name_of_tool",
+            "input": {
+            "key": "value"  // Use the correct parameter name for each tool
+            }
+        }
+        <<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>>
+
+        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
+        checking its output. 
+        You can only execute a single tool call or code execution at a time, then check its ouput
+        then proceed with the next call
+        Use the working directory as the current directory for all file operations unless otherwise specified.
+        
+        
+
+        These are the things that you learn't from the mistakes you made earlier :
+
+        - When given a data file and asked to understand data/do data analysis/ data visualisation or similar stuff
+        do not use file reader and read the whole data. Only use python code to do the analysis
+        - This is a standard Python environment, not a python notebook or a repl. previous execution
+         context is not preserved between executions.
+        - You have a get_user_input tool to ask user more context before, in between or after tasks
+
+"""

Env/__init__.py

@@ -0,0 +1,5 @@
+from .base_executor import BaseExecutor
+from .python_executor import PythonExecutor
+from .js_executor import JavaScriptExecutor
+
+__all__ = ["BaseExecutor", "PythonExecutor", "JavaScriptExecutor"]

Env/base_env.py

@@ -0,0 +1,44 @@
+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 execute(self):
+        raise NotImplementedError("This method should be overridden by subclasses")
+
+      
+def create_environment(language):
+    if language == "python":
+        return PythonExecutor()
+    elif language == "javascript":
+        return JavaScriptExecutor()
+    else:
+        raise ValueError(f"Unsupported language: {language}")
+    
+
+    # def stop(self):
+    #         """
+    #         Stops the execution of all active languages.
+    #         """        
+    #         for language in self._active_languages.values():
+    #             language.stop()
+
+    # def terminate(self):
+    #     """
+    #     Terminates all active language environments.
+    #     """        
+    #     for language_name in list(self._active_languages.keys()):
+    #         language = self._active_languages[language_name]
+    #         if (
+    #             language
+    #         ):  # Not sure why this is None sometimes. We should look into this
+    #             language.terminate()
+    #         del self._active_languages[language_name]

Env/base_executor.py

@@ -0,0 +1,24 @@
+# mostly won't need this. there will be only one base env
+
+from abc import ABC, abstractmethod
+
+class BaseExecutor(ABC):
+    pass
+#     def __init__(self):
+#         pass
+
+#     @abstractmethod
+#     def execute(self, code: str):
+#         """Executes the given code and returns the output or error."""
+#         pass
+
+#     def validate_code(self, code: str) -> bool:
+#         """Basic validation to ensure code is safe to execute. Override for custom rules."""
+#         # Implement basic checks (e.g., preventing dangerous commands)
+#         if "import os" in code or "import sys" in code:
+#             return False  # Disallow potentially unsafe imports for security
+#         return True
+
+#     def handle_error(self, error: Exception) -> str:
+#         """Handles errors during execution and returns a standardized error message."""
+#         return f"Execution failed: {str(error)}"

Env/env.py

@@ -0,0 +1,5 @@
+#mostly won't need this. There will be only one base env
+
+# import base_env from base_env
+
+# class env(base_env):

Env/js_executor.py

@@ -0,0 +1,63 @@
+import subprocess
+import sys
+
+class JavaScriptExecutor():
+    def __init__(self):
+        super().__init__()
+        self.node_installed = self.check_node_installed()
+
+    def check_node_installed(self) -> bool:
+        """Checks if Node.js is installed on the system."""
+        try:
+            subprocess.run(["node", "-v"], capture_output=True, check=True)
+            return True
+        except subprocess.CalledProcessError:
+            return False
+        except FileNotFoundError:
+            return False
+
+    def install_node(self) -> bool:
+        """Attempts to install Node.js based on the operating system."""
+        try:
+            if sys.platform.startswith("linux"):
+                # Try to install Node.js using apt-get (for Debian/Ubuntu-based systems)
+                subprocess.run(["sudo", "apt-get", "update"], check=True)
+                subprocess.run(["sudo", "apt-get", "install", "-y", "nodejs"], check=True)
+            elif sys.platform == "darwin":
+                # Try to install Node.js using Homebrew on macOS
+                subprocess.run(["brew", "install", "node"], check=True)
+            elif sys.platform == "win32":
+                # Check if Chocolatey is installed, and install Node.js
+                subprocess.run(["choco", "install", "nodejs", "-y"], check=True)
+            else:
+                return False  # Unsupported OS for automatic installation
+            return True
+        except subprocess.CalledProcessError:
+            return False  # Installation failed
+
+    def execute(self, code: str) -> str:
+        """Executes JavaScript code using Node.js and returns the result or an error message."""
+        # Check if Node.js is installed, attempt installation if not
+        if not self.node_installed:
+            if not self.install_node():
+                return "Node.js is not installed, and automatic installation failed. Please install Node.js manually."
+
+        # Recheck after attempted installation
+        self.node_installed = self.check_node_installed()
+        if not self.node_installed:
+            return "Node.js is required but not installed. Please install Node.js manually."
+
+        # Proceed with code execution if Node.js is available
+        # if not self.validate_code(code):
+        #     return "Code validation failed: Unsafe code detected."
+
+        try:
+            result = subprocess.run(
+                ["node", "-e", code],
+                capture_output=True,
+                text=True,
+                check=True
+            )
+            return result.stdout if result.stdout else "Code executed successfully."
+        except subprocess.CalledProcessError as e:
+            print(e)

Env/python_executor.py

@@ -0,0 +1,96 @@
+# from .base_executor import BaseExecutor
+
+# class PythonExecutor():
+#     def execute(self, code: str) -> str:
+#         """Executes Python code and returns the result or an error message."""
+
+#         # if not self.validate_code(code):
+#         #     return "Code validation failed: Unsafe code detected."
+
+#         local_vars = {}
+#         try:
+#             exec(code, {}, local_vars)  # Execute code in an isolated environment
+#             return local_vars.get("output", "Code executed successfully.")
+#         except Exception as e:
+#             # return self.handle_error(e)
+#             print("error in running python code", e)
+
+import subprocess
+import tempfile
+import os
+from typing import Dict
+import textwrap
+import sys
+
+class PythonExecutor:
+    def __init__(self):
+        self.forbidden_terms = [
+            'import os', 'import sys', 'import subprocess',
+            'open(', 'exec(', 'eval(',
+        ]
+
+    def basic_code_check(self, code: str) -> bool:
+        """Simple check for potentially dangerous code"""
+        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]:
+        """Executes Python code in a separate process and returns the result"""
+        
+        # Basic safety check
+        if not self.basic_code_check(code):
+            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"""
+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,
+                text=True,
+                timeout=30  # 30 second timeout
+            )
+
+            return {
+                'success': result.returncode == 0,
+                'output': result.stdout if result.returncode == 0 else result.stderr,
+                'error': result.stderr if result.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,
+                'output': f'Error: {str(e)}',
+                'error': str(e)
+            }
+        finally:
+            # Clean up the temporary file
+            try:
+                os.unlink(temp_file)
+            except:
+                pass  # Ignore cleanup errors
+    
+

Env/shell.py

@@ -0,0 +1,55 @@
+import subprocess
+
+class ShellExecutor:
+    def execute(self, command: str) -> dict:
+        """
+        Executes a shell command and captures its output, error, and success status.
+
+        Args:
+            command: The shell command to execute.
+
+        Returns:
+            A dictionary with the following keys:
+            - 'output': The captured standard output (string).
+            - 'error': The captured standard error (string).
+            - 'success': A boolean indicating whether the command executed successfully.
+        """
+        try:
+            process = subprocess.run(
+                command,
+                shell=True,
+                capture_output=True,
+                text=True,
+                check=False  # Don't raise an exception on non-zero exit codes
+            )
+            return {
+                "output": process.stdout,
+                "error": process.stderr,
+                "success": process.returncode == 0,
+            }
+        except Exception as e:
+            return {
+                "output": "",
+                "error": str(e),
+                "success": False,
+            }
+
+if __name__ == '__main__':
+    # Example usage (optional, for testing)
+    executor = ShellExecutor()
+
+    # Test case 1: Successful command
+    result1 = executor.execute("echo 'Hello, World!'")
+    print(f"Test Case 1 Result: {result1}")
+
+    # Test case 2: Command with an error
+    result2 = executor.execute("ls non_existent_directory")
+    print(f"Test Case 2 Result: {result2}")
+
+    # Test case 3: Command that succeeds but writes to stderr (e.g. some warnings)
+    result3 = executor.execute("echo 'Error output' >&2")
+    print(f"Test Case 3 Result: {result3}")
+
+    # Test case 4: Command that produces no output
+    result4 = executor.execute(":") # The ':' command is a no-op in bash
+    print(f"Test Case 4 Result: {result4}")