guidellm.backends.openai.request_handlers

@OpenAIRequestHandlerFactory.register(
    ["/v1/audio/transcriptions", "/v1/audio/translations"]
)
class AudioRequestHandler(ChatCompletionsRequestHandler):
    """
    Request handler for audio transcription and translation endpoints.

    Processes responses from audio processing APIs that convert speech to text,
    handling both transcription and translation services. Manages audio-specific
    usage metrics including audio tokens and processing duration.

    Example:
    ::
        handler = AudioResponseHandler()
        response = handler.compile_non_streaming(request, api_response)
    """

    def __init__(self):
        """
        Initialize the audio response handler.

        Sets up internal state for accumulating streaming response data including
        audio buffers, text chunks, and usage metrics.
        """
        super().__init__()
        self.streaming_buffer: bytearray = bytearray()

    def format(
        self,
        data: GenerationRequest,
        history: HistoryT[GenerationRequest, GenerationResponse] | None = None,
        **kwargs,
    ) -> GenerationRequestArguments:  # noqa: C901
        """
        Format the audio transcription generation request into the
        appropriate structure.

        :param data: The generation request to format
        :param **kwargs: Additional keyword arguments for request formatting
        :return: The formatted request arguments
        """
        if history:
            raise ValueError("AudioRequestHandler does not support multiturn.")

        arguments = GenerationRequestArguments(files={})
        arguments.body = {}

        # Add model
        if kwargs.get("model") is not None:
            arguments.body["model"] = kwargs["model"]

        # Configure streaming
        if kwargs.get("stream"):
            arguments.stream = True
            arguments.body["stream"] = True
            # NOTE: File upload endpoints use flattened stream options
            arguments.body["stream_include_usage"] = True
            arguments.body["stream_continuous_usage_stats"] = True

        # Apply extra arguments
        if kwargs.get("extras"):
            arguments.model_combine(kwargs["extras"])

        # Build audio input
        audio_columns = data.columns.get("audio_column", [])
        if len(audio_columns) != 1:
            raise ValueError(
                f"GenerativeAudioTranscriptionRequestFormatter expects exactly "
                f"one audio column, but got {len(audio_columns)}."
            )

        arguments.files = {
            "file": (
                audio_columns[0].get("file_name") or "audio_input",
                audio_columns[0].get("audio"),
                audio_columns[0].get("mimetype"),
            )
        }

        return arguments

    def compile_non_streaming(
        self,
        request: GenerationRequest,
        arguments: GenerationRequestArguments,
        response: dict,
    ) -> GenerationResponse:
        """
        Process a complete audio transcription or translation response.

        :param request: Original generation request
        :param arguments: The request arguments that were sent
        :param response: Complete API response containing top-level text and usage
        :return: Standardized GenerationResponse with extracted text and metrics
        """
        text = response.get("text")
        input_metrics, output_metrics = self.extract_metrics(
            response.get("usage"), text
        )

        return GenerationResponse(
            request_id=request.request_id,
            request_args=arguments.model_dump_json(),
            response_id=response.get("id"),
            text=text,
            input_metrics=input_metrics,
            output_metrics=output_metrics,
            response_metrics=self.response_metrics,
        )

    def post_validation(self, response: GenerationResponse) -> None:
        """Reject audio responses that omit the transcription text field.

        An empty string is a valid transcription for silent or VAD-filtered audio.

        :param response: The compiled audio response to validate.
        :raises ValueError: If the response has no transcription text field.
        """
        if response.text is None:
            raise ValueError(
                "[UNUSABLE_BACKEND_RESPONSE] backend resolved with empty response "
                "payload"
            )

    def extract_metrics(
        self, usage: dict[str, int | dict[str, int]] | None, text: str | None
    ) -> tuple[UsageMetrics, UsageMetrics]:
        """
        Extract input and output usage metrics from audio API response usage data.

        Handles audio-specific metrics including processing duration and audio tokens
        in addition to standard text token counts.

        :param usage: Usage data dictionary from audio API response
        :param text: Generated text for calculating word and character counts.
            None means text is not applicable (metrics will be None);
            empty string means text was applicable but empty (metrics will be 0).
        :return: Tuple of input_metrics and output_metrics as UsageMetrics objects
        """
        if text is None:
            # text not applicable (e.g. tool-call-only) — exclude from aggregation
            text_words = None
            text_chars = None
        else:
            text_words = len(text.split())
            text_chars = len(text)

        if not usage:
            return UsageMetrics(), UsageMetrics(
                text_words=text_words,
                text_characters=text_chars,
            )

        usage_metrics: dict[str, int] = cast("dict[str, int]", usage)

        return UsageMetrics(
            audio_tokens=(usage_metrics.get("prompt_tokens") or 0),
        ), UsageMetrics(
            text_tokens=(usage_metrics.get("completion_tokens") or 0),
            text_words=text_words,
            text_characters=text_chars,
        )

