Launching AI Workflows in E2B Sandboxes: A Beginner's Guide

"""
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("----------------
")

  # 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()