YG_FT/backend/app/modules/data_process/generation.py

"""数据处理任务的大模型生成适配器。"""

from __future__ import annotations

import hashlib
import json
import logging
import re
from collections.abc import Callable, Iterable, Mapping
from typing import Any
from urllib.parse import urlsplit, urlunsplit

import httpx

from app.modules.data_process.algorithms import normalize_text, stable_split_assignments
from app.modules.data_process.constants import (
    MAX_QA_PAIRS_PER_ITEM,
    MODEL_GENERATION_BATCH_SIZE,
)


class ModelGenerationError(ValueError):
    """模型配置、响应或调用失败。"""


class _TerminalModelGenerationError(ModelGenerationError):
    """使用相同参数重试也无法恢复的模型响应错误。"""


OUTPUT_TYPE_STANDARD = "standard"
OUTPUT_TYPE_REASONING = "reasoning"
SUPPORTED_OUTPUT_TYPES = {OUTPUT_TYPE_STANDARD, OUTPUT_TYPE_REASONING}
REASONING_DETAIL_NORMAL = "normal"
REASONING_DETAIL_DETAILED = "detailed"
SUPPORTED_REASONING_DETAILS = {
    REASONING_DETAIL_NORMAL,
    REASONING_DETAIL_DETAILED,
}
MINIMAX_M3_API_HOSTS = {"api.minimax.io", "api.minimaxi.com"}
MINIMAX_M3_MIN_COMPLETION_TOKENS = 4096
logger = logging.getLogger(__name__)


def _is_retryable_generation_error(exc: Exception) -> bool:
    if isinstance(exc, _TerminalModelGenerationError):
        return False
    if isinstance(exc, httpx.HTTPStatusError):
        status_code = exc.response.status_code
        return status_code in {408, 425, 429} or status_code >= 500
    if isinstance(exc, httpx.RequestError):
        return True
    return isinstance(exc, (json.JSONDecodeError, ModelGenerationError))


def _is_official_minimax_m3(endpoint: str, model_name: str) -> bool:
    host = (urlsplit(endpoint).hostname or "").casefold()
    return host in MINIMAX_M3_API_HOSTS and model_name.casefold() == "minimax-m3"


def chat_completions_url(value: str) -> str:
    """把域名、基础 URL 或完整地址统一为 chat completions 地址。"""

    raw = (value or "").strip()
    if not raw:
        raise ModelGenerationError("generation model api_url is required")
    if "://" not in raw:
        raw = f"https://{raw}"
    parsed = urlsplit(raw)
    if parsed.scheme not in {"http", "https"} or not parsed.hostname:
        raise ModelGenerationError("generation model api_url must be an HTTP(S) host or URL")
    if parsed.username or parsed.password:
        raise ModelGenerationError("generation model api_url must not contain credentials")

    path = parsed.path.rstrip("/")
    if path.endswith("/chat/completions"):
        target_path = path
    elif path.endswith("/v1"):
        target_path = f"{path}/chat/completions"
    elif not path:
        target_path = "/v1/chat/completions"
    else:
        target_path = f"{path}/v1/chat/completions"
    return urlunsplit((parsed.scheme, parsed.netloc, target_path, "", ""))


def _response_choice(payload: Mapping[str, Any]) -> Mapping[str, Any]:
    try:
        choice = payload["choices"][0]
    except (KeyError, IndexError, TypeError) as exc:
        raise ModelGenerationError("模型响应缺少 choices[0]") from exc
    if not isinstance(choice, Mapping):
        raise ModelGenerationError("模型响应 choices[0] 不是对象")
    return choice


def _response_finish_reason(payload: Mapping[str, Any]) -> str:
    try:
        return str(_response_choice(payload).get("finish_reason") or "").strip().lower()
    except ModelGenerationError:
        return ""


def _response_content_length(payload: Mapping[str, Any]) -> int:
    try:
        message = _response_choice(payload).get("message")
        if not isinstance(message, Mapping):
            return 0
        content = message.get("content")
        if isinstance(content, str):
            return len(content)
        if isinstance(content, list):
            return sum(
                len(str(item.get("text") or ""))
                for item in content
                if isinstance(item, Mapping)
            )
    except ModelGenerationError:
        pass
    return 0


