Skip to content

fix(translation): recover Responses terminal streams #796

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
ajcasagrande wants to merge 1 commit into from
+446 −132
Closed

fix(translation): recover Responses terminal streams #796

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
141 changes: 81 additions & 60 deletions crates/switchyard-translation/src/codecs/responses/stream.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,44 @@ struct ResponsesStreamResponse {
usage: Value,
}

fn decode_responses_identity(
state: &mut StreamTranslationState,
response: Option<&serde_json::Map<String, Value>>,
) -> Vec<LlmResponseChunk> {
let id = response
.and_then(|response| response.get("id"))
.and_then(Value::as_str)
.filter(|id| !id.is_empty())
.map(ToOwned::to_owned);
let model = response
.and_then(|response| response.get("model"))
.and_then(Value::as_str)
.filter(|model| !model.is_empty())
.map(ToOwned::to_owned);
let identity_changed = id
.as_deref()
.is_some_and(|id| state.message_id.as_deref() != Some(id))
|| model
.as_deref()
.is_some_and(|model| state.model.as_deref() != Some(model));

if id.is_some() {
state.message_id = id;
}
if model.is_some() {
state.model = model;
}
if !state.saw_message_start || identity_changed {
state.saw_message_start = true;
vec![LlmResponseChunk::MessageStart {
id: state.message_id.clone(),
model: state.model.clone(),
}]
} else {
Vec::new()
}
}

