Lanzando flujos de trabajo de IA en sandboxes de E2B: una guía para principiantes

export E2B_API_KEY=your_key

export DEEPSEEK_API_KEY=your_key

export MISTRAL_API_KEY=your_key

"""
Script Name: e2b_llm_integration.py

Description:
This script facilitates the generation and execution of Python code using AI models such as Mistral and DeepSeek. It accepts a user-provided prompt, sends it to the specified model (default is DeepSeek), retrieves the generated code, and executes it within a sandboxed environment provided by E2B. The script supports both Mistral and DeepSeek APIs, allowing for flexible API selection, output file path specification, and input prompt sourcing from either direct input or a text file.

Key Features:
- Extracts Python code blocks from the model's response using regex.
- Executes the generated code in a secure, sandboxed environment using E2B.
- Supports command-line arguments for user prompts, API selection, output file paths, and input prompt files.

Usage:
- Default API (DeepSeek):
  python e2b_llm_integration.py "<your prompt>"

- Specify API:
  python e2b_llm_integration.py "<your prompt>" --api <mistral|deepseek>

- Specify output file path:
  python e2b_llm_integration.py "<your prompt>" --output <path/to/output/file>

- Use a prompt from a text file:
  python e2b_llm_integration.py --file <path/to/prompt_file.txt>

Dependencies:
- Python 3.6+
- mistralai: Python client for the Mistral API.
- e2b-code-interpreter: Python client for the E2B Sandbox API.
- openai: Python client for the OpenAI API, used for DeepSeek.

Environment Variables:
- MISTRAL_API_KEY: API key for accessing the Mistral API.
- DEEPSEEK_API_KEY: API key for accessing the DeepSeek API.
- E2B_API_KEY: API key for accessing the E2B Sandbox API.

Arguments:
- prompt (str): User prompt for the AI model.
- --file (str): Optional; Path to a text file containing the prompt.
- --api (str): Optional; Specify the API to use ('mistral' or 'deepseek'). Default is 'deepseek'.
- --output (str): Optional; Local path to save the generated file. Default is 'plot.jpg'.

Version: 1.0.0
"""

import os
import sys
import re
import logging
import argparse
from typing import Optional

from mistralai import Mistral
from e2b_code_interpreter import Sandbox
from openai import OpenAI

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

class CodeGenerator:
    """
    A class to handle the generation of Python code using specified AI models.

    Attributes:
        api_option (str): The API to use for code generation ('mistral' or 'deepseek').
        prompt (str): The user-provided prompt for code generation.
        api_key (str): The API key for accessing the specified API.
    """
    def __init__(self, api_option: str, prompt: str):
        """
        Initializes the CodeGenerator with the specified API option and prompt.

        Args:
            api_option (str): The API to use for code generation ('mistral' or 'deepseek').
            prompt (str): The user-provided prompt for code generation.
        """
        self.api_option = api_option
        self.prompt = prompt
        self.api_key = os.environ.get("MISTRAL_API_KEY" if api_option == "mistral" else "DEEPSEEK_API_KEY")
        if not self.api_key:
            print(f"Error: The environment variable {'MISTRAL_API_KEY' if api_option == 'mistral' else 'DEEPSEEK_API_KEY'} is not set.")
            sys.exit(1)

    def get_client(self):
        """
        Returns a client instance for the specified API.

        Returns:
            An instance of the client for the specified API.
        """
        if self.api_option == "mistral":
            return Mistral(api_key=self.api_key)
        elif self.api_option == "deepseek":
            return OpenAI(api_key=self.api_key, base_url="https://api.deepseek.com")

    def generate_code(self):
        """
        Generates Python code based on the user-provided prompt using the specified API.

        Returns:
            str: The generated Python code.
        """
        client = self.get_client()
        system_prompt = (
            """
You are a Python coding assistant that only outputs Python code without any explanations or comments. For all modules or packages you use in the script, do not assume the packages are installed. Include the necessary installation commands before the import statements. This is an example for the module SomeClass:

import subprocess
import sys

try:
    from some_module import SomeClass
except ImportError:
    subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'some_module'])
    from some_module import SomeClass

You will be asked to create a method and call it. The call to the method has to be with if __name__ == "__main__":

Document well the method in the header including:
- Arguments
- Purpose of the method
- Return values

Very important: write the Python code like this:
```python
code
```
            """
        )

        try:
            if self.api_option == "mistral":
                response = client.chat.complete(
                    model="codestral-latest",
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": self.prompt}
                    ],
                    temperature=0.0
                )
            elif self.api_option == "deepseek":
                response = client.chat.completions.create(
                    model="deepseek-chat",
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": self.prompt}
                    ],
                    stream=False,
                    temperature=0.0
                )
            return response.choices[0].message.content.strip()
        except Exception as e:
            print(f"Error communicating with {self.api_option.capitalize()} API: {e}")
            sys.exit(1)