def _raise_for_terminal_response(payload: Mapping[str, Any]) -> Mapping[str, Any]:
    choice = _response_choice(payload)
    base_response = payload.get("base_resp")
    status_code: Any = None
    status_message = ""
    if isinstance(base_response, Mapping):
        status_code = base_response.get("status_code")
        status_message = re.sub(
            r"\s+", " ", str(base_response.get("status_msg") or "")
        ).strip()[:200]

    if bool(payload.get("input_sensitive")) or status_code in {1026, "1026"}:
        raise _TerminalModelGenerationError(
            f"模型输入触发内容安全拦截（code={status_code or 1026}）"
        )
    if bool(payload.get("output_sensitive")) or status_code in {1027, "1027"}:
        raise _TerminalModelGenerationError(
            f"模型输出触发内容安全拦截（code={status_code or 1027}）"
        )

    finish_reason = str(choice.get("finish_reason") or "").strip().lower()
    if finish_reason == "length":
        raise _TerminalModelGenerationError(
            "模型输出因达到 Token 上限被截断（finish_reason=length），"
            "请提高最大输出长度后重试"
        )
    if finish_reason == "content_filter":
        raise _TerminalModelGenerationError(
            "模型输出被内容安全策略拦截（finish_reason=content_filter）"
        )
    if finish_reason in {"tool_calls", "function_call"}:
        raise _TerminalModelGenerationError(
            f"模型返回了当前生成任务不支持的工具调用（finish_reason={finish_reason}）"
        )
    if status_code not in {None, "", 0, "0"}:
        detail = f"：{status_message}" if status_message else ""
        raise _TerminalModelGenerationError(
            f"模型服务返回业务错误（code={status_code}）{detail}"
        )
    return choice


def _message_content(payload: Mapping[str, Any]) -> str:
    choice = _raise_for_terminal_response(payload)
    message = choice.get("message")
    if not isinstance(message, Mapping):
        raise ModelGenerationError("模型响应缺少 choices[0].message")
    content = message.get("content")
    if isinstance(content, str):
        result = content
    elif isinstance(content, list):
        parts = [
            str(item.get("text") or "")
            for item in content
            if isinstance(item, Mapping) and item.get("type") in {None, "text", "output_text"}
        ]
        result = "".join(parts)
    elif content is None:
        result = ""
    else:
        raise ModelGenerationError("模型响应 content 必须是文本")
    if not result.strip():
        raise ModelGenerationError("模型返回的最终内容为空，未生成可解析的 JSON")
    return result


def _json_documents(content: str) -> list[Any]:
    decoder = json.JSONDecoder()
    documents: list[Any] = []
    cursor = 0
    while cursor < len(content):
        match = re.search(r"[\[{]", content[cursor:])
        if not match:
            break
        start = cursor + match.start()
        try:
            value, end = decoder.raw_decode(content[start:])
        except json.JSONDecodeError:
            cursor = start + 1
            continue
        if isinstance(value, (Mapping, list)):
            documents.append(value)
        cursor = start + max(end, 1)
    return documents


def _json_payload(content: str) -> Any:
    # 只移除模型在 JSON 之前自行输出的思考过程，不能破坏 JSON 字段中的训练内容。
    cleaned = content.strip()
    if re.match(r"^\s*<think>", cleaned, flags=re.IGNORECASE) and not re.match(
        r"^\s*<think>[\s\S]*?</think>", cleaned, flags=re.IGNORECASE
    ):
        raise ModelGenerationError("模型思考内容未闭合，响应可能已被截断")
    cleaned = re.sub(
        r"^\s*(?:<think>[\s\S]*?</think>\s*)+",
        "",
        cleaned,
        count=1,
        flags=re.IGNORECASE,
    ).strip()
    fenced = re.fullmatch(r"```(?:json)?\s*([\s\S]*?)\s*```", cleaned, flags=re.IGNORECASE)
    if fenced:
        cleaned = fenced.group(1).strip()
    try:
        return json.loads(cleaned)
    except json.JSONDecodeError as direct_error:
        documents = _json_documents(cleaned)
        if len(documents) == 1:
            return documents[0]
        if len(documents) > 1:
            raise ModelGenerationError("模型响应包含多个 JSON 对象，无法确定应使用哪一个")
        raise ModelGenerationError(
            "模型响应中没有找到唯一且完整的 JSON 对象"
            f"（第 {direct_error.lineno} 行，第 {direct_error.colno} 列）"
        ) from direct_error


