DockerでAIエージェント用セキュア実行環境を構築する方法

# 作業ディレクトリの作成
mkdir agent-sandbox
cd agent-sandbox

# Dockerfileの作成
cat <<EOF > Dockerfile
FROM python:3.11-slim

# セキュリティのため非ルートユーザーを作成
RUN useradd -m sandboxuser
WORKDIR /home/sandboxuser

# 必要最小限のライブラリをインストール
RUN pip install --no-cache-dir pandas numpy matplotlib seaborn

# ユーザーを切り替え
USER sandboxuser

# 実行コマンドを待機状態にする（コンテナを即終了させないため）
CMD ["tail", "-f", "/dev/null"]
EOF

# イメージのビルド
docker build -t ai-agent-sandbox .

pip install docker

import docker
import os
import time

class SandboxManager:
    def __init__(self, image_name="ai-agent-sandbox"):
        self.client = docker.from_env()
        self.image_name = image_name
        self.container = None

    def start(self):
        # 既存のコンテナがあれば削除
        self.stop()
        # コンテナの起動（CPUとメモリを制限）
        self.container = self.client.containers.run(
            self.image_name,
            detach=True,
            mem_limit="512m",  # メモリ使用量を512MBに制限
            nano_cpus=1000000000,  # CPU 1コア相当に制限
            network_disabled=True  # 外部ネットワークへのアクセスを禁止
        )
        return self.container.id

    def execute_code(self, code):
        if not self.container:
            raise Exception("コンテナが起動していません。")

        # 実行するコードを一時ファイルとして書き込む（コンテナ内）
        with open("temp_script.py", "w") as f:
            f.write(code)

        # ファイルをコンテナにコピー（簡易的な実装としてcatを使用）
        os.system(f"docker cp temp_script.py {self.container.id}:/home/sandboxuser/script.py")

        # コードの実行
        result = self.container.exec_run("python /home/sandboxuser/script.py")

        # 実行後、スクリプトは削除
        os.remove("temp_script.py")

        return {
            "exit_code": result.exit_code,
            "output": result.output.decode("utf-8")
        }

    def stop(self):
        if self.container:
            try:
                self.container.stop()
                self.container.remove()
            except:
                pass
            self.container = None

# 使用例
if __name__ == "__main__":
    sandbox = SandboxManager()
    sandbox.start()
    print("サンドボックス起動完了")

sandbox = SandboxManager()
sandbox.start()

test_code = """
import pandas as pd
import numpy as np

# ダミーデータの作成
df = pd.DataFrame({
    'A': np.random.randn(100),
    'B': np.random.randn(100)
})

# 集計処理
summary = df.describe()
print(summary)
"""

result = sandbox.execute_code(test_code)

print(f"終了コード: {result['exit_code']}")
print("--- 出力 ---")
print(result['output'])

sandbox.stop()

終了コード: 0
--- 出力 ---
                A           B
count  100.000000  100.000000
mean     0.123456   -0.012345
std      1.012345    0.987654
min     -2.345678   -2.123456
25%     -0.567890   -0.678901
50%      0.234567    0.123456
75%      0.890123    0.789012
max      2.567890    2.345678

import openai

def ask_agent_and_run(prompt, sandbox):
    client = openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"])

    # 1. LLMにコードを生成させる
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "あなたはPythonコードのみを出力するエンジニアです。マークダウンのコードブロックで囲んでください。"},
            {"role": "user", "content": prompt}
        ]
    )

    raw_content = response.choices[0].message.content
    # コードブロックの抽出（簡易版）
    code = raw_content.split("```python")[1].split("```")[0].strip()

    print(f"実行するコード:\n{code}")

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

    # 3. エラーがあればLLMに修正させる（最大2回）
    attempt = 0
    while result["exit_code"] != 0 and attempt < 2:
        print(f"エラー発生。修正を依頼中... (試行 {attempt + 1})")
        error_msg = result["output"]

        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": "コードにエラーがありました。修正してください。"},
                {"role": "user", "content": f"コード:\n{code}\n\nエラー内容:\n{error_msg}"}
            ]
        )
        code = response.choices[0].message.content.split("```python")[1].split("```")[0].strip()
        result = sandbox.execute_code(code)
        attempt += 1

    return result

# 実行
sandbox = SandboxManager()
sandbox.start()
final_result = ask_agent_and_run("1から100までの素数をリストアップして表示して", sandbox)
print(f"最終結果:\n{final_result['output']}")
sandbox.stop()