// Decodes one OpenAI Responses event into neutral streaming events.
fn decode_responses_stream(
state: &mut StreamTranslationState,
Expand All @@ -110,25 +148,8 @@ fn decode_responses_stream(
.or_else(|| event.get("event"))
.and_then(Value::as_str);
match event_type {
Some("response.created") => {
state.saw_message_start = true;
let response = event.get("response").and_then(Value::as_object);
if let Some(model) = response
.and_then(|response| response.get("model"))
.and_then(Value::as_str)
{
state.model = Some(model.to_string());
}
if let Some(id) = response
.and_then(|response| response.get("id"))
.and_then(Value::as_str)
{
state.message_id = Some(id.to_string());
}
vec![LlmResponseChunk::MessageStart {
id: state.message_id.clone(),
model: state.model.clone(),
}]
Some("response.created" | "response.in_progress") => {
decode_responses_identity(state, event.get("response").and_then(Value::as_object))
}
Some("response.output_text.delta") => event
.get("delta")
Expand Down Expand Up @@ -257,48 +278,10 @@ fn decode_responses_stream(
})
.unwrap_or_default()
}
Some("response.completed") => {
let mut out = Vec::new();
// Some providers send output only in the final snapshot. Reconcile it with
// decoded deltas before emitting the stop, without repeating streamed content.
if let Some(items) = event
.get("response")
.and_then(|response| response.get("output"))
.and_then(Value::as_array)
{
for (position, item) in items.iter().enumerate() {
if let Some(item) = item.as_object() {
out.extend(decode_responses_completed_item(item, position, state));
if matches!(out.last(), Some(LlmResponseChunk::StreamError { .. })) {
return out;
}
}
}
}
if let Some(usage) = event
.get("response")
.and_then(Value::as_object)
.and_then(|response| response.get("usage"))
.and_then(Value::as_object)
{
let usage = responses_usage(usage);
state.usage = usage.clone();
state.saw_backend_usage = true;
out.push(LlmResponseChunk::Usage(usage));
}
// A completed response that produced a tool call ended the turn to run that
// tool, not because the assistant was done. The buffered decoder reports tool
// use for such output; the stream must too, or stop-reason-driven tool loops
// (Anthropic `tool_use`, Chat `tool_calls`) stop without running the tool.
// Carries the Anthropic spelling because every encoder already maps it.
let reason = state.decoded_tool_call.then(|| "tool_use".to_string());
out.push(LlmResponseChunk::MessageStop { reason });
out
Some("response.completed") => decode_responses_terminal_snapshot(state, event, None),
Some("response.incomplete") => {
decode_responses_terminal_snapshot(state, event, Some("max_tokens".to_string()))
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
// Carries the Anthropic spelling because every encoder already maps it.
Some("response.incomplete") => vec![LlmResponseChunk::MessageStop {
reason: Some("max_tokens".to_string()),
}],
Some("response.failed") => vec![LlmResponseChunk::StreamError {
message: event
.get("response")
Expand All @@ -319,6 +302,44 @@ fn decode_responses_stream(
}
}

fn decode_responses_terminal_snapshot(
state: &mut StreamTranslationState,
event: &Value,
explicit_stop_reason: Option<String>,
) -> Vec<LlmResponseChunk> {
let response = event.get("response").and_then(Value::as_object);
let mut out = decode_responses_identity(state, response);

if let Some(items) = response
.and_then(|response| response.get("output"))
.and_then(Value::as_array)
{
for (position, item) in items.iter().enumerate() {
if let Some(item) = item.as_object() {
out.extend(decode_responses_completed_item(item, position, state));
if matches!(out.last(), Some(LlmResponseChunk::StreamError { .. })) {
return out;
}
}
}
}

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

let reason =
explicit_stop_reason.or_else(|| state.decoded_tool_call.then(|| "tool_use".to_string()));
out.push(LlmResponseChunk::MessageStop { reason });
out
}

// Encodes neutral streaming events into OpenAI Responses events.
fn encode_responses_stream(
state: &mut StreamTranslationState,
Expand Down
76 changes: 74 additions & 2 deletions crates/switchyard-translation/src/helpers.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -288,7 +288,8 @@ where
break;
}
sse::SseFrame::Data(value) => {
saw_terminal |= sse::is_terminal_event(source, &value);
let terminal = sse::terminal_behavior(source, &value);
saw_terminal |= terminal.is_terminal();
let normalized = codec.decode_event(&mut state, &value);
saw_error |= normalized.iter().any(|chunk| matches!(
chunk,
Expand All @@ -300,6 +301,9 @@ where
value,
normalized,
);
if terminal.is_stream_ending() {
break;
}
}
}
} else {
Expand All @@ -319,7 +323,8 @@ where
saw_terminal |= source != WireFormat::AnthropicMessages;
}
sse::SseFrame::Data(value) => {
saw_terminal |= sse::is_terminal_event(source, &value);
let terminal = sse::terminal_behavior(source, &value);
saw_terminal |= terminal.is_terminal();
let normalized = codec.decode_event(&mut state, &value);
saw_error |= normalized.iter().any(|chunk| matches!(
chunk,
Expand Down Expand Up @@ -845,6 +850,73 @@ mod tests {
Ok(())
}

#[test]
fn decode_stream_stops_polling_after_explicit_responses_terminals() -> Result<(), BoxError> {
let cases = [
(
"response.incomplete",
b"data: {\"type\":\"response.incomplete\",\"response\":{\"status\":\"incomplete\"}}\n\n"
.to_vec(),
),
(
"error",
b"data: {\"type\":\"error\",\"message\":\"boom\"}\n\n".to_vec(),
),
];
for (event_type, terminal) in cases {
let bytes =
stream::iter([Ok::<Vec<u8>, LlmClientError>(terminal)]).chain(stream::poll_fn(
move |_| panic!("decode_stream polled upstream after explicit {event_type}"),
));

let events = decode_all(bytes, WireFormat::OpenAiResponses)?;

assert_eq!(events.len(), 1);
match event_type {
"response.incomplete" => assert!(matches!(
events[0].normalized().last(),
Some(LlmResponseChunk::MessageStop { reason: Some(reason) })
if reason == "max_tokens"
)),
"error" => assert!(matches!(
events[0].normalized().last(),
Some(LlmResponseChunk::StreamError { message }) if message == "boom"
)),
_ => unreachable!(),
}
}
Ok(())
}

#[test]
fn decode_stream_keeps_chat_usage_after_finish_reason() -> Result<(), BoxError> {
let sse = concat!(
"data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":2,\"completion_tokens\":3,\"total_tokens\":5}}\n\n",
"data: [DONE]\n\n",
)
.as_bytes()
.to_vec();
let bytes =
stream::iter([Ok::<Vec<u8>, LlmClientError>(sse)]).chain(stream::poll_fn(|_| {
panic!("decode_stream polled upstream after [DONE]")
}));

let events = decode_all(bytes, WireFormat::OpenAiChat)?;
let usage = events
.iter()
.flat_map(LlmResponseStreamEvent::normalized)
.find_map(|chunk| match chunk {
LlmResponseChunk::Usage(usage) => Some(usage),
_ => None,
})
.ok_or("missing trailing Chat usage")?;
assert_eq!(usage.input_tokens, Some(2));
assert_eq!(usage.output_tokens, Some(3));
assert_eq!(usage.total_tokens, Some(5));
Ok(())
}

#[test]
fn responses_failed_before_done_remains_a_stream_error() -> Result<(), BoxError> {
let sse = b"data: {\"type\":\"response.failed\",\"response\":{\"error\":{\"message\":\"upstream failed\"}}}\n\ndata: [DONE]\n\n".to_vec();
Expand Down
Loading
Loading