AIにコードを生成するコードを作らせて、再現性を担保してみたい

TL;DR

この記事の前提の課題
- マイグレーションを直接AI Agentに行わせようとした
- 自然言語のプロンプトでコード完成物を生成させるのは難易度が高かった
- マイグレーション後のコードに変換するためのトランスパイルコードをLLMで書かせた
そのために具体的にやったこと
- AI Agentなどに任せっきりで正しいマイグレーションは難しい（構文的・変数参照など）
- 適用ファイルなどの指定や実際の処理など、一度うまく行った内容でもモデルだったりIDEの変化によって再現が出来ない
- 与えたプロンプトを元にAIが行う振る舞いをコード化することで、再現性を担保することが出来る例を示した
結果的に得られるもの
- コードを作るコードで再現性を担保出来た
- 型との組合せで変更パターンを列挙しやすくなった

はじめに

こんにちは、サポーターズ所属の @y_chu5 です。

古いコードが含まれるアプリケーションをmigrationしたい時、AI Agentにやらせるのが楽そうというのは皆さん体感としてあると思います。
しかし実際にAIに任せてみた結果、「なんか上手くいかない」だったり、「いやそこはそうじゃなくて…」と違和感を感じたことがあるんじゃないでしょうか。

上手くいかなかった、と思うのは例えば以下があってだと思います。

型があっていない
Syntaxがそもそもあっていない
AIに指示して直したものを再度他のファイルに適用していったら再現しなかった

私の所属しているサポーターズではMaterial UIを使用しているのですが、プロダクションコード内に非推奨となったmakeStyles関数が大量に含まれています。これをmigrationしたいのですが、変更対象が数百ファイルにおよび手作業では到底出来そうもなく、AIに任せてmigrationを実施したことがありました。

しかし使用するモデルや実施した時期によっても結果が異なり、段階的なmigrationが非常に難しい状態でした。

どうにかして継続的に、再現性を担保し、正しいSyntaxでmigration出来ないかと頭を悩ませた結果、以下の仮説にたどり着きました。

「コードのmigration方法をコードにおこして使う…それをAIに作らせればいいのでは？」

具体的に以下のようなステップを踏みます。

マイグレーション手段の提示
レビューを受けて改善
実際にマイグレーションスクリプトを実装

この方法を使うことで、結果的に以下2つのメリットを得られました。

コードを作るコードで再現性を担保出来た
型との組合せで変更パターンを列挙しやすくなった

やったこと

プロダクションコードの把握

まず初めに、コードにどのような変更対象が眠っているのかの把握を行いました。
この方法でmigrationしてほしいとAIに頼むのであれば、まず自分が変換したいコードのパターンと種類、そして変換の方向性や方法を知らなければなりません。

そこで、ざっくり関数名でコードを洗い出したり、GitHub Copilotのworkspace検索でパターンの抽出などを行いました。
すると例えばコードでは以下のような使われ方がしていることが判明します。

const useStyles = makeStyles((theme) => ({
    // theme非依存
    originalInput: {
        display: 'none',
    },
    // theme依存
    fileInfo: {
        padding: theme.spacing(1, 2),
        backgroundColor: theme.palette.background.default,
    },
}));

この作成されたカスタムフックをReactコンポーネント内で利用しているようです。
このコードを、自分の手で書き換えるとしたら、以下のような形に書き換えます。

const originalInputStyles: SxProps = {
    display: 'none',
};

const fileInfoStyles: SxProps<Theme> = (theme) => ({
    padding: theme.spacing(1, 2),
    backgroundColor: theme.palette.background.default,
});

このように、入力と出力のパターンをいくつか洗い出しました。

あくまでも一例ですが、簡単のため以後この例のみを対象として進めていきます。

実装のプランニング

洗い出された入出力のパターンを元に、書き換えのためのコード実装のプランニングを行います。
今回このプランニングを行うにあたって、一番気をつけたのは実装の1ステップ毎の粒度です。

ファイルの抽出から書き換えまでを一気通貫で行うコードを書かせることも可能ではあります。
しかし、そうした場合デバッグが非常に困難になるのは今までの経験からわかっているため、1ステップを刻みました。

ステップを刻むことで、自分の見逃していたパターンなども見つかることもあり、デバッグも容易になるためです。

実際の流れとしては以下のようなやり取りになりました。

