coSTAR: DatabricksでAIエージェントを迅速かつ安全に展開する方法

テストスイートなしでコーディングアシスタントにコードベースのリファクタリングを任せることはないでしょう。テストがなければ、アシスタントは手探り状態になります。ある関数を修正しても、知らないうちに他の3つを壊してしまうかもしれません。テストはループを完結させるものです。テストを実行し、失敗を観察し、コードを修正し、再度実行する。テストがなければ、自信もありません。

Databricksでは、Databricksプラットフォームの新機能（例：Genie Codeのデータエンジニアリング、トレース分析、機械学習機能）から、OSSプロジェクト（例：MLflowアシスタント）、社内エンジニアリングワークフロー（例：オンコールサポートや自動コードレビューア）まで、幅広い機能をカバーするエージェントを継続的に開発・デプロイしています。これらのエージェントは、時間のかかるタスクを実行し、何千行ものコードを生成し、新しいデータおよびAIアセットを作成するなど、さまざまなことができます。初期段階では基本的なチェック体制はありましたが、自信を持って反復できるような包括的で自動化されたテストスイートが不足していました。この投稿では、MLflowを使用してそのギャップをどのように埋めたか、そしてその周りに構築したベストプラクティスであるcoSTAR（結合されたシナリオ、トレース、評価、改善）手法について説明します。coSTARは2つの結合されたループを実行します。1つは審査員を人間の専門家の判断に合わせ、信頼できるようにするループ、もう1つはそれらの信頼できる審査員を使用して、エージェントがすべてのテストシナリオに合格するまで自動的に改善するループです。

図: coSTARフレームワークは、2つのミラーリングされたSTARループ（シナリオ → トレース → 評価 → 改善）を実行します。エージェントループ（青）は審査員を使用してトレースを自動採点し、審査員に合わせてエージェントを改善します。審査員ループ（オレンジ）は人間の専門家がトレースを採点し、専門家の評価に合わせて審査員を改善します。両方のループは同じシナリオとトレースを共有します。

問題：テストなしのコーディング

初期段階では、私たちの開発ループは次のようでした。エージェントを実行し、その出力を手動でレビューし、欠陥を見つけ、コーディングアシスタントに修正を指示する。これを繰り返していました。

これが、テストなしでコードを書き、すべての変更を手動でQAすることに似ていると感じるなら、まさにその通りでした。そして、予測どおりに失敗しました。明らかな反応は「ではテストを書けばいい」というものです。しかし、エージェントのテストは、決定論的な関数のテストとは構造的に異なり、いくつかの課題が同時に複合します。

• 非決定性。同じ実装、同じ入力でも、異なる実行で異なる出力を生成することがあります。テストは正確な出力をアサートするのではなく、出力の特性を評価する必要があります。
• 遅いフィードバックループ。単一のエージェント実行には数十分かかることがあります。1秒未満のテストスイートが許容するような反復はできません。すべての評価サイクルはコストが高いです。
• 連鎖的なエラー。ステップ3での誤った決定がステップ7での失敗を引き起こします。症状が表面化する頃には、根本原因がエージェント実行の数ステップ前に埋もれています。
• 主観的な品質。多くのテストの側面（この特徴量エンジニアリングコードは良いか？このデータクレンジングアプローチは適切か？）には正解がありません。これらの側面を判断するにはドメイン専門知識が必要です。

これらの制約が、その後のすべての設計決定を形作りました。また、これがこの問題を興味深いものにしている理由でもあります。私たちは単にテストランナーを構築しているのではなく、「正しい」が判断に委ねられる確率的で、長時間実行される、多段階のプロセスのための自動化された最適化手法を構築しているのです。

私たちのアプローチを導くアナロジー

注意深く見れば、エージェント開発は、すべてのエンジニアがすでに知っている開発ループにきれいにマッピングされます。

