Kotlin Tutorials

AI エージェントを Kotlin で構築する － 第 2 回：ツールの深掘り 

前回の記事では、リスト、読み取り、書き込み、および編集の機能を備えた基本的なコーディングエージェントのビルド方法を確認しました。今回の記事では、Koog フレームワーク内で追加のツールを作成することで、このエージェントの機能を拡張する方法を深掘りしていきましょう。ExecuteShellCommandTool を例に挙げてビルドし、エージェントにコードを実行し、実際のエンジニアリングで必要となるフィードバックループ（コードの実行、失敗の確認、および実際の出力に応じたコードの改善）の終了方法を教えます。

LLM は構文エラーを適切に回避するものの、統合に関する問題の対処は苦手な傾向があります。たとえば、存在しないメソッドを呼び出す、インポートを省略する、インターフェースを部分的に実装するなどの傾向です。このような問題は従来のやり方でコードをコンパイルして実行すればすぐに見つかるものですが、 ちょっとしたプロンプトを追加することで、LLM に小さなテストを実行し、このような各動作を検証させることができます。

では、そのようなツールをどうやってビルドすればいいのでしょうか？ まずは基本から見ていきましょう。

Koog ツールの構造

まず、抽象的な ai.koog.agents.core.tools.Tool クラスを継承することから始めましょう。以下の情報を提供する必要があります。

1. 名前: このチームでは名前を 2 つのアンダースコアで囲む snake_casing 命名規則に従いますが、これは個人の好みです。
2. 説明: このフィールドは LLM のメインのドキュメントとして機能し、ツールの実行内容と呼び出される理由を説明するためのものです。
3. Args クラス: このクラスでは、LLM がツールを呼び出す際に提供する必要のある、または提供できるパラメーターを記載します。
4. Result クラス: このクラスでは、LLM へのメッセージに整形されるデータを定義します。データを LLM が読み取れる文字列に整形する textForLLM() 関数を含めることができます。Result クラスは主に開発者の便宜を図るため、ツールの結果を記録または UI に表示しやすくするために存在するものです。エージェント自体は整形済みの文字列のみを必要とします。
5. Execute() メソッド: このメソッドは Args のインスタンスを引数に取り、Result のインスタンスを返します。LLM がツールを呼び出す際に起こる内容のロジックを定義します。

これらのコンポーネントはデータベースコネクター、API クライアント、またはこの記事でご紹介するシェルコマンド実行ツールなど、何をビルドする場合もすべての Koog ツールの基盤になります。では、ExecuteShellCommandTool のビルドを詳しく見ながら、これらの原則を実際に確認してみましょう。

安全に関する注意喚起

ExecuteShellCommandTool の詳細を深掘りする前に、安全に関する考慮事項をいくつかお伝えしておきます。

LLM は意図的に問題を引き起こそうとする脅威アクターではありませんが、将来的に問題につながる可能性のある予期しないミスを起こすことがあります。また、エージェントにコマンドラインを実行する権限を与えている場合、このようなミスは深刻な結果をもたらす可能性があります。

コマンドの実行を隔離環境にサンドボックス化したり、エージェントに与える権限を制限したりするなど、このリスクを緩和する手段はいくつかありますが、 それらを実施するのは非常に面倒ですし、この連載は実際のツールを作成するところに重きを置いています。

このことに留意しながら、非常に単純な 2 つのリスク緩和手段を以下に示します。

1. コマンドごとにコマンドの実行を確認する。LLM が実行する各コマンドの確認に時間をかけられるのなら、この手段が最も安全です。ただし、この場合はエージェントの自律性を大きく制限することになります。
2. 各コマンドが自動的に承認される Brave モードを使用する。このモードは適切にサンボドックス化されていないマシンをリスクにさらす可能性がありますが、この記事では悪影響を及ぼすことなく破棄できる隔離環境で実行される独自のベンチマーク内のみでこのモードを使用します。

ExecuteShellCommandTool を実装する

この例で使用する ExecuteShellCommandTool の場合、これらのコンポーネントは以下のようになっています。

