chatcmpl_stream_handler.py

from __future__ import annotations

from collections.abc import AsyncIterator
from dataclasses import dataclass, field
from typing import Any

from openai import AsyncStream
from openai.types.chat import ChatCompletionChunk
from openai.types.completion_usage import CompletionUsage
from openai.types.responses import (
    Response,
    ResponseCompletedEvent,
    ResponseContentPartAddedEvent,
    ResponseContentPartDoneEvent,
    ResponseCreatedEvent,
    ResponseFunctionCallArgumentsDeltaEvent,
    ResponseFunctionToolCall,
    ResponseOutputItem,
    ResponseOutputItemAddedEvent,
    ResponseOutputItemDoneEvent,
    ResponseOutputMessage,
    ResponseOutputRefusal,
    ResponseOutputText,
    ResponseReasoningItem,
    ResponseReasoningSummaryPartAddedEvent,
    ResponseReasoningSummaryPartDoneEvent,
    ResponseReasoningSummaryTextDeltaEvent,
    ResponseRefusalDeltaEvent,
    ResponseTextDeltaEvent,
    ResponseUsage,
)
from openai.types.responses.response_reasoning_item import Content, Summary
from openai.types.responses.response_reasoning_summary_part_added_event import (
    Part as AddedEventPart,
)
from openai.types.responses.response_reasoning_summary_part_done_event import Part as DoneEventPart
from openai.types.responses.response_reasoning_text_delta_event import (
    ResponseReasoningTextDeltaEvent,
)
from openai.types.responses.response_reasoning_text_done_event import (
    ResponseReasoningTextDoneEvent,
)
from openai.types.responses.response_usage import InputTokensDetails, OutputTokensDetails

from ..items import TResponseStreamEvent
from .fake_id import FAKE_RESPONSES_ID


# Define a Part class for internal use
class Part:
    def __init__(self, text: str, type: str):
        self.text = text
        self.type = type


@dataclass
class StreamingState:
    started: bool = False
    text_content_index_and_output: tuple[int, ResponseOutputText] | None = None
    refusal_content_index_and_output: tuple[int, ResponseOutputRefusal] | None = None
    reasoning_content_index_and_output: tuple[int, ResponseReasoningItem] | None = None
    function_calls: dict[int, ResponseFunctionToolCall] = field(default_factory=dict)
    # Fields for real-time function call streaming
    function_call_streaming: dict[int, bool] = field(default_factory=dict)
    function_call_output_idx: dict[int, int] = field(default_factory=dict)
    # Store accumulated thinking text and signature for Anthropic compatibility
    thinking_text: str = ""
    thinking_signature: str | None = None
    # Store thought signatures for Gemini function calls (indexed by tool call index)
    function_call_thought_signatures: dict[int, str] = field(default_factory=dict)


class SequenceNumber:
    def __init__(self):
        self._sequence_number = 0

    def get_and_increment(self) -> int:
        num = self._sequence_number
        self._sequence_number += 1
        return num