このアナロジーは単なる例示ではありません。これは、私たちがcoSTARと呼ぶシステムの文字通りのアーキテクチャです。シナリオ定義をテストフィクスチャとして、トレースキャプチャをテストハーネスとして、審査員による評価をテストスイートとして、そして赤緑ループとして改善する、2つの結合されたループです。それぞれの要素を見ていきましょう。

S - シナリオ定義

従来のテストでは、テストフィクスチャが事前条件を設定します。データベースを作成し、データをシードし、環境を構成するなどです。私たちの同等物はシナリオ定義であり、初期状態、ユーザープロンプト、および期待される結果の構造化された記述です。

以下は、乱雑なデータセットに対してデータアナリストエージェントをテストするための簡略化されたシナリオです。

各シナリオは、テストフィクスチャのように、セットアップ、入力、成功基準を1か所にまとめています。私たちは、さまざまなエージェントにわたってこれらのスイートを維持しており、一般的なケース、エッジケース、既知の過去の失敗をカバーしています。新しい障害モードを発見するにつれて、スイートは時間とともに成長します。本番環境で見つかるすべてのバグが新しいシナリオになるのは、すべての本番バグが回帰テストになるべきであるのと同じです。

なぜこの構造にこだわるのか？エージェントの実行はコストが高いからです。単一のシナリオの実行には数分かかります。私たちはテストする内容について意図的である必要があり、シナリオ定義がポータブルである必要があります。つまり、同じシナリオを異なるエージェント実装や同じエージェントの異なるバージョンに対して実行できる必要があるのです。

T - トレースキャプチャ

テストスイートを実行するために、各シナリオのプロンプトをテスト対象エージェント（AUT）に送信するハーネスを使用します。各実行はMLflow トレースとしてキャプチャされます。これは、すべてのツール呼び出し、すべての中間出力、およびエージェントが生成するすべてのアーティファクトの構造化されたログです。フライトレコーダーのように考えてください。エージェントが行ったすべてを順番にキャプチャするため、実行の任意の部分を後から検査できます。

重要なアーキテクチャ上の決定は、実行と採点を分離することです。テストハーネスはトレースを生成し、審査員（次に紹介します）がそれらを採点します。これらは別々のステップです。トレースを永続化することで、シナリオを再実行せずに審査員を反復できます。しきい値を調整する？記録されたトレースを数秒で再採点します。新しい審査員を追加する？これまでに収集したすべてのトレースに対して実行します。審査員が間違っていると疑う？その判断を記録と比較し、オフラインでデバッグします。1回のコストの高いエージェント実行で、後で審査員を整合させるために使用するゴールデンセットの候補として含むなど、何度も再利用されるデータが生成されます。

A - 審査員による評価

審査員はトレースを操作し、実行の特性について推論します。エージェントは有効なコードを生成したか？出力は品質しきい値を満たしたか？エージェントは正しいプロセスに従ったか？前述のとおり、この評価は従来の単体テストとは異なります。エージェントの出力は非決定論的で豊富であるため、正確な出力をアサートすることは本質的に無意味です。

これらの審査員を実装する標準的なアプローチは「LLMを審査員として」です。つまり、完全なトレースをモデルに与え、スコアと、同様に重要なそのスコアの根拠を尋ねるものです。しかし、それはプログラム全体の状態をアサーションにダンプするテストを書くようなものです。コストが高く、脆弱で、デバッグが難しいです。私たちのエージェントの場合、単一のトレースが何千行にもなることがあります。それを審査員のコンテキストウィンドウに詰め込むと、判断の品質が低下します。

代わりに、MLflowのエージェント型審査員を使用します。これらはそれ自体がエージェントであり、トレースを選択的に探索するためのツールを備えた審査員です。よく書かれたテストが特定の関数を呼び出し、特定の戻り値をチェックするように、エージェント型審査員はトレース上の特定のツールを呼び出し、特定のプロパティをチェックします。

以下に、私たちがエージェント全体で使用してきた審査員の例をいくつか紹介します。

スキル呼び出し審査員は、トレースを探索し、エージェントがシナリオによってターゲットとされたスキルを呼び出したかどうかを識別します（そうでない場合、スキルの目的がAUTに明確ではないことになります）。

