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 super::buffered::{decode_openai_usage, encode_openai_usage};
use crate::LlmResponseChunk;
use crate::codecs::common::{first_nonempty_string, reasoning_text_from_details};
use crate::codecs::stream::{
    StreamCodec, StreamTranslationState, record_source_identity, state_source_is, string_field,
    target_model_or_source_model,
};
use crate::format::{FormatId, WireFormat};

/// 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();
    let mut identity_changed = false;
    if state.model.is_none() {
        state.model = string_field(object, "model");
        identity_changed |= state.model.is_some();
    }
    if state.message_id.is_none() {
        state.message_id = string_field(object, "id");
        identity_changed |= state.message_id.is_some();
    }
    if !state.saw_message_start || identity_changed {
        state.saw_message_start = true;
        out.push(LlmResponseChunk::MessageStart {
            id: state.message_id.clone(),
            model: state.model.clone(),
        });
    }

    if let Some(usage) = object.get("usage").filter(|usage| usage.is_object()) {
        let usage = decode_openai_usage(Some(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(details) = delta
                .get("reasoning_details")
                .and_then(Value::as_array)
                .filter(|details| !details.is_empty())
            {
                let text = reasoning_text_from_details(details)
                    .or_else(|| {
                        first_nonempty_string(delta, &["reasoning_content", "reasoning"])
                            .map(ToOwned::to_owned)
                    })
                    .unwrap_or_default();
                out.push(LlmResponseChunk::ReasoningDetailsDelta {
                    index: 0,
                    details: details.clone(),
                    text,
                });
            } else {
                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(text) = delta.get("content").and_then(Value::as_str)
                && !text.is_empty()
            {
                out.push(LlmResponseChunk::TextDelta {
                    index: 0,
                    text: text.to_string(),
                });
            }
            if let Some(text) = delta.get("refusal").and_then(Value::as_str)
                && !text.is_empty()
            {
                // Refusal deltas mark the whole message as a content-filter stop. This state
                // persists across chunks so a later `finish_reason: "stop"` cannot overwrite it.
                state.stop_reason = Some("content_filter".to_string());
                out.push(LlmResponseChunk::TextDelta {
                    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);
                        let index =
                            tool_call.get("index").and_then(Value::as_u64).unwrap_or(0) as usize;
                        if let Some(name) = function
                            .and_then(|function| function.get("name"))
                            .and_then(Value::as_str)
                            .filter(|name| !name.is_empty())
                        {
                            state
                                .pending_chat_tool_names
                                .entry(index)
                                .or_default()
                                .push_str(name);
                        }
                        let arguments_delta = function
                            .and_then(|function| function.get("arguments"))
                            .and_then(Value::as_str)
                            .map(ToOwned::to_owned);
                        out.push(LlmResponseChunk::ToolCallDelta {
                            index,
                            id: tool_call
                                .get("id")
                                .and_then(Value::as_str)
                                .map(ToOwned::to_owned),
                            // Names may continue after arguments begin; emit them at finish_reason.
                            name: None,
                            arguments_delta,
                        });
                    }
                }
            }
        }
        if let Some(reason) = choice.get("finish_reason").and_then(Value::as_str) {
            for (index, name) in std::mem::take(&mut state.pending_chat_tool_names) {
                out.push(LlmResponseChunk::ToolCallDelta {
                    index,
                    id: None,
                    name: Some(name),
                    arguments_delta: None,
                });
            }
            let reason =
                if reason == "stop" && state.stop_reason.as_deref() == Some("content_filter") {
                    "content_filter"
                } else {
                    reason
                };
            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::ReasoningDetailsDelta { details, text, .. } => {
            // Chat cannot replay Anthropic signature fragments or an encrypted item's
            // id announcement without its payload.
            let details: Vec<Value> = details
                .into_iter()
                .filter(|detail| match detail.get("type").and_then(Value::as_str) {
                    Some("anthropic.signature_delta") => false,
                    Some("reasoning.encrypted") => detail.get("data").is_some(),
                    _ => true,
                })
                .collect();
            if details.is_empty() && text.is_empty() {
                return Vec::new();
            }
            let details_text = reasoning_text_from_details(&details);
            let mut delta = json!({"reasoning_details": details});
            if !text.is_empty() && details_text.as_deref() != Some(text.as_str()) {
                delta["reasoning"] = Value::String(text);
            }
            vec![openai_stream_chunk(state, delta, None, None)]
        }
        LlmResponseChunk::ToolCallDelta {
            index,
            id,
            name,
            arguments_delta,
        } => {
            let tool = state.tool_states.entry(index).or_default();
            // Repeated full names would be concatenated by Chat clients as new fragments.
            let name = name.filter(|name| tool.name.as_ref() != Some(name));
            if name.is_some() {
                tool.name = name.clone();
            }
            // The source index counts every content block (Anthropic) or output item
            // (Responses), so text ahead of the first tool call shifts it. Chat clients use
            // the index as a subscript into `tool_calls`, so number calls in that array
            // instead, in order of first appearance.
            let chat_index = match tool.chat_tool_index {
                Some(chat_index) => chat_index,
                None => {
                    let chat_index = state.next_chat_tool_index;
                    state.next_chat_tool_index += 1;
                    state.tool_states.entry(index).or_default().chat_tool_index = Some(chat_index);
                    chat_index
                }
            };
            vec![openai_tool_call_chunk(
                state,
                chat_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(|| encode_openai_usage(&state.usage)),
            )]
        }
        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(|| encode_openai_usage(&state.usage)),
    )]
}

// 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(encode_openai_usage(&state.usage)),
    );
    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,
    )
}

// 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("refusal") => "content_filter".to_string(),
        Some(other) => other.to_string(),
    }
}