AIエージェント用サンドボックス構築入門：Dockerで安全なコード実行環境を作る方法

pip install docker anthropic python-dotenv

import docker
import os
import time
from anthropic import Anthropic
from dotenv import load_dotenv

load_dotenv()

class AISandbox:
    def __init__(self, image="python:3.10-slim"):
        # Dockerクライアントの初期化。デフォルトでローカルのDocker Engineに接続します。
        self.client = docker.from_env()
        self.image = image

    def execute_code(self, code, timeout=10):
        # 実行するPythonコードをラップ。
        # print文がないとLLMが結果を受け取れないため、最後に変数を表示させる工夫が必要です。
        formatted_code = f"import sys\n{code}"

        container = None
        try:
            # コンテナの作成と実行
            # network_mode="none" で外部通信を遮断し、AIが勝手に外部へデータを送るのを防ぎます。
            # mem_limit="128m" でメモリを制限し、メモリ消費バグによるホストのハングを防ぎます。
            container = self.client.containers.run(
                image=self.image,
                command=['python', '-c', formatted_code],
                detach=True,
                network_mode="none",
                mem_limit="128m",
                cpu_quota=50000, # CPU使用率を最大50%に制限
            )

            # 終了を待機。設定したtimeoutを超えたら強制終了させます。
            start_time = time.time()
            while container.status != 'exited':
                container.reload()
                if time.time() - start_time > timeout:
                    container.kill()
                    return "Error: Timeout"
                time.sleep(0.5)

            # ログ（実行結果）の取得
            logs = container.logs().decode('utf-8')
            return logs

        except Exception as e:
            return f"Error: {str(e)}"
        finally:
            # 使い終わったコンテナは即座に削除。これを忘れるとディスクが死にます。
            if container:
                container.remove(force=True)

def ask_ai_to_solve(task):
    anthropic = Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

    prompt = f"""
    以下のタスクを解決するPythonコードを書いてください。
    出力はコードのみとし、マークダウンのコードブロックで囲んでください。
    最後に必ず結果をprint関数で出力してください。

    タスク: {task}
    """

    response = anthropic.messages.create(
        model="claude-3-5-sonnet-20240620",
        max_tokens=1000,
        messages=[{"role": "user", "content": prompt}]
    )

    # AIの回答からコード部分を抽出（簡易的な実装）
    full_text = response.content[0].text
    if "```python" in full_text:
        code = full_text.split("```python")[1].split("```")[0].strip()
    else:
        code = full_text.strip("`")

    print(f"--- 生成されたコード ---\n{code}\n----------------------")

    sandbox = AISandbox()
    result = sandbox.execute_code(code)
    return result

# 実際に動かしてみる
if __name__ == "__main__":
    # 1から100までの素数を計算させる
    task_description = "1から100までの素数をリストアップして、その合計を表示してください。"
    output = ask_ai_to_solve(task_description)
    print(f"--- 実行結果 ---\n{output}")

--- 生成されたコード ---
def is_prime(n):
    if n < 2: return False
    for i in range(2, int(n**0.5) + 1):
        if n % i == 0: return False
    return True

primes = [n for n in range(1, 101) if is_prime(n)]
print(f"Primes: {primes}")
print(f"Total: {sum(primes)}")

--- 実行結果 ---
Primes: [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97]
Total: 1060

# 実用的なDockerfileの例
"""
FROM python:3.10-slim
RUN pip install pandas numpy scipy scikit-learn
WORKDIR /workspace
"""

# コード実行部を拡張して、ローカルファイルを読み込ませる方法
def execute_with_data(self, code, local_file_path, timeout=15):
    # ホストのファイルをコンテナの /workspace にマウントして実行
    # 読み取り専用(ro)にすることで、元データを破壊されるリスクを抑えます。
    container = self.client.containers.run(
        image="my-ai-analysis-image",
        command=['python', '-c', code],
        volumes={os.path.abspath(local_file_path): {'bind': '/workspace/data.csv', 'mode': 'ro'}},
        working_dir='/workspace',
        detach=True,
        network_mode="none",
        mem_limit="512m"
    )
    # ...（以下、ログ回収と削除処理は同じ）