diff --git a/cortex_on/Dockerfile b/cortex_on/Dockerfile index 8d7373e..8c74c4f 100644 --- a/cortex_on/Dockerfile +++ b/cortex_on/Dockerfile @@ -4,20 +4,40 @@ WORKDIR /app COPY requirements.txt . RUN pip install uv + +# Install build tools and Docker RUN apt-get update && apt-get install -y \ build-essential \ cmake \ g++ \ + apt-transport-https \ + ca-certificates \ + curl \ + gnupg \ + lsb-release \ + && rm -rf /var/lib/apt/lists/* + +# Install Docker CLI +RUN curl -fsSL https://download.docker.com/linux/debian/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg \ + && echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/debian \ + $(lsb_release -cs) stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null \ + && apt-get update && apt-get install -y docker-ce-cli \ && rm -rf /var/lib/apt/lists/* RUN export PYTHONPATH=/app -RUN apt-get update -y && apt-get install build-essential -y # Add the --system flag to uv pip install RUN uv pip install --system --no-cache-dir -r requirements.txt COPY . . +# Set environment variables +ENV PYTHONPATH=/app +ENV ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY} +ENV ANTHROPIC_MODEL_NAME=${ANTHROPIC_MODEL_NAME:-claude-3-sonnet-20240229} + EXPOSE 8081 +EXPOSE 3001 +# Run only the main API - MCP server will be started programmatically CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8081"] \ No newline at end of file diff --git a/cortex_on/agents/code_agent.py b/cortex_on/agents/code_agent.py index 43fa047..a86369f 100644 --- a/cortex_on/agents/code_agent.py +++ b/cortex_on/agents/code_agent.py @@ -4,7 +4,8 @@ import shlex import subprocess from dataclasses import asdict, dataclass -from typing import Any, Callable, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional, Tuple +import uuid # Third-party imports from dotenv import load_dotenv @@ -13,10 +14,13 @@ from pydantic import BaseModel, Field from pydantic_ai import Agent, RunContext from pydantic_ai.models.anthropic import AnthropicModel +from pydantic_ai.providers.anthropic import AnthropicProvider # Local application imports from utils.ant_client import get_client from utils.stream_response_format import StreamResponse +from utils.code_formatter import format_execution_result +from utils.docker_executor import run_code load_dotenv() @@ -25,71 +29,20 @@ class CoderAgentDeps: websocket: Optional[WebSocket] = None stream_output: Optional[StreamResponse] = None - -# Constants -ALLOWED_COMMANDS = { - "ls", "dir", "cat", "echo", "python", "pip", - "mkdir", "touch", "rm", "cp", "mv" -} - -# Message templates - Replace elif ladders with lookup dictionaries -OPERATION_MESSAGES = { - "ls": lambda cmd, args: "Listing files in directory", - "dir": lambda cmd, args: "Listing files in directory", - "cat": lambda cmd, args: ( - f"Creating file {cmd.split('>', 1)[1].strip().split(' ', 1)[0]}" - if "<<" in cmd and ">" in cmd - else f"Reading file {args[1] if len(args) > 1 else 'file'}" - ), - "echo": lambda cmd, args: f"Creating file {cmd.split('>', 1)[1].strip()}" if ">" in cmd else "Echo command", - "python": lambda cmd, args: f"Running Python script {args[1] if len(args) > 1 else 'script'}", - "pip": lambda cmd, args: ( - f"Installing package(s): {cmd.split('install ', 1)[1]}" - if "install " in cmd - else "Managing Python packages" - ), - "mkdir": lambda cmd, args: f"Creating directory {args[1] if len(args) > 1 else 'directory'}", - "touch": lambda cmd, args: f"Creating empty file {args[1] if len(args) > 1 else 'file'}", - "rm": lambda cmd, args: f"Removing {args[1] if len(args) > 1 else 'file'}", - "cp": lambda cmd, args: f"Copying {args[1]} to {args[2]}" if len(args) >= 3 else "Copying file", - "mv": lambda cmd, args: f"Moving {args[1]} to {args[2]}" if len(args) >= 3 else "Moving file", -} - -EXECUTION_MESSAGES = { - "python": lambda cmd, args: f"Executing Python script {args[1] if len(args) > 1 else 'script'}", - "default": lambda cmd, args: "Executing operation" + session_id: Optional[str] = None # Add session_id for persistent Docker environment + +# Language-specific file extensions +LANGUAGE_EXTENSIONS = { + "python": ".py", + "java": ".java", + "cpp": ".cpp", + "javascript": ".js", + "typescript": ".ts", + "ruby": ".rb", + "go": ".go", + "rust": ".rs", + "php": ".php" } - -SUCCESS_MESSAGES = { - "ls": "Files listed successfully", - "dir": "Files listed successfully", - "cat": lambda cmd: "File created successfully" if "<<" in cmd else "File read successfully", - "echo": "File created successfully", - "python": "Python script executed successfully", - "pip": lambda cmd: "Package installation completed" if "install" in cmd else "Package operation completed", - "mkdir": "Directory created successfully", - "touch": "File created successfully", - "rm": "File removed successfully", - "cp": "File copied successfully", - "mv": "File moved successfully", - "default": "Operation completed successfully" -} - -FAILURE_MESSAGES = { - "ls": "Failed to list files", - "dir": "Failed to list files", - "cat": lambda cmd: "Failed to create file" if "<<" in cmd else "Failed to read file", - "echo": "Failed to create file", - "python": "Python script execution failed", - "pip": lambda cmd: "Package installation failed" if "install" in cmd else "Package operation failed", - "mkdir": "Failed to create directory", - "touch": "Failed to create file", - "rm": "Failed to remove file", - "cp": "Failed to copy file", - "mv": "Failed to move file", - "default": "Operation failed" -} - class CoderResult(BaseModel): dependencies: List = Field( description="All the packages name that has to be installed before the code execution" @@ -97,63 +50,81 @@ class CoderResult(BaseModel): content: str = Field(description="Response content in the form of code") code_description: str = Field(description="Description of the code") -coder_system_message = """You are a helpful AI assistant with coding capabilities. Solve tasks using your coding and language skills. +coder_system_message = """You are a helpful AI assistant with advanced coding capabilities. Solve tasks using your coding and language skills. - - You have access to a single shell tool that executes terminal commands and handles file operations. - - All commands will be executed in a restricted directory for security. - - Do NOT write code that attempts to access directories outside your working directory. - - Do NOT provide test run snippets that print unnecessary output. + - You have access to a secure Docker-based code execution system that runs your code in isolated containers. + - Each programming language has its own dedicated persistent container. + - All code executes in a secure, isolated environment with limited resources. - Never use interactive input functions like 'input()' in Python or 'read' in Bash. - All code must be non-interactive and should execute completely without user interaction. - Use command line arguments, environment variables, or file I/O instead of interactive input. -(restricted to your working directory which means you are already in the ./code_files directory) -When solving tasks, use your provided shell tool for all operations: +You have access to the following tools for code execution and file management: -- execute_shell(command: str) - Execute terminal commands including: - - File operations: Use 'cat' to read files, 'echo' with redirection (>) to write files - - Directory operations: 'ls', 'mkdir', etc. - - Code execution: 'python' for running Python scripts - - Package management: 'pip install' for dependencies +1. execute_code(language: str, code: str) - Execute code directly in the appropriate language container + - The code is saved to a file named program. and executed + - Supported languages: python, java, cpp, javascript, typescript, ruby, go, rust, php + - Resources: 0.5 CPU core, 512MB RAM, 30 second timeout -Allowed commands for execute_shell tool are as follows : ls, dir, cat, echo, python, pip, mkdir, touch, rm, cp, mv +2. create_file(filename: str, content: str, language: str = None) - Create a new file in the container + - Filename should include appropriate extension (e.g., 'utils.py', 'data.json') + - Language is optional and will be detected from the file extension -For Python code, don't use python3, just use python for execution. +3. read_file(filename: str) - Read the content of an existing file in the container + - Returns the content of the specified file -Follow this workflow: -1. First, explain your plan and approach to solving the task. -2. Use shell commands to gather information when needed (e.g., 'cat file.py', 'ls'). -3. Write code to files using echo with redirection (e.g., 'echo "print('hello')" > script.py'). - - For multi-line files, use the here-document syntax with 'cat' (e.g., 'cat > file.py << 'EOF'\\ncode\\nEOF'). -4. Execute the code using 'python script.py'. -5. After each execution, verify the results and fix any errors. -6. Continue this process until the task is complete. +4. list_files() - List all files currently in the container + - Shows what files you've created and can access -Code guidelines: -- Always specify the script type in code blocks (e.g., ```python, ```sh) -- For files that need to be saved, include "# filename: " as the first line -- Provide complete, executable code that doesn't require user modification -- Include only one code block per response -- Use print statements appropriately for output, not for debugging +5. execute_file(filename: str, language: str = None) - Execute a specific file in the container + - Use this to run files you've previously created + - Language is optional and will be detected from the file extension + +Each language container persists during your session, so you can: +- Create multiple files that work together +- Build more complex applications with separate modules +- Execute different files as needed +- Modify files based on execution results + +Follow this workflow for efficient coding: +1. Break down complex problems into manageable components +2. Create separate files for different modules when appropriate +3. Execute code to test and verify your implementation +4. Organize your code according to best practices for the language + +Supported programming languages: + +1. Python (.py) - Python 3.11 with numpy, pandas, matplotlib +2. Java (.java) - OpenJDK 17 +3. C++ (.cpp) - GCC 11 +4. JavaScript (.js) - Node.js 18 with axios +5. TypeScript (.ts) - Node.js 18 with typescript +6. Ruby (.rb) - Ruby 3.2 +7. Go (.go) - Go 1.20 +8. Rust (.rs) - Rust 1.70 +9. PHP (.php) - PHP 8.2 -Self-verification: -- After executing code, analyze the output to verify correctness -- If errors occur, fix them and try again with improved code -- If your approach isn't working after multiple attempts, reconsider your strategy +Code guidelines: +- Provide clean, well-structured code that follows language conventions +- Include appropriate error handling +- Use clear naming conventions and add comments for complex logic +- Structure multi-file projects appropriately based on language best practices +- For languages that require the filenames same as the class names, make sure to create the files with the same name as the class name. + +Example multi-file workflow: +1. Create a main file with core functionality +2. Create utility files for helper functions +3. Import/include utilities in the main file +4. Execute the main file to run the complete application Output explanation guidelines: - After code execution, structure your explanation according to the CoderResult format - For each code solution, explain: 1. Dependencies: List all packages that must be installed before executing the code - 2. Content: The actual code that solves the problem - 3. Code description: A clear explanation of how the code works, its approach, and key components - -When presenting results, format your explanation to match the CoderResult class structure: -- First list dependencies (even if empty) -- Then provide the complete code content -- Finally, include a detailed description of the code's functionality and implementation details + 2. Content: The actual code across all files you created + 3. Code description: A clear explanation of how the code works, its approach, and file relationships Example structure: Dependencies: @@ -161,73 +132,15 @@ class CoderResult(BaseModel): - pandas Content: -[The complete code solution] +[The complete code solution, with file relationships explained] Code Description: -This solution implements [approach] to solve [problem]. The code first [key step 1], -then [key step 2], and finally [produces result]. The implementation handles [edge cases] -by [specific technique]. Key functions include [function 1] which [purpose], -and [function 2] which [purpose]. +This solution implements [approach] to solve [problem] using [N] files: +- main.py: Handles the core functionality, including [key components] +- utils.py: Contains helper functions for [specific tasks] +The implementation handles [edge cases] by [specific technique]. """ -# Helper functions -def get_message_from_dict( - message_dict: Dict[str, Any], - command: str, - base_command: str -) -> str: - """Get the appropriate message from a dictionary based on the command.""" - args = command.split() - - if base_command in message_dict: - msg_source = message_dict[base_command] - if callable(msg_source): - return msg_source(command, args) - return msg_source - - # Use default message if available, otherwise a generic one - if "default" in message_dict: - default_source = message_dict["default"] - if callable(default_source): - return default_source(command, args) - return default_source - - return f"Operation: {base_command}" - -def get_high_level_operation_message(command: str, base_command: str) -> str: - """Returns a high-level description of the operation being performed""" - args = command.split() - return OPERATION_MESSAGES.get( - base_command, - lambda cmd, args: f"Executing operation: {base_command}" - )(command, args) - -def get_high_level_execution_message(command: str, base_command: str) -> str: - """Returns a high-level execution message for the command""" - args = command.split() - return EXECUTION_MESSAGES.get( - base_command, - EXECUTION_MESSAGES["default"] - )(command, args) - -def get_success_message(command: str, base_command: str) -> str: - """Returns a success message based on the command type""" - msg_source = SUCCESS_MESSAGES.get(base_command, SUCCESS_MESSAGES["default"]) - - if callable(msg_source): - return msg_source(command) - - return msg_source - -def get_failure_message(command: str, base_command: str) -> str: - """Returns a failure message based on the command type""" - msg_source = FAILURE_MESSAGES.get(base_command, FAILURE_MESSAGES["default"]) - - if callable(msg_source): - return msg_source(command) - - return msg_source - async def send_stream_update(ctx: RunContext[CoderAgentDeps], message: str) -> None: """Helper function to send websocket updates if available""" if ctx.deps.websocket and ctx.deps.stream_output: @@ -237,9 +150,11 @@ async def send_stream_update(ctx: RunContext[CoderAgentDeps], message: str) -> N logfire.debug("WebSocket message sent: {stream_output_json}", stream_output_json=stream_output_json) # Initialize the model +provider = AnthropicProvider(api_key=os.environ.get("ANTHROPIC_API_KEY")) + model = AnthropicModel( model_name=os.environ.get("ANTHROPIC_MODEL_NAME"), - anthropic_client=get_client() + provider = provider ) # Initialize the agent @@ -252,164 +167,499 @@ async def send_stream_update(ctx: RunContext[CoderAgentDeps], message: str) -> N ) @coder_agent.tool -async def execute_shell(ctx: RunContext[CoderAgentDeps], command: str) -> str: +async def execute_code(ctx: RunContext[CoderAgentDeps], language: str, code: str) -> str: """ - Executes a shell command within a restricted directory and returns the output. - This consolidated tool handles terminal commands and file operations. + Executes code in a secure Docker container with resource limits and isolation. + This tool handles various programming languages with appropriate execution environments. + + Args: + language: The programming language of the code (python, java, cpp) + code: The source code to execute + + Returns: + The execution results, including stdout, stderr, and status """ try: - # Extract base command for security checks and messaging - base_command = command.split()[0] if command.split() else "" + # Normalize language name + language = language.lower().strip() + + # Map language aliases to standard names + language_mapping = { + "python3": "python", + "py": "python", + "c++": "cpp", + "node": "javascript", + "nodejs": "javascript", + "js": "javascript", + "rb": "ruby", + "golang": "go", + "rust": "rust", + "php": "php", + "ts": "typescript", + } + + normalized_language = language_mapping.get(language, language) + + # Check if the language is supported + if normalized_language not in ["python", "java", "cpp", "javascript", "ruby", "go", "rust", "php", "typescript"]: + error_msg = f"Unsupported language: {normalized_language}." + await send_stream_update(ctx, error_msg) + return error_msg # Send operation description message - operation_message = get_high_level_operation_message(command, base_command) - await send_stream_update(ctx, operation_message) - - logfire.info("Executing shell command: {command}", command=command) - - # Setup restricted directory - base_dir = os.path.abspath(os.path.dirname(__file__)) - restricted_dir = os.path.join(base_dir, "code_files") - os.makedirs(restricted_dir, exist_ok=True) - - # Security check - if base_command not in ALLOWED_COMMANDS: - await send_stream_update(ctx, "Operation not permitted") - return f"Error: Command '{base_command}' is not allowed for security reasons." - - # Change to restricted directory for execution - original_dir = os.getcwd() - os.chdir(restricted_dir) - - try: - # Handle echo with redirection (file writing) - if ">" in command and base_command == "echo": - file_path = command.split(">", 1)[1].strip() - await send_stream_update(ctx, f"Writing content to {file_path}") - - # Parse command parts - parts = command.split(">", 1) - echo_cmd = parts[0].strip() - - # Extract content, removing quotes if present - content = echo_cmd[5:].strip() - if (content.startswith('"') and content.endswith('"')) or \ - (content.startswith("'") and content.endswith("'")): - content = content[1:-1] - - try: - with open(file_path, "w") as file: - file.write(content) - - await send_stream_update(ctx, f"File {file_path} created successfully") - return f"Successfully wrote to {file_path}" - except Exception as e: - error_msg = f"Error writing to file: {str(e)}" - await send_stream_update(ctx, f"Failed to create file {file_path}") - logfire.error(error_msg, exc_info=True) - return error_msg + await send_stream_update(ctx, f"Executing {normalized_language} code in secure container") + + logfire.info(f"Executing {normalized_language} code in Docker container") + + # Store the source code in the StreamResponse + if ctx.deps.stream_output: + ctx.deps.stream_output.source_code = code + ctx.deps.stream_output.metadata = {"language": normalized_language} + + # Run the code in a Docker container - we don't need session_id anymore with the new language-based approach + result = await run_code(normalized_language, code) + + # If there was an error with the Docker execution itself + if "error" in result: + error_message = result["error"] + await send_stream_update(ctx, f"Code execution failed: {error_message}") + logfire.error(f"Code execution failed: {error_message}") - # Handle cat with here-document for multiline file writing - elif "<<" in command and base_command == "cat": - cmd_parts = command.split("<<", 1) - cat_part = cmd_parts[0].strip() - - # Extract filename for status message if possible - file_path = None - if ">" in cat_part: - file_path = cat_part.split(">", 1)[1].strip() - await send_stream_update(ctx, f"Creating file {file_path}") - - try: - # Parse heredoc parts - doc_part = cmd_parts[1].strip() - - # Extract filename - if ">" in cat_part: - file_path = cat_part.split(">", 1)[1].strip() - else: - await send_stream_update(ctx, "Invalid file operation") - return "Error: Invalid cat command format. Must include redirection." - - # Parse the heredoc content and delimiter - if "\n" in doc_part: - delimiter_and_content = doc_part.split("\n", 1) - delimiter = delimiter_and_content[0].strip("'").strip('"') - content = delimiter_and_content[1] - - # Find the end delimiter and extract content - if f"\n{delimiter}" in content: - content = content.split(f"\n{delimiter}")[0] - - # Write to file - with open(file_path, "w") as file: - file.write(content) - - await send_stream_update(ctx, f"File {file_path} created successfully") - return f"Successfully wrote multiline content to {file_path}" - else: - await send_stream_update(ctx, "File content format error") - return "Error: End delimiter not found in heredoc" - else: - await send_stream_update(ctx, "File content format error") - return "Error: Invalid heredoc format" - except Exception as e: - error_msg = f"Error processing cat with heredoc: {str(e)}" - file_path_str = file_path if file_path else 'file' - await send_stream_update(ctx, f"Failed to create file {file_path_str}") - logfire.error(error_msg, exc_info=True) - return error_msg + # Create a manually crafted formatted output with error + formatted_output = f"```{normalized_language}\n{code}\n```\n\n" + formatted_output += f"## Errors\n\n```\n{error_message}\n```\n\n" + formatted_output += "## Status\n\n**❌ Execution failed**" + + # Update the StreamResponse with both code and formatted error + if ctx.deps.websocket and ctx.deps.stream_output: + ctx.deps.stream_output.output = formatted_output + ctx.deps.stream_output.status_code = 500 + await ctx.deps.websocket.send_text(json.dumps(asdict(ctx.deps.stream_output))) + + return f"Error: {error_message}" + + # Ensure stdout and stderr are strings + if "stdout" not in result or result["stdout"] is None: + result["stdout"] = "" + if "stderr" not in result or result["stderr"] is None: + result["stderr"] = "" - # Execute standard commands + # Format the execution results for console output + output = f"Execution results:\n\n" + + # Add stdout if available + if result.get("stdout"): + output += f"--- Output ---\n{result['stdout']}\n\n" + else: + output += "--- No Output ---\n\n" + + # Add stderr if there were errors + if result.get("stderr"): + output += f"--- Errors ---\n{result['stderr']}\n\n" + + # Add execution status + if result.get("success", False): + await send_stream_update(ctx, f"{normalized_language.capitalize()} code executed successfully") + output += "Status: Success\n" + else: + await send_stream_update(ctx, f"{normalized_language.capitalize()} code execution failed") + output += f"Status: Failed (Exit code: {result.get('exit_code', 'unknown')})\n" + + # Create a manually crafted formatted output for UI display + formatted_output = "" + + # Always add code section first with proper language syntax highlighting + formatted_output += f"## Code\n\n```{normalized_language}\n{code}\n```\n\n" + + # Add execution results section + formatted_output += "## Output\n\n" + if result.get("stdout"): + formatted_output += f"```\n{result['stdout']}\n```\n\n" + else: + formatted_output += "*No output captured*\n\n" + + # Add errors section if needed + if result.get("stderr"): + formatted_output += f"## Errors\n\n```\n{result['stderr']}\n```\n\n" + + # Add status section + if result.get("success", False): + formatted_output += "## Status\n\n**✅ Execution completed successfully**" + else: + exit_code = result.get("exit_code", "unknown") + formatted_output += f"## Status\n\n**❌ Execution failed** (Exit code: {exit_code})" + + # Update the StreamResponse with both code and results + if ctx.deps.websocket and ctx.deps.stream_output: + ctx.deps.stream_output.output = formatted_output + ctx.deps.stream_output.status_code = 200 if result.get("success", False) else 500 + ctx.deps.stream_output.metadata = { + "language": normalized_language, + "success": result.get("success", False), + "exit_code": result.get("exit_code", "unknown") + } + await ctx.deps.websocket.send_text(json.dumps(asdict(ctx.deps.stream_output))) + + logfire.info(f"Code execution completed with status: {result.get('success', False)}") + return output + + except Exception as e: + error_msg = f"Error during code execution: {str(e)}" + await send_stream_update(ctx, "Code execution failed") + logfire.error(error_msg, exc_info=True) + + # Create a manually crafted formatted output with error + formatted_error = f"```{language}\n{code}\n```\n\n" + formatted_error += f"## Errors\n\n```\n{error_msg}\n```\n\n" + formatted_error += "## Status\n\n**❌ Execution failed**" + + # Update the StreamResponse with both code and formatted error + if ctx.deps.websocket and ctx.deps.stream_output: + ctx.deps.stream_output.output = formatted_error + ctx.deps.stream_output.status_code = 500 + await ctx.deps.websocket.send_text(json.dumps(asdict(ctx.deps.stream_output))) + + return error_msg + +@coder_agent.tool +async def create_file(ctx: RunContext[CoderAgentDeps], filename: str, content: str, language: str = None) -> str: + """ + Creates a file in the persistent Docker environment. + + Args: + filename: Name of the file to create + content: Content to write to the file + language: Optional programming language for syntax highlighting + + Returns: + Result of the file creation operation + """ + try: + # Detect language from filename extension if not provided + if not language and "." in filename: + ext = os.path.splitext(filename)[1].lower() + language_map = {v: k for k, v in LANGUAGE_EXTENSIONS.items()} + language = language_map.get(ext, None) + + # Send operation description message + await send_stream_update(ctx, f"Creating file: {filename}") + + logfire.info(f"Creating file {filename} in Docker environment") + + # Get Docker environment + from utils.docker_executor import get_environment + env = get_environment(language or "python") + + # Connect to Docker environment + connect_result = await env.connect() + if not connect_result.get("success", False): + error_message = connect_result.get("error", "Unable to connect to Docker environment") + await send_stream_update(ctx, f"Failed to connect to environment: {error_message}") + return f"Error: {error_message}" + + # Write file to Docker environment + result = await env.write_file(filename, content) + + if result.get("success", False): + message = f"File {filename} created successfully" + await send_stream_update(ctx, message) + + # Format output for frontend display + formatted_output = f"## File Creation\n\n**{filename}** has been created successfully.\n\n" + if language: + formatted_output += f"```{language}\n{content}\n```" else: - # Send execution message - execution_msg = get_high_level_execution_message(command, base_command) - await send_stream_update(ctx, execution_msg) - - # Execute the command using subprocess - try: - args = shlex.split(command) - result = subprocess.run( - args, - shell=True, - capture_output=True, - text=True, - timeout=60, - ) - - logfire.info(f"Command executed: {result.args}") - - # Handle success - if result.returncode == 0: - success_msg = get_success_message(command, base_command) - await send_stream_update(ctx, success_msg) - logfire.info(f"Command executed successfully: {result.stdout}") - return result.stdout - - # Handle failure - else: - files = os.listdir('.') - error_msg = f"Command failed with error code {result.returncode}:\n{result.stderr}\n\nFiles in directory: {files}" - failure_msg = get_failure_message(command, base_command) - await send_stream_update(ctx, failure_msg) - return error_msg - - except subprocess.TimeoutExpired: - await send_stream_update(ctx, "Operation timed out") - return "Command execution timed out after 60 seconds" + formatted_output += f"```\n{content}\n```" + + # Update StreamResponse with formatted result + if ctx.deps.websocket and ctx.deps.stream_output: + ctx.deps.stream_output.output = formatted_output + ctx.deps.stream_output.status_code = 200 + await ctx.deps.websocket.send_text(json.dumps(asdict(ctx.deps.stream_output))) + + return message + else: + error_message = result.get("error", "Unknown error") + await send_stream_update(ctx, f"Failed to create file: {error_message}") + + # Update StreamResponse with error + if ctx.deps.websocket and ctx.deps.stream_output: + ctx.deps.stream_output.output = f"## Error\n\nFailed to create file {filename}: {error_message}" + ctx.deps.stream_output.status_code = 500 + await ctx.deps.websocket.send_text(json.dumps(asdict(ctx.deps.stream_output))) + + return f"Error: {error_message}" + + except Exception as e: + error_msg = f"Error creating file {filename}: {str(e)}" + await send_stream_update(ctx, f"File creation failed: {str(e)}") + logfire.error(error_msg, exc_info=True) + return error_msg + +@coder_agent.tool +async def read_file(ctx: RunContext[CoderAgentDeps], filename: str) -> str: + """ + Reads a file from the persistent Docker environment. + + Args: + filename: Name of the file to read + + Returns: + Content of the file or error message + """ + try: + # Detect language from filename extension for environment selection + language = "python" # Default + if "." in filename: + ext = os.path.splitext(filename)[1].lower() + language_map = {v: k for k, v in LANGUAGE_EXTENSIONS.items()} + detected_lang = language_map.get(ext, None) + if detected_lang: + language = detected_lang + + # Send operation description message + await send_stream_update(ctx, f"Reading file: {filename}") + + logfire.info(f"Reading file {filename} from Docker environment") + + # Get Docker environment + from utils.docker_executor import get_environment + env = get_environment(language) + + # Connect to Docker environment + connect_result = await env.connect() + if not connect_result.get("success", False): + error_message = connect_result.get("error", "Unable to connect to Docker environment") + await send_stream_update(ctx, f"Failed to connect to environment: {error_message}") + return f"Error: {error_message}" + + # Read file from Docker environment + result = await env.read_file(filename) + + if result.get("success", False): + content = result.get("content", "") + await send_stream_update(ctx, f"File {filename} read successfully") + + # Format output for frontend display + formatted_output = f"## File: {filename}\n\n" + if language: + formatted_output += f"```{language}\n{content}\n```" + else: + formatted_output += f"```\n{content}\n```" + + # Update StreamResponse with formatted result + if ctx.deps.websocket and ctx.deps.stream_output: + ctx.deps.stream_output.output = formatted_output + ctx.deps.stream_output.status_code = 200 + await ctx.deps.websocket.send_text(json.dumps(asdict(ctx.deps.stream_output))) + + return content + else: + error_message = result.get("error", "Unknown error") + await send_stream_update(ctx, f"Failed to read file: {error_message}") + + # Update StreamResponse with error + if ctx.deps.websocket and ctx.deps.stream_output: + ctx.deps.stream_output.output = f"## Error\n\nFailed to read file {filename}: {error_message}" + ctx.deps.stream_output.status_code = 500 + await ctx.deps.websocket.send_text(json.dumps(asdict(ctx.deps.stream_output))) + + return f"Error: {error_message}" + + except Exception as e: + error_msg = f"Error reading file {filename}: {str(e)}" + await send_stream_update(ctx, f"File reading failed: {str(e)}") + logfire.error(error_msg, exc_info=True) + return error_msg + +@coder_agent.tool +async def list_files(ctx: RunContext[CoderAgentDeps]) -> str: + """ + Lists all files in the persistent Docker environment. + + Returns: + List of files or error message + """ + try: + # Send operation description message + await send_stream_update(ctx, "Listing files in environment") + + logfire.info("Listing files in Docker environment") + + # Get Docker environment (use python as default) + from utils.docker_executor import get_environment + env = get_environment("python") + + # Connect to Docker environment + connect_result = await env.connect() + if not connect_result.get("success", False): + error_message = connect_result.get("error", "Unable to connect to Docker environment") + await send_stream_update(ctx, f"Failed to connect to environment: {error_message}") + return f"Error: {error_message}" + + # List files in Docker environment + result = await env.list_files() + + if result.get("success", False): + files = result.get("files", []) + await send_stream_update(ctx, f"Found {len(files)} files") + + # Format output for frontend display + if files: + formatted_output = "## Files in Environment\n\n" + formatted_output += "| Filename |\n|----------|\n" + for filename in files: + formatted_output += f"| `{filename}` |\n" + else: + formatted_output = "## Files in Environment\n\nNo files found." + + # Update StreamResponse with formatted result + if ctx.deps.websocket and ctx.deps.stream_output: + ctx.deps.stream_output.output = formatted_output + ctx.deps.stream_output.status_code = 200 + await ctx.deps.websocket.send_text(json.dumps(asdict(ctx.deps.stream_output))) + + if files: + return f"Files in environment: {', '.join(files)}" + else: + return "No files found in environment." + else: + error_message = result.get("error", "Unknown error") + await send_stream_update(ctx, f"Failed to list files: {error_message}") + + # Update StreamResponse with error + if ctx.deps.websocket and ctx.deps.stream_output: + ctx.deps.stream_output.output = f"## Error\n\nFailed to list files: {error_message}" + ctx.deps.stream_output.status_code = 500 + await ctx.deps.websocket.send_text(json.dumps(asdict(ctx.deps.stream_output))) + + return f"Error: {error_message}" - except Exception as e: - error_msg = f"Error executing command: {str(e)}" - await send_stream_update(ctx, "Operation failed") - logfire.error(error_msg, exc_info=True) - return error_msg - - finally: - # Always return to the original directory - os.chdir(original_dir) + except Exception as e: + error_msg = f"Error listing files: {str(e)}" + await send_stream_update(ctx, f"File listing failed: {str(e)}") + logfire.error(error_msg, exc_info=True) + return error_msg + +@coder_agent.tool +async def execute_file(ctx: RunContext[CoderAgentDeps], filename: str, language: str = None) -> str: + """ + Executes a file in the persistent Docker environment. + + Args: + filename: Name of the file to execute + language: Optional programming language (detected from extension if not specified) + + Returns: + Execution results including stdout, stderr, and status + """ + try: + # Detect language from filename extension if not provided + if not language and "." in filename: + ext = os.path.splitext(filename)[1].lower() + language_map = {v: k for k, v in LANGUAGE_EXTENSIONS.items()} + language = language_map.get(ext, None) + + if not language: + return "Error: Could not determine language for execution. Please specify language parameter." + + # Send operation description message + await send_stream_update(ctx, f"Executing file: {filename}") + + logfire.info(f"Executing file {filename} in Docker environment with language {language}") + + # Get Docker environment for the specific language + from utils.docker_executor import get_environment + env = get_environment(language) + + # Connect to Docker environment + connect_result = await env.connect() + if not connect_result.get("success", False): + error_message = connect_result.get("error", "Unable to connect to Docker environment") + await send_stream_update(ctx, f"Failed to connect to environment: {error_message}") + return f"Error: {error_message}" + + # Read file content for display before execution + file_content = "" + file_result = await env.read_file(filename) + if file_result.get("success", False): + file_content = file_result.get("content", "") + logfire.debug(f"File content to execute: {file_content}") + + # Execute file in Docker environment + result = await env.execute_code(filename) + + # Ensure stdout and stderr are strings + if "stdout" not in result or result["stdout"] is None: + result["stdout"] = "" + if "stderr" not in result or result["stderr"] is None: + result["stderr"] = "" + + # Format the execution results for console output + output = f"Execution results for {filename}:\n\n" + + # Add stdout if available + if result.get("stdout"): + output += f"--- Output ---\n{result['stdout']}\n\n" + else: + output += "--- No Output ---\n\n" + + # Add stderr if there were errors + if result.get("stderr"): + output += f"--- Errors ---\n{result['stderr']}\n\n" + + # Add execution status + if result.get("success", False): + await send_stream_update(ctx, f"File {filename} executed successfully") + output += "Status: Success\n" + else: + await send_stream_update(ctx, f"File {filename} execution failed") + output += f"Status: Failed (Exit code: {result.get('exit_code', 'unknown')})\n" + + # Create a manually crafted formatted output for UI display + formatted_output = "" + + # Always add code section first with proper language syntax highlighting + formatted_output += f"## File: {filename}\n\n```{language}\n{file_content}\n```\n\n" + + # Add execution results section + formatted_output += "## Output\n\n" + if result.get("stdout"): + formatted_output += f"```\n{result['stdout']}\n```\n\n" + else: + formatted_output += "*No output captured*\n\n" + + # Add errors section if needed + if result.get("stderr"): + formatted_output += f"## Errors\n\n```\n{result['stderr']}\n```\n\n" + + # Add status section + if result.get("success", False): + formatted_output += "## Status\n\n**✅ Execution completed successfully**" + else: + exit_code = result.get("exit_code", "unknown") + formatted_output += f"## Status\n\n**❌ Execution failed** (Exit code: {exit_code})" + + # Update StreamResponse with our manually crafted format + if ctx.deps.websocket and ctx.deps.stream_output: + ctx.deps.stream_output.output = formatted_output + ctx.deps.stream_output.source_code = file_content + ctx.deps.stream_output.status_code = 200 if result.get("success", False) else 500 + ctx.deps.stream_output.metadata = { + "language": language, + "filename": filename, + "success": result.get("success", False), + "exit_code": result.get("exit_code", "unknown") + } + await ctx.deps.websocket.send_text(json.dumps(asdict(ctx.deps.stream_output))) + + # Log the full output for debugging + logfire.debug(f"Execution output for {filename}: {output}") + + return output except Exception as e: - error_msg = f"Error executing command: {str(e)}" - await send_stream_update(ctx, "Operation failed") + error_msg = f"Error executing file {filename}: {str(e)}" + await send_stream_update(ctx, f"File execution failed: {str(e)}") logfire.error(error_msg, exc_info=True) return error_msg \ No newline at end of file diff --git a/cortex_on/agents/mcp_server.py b/cortex_on/agents/mcp_server.py new file mode 100644 index 0000000..2469b7a --- /dev/null +++ b/cortex_on/agents/mcp_server.py @@ -0,0 +1,30 @@ +from mcp.server.fastmcp import FastMCP +from pydantic_ai import Agent +from pydantic_ai.models.anthropic import AnthropicModel +from pydantic_ai import RunContext +from fastapi import WebSocket +import os +from typing import List, Optional, Dict, Any, Union, Tuple +import json +from dataclasses import asdict +from utils.ant_client import get_client +from utils.stream_response_format import StreamResponse +from agents.planner_agent import planner_agent +from agents.code_agent import coder_agent +from agents.code_agent import coder_agent, CoderAgentDeps +from agents.orchestrator_agent import orchestrator_deps +from agents.web_surfer import WebSurfer +import logfire + +# Initialize the single MCP server +server = FastMCP("CortexON MCP Server", host="0.0.0.0", port=3001) + +# Note: All tools are now dynamically registered in instructor.py +# This avoids the problem of websocket not being available when tools are defined + +def run_server(): + """Run the MCP server""" + server.run(transport="sse") + +if __name__ == "__main__": + run_server() \ No newline at end of file diff --git a/cortex_on/agents/orchestrator_agent.py b/cortex_on/agents/orchestrator_agent.py index 6b001ba..7c570c3 100644 --- a/cortex_on/agents/orchestrator_agent.py +++ b/cortex_on/agents/orchestrator_agent.py @@ -9,12 +9,14 @@ from fastapi import WebSocket from dotenv import load_dotenv from pydantic_ai.models.anthropic import AnthropicModel +from pydantic_ai.providers.anthropic import AnthropicProvider from pydantic_ai import Agent, RunContext -from agents.web_surfer import WebSurfer +from pydantic_ai.mcp import MCPServerHTTP, MCPServerStdio from utils.stream_response_format import StreamResponse from agents.planner_agent import planner_agent, update_todo_status from agents.code_agent import coder_agent, CoderAgentDeps from utils.ant_client import get_client +load_dotenv() @dataclass class orchestrator_deps: @@ -23,6 +25,7 @@ class orchestrator_deps: # Add a collection to track agent-specific streams agent_responses: Optional[List[StreamResponse]] = None + orchestrator_system_prompt = """You are an AI orchestrator that manages a team of agents to solve tasks. You have access to tools for coordinating the agents and managing the task flow. [AGENT CAPABILITIES] @@ -158,194 +161,26 @@ class orchestrator_deps: - Format: "Task description (agent_name)" """ +# Initialize MCP Server +# server = MCPServerStdio('python', ["-m", "agents.mcp_server"]) +server = MCPServerHTTP(url='http://localhost:3001/sse') + +# Initialize Anthropic provider with API key +provider = AnthropicProvider(api_key=os.environ.get("ANTHROPIC_API_KEY")) + model = AnthropicModel( model_name=os.environ.get("ANTHROPIC_MODEL_NAME"), - anthropic_client=get_client() + provider=provider ) orchestrator_agent = Agent( model=model, name="Orchestrator Agent", system_prompt=orchestrator_system_prompt, - deps_type=orchestrator_deps + deps_type=orchestrator_deps, + mcp_servers=[server], ) -@orchestrator_agent.tool -async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: - """Plans the task and assigns it to the appropriate agents""" - try: - logfire.info(f"Planning task: {task}") - - # Create a new StreamResponse for Planner Agent - planner_stream_output = StreamResponse( - agent_name="Planner Agent", - instructions=task, - steps=[], - output="", - status_code=0 - ) - - # Add to orchestrator's response collection if available - if ctx.deps.agent_responses is not None: - ctx.deps.agent_responses.append(planner_stream_output) - - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # Update planner stream - planner_stream_output.steps.append("Planning task...") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # Run planner agent - planner_response = await planner_agent.run(user_prompt=task) - - # Update planner stream with results - plan_text = planner_response.data.plan - planner_stream_output.steps.append("Task planned successfully") - planner_stream_output.output = plan_text - planner_stream_output.status_code = 200 - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # Also update orchestrator stream - ctx.deps.stream_output.steps.append("Task planned successfully") - await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) - - return f"Task planned successfully\nTask: {plan_text}" - except Exception as e: - error_msg = f"Error planning task: {str(e)}" - logfire.error(error_msg, exc_info=True) - - # Update planner stream with error - if planner_stream_output: - planner_stream_output.steps.append(f"Planning failed: {str(e)}") - planner_stream_output.status_code = 500 - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # Also update orchestrator stream - if ctx.deps.stream_output: - ctx.deps.stream_output.steps.append(f"Planning failed: {str(e)}") - await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) - - return f"Failed to plan task: {error_msg}" - -@orchestrator_agent.tool -async def coder_task(ctx: RunContext[orchestrator_deps], task: str) -> str: - """Assigns coding tasks to the coder agent""" - try: - logfire.info(f"Assigning coding task: {task}") - - # Create a new StreamResponse for Coder Agent - coder_stream_output = StreamResponse( - agent_name="Coder Agent", - instructions=task, - steps=[], - output="", - status_code=0 - ) - - # Add to orchestrator's response collection if available - if ctx.deps.agent_responses is not None: - ctx.deps.agent_responses.append(coder_stream_output) - - # Send initial update for Coder Agent - await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) - - # Create deps with the new stream_output - deps_for_coder_agent = CoderAgentDeps( - websocket=ctx.deps.websocket, - stream_output=coder_stream_output - ) - - # Run coder agent - coder_response = await coder_agent.run( - user_prompt=task, - deps=deps_for_coder_agent - ) - - # Extract response data - response_data = coder_response.data.content - - # Update coder_stream_output with coding results - coder_stream_output.output = response_data - coder_stream_output.status_code = 200 - coder_stream_output.steps.append("Coding task completed successfully") - await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) - - # Add a reminder in the result message to update the plan using planner_agent_update - response_with_reminder = f"{response_data}\n\nReminder: You must now call planner_agent_update with the completed task description: \"{task} (coder_agent)\"" - - return response_with_reminder - except Exception as e: - error_msg = f"Error assigning coding task: {str(e)}" - logfire.error(error_msg, exc_info=True) - - # Update coder_stream_output with error - coder_stream_output.steps.append(f"Coding task failed: {str(e)}") - coder_stream_output.status_code = 500 - await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) - - return f"Failed to assign coding task: {error_msg}" - -@orchestrator_agent.tool -async def web_surfer_task(ctx: RunContext[orchestrator_deps], task: str) -> str: - """Assigns web surfing tasks to the web surfer agent""" - try: - logfire.info(f"Assigning web surfing task: {task}") - - # Create a new StreamResponse for WebSurfer - web_surfer_stream_output = StreamResponse( - agent_name="Web Surfer", - instructions=task, - steps=[], - output="", - status_code=0, - live_url=None - ) - - # Add to orchestrator's response collection if available - if ctx.deps.agent_responses is not None: - ctx.deps.agent_responses.append(web_surfer_stream_output) - - await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) - - # Initialize WebSurfer agent - web_surfer_agent = WebSurfer(api_url="http://localhost:8000/api/v1/web/stream") - - # Run WebSurfer with its own stream_output - success, message, messages = await web_surfer_agent.generate_reply( - instruction=task, - websocket=ctx.deps.websocket, - stream_output=web_surfer_stream_output - ) - - # Update WebSurfer's stream_output with final result - if success: - web_surfer_stream_output.steps.append("Web search completed successfully") - web_surfer_stream_output.output = message - web_surfer_stream_output.status_code = 200 - - # Add a reminder to update the plan - message_with_reminder = f"{message}\n\nReminder: You must now call planner_agent_update with the completed task description: \"{task} (web_surfer_agent)\"" - else: - web_surfer_stream_output.steps.append(f"Web search completed with issues: {message[:100]}") - web_surfer_stream_output.status_code = 500 - message_with_reminder = message - - await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) - - web_surfer_stream_output.steps.append(f"WebSurfer completed: {'Success' if success else 'Failed'}") - await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) - - return message_with_reminder - except Exception as e: - error_msg = f"Error assigning web surfing task: {str(e)}" - logfire.error(error_msg, exc_info=True) - - # Update WebSurfer's stream_output with error - web_surfer_stream_output.steps.append(f"Web search failed: {str(e)}") - web_surfer_stream_output.status_code = 500 - await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) - return f"Failed to assign web surfing task: {error_msg}" - @orchestrator_agent.tool async def ask_human(ctx: RunContext[orchestrator_deps], question: str) -> str: """Sends a question to the frontend and waits for human input""" @@ -393,122 +228,346 @@ async def ask_human(ctx: RunContext[orchestrator_deps], question: str) -> str: return f"Failed to get human input: {error_msg}" -@orchestrator_agent.tool -async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_task: str) -> str: - """ - Updates the todo.md file to mark a task as completed and returns the full updated plan. +async def _safe_websocket_send(socket: WebSocket, message: Any) -> bool: + """Safely send message through websocket with error handling""" + try: + if socket and socket.client_state.CONNECTED: + await socket.send_text(json.dumps(asdict(message))) + logfire.debug("WebSocket message sent (_safe_websocket_send): {message}", message=message) + return True + return False + except Exception as e: + logfire.error(f"WebSocket send failed: {str(e)}") + return False +# @orchestrator_agent.tool +# async def plan_task(ctx: RunContext[orchestrator_deps], task: str) -> str: +# """Plans the task and assigns it to the appropriate agents""" +# try: +# logfire.info(f"Planning task: {task}") + +# # Create a new StreamResponse for Planner Agent +# planner_stream_output = StreamResponse( +# agent_name="Planner Agent", +# instructions=task, +# steps=[], +# output="", +# status_code=0 +# ) + +# # Add to orchestrator's response collection if available +# if ctx.deps.agent_responses is not None: +# ctx.deps.agent_responses.append(planner_stream_output) + +# await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + +# # Update planner stream +# planner_stream_output.steps.append("Planning task...") +# await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + +# # Run planner agent +# planner_response = await planner_agent.run(user_prompt=task) + +# # Update planner stream with results +# plan_text = planner_response.data.plan +# planner_stream_output.steps.append("Task planned successfully") +# planner_stream_output.output = plan_text +# planner_stream_output.status_code = 200 +# await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + +# # Also update orchestrator stream +# ctx.deps.stream_output.steps.append("Task planned successfully") +# await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) + +# return f"Task planned successfully\nTask: {plan_text}" +# except Exception as e: +# error_msg = f"Error planning task: {str(e)}" +# logfire.error(error_msg, exc_info=True) + +# # Update planner stream with error +# if planner_stream_output: +# planner_stream_output.steps.append(f"Planning failed: {str(e)}") +# planner_stream_output.status_code = 500 +# await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + +# # Also update orchestrator stream +# if ctx.deps.stream_output: +# ctx.deps.stream_output.steps.append(f"Planning failed: {str(e)}") +# await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) + +# return f"Failed to plan task: {error_msg}" + +# @orchestrator_agent.tool +# async def coder_task(ctx: RunContext[orchestrator_deps], task: str) -> str: +# """Assigns coding tasks to the coder agent""" +# try: +# logfire.info(f"Assigning coding task: {task}") + +# # Create a new StreamResponse for Coder Agent +# coder_stream_output = StreamResponse( +# agent_name="Coder Agent", +# instructions=task, +# steps=[], +# output="", +# status_code=0 +# ) + +# # Add to orchestrator's response collection if available +# if ctx.deps.agent_responses is not None: +# ctx.deps.agent_responses.append(coder_stream_output) + +# # Send initial update for Coder Agent +# await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) + +# # Create deps with the new stream_output +# deps_for_coder_agent = CoderAgentDeps( +# websocket=ctx.deps.websocket, +# stream_output=coder_stream_output +# ) + +# # Run coder agent +# coder_response = await coder_agent.run( +# user_prompt=task, +# deps=deps_for_coder_agent +# ) + +# # Extract response data +# response_data = coder_response.data.content + +# # Update coder_stream_output with coding results +# coder_stream_output.output = response_data +# coder_stream_output.status_code = 200 +# coder_stream_output.steps.append("Coding task completed successfully") +# await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) + +# # Add a reminder in the result message to update the plan using planner_agent_update +# response_with_reminder = f"{response_data}\n\nReminder: You must now call planner_agent_update with the completed task description: \"{task} (coder_agent)\"" + +# return response_with_reminder +# except Exception as e: +# error_msg = f"Error assigning coding task: {str(e)}" +# logfire.error(error_msg, exc_info=True) + +# # Update coder_stream_output with error +# coder_stream_output.steps.append(f"Coding task failed: {str(e)}") +# coder_stream_output.status_code = 500 +# await _safe_websocket_send(ctx.deps.websocket, coder_stream_output) + +# return f"Failed to assign coding task: {error_msg}" + +# @orchestrator_agent.tool +# async def web_surfer_task(ctx: RunContext[orchestrator_deps], task: str) -> str: +# """Assigns web surfing tasks to the web surfer agent""" +# try: +# logfire.info(f"Assigning web surfing task: {task}") + +# # Create a new StreamResponse for WebSurfer +# web_surfer_stream_output = StreamResponse( +# agent_name="Web Surfer", +# instructions=task, +# steps=[], +# output="", +# status_code=0, +# live_url=None +# ) + +# # Add to orchestrator's response collection if available +# if ctx.deps.agent_responses is not None: +# ctx.deps.agent_responses.append(web_surfer_stream_output) + +# await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) + +# # Initialize WebSurfer agent +# web_surfer_agent = WebSurfer(api_url="http://localhost:8000/api/v1/web/stream") + +# # Run WebSurfer with its own stream_output +# success, message, messages = await web_surfer_agent.generate_reply( +# instruction=task, +# websocket=ctx.deps.websocket, +# stream_output=web_surfer_stream_output +# ) + +# # Update WebSurfer's stream_output with final result +# if success: +# web_surfer_stream_output.steps.append("Web search completed successfully") +# web_surfer_stream_output.output = message +# web_surfer_stream_output.status_code = 200 + +# # Add a reminder to update the plan +# message_with_reminder = f"{message}\n\nReminder: You must now call planner_agent_update with the completed task description: \"{task} (web_surfer_agent)\"" +# else: +# web_surfer_stream_output.steps.append(f"Web search completed with issues: {message[:100]}") +# web_surfer_stream_output.status_code = 500 +# message_with_reminder = message + +# await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) + +# web_surfer_stream_output.steps.append(f"WebSurfer completed: {'Success' if success else 'Failed'}") +# await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) + +# return message_with_reminder +# except Exception as e: +# error_msg = f"Error assigning web surfing task: {str(e)}" +# logfire.error(error_msg, exc_info=True) + +# # Update WebSurfer's stream_output with error +# web_surfer_stream_output.steps.append(f"Web search failed: {str(e)}") +# web_surfer_stream_output.status_code = 500 +# await _safe_websocket_send(ctx.deps.websocket, web_surfer_stream_output) +# return f"Failed to assign web surfing task: {error_msg}" + +# @orchestrator_agent.tool +# async def ask_human(ctx: RunContext[orchestrator_deps], question: str) -> str: +# """Sends a question to the frontend and waits for human input""" +# try: +# logfire.info(f"Asking human: {question}") + +# # Create a new StreamResponse for Human Input +# human_stream_output = StreamResponse( +# agent_name="Human Input", +# instructions=question, +# steps=[], +# output="", +# status_code=0 +# ) + +# # Add to orchestrator's response collection if available +# if ctx.deps.agent_responses is not None: +# ctx.deps.agent_responses.append(human_stream_output) + +# # Send the question to frontend +# await _safe_websocket_send(ctx.deps.websocket, human_stream_output) + +# # Update stream with waiting message +# human_stream_output.steps.append("Waiting for human input...") +# await _safe_websocket_send(ctx.deps.websocket, human_stream_output) + +# # Wait for response from frontend +# response = await ctx.deps.websocket.receive_text() + +# # Update stream with response +# human_stream_output.steps.append("Received human input") +# human_stream_output.output = response +# human_stream_output.status_code = 200 +# await _safe_websocket_send(ctx.deps.websocket, human_stream_output) + +# return response +# except Exception as e: +# error_msg = f"Error getting human input: {str(e)}" +# logfire.error(error_msg, exc_info=True) + +# # Update stream with error +# human_stream_output.steps.append(f"Failed to get human input: {str(e)}") +# human_stream_output.status_code = 500 +# await _safe_websocket_send(ctx.deps.websocket, human_stream_output) + +# return f"Failed to get human input: {error_msg}" + +# @orchestrator_agent.tool +# async def planner_agent_update(ctx: RunContext[orchestrator_deps], completed_task: str) -> str: +# """ +# Updates the todo.md file to mark a task as completed and returns the full updated plan. - Args: - completed_task: Description of the completed task including which agent performed it +# Args: +# completed_task: Description of the completed task including which agent performed it - Returns: - The complete updated todo.md content with tasks marked as completed - """ - try: - logfire.info(f"Updating plan with completed task: {completed_task}") - - # Create a new StreamResponse for Planner Agent update - planner_stream_output = StreamResponse( - agent_name="Planner Agent", - instructions=f"Update todo.md to mark as completed: {completed_task}", - steps=[], - output="", - status_code=0 - ) - - # Send initial update - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # Directly read and update the todo.md file - base_dir = os.path.abspath(os.path.dirname(os.path.dirname(__file__))) - planner_dir = os.path.join(base_dir, "agents", "planner") - todo_path = os.path.join(planner_dir, "todo.md") - - planner_stream_output.steps.append("Reading current todo.md...") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - - # Make sure the directory exists - os.makedirs(planner_dir, exist_ok=True) - - try: - # Check if todo.md exists - if not os.path.exists(todo_path): - planner_stream_output.steps.append("No todo.md file found. Will create new one after task completion.") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) +# Returns: +# The complete updated todo.md content with tasks marked as completed +# """ +# try: +# logfire.info(f"Updating plan with completed task: {completed_task}") + +# # Create a new StreamResponse for Planner Agent update +# planner_stream_output = StreamResponse( +# agent_name="Planner Agent", +# instructions=f"Update todo.md to mark as completed: {completed_task}", +# steps=[], +# output="", +# status_code=0 +# ) + +# # Send initial update +# await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + +# # Directly read and update the todo.md file +# base_dir = os.path.abspath(os.path.dirname(os.path.dirname(__file__))) +# planner_dir = os.path.join(base_dir, "agents", "planner") +# todo_path = os.path.join(planner_dir, "todo.md") + +# planner_stream_output.steps.append("Reading current todo.md...") +# await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) + +# # Make sure the directory exists +# os.makedirs(planner_dir, exist_ok=True) + +# try: +# # Check if todo.md exists +# if not os.path.exists(todo_path): +# planner_stream_output.steps.append("No todo.md file found. Will create new one after task completion.") +# await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - # We'll directly call planner_agent.run() to create a new plan first - plan_prompt = f"Create a simple task plan based on this completed task: {completed_task}" - plan_response = await planner_agent.run(user_prompt=plan_prompt) - current_content = plan_response.data.plan - else: - # Read existing todo.md - with open(todo_path, "r") as file: - current_content = file.read() - planner_stream_output.steps.append(f"Found existing todo.md ({len(current_content)} bytes)") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) +# # We'll directly call planner_agent.run() to create a new plan first +# plan_prompt = f"Create a simple task plan based on this completed task: {completed_task}" +# plan_response = await planner_agent.run(user_prompt=plan_prompt) +# current_content = plan_response.data.plan +# else: +# # Read existing todo.md +# with open(todo_path, "r") as file: +# current_content = file.read() +# planner_stream_output.steps.append(f"Found existing todo.md ({len(current_content)} bytes)") +# await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - # Now call planner_agent.run() with specific instructions to update the plan - update_prompt = f""" - Here is the current todo.md content: +# # Now call planner_agent.run() with specific instructions to update the plan +# update_prompt = f""" +# Here is the current todo.md content: - {current_content} +# {current_content} - Please update this plan to mark the following task as completed: {completed_task} - Return ONLY the fully updated plan with appropriate tasks marked as [x] instead of [ ]. - """ +# Please update this plan to mark the following task as completed: {completed_task} +# Return ONLY the fully updated plan with appropriate tasks marked as [x] instead of [ ]. +# """ - planner_stream_output.steps.append("Asking planner to update the plan...") - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) +# planner_stream_output.steps.append("Asking planner to update the plan...") +# await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - updated_plan_response = await planner_agent.run(user_prompt=update_prompt) - updated_plan = updated_plan_response.data.plan +# updated_plan_response = await planner_agent.run(user_prompt=update_prompt) +# updated_plan = updated_plan_response.data.plan - # Write the updated plan back to todo.md - with open(todo_path, "w") as file: - file.write(updated_plan) +# # Write the updated plan back to todo.md +# with open(todo_path, "w") as file: +# file.write(updated_plan) - planner_stream_output.steps.append("Plan updated successfully") - planner_stream_output.output = updated_plan - planner_stream_output.status_code = 200 - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) +# planner_stream_output.steps.append("Plan updated successfully") +# planner_stream_output.output = updated_plan +# planner_stream_output.status_code = 200 +# await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - # Update orchestrator stream - if ctx.deps.stream_output: - ctx.deps.stream_output.steps.append(f"Plan updated to mark task as completed: {completed_task}") - await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) +# # Update orchestrator stream +# if ctx.deps.stream_output: +# ctx.deps.stream_output.steps.append(f"Plan updated to mark task as completed: {completed_task}") +# await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) - return updated_plan +# return updated_plan - except Exception as e: - error_msg = f"Error during plan update operations: {str(e)}" - logfire.error(error_msg, exc_info=True) +# except Exception as e: +# error_msg = f"Error during plan update operations: {str(e)}" +# logfire.error(error_msg, exc_info=True) - planner_stream_output.steps.append(f"Plan update failed: {str(e)}") - planner_stream_output.status_code = a500 - await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) +# planner_stream_output.steps.append(f"Plan update failed: {str(e)}") +# planner_stream_output.status_code = 500 +# await _safe_websocket_send(ctx.deps.websocket, planner_stream_output) - return f"Failed to update the plan: {error_msg}" +# return f"Failed to update the plan: {error_msg}" - except Exception as e: - error_msg = f"Error updating plan: {str(e)}" - logfire.error(error_msg, exc_info=True) +# except Exception as e: +# error_msg = f"Error updating plan: {str(e)}" +# logfire.error(error_msg, exc_info=True) - # Update stream output with error - if ctx.deps.stream_output: - ctx.deps.stream_output.steps.append(f"Failed to update plan: {str(e)}") - await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) +# # Update stream output with error +# if ctx.deps.stream_output: +# ctx.deps.stream_output.steps.append(f"Failed to update plan: {str(e)}") +# await _safe_websocket_send(ctx.deps.websocket, ctx.deps.stream_output) - return f"Failed to update plan: {error_msg}" +# return f"Failed to update plan: {error_msg}" + +# # Helper function for sending WebSocket messages -# Helper function for sending WebSocket messages -async def _safe_websocket_send(websocket: Optional[WebSocket], message: Any) -> bool: - """Safely send message through websocket with error handling""" - try: - if websocket and websocket.client_state.CONNECTED: - await websocket.send_text(json.dumps(asdict(message))) - logfire.debug("WebSocket message sent (_safe_websocket_send): {message}", message=message) - return True - return False - except Exception as e: - logfire.error(f"WebSocket send failed: {str(e)}") - return False \ No newline at end of file diff --git a/cortex_on/agents/planner_agent.py b/cortex_on/agents/planner_agent.py index 897c22f..27db51e 100644 --- a/cortex_on/agents/planner_agent.py +++ b/cortex_on/agents/planner_agent.py @@ -8,6 +8,7 @@ from pydantic import BaseModel, Field from pydantic_ai import Agent from pydantic_ai.models.anthropic import AnthropicModel +from pydantic_ai.providers.anthropic import AnthropicProvider # Local application imports from utils.ant_client import get_client @@ -176,9 +177,11 @@ class PlannerResult(BaseModel): plan: str = Field(description="The generated or updated plan in string format - this should be the complete plan text") +provider = AnthropicProvider(api_key=os.environ.get("ANTHROPIC_API_KEY")) + model = AnthropicModel( model_name=os.environ.get("ANTHROPIC_MODEL_NAME"), - anthropic_client=get_client() + provider = provider ) planner_agent = Agent( diff --git a/cortex_on/agents/web_surfer.py b/cortex_on/agents/web_surfer.py index 34e2cfd..3e54610 100644 --- a/cortex_on/agents/web_surfer.py +++ b/cortex_on/agents/web_surfer.py @@ -11,15 +11,7 @@ from dotenv import load_dotenv from fastapi import WebSocket import logfire -from pydantic_ai.messages import ( - ArgsJson, - ModelRequest, - ModelResponse, - ToolCallPart, - ToolReturnPart, - UserPromptPart, -) - +from pydantic_ai.messages import ModelResponse, ModelRequest, ToolReturnPart, UserPromptPart, ToolCallPart # Local application imports from utils.stream_response_format import StreamResponse diff --git a/cortex_on/instructor.py b/cortex_on/instructor.py index b4f0efb..3362e3b 100644 --- a/cortex_on/instructor.py +++ b/cortex_on/instructor.py @@ -1,10 +1,13 @@ # Standard library imports import json import os +import asyncio import traceback from dataclasses import asdict from datetime import datetime from typing import Any, Dict, List, Optional, Tuple, Union +import uuid +import threading # Third-party imports from dotenv import load_dotenv @@ -16,26 +19,335 @@ from pydantic_ai.models.anthropic import AnthropicModel # Local application imports -from agents.code_agent import coder_agent +from agents.code_agent import CoderAgentDeps, coder_agent from agents.orchestrator_agent import orchestrator_agent, orchestrator_deps from agents.planner_agent import planner_agent from agents.web_surfer import WebSurfer from utils.ant_client import get_client from utils.stream_response_format import StreamResponse - +from agents.mcp_server import server load_dotenv() +# Flag to track if MCP server is running +_mcp_server_running = False +def start_mcp_server_in_thread(): + """Start the MCP server in a separate thread""" + global _mcp_server_running + if _mcp_server_running: + return + + _mcp_server_running = True + + def run_server(): + logfire.info("Starting MCP server...") + server.run(transport="sse") + + # Start in a separate thread + thread = threading.Thread(target=run_server, daemon=True) + thread.start() + logfire.info("MCP server thread started") class DateTimeEncoder(json.JSONEncoder): - """Custom JSON encoder that can handle datetime objects""" + """Custom JSON encoder that can handle datetime objects and Pydantic models""" def default(self, obj): if isinstance(obj, datetime): return obj.isoformat() + if isinstance(obj, BaseModel): + # Handle both Pydantic v1 and v2 + if hasattr(obj, 'model_dump'): + return obj.model_dump() + elif hasattr(obj, 'dict'): + return obj.dict() + # Fallback for any other Pydantic structure + return {k: v for k, v in obj.__dict__.items() if not k.startswith('_')} return super().default(obj) +def register_tools(websocket: WebSocket) -> None: + """ + Dynamically register MCP server tools with the provided WebSocket. + This ensures all tools have access to the active WebSocket connection. + """ + # First, unregister existing tools if they exist + tool_names = ["plan_task", "code_task", "web_surf_task", "ask_human", "planner_agent_update"] + for tool_name in tool_names: + if tool_name in server._tool_manager._tools: + del server._tool_manager._tools[tool_name] + + logfire.info("Registering MCP tools with WebSocket connection") + + # Function to create each tool with the websocket in closure + async def plan_task(task: str) -> str: + """Plans the task and assigns it to the appropriate agents""" + try: + logfire.info(f"Planning task: {task}") + print(f"Planning task: {task}") + # Create a new StreamResponse for Planner Agent + planner_stream_output = StreamResponse( + agent_name="Planner Agent", + instructions=task, + steps=[], + output="", + status_code=0 + ) + + await _safe_websocket_send(websocket, planner_stream_output) + + # Update planner stream + planner_stream_output.steps.append("Planning task...") + await _safe_websocket_send(websocket, planner_stream_output) + + # Run planner agent + planner_response = await planner_agent.run(user_prompt=task) + + # Update planner stream with results + plan_text = planner_response.data.plan + planner_stream_output.steps.append("Task planned successfully") + planner_stream_output.output = plan_text + planner_stream_output.status_code = 200 + await _safe_websocket_send(websocket, planner_stream_output) + + return f"Task planned successfully\nTask: {plan_text}" + except Exception as e: + error_msg = f"Error planning task: {str(e)}" + logfire.error(error_msg, exc_info=True) + + # Update planner stream with error + if 'planner_stream_output' in locals(): + planner_stream_output.steps.append(f"Planning failed: {str(e)}") + planner_stream_output.status_code = 500 + await _safe_websocket_send(websocket, planner_stream_output) + + return f"Failed to plan task: {error_msg}" + + async def code_task(task: str) -> str: + """Assigns coding tasks to the coder agent""" + try: + logfire.info(f"Assigning coding task: {task}") + print(f"Assigning coding task: {task}") + # Create a new StreamResponse for Coder Agent + coder_stream_output = StreamResponse( + agent_name="Coder Agent", + instructions=task, + steps=[], + output="", + status_code=0 + ) + + await _safe_websocket_send(websocket, coder_stream_output) + + # Create deps with the new stream_output + deps_for_coder_agent = CoderAgentDeps( + websocket=websocket, + stream_output=coder_stream_output + ) + + # Run coder agent + coder_response = await coder_agent.run( + user_prompt=task, + deps=deps_for_coder_agent + ) + + + # Update coder_stream_output with coding results + coder_stream_output.status_code = 200 + coder_stream_output.steps.append("Coding task completed successfully") + await _safe_websocket_send(websocket, coder_stream_output) + + # Add a reminder in the result message to update the plan using planner_agent_update + response_with_reminder = f"{coder_response.data.content}\n\nReminder: You must now call planner_agent_update with the completed task description: \"{task} (coder_agent)\"" + + return response_with_reminder + except Exception as e: + error_msg = f"Error assigning coding task: {str(e)}" + logfire.error(error_msg, exc_info=True) + + # Update coder_stream_output with error + if 'coder_stream_output' in locals(): + coder_stream_output.steps.append(f"Coding task failed: {str(e)}") + coder_stream_output.status_code = 500 + await _safe_websocket_send(websocket, coder_stream_output) + + return f"Failed to assign coding task: {error_msg}" + + async def web_surf_task(task: str) -> str: + """Assigns web surfing tasks to the web surfer agent""" + try: + logfire.info(f"Assigning web surfing task: {task}") + + # Create a new StreamResponse for WebSurfer + web_surfer_stream_output = StreamResponse( + agent_name="Web Surfer", + instructions=task, + steps=[], + output="", + status_code=0, + live_url=None + ) + + await _safe_websocket_send(websocket, web_surfer_stream_output) + + # Initialize WebSurfer agent + web_surfer_agent = WebSurfer(api_url="http://localhost:8000/api/v1/web/stream") + + # Run WebSurfer with its own stream_output + success, message, messages = await web_surfer_agent.generate_reply( + instruction=task, + websocket=websocket, + stream_output=web_surfer_stream_output + ) + + # Update WebSurfer's stream_output with final result + if success: + web_surfer_stream_output.steps.append("Web search completed successfully") + web_surfer_stream_output.output = message + web_surfer_stream_output.status_code = 200 + + # Add a reminder to update the plan + message_with_reminder = f"{message}\n\nReminder: You must now call planner_agent_update with the completed task description: \"{task} (web_surfer_agent)\"" + else: + web_surfer_stream_output.steps.append(f"Web search completed with issues: {message[:100]}") + web_surfer_stream_output.status_code = 500 + message_with_reminder = message + + await _safe_websocket_send(websocket, web_surfer_stream_output) + + web_surfer_stream_output.steps.append(f"WebSurfer completed: {'Success' if success else 'Failed'}") + await _safe_websocket_send(websocket, web_surfer_stream_output) + + return message_with_reminder + except Exception as e: + error_msg = f"Error assigning web surfing task: {str(e)}" + logfire.error(error_msg, exc_info=True) + + # Update WebSurfer's stream_output with error + if 'web_surfer_stream_output' in locals(): + web_surfer_stream_output.steps.append(f"Web search failed: {str(e)}") + web_surfer_stream_output.status_code = 500 + await _safe_websocket_send(websocket, web_surfer_stream_output) + return f"Failed to assign web surfing task: {error_msg}" + + async def planner_agent_update(completed_task: str) -> str: + """ + Updates the todo.md file to mark a task as completed and returns the full updated plan. + """ + try: + logfire.info(f"Updating plan with completed task: {completed_task}") + print(f"Updating plan with completed task: {completed_task}") + # Create a new StreamResponse for Planner Agent update + planner_stream_output = StreamResponse( + agent_name="Planner Agent", + instructions=f"Update todo.md to mark as completed: {completed_task}", + steps=[], + output="", + status_code=0 + ) + + # Send initial update + await _safe_websocket_send(websocket, planner_stream_output) + + # Directly read and update the todo.md file + base_dir = os.path.abspath(os.path.dirname(os.path.dirname(__file__))) + planner_dir = os.path.join(base_dir, "agents", "planner") + todo_path = os.path.join(planner_dir, "todo.md") + + planner_stream_output.steps.append("Reading current todo.md...") + await _safe_websocket_send(websocket, planner_stream_output) + + # Make sure the directory exists + os.makedirs(planner_dir, exist_ok=True) + + try: + # Check if todo.md exists + if not os.path.exists(todo_path): + planner_stream_output.steps.append("No todo.md file found. Will create new one after task completion.") + await _safe_websocket_send(websocket, planner_stream_output) + + # We'll directly call planner_agent.run() to create a new plan first + plan_prompt = f"Create a simple task plan based on this completed task: {completed_task}" + plan_response = await planner_agent.run(user_prompt=plan_prompt) + current_content = plan_response.data.plan + else: + # Read existing todo.md + with open(todo_path, "r") as file: + current_content = file.read() + planner_stream_output.steps.append(f"Found existing todo.md ({len(current_content)} bytes)") + await _safe_websocket_send(websocket, planner_stream_output) + + # Now call planner_agent.run() with specific instructions to update the plan + update_prompt = f""" + Here is the current todo.md content: + + {current_content} + + Please update this plan to mark the following task as completed: {completed_task} + Return ONLY the fully updated plan with appropriate tasks marked as [x] instead of [ ]. + """ + + planner_stream_output.steps.append("Asking planner to update the plan...") + await _safe_websocket_send(websocket, planner_stream_output) + + updated_plan_response = await planner_agent.run(user_prompt=update_prompt) + updated_plan = updated_plan_response.data.plan + + # Write the updated plan back to todo.md + with open(todo_path, "w") as file: + file.write(updated_plan) + + planner_stream_output.steps.append("Plan updated successfully") + planner_stream_output.output = updated_plan + planner_stream_output.status_code = 200 + await _safe_websocket_send(websocket, planner_stream_output) + + return updated_plan + + except Exception as e: + error_msg = f"Error during plan update operations: {str(e)}" + logfire.error(error_msg, exc_info=True) + + planner_stream_output.steps.append(f"Plan update failed: {str(e)}") + planner_stream_output.status_code = 500 + await _safe_websocket_send(websocket, planner_stream_output) + + return f"Failed to update the plan: {error_msg}" + + except Exception as e: + error_msg = f"Error updating plan: {str(e)}" + logfire.error(error_msg, exc_info=True) + + return f"Failed to update plan: {error_msg}" + + # Helper function for websocket communication + async def _safe_websocket_send(socket: WebSocket, message: Any) -> bool: + """Safely send message through websocket with error handling""" + try: + if socket and socket.client_state.CONNECTED: + await socket.send_text(json.dumps(asdict(message))) + logfire.debug("WebSocket message sent (_safe_websocket_send): {message}", message=message) + return True + return False + except Exception as e: + logfire.error(f"WebSocket send failed: {str(e)}") + return False + + # Now register all the generated tools with the MCP server + tool_definitions = { + "plan_task": (plan_task, "Plans the task and assigns it to the appropriate agents"), + "code_task": (code_task, "Assigns coding tasks to the coder agent"), + "web_surf_task": (web_surf_task, "Assigns web surfing tasks to the web surfer agent"), + "planner_agent_update": (planner_agent_update, "Updates the todo.md file to mark a task as completed") + } + + # Register each tool + for name, (fn, desc) in tool_definitions.items(): + server._tool_manager.add_tool(fn, name=name, description=desc) + + logfire.info(f"Successfully registered {len(tool_definitions)} tools with the MCP server") + + + # Main Orchestrator Class class SystemInstructor: def __init__(self): @@ -84,18 +396,28 @@ async def run(self, task: str, websocket: WebSocket) -> List[Dict[str, Any]]: ) try: + # Register tools first - before MCP server starts or is accessed by orchestrator + register_tools(websocket=self.websocket) + + # Start the MCP server if it's not already running + start_mcp_server_in_thread() + + # Give MCP server a moment to initialize + await asyncio.sleep(1) + # Initialize system await self._safe_websocket_send(stream_output) stream_output.steps.append("Agents initialized successfully") await self._safe_websocket_send(stream_output) - orchestrator_response = await orchestrator_agent.run( - user_prompt=task, - deps=deps_for_orchestrator - ) - stream_output.output = orchestrator_response.data + async with orchestrator_agent.run_mcp_servers(): + orchestrator_response = await orchestrator_agent.run( + user_prompt=task, + deps=deps_for_orchestrator + ) + stream_output.output = orchestrator_response.output stream_output.status_code = 200 - logfire.debug(f"Orchestrator response: {orchestrator_response.data}") + logfire.debug(f"Orchestrator response: {orchestrator_response.output}") await self._safe_websocket_send(stream_output) logfire.info("Task completed successfully") @@ -112,11 +434,17 @@ async def run(self, task: str, websocket: WebSocket) -> List[Dict[str, Any]]: await self._safe_websocket_send(stream_output) # Even in case of critical error, return what we have - return [asdict(i) for i in self.orchestrator_response] + try: + return [json.loads(json.dumps(asdict(i), cls=DateTimeEncoder)) for i in self.orchestrator_response] + except Exception as serialize_error: + logfire.error(f"Failed to serialize response: {str(serialize_error)}") + # Last resort - return a simple error message + return [{"error": error_msg, "status_code": 500}] finally: logfire.info("Orchestration process complete") # Clear any sensitive data + async def shutdown(self): """Clean shutdown of orchestrator""" try: diff --git a/cortex_on/multi_lang_env/Dockerfile b/cortex_on/multi_lang_env/Dockerfile new file mode 100644 index 0000000..7d32e46 --- /dev/null +++ b/cortex_on/multi_lang_env/Dockerfile @@ -0,0 +1,115 @@ +FROM ubuntu:22.04 + +# Prevent interactive prompts during installation +ENV DEBIAN_FRONTEND=noninteractive + +# Set labels +LABEL maintainer="CortexON Team" +LABEL description="Multi-language execution environment for CortexON" + +# Install common dependencies +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + curl \ + wget \ + git \ + gnupg \ + software-properties-common \ + unzip \ + vim \ + nano \ + && rm -rf /var/lib/apt/lists/* + +# Set up language directories +RUN mkdir -p /environments/python \ + /environments/java \ + /environments/cpp \ + /environments/javascript \ + /environments/typescript \ + /environments/ruby \ + /environments/go \ + /environments/rust \ + /environments/php + +# Python setup +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3 \ + python3-pip \ + python3-venv \ + && ln -s /usr/bin/python3 /usr/bin/python \ + && pip3 install --no-cache-dir numpy pandas matplotlib \ + && rm -rf /var/lib/apt/lists/* + +# Java setup +RUN apt-get update && apt-get install -y --no-install-recommends \ + openjdk-17-jdk \ + && rm -rf /var/lib/apt/lists/* + +# C++ setup - already included in build-essential + +# JavaScript/TypeScript setup +RUN curl -fsSL https://deb.nodesource.com/setup_18.x | bash - \ + && apt-get install -y nodejs \ + && npm install -g typescript axios \ + && rm -rf /var/lib/apt/lists/* + +# Ruby setup +RUN apt-get update && apt-get install -y --no-install-recommends \ + ruby-full \ + && gem install bundler \ + && rm -rf /var/lib/apt/lists/* + +# Go setup +RUN wget https://golang.org/dl/go1.20.linux-amd64.tar.gz \ + && tar -C /usr/local -xzf go1.20.linux-amd64.tar.gz \ + && rm go1.20.linux-amd64.tar.gz +ENV PATH="/usr/local/go/bin:${PATH}" + +# Rust setup +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \ + && . "$HOME/.cargo/env" \ + && rustup component add rustfmt +ENV PATH="/root/.cargo/bin:${PATH}" + +# PHP setup +RUN apt-get update && apt-get install -y --no-install-recommends \ + php-cli \ + && rm -rf /var/lib/apt/lists/* + +# Create language-specific work directories +RUN mkdir -p /app/python \ + /app/java \ + /app/cpp \ + /app/javascript \ + /app/typescript \ + /app/ruby \ + /app/go \ + /app/rust \ + /app/php + +# Create activation scripts for different language environments +COPY setup/*.sh /setup/ +RUN chmod +x /setup/*.sh + +# Set working directory +WORKDIR /app + +# Set up environment selection script +RUN echo '#!/bin/bash\n\ +case "$1" in\n\ + python) source /setup/activate_python.sh ;;\n\ + java) source /setup/activate_java.sh ;;\n\ + cpp) source /setup/activate_cpp.sh ;;\n\ + javascript) source /setup/activate_javascript.sh ;;\n\ + typescript) source /setup/activate_typescript.sh ;;\n\ + ruby) source /setup/activate_ruby.sh ;;\n\ + go) source /setup/activate_go.sh ;;\n\ + rust) source /setup/activate_rust.sh ;;\n\ + php) source /setup/activate_php.sh ;;\n\ + *) echo "Unknown language: $1" ;;\n\ +esac\n\ +' > /usr/local/bin/use_env && \ + chmod +x /usr/local/bin/use_env + +CMD ["tail", "-f", "/dev/null"] \ No newline at end of file diff --git a/cortex_on/multi_lang_env/README.md b/cortex_on/multi_lang_env/README.md new file mode 100644 index 0000000..db22243 --- /dev/null +++ b/cortex_on/multi_lang_env/README.md @@ -0,0 +1,57 @@ +# CortexON Multi-Language Environment + +This directory contains the configuration for a consolidated multi-language execution environment for CortexON. Instead of running separate containers for each programming language, we use a single container with all language runtimes installed and provide mechanisms to switch between them. + +## How it Works + +The multi-language container includes: + +1. All necessary language runtimes (Python, Java, C++, JavaScript, TypeScript, Ruby, Go, Rust, PHP) +2. Language-specific directories under `/app/` for code execution +3. Environment switching scripts that set up the appropriate context for each language + +## Benefits + +- **Reduced resource usage**: A single container instead of multiple containers +- **Simplified management**: Only one container to monitor and maintain +- **Easy scaling**: Add new languages by extending a single container + +## Implementation Details + +- The container is built from the Dockerfile in this directory +- Language activation scripts in `/setup/` handle environment switching +- Each language has a dedicated workspace in `/app/` +- The main `use_env` script allows changing language contexts + +## Usage + +The container is primarily managed through the `DockerEnvironment` class in `cortex_on/utils/docker_executor.py`, which has been updated to: + +1. Connect to a single container instead of multiple containers +2. Switch language environments as needed +3. Execute code in the appropriate language context + +## Adding a New Language + +To add support for a new language: + +1. Update the Dockerfile to install the required runtime and tools +2. Create an activation script in the `setup/` directory +3. Add the language configuration to `SUPPORTED_LANGUAGES` in `docker_executor.py` +4. Add an entry in the `use_env` script in the Dockerfile + +## Building the Container + +The container is built automatically as part of the main docker-compose setup: + +```bash +docker-compose build multi_language_env +``` + +## Running the Container Standalone + +If needed, you can run the container standalone: + +```bash +docker-compose up multi_language_env +``` \ No newline at end of file diff --git a/cortex_on/multi_lang_env/setup/activate_cpp.sh b/cortex_on/multi_lang_env/setup/activate_cpp.sh new file mode 100644 index 0000000..39ec052 --- /dev/null +++ b/cortex_on/multi_lang_env/setup/activate_cpp.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +# Activate C++ environment +echo "Activating C++ environment" +export LANG_ENV="cpp" +cd /app/cpp +export PATH="/usr/bin:$PATH" +echo "C++ $(g++ --version | head -n 1) environment activated" \ No newline at end of file diff --git a/cortex_on/multi_lang_env/setup/activate_go.sh b/cortex_on/multi_lang_env/setup/activate_go.sh new file mode 100644 index 0000000..725c107 --- /dev/null +++ b/cortex_on/multi_lang_env/setup/activate_go.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +# Activate Go environment +echo "Activating Go environment" +export LANG_ENV="go" +cd /app/go +export PATH="/usr/local/go/bin:$PATH" +echo "Go $(go version) environment activated" \ No newline at end of file diff --git a/cortex_on/multi_lang_env/setup/activate_java.sh b/cortex_on/multi_lang_env/setup/activate_java.sh new file mode 100644 index 0000000..8f0abc3 --- /dev/null +++ b/cortex_on/multi_lang_env/setup/activate_java.sh @@ -0,0 +1,9 @@ +#!/bin/bash + +# Activate Java environment +echo "Activating Java environment" +export LANG_ENV="java" +cd /app/java +export JAVA_HOME=$(dirname $(dirname $(readlink -f $(which javac)))) +export PATH="$JAVA_HOME/bin:$PATH" +echo "Java $(java -version 2>&1 | head -n 1) environment activated" \ No newline at end of file diff --git a/cortex_on/multi_lang_env/setup/activate_javascript.sh b/cortex_on/multi_lang_env/setup/activate_javascript.sh new file mode 100644 index 0000000..dfd66f1 --- /dev/null +++ b/cortex_on/multi_lang_env/setup/activate_javascript.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +# Activate JavaScript environment +echo "Activating JavaScript environment" +export LANG_ENV="javascript" +cd /app/javascript +export PATH="/usr/bin:$PATH" +echo "Node.js $(node --version) environment activated" \ No newline at end of file diff --git a/cortex_on/multi_lang_env/setup/activate_php.sh b/cortex_on/multi_lang_env/setup/activate_php.sh new file mode 100644 index 0000000..5a8222e --- /dev/null +++ b/cortex_on/multi_lang_env/setup/activate_php.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +# Activate PHP environment +echo "Activating PHP environment" +export LANG_ENV="php" +cd /app/php +export PATH="/usr/bin:$PATH" +echo "PHP $(php --version | head -n 1) environment activated" \ No newline at end of file diff --git a/cortex_on/multi_lang_env/setup/activate_python.sh b/cortex_on/multi_lang_env/setup/activate_python.sh new file mode 100644 index 0000000..ca93c90 --- /dev/null +++ b/cortex_on/multi_lang_env/setup/activate_python.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +# Activate Python environment +echo "Activating Python environment" +export LANG_ENV="python" +cd /app/python +export PATH="/usr/bin:$PATH" +echo "Python $(python --version) environment activated" \ No newline at end of file diff --git a/cortex_on/multi_lang_env/setup/activate_ruby.sh b/cortex_on/multi_lang_env/setup/activate_ruby.sh new file mode 100644 index 0000000..d02ac8b --- /dev/null +++ b/cortex_on/multi_lang_env/setup/activate_ruby.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +# Activate Ruby environment +echo "Activating Ruby environment" +export LANG_ENV="ruby" +cd /app/ruby +export PATH="/usr/bin:$PATH" +echo "Ruby $(ruby --version) environment activated" \ No newline at end of file diff --git a/cortex_on/multi_lang_env/setup/activate_rust.sh b/cortex_on/multi_lang_env/setup/activate_rust.sh new file mode 100644 index 0000000..cd106ad --- /dev/null +++ b/cortex_on/multi_lang_env/setup/activate_rust.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +# Activate Rust environment +echo "Activating Rust environment" +export LANG_ENV="rust" +cd /app/rust +export PATH="$HOME/.cargo/bin:$PATH" +echo "Rust $(rustc --version) environment activated" \ No newline at end of file diff --git a/cortex_on/multi_lang_env/setup/activate_typescript.sh b/cortex_on/multi_lang_env/setup/activate_typescript.sh new file mode 100644 index 0000000..cdac3de --- /dev/null +++ b/cortex_on/multi_lang_env/setup/activate_typescript.sh @@ -0,0 +1,8 @@ +#!/bin/bash + +# Activate TypeScript environment +echo "Activating TypeScript environment" +export LANG_ENV="typescript" +cd /app/typescript +export PATH="/usr/bin:$PATH" +echo "TypeScript $(tsc --version) environment activated" \ No newline at end of file diff --git a/cortex_on/requirements.txt b/cortex_on/requirements.txt index a635c7e..e656aee 100644 --- a/cortex_on/requirements.txt +++ b/cortex_on/requirements.txt @@ -2,7 +2,7 @@ aiohappyeyeballs==2.4.4 aiohttp==3.11.11 aiosignal==1.3.2 annotated-types==0.7.0 -anthropic==0.42.0 +anthropic==0.49.0 anyio==4.7.0 asyncio-atexit==1.0.1 attrs==24.3.0 @@ -25,7 +25,7 @@ frozenlist==1.5.0 google-auth==2.37.0 googleapis-common-protos==1.66.0 griffe==1.5.4 -groq==0.13.1 +groq==0.15.0 h11==0.14.0 httpcore==1.0.7 httpx==0.27.2 @@ -44,7 +44,7 @@ mistralai==1.2.5 multidict==6.1.0 mypy-extensions==1.0.0 numpy==2.2.1 -openai==1.58.1 +openai==1.74.0 opentelemetry-api==1.29.0 opentelemetry-exporter-otlp-proto-common==1.29.0 opentelemetry-exporter-otlp-proto-http==1.29.0 @@ -65,9 +65,12 @@ pyasn1_modules==0.4.1 pycparser==2.22 pycryptodome==3.21.0 pydantic==2.10.4 -pydantic-ai==0.0.17 -pydantic-ai-slim==0.0.17 +pydantic-ai==0.1.2 +pydantic-ai-slim==0.1.2 pydantic_core==2.27.2 +pydantic-ai==0.1.2 +pydantic-ai-slim==0.1.2 +mcp==1.6.0 Pygments==2.18.0 python-dateutil==2.9.0.post0 python-dotenv==1.0.1 diff --git a/cortex_on/test_code_formatter.py b/cortex_on/test_code_formatter.py new file mode 100644 index 0000000..fad453a --- /dev/null +++ b/cortex_on/test_code_formatter.py @@ -0,0 +1,89 @@ +import json +from utils.code_formatter import ( + format_code_for_frontend, + format_output_for_frontend, + format_execution_result +) + +def test_code_formatting(): + # Test Python code formatting + python_code = """ +def hello_world(): + print("Hello, World!") + +for i in range(5): + print(f"Number: {i}") + +hello_world() + + + +# Too many blank lines above +""" + + formatted_python = format_code_for_frontend(python_code, "python") + print("Python code formatting:") + print(formatted_python) + print("\n" + "-"*50 + "\n") + + # Test JavaScript code formatting + js_code = """ +function helloWorld() { + console.log("Hello, World!"); +} + +for (let i = 0; i < 5; i++) { + console.log(`Number: ${i}`); +} + +helloWorld(); +""" + + formatted_js = format_code_for_frontend(js_code, "js") + print("JavaScript code formatting:") + print(formatted_js) + print("\n" + "-"*50 + "\n") + + # Test output formatting + output = "Hello, World!\nNumber: 0\nNumber: 1\nNumber: 2\nNumber: 3\nNumber: 4" + formatted_output = format_output_for_frontend(output) + print("Output formatting:") + print(formatted_output) + print("\n" + "-"*50 + "\n") + + # Test error formatting + error_result = {"error": "Container timeout after 30 seconds"} + formatted_error = format_execution_result(python_code, "python", error_result) + print("Error formatting:") + print(formatted_error) + print("\n" + "-"*50 + "\n") + + # Test successful execution result formatting + success_result = { + "execution_id": "12345", + "language": "python", + "stdout": "Hello, World!\nNumber: 0\nNumber: 1\nNumber: 2\nNumber: 3\nNumber: 4", + "stderr": "", + "exit_code": 0, + "success": True + } + formatted_success = format_execution_result(python_code, "python", success_result) + print("Success result formatting:") + print(formatted_success) + print("\n" + "-"*50 + "\n") + + # Test failed execution result formatting + failed_result = { + "execution_id": "12345", + "language": "python", + "stdout": "", + "stderr": "NameError: name 'undefined_variable' is not defined", + "exit_code": 1, + "success": False + } + formatted_failure = format_execution_result(python_code, "python", failed_result) + print("Failed result formatting:") + print(formatted_failure) + +if __name__ == "__main__": + test_code_formatting() \ No newline at end of file diff --git a/cortex_on/test_docker.py b/cortex_on/test_docker.py new file mode 100644 index 0000000..a669e6f --- /dev/null +++ b/cortex_on/test_docker.py @@ -0,0 +1,57 @@ +import asyncio +import logging +from utils.docker_executor import get_or_create_environment, cleanup_environments + +# Configure logging +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +async def test_docker_environment(): + try: + # Create test environment + env_id = "test123" + env = get_or_create_environment(env_id, 'python') + + # Start the environment + logger.info("Starting Docker environment...") + start_result = await env.start() + logger.info(f"Start result: {start_result}") + + # Write a test file + test_content = 'print("Hello from Docker!")' + logger.info("Writing test file...") + write_result = await env.write_file('test.py', test_content) + logger.info(f"Write result: {write_result}") + + # List files + logger.info("Listing files...") + list_result = await env.list_files() + logger.info(f"List result: {list_result}") + + # Read the file back + logger.info("Reading file...") + read_result = await env.read_file('test.py') + logger.info(f"Read result: {read_result}") + + # Execute the file + logger.info("Executing file...") + exec_result = await env.execute_code('python', 'test.py') + logger.info(f"Execution result: {exec_result}") + + # Clean up + logger.info("Stopping environment...") + stop_result = await env.stop() + logger.info(f"Stop result: {stop_result}") + + return True + except Exception as e: + logger.error(f"Error in test: {str(e)}", exc_info=True) + return False + finally: + # Ensure cleanup + await cleanup_environments() + +if __name__ == "__main__": + logger.info("Starting Docker environment test...") + asyncio.run(test_docker_environment()) + logger.info("Test completed.") \ No newline at end of file diff --git a/cortex_on/test_docker_executor.py b/cortex_on/test_docker_executor.py new file mode 100644 index 0000000..502eb89 --- /dev/null +++ b/cortex_on/test_docker_executor.py @@ -0,0 +1,61 @@ +import asyncio +import json +from utils.docker_executor import run_docker_container + +async def test_python(): + code = """ +print("Hello from Python!") +for i in range(5): + print(f"Number: {i}") +""" + result = await run_docker_container("python", code) + print("Python Test Results:") + print(json.dumps(result, indent=2)) + +async def test_javascript(): + code = """ +console.log("Hello from JavaScript!"); +for (let i = 0; i < 5; i++) { + console.log(`Number: ${i}`); +} +""" + result = await run_docker_container("javascript", code) + print("\nJavaScript Test Results:") + print(json.dumps(result, indent=2)) + +async def test_cpp(): + code = """ +#include +using namespace std; + +int main() { + cout << "Hello from C++!" << endl; + for (int i = 0; i < 5; i++) { + cout << "Number: " << i << endl; + } + return 0; +} +""" + result = await run_docker_container("cpp", code) + print("\nC++ Test Results:") + print(json.dumps(result, indent=2)) + +async def test_infinite_loop(): + code = """ +# This should be killed after the timeout +while True: + pass +""" + result = await run_docker_container("python", code) + print("\nInfinite Loop Test Results:") + print(json.dumps(result, indent=2)) + +async def main(): + print("Testing Docker Executor...") + await test_python() + await test_javascript() + await test_cpp() + await test_infinite_loop() + +if __name__ == "__main__": + asyncio.run(main()) \ No newline at end of file diff --git a/cortex_on/test_prime.py b/cortex_on/test_prime.py new file mode 100644 index 0000000..37aa4b2 --- /dev/null +++ b/cortex_on/test_prime.py @@ -0,0 +1,49 @@ +import asyncio +import logging +from utils.docker_executor import run_docker_container + +# Set up logging +logging.basicConfig(level=logging.INFO) + +# Test prime function code +prime_test_code = """ +def is_prime(n): + \"\"\"Check if a number is prime.\"\"\" + if n <= 1: + return False + if n <= 3: + return True + if n % 2 == 0 or n % 3 == 0: + return False + i = 5 + while i * i <= n: + if n % i == 0 or n % (i + 2) == 0: + return False + i += 6 + return True + +# Test cases +test_cases = [2, 3, 4, 5, 15, 17, 20, 97] +print("Testing prime numbers:") +for num in test_cases: + result = is_prime(num) + print(f"{num} is {'prime' if result else 'not prime'}") +""" + +async def test_prime_function(): + print("Running Docker test for prime function...") + result = await run_docker_container("python", prime_test_code) + + print("\n--- Execution Result ---") + print(f"Execution ID: {result.get('execution_id')}") + print(f"Success: {result.get('success')}") + print(f"Exit Code: {result.get('exit_code')}") + print("\n--- Output ---") + print(result.get('stdout')) + + if result.get('stderr'): + print("\n--- Errors ---") + print(result.get('stderr')) + +if __name__ == "__main__": + asyncio.run(test_prime_function()) \ No newline at end of file diff --git a/cortex_on/utils/code_formatter.py b/cortex_on/utils/code_formatter.py new file mode 100644 index 0000000..3175723 --- /dev/null +++ b/cortex_on/utils/code_formatter.py @@ -0,0 +1,204 @@ +import re +import logging +from typing import Dict, Optional + +logger = logging.getLogger(__name__) + +# Language syntax highlighting mappings +LANGUAGE_HIGHLIGHTS = { + "python": "python", + "javascript": "javascript", + "typescript": "typescript", + "java": "java", + "cpp": "cpp", + "c++": "cpp", + "c": "c", + # Additional languages + "ruby": "ruby", + "go": "go", + "rust": "rust", + "php": "php", + "csharp": "csharp", + "kotlin": "kotlin", + "swift": "swift", + "r": "r", + "scala": "scala", + "perl": "perl", + "dart": "dart", + "julia": "julia" +} + +def format_code_for_frontend(code: str, language: str) -> str: + """ + Format code for frontend display with proper syntax highlighting. + + Args: + code: The source code to format + language: The programming language + + Returns: + Markdown-formatted code ready for frontend display + """ + # Normalize language name + normalized_language = normalize_language(language) + + # Get proper language identifier for syntax highlighting + highlight_lang = LANGUAGE_HIGHLIGHTS.get(normalized_language, normalized_language) + + # Clean up code (remove excessive newlines, normalize spacing) + cleaned_code = clean_code(code) + + # Format as markdown code block with language syntax highlighting + return f"```{highlight_lang}\n{cleaned_code}\n```" + +def format_output_for_frontend(output: str) -> str: + """ + Format execution output for frontend display. + + Args: + output: The execution output text + + Returns: + Formatted output text ready for frontend display + """ + # Log the raw output for debugging + logger.info(f"Raw output before formatting: [{output}]") + + # Check if output is None or empty + if output is None: + logger.info("Output is None") + return "*No output produced*" + + # Clean up output + cleaned_output = output.strip() + + if not cleaned_output: + logger.info("Output is empty after stripping") + return "*No output produced*" + + # Strip excessive blank lines + cleaned_output = re.sub(r'\n{3,}', '\n\n', cleaned_output) + + # Log the final formatted output + logger.info(f"Formatted output: [{cleaned_output}]") + + # Format terminal output section - using plain backticks for cleaner display + return f"```\n{cleaned_output}\n```" + +def format_execution_result(code: str, language: str, result: Dict) -> str: + """ + Create a complete formatted output with both code and execution results. + + Args: + code: The source code + language: The programming language + result: The execution result dictionary + + Returns: + A formatted string containing both code and execution results + """ + # Format the code section + formatted_code = format_code_for_frontend(code, language) + + # Check if there was an error in execution setup (not in the code itself) + if "error" in result: + error_message = result["error"] + logger.debug(f"Formatting error output: {error_message}") + return f"{formatted_code}\n\n## Errors\n\n{format_output_for_frontend(error_message)}\n\n## Status\n\n**❌ Execution failed**" + + # Process stdout and stderr + stdout = result.get("stdout", "").strip() + stderr = result.get("stderr", "").strip() + + # Log output for debugging + logger.info(f"Formatting stdout: {stdout[:200]}{'...' if len(stdout) > 200 else ''}") + logger.info(f"Formatting stderr: {stderr[:200]}{'...' if len(stderr) > 200 else ''}") + + # Format sections for the frontend + sections = [] + + # Always add the code section first + sections.append(formatted_code) + + # Add output section if stdout exists or explicitly note if no output + if stdout: + sections.append(f"## Output\n\n{format_output_for_frontend(stdout)}") + else: + sections.append("## Output\n\n*No output produced*") + + # Add errors section if stderr exists + if stderr: + sections.append(f"## Errors\n\n{format_output_for_frontend(stderr)}") + + # Add execution status with emoji for better visibility + if result.get("success", False): + status = "**✅ Execution completed successfully**" + else: + exit_code = result.get("exit_code", "unknown") + status = f"**❌ Execution failed** (Exit code: {exit_code})" + + sections.append(f"## Status\n\n{status}") + + # Join all sections with double newlines for proper separation + return "\n\n".join(sections) + +def normalize_language(language: str) -> str: + """ + Normalize language name. + + Args: + language: The programming language name to normalize + + Returns: + Normalized language name + """ + # Convert to lowercase and strip whitespace + normalized = language.lower().strip() + + # Handle common aliases + language_aliases = { + "python3": "python", + "py": "python", + "js": "javascript", + "ts": "typescript", + "node": "javascript", + "nodejs": "javascript", + "java": "java", + "c++": "cpp", + "c": "c", + # Additional languages + "rb": "ruby", + "golang": "go", + "rs": "rust", + "kt": "kotlin", + "dotnet": "csharp", + "c#": "csharp", + "dot-net": "csharp", + "pl": "perl", + "php7": "php", + "php8": "php", + "jl": "julia", + "dart2": "dart", + "scala3": "scala", + "r-lang": "r" + } + + return language_aliases.get(normalized, normalized) + +def clean_code(code: str) -> str: + """ + Clean up code by normalizing whitespace, indentation, etc. + + Args: + code: The source code to clean + + Returns: + Cleaned code + """ + # Remove leading/trailing whitespace + cleaned = code.strip() + + # Remove multiple consecutive blank lines (more than 2) + cleaned = re.sub(r'\n{3,}', '\n\n', cleaned) + + return cleaned \ No newline at end of file diff --git a/cortex_on/utils/docker_executor.py b/cortex_on/utils/docker_executor.py new file mode 100644 index 0000000..fa821ed --- /dev/null +++ b/cortex_on/utils/docker_executor.py @@ -0,0 +1,845 @@ +import os +import uuid +import logging +import tempfile +import time +import json +import io +import tarfile +import docker +from typing import Dict, Optional, List, Union, Tuple, Any + +# Configure logger with more detailed format +logger = logging.getLogger(__name__) +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' +) + +# Language configurations +SUPPORTED_LANGUAGES = { + "python": { + "container_name": "cortexon_multi_env", + "file_extension": ".py", + "execute_cmd": lambda filename: f"python {filename}", + "work_dir": "/app/python" + }, + "java": { + "container_name": "cortexon_multi_env", + "file_extension": ".java", + "execute_cmd": lambda filename: f"javac {filename} && java -cp . {os.path.splitext(os.path.basename(filename))[0]}", + "work_dir": "/app/java" + }, + "cpp": { + "container_name": "cortexon_multi_env", + "file_extension": ".cpp", + "execute_cmd": lambda filename: f"g++ {filename} -o /tmp/program", + "work_dir": "/app/cpp" + }, + "javascript": { + "container_name": "cortexon_multi_env", + "file_extension": ".js", + "execute_cmd": lambda filename: f"node {filename}", + "work_dir": "/app/javascript" + }, + "typescript": { + "container_name": "cortexon_multi_env", + "file_extension": ".ts", + "execute_cmd": lambda filename: f"tsc {filename} --outFile /tmp/out.js && node /tmp/out.js", + "work_dir": "/app/typescript" + }, + "ruby": { + "container_name": "cortexon_multi_env", + "file_extension": ".rb", + "execute_cmd": lambda filename: f"ruby {filename}", + "work_dir": "/app/ruby" + }, + "go": { + "container_name": "cortexon_multi_env", + "file_extension": ".go", + "execute_cmd": lambda filename: f"cd {os.path.dirname(filename) or '.'} && go run {os.path.basename(filename)}", + "work_dir": "/app/go" + }, + "rust": { + "container_name": "cortexon_multi_env", + "file_extension": ".rs", + "execute_cmd": lambda filename: f"rustc {filename} -o /tmp/program && /tmp/program", + "work_dir": "/app/rust" + }, + "php": { + "container_name": "cortexon_multi_env", + "file_extension": ".php", + "execute_cmd": lambda filename: f"php {filename}", + "work_dir": "/app/php" + } +} + +# Language aliases mapping +LANGUAGE_ALIASES = { + "python3": "python", + "py": "python", + "c++": "cpp", + "node": "javascript", + "nodejs": "javascript", + "js": "javascript", + "rb": "ruby", + "golang": "go", + "rs": "rust", + "ts": "typescript", + "php": "php", + "java": "java" +} + +class DockerEnvironment: + """ + Connects to a persistent Docker container for code execution. + These containers should be defined in the docker-compose.yml. + """ + def __init__( + self, + language: str = "python", + work_dir: str = None + ): + """ + Initialize a connection to a Docker environment + + Args: + language: The primary programming language for this environment + work_dir: Working directory in the container (optional, will use language-specific dir if not provided) + """ + self.client = docker.from_env() + self.language = language + self.container_name = SUPPORTED_LANGUAGES[language]["container_name"] + self.work_dir = work_dir if work_dir else SUPPORTED_LANGUAGES[language]["work_dir"] + self.files = {} # Keep track of files in the container + self.active = False + self.container = None + + logger.info(f"Initialized Docker environment for {self.language}") + + async def connect(self) -> Dict[str, Any]: + """ + Connect to the persistent Docker container for this environment + + Returns: + Status dictionary with success/error information + """ + if self.active: + logger.info(f"Already connected to container {self.container_name}") + return {"success": True, "message": "Already connected"} + + try: + logger.info(f"Connecting to container {self.container_name}") + + # Get container by name + self.container = self.client.containers.get(self.container_name) + + # Check if container is running + if self.container.status != "running": + logger.info(f"Container {self.container_name} is not running, attempting to start") + self.container.start() + + self.active = True + logger.info(f"Successfully connected to container {self.container_name}") + + # Create workspace directory if it doesn't exist + self._exec_command(f"mkdir -p {self.work_dir}") + + # Activate the appropriate language environment + self._activate_language_environment() + + return { + "success": True, + "container_id": self.container.id, + "message": f"Successfully connected to container {self.container_name} and activated {self.language} environment" + } + + except docker.errors.NotFound: + error_msg = f"Container {self.container_name} not found. Make sure it's defined in docker-compose.yml" + logger.error(error_msg) + return {"success": False, "error": error_msg} + + except Exception as e: + error_msg = f"Failed to connect to container: {str(e)}" + logger.error(error_msg, exc_info=True) + return {"success": False, "error": error_msg} + + def _activate_language_environment(self): + """ + Activate the appropriate language environment in the container + """ + try: + logger.info(f"Activating {self.language} environment in container") + + # Use the use_env script to activate the environment + exit_code, stdout, stderr = self._exec_command(f"use_env {self.language}") + + if exit_code != 0: + logger.error(f"Failed to activate {self.language} environment: {stderr}") + else: + logger.info(f"Successfully activated {self.language} environment: {stdout}") + + except Exception as e: + logger.error(f"Error activating {self.language} environment: {str(e)}") + + def _exec_command(self, cmd: str) -> Tuple[int, str, str]: + """ + Execute a command in the container + + Args: + cmd: Command to execute + + Returns: + Tuple of (exit_code, stdout, stderr) + """ + if not self.active: + logger.error("Cannot execute command: Not connected to container") + return (1, "", "Not connected to container") + + try: + # Always wrap commands in 'bash -c' but ensure they're simple + shell_cmd = ['bash', '-c', cmd] + + logger.info(f"Running command: {cmd}") + + # Execute command in container with TTY disabled for proper output capture + exec_result = self.container.exec_run( + cmd=shell_cmd, + workdir=self.work_dir, + demux=True, # Split stdout and stderr + tty=False, # Disable TTY to ensure proper output capture + stream=False # Don't stream output + ) + + exit_code = exec_result.exit_code + + # Process stdout and stderr + stdout, stderr = "", "" + if isinstance(exec_result.output, tuple) and len(exec_result.output) == 2: + stdout_bytes, stderr_bytes = exec_result.output + if stdout_bytes: + stdout = stdout_bytes.decode('utf-8', errors='replace') + if stderr_bytes: + stderr = stderr_bytes.decode('utf-8', errors='replace') + + # Log the output + logger.info(f"Command exit code: {exit_code}") + logger.info(f"Command stdout: [{stdout}]") + logger.info(f"Command stderr: [{stderr}]") + + # Try alternate output capture method if output is empty + if not stdout and not stderr and exit_code == 0: + logger.info("No output captured with primary method, trying alternate method") + # Use output redirection to a file and then read it + output_file = f"/tmp/output_{int(time.time())}.txt" + + # Run the command and redirect output to file, then read file + alt_cmd1 = f"{cmd} > {output_file} 2>> {output_file}" + self.container.exec_run( + cmd=['bash', '-c', alt_cmd1], + workdir=self.work_dir + ) + + # Read the output file + alt_cmd2 = f"cat {output_file}" + alt_result = self.container.exec_run( + cmd=['bash', '-c', alt_cmd2], + workdir=self.work_dir + ) + + if alt_result.exit_code == 0 and alt_result.output: + stdout = alt_result.output.decode('utf-8', errors='replace') + logger.info(f"Alternate method stdout: [{stdout}]") + + # Clean up + alt_cmd3 = f"rm -f {output_file}" + self.container.exec_run( + cmd=['bash', '-c', alt_cmd3], + workdir=self.work_dir + ) + + return (exit_code, stdout, stderr) + + except Exception as e: + error_msg = f"Command execution failed: {str(e)}" + logger.error(error_msg, exc_info=True) + return (1, "", error_msg) + + async def write_file(self, filename: str, content: str) -> Dict[str, Any]: + """ + Write content to a file in the container + + Args: + filename: Name of the file to create/write + content: Content to write to the file + + Returns: + Status dictionary with success/error information + """ + if not self.active: + await self.connect() + + try: + # Create a temporary directory for the file + with tempfile.TemporaryDirectory() as temp_dir: + # Write content to a local file + temp_file_path = os.path.join(temp_dir, os.path.basename(filename)) + with open(temp_file_path, 'w', encoding='utf-8') as f: + f.write(content) + + # Read the file as binary + with open(temp_file_path, 'rb') as f: + data = f.read() + + # Create archive containing the file + import tarfile + import io + + # Create tar archive in memory + tar_stream = io.BytesIO() + with tarfile.open(fileobj=tar_stream, mode='w') as tar: + tarinfo = tarfile.TarInfo(name=os.path.basename(filename)) + tarinfo.size = len(data) + tar.addfile(tarinfo, io.BytesIO(data)) + + tar_stream.seek(0) + tar_data = tar_stream.read() + + # Create any necessary directories in the container + dir_name = os.path.dirname(filename) + if dir_name: + # Create directory if needed + logger.info(f"Creating directory in container: {os.path.join(self.work_dir, dir_name)}") + self._exec_command(f"mkdir -p {os.path.join(self.work_dir, dir_name)}") + + # Path where to extract the archive + extract_path = self.work_dir + if dir_name: + extract_path = os.path.join(self.work_dir, dir_name) + + logger.info(f"Extracting file to container path: {extract_path}") + + # Copy the tar archive to the container + result = self.container.put_archive(path=extract_path, data=tar_data) + + if not result: + error_msg = "Failed to copy file to container" + logger.error(error_msg) + return {"success": False, "error": error_msg} + + # Verify the file was created - construct full path for verification + full_path = os.path.join(self.work_dir, filename) + logger.info(f"Verifying file existence at: {full_path}") + check_cmd = f"test -f '{full_path}' && echo 'success' || echo 'not found'" + exit_code, stdout, stderr = self._exec_command(check_cmd) + + # List directory contents for debugging + ls_cmd = f"ls -la {os.path.dirname(full_path) or '.'}" + self._exec_command(ls_cmd) + + if "not found" in stdout: + error_msg = f"File verification failed: {full_path} not found" + logger.error(error_msg) + return {"success": False, "error": error_msg} + + # Add to files dictionary + self.files[filename] = { + "path": full_path, + "size": len(content), + "last_modified": time.time() + } + + logger.info(f"File {filename} written to container {self.container_name}") + return { + "success": True, + "filename": filename, + "size": len(content), + "message": f"File {filename} created successfully" + } + + except Exception as e: + error_msg = f"Failed to write file {filename}: {str(e)}" + logger.error(error_msg, exc_info=True) + return {"success": False, "error": error_msg} + + async def read_file(self, filename: str) -> Dict[str, Any]: + """ + Read content from a file in the container + + Args: + filename: Name of the file to read + + Returns: + Dictionary with file content and success status + """ + if not self.active: + return {"success": False, "error": "Not connected to container"} + + try: + # Ensure we're in the correct language environment + self._activate_language_environment() + + # Check if file exists using a shell-compatible command + exit_code, stdout, stderr = self._exec_command(f"test -f {filename} && echo 'exists' || echo 'not_exists'") + + if "not_exists" in stdout: + return {"success": False, "error": f"File {filename} not found"} + + # Read file content + exit_code, stdout, stderr = self._exec_command(f"cat {filename}") + + if exit_code != 0: + return {"success": False, "error": f"Failed to read file: {stderr}"} + + return { + "success": True, + "filename": filename, + "content": stdout, + "size": len(stdout) + } + + except Exception as e: + error_msg = f"Failed to read file {filename}: {str(e)}" + logger.error(error_msg, exc_info=True) + return {"success": False, "error": error_msg} + + async def delete_file(self, filename: str) -> Dict[str, Any]: + """ + Delete a file from the container + + Args: + filename: Name of the file to delete + + Returns: + Status dictionary with success/error information + """ + if not self.active: + return {"success": False, "error": "Not connected to container"} + + try: + # Ensure we're in the correct language environment + self._activate_language_environment() + + # Delete the file + exit_code, stdout, stderr = self._exec_command(f"rm -f {filename}") + + if exit_code != 0: + return {"success": False, "error": f"Failed to delete file: {stderr}"} + + # Remove from files dictionary + if filename in self.files: + del self.files[filename] + + return { + "success": True, + "filename": filename, + "message": f"File {filename} deleted successfully" + } + + except Exception as e: + error_msg = f"Failed to delete file {filename}: {str(e)}" + logger.error(error_msg, exc_info=True) + return {"success": False, "error": error_msg} + + async def list_files(self) -> Dict[str, Any]: + """ + List all files in the container's working directory + + Returns: + Dictionary with file listing and success status + """ + if not self.active: + return {"success": False, "error": "Not connected to container"} + + try: + # Ensure we're in the correct language environment + self._activate_language_environment() + + # List files - Using a simpler find command that works correctly + exit_code, stdout, stderr = self._exec_command(f"find '{self.work_dir}' -type f -not -path '*/\\.*'") + + if exit_code != 0: + return {"success": False, "error": f"Failed to list files: {stderr}"} + + # Process file list + file_list = [] + if stdout: + # Get more detailed info for each file + for file_path in stdout.strip().split('\n'): + if file_path: + # Get file information + name = os.path.basename(file_path) + file_list.append(name) + + # Update files dictionary + if name not in self.files: + self.files[name] = { + "path": file_path, + "last_modified": time.time() + } + + return { + "success": True, + "files": file_list, + "count": len(file_list) + } + + except Exception as e: + error_msg = f"Failed to list files: {str(e)}" + logger.error(error_msg, exc_info=True) + return {"success": False, "error": error_msg} + + async def execute_code(self, filename: str) -> Dict[str, Any]: + """ + Execute a file in the container + + Args: + filename: Name of the file to execute + + Returns: + Dictionary with execution results + """ + if not self.active: + await self.connect() + + try: + # Check if file exists using simple test command + exit_code, stdout, stderr = self._exec_command(f"test -f {filename}") + if exit_code != 0: + return {"success": False, "error": f"File {filename} not found"} + + # Ensure the correct language environment is activated + self._activate_language_environment() + + # Get execution command for this language + exec_cmd_generator = SUPPORTED_LANGUAGES[self.language]["execute_cmd"] + if not exec_cmd_generator: + return {"success": False, "error": f"No execution command defined for {self.language}"} + + # Special handling for C++ to separate compile and run steps + if self.language == "cpp": + logger.info(f"Compiling C++ file: {filename}") + + # First compile + compile_cmd = exec_cmd_generator(filename) + compile_exit_code, compile_stdout, compile_stderr = self._exec_command(compile_cmd) + + # If compilation failed, return the error + if compile_exit_code != 0: + return { + "execution_id": str(uuid.uuid4()), + "language": self.language, + "filename": filename, + "stdout": compile_stdout, + "stderr": compile_stderr, + "exit_code": compile_exit_code, + "success": False + } + + logger.info(f"C++ compilation successful, running: /tmp/program") + + # Then run the compiled program + run_cmd = "/tmp/program" + exit_code, stdout, stderr = self._exec_command(run_cmd) + else: + # For other languages, execute directly + if callable(exec_cmd_generator): + exec_cmd = exec_cmd_generator(filename) + else: + exec_cmd = f"{exec_cmd_generator} {filename}" + + logger.info(f"Executing {filename} with command: {exec_cmd}") + + # Execute command + exit_code, stdout, stderr = self._exec_command(exec_cmd) + + # If no output, try with explicit redirection to a file then read it + if exit_code == 0 and not stdout and not stderr: + logger.info("No output from direct execution, trying with file redirection") + output_file = f"/tmp/output_{uuid.uuid4().hex}.txt" + + if self.language == "cpp": + # For C++, redirect the compiled program output + redirect_cmd = f"/tmp/program > {output_file} 2>> {output_file}" + self._exec_command(redirect_cmd) + else: + # For other languages + if callable(exec_cmd_generator): + redirect_cmd = f"{exec_cmd} > {output_file} 2>> {output_file}" + self._exec_command(redirect_cmd) + + # Read the output file + cat_cmd = f"cat {output_file}" + cat_result = self._exec_command(cat_cmd) + if cat_result[0] == 0 and cat_result[1]: + stdout = cat_result[1] + + # Clean up + rm_cmd = f"rm -f {output_file}" + self._exec_command(rm_cmd) + + # Return execution results + execution_id = str(uuid.uuid4()) + result = { + "execution_id": execution_id, + "language": self.language, + "filename": filename, + "stdout": stdout, + "stderr": stderr, + "exit_code": exit_code, + "success": exit_code == 0 + } + + logger.info(f"Execution completed with status: {result['success']}") + return result + + except Exception as e: + error_msg = f"Execution failed: {str(e)}" + logger.error(error_msg, exc_info=True) + return {"success": False, "error": error_msg} + + async def disconnect(self) -> Dict[str, Any]: + """ + Disconnect from the container (does not stop it) + + Returns: + Status dictionary + """ + if not self.active: + return {"success": True, "message": "Already disconnected"} + + try: + self.active = False + self.container = None + logger.info(f"Disconnected from container {self.container_name}") + + return { + "success": True, + "message": f"Disconnected from container {self.container_name}" + } + + except Exception as e: + error_msg = f"Failed to disconnect: {str(e)}" + logger.error(error_msg, exc_info=True) + return {"success": False, "error": error_msg} + +# Global registry to track active Docker environments - indexed by language +docker_environments = {} + +def get_environment(language: str) -> DockerEnvironment: + """ + Get an existing Docker environment or create a new connection + + Args: + language: Programming language for this environment + + Returns: + DockerEnvironment instance + """ + global docker_environments + + # Normalize language name + language = language.lower().strip() + + # Map language aliases to standard names + language = LANGUAGE_ALIASES.get(language, language) + + # Check if language is supported + if language not in SUPPORTED_LANGUAGES: + logger.warning(f"Unsupported language: {language}, falling back to Python") + language = "python" + + # Get or create environment for this language + if language in docker_environments: + env = docker_environments[language] + logger.info(f"Reusing existing environment for {language}") + return env + + logger.info(f"Creating new environment for language: {language}") + env = DockerEnvironment(language=language) + docker_environments[language] = env + return env + +async def run_code(language: str, code: str) -> Dict: + """ + Execute code in a Docker container + + Args: + language: Programming language (python, java, cpp, etc.) + code: Source code to execute + + Returns: + Dictionary with execution results + """ + # Normalize language name + language = language.lower().strip() + + # Map language aliases to standard names + language = LANGUAGE_ALIASES.get(language, language) + + # Check if language is supported + if language not in SUPPORTED_LANGUAGES: + logger.warning(f"Unsupported language: {language}, falling back to Python") + language = "python" + + # Get Docker environment + env = get_environment(language) + + # Connect to container + if not env.active: + connect_result = await env.connect() + if not connect_result.get("success", False): + return {"error": connect_result.get("error", "Failed to connect to container")} + + # Explicitly activate the language environment + env._activate_language_environment() + + # Write code to a file with appropriate extension + extension = SUPPORTED_LANGUAGES[env.language]["file_extension"] + + # Special handling for Java - use class name as filename + if language.lower() == 'java': + # For Java, we need to use the class name as the filename + try: + # Look for the main class name + # This is a simple check for "public class X" without using regex + lines = code.split('\n') + class_name = None + for line in lines: + line = line.strip() + if line.startswith('public class '): + parts = line.split('public class ', 1)[1].split('{')[0].strip() + class_name = parts.split()[0].strip() + break + + if class_name: + filename = f"{class_name}{extension}" + logger.info(f"Using Java class name as filename: {filename}") + else: + filename = f"program{extension}" + logger.info(f"No Java class name found, using default filename: {filename}") + except Exception as e: + logger.error(f"Error extracting Java class name: {str(e)}") + filename = f"program{extension}" + else: + filename = f"program{extension}" + + write_result = await env.write_file(filename, code) + + if not write_result.get("success", False): + return {"error": write_result.get("error", "Failed to write code file")} + + # Execute the code + return await env.execute_code(filename) + +# Function to generate docker-compose config for language environments +def generate_docker_compose_config() -> str: + """ + Generate docker-compose configuration for all language environments + + Returns: + docker-compose.yml content for language environments + """ + # Start with version and services + config = """version: '3' + +services: +""" + + # Add each language environment + for language, info in SUPPORTED_LANGUAGES.items(): + if language == "python": + image = "python:3.11-slim" + setup_cmds = "pip install numpy pandas matplotlib" + elif language == "java": + image = "openjdk:17-slim" + setup_cmds = "apt-get update && apt-get install -y --no-install-recommends ca-certificates-java" + elif language == "cpp": + image = "gcc:11-bullseye" + setup_cmds = "apt-get update && apt-get install -y --no-install-recommends build-essential" + elif language == "javascript": + image = "node:18-slim" + setup_cmds = "npm install -g axios" + elif language == "typescript": + image = "node:18-slim" + setup_cmds = "npm install -g typescript axios" + elif language == "ruby": + image = "ruby:3.2-slim" + setup_cmds = "gem install bundler" + elif language == "go": + image = "golang:1.20-bullseye" + setup_cmds = "go get -u github.com/gorilla/mux" + elif language == "rust": + image = "rust:1.70-slim" + setup_cmds = "rustup component add rustfmt" + elif language == "php": + image = "php:8.2-cli" + setup_cmds = "apt-get update && apt-get install -y --no-install-recommends php-cli" + else: + continue # Skip unknown languages + + # Generate configuration for this language + container_name = info["container_name"] + + config += f""" {language}_env: + container_name: {container_name} + image: {image} + command: tail -f /dev/null + volumes: + - {language}_code:/app + working_dir: /app + restart: unless-stopped + deploy: + resources: + limits: + cpus: '0.5' + memory: 512M +""" + + # Add init script + setup_script = f"""echo "Setting up {language} environment..." +{setup_cmds} +echo "{language} environment ready!" +""" + config += f""" healthcheck: + test: ["CMD", "echo", "healthy"] + interval: 30s + timeout: 10s + retries: 3 + start_period: 10s + environment: + - SETUP_SCRIPT={setup_script} + +""" + + # Add volumes section + config += "\nvolumes:\n" + for language in SUPPORTED_LANGUAGES.keys(): + config += f" {language}_code:\n" + + return config + +# Save docker-compose configuration to a file +def save_docker_compose_config(output_path: str = "docker-compose.lang-env.yml") -> bool: + """ + Save docker-compose configuration for language environments to a file + + Args: + output_path: Path to save the configuration + + Returns: + True if successful, False otherwise + """ + try: + config = generate_docker_compose_config() + + with open(output_path, "w") as f: + f.write(config) + + logger.info(f"Docker Compose configuration saved to {output_path}") + logger.info(f"Run 'docker-compose -f {output_path} up -d' to start all language environments") + return True + + except Exception as e: + logger.error(f"Failed to save Docker Compose configuration: {str(e)}") + return False \ No newline at end of file diff --git a/cortex_on/utils/executors/__init__.py b/cortex_on/utils/executors/__init__.py deleted file mode 100644 index fd56669..0000000 --- a/cortex_on/utils/executors/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ - -from .local_code_executor import LocalCommandLineCodeExecutor - -__all__ = ["LocalCommandLineCodeExecutor"] diff --git a/cortex_on/utils/executors/executor_utils/__init__.py b/cortex_on/utils/executors/executor_utils/__init__.py deleted file mode 100644 index f1789a5..0000000 --- a/cortex_on/utils/executors/executor_utils/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -from ._base import CodeBlock, CodeExecutor, CodeResult -from ._func_with_reqs import ( - Alias, - FunctionWithRequirements, - FunctionWithRequirementsStr, - Import, - ImportFromModule, - with_requirements, -) - -__all__ = [ - "CodeBlock", - "CodeExecutor", - "CodeResult", - "Alias", - "ImportFromModule", - "Import", - "FunctionWithRequirements", - "FunctionWithRequirementsStr", - "with_requirements", -] diff --git a/cortex_on/utils/executors/executor_utils/_base.py b/cortex_on/utils/executors/executor_utils/_base.py deleted file mode 100644 index a9149a3..0000000 --- a/cortex_on/utils/executors/executor_utils/_base.py +++ /dev/null @@ -1,55 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import List, Protocol, runtime_checkable - -from utils import CancellationToken - -@dataclass -class CodeBlock: - """A code block extracted fromm an agent message.""" - - code: str - packages: List - language: str - human_input_or_command_line_args:str - -@dataclass -class CodeResult: - """Result of a code execution.""" - - exit_code: int - output: str - -@runtime_checkable -class CodeExecutor(Protocol): - """Executes code blocks and returns the result.""" - - async def execute_code_blocks( - self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken - ) -> CodeResult: - """Execute code blocks and return the result. - - This method should be implemented by the code executor. - - Args: - code_blocks (List[CodeBlock]): The code blocks to execute. - - Returns: - CodeResult: The result of the code execution. - - Raises: - ValueError: Errors in user inputs - asyncio.TimeoutError: Code execution timeouts - asyncio.CancelledError: CancellationToken evoked during execution - """ - ... - - async def restart(self) -> None: - """Restart the code executor. - - This method should be implemented by the code executor. - - This method is called when the agent is reset. - """ - ... diff --git a/cortex_on/utils/executors/executor_utils/_common.py b/cortex_on/utils/executors/executor_utils/_common.py deleted file mode 100644 index bf529c2..0000000 --- a/cortex_on/utils/executors/executor_utils/_common.py +++ /dev/null @@ -1,197 +0,0 @@ -import inspect -import re -from dataclasses import dataclass -from pathlib import Path -from textwrap import dedent, indent -from typing import Any, Callable, Optional, Sequence, Set, TypeVar, Union - -from ..executor_utils import ( - Alias, - CodeResult, - FunctionWithRequirements, - FunctionWithRequirementsStr, - Import, -) -from typing_extensions import ParamSpec - -@dataclass -class CommandLineCodeResult(CodeResult): - """A code result class for command line code executor.""" - - code_file: Optional[str] - -T = TypeVar("T") -P = ParamSpec("P") - -def _to_code( - func: Union[ - FunctionWithRequirements[T, P], Callable[P, T], FunctionWithRequirementsStr - ] -) -> str: - if isinstance(func, FunctionWithRequirementsStr): - return func.func - - code = inspect.getsource(func) - # Strip the decorator - if code.startswith("@"): - code = code[code.index("\n") + 1 :] - return code - -def _import_to_str(im: Import) -> str: - if isinstance(im, str): - return f"import {im}" - elif isinstance(im, Alias): - return f"import {im.name} as {im.alias}" - else: - - def to_str(i: Union[str, Alias]) -> str: - if isinstance(i, str): - return i - else: - return f"{i.name} as {i.alias}" - imports = ", ".join(map(to_str, im.imports)) - return f"from {im.module} import {imports}" - -def build_python_functions_file( - funcs: Sequence[ - Union[ - FunctionWithRequirements[Any, P], - Callable[..., Any], - FunctionWithRequirementsStr, - ] - ], -) -> str: - """:meta private:""" - # First collect all global imports - global_imports: Set[Import] = set() - for func in funcs: - if isinstance(func, (FunctionWithRequirements, FunctionWithRequirementsStr)): - global_imports.update(func.global_imports) - - content = "\n".join(map(_import_to_str, global_imports)) + "\n\n" - - for func in funcs: - content += _to_code(func) + "\n\n" - - return content - -def to_stub(func: Union[Callable[..., Any], FunctionWithRequirementsStr]) -> str: - """Generate a stub for a function as a string - - Args: - func (Callable[..., Any]): The function to generate a stub for - - Returns: - str: The stub for the function - """ - if isinstance(func, FunctionWithRequirementsStr): - return to_stub(func.compiled_func) - - content = f"def {func.__name__}{inspect.signature(func)}:\n" - docstring = func.__doc__ - - if docstring: - docstring = dedent(docstring) - docstring = '"""' + docstring + '"""' - docstring = indent(docstring, " ") - content += docstring + "\n" - - content += " ..." - return content - -# Raises ValueError if the file is not in the workspace -def get_file_name_from_content(code: str, workspace_path: Path) -> Optional[str]: - first_line = code.split("\n")[0] - # TODO - support other languages - if first_line.startswith("# filename:"): - filename = first_line.split(":")[1].strip() - - # Handle relative paths in the filename - path = Path(filename) - if not path.is_absolute(): - path = workspace_path / path - path = path.resolve() - # Throws an error if the file is not in the workspace - relative = path.relative_to(workspace_path.resolve()) - return str(relative) - - return None - -def silence_pip(code: str, lang: str) -> str: - """Apply -qqq flag to pip install commands.""" - if lang == "python": - regex = r"^pip install" - elif lang in ["bash", "shell", "sh", "pwsh", "powershell", "ps1"]: - regex = r"^pip install" - else: - return code - - # Find lines that start with pip install and make sure "-qqq" flag is added. - lines = code.split("\n") - for i, line in enumerate(lines): - # use regex to find lines that start with pip install. - match = re.search(regex, line) - if match is not None: - # print(line) - if "-qqq" not in line: - lines[i] = line.replace(match.group(0), match.group(0) + " -qqq") - return "\n".join(lines) - -def get_required_packages(code: str, lang: str) -> set[str]: - ret: set[str] = set() - if lang == "python": - regex = r"^! ?pip install(.*)$" - else: - return ret - - # Find lines that start with pip install and make sure "-qqq" flag is added. - lines = code.split("\n") - for _, line in enumerate(lines): - # use regex to find lines that start with pip install. - match = re.search(regex, line) - if match is not None: - reqs = match.group(1).split(",") - ret = {req.strip(" ") for req in reqs} - return ret - -PYTHON_VARIANTS = ["python", "Python", "py"] - -def lang_to_cmd(lang: str) -> str: - if lang in PYTHON_VARIANTS: - return "python" - if lang.startswith("python") or lang in ["bash", "sh"]: - return lang - if lang in ["shell"]: - return "sh" - else: - raise ValueError(f"Unsupported language: {lang}") - -# Regular expression for finding a code block -# ```[ \t]*(\w+)?[ \t]*\r?\n(.*?)[ \t]*\r?\n``` Matches multi-line code blocks. -# The [ \t]* matches the potential spaces before language name. -# The (\w+)? matches the language, where the ? indicates it is optional. -# The [ \t]* matches the potential spaces (not newlines) after language name. -# The \r?\n makes sure there is a linebreak after ```. -# The (.*?) matches the code itself (non-greedy). -# The \r?\n makes sure there is a linebreak before ```. -# The [ \t]* matches the potential spaces before closing ``` (the spec allows indentation). -CODE_BLOCK_PATTERN = r"```[ \t]*(\w+)?[ \t]*\r?\n(.*?)\r?\n[ \t]*```" - -def infer_lang(code: str) -> str: - """infer the language for the code. - TODO: make it robust. - """ - if ( - code.startswith("python ") - or code.startswith("pip") - or code.startswith("python3 ") - ): - return "sh" - - # check if code is a valid python code - try: - compile(code, "test", "exec") - return "python" - except SyntaxError: - # not a valid python code - return "unknown" diff --git a/cortex_on/utils/executors/executor_utils/_func_with_reqs.py b/cortex_on/utils/executors/executor_utils/_func_with_reqs.py deleted file mode 100644 index 2df1e0d..0000000 --- a/cortex_on/utils/executors/executor_utils/_func_with_reqs.py +++ /dev/null @@ -1,225 +0,0 @@ -from __future__ import annotations - -import functools -import inspect -from dataclasses import dataclass, field -from importlib.abc import SourceLoader -from importlib.util import module_from_spec, spec_from_loader -from textwrap import dedent, indent -from typing import Any, Callable, Generic, List, Sequence, Set, Tuple, TypeVar, Union - -from typing_extensions import ParamSpec - -T = TypeVar("T") -P = ParamSpec("P") - -def _to_code( - func: Union[ - FunctionWithRequirements[T, P], Callable[P, T], FunctionWithRequirementsStr - ] -) -> str: - if isinstance(func, FunctionWithRequirementsStr): - return func.func - - if isinstance(func, FunctionWithRequirements): - code = inspect.getsource(func.func) - else: - code = inspect.getsource(func) - # Strip the decorator - if code.startswith("@"): - code = code[code.index("\n") + 1 :] - return code - -@dataclass(frozen=True) -class Alias: - name: str - alias: str - -@dataclass(frozen=True) -class ImportFromModule: - module: str - imports: Tuple[Union[str, Alias], ...] - - ## backward compatibility - def __init__( - self, - module: str, - imports: Union[Tuple[Union[str, Alias], ...], List[Union[str, Alias]]], - ): - object.__setattr__(self, "module", module) - if isinstance(imports, list): - object.__setattr__(self, "imports", tuple(imports)) - else: - object.__setattr__(self, "imports", imports) - -Import = Union[str, ImportFromModule, Alias] - -def _import_to_str(im: Import) -> str: - if isinstance(im, str): - return f"import {im}" - elif isinstance(im, Alias): - return f"import {im.name} as {im.alias}" - else: - - def to_str(i: Union[str, Alias]) -> str: - if isinstance(i, str): - return i - else: - return f"{i.name} as {i.alias}" - imports = ", ".join(map(to_str, im.imports)) - return f"from {im.module} import {imports}" - -class _StringLoader(SourceLoader): - def __init__(self, data: str): - self.data = data - - def get_source(self, fullname: str) -> str: - return self.data - - def get_data(self, path: str) -> bytes: - return self.data.encode("utf-8") - - def get_filename(self, fullname: str) -> str: - return "/" + fullname + ".py" - -@dataclass -class FunctionWithRequirementsStr: - func: str - compiled_func: Callable[..., Any] - _func_name: str - python_packages: Sequence[str] = field(default_factory=list) - global_imports: Sequence[Import] = field(default_factory=list) - - def __init__( - self, - func: str, - python_packages: Sequence[str] = [], - global_imports: Sequence[Import] = [], - ): - self.func = func - self.python_packages = python_packages - self.global_imports = global_imports - - module_name = "func_module" - loader = _StringLoader(func) - spec = spec_from_loader(module_name, loader) - if spec is None: - raise ValueError("Could not create spec") - module = module_from_spec(spec) - if spec.loader is None: - raise ValueError("Could not create loader") - - try: - spec.loader.exec_module(module) - except Exception as e: - raise ValueError(f"Could not compile function: {e}") from e - - functions = inspect.getmembers(module, inspect.isfunction) - if len(functions) != 1: - raise ValueError("The string must contain exactly one function") - - self._func_name, self.compiled_func = functions[0] - - def __call__(self, *args: Any, **kwargs: Any) -> None: - raise NotImplementedError( - "String based function with requirement objects are not directly callable" - ) - -@dataclass -class FunctionWithRequirements(Generic[T, P]): - func: Callable[P, T] - python_packages: Sequence[str] = field(default_factory=list) - global_imports: Sequence[Import] = field(default_factory=list) - - @classmethod - def from_callable( - cls, - func: Callable[P, T], - python_packages: Sequence[str] = [], - global_imports: Sequence[Import] = [], - ) -> FunctionWithRequirements[T, P]: - return cls( - python_packages=python_packages, global_imports=global_imports, func=func - ) - - @staticmethod - def from_str( - func: str, - python_packages: Sequence[str] = [], - global_imports: Sequence[Import] = [], - ) -> FunctionWithRequirementsStr: - return FunctionWithRequirementsStr( - func=func, python_packages=python_packages, global_imports=global_imports - ) - # Type this based on F - def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T: - return self.func(*args, **kwargs) - -def with_requirements( - python_packages: Sequence[str] = [], global_imports: Sequence[Import] = [] -) -> Callable[[Callable[P, T]], FunctionWithRequirements[T, P]]: - """Decorate a function with package and import requirements - - Args: - python_packages (List[str], optional): Packages required to function. Can include version info.. Defaults to []. - global_imports (List[Import], optional): Required imports. Defaults to []. - - Returns: - Callable[[Callable[P, T]], FunctionWithRequirements[T, P]]: The decorated function - """ - - def wrapper(func: Callable[P, T]) -> FunctionWithRequirements[T, P]: - func_with_reqs = FunctionWithRequirements( - python_packages=python_packages, global_imports=global_imports, func=func - ) - - functools.update_wrapper(func_with_reqs, func) - return func_with_reqs - return wrapper - -def build_python_functions_file( - funcs: Sequence[ - Union[ - FunctionWithRequirements[Any, P], - Callable[..., Any], - FunctionWithRequirementsStr, - ] - ], -) -> str: - """:meta private:""" - # First collect all global imports - global_imports: Set[Import] = set() - for func in funcs: - if isinstance(func, (FunctionWithRequirements, FunctionWithRequirementsStr)): - global_imports.update(func.global_imports) - - content = "\n".join(map(_import_to_str, global_imports)) + "\n\n" - - for func in funcs: - content += _to_code(func) + "\n\n" - - return content - -def to_stub(func: Union[Callable[..., Any], FunctionWithRequirementsStr]) -> str: - """Generate a stub for a function as a string - - Args: - func (Callable[..., Any]): The function to generate a stub for - - Returns: - str: The stub for the function - """ - if isinstance(func, FunctionWithRequirementsStr): - return to_stub(func.compiled_func) - - content = f"def {func.__name__}{inspect.signature(func)}:\n" - docstring = func.__doc__ - - if docstring: - docstring = dedent(docstring) - docstring = '"""' + docstring + '"""' - docstring = indent(docstring, " ") - content += docstring + "\n" - - content += " ..." - return content diff --git a/cortex_on/utils/executors/executor_utils/extract_command_line_args.py b/cortex_on/utils/executors/executor_utils/extract_command_line_args.py deleted file mode 100644 index 671f359..0000000 --- a/cortex_on/utils/executors/executor_utils/extract_command_line_args.py +++ /dev/null @@ -1,19 +0,0 @@ -import re - -def extract_command_line_args(lang, filename, human_input_or_command_line_args): - human_input_or_command_line_args = " ".join(human_input_or_command_line_args).strip() - - extension = filename.split('.')[-1] if '.' in filename else 'py' if lang.startswith('python') else lang - - # Define prefixes to remove - prefixes = [f"{lang} {filename}", f"{lang}", f"{filename}"] - for prefix in prefixes: - if human_input_or_command_line_args.startswith(prefix): - human_input_or_command_line_args = human_input_or_command_line_args[len(prefix):].strip() - break - - # Split into arguments and filter out matches of *.extension - args = human_input_or_command_line_args.split() - args = [arg for arg in args if not re.fullmatch(rf".*\.{extension}", arg)] - - return args \ No newline at end of file diff --git a/cortex_on/utils/executors/local_code_executor.py b/cortex_on/utils/executors/local_code_executor.py deleted file mode 100644 index 099ea72..0000000 --- a/cortex_on/utils/executors/local_code_executor.py +++ /dev/null @@ -1,530 +0,0 @@ -import asyncio -import logging -import os -import sys -import warnings -from hashlib import sha256 -from pathlib import Path -from string import Template -from types import SimpleNamespace -from typing import Any, Callable, ClassVar, List, Optional, Sequence, Union -import venv -from utils import CancellationToken -from .executor_utils import ( - CodeBlock, - CodeExecutor, - FunctionWithRequirements, - FunctionWithRequirementsStr, -) -from typing_extensions import ParamSpec -from dataclasses import asdict -import json -from utils.stream_response_format import StreamResponse -from fastapi import WebSocket -from .executor_utils._common import ( - PYTHON_VARIANTS, - CommandLineCodeResult, - build_python_functions_file, - get_file_name_from_content, - lang_to_cmd, - silence_pip, - to_stub, -) -from utils.executors.executor_utils.extract_command_line_args import extract_command_line_args - -__all__ = ("LocalCommandLineCodeExecutor",) - -A = ParamSpec("A") - -class LocalCommandLineCodeExecutor(CodeExecutor): - """A code executor class that executes code through a local command line - environment. - - .. danger:: - - This will execute code on the local machine. If being used with LLM generated code, caution should be used. - - Each code block is saved as a file and executed in a separate process in - the working directory, and a unique file is generated and saved in the - working directory for each code block. - The code blocks are executed in the order they are received. - Command line code is sanitized using regular expression match against a list of dangerous commands in order to prevent self-destructive - commands from being executed which may potentially affect the users environment. - Currently the only supported languages is Python and shell scripts. - For Python code, use the language "python" for the code block. - For shell scripts, use the language "bash", "shell", or "sh" for the code - block. - - Args: - timeout (int): The timeout for the execution of any single code block. Default is 60. - work_dir (str): The working directory for the code execution. If None, - a default working directory will be used. The default working - directory is the current directory ".". - functions (List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]]): A list of functions that are available to the code executor. Default is an empty list. - functions_module (str, optional): The name of the module that will be created to store the functions. Defaults to "functions". - virtual_env_context (Optional[SimpleNamespace], optional): The virtual environment context. Defaults to None. - - Example: - - How to use `LocalCommandLineCodeExecutor` with a virtual environment different from the one used to run the application: - Set up a virtual environment using the `venv` module, and pass its context to the initializer of `LocalCommandLineCodeExecutor`. This way, the executor will run code within the new environment. - - .. code-block:: python - - import venv - from pathlib import Path - import asyncio - - - async def example(): - work_dir = Path("coding") - work_dir.mkdir(exist_ok=True) - - venv_dir = work_dir / ".venv" - venv_builder = venv.EnvBuilder(with_pip=True) - venv_builder.create(venv_dir) - venv_context = venv_builder.ensure_directories(venv_dir) - - local_executor = LocalCommandLineCodeExecutor(work_dir=work_dir, virtual_env_context=venv_context) - await local_executor.execute_code_blocks( - code_blocks=[ - CodeBlock(language="bash", code="pip install matplotlib"), - ], - cancellation_token=CancellationToken(), - ) - - - asyncio.run(example()) - - """ - - SUPPORTED_LANGUAGES: ClassVar[List[str]] = [ - "bash", - "shell", - "sh", - "pwsh", - "powershell", - "ps1", - "python", - ] - FUNCTION_PROMPT_TEMPLATE: ClassVar[ - str - ] = """You have access to the following user defined functions. They can be accessed from the module called `$module_name` by their function names. - -For example, if there was a function called `foo` you could import it by writing `from $module_name import foo` - -$functions""" - - def __init__( - self, - timeout: int = 60, - work_dir: Union[Path, str] = Path("./code_files"), - functions: Sequence[ - Union[ - FunctionWithRequirements[Any, A], - Callable[..., Any], - FunctionWithRequirementsStr, - ] - ] = [], - functions_module: str = "functions", - virtual_env_context: Optional[SimpleNamespace] = None, - ): - if timeout < 1: - raise ValueError("Timeout must be greater than or equal to 1.") - - if isinstance(work_dir, str): - work_dir = Path(work_dir) - - if not functions_module.isidentifier(): - raise ValueError("Module name must be a valid Python identifier") - - self._functions_module = functions_module - - work_dir.mkdir(exist_ok=True) - - self._timeout = timeout - self._work_dir: Path = work_dir - print("functions in init", functions) - self._functions = functions - # Setup could take some time so we intentionally wait for the first code block to do it. - # if len(functions) > 0: - self._setup_functions_complete = False - # else: - # self._setup_functions_complete = True - # if(virtual_env_context==None): - # self._virtual_env_context: Optional[SimpleNamespace] = self.create_venv(work_dir) - # else: - self._virtual_env_context: Optional[SimpleNamespace] = virtual_env_context - self.websocket:Optional[WebSocket]= None - self.stream_output:Optional[StreamResponse] = None - - def format_functions_for_prompt( - self, prompt_template: str = FUNCTION_PROMPT_TEMPLATE - ) -> str: - """(Experimental) Format the functions for a prompt. - - The template includes two variables: - - `$module_name`: The module name. - - `$functions`: The functions formatted as stubs with two newlines between each function. - - Args: - prompt_template (str): The prompt template. Default is the class default. - - Returns: - str: The formatted prompt. - """ - - template = Template(prompt_template) - return template.substitute( - module_name=self._functions_module, - functions="\n\n".join([to_stub(func) for func in self._functions]), - ) - - @property - def functions_module(self) -> str: - """(Experimental) The module name for the functions.""" - return self._functions_module - - @property - def functions(self) -> List[str]: - raise NotImplementedError - - @property - def timeout(self) -> int: - """(Experimental) The timeout for code execution.""" - return self._timeout - - @property - def work_dir(self) -> Path: - """(Experimental) The working directory for the code execution.""" - return self._work_dir - - async def create_venv(self, work_dir): - if self.stream_output and self.websocket: - self.stream_output.steps.append( - "Creating a secure environment for the code to be executed" - ) - await self.websocket.send_text( - json.dumps(asdict(self.stream_output)) - ) - - venv_dir = work_dir / ".venv" - venv_builder = venv.EnvBuilder(with_pip=True) - venv_builder.create(venv_dir) - venv_context = venv_builder.ensure_directories(venv_dir) - print("created venv") - return venv_context - - async def _setup_functions( - self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken - ) -> None: - print("functions", self._functions) - print("code block", code_blocks) - required_packages = code_blocks[0].packages - print("required packages", required_packages) - if len(required_packages) > 0: - log="Ensuring packages are installed in executor." - logging.info(log) - if self.stream_output and self.websocket: - self.stream_output.steps.append( - log - ) - await self.websocket.send_text( - json.dumps(asdict(self.stream_output)) - ) - - cmd_args = ["-m", "pip", "install"] - cmd_args.extend(required_packages) - print("cmd args", cmd_args) - if self._virtual_env_context: - py_executable = self._virtual_env_context.env_exe - print("py executable already initialized", py_executable) - - else: - self._virtual_env_context = await self.create_venv(self.work_dir) - py_executable = self._virtual_env_context.env_exe - print("py executable initialized", py_executable) - - # py_executable = sys.executable - - task = asyncio.create_task( - asyncio.create_subprocess_exec( - py_executable, - *cmd_args, - cwd=Path("./"), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - ) - print("task created", task) - cancellation_token.link_future(task) - proc = None - try: - if self.stream_output and self.websocket: - self.stream_output.steps.append( - "Installing the code dependencies in your local environment before the code execution" - ) - await self.websocket.send_text( - json.dumps(asdict(self.stream_output)) - ) - proc = await task - stdout, stderr = await asyncio.wait_for( - proc.communicate(), self._timeout - ) - print("task completed") - except asyncio.TimeoutError as e: - raise ValueError("Pip install timed out") from e - except asyncio.CancelledError as e: - raise ValueError("Pip install was cancelled") from e - except Exception as e: - print("error", e) - if proc.returncode is not None and proc.returncode != 0: - raise ValueError( - f"Pip install failed. {stdout.decode()}, {stderr.decode()}" - ) - - # Attempt to load the function file to check for syntax errors, imports etc. - # exec_result = await self._execute_code_dont_check_setup( - # [CodeBlock(code=func_file_content, language="python")], cancellation_token - # ) - # exec_result = await self._execute_code_dont_check_setup( - # code_blocks, cancellation_token - # ) - - # if exec_result.exit_code != 0: - # raise ValueError(f"Functions failed to load: {exec_result.output}") - - self._setup_functions_complete = True - - async def execute_code_blocks( - self, code_blocks: List[CodeBlock],websocket:WebSocket,stream_output:StreamResponse, cancellation_token: CancellationToken - ) -> CommandLineCodeResult: - """(Experimental) Execute the code blocks and return the result. - - Args: - code_blocks (List[CodeBlock]): The code blocks to execute. - cancellation_token (CancellationToken): a token to cancel the operation - - Returns: - CommandLineCodeResult: The result of the code execution.""" - - self.websocket=websocket - self.stream_output=stream_output - if not self._setup_functions_complete: - print("setting up functions") - await self._setup_functions(code_blocks, cancellation_token) - return await self._execute_code_dont_check_setup( - code_blocks, cancellation_token - ) - - async def _execute_code_dont_check_setup( - self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken - ) -> CommandLineCodeResult: - logs_all: str = "" - file_names: List[Path] = [] - exitcode = 0 - for code_block in code_blocks: - lang, code, packages,human_input_or_command_line_args = ( - code_block.language, - code_block.code, - code_block.packages, - code_block.human_input_or_command_line_args - ) - lang = lang.lower() - - code = silence_pip(code, lang) - - if lang in PYTHON_VARIANTS: - lang = "python" - - if lang not in self.SUPPORTED_LANGUAGES: - # In case the language is not supported, we return an error message. - exitcode = 1 - logs_all += "\n" + f"unknown language {lang}" - break - - try: - # Check if there is a filename comment - filename = get_file_name_from_content(code, self._work_dir) - except ValueError: - return CommandLineCodeResult( - exit_code=1, - output="Filename is not in the workspace", - code_file=None, - ) - if self.stream_output and self.websocket: - self.stream_output.steps.append( - f"Saving the code in a file under the directory: {self._work_dir}" - ) - await self.websocket.send_text( - json.dumps(asdict(self.stream_output)) - ) - if filename is None: - # create a file with an automatically generated name - code_hash = sha256(code.encode()).hexdigest() - filename = f"tmp_code_{code_hash}.{'py' if lang.startswith('python') else lang}" - - command_line_args = extract_command_line_args(lang, filename, human_input_or_command_line_args) - print("extracted command_line_args", command_line_args) - - written_file = (self._work_dir / filename).resolve() - with written_file.open("w", encoding="utf-8") as f: - f.write(code) - file_names.append(written_file) - - env = os.environ.copy() - - if self._virtual_env_context: - virtual_env_exe_abs_path = os.path.abspath( - self._virtual_env_context.env_exe - ) - virtual_env_bin_abs_path = os.path.abspath( - self._virtual_env_context.bin_path - ) - env["PATH"] = f"{virtual_env_bin_abs_path}{os.pathsep}{env['PATH']}" - program = ( - virtual_env_exe_abs_path - if lang.startswith("python") - else lang_to_cmd(lang) - ) - print("program", program) - else: - program = ( - sys.executable if lang.startswith("python") else lang_to_cmd(lang) - ) - - # Wrap in a task to make it cancellable - - # if(lang.startswith("python") and len(packages)!=0): - # process=await asyncio.create_subprocess_exec('pip','install',*packages,stdout=asyncio.subprocess.PIPE, - # stderr=asyncio.subprocess.PIPE) - # stdout,stderr = await process.communicate() - - # if process.returncode==0: - # print("packages installed successfully") - # else: - # print("error installing packages") - task = asyncio.create_task( - asyncio.create_subprocess_exec( - program, - str(written_file.absolute()), - *command_line_args, - cwd=self._work_dir, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - stdin=asyncio.subprocess.PIPE, - env=env, - ) - ) - cancellation_token.link_future(task) - if self.stream_output and self.websocket: - self.stream_output.steps.append( - "Executing the generated code in your safe environment" - ) - await self.websocket.send_text( - json.dumps(asdict(self.stream_output)) - ) - proc = await task - - if(len(command_line_args) == 0): - try: - stdout, stderr = await asyncio.wait_for( - proc.communicate(b""), self._timeout - ) - logs_all += stderr.decode() - logs_all += stdout.decode() - except asyncio.TimeoutError: - logs_all += "\n Timeout" - exitcode = 124 # Exit code for timeout - except asyncio.CancelledError: - logs_all += "\n Cancelled" - exitcode = 125 # Exit code for operation canceled - except Exception as e: - logs_all += f"\n Error: {e}" - exitcode = 1 # Generic error code - elif(len(command_line_args) == 1): - try: - stdout, stderr = await asyncio.wait_for( - proc.communicate(command_line_args[0].encode()), self._timeout - ) - logs_all += stderr.decode() - logs_all += stdout.decode() - except asyncio.TimeoutError: - logs_all += "\n Timeout" - exitcode = 124 # Exit code for timeout - except asyncio.CancelledError: - logs_all += "\n Cancelled" - exitcode = 125 # Exit code for operation canceled - except Exception as e: - logs_all += f"\n Error: {e}" - exitcode = 1 # Generic error code - else: - for index, cmd_arg in enumerate(command_line_args): - try: - # Send the input to the subprocess - proc.stdin.write(f"{cmd_arg}\n".encode()) - await proc.stdin.drain() # Ensure the input is sent - - timeout = self._timeout - if index != len(command_line_args) - 1: - timeout = 5 - - # Read the output (if any) - stdout = await asyncio.wait_for(proc.stdout.readline(), timeout) - stderr = await asyncio.wait_for(proc.stderr.readline(), timeout) - - logs_all += stderr.decode() - logs_all += stdout.decode() - except asyncio.TimeoutError: - if(index == len(command_line_args) - 1): - logs_all += "\n Timeout" - exitcode = 124 # Exit code for timeout - break - except asyncio.CancelledError: - logs_all += "\n Cancelled" - exitcode = 125 # Exit code for operation canceled - break - except ConnectionResetError: # No human input needed, command line args were needed - pass - except Exception as e: - logs_all += f"\n Error: {e}" - exitcode = 1 # Generic error code - break - - try: - stdout, stderr = await asyncio.wait_for( - proc.communicate(b""), self._timeout - ) - logs_all += stderr.decode() - logs_all += stdout.decode() - except asyncio.TimeoutError: - logs_all += "\n Timeout" - exitcode = 124 # Exit code for timeout - except asyncio.CancelledError: - logs_all += "\n Cancelled" - exitcode = 125 # Exit code for operation canceled - except Exception as e: - logs_all += f"\n Error: {e}" - exitcode = 1 # Generic error code - - print("exit code", exitcode) - print("logs all", logs_all) - - self._running_cmd_task = None - proc.stdin.close() - await proc.wait() - exitcode = proc.returncode or exitcode - - if exitcode != 0: - break - code_file = str(file_names[0]) if len(file_names) > 0 else None - return CommandLineCodeResult( - exit_code=exitcode, output=logs_all, code_file=code_file - ) - - async def restart(self) -> None: - """(Experimental) Restart the code executor.""" - warnings.warn( - "Restarting local command line code executor is not supported. No action is taken.", - stacklevel=2, - ) diff --git a/cortex_on/utils/stream_response_format.py b/cortex_on/utils/stream_response_format.py index d99ac9a..b7e986a 100644 --- a/cortex_on/utils/stream_response_format.py +++ b/cortex_on/utils/stream_response_format.py @@ -1,5 +1,5 @@ from dataclasses import dataclass -from typing import List, Optional +from typing import List, Optional, Dict @dataclass class StreamResponse: @@ -9,3 +9,5 @@ class StreamResponse: status_code: int output: str live_url: Optional[str] = None + source_code: Optional[str] = None + metadata: Optional[Dict] = None diff --git a/docker-compose.yaml b/docker-compose.yaml index e9f951a..7ae98b4 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -7,10 +7,14 @@ services: dockerfile: Dockerfile volumes: - ./cortex_on:/app + - /var/run/docker.sock:/var/run/docker.sock env_file: - .env restart: always network_mode: host + privileged: true + depends_on: + - multi_language_env agentic_browser: build: @@ -37,3 +41,31 @@ services: - agentic_browser restart: always network_mode: host + + # Multi-language environment container + multi_language_env: + container_name: cortexon_multi_env + build: + context: ./cortex_on/multi_lang_env + dockerfile: Dockerfile + volumes: + - multi_language_code:/app + - ./cortex_on/multi_lang_env/setup:/setup + working_dir: /app + restart: unless-stopped + network_mode: host + deploy: + resources: + limits: + cpus: '1.0' + memory: 2G + healthcheck: + test: ["CMD", "echo", "healthy"] + interval: 30s + timeout: 10s + retries: 3 + start_period: 60s + command: tail -f /dev/null + +volumes: + multi_language_code: diff --git a/frontend/src/components/home/ChatList.tsx b/frontend/src/components/home/ChatList.tsx index 9d8cb21..4e89d8c 100644 --- a/frontend/src/components/home/ChatList.tsx +++ b/frontend/src/components/home/ChatList.tsx @@ -133,7 +133,7 @@ const ChatList = ({isLoading, setIsLoading}: ChatListPageProps) => { const {agent_name, instructions, steps, output, status_code, live_url} = lastJsonMessage as SystemMessage; - console.log(lastJsonMessage); + console.log("Received message:", lastJsonMessage); if (live_url && liveUrl.length === 0) { setCurrentOutput(null); @@ -195,34 +195,30 @@ const ChatList = ({isLoading, setIsLoading}: ChatListPageProps) => { setIsLoading(false); } - if (status_code === 200) { - setOutputsList((prevList) => { - const existingIndex = prevList.findIndex( - (item) => item.agent === agent_name - ); - - let newList; - let newOutputIndex; - - if (existingIndex >= 0) { - newList = [...prevList]; - newList[existingIndex] = {agent: agent_name, output}; - newOutputIndex = existingIndex; - } else { - newList = [...prevList, {agent: agent_name, output}]; - newOutputIndex = newList.length - 1; - } - - setAnimateOutputEntry(false); - - setTimeout(() => { - setCurrentOutput(newOutputIndex); - setAnimateOutputEntry(true); - }, 300); - - return newList; - }); - } + // Update outputs list and show the output immediately + setOutputsList((prevList) => { + const existingIndex = prevList.findIndex( + (item) => item.agent === agent_name + ); + + let newList; + let newOutputIndex; + + if (existingIndex >= 0) { + newList = [...prevList]; + newList[existingIndex] = {agent: agent_name, output}; + newOutputIndex = existingIndex; + } else { + newList = [...prevList, {agent: agent_name, output}]; + newOutputIndex = newList.length - 1; + } + + // Immediately show the output + setCurrentOutput(newOutputIndex); + setAnimateOutputEntry(true); + + return newList; + }); } if (agent_name === "Human Input") { @@ -231,9 +227,7 @@ const ChatList = ({isLoading, setIsLoading}: ChatListPageProps) => { } else { setIsLoading(false); } - setTimeout(() => { - setCurrentOutput(null); - }, 300); + setCurrentOutput(null); } const updatedMessages = [ @@ -504,7 +498,7 @@ const ChatList = ({isLoading, setIsLoading}: ChatListPageProps) => { const chatContainerWidth = liveUrl || currentOutput !== null ? "50%" : "65%"; const outputPanelClasses = `border-2 rounded-xl w-[50%] flex flex-col h-[95%] justify-between items-center transition-all duration-700 ease-in-out ${ - animateOutputEntry + animateOutputEntry && currentOutput !== null ? "opacity-100 translate-x-0 animate-fade-in animate-once animate-duration-1000" : "opacity-0 translate-x-2" }`; diff --git a/frontend/src/components/home/CodeBlock.tsx b/frontend/src/components/home/CodeBlock.tsx index 3f1d619..0108dc1 100644 --- a/frontend/src/components/home/CodeBlock.tsx +++ b/frontend/src/components/home/CodeBlock.tsx @@ -1,22 +1,32 @@ import {useRef, useState} from "react"; import Markdown from "react-markdown"; import {Prism as SyntaxHighlighter} from "react-syntax-highlighter"; -import {tomorrow} from "react-syntax-highlighter/dist/esm/styles/prism"; +import {vscDarkPlus} from "react-syntax-highlighter/dist/esm/styles/prism"; import rehypeRaw from "rehype-raw"; import remarkBreaks from "remark-breaks"; export const CodeBlock = ({content}: {content: string}) => { - const codeBlock = content.includes("content='") - ? content.split("content='")[1] - : content; - const [isCopied, setIsCopied] = useState(false); const codeRef = useRef(null); const handleCopyClick = () => { if (codeRef.current) { + // Find all code blocks and join their text content + const codeElements = codeRef.current.querySelectorAll('pre code'); + let textToCopy = ''; + + if (codeElements.length > 0) { + // Get text from syntax highlighted blocks + codeElements.forEach(el => { + textToCopy += el.textContent + '\n\n'; + }); + } else { + // Fallback to all text content + textToCopy = codeRef.current.innerText; + } + navigator.clipboard - .writeText(codeRef.current.innerText) + .writeText(textToCopy.trim()) .then(() => { setIsCopied(true); // Add a visual pulse effect @@ -34,17 +44,32 @@ export const CodeBlock = ({content}: {content: string}) => { } }; + // Process the content to handle any special formatting + const processedContent = content.replace(/\\n/g, "\n"); + return ( -
+
-
+
{ {String(children).replace(/\n$/, "")} ) : ( - + {children} ); }, + // Add heading styles + h1: ({children}) => ( +