1. 名前: __execute_shell_command__
2. 説明: 「シェルコマンドを実行してその出力を返す」など。
3. Args クラス: command、timeoutSeconds、workingDirectory
4. Result クラス: output、exitCode、command（意外かもしれませんが、ログや UI にどのコマンドが実行されたかを出力すると便利です）。
5. execute() メソッド: 確認をリクエストしてからコマンドを実行します（以下で実装を詳細に確認します）。

execute() メソッドの実装

では、execute() メソッドがどのように実装されているかを見てみましょう。ここでは以下のように基本ロジックをヘルパーメソッドに委譲することで、メソッドを単純化しています。

override suspend fun execute(args: Args): Result = when (
        val confirmation = confirmationHandler.requestConfirmation(args)
    ) {
        is ShellCommandConfirmation.Approved -> try {
            val result = executor.execute(
                args.command, args.workingDirectory, args.timeoutSeconds
            )
            Result(args.command, result.exitCode, result.output)
        } catch (e: CancellationException) {
            throw e
        } catch (e: Exception) {
            Result(
               args.command, null, "Failed to execute command: ${e.message}"
            )
        }

        is ShellCommandConfirmation.Denied ->
            Result(
                args.command, null, 
                "Command execution denied with user response: ${confirmation.userResponse}"
            )
    }

これは簡潔なフローです。ユーザーにコマンドを実行するための確認をリクエストし、コマンド実行処理を介してリクエストが承認されれば実行し、そうでなければ拒否されたことを伝えるメッセージを LLM に返します。

この実装によって例外のキャッチとエラーメッセージの LLM への転送が可能になり、やり方を調整したり、別のやり方を試したりできるようにもなります。

ConfirmationHandler の構成

ConfirmationHandler は ExecuteShellCommandTool を作成する場合に構成可能になり、さまざまな実装を可能にします。現時点では、以下の 2 つの実装が提供されています。

1. PrintShellCommandConfirmationHandler: コマンドラインを介してユーザーにプロンプトを表示します。
2. BraveModeConfirmationHandler: すべてを自動的に承認します。

2 つ目の実装は無条件で単に承認するものですが、1 つ目の実装には若干興味深い特徴があります。

override suspend fun requestConfirmation(
        args: ExecuteShellCommandTool.Args
    ): ShellCommandConfirmation {
        println("Agent wants to execute: ${args.command}")
        args.workingDirectory?.let { println("In: $it") }
        println("Timeout: ${args.timeoutSeconds}s")
        print("Confirm (y / n / reason-for-denying): ")

        val userResponse = readln().lowercase()
        return when (userResponse) {
            "y", "yes" -> ShellCommandConfirmation.Approved
            else -> ShellCommandConfirmation.Denied(userResponse)
        }
    }

ユーザーには 3 つの選択肢（承認、拒否、理由付きの拒否）が与えられますが、実装上は両タイプの拒否が同じように処理されています。どちらも LLM に返され、各タイプは LLM によって適切に解釈されて処理されます。

CommandExecutor の構成

ConfirmationHandler と同様に CommandExecutor も構成可能ですが、現時点では JVM 実装のみを提供しています。理屈の上では Android、iOS、WebAssembly、およびその他のプラットフォームに対応する実装を作成できるかもしれませんが、明確な需要がないため、それらについてはとりあえず後回しにします。

タイムアウトの処理方法

使用されるプラットフォームを問わず、CommandExecutor の 1 つの側面であるタイムアウトの処理には特段の注意が必要です。現在の実装では、エージェントは実行に時間がかかるコマンドを中断することはできません。

人は直感的に焦って行動することが多く、その結果 Ctrl+C を使用して実行をキャンセルすることがあります。しかし、その行動はマルチスレッド化されたエージェントが認識し、焦燥感という概念が根底にあることを前提としています。とはいえ、この状況にはるかに適した、より単純かつ直感的な方法はあります。

LLM に最大実行時間を指定させることで、この制限時間を超えるコマンドを安全に中断することができます。このタイムアウト値はユーザーに表示されるため、ユーザーは要求された期間が合理的でないか、過剰であると判断した場合に実行を拒否することができます。

単にプロセスを終了して一般的なタイムアウトメッセージを返すのではなく、できるだけ多くの出力を保持できるようにする必要があります。不完全な結果であっても、LLM が有用な情報を抽出したり、少なくともタイムアウトが発生した箇所を理解したりするのには役立ちます。これは、以下のような慎重な実装によって実現できます。

