code/02-prompting/json_output.py

72 Zeilen · 2.8 KB

Code und Programmausgaben stehen genau so da, wie sie gelaufen sind – Kommentare und Ausgaben sind daher auf Chinesisch.

"""从用户的求助原话里提取结构化信息:JSON 模式 + Pydantic 校验 + 失败重试。"""
import json
import os

from openai import BadRequestError, OpenAI
from pydantic import BaseModel, ValidationError

client = OpenAI(
    api_key=os.environ["LLM_API_KEY"],
    base_url=os.environ.get("LLM_BASE_URL", "https://api.deepseek.com"),
)
MODEL = os.environ.get("LLM_MODEL", "deepseek-flash")


class BugReport(BaseModel):
    title: str
    httpx_version: str | None  # 原话里没提就是 null
    python_version: str | None
    os: str | None
    error: str | None
    missing_info: list[str]


SYSTEM = """从用户的求助原话中提取信息,输出 JSON。原话里没有的字段填 null,不要猜。
JSON 格式示例:
{"title": "一句话概括问题", "httpx_version": "0.27", "python_version": "3.11",
 "os": "macOS", "error": "报错类型或信息", "missing_info": ["排查还需要知道的信息"]}"""

REPORT = """用httpx下载大文件老是断 我用的stream 下到一半就报错了 ReadTimeout
文件2个G左右,网不太好,公司内网。版本是0.27 python3.11 mac"""


def extract(report, max_attempts=3):
    messages = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": report}]
    for attempt in range(1, max_attempts + 1):
        response = client.chat.completions.create(
            model=MODEL,
            messages=messages,
            response_format={"type": "json_object"},
            max_tokens=1000,
            extra_body={"thinking": {"type": "disabled"}},
        )
        content = response.choices[0].message.content
        try:
            return BugReport.model_validate_json(content), attempt
        except ValidationError as e:
            # 把错误原样告诉模型,让它改正。json 语法错误和字段错误都会走到这里
            print(f"第 {attempt} 次校验失败:{e.errors()[0]['msg']}")
            messages += [
                {"role": "assistant", "content": content or ""},
                {"role": "user", "content": f"你的输出没有通过校验:{e}\n请重新输出完整、正确的 JSON。"},
            ]
    raise RuntimeError(f"{max_attempts} 次都没有得到合法的输出")


report, attempts = extract(REPORT)
print(f"第 {attempts} 次成功:")
print(json.dumps(report.model_dump(), ensure_ascii=False, indent=2))

# 实验:提示词里不写 json 这个词,会怎样?
print()
try:
    client.chat.completions.create(
        model=MODEL,
        messages=[{"role": "user", "content": "提取这段话里的版本号:" + REPORT}],
        response_format={"type": "json_object"},
        extra_body={"thinking": {"type": "disabled"}},
    )
    print("提示词里没有 json,也成功了")
except BadRequestError as e:
    print("提示词里没有 json,报错:", e.message)