def _result_items(payload: Any) -> list[Mapping[str, Any]]:
    if isinstance(payload, list):
        values = payload
    elif isinstance(payload, Mapping):
        nested = next(
            (
                payload[key]
                for key in ("items", "results", "data", "records")
                if isinstance(payload.get(key), list)
            ),
            None,
        )
        values = nested if isinstance(nested, list) else [payload]
    else:
        raise ModelGenerationError("model JSON must be an object or array")
    items = [item for item in values if isinstance(item, Mapping)]
    if not items:
        raise ModelGenerationError("model JSON does not contain result objects")
    return items


def _prompt_messages(
    prompt: str,
    content: str,
    count: int,
    *,
    start_index: int,
    total_count: int,
    output_type: str,
    reasoning_detail: str,
) -> list[dict[str, str]]:
    end_index = start_index + count - 1
    if output_type == OUTPUT_TYPE_REASONING:
        schema = '{"items":[{"instruction":"...","input":"...","reasoning":"...","answer":"..."}]}'
        detail_rule = (
            "推理详细程度为“详细”：完整展开问题条件、来源依据、中间计算或推导，"
            "并在得出答案前核对结论；每一步都必须能从来源内容中验证。"
            if reasoning_detail == REASONING_DETAIL_DETAILED
            else
            "推理详细程度为“普通”：只保留得出答案所需的关键依据和必要步骤，"
            "避免冗长复述、套话和无依据扩展。"
        )
        output_rule = (
            "你正在生成用于训练推理模型的思维链数据，而不是普通问答数据。"
            "instruction、reasoning 和 answer 均不得为空；reasoning 必须是基于来源内容、"
            f"可核对的推理过程，answer 只写最终答案。{detail_rule}"
            "这是思维链输出模式，即使其他提示语要求省略分析，也不得省略 reasoning。"
            "不要自行添加 <think> 标签，系统会在保存时统一组装。"
        )
    else:
        schema = '{"items":[{"instruction":"...","input":"...","output":"..."}]}'
        output_rule = (
            "你正在生成标准监督微调问答数据。instruction 和 output 不得为空；"
            "output 只写最终答案，禁止输出分析、推理过程或 <think> 标签。"
        )
    schema_instruction = (
        f"必须只返回 JSON 对象，格式为 {schema}；items 必须包含 {count} 条。"
        f"这是总计 {total_count} 条中的第 {start_index}-{end_index} 条，"
        "各条必须使用不同的提问角度和表述，避免重复。"
        f"{output_rule}不要输出 Markdown 代码围栏或 JSON 之外的说明。"
    )
    base_prompt = normalize_text(prompt) or "请根据来源内容生成可用于监督微调的问答数据。"
    if "{{ content }}" in base_prompt:
        user_prompt = base_prompt.replace("{{ content }}", content)
        return [
            {"role": "system", "content": schema_instruction},
            {"role": "user", "content": user_prompt},
        ]
    return [
        {"role": "system", "content": f"{base_prompt}\n{schema_instruction}"},
        {"role": "user", "content": f"来源内容：\n{content}"},
    ]


