【SmolAgents】Dockerコンテナをサンドボックスとして使用した場合に各Step毎に結果を出力させる方法

はじめに

以前SmolAgentsの記事を書きましたが、いずれもすべてのStepが終了した後にまとめて結果が出力されていました。
touch-sp.hatenablog.com
touch-sp.hatenablog.com
サンドボックスを使わなければそんなことは起きないのですが、サンドボックスを使った時はそうなっていました。

解決方法をClaude 3.7 Sonnetに教えてもらったので参考になるよう記録を残しておきます。

修正後

Dockerイメージ（名前：agent-sandbox）は前回と同様の方法ですでに作ってある前提です。

import docker
from typing import Generator

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

    def create_container(self):
        try:
            # コンテナを作成
            self.container = self.client.containers.run(
                self.image_name,
                command="tail -f /dev/null",  # コンテナを実行状態に保つ
                detach=True,
                extra_hosts={"host.docker.internal": "host-gateway"},
                network_mode="bridge",
                volumes={
                    "/home/hoge/data": {"bind": "/app/data", "mode": "rw"}
                }
            )
            print(f"コンテナ作成: ID={self.container.id}")
        except Exception as e:
            raise Exception(f"コンテナ作成エラー: {e}")

    def run_code_stream(self, code: str) -> Generator[str, None, None]:
        """低レベルDocker APIを使用してコードを実行し、出力をストリーミングする"""
        if not self.container:
            self.create_container()

        try:
            # コンテナの状態を確認
            self.container.reload()
            print(f"コンテナ状態: {self.container.status}")
            
            # ExecConfigを作成
            exec_config = self.client.api.exec_create(
                container=self.container.id,
                cmd=["python3", "-c", code],
                stdout=True,
                stderr=True
            )
            exec_id = exec_config["Id"]
            
            # 低レベルAPIを使ってストリーミングを開始
            output_stream = self.client.api.exec_start(
                exec_id=exec_id,
                stream=True
            )
            
            # デバッグ情報: ストリームのタイプを表示
            print(f"ストリームタイプ: {type(output_stream)}")
            
            # 各チャンクを処理
            for chunk in output_stream:
                if chunk:
                    try:
                        text = chunk.decode('utf-8')
                        yield text
                    except Exception as e:
                        error_msg = f"[デコードエラー: {e}]"
                        print(error_msg)
                        yield error_msg
            
            # 終了ステータスを確認
            inspect_result = self.client.api.exec_inspect(exec_id)
            exit_code = inspect_result.get("ExitCode")
            # 正常終了を示すメッセージを送信
            yield f"\n[実行完了: 終了コード={exit_code}]\n"
            
        except Exception as e:
            error_msg = f"[実行エラー: {str(e)}]"
            print(error_msg)
            yield error_msg

    def cleanup(self):
        if self.container:
            try:
                self.container.stop()
                self.container.remove()
                print("コンテナ停止・削除完了")
            except Exception as e:
                print(f"クリーンアップエラー: {e}")
            finally:
                self.container = None

from sandbox import DockerSandbox

# サンドボックスのインスタンスを作成
sandbox = DockerSandbox()

try:
    # エージェントのコードを定義 ({}がプロンプト変数と衝突しないように調整)
    agent_code = """
import os
import sys
import time

from smolagents import CodeAgent, ToolCollection, OpenAIServerModel
from mcp import StdioServerParameters

model = OpenAIServerModel(
    model_id="gemma-3-12b-it-4bit",
    api_base="http://192.168.10.12:8080",
    api_key="EMPTY"
)

server_parameters = StdioServerParameters(
    command="npx",
    args=["-y", "@modelcontextprotocol/server-filesystem","/app/data"],
    env={{}},  # 二重中括弧でエスケープ
)

with ToolCollection.from_mcp(server_parameters) as tool_collection:
    agent = CodeAgent(
        model=model,
        tools=[*tool_collection.tools],
        add_base_tools=True
    )
    
    # エージェントの実行
    response = agent.run("{0}")  # インデックスを使用してformat
    
"""

    prompt = "ローカルPCにアクセスした上で、ファイルの内容を要約して下さい。"
    
    print("===== LLMエージェントの実行を開始します =====")
    
    # 出力をリアルタイムで取得して表示
    for output in sandbox.run_code_stream(agent_code.format(prompt)):
        # 出力を即時表示
        print(output, end="", flush=True)
    
    print("===== エージェントの実行が完了しました =====")

finally:
    # 常に最後にクリーンアップする
    sandbox.cleanup()