{children}

+ ), + h2: ({children}) => ( +

+ {children} +

+ ), + h3: ({children}) => ( +

{children}

+ ), + // Add paragraph styles + p: ({children}) => ( +

{children}

+ ), + // Style output code blocks + pre: ({children}) => ( + 
{children}
+ ), + // Style emojis and status indicators + strong: ({children}) => { + const text = String(children); + if (text.includes("✅")) { + return {children}; + } else if (text.includes("❌")) { + return {children}; + } + return {children}; + }, + // Style links + a: ({children, href}) => ( + + {children} + + ), + // Style lists + ul: ({children}) => ( +
    {children}
+ ), + ol: ({children}) => ( +
    {children}
+ ), + li: ({children}) => ( +
  • {children}
    • + ), }} />
    diff --git a/ta-browser/core/orchestrator.py b/ta-browser/core/orchestrator.py index 9dbf16e..f112130 100644 --- a/ta-browser/core/orchestrator.py +++ b/ta-browser/core/orchestrator.py @@ -676,7 +676,7 @@ async def run(self, command): self.log_token_usage( agent_type='planner', - usage=planner_response._usage, + usage=planner_response.usage, step=self.iteration_counter ) @@ -719,7 +719,7 @@ async def run(self, command): self.log_token_usage( agent_type='browser', - usage=browser_response._usage, + usage=browser_response.usage, step=self.iteration_counter ) @@ -780,7 +780,7 @@ async def run(self, command): self.log_token_usage( agent_type='critique', - usage=critique_response._usage, + usage=critique_response.usage, step=self.iteration_counter ) diff --git a/ta-browser/core/skills/final_response.py b/ta-browser/core/skills/final_response.py index fe3e3a9..c658b8d 100644 --- a/ta-browser/core/skills/final_response.py +++ b/ta-browser/core/skills/final_response.py @@ -50,14 +50,14 @@ def get_final_response_provider(): from core.utils.anthropic_client import get_client as get_anthropic_client from pydantic_ai.models.anthropic import AnthropicModel client = get_anthropic_client() - model = AnthropicModel(model_name=model_name, anthropic_client=client) + model = AnthropicModel(model_name=model_name, provider = "anthropic") provider = "anthropic" else: # OpenAI provider (default) from core.utils.openai_client import get_client as get_openai_client from pydantic_ai.models.openai import OpenAIModel client = get_openai_client() - model = OpenAIModel(model_name=model_name, openai_client=client) + model = OpenAIModel(model_name=model_name, provider = "openai") provider = "openai" return provider, client, model diff --git a/ta-browser/core/utils/init_client.py b/ta-browser/core/utils/init_client.py index 7d170c6..d33fa37 100644 --- a/ta-browser/core/utils/init_client.py +++ b/ta-browser/core/utils/init_client.py @@ -34,7 +34,7 @@ async def initialize_client(): # Create model instance from pydantic_ai.models.anthropic import AnthropicModel - model_instance = AnthropicModel(model_name=model_name, anthropic_client=client_instance) + model_instance = AnthropicModel(model_name=model_name, provider = "anthropic") logger.info(f"Anthropic client initialized successfully with model: {model_name}") return client_instance, model_instance diff --git a/ta-browser/requirements.txt b/ta-browser/requirements.txt index af8c9b0..7d51ddb 100644 --- a/ta-browser/requirements.txt +++ b/ta-browser/requirements.txt @@ -6,7 +6,7 @@ aiosignal==1.3.2 aiosmtplib==3.0.2 alembic==1.14.1 annotated-types==0.7.0 -anthropic==0.42.0 +anthropic==0.49.0 anyio==4.8.0 asgiref==3.8.1 asyncpg==0.30.0 @@ -41,7 +41,7 @@ google-auth==2.37.0 googleapis-common-protos==1.66.0 greenlet==3.0.3 griffe==1.5.4 -groq==0.13.1 +groq==0.15.0 grpcio==1.67.0 grpcio-status==1.62.3 h11==0.14.0 @@ -69,7 +69,7 @@ mypy-extensions==1.0.0 nest-asyncio==1.6.0 nltk==3.8.1 numpy==1.26.4 -openai==1.59.3 +openai==1.74.0 opentelemetry-api==1.29.0 opentelemetry-exporter-otlp-proto-common==1.29.0 opentelemetry-exporter-otlp-proto-http==1.29.0 @@ -96,8 +96,8 @@ pyautogen==0.2.27 pycparser==2.22 pycryptodome==3.20.0 pydantic==2.10.4 -pydantic-ai==0.0.17 -pydantic-ai-slim==0.0.17 +pydantic-ai==0.1.0 +pydantic-ai-slim==0.1.0 pydantic-core==2.27.2 pyee==11.1.0 pygments==2.18.0 @@ -137,7 +137,6 @@ typing-inspect==0.9.0 uritemplate==4.1.1 urllib3==2.3.0 uvicorn==0.30.3 -uvloop==0.21.0 watchfiles==0.24.0 websockets==13.1 wrapt==1.17.0