DockerでAIエージェントのコード実行環境を隔離する方法

pip install docker openai python-dotenv

import os
import docker
from docker.errors import ContainerError, ImageNotFound
import time

class AgentSandbox:
    def __init__(self, image_name="python:3.10-slim"):
        # Dockerクライアントの初期化
        # ホストのDockerエンジンと通信を開始します
        try:
            self.client = docker.from_env()
        except Exception as e:
            print(f"Dockerへの接続に失敗しました: {e}")
            raise

        self.image_name = image_name
        self._ensure_image_exists()

    def _ensure_image_exists(self):
        # 指定されたイメージがない場合は自動でプルする
        try:
            self.client.images.get(self.image_name)
        except ImageNotFound:
            print(f"イメージ {self.image_name} をダウンロード中...")
            self.client.images.pull(self.image_name)

    def execute_code(self, code: str, timeout=10):
        # コードをコンテナ内で実行するメイン処理
        # なぜ python -c を使うのか：一時ファイルを作らずに文字列として実行できるため
        command = f'python -c "{code}"'

        try:
            # コンテナを起動して実行
            container = self.client.containers.run(
                image=self.image_name,
                command=["python", "-c", code],
                detach=True,           # バックグラウンドで実行
                network_disabled=True, # 外部への通信を遮断（重要！）
                mem_limit="128m",      # メモリ制限（OSフリーズ防止）
                cpu_quota=50000,       # CPU使用率を50%に制限
                remove=True            # 実行終了後にコンテナを自動削除
            )

            # タイムアウト監視
            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"Execution Error: {str(e)}"

if __name__ == "__main__":
    sandbox = AgentSandbox()

    # 正常なコードのテスト
    test_code = "print(sum([i for i in range(100)]))"
    print("--- 正常実行テスト ---")
    print(sandbox.execute_code(test_code))

    # 悪意ある（またはミスによる）無限ループのテスト
    bad_code = "import time\nwhile True: pass"
    print("--- 無限ループテスト（10秒で停止するはず） ---")
    print(sandbox.execute_code(bad_code, timeout=5))

--- 正常実行テスト ---
4950

--- 無限ループテスト（10秒で停止するはず） ---
Error: Timeout

import json
from openai import OpenAI

client = OpenAI(api_key="あなたのAPIキー")

def ask_agent(prompt: str):
    # エージェントに道具（サンドボックス）を定義
    tools = [{
        "type": "function",
        "function": {
            "name": "execute_python",
            "description": "Pythonコードを実行して結果を返す。計算やデータ処理に使う。",
            "parameters": {
                "type": "object",
                "properties": {
                    "code": {"type": "string", "description": "実行するPythonコード"}
                },
                "required": ["code"]
            }
        }
    }]

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": prompt}],
        tools=tools
    )

    tool_call = response.choices[0].message.tool_calls[0]
    if tool_call.function.name == "execute_python":
        args = json.loads(tool_call.function.arguments)
        code = args['code']

        print(f"AIが生成したコード:\n{code}")

        # サンドボックスで実行
        sandbox = AgentSandbox()
        result = sandbox.execute_code(code)

        print(f"実行結果: {result}")

# 実行例
ask_agent("1から1000までの素数の数を数えて")