def extract_code_block(content: str) -> Optional[str]:
    """
    Extracts a Python code block from the given text data using regex.

    Args:
        content (str): The text data containing a Python code block.

    Returns:
        Optional[str]: The extracted Python code block, or None if not found.
    """
    if not content:
        logging.warning("Content is empty or None.")
        return None

    match = re.search(r'\`\`\`python(.*?)\`\`\`', content, re.DOTALL)
    if match:
        return match.group(1).strip()  # Return the Python code without the backticks
    logging.warning("No Python code block found in the content.")
    return None

def main():
    """
    Main function to parse command-line arguments, generate code, and execute it in a sandboxed environment.
    """
    # Command-line argument parsing
    parser = argparse.ArgumentParser(description="Generate and execute Python code using AI models.")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("prompt", type=str, nargs="?", help="Prompt for the AI model.")
    group.add_argument("--file", type=str, help="Path to a text file containing the prompt.")
    parser.add_argument(
        "--api", type=str, choices=["mistral", "deepseek"], default="deepseek",
        help="Specify the API to use (default: deepseek)."
    )
    parser.add_argument(
        "--output", type=str, default="plot.jpg",
        help="Optional: Local path to save the file (default: plot.jpg)."
    )
    args = parser.parse_args()

    # Check for E2B_API_KEY environment variable
    if not os.environ.get("E2B_API_KEY"):
        print("Error: The environment variable E2B_API_KEY is not set.")
        sys.exit(1)

    # Read the prompt from the provided argument or file
    if args.file:
        try:
            with open(args.file, 'r') as file:
                prompt = file.read().strip()
        except Exception as e:
            print(f"Error reading the prompt file: {e}")
            sys.exit(1)
    else:
        prompt = args.prompt

    api_option = args.api
    local_file_path = args.output
    sandbox_filename = os.path.basename(local_file_path)  # Extract only the filename for the sandbox

    print("Input Prompt:")
    print(prompt)
    print(f"Using model: {api_option.capitalize()}")
    print(f"Local file path: {local_file_path}")

    # Generate code
    code_generator = CodeGenerator(api_option, prompt)
    raw_code = code_generator.generate_code()
    code = extract_code_block(raw_code)
    if not code:
        print("Error: Received empty code from the model.")
        sys.exit(1)

    print("Generated Code:")
    print("----------------")
    print("```python")
    print(code)
    print("```")
    print("----------------\n")

    # Execute code in E2B Sandbox
    try:
        with Sandbox() as sandbox:
            execution = sandbox.run_code(code)
            result = execution.text
            logs = execution.logs

            try:
                # Attempt to download file from the sandbox
                content = sandbox.files.read(f"/home/user/{sandbox_filename}", format="bytes")

                # Write file to the specified local path
                with open(local_file_path, "wb") as file:
                    file.write(content)

                logging.info(f"File successfully downloaded and saved to {local_file_path}")
            except Exception as e:
                logging.error(f"Failed to download file from sandbox: {e}")

    except Exception as e:
        logging.error(f"Error executing code in Sandbox: {e}")
        sys.exit(1)

    logging.info("Execution Output:")
    logging.info("Stdout Logs:")
    for log in logs.stdout:
        logging.info(log)
    logging.info("Stderr Logs:")
    for log in logs.stderr:
        logging.info(log)

    logging.info("Execution Result:")
    logging.info(result)

if __name__ == "__main__":
    main()