ai-ticket-platform/backend/apps/modelhub/structured_output.py

import hashlib
import json
import re
from copy import deepcopy
from typing import Any

CHAT_CAPABILITIES = {"text_chat", "vision_chat"}
STRUCTURED_OUTPUT_MODES = {
    "native_json_schema",
    "json_object",
    "prompt_schema",
}
MODE_STRENGTH = {
    "prompt_schema": 1,
    "json_object": 2,
    "native_json_schema": 3,
}
ENFORCEMENT_STRENGTH = {
    "validated_json": 1,
    "json_object": 2,
    "native_json_schema": 3,
}
NATIVE_JSON_SCHEMA_UNSUPPORTED_KEYWORDS = {
    "allOf",
    "dependentRequired",
    "dependentSchemas",
    "else",
    "if",
    "not",
    "then",
    "uniqueItems",
}

STRUCTURED_OUTPUT_TEST_SCHEMA = {
    "type": "object",
    "properties": {
        "status": {"type": "string", "enum": ["ok"]},
        "items": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "code": {"type": "string", "enum": ["A1"]},
                    "score": {"type": "integer"},
                },
                "required": ["code", "score"],
                "additionalProperties": False,
            },
        },
    },
    "required": ["status", "items"],
    "additionalProperties": False,
}
STRUCTURED_OUTPUT_TEST_EXPECTED = {
    "status": "ok",
    "items": [{"code": "A1", "score": 1}],
}
STRUCTURED_OUTPUT_TEST_PROMPT = (
    "返回结构化能力测试结果：status 必须为 ok，items 只包含一项，"
    "该项 code 为 A1、score 为 1。不要输出解释。"
)


class StructuredOutputError(ValueError):
    def __init__(self, code: str, message: str):
        super().__init__(message)
        self.code = code


def validate_mode_for_capability(capability: str, mode: str) -> None:
    if capability in CHAT_CAPABILITIES:
        if mode not in STRUCTURED_OUTPUT_MODES:
            raise StructuredOutputError(
                "structured_output_not_configured",
                "文本或视觉模型必须配置结构化输出方式",
            )
    # Embedding 与重排不产生业务 JSON，结构化输出不属于它们的运行能力边界。
    # 持久化层会将这两类模型统一归一为 none，此处不以该字段阻断保存或使用。


def ensure_compatible(
    *,
    mode: str,
    test_status: str,
    enforcement: str,
) -> None:
    if mode not in MODE_STRENGTH:
        raise StructuredOutputError(
            "structured_output_not_configured",
            "模型尚未配置结构化输出方式",
        )
    if test_status != "passed":
        raise StructuredOutputError(
            "structured_output_not_verified",
            "模型结构化输出能力尚未验证通过",
        )
    if MODE_STRENGTH[mode] < ENFORCEMENT_STRENGTH.get(enforcement, 1):
        raise StructuredOutputError(
            "structured_output_incompatible",
            f"模型输出模式 {mode} 不满足 Skill 的 {enforcement} 要求",
        )


def schema_hash(schema: dict[str, Any]) -> str:
    canonical = json.dumps(
        schema,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
    )
    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()


def normalize_schema_name(value: str) -> str:
    normalized = re.sub(r"[^A-Za-z0-9_-]+", "_", value).strip("_")
    return (normalized or "structured_output")[:64]


def validate_native_schema_subset(
    schema: dict[str, Any],
    *,
    path: str = "$",
) -> None:
    """校验平台原生 Structured Outputs 采用的可移植 JSON Schema 子集。"""

    for keyword in sorted(NATIVE_JSON_SCHEMA_UNSUPPORTED_KEYWORDS):
        if keyword in schema:
            raise StructuredOutputError(
                "structured_output_schema_unsupported",
                f"原生结构化输出不支持 Schema 关键字 {keyword}: {path}.{keyword}",
            )

    for container in ("properties", "patternProperties", "$defs", "definitions"):
        children = schema.get(container)
        if isinstance(children, dict):
            for name, child in children.items():
                if isinstance(child, dict):
                    validate_native_schema_subset(
                        child,
                        path=f"{path}.{container}.{name}",
                    )

    for keyword in ("additionalProperties", "contains", "items", "propertyNames"):
        child = schema.get(keyword)
        if isinstance(child, dict):
            validate_native_schema_subset(child, path=f"{path}.{keyword}")

    for keyword in ("anyOf", "oneOf", "prefixItems"):
        children = schema.get(keyword)
        if isinstance(children, list):
            for index, child in enumerate(children):
                if isinstance(child, dict):
                    validate_native_schema_subset(
                        child,
                        path=f"{path}.{keyword}[{index}]",
                    )


def build_structured_request(
    messages: list[dict[str, Any]],
    *,
    mode: str,
    output_schema: dict[str, Any],
    schema_name: str,
) -> tuple[list[dict[str, Any]], dict[str, Any] | None, dict[str, Any]]:
    request_messages = deepcopy(messages)
    normalized_name = normalize_schema_name(schema_name)
    digest = schema_hash(output_schema)
    response_format: dict[str, Any] | None = None
    schema_in_system_message = False

    if mode == "native_json_schema":
        validate_native_schema_subset(output_schema)
        response_format = {
            "type": "json_schema",
            "json_schema": {
                "name": normalized_name,
                "strict": True,
                "schema": deepcopy(output_schema),
            },
        }
    elif mode in {"json_object", "prompt_schema"}:
        if mode == "json_object":
            response_format = {"type": "json_object"}
        compact_schema = json.dumps(
            output_schema,
            ensure_ascii=False,
            sort_keys=True,
            separators=(",", ":"),
        )
        request_messages.insert(
            0,
            {
                "role": "system",
                "content": (
                    "你必须只输出一个符合以下 JSON Schema 的 JSON 对象，不得增加解释、"
                    "Markdown 或 Schema 外字段。\n"
                    f"<output_schema>{compact_schema}</output_schema>"
                ),
            },
        )
        schema_in_system_message = True
    else:
        raise StructuredOutputError(
            "structured_output_not_configured",
            f"不支持的结构化输出模式：{mode}",
        )

    metadata = {
        "mode": mode,
        "schema_name": normalized_name,
        "schema_hash": digest,
        "response_format_included": response_format is not None,
        "schema_in_system_message": schema_in_system_message,
    }
    return request_messages, response_format, metadata