@OpenAIRequestHandlerFactory.register("/v1/chat/completions")
class ChatCompletionsRequestHandler(TextCompletionsRequestHandler):
    """
    Request handler for OpenAI-style chat completion endpoints.

    Extends TextCompletionsResponseHandler to handle chat completion requests where
    generated text is nested within message objects in the choices array. Processes
    both streaming and non-streaming chat completion responses, including tool call
    responses where the model outputs ``tool_calls`` instead of text content.
    """

    def __init__(self):
        super().__init__()
        # Full tool call payloads accumulated across streaming deltas,
        # keyed by the delta ``index`` field.  Needed for multi-turn tool
        # calling so the response carries the id/name/arguments of each call.
        self.streaming_tool_calls: dict[int, ToolCall] = {}
        self.streaming_reasoning_texts: list[str] = []
        self._last_iteration_had_content: bool = False

    @property
    def last_iteration_had_content(self) -> bool:
        """
        :return: True if the last chunk carried output (text/tool-call) tokens,
            not solely reasoning tokens.
        """
        return self._last_iteration_had_content

    @staticmethod
    def _ensure_tool_format(tool: dict[str, Any]) -> dict[str, Any]:
        """Normalise a single tool definition to Chat Completions format.

        Chat Completions expects
        ``{"type": "function", "function": {"name": ..., ...}}``.
        If the tool is already in that format it is returned as-is.  If it is in the
        flat Responses API format (top-level ``name``, no ``function`` key) the detail
        fields are wrapped into a nested ``function`` dict.

        :param tool: A single tool definition dict in either format.
        :return: The tool in Chat Completions format.
        """
        if "name" in tool and "function" not in tool:
            fn = {k: tool[k] for k in _FUNCTION_DETAIL_KEYS if k in tool}
            return {"type": tool.get("type", "function"), "function": fn}
        return tool

    def _format_prompts(
        self,
        column_data: list,
        column_type: str,
        content_extras: dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        """
        Helper method to format different types of data columns
        into the appropriate structure for chat messages.
        """
        formatted_data = []
        for item in column_data:
            if column_type == "text_column":
                content = {"type": "text", "text": item}
                if content_extras:
                    content.update(content_extras)
                formatted_data.append(content)
            elif column_type == "image_column":
                formatted_data.append(
                    {
                        "type": "image_url",
                        "image_url": {"url": item.get("image")},
                    }
                )
            elif column_type == "video_column":
                formatted_data.append(
                    {
                        "type": "video_url",
                        "video_url": {"url": item.get("video")},
                    }
                )
            elif column_type == "audio_column":
                formatted_data.append(
                    {
                        "type": "input_audio",
                        "input_audio": {
                            "data": base64.b64encode(item.get("audio", b"")).decode(
                                "utf-8"
                            ),
                            "format": item.get("format"),
                        },
                    }
                )
            else:
                raise ValueError(f"Unsupported column type: {column_type}")

        return formatted_data

    @staticmethod
    def _build_tool_response_messages(
        tool_calls: list[ToolCall],
        tool_response_columns: list[Any],
    ) -> list[dict[str, Any]]:
        """Build synthetic ``role: "tool"`` messages for each tool call.

        Uses per-request tool response content from ``tool_response_columns``
        when available, falling back to
        :attr:`settings.default_synthetic_tool_response`.

        :param tool_calls: The tool call objects from the prior assistant response.
        :param tool_response_columns: Per-tool-call response content from the
            dataset, which may be ``str`` or ``bytes`` (orjson).
        :return: List of tool-role message dicts ready to append to messages.
        """
        messages: list[dict[str, Any]] = []
        for idx, tc in enumerate(tool_calls):
            # The OpenAI spec allows the server to do multiple tool calls.
            # If the quantity of calls exceeds those in the dataset, inject the
            # default synthetic tool response.
            raw_content = (
                tool_response_columns[idx]
                if idx < len(tool_response_columns)
                else settings.default_synthetic_tool_response
            )
            # The project JSON utils can return bytes or string; ensure string.
            content = (
                raw_content.decode("utf-8")
                if isinstance(raw_content, bytes)
                else raw_content
            )
            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": tc.id,
                    "content": content,
                }
            )
        return messages

    @staticmethod
    def _apply_tool_call_overrides(
        body: dict[str, Any],
        data: GenerationRequest,
    ) -> None:
        """Inject tool definitions and constrain the request body for tool calling.

        Handles three concerns:

        1. Deserializes and injects tool definitions from dataset columns.
        2. Sets ``tool_choice`` from ``tool_choice_column`` when present,
           otherwise ``"required"``, or ``"none"`` on turns that must not
           emit tool calls.
        3. Removes body keys that are incompatible with tool calling
           (``ignore_eos``, ``stop``, and token-limit keys on tool-call turns).

        :param body: The mutable request body dict being built.
        :param data: The current generation request.
        """
        tools_column = data.columns.get("tools_column", [])
        if tools_column:
            tools_value = tools_column[0]
            # JSON-serialized tool definitions (e.g. from synthetic data
            # generators that store tools as strings for HuggingFace Features
            # compatibility). The project JSON utils can return bytes or str.
            if isinstance(tools_value, str | bytes):
                tools_value = json.loads(tools_value)
            if isinstance(tools_value, list):
                body["tools"] = [
                    ChatCompletionsRequestHandler._ensure_tool_format(t)
                    for t in tools_value
                ]
                body.setdefault("tool_choice", _dataset_tool_choice(data))

        _constrain_tool_call_body(body, data, "max_completion_tokens")

    def _append_request_content_messages(
        self,
        messages: list[dict[str, Any]],
        req: GenerationRequest,
        **kwargs,
    ) -> None:
        """Append current-turn content: recorded chat dicts or text/multimodal.

        When ``raw_messages_column`` is set, those dicts are extended as-is so
        multi-turn replay can send user, assistant, and tool messages without
        wrapping them as a single user string.

        :param messages: Message list to extend in place.
        :param req: Request whose columns supply the current turn.
        :param kwargs: Forwarded format kwargs (``extras`` for multimodal).
        """
        recorded = _recorded_chat_messages(req.columns)
        if recorded is not None:
            messages.extend(recorded)
            return

        prefix = " ".join(req.columns.get("prefix_column", []))
        if prefix:
            messages.append({"role": "system", "content": prefix})

        extras = kwargs.get("extras")
        content_extras = extras.content if extras is not None else None
        prompts = [
            self._format_prompts(
                req.columns.get(col, []),
                col,
                content_extras,
            )
            for col in (
                "text_column",
                "image_column",
                "video_column",
                "audio_column",
            )
        ]
        user_content = list(roundrobin(*prompts))
        if user_content:
            messages.append({"role": "user", "content": user_content})

    def _build_history_messages(
        self,
        history: HistoryT[GenerationRequest, GenerationResponse],
        **kwargs,
    ) -> list[dict[str, Any]]:
        """Build the messages array from completed conversation turns.

        :param history: Completed (request, response) pairs.
        :param kwargs: Forwarded config (``multiturn_reasoning``, etc.).
        :return: Flat list of message dicts ready to extend into the body.
        """
        messages: list[dict[str, Any]] = []
        for idx, (req, res) in enumerate(history):
            # Passes in the prior response so that past data can be included,
            # like for tool calling.
            prior_response = history[idx - 1][1] if idx > 0 else None
            messages.extend(
                self._build_turn_messages(req, res, prior_response, **kwargs)
            )
        return messages

    def _build_turn_messages(  # noqa: C901
        self,
        req: GenerationRequest,
        res: GenerationResponse | None,
        prior_response: GenerationResponse | None,
        **kwargs,
    ) -> list[dict[str, Any]]:
        """Build messages for a single history turn.

        Dispatches on ``req.turn_type``:

        * ``"client_tool_call"``: user content + assistant tool_calls (no tool
          response messages — those come from the following injection turn).
        * ``"tool_response_injection"``: tool-role messages built from
          ``prior_response.tool_calls`` IDs + this turn's
          ``tool_response_column``, then the assistant text response.
        * ``"standard"``: user content + assistant text response.

        :param req: The request for this history turn.
        :param res: The response the server gave for this turn.
        :param prior_response: The response from the immediately preceding
            history turn (used by injection turns for ``tool_call_id``s).
        :param kwargs: Forwarded config (``multiturn_reasoning``, etc.).
        :return: List of message dicts for this turn.
        """
        messages: list[dict[str, Any]] = []
        multiturn_reasoning = kwargs.get("multiturn_reasoning", False)

        if req.turn_type == "tool_response_injection":
            # Injection turn: tool-role messages then assistant text.
            if prior_response and prior_response.tool_calls:
                tool_response_columns = req.columns.get("tool_response_column", [])
                messages.extend(
                    self._build_tool_response_messages(
                        prior_response.tool_calls, tool_response_columns
                    )
                )
            if res is not None and res.text is not None:
                wrapped = _wrap_reasoning(res.reasoning_text, multiturn_reasoning)
                content = res.text
                if wrapped:
                    content = wrapped + content
                messages.append({"role": "assistant", "content": content})
        else:
            # Standard or tool_call turn: recorded messages, or system + user content.
            self._append_request_content_messages(messages, req, **kwargs)

            # Assistant response for history replay.
            wrapped = _wrap_reasoning(
                res.reasoning_text if res else None, multiturn_reasoning
            )
            if res is not None:
                if res.tool_calls:
                    # Tool-call turn: assistant with tool_calls only.
                    # Tool response messages come from the injection turn.
                    assistant_content = res.text
                    if wrapped:
                        assistant_content = wrapped + (assistant_content or "")
                    messages.append(
                        {
                            "role": "assistant",
                            "content": assistant_content,
                            "tool_calls": [tc.model_dump() for tc in res.tool_calls],
                        }
                    )
                elif res.text is not None or wrapped:
                    content = res.text or ""
                    if wrapped:
                        content = wrapped + content
                    messages.append({"role": "assistant", "content": content})

        return messages

    def format(  # noqa: C901, PLR0912, PLR0915
        self,
        data: GenerationRequest,
        history: HistoryT[GenerationRequest, GenerationResponse] | None = None,
        **kwargs,
    ) -> GenerationRequestArguments:
        """
        Format the chat completion generation request into the appropriate structure.

        :param data: The generation request to format
        :param history: Prior (request, response) pairs in the conversation
        :param **kwargs: Additional keyword arguments for request formatting
        :return: The formatted request arguments
        :raises ValueError: If ``messages`` cannot be built from columns, extras,
            history, or a tool-injection turn.
        """
        arguments = GenerationRequestArguments()
        arguments.body = {}  # The type checker works best with body assigned here

        # Add model
        if kwargs.get("model") is not None:
            arguments.body["model"] = kwargs["model"]

        # Configure streaming
        if kwargs.get("stream"):
            arguments.stream = True
            arguments.body["stream"] = True
            arguments.body["stream_options"] = {
                "include_usage": True,
                "continuous_usage_stats": True,
            }

        # Handle output tokens
        if data.output_metrics.text_tokens:
            arguments.body.update(
                {
                    "max_completion_tokens": data.output_metrics.text_tokens,
                    "stop": None,
                    "ignore_eos": True,
                }
            )
        elif kwargs.get("max_tokens") is not None:
            arguments.body["max_completion_tokens"] = kwargs["max_tokens"]

        # Apply extra arguments
        if kwargs.get("extras"):
            arguments.model_combine(kwargs["extras"])

        # Build messages from history
        arguments.body["messages"] = (
            self._build_history_messages(history, **kwargs) if history else []
        )

        # Build the current turn's messages
        if data.turn_type == "tool_response_injection":
            # Injection turn: send tool results back to the server.
            # tool_call_ids come from the last history entry's response.
            prior_response = history[-1][1] if history else None
            if prior_response and prior_response.tool_calls:
                tool_response_columns = data.columns.get("tool_response_column", [])
                arguments.body["messages"].extend(
                    self._build_tool_response_messages(
                        prior_response.tool_calls, tool_response_columns
                    )
                )
        else:
            # Standard or tool_call turn: recorded messages, or system + user content.
            self._append_request_content_messages(
                arguments.body["messages"], data, **kwargs
            )

        # Inject tool definitions and apply tool-call-specific overrides.
        self._apply_tool_call_overrides(arguments.body, data)

        _require_request_payload(
            arguments,
            endpoint="/v1/chat/completions",
            field="messages",
            sources=(
                "raw_messages_column, text_column, prefix_column, or media columns"
            ),
            columns=data.columns,
        )
        return arguments

    def compile_non_streaming(
        self,
        request: GenerationRequest,
        arguments: GenerationRequestArguments,
        response: dict,
    ) -> GenerationResponse:
        """
        Process a complete chat completion response.

        Extracts content from the message object within choices, handling the nested
        structure specific to chat completion endpoints.

        :param request: Original generation request
        :param arguments: The request arguments that were sent
        :param response: Complete API response containing choices and usage data
        :return: Standardized GenerationResponse with extracted content and metrics
        """
        choices, usage = self.extract_choices_and_usage(response)
        choice: dict[str, dict] = choices[0] if choices else {}
        message = choice.get("message", {})
        text = message.get("content")
        reasoning_text = (
            message.get("reasoning") or message.get("reasoning_content") or None
        )
        raw_tool_calls = message.get("tool_calls")
        if text is None and not raw_tool_calls:
            text = ""  # Edge case: null content and no tools
        input_metrics, output_metrics = self.extract_metrics(usage, text)

        tool_calls: list[ToolCall] | None = None
        if raw_tool_calls:
            tool_calls = [ToolCall.model_validate(tc) for tc in raw_tool_calls]
        _apply_tool_call_metrics(
            output_metrics, len(tool_calls) if tool_calls else 0, text
        )

        return GenerationResponse(
            request_id=request.request_id,
            request_args=arguments.model_dump_json(),
            response_id=response.get("id"),  # use vLLM ID if available
            text=text,
            reasoning_text=reasoning_text,
            tool_calls=tool_calls,
            input_metrics=input_metrics,
            output_metrics=output_metrics,
            response_metrics=self.response_metrics,
        )

    def add_streaming_line(self, line: str) -> int | None:
        """
        Process a single line from a chat completion streaming response.

        Handles the chat completion specific delta structure where content is nested
        within delta objects in the streaming response chunks. Also accumulates
        ``tool_calls`` deltas when the model streams function call output.

        :param line: Raw SSE line from the streaming response
        :return: 1 if content was extracted, 0 if line ignored, None if done
        """
        if not (data := self.extract_line_data(line)):
            return None if data is None else 0

        if "id" in data and self.streaming_response_id is None:
            self.streaming_response_id = data["id"]

        updated = False
        # Tracks whether this iteration produced user-visible content (not
        # reasoning-only). Used by the HTTP layer to set TTFOT timing.
        had_content = False
        choices, usage = self.extract_choices_and_usage(data)
        choice: dict[str, dict] = choices[0] if choices else {}
        delta = choice.get("delta", {}) if choices else {}

        # Reasoning tokens trigger TTFT (updated=True) but are not
        # considered "content" for the TTFOT metric.
        if reasoning := (delta.get("reasoning") or delta.get("reasoning_content")):
            self.streaming_reasoning_texts.append(reasoning)
            updated = True
        if content := delta.get("content"):
            self.streaming_texts.append(content)
            updated = True
            had_content = True

        # Accumulate streamed tool_calls deltas.  Each tool call may be split
        # across multiple chunks; we reassemble by ``index``.
        # ``tool_calls`` could be either missing or ``null`` in the delta
        # (some OpenAI-compatible servers emit this), handle both cases
        for tc_delta in delta.get("tool_calls") or []:
            self._accumulate_tool_call_delta(tc_delta)
            updated = True
            had_content = True

        if usage:
            self.streaming_usage = usage

        # Only update the flag when we processed a token-bearing iteration;
        # non-updating lines should not reset the flag.
        if updated:
            self._last_iteration_had_content = had_content
        return 1 if updated else 0

    def _accumulate_tool_call_delta(self, tc_delta: dict[str, Any]) -> None:
        """Merge a single streaming tool_call delta into accumulated state.

        Each tool call is split across multiple SSE chunks.  This method
        creates or updates the :class:`ToolCall` entry keyed by the
        delta's ``index`` field.

        :param tc_delta: A single element from the ``tool_calls`` array in a
            streaming chat completion delta.
        """
        idx = tc_delta["index"]

        if idx not in self.streaming_tool_calls:
            self.streaming_tool_calls[idx] = ToolCall(
                id=tc_delta.get("id", ""),
                type=tc_delta.get("type", "function"),
            )

        tc = self.streaming_tool_calls[idx]
        fn_delta = tc_delta.get("function", {})

        if fn_id := tc_delta.get("id"):
            tc.id = fn_id
        if fn_name := fn_delta.get("name"):
            tc.function.name += fn_name
        if fn_args := fn_delta.get("arguments"):
            tc.function.arguments += fn_args

    def compile_streaming(
        self, request: GenerationRequest, arguments: GenerationRequestArguments
    ) -> GenerationResponse:
        """
        Compile accumulated streaming chat completion content into a final response.

        :param request: Original generation request
        :return: Standardized GenerationResponse with concatenated content and metrics
        """
        return _compile_streaming_response(
            request,
            arguments,
            self.streaming_texts,
            self.streaming_tool_calls,
            self.streaming_usage,
            self.streaming_response_id,
            self.extract_metrics,
            streaming_reasoning_texts=self.streaming_reasoning_texts,
            response_metrics=self.response_metrics,
        )