Bedrock Tool Useの構造化出力をSSEでトークン単位ストリーミングする — 増分JSONパーサの状態機械設計

{
  "response_type": "results",
  "general_advice": "この問題はVPN接続の認証設定に起因する可能性が...",
  "results": [
    {"issue_id": 12345, "subject": "VPN接続エラー", "summary": "..."},
    {"issue_id": 12346, "subject": "認証設定変更手順", "summary": "..."}
  ],
  "reasoning": "ユーザーの症状からVPN関連のチケットが最も関連..."
}

class ChatToolInput(BaseModel):
    """フィールド順序が重要: Bedrock Tool Useは宣言順にJSONをストリームする。
    general_adviceを先頭近くに置くことで、テキストトークンが即座にフロントエンドに届く。
    """

    response_type: Literal["clarification", "results"]
    general_advice: str = Field(default="", description="担当者向けの回答（マークダウン形式）")
    question: str = ""
    options: list[str] = Field(default_factory=list)
    results: list[TicketResult] = []  # 配列は後ろに配置 → テキスト完了後にストリーム
    reasoning: str = Field(default="")

# JSON escape sequences inside string values
_JSON_ESCAPE_MAP: dict[str, str] = {
    "n": "\n", "t": "\t", "r": "\r",
    '"': '"', "\\": "\\", "/": "/",
    "b": "\b", "f": "\f",
}

# 状態定義
SCANNING = 0           # "general_advice": を探索中
IN_VALUE = 1           # general_advice の文字列値をストリーミング中
SCANNING_RESULTS = 2   # "results":[ を探索中
IN_RESULTS_ARRAY = 3   # results配列内で { または ] を待機中
IN_TICKET_OBJ = 4      # チケットオブジェクト {...} 内、ブレース深度を追跡中
TICKETS_DONE = 5       # results の ] を通過 — チケット処理終了

_MARKER = '"general_advice":'
_RESULTS_MARKER = '"results":['

async def generate_chat_answer(
    bedrock_messages: list[dict],
    source_chunks_count: int,
):
    """Tool Useストリーミングの増分パーサ。

    str（general_adviceの1文字）と TicketResult（完全なチケットオブジェクト）を
    非対称にyieldする。
    """
    tool_schema = ChatToolInput.model_json_schema()

    full_tool_input = ""
    streamed_advice = ""  # max_tokens到達時のフォールバック用

    state = SCANNING
    key_buf = ""
    escape_next = False
    unicode_buf: str | None = None

    # チケットストリーミング用の状態
    results_key_buf = ""
    ticket_buf = ""
    ticket_depth = 0
    ticket_in_str = False
    ticket_esc = False

    async with client.messages.stream(
        model=MODEL_ID,
        messages=sdk_messages,
        max_tokens=8192,
        tools=[{
            "name": "submit_response",
            "description": "Submit the structured helpdesk response",
            "input_schema": tool_schema,
        }],
        tool_choice={"type": "tool", "name": "submit_response"},
    ) as stream:
        async for event in stream:
            if event.type != "input_json":
                continue
            chunk = event.partial_json
            if not chunk:
                continue

            full_tool_input += chunk

            for ch in chunk:
                # ... 状態機械のロジック（後述）
                pass

elif state == IN_VALUE:
    # \uXXXX ユニコードエスケープの処理
    if unicode_buf is not None:
        unicode_buf += ch
        if len(unicode_buf) == 4:
            decoded = chr(int(unicode_buf, 16))
            yield decoded
            streamed_advice += decoded
            unicode_buf = None
        continue

    if escape_next:
        if ch == "u":
            unicode_buf = ""  # \uXXXX の開始
        else:
            decoded = _JSON_ESCAPE_MAP.get(ch, ch)
            yield decoded
            streamed_advice += decoded
        escape_next = False
    elif ch == "\\":
        escape_next = True
    elif ch == '"':
        state = SCANNING_RESULTS  # general_advice完了 → チケット探索へ
    else:
        yield ch
        streamed_advice += ch

