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

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

from __future__ import annotations

import hashlib
import json
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


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


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 _message_content(payload: Mapping[str, Any]) -> str:
    try:
        content = payload["choices"][0]["message"]["content"]
    except (KeyError, IndexError, TypeError) as exc:
        raise ModelGenerationError("model response does not contain choices[0].message.content") from exc
    if isinstance(content, str):
        return content
    if 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"}
        ]
        if parts:
            return "".join(parts)
    raise ModelGenerationError("model response content must be text")


def _json_payload(content: str) -> Any:
    cleaned = re.sub(r"<think>[\s\S]*?</think>", "", content, 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 exc:
        raise ModelGenerationError(
            f"model response is not valid JSON at line {exc.lineno}, column {exc.colno}"
        ) from exc


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) -> list[dict[str, str]]:
    schema_instruction = (
        f"必须只返回 JSON 对象，格式为 {{\"items\":[{{\"instruction\":\"...\","
        f"\"input\":\"...\",\"output\":\"...\"}}]}}；items 必须包含 {count} 条。"
        "instruction 和 output 不得为空，不要输出 Markdown 代码围栏或分析过程。"
    )
    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 <= 5:
        raise ModelGenerationError("qa_pairs_per_item must be in [1, 5]")
    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")

    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 "")
            )
            request_payload: dict[str, Any] = {
                "model": model_name,
                "messages": _prompt_messages(
                    str(config.get("generation_prompt") or ""),
                    content,
                    qa_pairs_per_item,
                ),
                "temperature": temperature,
                "max_tokens": max_tokens,
            }
            if bool(config.get("json_mode", False)):
                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")
                    generated_items = _result_items(_json_payload(_message_content(body)))
                    break
                except (httpx.HTTPError, json.JSONDecodeError, ModelGenerationError) as exc:
                    last_error = exc

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

            for variant_index, value in enumerate(generated_items[:qa_pairs_per_item]):
                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 ""))
                output = normalize_text(
                    str(
                        value.get("output")
                        or value.get("answer")
                        or value.get("response")
                        or ""
                    )
                )
                raw_id = f"{preview_id}:{variant_index + 1}:{instruction}:{output}"
                result_id = f"result_{hashlib.sha256(raw_id.encode()).hexdigest()[:16]}"
                valid = bool(instruction and output)
                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 "model result is missing instruction or output",
                        "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"]