switchyard_translation/codecs/openai_chat/

stream.rs

// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Streaming codec for OpenAI Chat Completions chunks.

use serde_json::{Map, Value, json};

use crate::LlmResponseChunk;
use crate::codecs::stream::{
    StreamCodec, StreamTranslationState, record_source_identity, state_source_is, string_field,
    target_model_or_source_model,
};
use crate::format::{FormatId, WireFormat};
use crate::llm::Usage;

/// Stream codec for OpenAI Chat Completions chunks.
pub struct OpenAiChatStreamCodec;

impl StreamCodec for OpenAiChatStreamCodec {
    fn format(&self) -> FormatId {
        WireFormat::OpenAiChat.into()
    }

    fn decode_event(
        &self,
        state: &mut StreamTranslationState,
        event: &Value,
    ) -> Vec<LlmResponseChunk> {
        decode_openai_chat_stream(state, event)
    }

    fn encode_event(
        &self,
        state: &mut StreamTranslationState,
        event: LlmResponseChunk,
    ) -> Vec<Value> {
        encode_openai_chat_stream(state, event)
    }

    fn finish(&self, state: &mut StreamTranslationState) -> Vec<Value> {
        finish_openai_chat_stream(state)
    }
}

// Decodes one OpenAI Chat chunk into neutral streaming events.
fn decode_openai_chat_stream(
    state: &mut StreamTranslationState,
    event: &Value,
) -> Vec<LlmResponseChunk> {
    let Some(object) = event.as_object() else {
        return vec![LlmResponseChunk::DecodeError {
            message: "OpenAI stream event is not an object".to_string(),
        }];
    };

    // Upstream error frames carry no choices, so without this they would decode to a bare
    // `MessageStart` and the error text would be dropped.
    if let Some(error) = object.get("error") {
        return vec![LlmResponseChunk::StreamError {
            message: error
                .get("message")
                .and_then(Value::as_str)
                .unwrap_or("unknown OpenAI stream error")
                .to_string(),
        }];
    }

    let mut out = Vec::new();
    if !state.saw_message_start {
        state.saw_message_start = true;
        if let Some(model) = string_field(object, "model") {
            state.model = Some(model);
        }
        if let Some(id) = string_field(object, "id") {
            state.message_id = Some(id);
        }
        out.push(LlmResponseChunk::MessageStart {
            id: state.message_id.clone(),
            model: state.model.clone(),
        });
    }

    if let Some(usage) = object.get("usage").and_then(Value::as_object) {
        let usage = openai_usage(usage);
        state.usage = usage.clone();
        state.saw_backend_usage = true;
        out.push(LlmResponseChunk::Usage(usage));
    }

    for choice in object
        .get("choices")
        .and_then(Value::as_array)
        .into_iter()
        .flatten()
    {
        let Some(choice) = choice.as_object() else {
            continue;
        };
        if let Some(delta) = choice.get("delta").and_then(Value::as_object) {
            if let Some(text) = delta.get("content").and_then(Value::as_str)
                && !text.is_empty()
            {
                out.push(LlmResponseChunk::TextDelta {
                    index: 0,
                    text: text.to_string(),
                });
            }
            for reasoning_key in ["reasoning_content", "reasoning"] {
                if let Some(text) = delta.get(reasoning_key).and_then(Value::as_str)
                    && !text.is_empty()
                {
                    out.push(LlmResponseChunk::ReasoningDelta {
                        index: 0,
                        text: text.to_string(),
                    });
                }
            }
            if let Some(tool_calls) = delta.get("tool_calls").and_then(Value::as_array) {
                for tool_call in tool_calls {
                    if let Some(tool_call) = tool_call.as_object() {
                        let function = tool_call.get("function").and_then(Value::as_object);
                        out.push(LlmResponseChunk::ToolCallDelta {
                            index: tool_call.get("index").and_then(Value::as_u64).unwrap_or(0)
                                as usize,
                            id: tool_call
                                .get("id")
                                .and_then(Value::as_str)
                                .map(ToOwned::to_owned),
                            name: function
                                .and_then(|function| function.get("name"))
                                .and_then(Value::as_str)
                                .map(ToOwned::to_owned),
                            arguments_delta: function
                                .and_then(|function| function.get("arguments"))
                                .and_then(Value::as_str)
                                .map(ToOwned::to_owned),
                        });
                    }
                }
            }
        }
        if let Some(reason) = choice.get("finish_reason").and_then(Value::as_str) {
            out.push(LlmResponseChunk::MessageStop {
                reason: Some(reason.to_string()),
            });
        }
    }
    out
}

