File size: 2,939 Bytes
b4f2238 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | """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,
]
)
|