def generate_model_records(
    preview_items: Iterable[Mapping[str, Any]],
    *,
    model: Mapping[str, Any],
    config: Mapping[str, Any],
    task_id: str,
    split: Mapping[str, int],
    qa_pairs_per_item: int,
    client: httpx.Client | None = None,
    on_progress: Callable[[int, int], None] | None = None,
) -> list[dict[str, Any]]:
    """调用 OpenAI 兼容接口，将预览切片生成标准训练记录。

    每个切片按安全批次调用模型；失败批次会产生一条可人工修复的
    invalid 结果，已经成功的批次不会丢失。
    """

    if not 1 <= qa_pairs_per_item <= MAX_QA_PAIRS_PER_ITEM:
        raise ModelGenerationError(f"qa_pairs_per_item must be in [1, {MAX_QA_PAIRS_PER_ITEM}]")
    output_type = str(config.get("output_type") or OUTPUT_TYPE_STANDARD).strip().lower()
    if output_type not in SUPPORTED_OUTPUT_TYPES:
        raise ModelGenerationError(f"output_type must be one of {sorted(SUPPORTED_OUTPUT_TYPES)}")
    reasoning_detail = str(
        config.get("reasoning_detail") or REASONING_DETAIL_NORMAL
    ).strip().lower()
    if reasoning_detail not in SUPPORTED_REASONING_DETAILS:
        raise ModelGenerationError(
            f"reasoning_detail must be one of {sorted(SUPPORTED_REASONING_DETAILS)}"
        )
    endpoint = chat_completions_url(str(model.get("api_url") or ""))
    model_name = str(model.get("online_model_name") or model.get("name") or "").strip()
    if not model_name:
        raise ModelGenerationError("generation model name is required")
    is_minimax_m3 = _is_official_minimax_m3(endpoint, model_name)

    temperature = float(config.get("temperature", 0.7))
    max_tokens = int(config.get("max_tokens", 1024))
    timeout = max(1.0, min(120.0, float(config.get("request_timeout_seconds", 60))))
    retries = max(0, min(5, int(config.get("generation_retries", 2))))
    headers = {"Content-Type": "application/json"}
    api_key = str(model.get("api_key") or "").strip()
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    owns_client = client is None
    http_client = client or httpx.Client(timeout=timeout)
    results: list[dict[str, Any]] = []
    try:
        preview_list = list(preview_items)
        total_items = len(preview_list)
        for item_index, item in enumerate(preview_list):
            preview_id = str(item.get("id") or f"preview-{item_index + 1}")
            content = normalize_text(
                str(item.get("edited_content") or item.get("original_content") or "")
            )
            for batch_offset in range(0, qa_pairs_per_item, MODEL_GENERATION_BATCH_SIZE):
                batch_count = min(
                    MODEL_GENERATION_BATCH_SIZE,
                    qa_pairs_per_item - batch_offset,
                )
                batch_start = batch_offset + 1
                batch_end = batch_offset + batch_count
                request_payload: dict[str, Any] = {
                    "model": model_name,
                    "messages": _prompt_messages(
                        str(config.get("generation_prompt") or ""),
                        content,
                        batch_count,
                        start_index=batch_start,
                        total_count=qa_pairs_per_item,
                        output_type=output_type,
                        reasoning_detail=reasoning_detail,
                    ),
                    "temperature": temperature,
                }
                if is_minimax_m3:
                    request_payload.update(
                        reasoning_split=True,
                        max_completion_tokens=max(
                            max_tokens,
                            MINIMAX_M3_MIN_COMPLETION_TOKENS,
                        ),
                    )
                else:
                    request_payload["max_tokens"] = max_tokens
                if bool(config.get("json_mode", False)) and not is_minimax_m3:
                    request_payload["response_format"] = {"type": "json_object"}

                last_error: Exception | None = None
                generated_items: list[Mapping[str, Any]] | None = None
                for _ in range(retries + 1):
                    try:
                        response = http_client.post(
                            endpoint,
                            headers=headers,
                            json=request_payload,
                        )
                        response.raise_for_status()
                        body = response.json()
                        if not isinstance(body, Mapping):
                            raise ModelGenerationError("model response body must be a JSON object")
                        try:
                            candidate_items = _result_items(
                                _json_payload(_message_content(body))
                            )
                        except ModelGenerationError as exc:
                            logger.warning(
                                "data process model response rejected task_id=%s model=%s "
                                "finish_reason=%s response_chars=%s input_sensitive=%s "
                                "output_sensitive=%s reason=%s",
                                task_id,
                                model_name,
                                _response_finish_reason(body) or "missing",
                                _response_content_length(body),
                                bool(body.get("input_sensitive")),
                                bool(body.get("output_sensitive")),
                                str(exc),
                            )
                            raise
                        if len(candidate_items) < batch_count:
                            raise ModelGenerationError(
                                "model response contains fewer result objects than requested: "
                                f"expected {batch_count}, got {len(candidate_items)}"
                            )
                        generated_items = candidate_items
                        break
                    except (
                        httpx.HTTPError,
                        json.JSONDecodeError,
                        ModelGenerationError,
                    ) as exc:
                        last_error = exc
                        if not _is_retryable_generation_error(exc):
                            break

                if generated_items is None:
                    error_message = str(last_error or "model generation failed")[:2000]
                    failure_instruction = (
                        f"模型生成失败，请人工补充（第 {batch_start}-{batch_end} 条）"
                    )
                    result_id = (
                        "result_"
                        f"{hashlib.sha256(f'{preview_id}:error:{batch_start}'.encode()).hexdigest()[:16]}"
                    )
                    results.append(
                        {
                            "id": result_id,
                            "preview_item_id": preview_id,
                            "instruction": failure_instruction,
                            "input": content,
                            "output": "",
                            "original_instruction": failure_instruction,
                            "original_input": content,
                            "original_output": "",
                            "status": "invalid",
                            "error": error_message,
                            "split": "train",
                        }
                    )
                    continue

                for batch_index, value in enumerate(generated_items[:batch_count]):
                    variant_index = batch_offset + batch_index
                    instruction = normalize_text(
                        str(value.get("instruction") or value.get("question") or "")
                    )
                    input_text = normalize_text(
                        str(value.get("input") or value.get("context") or "")
                    )
                    if output_type == OUTPUT_TYPE_REASONING:
                        reasoning = normalize_text(
                            re.sub(
                                r"</?think>",
                                "",
                                str(value.get("reasoning") or value.get("analysis") or ""),
                                flags=re.IGNORECASE,
                            )
                        )
                        answer = normalize_text(
                            re.sub(
                                r"</?think>",
                                "",
                                str(
                                    value.get("answer")
                                    or value.get("final_answer")
                                    or value.get("output")
                                    or ""
                                ),
                                flags=re.IGNORECASE,
                            )
                        )
                        output = (
                            f"<think>\n{reasoning}\n</think>\n{answer}"
                            if reasoning and answer
                            else answer or (f"<think>\n{reasoning}\n</think>" if reasoning else "")
                        )
                        valid = bool(instruction and reasoning and answer)
                        missing_error = "model result is missing instruction, reasoning or answer"
                    else:
                        output = normalize_text(
                            str(
                                value.get("output")
                                or value.get("answer")
                                or value.get("response")
                                or ""
                            )
                        )
                        output = normalize_text(
                            re.sub(
                                r"<think>[\s\S]*?(?:</think>|$)",
                                "",
                                output,
                                flags=re.IGNORECASE,
                            )
                        )
                        valid = bool(instruction and output)
                        missing_error = "model result is missing instruction or output"
                    raw_id = f"{preview_id}:{variant_index + 1}:{instruction}:{output}"
                    result_id = f"result_{hashlib.sha256(raw_id.encode()).hexdigest()[:16]}"
                    results.append(
                        {
                            "id": result_id,
                            "preview_item_id": preview_id,
                            "instruction": instruction,
                            "input": input_text,
                            "output": output,
                            "original_instruction": instruction,
                            "original_input": input_text,
                            "original_output": output,
                            "status": "valid" if valid else "invalid",
                            "error": (None if valid else missing_error),
                            "split": "train",
                        }
                    )
            if on_progress:
                on_progress(item_index + 1, total_items)
    finally:
        if owns_client:
            http_client.close()
    assignments = stable_split_assignments(
        [str(result["id"]) for result in results],
        split,
        seed=task_id,
    )
    for result, assignment in zip(results, assignments, strict=True):
        result["split"] = assignment
    return results


__all__ = ["ModelGenerationError", "chat_completions_url", "generate_model_records"]