AIが返す文字列の中のJSONとして抽出可能な部分をデコードして返す関数、が意外と便利

はじめに

生成AIが返すテキストで「JSON形式で こういう 感じで返して」と指示をしたときに、JSON部分はあっているんだけど ```json ~ ``` とかコードフェンスで囲ったり、はい とか余計な一言を付けたりすることがよくあります。もちろん、response_format などを指定すればそういう心配はないですが、Modelによっては対応していない場合があったり。
JSONの中身自体は意図通りなので惜しいなぁ、、、と思うこともあるんじゃないでしょうか。

それの対策として、文字列中にあるJSONっぽい部分をよしなに抽出する関数、を作ったのですが思いの外便利に使えることも多いので共有します。

文字列の中のJSONとして抽出可能な部分をデコードして返す関数

シンプルにこんな感じで、想定される応答が dict か list かを指定して、一番大きくJSONとして解釈できる部分を解釈して返す、というものです。

import json
from typing import Union, Dict, List

def extract_json_segment(data: str, target: str = "dict") -> Union[Dict, List]:
    """文字列の中のJSONとして抽出可能な部分をデコードして返す。
    target が "dict" の場合は `{` で始まり `}` で終わる辞書を、
    target が "list" の場合は `[` で始まり `]` で終わる配列を対象とする。
    
    Args:
        data: パース対象の文字列
        target: "dict" または "list" を指定。デフォルトは "dict"
    
    Returns:
        デコードされたJSONオブジェクト。失敗時は空の辞書または空の配列
    """
    if target not in ["dict", "list"]:
        raise ValueError("target must be either 'dict' or 'list'")
    
    start_char = "{" if target == "dict" else "["
    end_char = "}" if target == "dict" else "]"
    default_value = {} if target == "dict" else []
    
    start = 0
    while start < len(data):
        if data[start] == start_char:
            end = len(data)
            while end > start:
                if data[end - 1] == end_char:
                    try:
                        return json.loads(data[start:end])
                    except json.JSONDecodeError:
                        pass
                end -= 1
        start += 1
    return default_value

生成AIとのやりとり時間と比較すれば無視できる実行時間ですし、大抵そんな複雑な話ではないので 問題なくこれで抽出できます。

さいごに

一つ悩みが減りました。