val stdoutJob = launch {
    process.inputStream.bufferedReader().useLines { lines ->
        try {
            lines.forEach { stdoutBuilder.appendLine(it) }
        } catch (_: IOException) {
            // Ignore IO exception if the stream is closed and silently stop stream collection
        }
    }
}

val isCompleted = withTimeoutOrNull(timeoutSeconds * 1000L) {
    process.onExit().await()
} != null

if (!isCompleted) {
    process.destroyForcibly()
}

stdoutJob.join()

エージェントレベルでの変更内容

エージェントレベルでの変更は最小限であり、10 行程度のコードです。差分で示されているように、この更新の大部分にはシステムプロンプトの拡張が伴っています。

単にツールを追加して変更をさらに小規模にすることも可能でしたが、そうではなく、エージェントにさらに 2 つの変更も加えました。とはいえ、これらも些細な変更であることは変わりありません。

A）BRAVE_MODE の切り替え処理

この切り替え処理の実装は、BRAVE_MODE 環境変数をチェックするだけの比較的単純なものです。自分で実装するのが非常に簡単なことを示すために Brave モードの ConfirmationHandler にラムダ関数も使用しましたが、前述の BraveModeConfirmationHandler を使用することもできます。 

fun createExecuteShellCommandToolFromEnv(): ExecuteShellCommandTool {
    return if (System.getenv("BRAVE_MODE")?.lowercase() == "true") {
        ExecuteShellCommandTool(JvmShellCommandExecutor()) {
            _ -> ShellCommandConfirmation.Approved 
        }
    } else {
        ExecuteShellCommandTool(
            JvmShellCommandExecutor(), 
            PrintShellCommandConfirmationHandler(),
        )
    }
}

B）プロンプトの「完了の定義」の拡張

LLM がこの新機能を確実に利用してコードを効果的に実行できるようにするため、テストの作成と実行を強く促す「完了の定義」を使用してプロンプトを拡張しました。

"""
... // Previous prompt from step 01

Production-ready means verified to work—your changes must be proven correct and not introduce regressions.

You have shell access to execute commands and run tests. Use this to work with executable feedback instead of assumptions. Establish what correct behavior looks like through tests, then iterate your implementation until tests pass. Validate that existing functionality remains intact. Production-ready means proven through green tests—that's your definition of done.
"""

完成です。完全な ExecuteShellCommandTool コンポーネントをビルドし、それを最小限の変更でエージェントに組み込みました。ただし、本当の疑問は「それで実際にパフォーマンスが改善するのか？」ということです。

ベンチマークテストの結果

両方のバージョンを SWE-bench-verified セットに対して実行することで、実行機能が明確なメリットを提供することが証明されました。エージェントのパフォーマンスが改善され、成功件数が 249/500 件（50%）から 279/500（56%）件に増加しました。リーダーボードスコアは約 70% を達成しましたが、この結果はエージェントにコードを実行して検証する機能を与えることが正しい方向への一歩であることを示しています。ただし、エージェントがまだうまく処理できない部分や次の改善点を理解するには、その動作をさらに詳細に把握する必要があります。そこで欠かせないのがログ記録です。

まとめ: Koog でのツールのビルド

この記事全体を通し、どのツールにも名前、説明、および execute() メソッドが必要だという Koog のツール構造の仕組みを確認しました。ConfirmationHandler コンポーネントと CommandExecutor コンポーネントを介して ExecuteShellCommandTool を構成可能にし、複雑なロジックを置換可能な実装に委譲する方法を示しました。

これと同じパターンは、データベースコネクター、API クライアント、ファイルプロセッサー、カスタム統合など、どんなツールをビルドする場合も適用されます。このフレームワークには LLM と対話するための構造が備わっており、開発者は固有の機能を提供することができます。 

次回の記事では、エージェントのログを記録してトレースする処理を追加する方法を取り上げる予定です。そうすることで、エージェントの動作を理解して改善するのに必要な情報を得られるでしょう。エージェントがどのように判断を下すのかを理解しておくことは、イテレーションを実施する上では非常に重要です。また、可観測性に関する適切なツールもエージェント自体に提供するツールと同じくらい重要です。

オリジナル（英語）ブログ投稿記事の作者：