// Encodes neutral streaming events into OpenAI Chat chunks.
fn encode_openai_chat_stream(
    state: &mut StreamTranslationState,
    event: LlmResponseChunk,
) -> Vec<Value> {
    // An in-band error is terminal: once the error is emitted, drop every later chunk.
    if state.errored {
        return Vec::new();
    }
    match event {
        LlmResponseChunk::MessageStart { id, model } => {
            record_source_identity(state, id, model);
            // `finish_openai_chat_stream` uses this flag to decide whether a clean EOF needs a
            // synthesized terminal chunk. Set it on the encode path as well as the decode path.
            state.saw_message_start = true;
            if state.emitted_message_start
                || (!state_source_is(state, WireFormat::AnthropicMessages)
                    && !state_source_is(state, WireFormat::OpenAiResponses))
            {
                Vec::new()
            } else {
                state.emitted_message_start = true;
                vec![openai_stream_chunk(
                    state,
                    json!({"role": "assistant"}),
                    None,
                    None,
                )]
            }
        }
        LlmResponseChunk::TextDelta { text, .. } => {
            vec![openai_stream_chunk(
                state,
                json!({"content": text}),
                None,
                None,
            )]
        }
        LlmResponseChunk::ReasoningDelta { text, .. } => {
            vec![openai_stream_chunk(
                state,
                json!({"reasoning_content": text}),
                None,
                None,
            )]
        }
        LlmResponseChunk::ToolCallDelta {
            index,
            id,
            name,
            arguments_delta,
        } => vec![openai_tool_call_chunk(
            state,
            index,
            id,
            name,
            arguments_delta,
        )],
        LlmResponseChunk::Usage(usage) => {
            state.usage = usage;
            state.saw_backend_usage = true;
            if state.finished {
                vec![openai_usage_chunk(state)]
            } else {
                Vec::new()
            }
        }
        LlmResponseChunk::MessageStop { reason } => {
            if state.finished {
                return Vec::new();
            }
            state.finished = true;
            vec![openai_stream_chunk(
                state,
                json!({}),
                Some(openai_finish_reason(reason.as_deref())),
                state.saw_backend_usage.then(|| openai_usage_value(state)),
            )]
        }
        LlmResponseChunk::DecodeError { message } | LlmResponseChunk::StreamError { message } => {
            // An in-band error is terminal: emit the error, then nothing further.
            state.finished = true; // finish() adds no success events
            state.errored = true; // the entry guard drops any later chunk
            vec![json!({"error": {"message": message}})]
        }
    }
}

// Emits a terminal chunk if the source stream ended before a stop event arrived.
fn finish_openai_chat_stream(state: &mut StreamTranslationState) -> Vec<Value> {
    if state.finished || !state.saw_message_start {
        return Vec::new();
    }
    state.finished = true;
    vec![openai_stream_chunk(
        state,
        json!({}),
        Some(openai_finish_reason(state.stop_reason.as_deref())),
        state.saw_backend_usage.then(|| openai_usage_value(state)),
    )]
}

// Normalizes OpenAI token usage fields.
fn openai_usage(usage: &Map<String, Value>) -> Usage {
    let cached_input_tokens = usage
        .get("prompt_tokens_details")
        .and_then(|details| details.get("cached_tokens"))
        .and_then(Value::as_u64);
    let cache_creation_input_tokens = usage
        .get("prompt_tokens_details")
        .and_then(|details| {
            details
                .get("cache_write_tokens")
                .or_else(|| details.get("cache_creation_tokens"))
        })
        .and_then(Value::as_u64);
    Usage {
        input_tokens: usage
            .get("prompt_tokens")
            .and_then(Value::as_u64)
            .map(|tokens| {
                tokens
                    .saturating_sub(cached_input_tokens.unwrap_or(0))
                    .saturating_sub(cache_creation_input_tokens.unwrap_or(0))
            }),
        cache: Usage::cache_details(cached_input_tokens, cache_creation_input_tokens),
        output_tokens: usage.get("completion_tokens").and_then(Value::as_u64),
        total_tokens: usage.get("total_tokens").and_then(Value::as_u64),
        reasoning_tokens: usage
            .get("completion_tokens_details")
            .and_then(|details| details.get("reasoning_tokens"))
            .or_else(|| {
                usage
                    .get("output_tokens_details")
                    .and_then(|details| details.get("reasoning_tokens"))
            })
            .and_then(Value::as_u64),
    }
}

