switchyard_translation/codecs/
stream.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Shared streaming codec contracts, registry, and stream state.
5
6use std::collections::BTreeMap;
7use std::sync::Arc;
8
9use serde::{Deserialize, Serialize};
10use serde_json::{json, Map, Value};
11
12use crate::codecs::anthropic::AnthropicMessagesStreamCodec;
13use crate::codecs::openai_chat::OpenAiChatStreamCodec;
14use crate::codecs::responses::OpenAiResponsesStreamCodec;
15use crate::engine::{FormatRegistry, TranslationEngine};
16use crate::error::{Result, TranslationError};
17use crate::format::{FormatId, WireFormat};
18use crate::llm::{LlmStreamEvent, Usage};
19
20/// Mutable state accumulated while translating one streaming response.
21#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
22pub struct StreamTranslationState {
23    pub source: Option<FormatId>,
24    pub target: Option<FormatId>,
25    /// Model name observed on the source provider stream.
26    pub model: Option<String>,
27    /// Message/response ID observed on the source provider stream.
28    pub message_id: Option<String>,
29    /// Optional model name the target stream should expose to the client.
30    pub target_model: Option<String>,
31    /// Optional message/response ID the target stream should expose to the client.
32    pub target_message_id: Option<String>,
33    pub saw_message_start: bool,
34    pub emitted_message_start: bool,
35    pub finished: bool,
36    pub usage: Usage,
37
38    pub(crate) output_tokens_seen: u64,
39    pub(crate) saw_backend_usage: bool,
40    pub(crate) usage_extras: BTreeMap<String, u64>,
41    pub(crate) stop_reason: Option<String>,
42
43    pub(crate) next_content_index: usize,
44    pub(crate) text_block_index: Option<usize>,
45    pub(crate) text_block_started: bool,
46    pub(crate) emitted_content_block: bool,
47    pub(crate) tool_states: BTreeMap<usize, StreamToolState>,
48
49    pub(crate) response_created: bool,
50    pub(crate) response_text_started: bool,
51    pub(crate) response_text_output_index: Option<usize>,
52    pub(crate) response_text: String,
53    pub(crate) response_reasoning_started: bool,
54    pub(crate) response_reasoning_output_index: Option<usize>,
55    pub(crate) response_reasoning_text: String,
56    pub(crate) next_response_output_index: usize,
57
58    pub(crate) reasoning_block_index: Option<usize>,
59    pub(crate) reasoning_block_started: bool,
60}
61
62// Tracks an in-progress streamed tool call across provider-specific deltas.
63#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
64pub(crate) struct StreamToolState {
65    pub(crate) id: Option<String>,
66    pub(crate) name: Option<String>,
67    pub(crate) arguments: String,
68    pub(crate) pending_arguments: String,
69    pub(crate) started: bool,
70    pub(crate) content_index: Option<usize>,
71    pub(crate) response_output_index: Option<usize>,
72    pub(crate) response_item_id: Option<String>,
73}
74
75impl StreamTranslationState {
76    /// Creates stream state with source and target formats already attached.
77    pub fn new(source: impl Into<FormatId>, target: impl Into<FormatId>) -> Self {
78        Self {
79            source: Some(source.into()),
80            target: Some(target.into()),
81            ..Self::default()
82        }
83    }
84}
85
86/// Registry-backed streaming translator.
87#[derive(Default)]
88pub struct StreamTranslationEngine {
89    engine: TranslationEngine,
90}
91
92/// Codec contract for one provider streaming event format.
93pub trait StreamCodec: Send + Sync {
94    /// Returns the stream format handled by this codec.
95    fn format(&self) -> FormatId;
96
97    /// Decodes one provider event into zero or more neutral events.
98    fn decode_event(
99        &self,
100        state: &mut StreamTranslationState,
101        event: &Value,
102    ) -> Vec<LlmStreamEvent>;
103
104    /// Encodes one neutral event into zero or more provider events.
105    fn encode_event(&self, state: &mut StreamTranslationState, event: LlmStreamEvent)
106        -> Vec<Value>;
107
108    /// Emits any terminal provider events needed after the source stream ends.
109    ///
110    /// This is intentionally required on every codec. Some target formats
111    /// need explicit terminal events after the source closes (for example,
112    /// Anthropic ``message_delta``/``message_stop`` or Responses
113    /// ``response.completed``). Formats that have no source-close work should
114    /// return an empty vector explicitly so the no-op behavior is a conscious
115    /// codec-level choice.
116    fn finish(&self, state: &mut StreamTranslationState) -> Vec<Value>;
117}
118
119/// Registry mapping stream wire formats to stream codecs.
120#[derive(Default)]
121pub struct StreamCodecRegistry {
122    codecs: BTreeMap<FormatId, Arc<dyn StreamCodec>>,
123}
124
125impl StreamCodecRegistry {
126    /// Creates an empty stream codec registry.
127    pub fn new() -> Self {
128        Self::default()
129    }
130
131    /// Creates a registry populated with built-in stream codecs.
132    pub fn with_builtins() -> Self {
133        let mut registry = Self::new();
134        registry.register(OpenAiChatStreamCodec);
135        registry.register(AnthropicMessagesStreamCodec);
136        registry.register(OpenAiResponsesStreamCodec);
137        registry
138    }
139
140    /// Registers or replaces a stream codec for its declared format.
141    pub fn register(&mut self, codec: impl StreamCodec + 'static) {
142        self.codecs.insert(codec.format(), Arc::new(codec));
143    }
144
145    /// Looks up a stream codec by format identifier.
146    pub fn codec(&self, format: impl Into<FormatId>) -> Result<Arc<dyn StreamCodec>> {
147        let format = format.into();
148        self.codecs.get(&format).cloned().ok_or_else(|| {
149            TranslationError::Other(format!("no stream codec registered for {format}"))
150        })
151    }
152}
153
154impl StreamTranslationEngine {
155    /// Creates a streaming engine from an explicit codec registry.
156    pub fn new(registry: StreamCodecRegistry) -> Self {
157        Self {
158            engine: TranslationEngine::with_registries(FormatRegistry::with_builtins(), registry),
159        }
160    }
161
162    /// Translates one source provider event into target provider events.
163    pub fn translate_event(
164        &self,
165        state: &mut StreamTranslationState,
166        source: impl Into<FormatId>,
167        target: impl Into<FormatId>,
168        event: &Value,
169    ) -> Result<Vec<Value>> {
170        self.engine.translate_event(state, source, target, event)
171    }
172
173    /// Finishes target-provider stream emission after the source stream closes.
174    pub fn finish(
175        &self,
176        state: &mut StreamTranslationState,
177        target: impl Into<FormatId>,
178    ) -> Result<Vec<Value>> {
179        self.engine.finish_stream(state, target)
180    }
181
182    /// Convenience helper using built-in codecs and error events instead of `Result`.
183    pub fn translate_event_with_builtins(
184        state: &mut StreamTranslationState,
185        source: impl Into<FormatId>,
186        target: impl Into<FormatId>,
187        event: &Value,
188    ) -> Vec<Value> {
189        Self::default()
190            .translate_event(state, source, target, event)
191            .unwrap_or_else(|error| vec![json!({"error": {"message": error.to_string()}})])
192    }
193}
194
195/// Decodes one provider stream event with the built-in codec registry.
196pub fn decode_stream_event(
197    state: &mut StreamTranslationState,
198    source: impl Into<FormatId>,
199    event: &Value,
200) -> Vec<LlmStreamEvent> {
201    let source = source.into();
202    StreamCodecRegistry::with_builtins()
203        .codec(source)
204        .map(|codec| codec.decode_event(state, event))
205        .unwrap_or_else(|error| {
206            vec![LlmStreamEvent::Error {
207                message: error.to_string(),
208            }]
209        })
210}
211
212/// Encodes one neutral stream event with the built-in codec registry.
213pub fn encode_stream_event(
214    state: &mut StreamTranslationState,
215    target: impl Into<FormatId>,
216    event: LlmStreamEvent,
217) -> Vec<Value> {
218    StreamCodecRegistry::with_builtins()
219        .codec(target)
220        .map(|codec| codec.encode_event(state, event))
221        .unwrap_or_else(|error| vec![json!({"error": {"message": error.to_string()}})])
222}
223
224// Records source-provider identity carried by decoded stream events.
225pub(crate) fn record_source_identity(
226    state: &mut StreamTranslationState,
227    id: Option<String>,
228    model: Option<String>,
229) {
230    if id.is_some() {
231        state.message_id = id;
232    }
233    if model.is_some() {
234        state.model = model;
235    }
236}
237
238// Returns the source model observed from the upstream stream.
239pub(crate) fn source_model_or_unknown(state: &StreamTranslationState) -> String {
240    state.model.clone().unwrap_or_else(|| "unknown".to_string())
241}
242
243// Returns the target/client model when supplied, otherwise the upstream model.
244pub(crate) fn target_model_or_source_model(state: &StreamTranslationState) -> String {
245    state
246        .target_model
247        .clone()
248        .or_else(|| state.model.clone())
249        .unwrap_or_else(|| "unknown".to_string())
250}
251
252// Returns the target/client ID when supplied, otherwise the upstream ID.
253pub(crate) fn target_message_id_or_source_message_id(
254    state: &StreamTranslationState,
255) -> Option<&str> {
256    state
257        .target_message_id
258        .as_deref()
259        .or(state.message_id.as_deref())
260}
261
262// Checks whether the current source format matches a built-in format.
263pub(crate) fn state_source_is(state: &StreamTranslationState, format: WireFormat) -> bool {
264    let format_id: FormatId = format.into();
265    match &state.source {
266        Some(source) => source == &format_id,
267        None => false,
268    }
269}
270
271// Reads a non-empty string field from an event object.
272pub(crate) fn string_field(object: &Map<String, Value>, key: &str) -> Option<String> {
273    object
274        .get(key)
275        .and_then(Value::as_str)
276        .filter(|value| !value.is_empty())
277        .map(ToOwned::to_owned)
278}
switchyard_translation/codecs/stream.rs

switchyard_translation/codecs/
stream.rs