ベストプラクティス審査員は、出力がDatabricksの公式ドキュメントに従ってベストプラクティスに従っているかどうかを探索します。

Outcome Judge は、出力アセットのトレースを検査し、特定のプロパティをアサートします。データアナリストの例に戻り、エンジニアリングコードが作成されたトレースの部分を特定し、そのコードが現在のタスクに適しているかどうかを評価します。

このジャッジは、主観的な品質問題に正面から取り組むため興味深いものです。優れた特徴量エンジニアリングとは何かは、ドメインの専門知識に依存します。LLMジャッジは、すぐにこれを正しく判断することはできません。ジャッジのプロンプトに完全な基準を書き出そうとすることは魅力的です。「歪んだ分布には平均よりも中央値補完を優先する、距離ベースのモデルの前には常に特徴量をスケーリングする、...」しかし、ドメインエキスパートの完全な判断をプロンプトにエンコードするのは骨の折れる作業であり、脆いです。人間が例を見て「これは良い」または「これは悪い」と言う方が、完全な仕様を書き出すよりもはるかに簡単です。これがアライメントが機能する理由であり、これについては後ほど説明します。

一般的に、単一エージェントのテストスイートには、いくつかのカテゴリにわたるジャッジが含まれています。

決定論的チェック（機械的に検証可能で、LLMは不要なもの）：

• 生成されたコードの構文/リンティング
• 出力スキーマの検証（期待されるテーブルは存在するか？カラムの型は正しいか？）
• ツールシーケンスのリンティング（エージェントは問題を修正しようとする前にエラーログを読んだか、それとも直接コード編集にスキップしたか？）

LLMベースのチェック（コンテキストの理解を必要とする判断）：

• コード差分ガイドライン（エージェントは無関係な行を変更したか？非推奨のAPIを導入したか？）
• ベストプラクティスへの準拠（生成されたコードはこのドメインの慣例に従っているか？）

運用メトリクス（個々には合否を判断しないが、時間の経過とともに健全性を追跡するシグナル）：

• トークン使用量（高いトークン数は、エージェントが苦戦している、再試行している、後戻りしている、または堂々巡りしていることを示すことが多い）
• ツール呼び出し回数と失敗率（失敗したツール呼び出しの急増は、何かが間違っていることを示す）
• レイテンシー（エージェントがタスクを完了するまでの実時間）

運用メトリクスには注意が必要です。これらは合否を判断するジャッジのようにリリースを阻止するものではありませんが、コスト管理と早期警告には不可欠です。変更後にトークン使用量が2倍になった場合、すべてのジャッジが合格したとしても何かが間違っています。エージェントはおそらく必要以上に多くの作業を行っているでしょう。私たちはこれらを時間の経過とともに追跡し、異常を警告します。

テストスイートを時間とともに成長させる

テストスイートは一度に作成されるものではありません。時間の経過とともに進化します。最初は、出力が存在するか？解析できるか？といった最も単純なチェックから始まります。次に、出力が正しいスキーマ、正しいカラム、正しい型を持っているか？といった構造的なチェックが続きます。その後になって初めて、エンドツーエンドのデータ検証ジャッジが登場します。つまり、実行したときに出力が実際に正しい結果を生成するか？というものです。

これは、従来のソフトウェアにおけるテストスイートの成熟の仕方を反映しています。網羅的な統合テストは初日から存在するわけではありません。最初はスモークテストから始まり、障害モードが出現するにつれて単体テストが追加され、時間の経過とともにエンドツーエンドのカバレッジへと構築されていきます。重要なのは、インフラストラクチャが新しいジャッジを安価に追加することをサポートしているため、テストスイートがエージェントとともに成長することです。

テストのテスト：ジャッジのアライメント

