| """Helpers for extracting structured JSON output from LLM responses.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| import re |
| from typing import Any, TypeVar, cast |
|
|
| from pydantic import BaseModel |
|
|
| TModel = TypeVar("TModel", bound=BaseModel) |
|
|
|
|
| _CODE_FENCE_RE = re.compile(r"```(?:json)?\s*(.*?)\s*```", re.IGNORECASE | re.DOTALL) |
|
|
|
|
| def clean_text(text: str) -> str: |
| """Remove common wrappers around model output.""" |
| cleaned = text or "" |
| for marker in [ |
| "```json", |
| "```", |
| "<|begin_of_text|>", |
| "<|end_of_text|>", |
| "<|begin_of_box|>", |
| "<|end_of_box|>", |
| ]: |
| cleaned = cleaned.replace(marker, "") |
| return cleaned.strip() |
|
|
|
|
| def extract_json_text(text: str) -> str | None: |
| """Extract the first valid JSON object/array text from a model response.""" |
| cleaned = clean_text(text) |
| if not cleaned: |
| return None |
|
|
| for candidate in [cleaned, *_CODE_FENCE_RE.findall(cleaned)]: |
| snippet = (candidate or "").strip() |
| if not snippet: |
| continue |
| try: |
| json.loads(snippet) |
| return snippet |
| except json.JSONDecodeError: |
| pass |
|
|
| stack: list[str] = [] |
| start_idx: int | None = None |
| for idx, ch in enumerate(cleaned): |
| if ch in "[{": |
| if not stack: |
| start_idx = idx |
| stack.append(ch) |
| continue |
|
|
| if ch not in "]}": |
| continue |
|
|
| if not stack: |
| continue |
|
|
| opening = stack.pop() |
| if (opening == "{" and ch != "}") or (opening == "[" and ch != "]"): |
| stack = [] |
| start_idx = None |
| continue |
|
|
| if not stack and start_idx is not None: |
| snippet = cleaned[start_idx : idx + 1].strip() |
| try: |
| json.loads(snippet) |
| return snippet |
| except json.JSONDecodeError: |
| start_idx = None |
| continue |
|
|
| return None |
|
|
|
|
| def parse_json_content(text: str) -> Any: |
| """Parse JSON payload from model response text.""" |
| payload = extract_json_text(text) |
| if not payload: |
| raise ValueError("LLM response does not contain valid JSON") |
| return json.loads(payload) |
|
|
|
|
| def parse_model_content(text: str, model_cls: type[TModel]) -> TModel: |
| """Parse and validate structured model output from model response text.""" |
| parsed = parse_json_content(text) |
| return cast(TModel, model_cls.model_validate(parsed)) |
|
|
|
|
| def json_schema_instructions(model_cls: type[BaseModel]) -> str: |
| """Return concise JSON-schema instructions for structured generation.""" |
| schema_text = json.dumps(model_cls.model_json_schema(), ensure_ascii=False) |
| return "\n".join( |
| [ |
| "请只输出一个 JSON，不要输出 Markdown、解释或额外文本。", |
| "输出必须严格满足以下 JSON Schema:", |
| schema_text, |
| ] |
| ) |
|
|
