DockerでAIエージェント専用サンドボックスを構築してコード実行を安全にする方法

# Docker操作用とLLM用のライブラリをインストール
pip install docker openai python-dotenv

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

class CodeSandbox:
    def __init__(self, image_name="python:3.11-slim"):
        self.client = docker.from_env()
        self.image_name = image_name
        # 必要なイメージを事前にプルしておく
        try:
            self.client.images.get(self.image_name)
        except ImageNotFound:
            print(f"Pulling image {self.image_name}...")
            self.client.images.pull(self.image_name)

    def execute_code(self, code: str, timeout: int = 10):
        # コンテナを起動してコードを実行する
        # mem_limit: メモリを128MBに制限（暴走対策）
        # network_disabled: 外部への通信を遮断（セキュリティ対策）
        container = self.client.containers.run(
            image=self.image_name,
            command=['python', '-c', code],
            detach=True,
            mem_limit="128m",
            network_disabled=True
        )

        start_time = time.time()
        while container.status != 'exited':
            container.reload()
            if time.time() - start_time > timeout:
                container.kill()
                return "Error: Execution timed out."
            time.sleep(0.5)

        # 実行結果のログを取得
        logs = container.logs().decode('utf-8')
        container.remove() # コンテナを掃除
        return logs

# 動作確認用
if __name__ == "__main__":
    sandbox = CodeSandbox()
    test_code = "print(sum([i for i in range(100)]))"
    print(f"Result: {sandbox.execute_code(test_code)}")

import os
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()
client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

def ask_agent(prompt: str):
    system_prompt = """
    あなたはPythonエンジニアです。ユーザーの依頼に対して、実行可能なPythonコードのみを出力してください。
    解説は不要です。コードは必ず print() 関数で結果を出力するようにしてください。
    """

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

    code = response.choices[0].message.content.strip()
    # Markdownのコードブロックが入る場合があるので除去
    if code.startswith("```python"):
        code = code.split("```python")[1].split("```")[0].strip()
    elif code.startswith("```"):
        code = code.split("```")[1].split("```")[0].strip()

    return code

# メイン処理
sandbox = CodeSandbox()
user_input = "1から50までの素数をリストアップして合計を計算して"

print(f"依頼: {user_input}")
generated_code = ask_agent(user_input)

print("--- 生成されたコード ---")
print(generated_code)
print("------------------------")

result = sandbox.execute_code(generated_code)
print(f"実行結果:\n{result}")

依頼: 1から50までの素数をリストアップして合計を計算して
--- 生成されたコード ---
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, 51) if is_prime(n)]
print(f"Primes: {primes}")
print(f"Sum: {sum(primes)}")
------------------------
実行結果:
Primes: [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47]
Sum: 328

def run_autonomous_agent(prompt: str, max_retries=3):
    sandbox = CodeSandbox()
    current_prompt = prompt

    for i in range(max_retries):
        print(f"\n--- 試行 {i+1} ---")
        code = ask_agent(current_prompt)
        result = sandbox.execute_code(code)

        if "Error" in result or not result.strip():
            print(f"失敗しました。エラー内容: {result}")
            # エラー内容を次のプロンプトに含めて修正を依頼する
            current_prompt = f"以下のコードを実行したところエラーが出ました。\nコード:\n{code}\nエラー内容:\n{result}\n修正して再出力してください。"
        else:
            print("成功しました。")
            return result
    return "制限回数内に解決できませんでした。"

# 実行
# final_result = run_autonomous_agent("複雑なデータ処理の依頼...")