すべてのエンジニアが知っている問題があります。それは、不安定な、または間違ったテストスイートが不良なコードを承認し、自信を持ってバグを出荷してしまうことです。同様に、不適切な結果を承認するジャッジは、誤った安心感を与えます。ここでcoSTARフレームワークの第2のループが登場します。エージェントの洗練を促進するのと同じシナリオとトレースが、人間の専門家によるスコアを真実として、ジャッジの洗練も促進します。これは重要です。なぜなら、テストの正確性が検査によって検証できる従来のテストとは異なり、LLMジャッジは確率的であり、自然言語の基準を解釈する方法がずれる可能性があるからです。そのため、私たちはそれらを検証し、人間の専門家とアライメントを維持する方法が必要です。

このアライメントを行うために、まず、当社のエンジニアが手動で評価した、通常数十個のエージェント出力の例からなるゴールデンセットをキュレーションします。これがジャッジが同意しなければならない真実です。次に、MLflow のアライメント機能（GEPA や MemAlign のような技術を搭載）を活用して、ゴールデンセットに対してジャッジを自動的に洗練させます。これは、AUT 自体を洗練するために使用する STAR ループと構造的には同じですが、評価ステップは人間の専門家によって実行され、洗練ステップはジャッジに適用されることに注意してください。

R - 洗練 (Refine)

ジャッジループが人間の専門家の判断とアライメントしたジャッジがあれば、エージェントループを信頼できるようになります。コーディングアシスタントはエージェントをそのコードベースとして扱い、ジャッジをそのテストスイートとして扱います。失敗を読み取り、根本原因を診断し、エージェントをパッチし、すべてを再実行します。エンジニアは依然としてエージェントへの提案された変更のレビュー担当者であり最終決定者ですが、この自動化された反復により、エージェントの分析と改善における人間の労力を大幅に節約できます。

データアナリストエージェントの1回のイテレーションは次のようになりました。

赤（Red）。エージェントの初期バージョンをシナリオスイートに対して実行しました。ベストプラクティスジャッジは不一致を指摘しました。つまり、エージェントが生成した論理ビューのコードが、公式の推奨事項/ドキュメントと異なっていたのです。この不一致は正確性には影響しませんでしたが、生成されたコードの保守とデプロイに影響を与えました。これは、手動調査では発見が困難な、潜在的なリグレッションの一例です。

緑（Green）。コーディングアシスタントはジャッジのフィードバックを分析し、ギャップを特定しました。エージェントは、作成すべきビューのタイプ（一時的か永続的か）について指示的ではないスキルを使用していたのです。関連するガイダンスをスキルに追加した後、テストは正常に合格し、その変更が他のリグレッションを導入しないことが検証されました（他のテストシナリオに基づく）。

エージェントだけでなくインフラストラクチャのリグレッションテスト

これまで、ジャッジをエージェントのテストとして、エージェントの実装が変更されたときにリグレッションを検出するものとして説明してきました。しかし実際には、変更されるのはエージェント自体だけではありません。エージェントは外部ツールやインフラストラクチャに依存しており、それらも変更されます。

当社のエージェントは、データアクセス、コード実行、環境設定などのための標準化されたインターフェースであるMCPツールを呼び出します。これらのツールには独自の開発チームとリリースサイクルがあります。ツールがその実装を変更した場合（例えば、コード実行ツールが異なる形式でstderrを返すようになったり、データアクセスツールがnull値の処理方法を変更したりした場合）、エージェント自体はまったく変更されていなくても、エージェントの動作が壊れる可能性があります。

私たちは毎晩のビルドでジャッジを実行しているため、それらはエージェントの現在の実装だけでなく、フルスタックに対するリグレッションテストとして機能します。ツールチームが変更をリリースし、それが原因でエージェントがジャッジに失敗し始めた場合、私たちはそのエラーが顧客に届く前にすぐに検出します。さらに重要なのは、ジャッジの失敗が何が壊れたのか（リグレッションした特定の品質次元）を教えてくれるため、根本原因がエージェントにあるのか、それともエージェントが依存するツールにあるのかをトリアージするのがはるかに容易になることです。