elif state == IN_TICKET_OBJ:
    ticket_buf += ch  # 生のJSON文字列を蓄積（エスケープはそのまま）

    if ticket_esc:
        ticket_esc = False
    elif ticket_in_str:
        if ch == "\\":
            ticket_esc = True
        elif ch == '"':
            ticket_in_str = False
    else:
        if ch == '"':
            ticket_in_str = True
        elif ch == "{":
            ticket_depth += 1
        elif ch == "}":
            ticket_depth -= 1
            if ticket_depth == 0:
                # 完全なオブジェクトが揃った → バリデーション＆yield
                try:
                    ticket = TicketResult.model_validate_json(ticket_buf)
                    yield ticket
                except Exception:
                    logger.warning("Failed to parse streamed ticket: %.120s", ticket_buf)
                ticket_buf = ""
                state = IN_RESULTS_ARRAY  # 次のチケットを待機

async for item in generate_chat_answer(bedrock_messages, source_chunks_count):
    if isinstance(item, str):
        yield _sse("token", {"content": item})
    elif isinstance(item, TicketResult):
        yield _sse("ticket", item.model_dump())
    else:
        # ChatResponse（最終結果）
        yield _sse("result", item.model_dump())

SSE_EVENTS: dict[str, tuple[type, str]] = {
    "progress": (SseProgressData, "処理ステップの進捗（analyzing→retrieving→fetching→generating）"),
    "kb_query":  (SseKbQueryData, "Knowledge Base 検索に使用したクエリ文字列"),
    "kb_chunks": (SseKbChunksData, "Knowledge Base から取得したチャンクのプレビュー一覧"),
    "token":     (SseTokenData, "AIが生成するテキストチャンク（ストリーミング）"),
    "ticket":    (TicketResult, "検索結果チケット1件（段階的に配信）"),
    "result":    (ChatResponse, "最終結果（会話ターン完了時に1回送出）"),
    "error":     (SseErrorData, "エラー発生時の詳細メッセージ"),
}

const reader = res.body!.getReader();
const decoder = new TextDecoder();
let buffer = "";

while (true) {
  const { done, value } = await reader.read();
  if (done) break;

  buffer += decoder.decode(value, { stream: true });
  const parts = buffer.split("\n\n");
  buffer = parts.pop() ?? "";

  for (const part of parts) {
    if (!part.trim()) continue;

    let eventType = "message";
    let dataStr = "";
    for (const line of part.split("\n")) {
      if (line.startsWith("event: ")) eventType = line.slice(7).trim();
      else if (line.startsWith("data: ")) dataStr = line.slice(6);
    }

    if (!dataStr) continue;
    const data = JSON.parse(dataStr);

    // イベント種別に応じてUIバブルを更新
    switch (eventType) {
      case "progress": /* プログレスバーを更新 */ break;
      case "token":    /* ストリーミングバブルにテキスト追加 */ break;
      case "ticket":   /* 結果バブルにチケットカード追加 */ break;
      case "result":   /* 最終結果でUI確定 */ break;
      case "error":    /* エラー表示 */ break;
    }
  }
}

// 最初のticket到着時: streaming → results に遷移（テキストを引き継ぐ）
const streamingBubble = prev.find(
  (b) => b.type === "streaming" && b.queryId === queryId,
);
return [
  ...prev.filter((b) => !(b.type === "streaming" && b.queryId === queryId)),
  {
    type: "results",
    results: [ticket],
    general_advice: streamingBubble?.general_advice ?? "",
    // ...
  },
];

try:
    tool_input = ChatToolInput.model_validate_json(full_tool_input)
    result = ChatResponse(...)
except Exception:
    # max_tokens到達でJSONが不完全 → ストリーム済みテキストをフォールバックに使用
    logger.warning("Tool input validation failed (likely max_tokens hit); using streamed advice")
    result = ChatResponse(
        response_type="results",
        general_advice=streamed_advice or "回答の生成中にエラーが発生しました。",
        source_chunks_count=source_chunks_count,
    )
yield result