yamachu: ◯◯を使ってコードを変換するmigrationスクリプトを作成したいです。以下の内容を実現するコードに変換するスクリプトを作成するための流れを考えてください。後述しやすいように、ステップ毎に数字を振ってください。（ここに書き換えパターンや適用したいルール）
AI: 〜〜〜ここにめっちゃプラン
yamachu: では3のステップをもう少し細かいタスクに分解してください。そのステップ内にいくつかの条件分岐があるように思えます。
AI: 〜〜〜ここにいい感じに分解してくれた全体プラン

本事例はTypeScriptが変換対象のコードだったため、ts-morph を使用するように指示しました。
ASTを扱えるツールとして、以前使っていた jscodeshiftでも良かったのですが、ts-morphはTypeScript Compiler APIをヒューマンフレンドリーな形でラップしてくれているため、自分がレビューしたり書き直す際に読みやすい・書きやすいという理由で採用しました。TS Compiler APIに近いのであれば、もしかしたらAIの学習量が多くないかなという期待も込められています。

このようにやり取りを繰り返して、コードを生成するコードを生成するプランを立てました。

実装（させる、する）

あとはAIにだいたいお任せになります。

各ステップに対して、「例えばテストケースとして、〇〇の入力に対して、✗✗の様なものになるのが望まれます」という風に、入出力のペアを一緒に渡して実装のブレをなくすようにしていきます。その際、各ステップを1関数実装というレベルまで落としておき、その関数に対し仕様をコメントとして残すようにしておくと、よりブレがなくなります。

各ステップで生成されたコードのレビューを行い、入出力をテストしていきます。
その際に例外的なパターンもいくつか見つかったため、随時変更をお願いしたり、コード自体を手作業で直すなどの修正も共に行いました。本当に例外中の例外であれば、手でフィットするように直した方が早いこともあります。

例えばこれがステップごとの実装でなかった場合、どのステップでその例外が見つかったのかを判断するのが非常に難しかったと容易に想像できるため、やはり刻んでいくのが良いと考えています。

今回はts-morphなどのコード生成の機構を使っているため、基本的に文法から外れたコードは生まれないことを前提に出来ます。そのため、型チェックがおかしければ網羅できていないパターンが見つかったということになります。その見つけた考慮漏れを解消するために、またAIとおしゃべりして直していきます。

また今回は ts-morph を使っての書き換えを行なっていたため、実際にファイルを出力する前に、以下のコードで型チェックの結果を得ることが出来ました。

const diagnostics = project.getPreEmitDiagnostics();
const currentFileDiagnostics = diagnostics.filter(
  (v) => v.getSourceFile()?.getSymbol() === sourceFile.getSymbol()
);
if (currentFileDiagnostics.length > 0) {
  // ここをAgentに読ませるようにしている
  console.error(
    project.formatDiagnosticsWithColorAndContext(currentFileDiagnostics)
  );
  continue;
}

このコードがある状態で、以下のような指示をAI Agentに与えることで、型を考慮した自律的な修正のループを回すことが回すことが出来るようになっています。

各書き換えを行った後make runを実行して試してください。実行結果にTypeScriptのproblemsが出力されるので、それを元にコードの改善を行ってください。

このように、単純にコードを変換するタスクでは難しいコードの再現担保が、コードを生成させるコード生成によって解消できました。

最後に

この取り組みによって、以下2点のメリットを生み出すことが出来ました。

コードを作るコードで再現性を担保出来た
型との組合せで変更パターンを列挙しやすくなった

AIは完全にコードを理解しているわけでもないし、型を理解しているわけではないです。そこを人間が仕組みとして支援してあげることで再現性を持ってコードを生み出せることが示されたと思います。

例えば今回難しさの一つとなった型の問題は、Language Server自体がMCP Serverとなったら…みたいな解決策も考えられるかもしれません。最近ではScalaのLSP実装であるMetals がMCPサポートを行ったりしました。

AIを取り巻く状況が変われば、もしかしたら今回のような取り組みを経ずとも、再現性を担保する仕組みが出来たり、型をより厳密に見てコードを生成することが出来るかもしれません。しかし、今はそこまで到達していないため、AIの動きを矯正するガードレールを敷いてあげるのをエンジニアがいい感じに頑張っていかないといけないかなと思っています。

うちではこうやって再現性を担保したり、うまくAIと協調しているよという話があればぜひ聞かせてください！