これは、従来のソフトウェアにおける統合テストが提供する価値と同じです。つまり、コードとその依存関係との間の契約を保護します。唯一の違いは、ここでは「コード」がエージェントであり、「依存関係」がMCPツールであるということです。

評価から本番環境の監視へ

テストの類推には、もう一つ驚くほど価値のある拡張があります。それは、同じジャッジを本番環境のトラフィックで実行することです。

従来のソフトウェアでは、テストはCIで終わりません。本番環境も監視されます。エラー率、レイテンシーのパーセンタイル、ライブトラフィックのビジネスメトリクスなどです。開発環境でコードを検証するのと同じテストロジックが、本番環境ではヘルスチェックやアラートとして再登場することがよくあります。

私たちも同じことを行います。評価用に構築したジャッジは、評価シナリオだけでなく、あらゆるエージェントの会話をスコアリングするように設計されています。そのため、実際の運用会話でそれら（またはサンプリングされたサブセット）を実行します。これにより、以下のことが可能になります。

• ドリフトの早期警告。運用会話でジャッジの合格率が低下した場合、何かが変更されたことを意味します。モデルのアップグレードによって品質が低下したか、ユーザープロンプトがエージェントがうまく処理できない方法で変化したのかもしれません。ユーザーからの苦情が出る前に、ジャッジのスコアでそれを確認できます。
• テストスイートのリアルワールドシグナル。ジャッジが失敗としてフラグを立てた運用会話は、新しい評価シナリオの候補となります。このようにしてテストスイートは有機的に成長します。実際の失敗が評価にフィードバックされ、運用と開発の間のループが閉じられます。
• エージェントレベルでのコスト監視。運用会話におけるトークン使用量とツール呼び出し回数を追跡します。品質に影響を与えない変更であっても、コストが3倍になる場合は依然としてリグレッションです。

重要な洞察は、同じスコアリングインフラストラクチャ（ジャッジ、メトリクス、記録されたトレース）が二重の役割を果たすということです。評価のために一度構築すれば、本番環境の監視はその副次的な効果として得られます。

現状

私たちは、Databricksプラットフォームでリリースしたいくつかのエージェント（例：Genieのデータエンジニアリング、Machine Learning、トレース分析機能）、開発者生産性向上のための社内エージェント、そしてその他の顧客向けエージェント（例：AI Dev Kit、またはOSS MLflow Assistant）全体でこの手法を採用してきました。全体として、具体的なメリットが見られました。

• 手動評価と比較して、自動テストスイートは変更検証にかかる時間を2週間から数時間に短縮しました。これにより、私たちのチームはより高い速度で改善を出荷できるようになりました。
• いくつかのテストスイートは、エージェントごとに数百のテストシナリオにまで成長し、リグレッションを捕捉する信頼性を高めています。
• 統合テストは、依存するインフラストラクチャの変更を検出し、本番環境でのリグレッションを防ぐことを可能にしました。これらの変更の例には、基盤モデルにおけるTODO管理の挙動、レイテンシに影響を与える変更、またはモデルの変更が含まれます。

MLflowはGenAIテストプラットフォームとしても重要な役割を果たしており、エンジニアが手法を標準化し、テスト開発を加速し、チーム間でベストプラクティスを共有するのに役立っています。

まだ機能しないこと

ここでもテストのアナロジーは役立ちます。私たちの限界は、おなじみのテスト問題に当てはまります。

シナリオ生成は手動（テストケースの作成は高コスト）。 スコアリング、アライメント、最適化は自動化しましたが、シナリオ自体の生成は依然として人間の作業です。各シナリオには、現実的な初期状態、意味のあるプロンプト、そして正しい期待値の作成が必要です。これがテストスイートのサイズを制限するボトルネックであり、狭いテストスイートは次の問題に直結します。シナリオ生成の自動化（本番トラフィックパターンやエージェントの仕様から多様で現実的なテストケースを合成すること）は、私たちにとって活発な研究分野です。

