LLMのツール引数をPythonで検査、JSONの落とし穴

json.loads()を通ったツール引数、そのまま関数に渡して大丈夫？

自分はLLMの引数チェックを書いてて、今日、同じキーを二度入れたJSONを試した。例外は出ず、後ろの値だけが残った。辞書には重複の形跡がない。あ、ここで消えるのか。

ツール引数は、生のJSON文字列から検査したい。検索件数を指定するツールで、辞書化の前後に必要な検査を確かめる。外部接続は不要。CPython 3.12.13で実行した。

辞書になる前に、重複キーを止める

まずは標準の読み込みだけ。

import json

raw = '{"limit": 1, "limit": 9}'
print(json.loads(raw))
print(json.loads(raw, strict=True))

出力は両方とも同じだった。

{'limit': 9}
{'limit': 9}

strict=Trueは、文字列内のエスケープされていない制御文字を許すかどうかの指定。重複キーは拒否しない。この名前、設定画面に出すなら説明を添えたくなる。

Python公式ドキュメントにも、最後の値だけを使うとある。手元のLib/json/decoder.pyのJSONObjectでも、キーと値の組をdict(pairs)へ渡していた。辞書化した後のobject_hookでは、元のキー数に戻れない。

ここで使えるのがobject_pairs_hook。辞書化の前にキーと値の組を受け取り、同じキーを見つけたら例外にする。入れ子のオブジェクトにも呼ばれる。

検査は「生の文字列 → 読み込み時の検査 → 引数の型と範囲 → ツール実行」の順に置く。

SDKが辞書にして返す場合、再びJSON化しても重複は見つからない。生の引数を取得できる箇所か、SDK側のデコード処理を確認する必要がある。

NaNを拒否しても、1e400は別の入口を通る

Pythonのjsonは、標準設定だと仕様外のNaNやInfinityも受け入れる。

import json
import math

for token in ("NaN", "Infinity", "1e400"):
    value = json.loads(token)
    print(token, value, math.isfinite(value))

NaN nan False
Infinity inf False
1e400 inf False

NaNとInfinityはparse_constantで拒否できる。ところが1e400はJSONの数値として書けるので、このフックを通らない。floatへの変換でinfになった。

ここは別に止める。

デコーダーの仕様では、小数点や指数を含む数値はparse_floatが担当する。変換後にmath.isfinite()へ通す。巨大な数を保持したいならDecimalなどの設計が要るけれど、今回は拒否する。

ちなみに、allow_nan=Falseはdumps()側の指定。受信する引数を検査する場所に、そのまま持ってこないようにしたい。

検索ツールの入口を、動く形にする

limitは1〜10の整数に限定し、未知のキーも拒否する。以下は一つのファイルで動く。

import json
import math

def unique_pairs(pairs):
    result = {}
    for key, value in pairs:
        if key in result:
            raise ValueError(f"duplicate: {key}")
        result[key] = value
    return result

def reject_constant(token):
    raise ValueError(f"constant: {token}")

def finite_float(token):
    value = float(token)
    if not math.isfinite(value):
        raise ValueError(f"nonfinite: {token}")
    return value

def load_args(raw):
    args = json.loads(
        raw, object_pairs_hook=unique_pairs,
        parse_constant=reject_constant,
        parse_float=finite_float,
    )
    if type(args) is not dict or set(args) != {"limit"}:
        raise ValueError("keys: limit only")
    if type(args["limit"]) is not int:
        raise ValueError("limit: integer required")
    if not 1 <= args["limit"] <= 10:
        raise ValueError("limit: out of range")
    return args

cases = [
    '{"limit":3}', '{"limit":1,"limit":9}',
    '{"limit":NaN}', '{"limit":1e400}',
    '{"limit":true}', '{"limit":3.0}',
    '{"limit":11}', '{"limit":3,"extra":0}',
]
for raw in cases:
    try:
        print("OK", load_args(raw))
    except ValueError as exc:
        print("NG", exc)

実行結果。

OK {'limit': 3}
NG duplicate: limit
NG constant: NaN
NG nonfinite: 1e400
NG limit: integer required
NG limit: integer required
NG limit: out of range
NG keys: limit only

trueの拒否も意図どおり。boolはintのサブクラスなので、isinstance(value, int)だとTrueも通る。type(value) is intで型を絞った。3.0も拒否する方針だ。小数表記も許したいツールでは、受け付ける条件を別に決める。

この例では有限の小数も最後に拒否する。parse_floatは、小数を扱う別のツールでも使える検査として分けた。

ツール名を許可リストで選んだ後、関数を呼ぶ直前に専用の引数検査を挟む。検査に落ちたら呼ばない。構文エラーのJSONDecodeErrorもValueErrorのサブクラスなので、同じ経路で捕捉できる。

受信サイズの上限は通信側でも決める。このコードには深い入れ子への資源制限や認可は含めていない。検索してよいデータかどうかは、実行側で確かめる。

ツールを呼ぶ前に残すチェック

辞書を受け取った時点で、重複キーの情報は消えていた。生の文字列を受け取れる場所でobject_pairs_hookを使い、数値の変換経路も検査する。通過した引数にも、ツールごとの型と範囲のチェックをかける。検索件数の入口なら、この確認ループで受け付ける入力を具体的に決められる。