// Builds a single OpenAI Chat stream chunk payload.
fn openai_stream_chunk(
    state: &StreamTranslationState,
    delta: Value,
    finish_reason: Option<String>,
    usage: Option<Value>,
) -> Value {
    let mut payload = json!({
        "id": openai_stream_id(state),
        "object": "chat.completion.chunk",
        "created": 0,
        "model": target_model_or_source_model(state),
        "choices": [{
            "index": 0,
            "delta": delta,
            "finish_reason": finish_reason,
        }],
    });
    if let Some(usage) = usage {
        payload["usage"] = usage;
    }
    payload
}

// Builds the usage-only chunk OpenAI emits after the terminal choices chunk.
fn openai_usage_chunk(state: &StreamTranslationState) -> Value {
    let mut payload = openai_stream_chunk(state, json!({}), None, Some(openai_usage_value(state)));
    payload["choices"] = json!([]);
    payload
}

// Builds an OpenAI Chat tool-call delta chunk.
fn openai_tool_call_chunk(
    state: &StreamTranslationState,
    index: usize,
    id: Option<String>,
    name: Option<String>,
    arguments: Option<String>,
) -> Value {
    let mut function = Map::new();
    if let Some(name) = name {
        function.insert("name".to_string(), Value::String(name));
    }
    if let Some(arguments) = arguments {
        function.insert("arguments".to_string(), Value::String(arguments));
    }
    let mut tool_call = Map::new();
    tool_call.insert("index".to_string(), json!(index));
    tool_call.insert("type".to_string(), json!("function"));
    tool_call.insert("function".to_string(), Value::Object(function));
    if let Some(id) = id {
        tool_call.insert("id".to_string(), Value::String(id));
    }
    openai_stream_chunk(
        state,
        json!({"tool_calls": [Value::Object(tool_call)]}),
        None,
        None,
    )
}

// Builds OpenAI usage payloads from normalized and provider-extra state.
fn openai_usage_value(state: &StreamTranslationState) -> Value {
    let cache_creation_tokens = state.usage.cache_creation_input_tokens().unwrap_or(0);
    let cache_read_tokens = state.usage.cached_input_tokens().unwrap_or(0);
    let prompt_tokens =
        state.usage.input_tokens.unwrap_or(0) + cache_creation_tokens + cache_read_tokens;
    let completion_tokens = state.usage.output_tokens.unwrap_or(0);
    let mut usage = json!({
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "total_tokens": state.usage.total_tokens.unwrap_or(prompt_tokens + completion_tokens),
    });
    if let Some(reasoning_tokens) = state.usage.reasoning_tokens {
        usage["completion_tokens_details"] = json!({
            "reasoning_tokens": reasoning_tokens,
        });
    }
    if state.usage.cache_creation_input_tokens().is_some()
        || state.usage.cached_input_tokens().is_some()
    {
        usage["prompt_tokens_details"] = json!({
            "cached_tokens": cache_read_tokens,
            "cache_creation_tokens": cache_creation_tokens,
        });
    }
    usage
}

// Converts any upstream message ID into an OpenAI-looking stream ID.
fn openai_stream_id(state: &StreamTranslationState) -> String {
    let Some(id) = state.message_id.as_deref() else {
        return "chatcmpl_switchyard".to_string();
    };
    if id.starts_with("chatcmpl") {
        id.to_string()
    } else if let Some(rest) = id.strip_prefix("msg_").or_else(|| id.strip_prefix("resp_")) {
        format!("chatcmpl_{rest}")
    } else {
        format!("chatcmpl_{id}")
    }
}

// Maps provider stop reasons into OpenAI's finish-reason vocabulary.
fn openai_finish_reason(reason: Option<&str>) -> String {
    match reason {
        Some("end_turn") | Some("stop_sequence") | None => "stop".to_string(),
        Some("max_tokens") => "length".to_string(),
        Some("tool_use") => "tool_calls".to_string(),
        Some(other) => other.to_string(),
    }
}