src/openai.py

from __future__ import annotations

import json
from typing import Any

from src.errors import ApiError
from src.models import default_model_id


def extract_openai_payload(body: dict[str, Any]) -> dict[str, Any]:
    if not isinstance(body, dict):
        raise ApiError("invalid_request", "request body must be a JSON object", 400)

    messages = body.get("messages")
    if not messages:
        raise ApiError("invalid_request", "messages is required", 400)

    extra = body.get("extra_body") or {}
    if not isinstance(extra, dict):
        extra = {}

    payload: dict[str, Any] = {
        "model": body.get("model") or default_model_id(),
        "messages": messages,
        "max_tokens": body.get("max_tokens") or body.get("max_completion_tokens"),
        "thinking": extra.get("thinking", body.get("thinking")),
        "diffusion_visual": extra.get("diffusion_visual", False),
        "max_denoising_steps": extra.get("max_denoising_steps"),
        "diffusion_kv_cache": extra.get("diffusion_kv_cache"),
    }

    # Diffusion-specific optional knobs via extra_body.
    for key in [
        "diffusion_eb_t_max",
        "diffusion_eb_t_min",
        "diffusion_eb_entropy_bound",
        "diffusion_eb_confidence",
        "n_gpu_layers",
        "extra_cli_args",
        "prompt_mode",
    ]:
        if key in extra:
            payload[key] = extra[key]

    return payload


def make_chat_completion(request_id: str, created: int, model: str, content: str) -> dict[str, Any]:
    return {
        "id": request_id,
        "object": "chat.completion",
        "created": created,
        "model": model,
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": content},
                "finish_reason": "stop",
            }
        ],
        "usage": {"prompt_tokens": None, "completion_tokens": None, "total_tokens": None},
    }


def make_sse_chunk(
    *,
    request_id: str,
    created: int,
    model: str,
    delta: dict[str, Any],
    finish_reason: str | None,
) -> str:
    payload = {
        "id": request_id,
        "object": "chat.completion.chunk",
        "created": created,
        "model": model,
        "choices": [
            {
                "index": 0,
                "delta": delta,
                "finish_reason": finish_reason,
            }
        ],
    }
    return "data: " + json.dumps(payload, ensure_ascii=False) + "\n\n"


def make_sse_done() -> str:
    return "data: [DONE]\n\n"