コーディングアシスタントは過学習する可能性がある（テストスイートが狭すぎる）。 テストスイートが十分なケースをカバーしていない場合、コーディングアシスタントは特定の入力には優れているが、新しい入力では失敗するエージェント実装を設計してしまいます。これは、単体テストはパスするが本番環境で壊れるコードを書くことのエージェント版です。私たちは、本番環境での失敗を評価にフィードバックし、時間をかけてカバレッジを拡大することでこれを軽減していますが、シナリオ生成が自動化されるまでは、テストスイートの成長は望むよりも遅くなります。

ジャッジのアライメントは高コスト（テストの調整には人手が必要）。 ゴールデンセットの構築には、ドメインエキスパートが手動で出力を評価する必要があり、これこそが私たちが排除しようとしているボトルネックです。しかも、これは一度きりのコストではありません。エージェントが進化するにつれて、ジャッジの再調整が必要になります。私たちは、ジャッジの不確実性を測定し、ジャッジの仕様が不十分で人間のラベルが実際に曖昧さを解消する特定の例を特定することで、これをよりスマートにする方法を調査しています。目標は、ジャッジのアライメントのためのアクティブラーニングです。つまり、専門家にランダムなサンプルを評価してもらうのではなく、ジャッジが不確実であり、ドメインエキスパートの入力がその基準を最も明確にするであろう例のみを提示するのです。

多段階の失敗は原因特定が困難（根本原因分析）。 10ステップのパイプラインのステップ7でエージェントが失敗した場合、根本原因はステップ7にあったのか、それともステップ3にあったのか？私たちのジャッジは症状を捉えますが、コーディングアシスタントは間違った関数を変更してテストの失敗を修正するように、誤ったステップを修正してしまうことがあります。より良い因果トレースは活発な研究分野です。

新しい失敗モードは見過ごされる（カバレッジのギャップ）。 coSTARは、ジャッジがカバーする次元内で最適化を行います。どのジャッジもチェックしない新しい種類の失敗が発生した場合、それは見過ごされます。まるで、どのテストも実行しないコードのバグのようにです。coSTARはテストスイート内で改善しますが、テストスイート自体を拡張することはできません。人間が新しい失敗モードに気づき、ジャッジを追加する必要があります。

主要なポイント

1. エージェント開発にはテストの問題がある。 自動評価がなければ、テストなしでコーディングしていることになり、それ相応のリグレッションが発生するでしょう。
2. ジャッジにはトレースではなくツールを与える。 ターゲットとなるツールを呼び出すエージェント的なジャッジは、集中した単体テストのようなものです。完全なトレースをジャッジにダンプすることは、プログラムの状態をアサーションにダンプするようなものです。それはスケールしません。
3. テストをテストする。 LLMジャッジは確率的です。人間が評価したゴールデンセットに対して、テストスイートを仕様に対して検証するのと同じ方法でアライメントしてください。
4. ループを閉じる。 真の勝利は完全なcoSTARループです。信頼できるシナリオ、記録されたトレース、アライメントされたジャッジ、そしてテストがパスするまでエージェントを洗練するコーディングアシスタント。自動化された洗練なしの評価は、物語の半分に過ぎません。
5. 一度構築し、どこでも監視する。 評価で検証するのと同じジャッジが本番環境を監視できます。一つの投資で二つのリターン。
6. 結合が重要。 エージェントの洗練は、それを駆動するジャッジの信頼性にかかっています。coSTARの二つの結合されたループ — 一つはジャッジへの信頼を築き、もう一つはその信頼を使ってエージェントを洗練する — が、自動化された洗練を単に高速なだけでなく、意味のあるものにしています。

私たちはcoSTARをMLflowの一部として構築しています。もし同様の問題に取り組んでいらっしゃるなら、ぜひお聞かせください。

• coSTAR手法を使用して出荷した機能を見るには、Genie Codeをお試しください。
• 反復的なエージェントの洗練のためにLLMジャッジを定義し使用する方法を始めるには、MLflowのチュートリアルに従ってください。

(このブログ記事はAI翻訳ツールを使用して翻訳されています) 原文記事