class ChatCmplStreamHandler:
    @classmethod
    async def handle_stream(
        cls,
        response: Response,
        stream: AsyncStream[ChatCompletionChunk],
    ) -> AsyncIterator[TResponseStreamEvent]:
        usage: CompletionUsage | None = None
        state = StreamingState()
        sequence_number = SequenceNumber()
        async for chunk in stream:
            if not state.started:
                state.started = True
                yield ResponseCreatedEvent(
                    response=response,
                    type="response.created",
                    sequence_number=sequence_number.get_and_increment(),
                )

            # This is always set by the OpenAI API, but not by others e.g. LiteLLM
            usage = chunk.usage if hasattr(chunk, "usage") else None

            if not chunk.choices or not chunk.choices[0].delta:
                continue

            delta = chunk.choices[0].delta

            # Handle thinking blocks from Anthropic (for preserving signatures)
            if hasattr(delta, "thinking_blocks") and delta.thinking_blocks:
                for block in delta.thinking_blocks:
                    if isinstance(block, dict):
                        # Accumulate thinking text
                        thinking_text = block.get("thinking", "")
                        if thinking_text:
                            state.thinking_text += thinking_text
                        # Store signature if present
                        signature = block.get("signature")
                        if signature:
                            state.thinking_signature = signature

            # Handle reasoning content for reasoning summaries
            if hasattr(delta, "reasoning_content"):
                reasoning_content = delta.reasoning_content
                if reasoning_content and not state.reasoning_content_index_and_output:
                    state.reasoning_content_index_and_output = (
                        0,
                        ResponseReasoningItem(
                            id=FAKE_RESPONSES_ID,
                            summary=[Summary(text="", type="summary_text")],
                            type="reasoning",
                        ),
                    )
                    yield ResponseOutputItemAddedEvent(
                        item=ResponseReasoningItem(
                            id=FAKE_RESPONSES_ID,
                            summary=[Summary(text="", type="summary_text")],
                            type="reasoning",
                        ),
                        output_index=0,
                        type="response.output_item.added",
                        sequence_number=sequence_number.get_and_increment(),
                    )

                    yield ResponseReasoningSummaryPartAddedEvent(
                        item_id=FAKE_RESPONSES_ID,
                        output_index=0,
                        summary_index=0,
                        part=AddedEventPart(text="", type="summary_text"),
                        type="response.reasoning_summary_part.added",
                        sequence_number=sequence_number.get_and_increment(),
                    )

                if reasoning_content and state.reasoning_content_index_and_output:
                    # Ensure summary list has at least one element
                    if not state.reasoning_content_index_and_output[1].summary:
                        state.reasoning_content_index_and_output[1].summary = [
                            Summary(text="", type="summary_text")
                        ]

                    yield ResponseReasoningSummaryTextDeltaEvent(
                        delta=reasoning_content,
                        item_id=FAKE_RESPONSES_ID,
                        output_index=0,
                        summary_index=0,
                        type="response.reasoning_summary_text.delta",
                        sequence_number=sequence_number.get_and_increment(),
                    )

                    # Create a new summary with updated text
                    current_content = state.reasoning_content_index_and_output[1].summary[0]
                    updated_text = current_content.text + reasoning_content
                    new_content = Summary(text=updated_text, type="summary_text")
                    state.reasoning_content_index_and_output[1].summary[0] = new_content

            # Handle reasoning content from 3rd party platforms
            if hasattr(delta, "reasoning"):
                reasoning_text = delta.reasoning
                if reasoning_text and not state.reasoning_content_index_and_output:
                    state.reasoning_content_index_and_output = (
                        0,
                        ResponseReasoningItem(
                            id=FAKE_RESPONSES_ID,
                            summary=[],
                            content=[Content(text="", type="reasoning_text")],
                            type="reasoning",
                        ),
                    )
                    yield ResponseOutputItemAddedEvent(
                        item=ResponseReasoningItem(
                            id=FAKE_RESPONSES_ID,
                            summary=[],
                            content=[Content(text="", type="reasoning_text")],
                            type="reasoning",
                        ),
                        output_index=0,
                        type="response.output_item.added",
                        sequence_number=sequence_number.get_and_increment(),
                    )

                if reasoning_text and state.reasoning_content_index_and_output:
                    yield ResponseReasoningTextDeltaEvent(
                        delta=reasoning_text,
                        item_id=FAKE_RESPONSES_ID,
                        output_index=0,
                        content_index=0,
                        type="response.reasoning_text.delta",
                        sequence_number=sequence_number.get_and_increment(),
                    )

                    # Create a new summary with updated text
                    if not state.reasoning_content_index_and_output[1].content:
                        state.reasoning_content_index_and_output[1].content = [
                            Content(text="", type="reasoning_text")
                        ]
                    current_text = state.reasoning_content_index_and_output[1].content[0]
                    updated_text = current_text.text + reasoning_text
                    new_text_content = Content(text=updated_text, type="reasoning_text")
                    state.reasoning_content_index_and_output[1].content[0] = new_text_content

            # Handle regular content
            if delta.content is not None:
                if not state.text_content_index_and_output:
                    content_index = 0
                    if state.reasoning_content_index_and_output:
                        content_index += 1
                    if state.refusal_content_index_and_output:
                        content_index += 1

                    state.text_content_index_and_output = (
                        content_index,
                        ResponseOutputText(
                            text="",
                            type="output_text",
                            annotations=[],
                            logprobs=[],
                        ),
                    )
                    # Start a new assistant message stream
                    assistant_item = ResponseOutputMessage(
                        id=FAKE_RESPONSES_ID,
                        content=[],
                        role="assistant",
                        type="message",
                        status="in_progress",
                    )
                    # Notify consumers of the start of a new output message + first content part
                    yield ResponseOutputItemAddedEvent(
                        item=assistant_item,
                        output_index=state.reasoning_content_index_and_output
                        is not None,  # fixed 0 -> 0 or 1
                        type="response.output_item.added",
                        sequence_number=sequence_number.get_and_increment(),
                    )
                    yield ResponseContentPartAddedEvent(
                        content_index=state.text_content_index_and_output[0],
                        item_id=FAKE_RESPONSES_ID,
                        output_index=state.reasoning_content_index_and_output
                        is not None,  # fixed 0 -> 0 or 1
                        part=ResponseOutputText(
                            text="",
                            type="output_text",
                            annotations=[],
                            logprobs=[],
                        ),
                        type="response.content_part.added",
                        sequence_number=sequence_number.get_and_increment(),
                    )
                # Emit the delta for this segment of content
                yield ResponseTextDeltaEvent(
                    content_index=state.text_content_index_and_output[0],
                    delta=delta.content,
                    item_id=FAKE_RESPONSES_ID,
                    output_index=state.reasoning_content_index_and_output
                    is not None,  # fixed 0 -> 0 or 1
                    type="response.output_text.delta",
                    sequence_number=sequence_number.get_and_increment(),
                    logprobs=[],
                )
                # Accumulate the text into the response part
                state.text_content_index_and_output[1].text += delta.content

            # Handle refusals (model declines to answer)
            # This is always set by the OpenAI API, but not by others e.g. LiteLLM
            if hasattr(delta, "refusal") and delta.refusal:
                if not state.refusal_content_index_and_output:
                    refusal_index = 0
                    if state.reasoning_content_index_and_output:
                        refusal_index += 1
                    if state.text_content_index_and_output:
                        refusal_index += 1

                    state.refusal_content_index_and_output = (
                        refusal_index,
                        ResponseOutputRefusal(refusal="", type="refusal"),
                    )
                    # Start a new assistant message if one doesn't exist yet (in-progress)
                    assistant_item = ResponseOutputMessage(
                        id=FAKE_RESPONSES_ID,
                        content=[],
                        role="assistant",
                        type="message",
                        status="in_progress",
                    )
                    # Notify downstream that assistant message + first content part are starting
                    yield ResponseOutputItemAddedEvent(
                        item=assistant_item,
                        output_index=state.reasoning_content_index_and_output
                        is not None,  # fixed 0 -> 0 or 1
                        type="response.output_item.added",
                        sequence_number=sequence_number.get_and_increment(),
                    )
                    yield ResponseContentPartAddedEvent(
                        content_index=state.refusal_content_index_and_output[0],
                        item_id=FAKE_RESPONSES_ID,
                        output_index=(1 if state.reasoning_content_index_and_output else 0),
                        part=ResponseOutputRefusal(
                            refusal="",
                            type="refusal",
                        ),
                        type="response.content_part.added",
                        sequence_number=sequence_number.get_and_increment(),
                    )
                # Emit the delta for this segment of refusal
                yield ResponseRefusalDeltaEvent(
                    content_index=state.refusal_content_index_and_output[0],
                    delta=delta.refusal,
                    item_id=FAKE_RESPONSES_ID,
                    output_index=state.reasoning_content_index_and_output
                    is not None,  # fixed 0 -> 0 or 1
                    type="response.refusal.delta",
                    sequence_number=sequence_number.get_and_increment(),
                )
                # Accumulate the refusal string in the output part
                state.refusal_content_index_and_output[1].refusal += delta.refusal

            # Handle tool calls with real-time streaming support
            if delta.tool_calls:
                for tc_delta in delta.tool_calls:
                    if tc_delta.index not in state.function_calls:
                        state.function_calls[tc_delta.index] = ResponseFunctionToolCall(
                            id=FAKE_RESPONSES_ID,
                            arguments="",
                            name="",
                            type="function_call",
                            call_id="",
                        )
                        state.function_call_streaming[tc_delta.index] = False

                    tc_function = tc_delta.function

                    # Accumulate arguments as they come in
                    state.function_calls[tc_delta.index].arguments += (
                        tc_function.arguments if tc_function else ""
                    ) or ""

                    # Set function name directly (it's correct from the first function call chunk)
                    if tc_function and tc_function.name:
                        state.function_calls[tc_delta.index].name = tc_function.name

                    if tc_delta.id:
                        state.function_calls[tc_delta.index].call_id = tc_delta.id

                    # Capture thought_signature from Gemini (provider_specific_fields)
                    if (
                        hasattr(tc_delta, "provider_specific_fields")
                        and tc_delta.provider_specific_fields
                    ):
                        provider_fields = tc_delta.provider_specific_fields
                        if isinstance(provider_fields, dict):
                            thought_sig = provider_fields.get("thought_signature")
                            if thought_sig:
                                state.function_call_thought_signatures[tc_delta.index] = thought_sig

                    function_call = state.function_calls[tc_delta.index]

                    # Start streaming as soon as we have function name and call_id
                    if (
                        not state.function_call_streaming[tc_delta.index]
                        and function_call.name
                        and function_call.call_id
                    ):
                        # Calculate the output index for this function call
                        function_call_starting_index = 0
                        if state.reasoning_content_index_and_output:
                            function_call_starting_index += 1
                        if state.text_content_index_and_output:
                            function_call_starting_index += 1
                        if state.refusal_content_index_and_output:
                            function_call_starting_index += 1

                        # Add offset for already started function calls
                        function_call_starting_index += sum(
                            1 for streaming in state.function_call_streaming.values() if streaming
                        )

                        # Mark this function call as streaming and store its output index
                        state.function_call_streaming[tc_delta.index] = True
                        state.function_call_output_idx[tc_delta.index] = (
                            function_call_starting_index
                        )

                        # Send initial function call added event
                        yield ResponseOutputItemAddedEvent(
                            item=ResponseFunctionToolCall(
                                id=FAKE_RESPONSES_ID,
                                call_id=function_call.call_id,
                                arguments="",  # Start with empty arguments
                                name=function_call.name,
                                type="function_call",
                            ),
                            output_index=function_call_starting_index,
                            type="response.output_item.added",
                            sequence_number=sequence_number.get_and_increment(),
                        )

                    # Stream arguments if we've started streaming this function call
                    if (
                        state.function_call_streaming.get(tc_delta.index, False)
                        and tc_function
                        and tc_function.arguments
                    ):
                        output_index = state.function_call_output_idx[tc_delta.index]
                        yield ResponseFunctionCallArgumentsDeltaEvent(
                            delta=tc_function.arguments,
                            item_id=FAKE_RESPONSES_ID,
                            output_index=output_index,
                            type="response.function_call_arguments.delta",
                            sequence_number=sequence_number.get_and_increment(),
                        )

        if state.reasoning_content_index_and_output:
            if (
                state.reasoning_content_index_and_output[1].summary
                and len(state.reasoning_content_index_and_output[1].summary) > 0
            ):
                yield ResponseReasoningSummaryPartDoneEvent(
                    item_id=FAKE_RESPONSES_ID,
                    output_index=0,
                    summary_index=0,
                    part=DoneEventPart(
                        text=state.reasoning_content_index_and_output[1].summary[0].text,
                        type="summary_text",
                    ),
                    type="response.reasoning_summary_part.done",
                    sequence_number=sequence_number.get_and_increment(),
                )
            elif state.reasoning_content_index_and_output[1].content is not None:
                yield ResponseReasoningTextDoneEvent(
                    item_id=FAKE_RESPONSES_ID,
                    output_index=0,
                    content_index=0,
                    text=state.reasoning_content_index_and_output[1].content[0].text,
                    type="response.reasoning_text.done",
                    sequence_number=sequence_number.get_and_increment(),
                )
            yield ResponseOutputItemDoneEvent(
                item=state.reasoning_content_index_and_output[1],
                output_index=0,
                type="response.output_item.done",
                sequence_number=sequence_number.get_and_increment(),
            )

        function_call_starting_index = 0
        if state.reasoning_content_index_and_output:
            function_call_starting_index += 1

        if state.text_content_index_and_output:
            function_call_starting_index += 1
            # Send end event for this content part
            yield ResponseContentPartDoneEvent(
                content_index=state.text_content_index_and_output[0],
                item_id=FAKE_RESPONSES_ID,
                output_index=state.reasoning_content_index_and_output
                is not None,  # fixed 0 -> 0 or 1
                part=state.text_content_index_and_output[1],
                type="response.content_part.done",
                sequence_number=sequence_number.get_and_increment(),
            )

        if state.refusal_content_index_and_output:
            function_call_starting_index += 1
            # Send end event for this content part
            yield ResponseContentPartDoneEvent(
                content_index=state.refusal_content_index_and_output[0],
                item_id=FAKE_RESPONSES_ID,
                output_index=state.reasoning_content_index_and_output
                is not None,  # fixed 0 -> 0 or 1
                part=state.refusal_content_index_and_output[1],
                type="response.content_part.done",
                sequence_number=sequence_number.get_and_increment(),
            )

        # Send completion events for function calls
        for index, function_call in state.function_calls.items():
            if state.function_call_streaming.get(index, False):
                # Function call was streamed, just send the completion event
                output_index = state.function_call_output_idx[index]

                # Build function call kwargs with thought_signature if available
                func_call_kwargs: dict[str, Any] = {
                    "id": FAKE_RESPONSES_ID,
                    "call_id": function_call.call_id,
                    "arguments": function_call.arguments,
                    "name": function_call.name,
                    "type": "function_call",
                }

                # Add thought_signature from Gemini if present
                if index in state.function_call_thought_signatures:
                    func_call_kwargs["provider_specific_fields"] = {
                        "google": {
                            "thought_signature": state.function_call_thought_signatures[index]
                        }
                    }

                yield ResponseOutputItemDoneEvent(
                    item=ResponseFunctionToolCall(**func_call_kwargs),
                    output_index=output_index,
                    type="response.output_item.done",
                    sequence_number=sequence_number.get_and_increment(),
                )
            else:
                # Function call was not streamed (fallback to old behavior)
                # This handles edge cases where function name never arrived
                fallback_starting_index = 0
                if state.reasoning_content_index_and_output:
                    fallback_starting_index += 1
                if state.text_content_index_and_output:
                    fallback_starting_index += 1
                if state.refusal_content_index_and_output:
                    fallback_starting_index += 1

                # Add offset for already started function calls
                fallback_starting_index += sum(
                    1 for streaming in state.function_call_streaming.values() if streaming
                )

                # Build function call kwargs with thought_signature if available
                fallback_func_call_kwargs: dict[str, Any] = {
                    "id": FAKE_RESPONSES_ID,
                    "call_id": function_call.call_id,
                    "arguments": function_call.arguments,
                    "name": function_call.name,
                    "type": "function_call",
                }

                # Add thought_signature from Gemini if present
                if index in state.function_call_thought_signatures:
                    fallback_func_call_kwargs["provider_specific_fields"] = {
                        "google": {
                            "thought_signature": state.function_call_thought_signatures[index]
                        }
                    }

                # Send all events at once (backward compatibility)
                yield ResponseOutputItemAddedEvent(
                    item=ResponseFunctionToolCall(**fallback_func_call_kwargs),
                    output_index=fallback_starting_index,
                    type="response.output_item.added",
                    sequence_number=sequence_number.get_and_increment(),
                )
                yield ResponseFunctionCallArgumentsDeltaEvent(
                    delta=function_call.arguments,
                    item_id=FAKE_RESPONSES_ID,
                    output_index=fallback_starting_index,
                    type="response.function_call_arguments.delta",
                    sequence_number=sequence_number.get_and_increment(),
                )
                yield ResponseOutputItemDoneEvent(
                    item=ResponseFunctionToolCall(**fallback_func_call_kwargs),
                    output_index=fallback_starting_index,
                    type="response.output_item.done",
                    sequence_number=sequence_number.get_and_increment(),
                )

        # Finally, send the Response completed event
        outputs: list[ResponseOutputItem] = []

        # include Reasoning item if it exists
        if state.reasoning_content_index_and_output:
            reasoning_item = state.reasoning_content_index_and_output[1]
            # Store thinking text in content and signature in encrypted_content
            if state.thinking_text:
                # Add thinking text as a Content object
                if not reasoning_item.content:
                    reasoning_item.content = []
                reasoning_item.content.append(
                    Content(text=state.thinking_text, type="reasoning_text")
                )
            # Store signature in encrypted_content
            if state.thinking_signature:
                reasoning_item.encrypted_content = state.thinking_signature
            outputs.append(reasoning_item)

        # include text or refusal content if they exist
        if state.text_content_index_and_output or state.refusal_content_index_and_output:
            assistant_msg = ResponseOutputMessage(
                id=FAKE_RESPONSES_ID,
                content=[],
                role="assistant",
                type="message",
                status="completed",
            )
            if state.text_content_index_and_output:
                assistant_msg.content.append(state.text_content_index_and_output[1])
            if state.refusal_content_index_and_output:
                assistant_msg.content.append(state.refusal_content_index_and_output[1])
            outputs.append(assistant_msg)

            # send a ResponseOutputItemDone for the assistant message
            yield ResponseOutputItemDoneEvent(
                item=assistant_msg,
                output_index=state.reasoning_content_index_and_output
                is not None,  # fixed 0 -> 0 or 1
                type="response.output_item.done",
                sequence_number=sequence_number.get_and_increment(),
            )

        for index, function_call in state.function_calls.items():
            # Reconstruct function call with thought_signature if available
            if index in state.function_call_thought_signatures:
                func_call_with_signature = ResponseFunctionToolCall(
                    id=function_call.id,
                    call_id=function_call.call_id,
                    arguments=function_call.arguments,
                    name=function_call.name,
                    type="function_call",
                    provider_specific_fields={  # type: ignore[call-arg]
                        "google": {
                            "thought_signature": state.function_call_thought_signatures[index]
                        }
                    },
                )
                outputs.append(func_call_with_signature)
            else:
                outputs.append(function_call)

        final_response = response.model_copy()
        final_response.output = outputs
        final_response.usage = (
            ResponseUsage(
                input_tokens=usage.prompt_tokens or 0,
                output_tokens=usage.completion_tokens or 0,
                total_tokens=usage.total_tokens or 0,
                output_tokens_details=OutputTokensDetails(
                    reasoning_tokens=usage.completion_tokens_details.reasoning_tokens
                    if usage.completion_tokens_details
                    and usage.completion_tokens_details.reasoning_tokens
                    else 0
                ),
                input_tokens_details=InputTokensDetails(
                    cached_tokens=usage.prompt_tokens_details.cached_tokens
                    if usage.prompt_tokens_details and usage.prompt_tokens_details.cached_tokens
                    else 0
                ),
            )
            if usage
            else None
        )

        yield ResponseCompletedEvent(
            response=final_response,
            type="response.completed",
            sequence_number=sequence_number.get_and_increment(),
        )