ast-grepで関数本体を削ぎ落とす：LLM入力のトークン削減と精度改善
LLMにリポジトリのコードを渡すとき、とりあえずファイルを丸ごと貼っていませんか。関数本体まで含めると、質問に関係ない実装がコンテキストの大半を占め、コストが膨らみ、肝心な情報が中盤に埋もれます。本記事では、Tree-sitterベースの構造検索ツール ast-grep で、関数本体を潰したスケルトンと Git 差分だけを組み合わせる方法を紹介します。ルール定義、CLIパイプライン、削減効果の測り方まで、手元で再現できる形で示します。
丸ごとプロンプト注入の限界：トークン浪費と精度低下の実態
ファイル丸ごとの注入は、2つの面で効率が悪くなります。1つ目はトークン消費です。たとえば「この関数の呼び出し側を修正して」という依頼に必要なのは、呼び出し先の名前、引数、戻り値の型だけです。内部の実装はほぼ不要ですが、丸ごと注入すると数百行ぶんの本体も課金対象になります。
2つ目は精度です。Liu らの論文「Lost in the Middle」(arXiv 2307.03172) は、関連情報が長い入力の中央付近にあると、先頭や末尾にある場合より性能が落ちる傾向を報告しています。モデルや世代によって程度は変わりますが、不要な行を足すほど、必要な行が埋もれやすくなる点は変わりません。
つまり、渡すコードを減らすことは、コスト削減だけでなく品質施策にもなります。ただし、何でも削ればよいわけではありません。型定義やシグネチャは「他のコードとの契約」を表すので残し、実装の詳細は質問に応じて落とす、という切り分けが基本方針です。
なぜ正規表現ではなくast-grepなのか？構文木（CST）ベース抽出の強み
シグネチャ抽出を grep や正規表現で書くと、すぐ壊れます。典型的な失敗は次のとおりです。
- 戻り値型に
{ a: string }のような型リテラルがあり、本体の{と区別できない - 文字列やコメントの中にある
function foo()を誤検出する - ジェネリクスの
<T extends Foo<Bar>>や複数行にまたがる引数リストで、括弧の対応を追えない
ast-grep は Tree-sitter で構文解析し、CST (具象構文木) 上のノードに対してマッチします。ノードの種類 (kind) や親子関係 (inside、has) で条件を書けるので、「関数宣言の body フィールドにあたるブロック」のように構造で指定できます。括弧のネストや文字列内の偽陽性を、自分で考慮する必要がありません。
もうひとつの利点は、検索だけでなくルールに fix を書いて書き換えまで行える点です。今回はこの機能を使い、「抽出」ではなく「本体の置き換え」でスケルトンを作ります。戻り値型や async の有無に依存しないので、パターン列挙より堅牢です。
TypeScript/Pythonコードから型定義と関数シグネチャだけを抜き出すルール記述
方針は、関数やメソッドの本体ブロックだけを省略記号に置き換えることです。interface や type はそのまま残ります。TypeScript 用のルールを rules/strip-bodies.yml に書きます。YAML は --- で複数ルールを並べられます。
id: ts-strip-function-body
language: TypeScript
rule:
kind: statement_block
inside:
field: body
kind: function_declaration
fix: '{ /* ... */ }'
---
id: ts-strip-method-body
language: TypeScript
rule:
kind: statement_block
inside:
field: body
kind: method_definition
fix: '{ /* ... */ }'
---
id: ts-strip-arrow-body
language: TypeScript
rule:
kind: statement_block
inside:
field: body
kind: arrow_function
fix: '{ /* ... */ }'
---
id: py-strip-function-body
language: Python
rule:
kind: block
inside:
field: body
kind: function_definition
fix: '...'inside に field: body を併記すると、「親の body フィールドにあたる子」だけが対象になります。関数内にネストした if のブロックは該当しません。ネストした関数は、外側と内側の両方が置換対象になります。
注意点が3つあります。
- TSX ファイルには
language: Tsxのルールを別途用意します。TypeScript 用ルールは.tsxに適用されません。 - Python の docstring も本体と一緒に消えます。意図の説明が重要なコードベースでは、docstring を残すルールを追加するか、別の手段で補ってください。
- ルールのフィールド指定や
kind名はパーサの文法定義に依存します。手元でast-grep run --debug-query=cstを使い、実際のノード名を確認してから書くと安全です。
CLIパイプラインの構築：Git差分とast-grepを連携させて最小コンテキストを生成する
渡す情報は2種類に分けます。1つは変更箇所そのもの (Git差分)、もう1つは周辺コードの契約 (スケルトン) です。元ファイルを壊さないよう、対象ファイルを一時ディレクトリにコピーしてから --update-all でルールを適用します。
#!/usr/bin/env bash
set -euo pipefail
BASE="${1:-origin/main}"
WORK="$(mktemp -d)"
trap 'rm -rf "$WORK"' EXIT
# 変更ファイルが属するディレクトリを集める
git diff --name-only --diff-filter=AM "$BASE"...HEAD -- '*.ts' '*.py' \
| xargs -n1 dirname | sort -u > "$WORK/dirs.txt"
# 同じディレクトリのファイルをコピーしてスケルトン化
while read -r dir; do
mkdir -p "$WORK/skeleton/$dir"
find "$dir" -maxdepth 1 -type f \( -name '*.ts' -o -name '*.py' \) \
-exec cp {} "$WORK/skeleton/$dir/" \;
done < "$WORK/dirs.txt"
ast-grep scan -r rules/strip-bodies.yml --update-all "$WORK/skeleton" > /dev/null
{
echo "## Changed (diff)"
git diff -U5 "$BASE"...HEAD -- '*.ts' '*.py'
echo
echo "## Surrounding signatures"
find "$WORK/skeleton" -type f | sort | while read -r f; do
echo "### ${f#"$WORK/skeleton/"}"
cat "$f"
done
} > context.md生成した context.md をそのままプロンプトに貼るか、CLI型のエージェントやスクリプトから読み込ませます。差分は -U5 で前後5行の文脈を付けています。リファクタリングなら周辺を広く、レビュー用途なら狭く、と調整してください。
この構成には限界もあります。同一ディレクトリのファイルしか拾わないため、別ディレクトリの型を参照する変更では情報が足りなくなります。その場合は、import 文を ast-grep run -p 'import $$$X from $PATH' -l ts --json で拾って対象ファイルを追加する、といった拡張が必要です。
実プロジェクトでの削減効果検証とプロンプト品質の変化
削減率はリポジトリの性質で大きく変わります。実装が長いファイルが多ければ半分以下になる例もありますが、型定義や宣言ばかりのファイルでは削減は小さくなります。そのため「半減する」と決めつけず、自分のリポジトリで測ってください。トークン数は、使うモデルのトークナイザで数えるのが理想です。簡易的には tiktoken が使えます。
import sys
import tiktoken
enc = tiktoken.get_encoding("o200k_base") # 概算用。モデルごとに値は異なる
def count(path: str) -> int:
with open(path, encoding="utf-8") as f:
return len(enc.encode(f.read()))
for p in sys.argv[1:]:
print(p, count(p))丸ごと注入したプロンプトと context.md を同じタスクで比較します。見る指標は次の3つです。
- 入力トークン数と1回あたりのコスト
- 修正結果がビルドとテストを通った割合
- 存在しない関数名や引数を使った回数 (ハルシネーション)
3つ目は、スケルトンが逆に効くかどうかを判断する材料になります。シグネチャが明示されていると、呼び出し側の引数ミスが減る傾向を期待できます。ただし、これは自分のタスクで確認すべき仮説です。
最後に、使い分けの注意です。バグの原因調査や性能改善のように、実装の中身を読まないと答えが出ないタスクでは、スケルトンだけでは情報が足りません。この場合は2段階にします。まずスケルトンを渡し、モデルに「本体が必要なファイル名」を挙げさせます。次に、そのファイルだけを全文で追加します。この流れなら、コンテキストを最小限に保ちつつ、必要な箇所には実装を渡せます。
