diff --git a/CHANGELOG.md b/CHANGELOG.md index f99ff1ba5..9e03fed7f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -712,6 +712,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 longer merged into the first call. - `SigV4Fetch::with_clock` takes a plain function (`fn() -> DateTime`) instead of a closure. +- Gains `Transformer`, `TransformStreamController` and `pipe_through`, the + Rust form of the WHATWG `TransformStream` the AI SDK models build `doStream` + with (`start` / `transform` / `flush`; `enqueue`, `enqueue_error`, `error`, + `terminate`). `pipe_through` returns `PipeThrough`, a `Stream` that calls + `start` immediately and the transformer's other methods as it is polled; + `TransformStreamController::new`, `drain` and `is_closed` let a test drive + a transformer without a stream. +- `StreamingToolCallTracker`, `StreamingToolCallDelta`, `TypeValidation` and + `TrackerError` are removed. The tool-call matching they ported from + `@ai-sdk/provider-utils` 5.0.55 is now part of the OpenAI-style chat + streams in `aimux-providers` (OpenAI, OpenAI-compatible, DeepSeek, Groq, + Mistral), which emit the same parts as before. **Rust (aimux-providers)** @@ -741,6 +753,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `PresetDescriptor` gains a crate-private field (the row's usage converter), so code outside the crate can no longer build one with a struct literal; descriptors come from the registry JSON. +- Every language-model `doStream` (OpenAI chat and Responses, OpenAI-compatible + chat and completion, Groq, DeepSeek, Mistral, Anthropic, Bedrock, Cohere, + Google, Vertex, xAI, Hugging Face, Open Responses) is a named stream struct + driven by `pipe_through`; emitted parts are unchanged, and a stream still + stops reading the body at an Anthropic `message_stop` or a + Google/Vertex/xAI/Hugging Face error event. **Rust (provider factory, RFC-0036)** diff --git a/aimux-provider-utils/src/lib.rs b/aimux-provider-utils/src/lib.rs index e42cc14ac..48d76adc8 100644 --- a/aimux-provider-utils/src/lib.rs +++ b/aimux-provider-utils/src/lib.rs @@ -28,8 +28,7 @@ pub mod read_response_with_size_limit; pub mod resolvable; pub mod response_handler; pub mod sigv4_fetch; -mod streaming_tool_call_argument_state; -mod streaming_tool_call_tracker; +pub mod transform_stream; pub mod url; /// WebSocket client for realtime provider APIs (RFC-0028). Empty unless the /// `ws` feature is enabled. @@ -65,9 +64,5 @@ pub use response_handler::{ create_status_code_error_response_handler, invalid_response_api_call, stream_error_api_call, }; pub use sigv4_fetch::{AwsCredentials, SigV4Fetch, SigV4Signing}; -// As `@ai-sdk/provider-utils` (`index.ts:115-119`): the tracker is exported -// from the crate root, its argument state stays internal. -pub use streaming_tool_call_tracker::{ - StreamingToolCallDelta, StreamingToolCallTracker, TrackerError, TypeValidation, -}; +pub use transform_stream::{PipeThrough, TransformStreamController, Transformer, pipe_through}; pub use url::{validate_base_url, without_trailing_slash, without_trailing_slash_opt}; diff --git a/aimux-provider-utils/src/transform_stream.rs b/aimux-provider-utils/src/transform_stream.rs new file mode 100644 index 000000000..d8cf90bf1 --- /dev/null +++ b/aimux-provider-utils/src/transform_stream.rs @@ -0,0 +1,279 @@ +//! The Rust form of the WHATWG `TransformStream` that the AI SDK's models +//! build their `doStream` output with (`response.pipeThrough(new +//! TransformStream({ start, transform, flush }))`). +//! +//! The state the TypeScript closures capture lives in a [`Transformer`] +//! value instead, so a model's stream logic is a plain struct with methods +//! that can be driven and tested without any I/O. + +use std::collections::VecDeque; +use std::pin::Pin; +use std::task::{Context, Poll, ready}; + +use futures::Stream; + +use aimux_core::error::AiMuxError; + +/// What a transformer writes to: the `TransformStreamDefaultController`. +/// +/// [`pipe_through`] owns one and drains it after every call into the +/// transformer; tests can own one directly, call the transformer's methods +/// with it and [`drain`](Self::drain) what was emitted. +pub struct TransformStreamController { + queue: VecDeque>, + closed: bool, +} + +impl Default for TransformStreamController { + fn default() -> Self { + Self { + queue: VecDeque::new(), + closed: false, + } + } +} + +impl TransformStreamController { + /// An open controller with nothing queued. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Take everything emitted so far, in order. The buffer is kept for reuse. + pub fn drain(&mut self) -> impl Iterator> + '_ { + self.queue.drain(..) + } + + /// Whether [`error`](Self::error) or [`terminate`](Self::terminate) ended + /// the stream. + #[must_use] + pub fn is_closed(&self) -> bool { + self.closed + } + + /// Emit one value downstream. + pub fn enqueue(&mut self, value: O) { + if !self.closed { + self.queue.push_back(Ok(value)); + } + } + + /// Emit `error` and keep the stream open: the next chunk is transformed + /// as usual. For errors a consumer may skip (a malformed event inside an + /// otherwise valid stream). The AI SDK's models do this by enqueueing an + /// `{ type: 'error', error }` part and carrying on. + pub fn enqueue_error(&mut self, error: AiMuxError) { + if !self.closed { + self.queue.push_back(Err(error)); + } + } + + /// Emit `error` and end the stream: no further chunk is transformed and + /// `flush` is not called. WHATWG `controller.error(e)`. + pub fn error(&mut self, error: AiMuxError) { + if !self.closed { + self.queue.push_back(Err(error)); + self.closed = true; + } + } + + /// End the stream without an error: no further chunk is transformed and + /// `flush` is not called. WHATWG `controller.terminate()`. Anything + /// emitted after the stream ended is dropped. + pub fn terminate(&mut self) { + self.closed = true; + } +} + +/// The `{ start, transform, flush }` object of a `TransformStream`. +pub trait Transformer: Send + 'static { + type Input: Send + 'static; + type Output: Send + 'static; + + /// Called once, before the first chunk. + fn start(&mut self, controller: &mut TransformStreamController) { + let _ = controller; + } + + /// Called for every input chunk. + fn transform( + &mut self, + chunk: Self::Input, + controller: &mut TransformStreamController, + ); + + /// Called once when the input ends, unless the stream was ended earlier + /// with [`TransformStreamController::error`] or + /// [`TransformStreamController::terminate`]. + fn flush(self, controller: &mut TransformStreamController); +} + +/// `input.pipeThrough(new TransformStream(transformer))`. +/// +/// Like the WHATWG constructor, this calls [`Transformer::start`] right away; +/// every later call happens while the returned stream is polled. +pub fn pipe_through(input: S, mut transformer: T) -> PipeThrough +where + S: Stream + Send + 'static, + T: Transformer, +{ + let mut controller = TransformStreamController::new(); + transformer.start(&mut controller); + PipeThrough { + input: Box::pin(input), + transformer: Some(transformer), + controller, + } +} + +/// The stream [`pipe_through`] returns: it pulls a chunk from `input` only +/// when everything the transformer emitted so far has been handed out. +pub struct PipeThrough { + input: Pin>, + /// `None` once `flush` has run. + transformer: Option, + controller: TransformStreamController, +} + +// The transformer is never pinned: it is only reached through `&mut`. +impl Unpin for PipeThrough {} + +impl Stream for PipeThrough +where + S: Stream, + T: Transformer, +{ + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + loop { + if let Some(item) = this.controller.queue.pop_front() { + return Poll::Ready(Some(item)); + } + if this.controller.closed { + return Poll::Ready(None); + } + let Some(transformer) = this.transformer.as_mut() else { + return Poll::Ready(None); + }; + match ready!(this.input.as_mut().poll_next(cx)) { + Some(chunk) => transformer.transform(chunk, &mut this.controller), + None => { + if let Some(transformer) = this.transformer.take() { + transformer.flush(&mut this.controller); + } + } + } + } + } +} + +#[cfg(test)] +mod tests { + use futures::StreamExt; + + use super::*; + + /// Doubles each number, reports a skippable error for zero, stops with an + /// error on a negative one, and emits the running total on flush. + #[derive(Default)] + struct Doubler { + total: i32, + } + + impl Transformer for Doubler { + type Input = i32; + type Output = i32; + + fn start(&mut self, controller: &mut TransformStreamController) { + controller.enqueue(0); + } + + fn transform(&mut self, chunk: i32, controller: &mut TransformStreamController) { + if chunk == 0 { + controller.enqueue_error(AiMuxError::InvalidResponseData("zero".into())); + return; + } + if chunk < 0 { + controller.error(AiMuxError::InvalidResponseData("negative".into())); + return; + } + self.total += chunk; + controller.enqueue(chunk * 2); + } + + fn flush(self, controller: &mut TransformStreamController) { + controller.enqueue(self.total); + } + } + + async fn run(input: Vec) -> Vec> { + pipe_through(futures::stream::iter(input), Doubler::default()) + .map(|item| item.map_err(|error| error.to_string())) + .collect() + .await + } + + #[tokio::test] + async fn start_transform_then_flush_in_order() { + assert_eq!(run(vec![1, 2]).await, vec![Ok(0), Ok(2), Ok(4), Ok(3)]); + } + + #[tokio::test] + async fn enqueue_error_keeps_transforming_and_flushes() { + let out = run(vec![1, 0, 2]).await; + assert_eq!(out.len(), 5); + assert_eq!(out[..2], [Ok(0), Ok(2)]); + assert!(out[2].is_err()); + assert_eq!(out[3..], [Ok(4), Ok(3)]); + } + + #[tokio::test] + async fn error_ends_the_stream_without_flush() { + let out = run(vec![1, -1, 5]).await; + assert_eq!(out.len(), 3); + assert_eq!(out[..2], [Ok(0), Ok(2)]); + assert!(out[2].is_err()); + } + + #[test] + fn a_transformer_can_be_driven_without_a_stream() { + let mut controller = TransformStreamController::new(); + let mut doubler = Doubler::default(); + doubler.start(&mut controller); + doubler.transform(3, &mut controller); + let emitted: Vec = controller.drain().map(Result::unwrap).collect(); + assert_eq!(emitted, [0, 6]); + doubler.transform(-1, &mut controller); + assert!(controller.is_closed()); + assert!(controller.drain().next().unwrap().is_err()); + } + + #[test] + fn nothing_is_emitted_after_the_stream_ended() { + let mut controller = TransformStreamController::new(); + controller.enqueue(1); + controller.terminate(); + controller.enqueue(2); + controller.enqueue_error(AiMuxError::InvalidResponseData("late".into())); + controller.error(AiMuxError::InvalidResponseData("late".into())); + let emitted: Vec = controller.drain().map(Result::unwrap).collect(); + assert_eq!(emitted, [1]); + assert!(controller.is_closed()); + } + + #[tokio::test] + async fn waits_for_input_that_is_not_ready_yet() { + let (tx, rx) = futures::channel::mpsc::unbounded(); + let mut out = pipe_through(rx, Doubler::default()); + assert_eq!(out.next().await.unwrap().unwrap(), 0); + tx.unbounded_send(4).unwrap(); + assert_eq!(out.next().await.unwrap().unwrap(), 8); + drop(tx); + assert_eq!(out.next().await.unwrap().unwrap(), 4); + assert!(out.next().await.is_none()); + assert!(out.next().await.is_none()); + } +} diff --git a/aimux-provider-utils/tests/t3_unattributable_delta_parity_test.rs b/aimux-provider-utils/tests/t3_unattributable_delta_parity_test.rs deleted file mode 100644 index 9a4faa063..000000000 --- a/aimux-provider-utils/tests/t3_unattributable_delta_parity_test.rs +++ /dev/null @@ -1,153 +0,0 @@ -//! Differential test against the pinned AI SDK tracker. -//! -//! `fixtures/aisdk/provider-utils/t3-unattributable-deltas.json` holds the -//! deltas for each case plus what the `StreamingToolCallTracker` of the -//! `@ai-sdk/provider-utils` version pinned in `fixtures/aisdk/VERSIONS.json` -//! (5.0.55) emitted for them with the case's `typeValidation` (`'if-present'`: -//! openai chat; `'required'`: groq; `'none'`: openai-compatible), recorded by -//! `scripts/gen-t3-tracker-fixture.mjs`. Each case asserts aimux's tracker, -//! configured the same way, matches per delta, on the error, and on flush. - -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; - -use aimux_core::error::AiMuxError; -use aimux_core::stream_part::StreamPart; -use aimux_core::tool::RawToolCall; -use aimux_provider_utils::{StreamingToolCallDelta, StreamingToolCallTracker, TypeValidation}; -use serde_json::{Value, json}; - -const FIXTURE: &str = - include_str!("../../fixtures/aisdk/provider-utils/t3-unattributable-deltas.json"); - -fn to_aisdk(part: StreamPart) -> Value { - match part { - StreamPart::ToolInputStart { id, tool_name, .. } => { - json!({"type": "tool-input-start", "id": id, "toolName": tool_name}) - } - StreamPart::ToolInputDelta { id, delta, .. } => { - json!({"type": "tool-input-delta", "id": id, "delta": delta}) - } - StreamPart::ToolInputEnd { id, .. } => json!({"type": "tool-input-end", "id": id}), - StreamPart::ToolCall(RawToolCall { - tool_call_id, - tool_name, - input, - .. - }) => json!({ - "type": "tool-call", "toolCallId": tool_call_id, "toolName": tool_name, "input": input - }), - other => panic!("unexpected part: {other:?}"), - } -} - -fn opt_str<'a>(v: &'a Value, key: &str) -> Option<&'a str> { - v.get(key).and_then(Value::as_str) -} - -fn run_case(case: &Value) { - let name = case["name"].as_str().unwrap(); - let expected = &case["aisdk"]; - let type_validation = match case["typeValidation"].as_str().unwrap() { - "if-present" => TypeValidation::IfPresent, - "required" => TypeValidation::Required, - "none" => TypeValidation::None, - other => panic!("{name}: unknown typeValidation {other}"), - }; - // Ids are generated `gen-0`, `gen-1`, ... like the recording script. - let next_id = Arc::new(AtomicUsize::new(0)); - let counter = Arc::clone(&next_id); - let mut tracker = StreamingToolCallTracker::new() - .with_generate_id(move || format!("gen-{}", counter.fetch_add(1, Ordering::SeqCst))) - .with_type_validation(type_validation); - let mut steps = Vec::new(); - let mut error = Value::Null; - for d in case["deltas"].as_array().unwrap() { - let f = &d["function"]; - let delta = StreamingToolCallDelta { - index: d.get("index").and_then(Value::as_u64).map(|i| i as usize), - id: opt_str(d, "id"), - r#type: opt_str(d, "type"), - name: opt_str(f, "name"), - arguments: opt_str(f, "arguments"), - provider_metadata: None, - }; - let mut out = Vec::new(); - let result = tracker.process(delta, &mut out); - steps.push(json!({"parts": out.into_iter().map(to_aisdk).collect::>()})); - if let Err(e) = result { - // The provider surfaces this as AiMuxError::InvalidResponseData. - let AiMuxError::InvalidResponseData(message) = AiMuxError::from(e) else { - panic!("{name}: tracker error is not InvalidResponseData"); - }; - error = json!({"name": "AI_InvalidResponseDataError", "message": message}); - break; - } - } - let flush = if error.is_null() { - let mut out = Vec::new(); - tracker.finish(&mut out); - Value::Array(out.into_iter().map(to_aisdk).collect()) - } else { - Value::Null - }; - let actual = json!({"steps": steps, "error": error, "flush": flush}); - assert_eq!(&actual, expected, "{name} ({})", case["tracker"]); -} - -fn case(name: &str) -> Value { - let fixture: Value = serde_json::from_str(FIXTURE).unwrap(); - fixture["cases"] - .as_array() - .unwrap() - .iter() - .find(|c| c["name"] == name) - .unwrap_or_else(|| panic!("missing fixture case {name}")) - .clone() -} - -/// Every recorded case, including the typeValidation `required` (groq) and -/// `none` (openai-compatible) ones and the interleaved ones. -#[test] -fn every_fixture_case_matches_the_aisdk_tracker() { - let fixture: Value = serde_json::from_str(FIXTURE).unwrap(); - let cases = fixture["cases"].as_array().unwrap(); - assert!(!cases.is_empty()); - for case in cases { - run_case(case); - } -} - -/// streaming-tool-call-tracker.ts:144-146, :305-308. -#[test] -fn ambiguous_delta_across_open_calls_is_dropped() { - run_case(&case("ambiguous_delta_across_open_calls_is_dropped")); -} - -/// streaming-tool-call-tracker.ts:153-155. -#[test] -fn blank_name_matching_no_call_is_dropped() { - run_case(&case("blank_name_matching_no_call_is_dropped")); -} - -/// streaming-tool-call-tracker.ts:305-308, :384-389. -#[test] -fn nameless_delta_without_open_call_is_invalid_response_data() { - run_case(&case( - "nameless_delta_without_open_call_is_invalid_response_data", - )); -} - -/// streaming-tool-call-tracker.ts:375-381. -#[test] -fn non_function_type_is_invalid_response_data() { - run_case(&case("non_function_type_is_invalid_response_data")); -} - -/// The case 5.0.55 changed: an index-only continuation after the index is -/// reused goes to the only call there without complete structured arguments -/// (streaming-tool-call-tracker.ts:291-298, :345-347); 5.0.53 dropped it. -#[test] -fn index_only_continuation_goes_to_the_only_incomplete_call() { - run_case(&case("index_reuse_with_new_id")); -} diff --git a/aimux-providers/src/anthropic/stream.rs b/aimux-providers/src/anthropic/stream.rs index 219a205e6..d1a1bee78 100644 --- a/aimux-providers/src/anthropic/stream.rs +++ b/aimux-providers/src/anthropic/stream.rs @@ -25,14 +25,18 @@ use aimux_core::shared::provider_namespace; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; use aimux_core::types::{NamespacedKind, ProviderMetadata, Warning}; -use aimux_provider_utils::HttpRequest; +use aimux_provider_utils::{HttpRequest, Transformer, pipe_through}; use serde_json::{Value, json}; use super::config::AnthropicModelConfig; use super::convert::parse_stop_reason; use super::options::CANONICAL; use super::tool_name_mapping::ToolNameMapping; -use super::types::{AnthropicResponse, ContentBlock, StreamErrorData, StreamEvent, ToolCallCaller}; +use super::types::{ + AnthropicResponse, AnthropicUsage, ContentBlock, DeltaBlock, MessageDeltaBody, + MessageStartData, StreamErrorData, StreamEvent, ToolCallCaller, +}; +use crate::stream::{Controller, Emit, FINISH_OTHER, StreamContext}; pub(crate) fn anthropic_stream_error( error: &StreamErrorData, @@ -1067,10 +1071,7 @@ pub(crate) async fn anthropic_generate_core( .stop_reason .as_deref() .map(parse_stop_reason) - .unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }); + .unwrap_or(FINISH_OTHER); if is_json_response_from_tool && finish_reason.unified == FinishReasonUnified::ToolCalls { finish_reason.unified = FinishReasonUnified::Stop; @@ -1152,7 +1153,7 @@ pub(crate) async fn anthropic_stream_core( options: &CallOptions, used_custom_options_key: bool, uses_json_response_tool: bool, - mut citation_documents: Vec, + citation_documents: Vec, generate_id: std::sync::Arc String + Send + Sync>, ) -> Result { let tool_names = ToolNameMapping::new(options.tools.as_deref()); @@ -1194,487 +1195,747 @@ pub(crate) async fn anthropic_stream_core( response_headers, )); } - let stream_error_url = endpoint; - let stream_request_body = body.clone(); - let stream_response_headers = response_headers.clone(); - - let stream = async_stream::stream! { - // First part: StreamStart. - yield Ok(StreamPart::StreamStart { warnings }); - - let mut sse = futures::stream::iter(first_event.into_iter()).chain(sse_stream); - let mut blocks: HashMap = HashMap::new(); - let mut is_json_response_from_tool = false; - let mut final_usage = Usage::default(); - let mut final_finish_reason: Option = None; - // Result-level providerMetadata: the raw usage (`message_start`'s, - // updated by every `message_delta`'s), the stop sequence, the - // container and the context-management edits. - let mut raw_usage = Value::Object(serde_json::Map::new()); - let mut effective_usage = raw_usage.clone(); - let mut stop_sequence: Option = None; - let mut container: Option = None; - let mut context_management: Option = None; - let mut stop_details: Option = None; - let mut input_transformations: Option = None; - let mut safeguard_results: Option = None; - let mut response_meta_emitted = false; - let mut active_message_id: Option = None; - let mut message_stopped = false; - - // id → (tool name, server name), so `mcp_tool_result` can inherit them - // from the `mcp_tool_use` it answers. - let mut mcp_tool_calls: HashMap = HashMap::new(); - // tool_use_id → provider tool name. Both tool-search variants share - // one result block type, so the id is required to disambiguate aliases. - let mut server_tool_calls: HashMap = HashMap::new(); - - while let Some((raw, event)) = sse.next().await { - // Every chunk, valid or not, is forwarded first when raw chunks - // are requested (anthropic@4.0.71 `anthropic-language-model.ts:1956-1958`); - // a chunk that is not JSON has no raw value. - if include_raw_chunks - && (raw.is_some() - || event.as_ref().is_err_and(AiMuxError::is_recoverable_stream_error)) - { - yield Ok(StreamPart::Raw { - raw_value: raw.unwrap_or(Value::Null), + + let raw_usage = Value::Object(serde_json::Map::new()); + let stream = pipe_through( + futures::stream::iter(first_event).chain(sse_stream), + AnthropicMessagesStream { + ctx: StreamContext { + generate_id, + ..StreamContext::new(options, warnings, endpoint, &body, &response_headers) + }, + options_name, + tool_names, + used_custom_options_key, + uses_json_response_tool, + citation_documents, + effective_usage: raw_usage.clone(), + raw_usage, + ..AnthropicMessagesStream::default() + }, + ); + + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), + response: Some(aimux_core::shared::StreamResponseInfo { + headers: Some(response_headers), + }), + }) +} + +/// The `TransformStream` of `AnthropicMessagesLanguageModel.doStream`. +#[derive(Default)] +struct AnthropicMessagesStream { + ctx: StreamContext, + options_name: String, + tool_names: ToolNameMapping, + used_custom_options_key: bool, + uses_json_response_tool: bool, + citation_documents: Vec, + blocks: HashMap, + is_json_response_from_tool: bool, + final_usage: Usage, + final_finish_reason: Option, + /// Result-level providerMetadata: the raw usage (`message_start`'s, + /// updated by every `message_delta`'s), the stop sequence, the + /// container and the context-management edits. + raw_usage: Value, + effective_usage: Value, + stop_sequence: Option, + container: Option, + context_management: Option, + stop_details: Option, + input_transformations: Option, + safeguard_results: Option, + response_meta_emitted: bool, + active_message_id: Option, + /// Set by `message_stop`; later events are ignored and only then does + /// `flush` emit the finish part. + message_stopped: bool, + /// id → (tool name, server name), so `mcp_tool_result` can inherit them + /// from the `mcp_tool_use` it answers. + mcp_tool_calls: HashMap, + /// tool_use_id → provider tool name. Both tool-search variants share + /// one result block type, so the id is required to disambiguate aliases. + server_tool_calls: HashMap, +} + +impl AnthropicMessagesStream { + fn message_start(&mut self, message: MessageStartData, controller: &mut Controller) { + if let Some(active_id) = &self.active_message_id { + if active_id == &message.id { + return; + } + controller.enqueue(StreamPart::Error { + error: AiMuxError::InvalidResponseData(format!( + "Received message_start for message {:?} while message {:?} is still open.", + message.id, active_id, + )), + }); + controller.terminate(); + return; + } + self.active_message_id = Some(message.id.clone()); + self.input_transformations = message + .input_transformations + .or(self.input_transformations.take()); + self.container = message.container.or(self.container.take()); + if let Some(reason) = message.stop_reason.as_deref() { + self.final_finish_reason = Some(parse_stop_reason(reason)); + } + if let Some(usage) = &message.usage { + // RFC-0015 P0-2: full input side incl. cache + // fields + raw (Anthropic reports cache only + // in message_start). + self.final_usage = super::usage::usage_from_anthropic(usage); + if let Ok(raw_usage) = serde_json::to_value(usage) { + self.raw_usage = raw_usage; + } + self.effective_usage = self.raw_usage.clone(); + } + if !self.response_meta_emitted { + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: Some(message.id.clone()), + timestamp: None, + model_id: Some(message.model.clone()), + })); + self.response_meta_emitted = true; + } + } + + fn content_block_start( + &mut self, + index: usize, + content_block: ContentBlock, + controller: &mut Controller, + ) { + match content_block { + ContentBlock::Text { .. } if self.uses_json_response_tool => {} + ContentBlock::Text { .. } => { + self.blocks.insert( + index, + BlockState::Text { + citations: Vec::new(), + }, + ); + controller.text_start(index.to_string()); + } + ContentBlock::Compaction { content, signature } => { + self.blocks.insert( + index, + BlockState::Text { + citations: Vec::new(), + }, + ); + controller.enqueue(StreamPart::TextStart { + id: index.to_string(), + provider_metadata: Some(compaction_metadata(signature.as_deref())), }); + if signature.is_some() + && let Some(text) = content.filter(|text| !text.is_empty()) + { + controller.text_delta(index.to_string(), text); + } } - match event { - Ok(stream_event) => { - match stream_event { - StreamEvent::MessageStart { message } => { - if let Some(active_id) = &active_message_id { - if active_id == &message.id { - continue; - } - yield Ok(StreamPart::Error { - error: AiMuxError::InvalidResponseData(format!( - "Received message_start for message {:?} while message {:?} is still open.", - message.id, active_id, - )), - }); - return; - } - active_message_id = Some(message.id.clone()); - input_transformations = message.input_transformations.or(input_transformations); - container = message.container.or(container); - if let Some(reason) = message.stop_reason.as_deref() { - final_finish_reason = Some(parse_stop_reason(reason)); - } - if let Some(usage) = &message.usage { - // RFC-0015 P0-2: full input side incl. cache - // fields + raw (Anthropic reports cache only - // in message_start). - final_usage = super::usage::usage_from_anthropic(usage); - raw_usage = serde_json::to_value(usage).unwrap_or(raw_usage); - effective_usage = raw_usage.clone(); - } - if !response_meta_emitted { - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: Some(message.id.clone()), - timestamp: None, - model_id: Some(message.model.clone()), - })); - response_meta_emitted = true; - } - } - StreamEvent::ContentBlockStart { index, content_block } => { - match content_block { - ContentBlock::Text { .. } if uses_json_response_tool => {} - ContentBlock::Text { .. } => { - blocks.insert(index, BlockState::Text { citations: Vec::new() }); - yield Ok(StreamPart::TextStart { id: index.to_string(), provider_metadata: None }); - } - ContentBlock::Compaction { content, signature } => { - blocks.insert(index, BlockState::Text { citations: Vec::new() }); - yield Ok(StreamPart::TextStart { id: index.to_string(), provider_metadata: Some(compaction_metadata(signature.as_deref())) }); - if signature.is_some() - && let Some(text) = content.filter(|text| !text.is_empty()) { - yield Ok(StreamPart::TextDelta { id: index.to_string(), delta: text, provider_metadata: None }); - } - } - ContentBlock::Thinking { .. } => { - yield Ok(StreamPart::ReasoningStart { id: index.to_string(), provider_metadata: None }); - blocks.insert( - index, - BlockState::Thinking, - ); - } - ContentBlock::ToolUse { name, .. } if uses_json_response_tool && name == "json" => { - is_json_response_from_tool = true; - blocks.insert(index, BlockState::Text { citations: Vec::new() }); - yield Ok(StreamPart::TextStart { id: index.to_string(), provider_metadata: None }); - } - ContentBlock::ToolUse { - id, - name, - input, - toolset_name, - caller, - } => { - let custom_name = toolset_name.as_deref().map_or_else( - || name.clone(), - |name| tool_names.to_custom_tool_name(name).to_string(), - ); - let initial_input = initial_tool_input(&input); - yield Ok(StreamPart::ToolInputStart { - id: id.clone(), - tool_name: custom_name.clone(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - blocks.insert(index, BlockState::ToolUse { - id, - name: custom_name, - first_delta: initial_input.is_empty(), - accumulated_json: initial_input, - provider_executed: None, - dynamic: None, - provider_tool_name: None, - provider_tool_input_type: None, - provider_metadata: tool_call_metadata(caller.as_ref(), toolset_name.as_deref(), CANONICAL), - toolset_member_name: toolset_name.map(|_| name), - }); - } - // Server-side tool use follows the same input - // lifecycle as client tools because some code - // execution inputs arrive entirely via deltas. - ContentBlock::ServerToolUse { id, name, input, caller } => { - if is_tool_search_provider_name(&name) { - server_tool_calls.insert(id.clone(), name.clone()); - } - let provider_name = server_tool_provider_name(&name); - let custom_name = tool_names - .to_custom_tool_name(provider_name) - .to_string(); - let dynamic = (provider_name == "code_execution" - && tool_names.mark_code_execution_dynamic()) - .then_some(true); - let initial_input = initial_tool_input(&input); - yield Ok(StreamPart::ToolInputStart { - id: id.clone(), - tool_name: custom_name.clone(), - provider_executed: Some(true), - dynamic, - title: None, - provider_metadata: None, - }); - blocks.insert(index, BlockState::ToolUse { - id, - name: custom_name, - first_delta: initial_input.is_empty(), - accumulated_json: initial_input, - provider_executed: Some(true), - dynamic, - provider_tool_name: Some(provider_name.to_string()), - provider_tool_input_type: match name.as_str() { - "text_editor_code_execution" | "bash_code_execution" => { - Some(name) - } - "code_execution" => { - Some("programmatic-tool-call".to_string()) - } - _ => None, - }, - provider_metadata: tool_call_caller_metadata(caller.as_ref(), CANONICAL), - toolset_member_name: None, - }); - } - // MCP tool use — provider-executed + dynamic. - ContentBlock::McpToolUse { id, name, input, server_name } => { - mcp_tool_calls - .insert(id.clone(), (name.clone(), server_name.clone())); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id.clone(), - tool_name: name.clone(), - input: input.to_string(), - provider_executed: Some(true), - dynamic: Some(true), - provider_metadata: Some(provider_namespace(CANONICAL, json!({ - "type": "mcp-tool-use", - "serverName": server_name, - })).expect("provider metadata must be an object")), - })); - } - // Redacted thinking — emit as ReasoningStart. - ContentBlock::RedactedThinking { data } => { - let id = index.to_string(); - yield Ok(StreamPart::ReasoningStart { - id: id.clone(), - provider_metadata: Some(provider_namespace(CANONICAL, json!({ "redactedData": data })).expect("provider metadata must be an object")), - }); - blocks.insert( - index, - BlockState::Thinking, - ); - } - // Server-tool result blocks arrive whole on - // `content_block_start`, so they reuse the - // non-streaming payload mapping. Upstream - // mirrors its `doGenerate` switch here too - // (anthropic-language-model.ts:1901-2178). - other => { - if let ContentBlock::WebFetchToolResult { content, .. } = &other - && let Some(document) = web_fetch_document(content) { - citation_documents.push(document); - } - for part in stream_parts_for_result_block( - CANONICAL, - &other, - &tool_names, - &mcp_tool_calls, - &server_tool_calls, - generate_id.as_ref(), - ) { - yield Ok(part); - } - } - } - } - StreamEvent::ContentBlockDelta { index, delta } => { - if !uses_json_response_tool && let Some(text) = delta.text { - yield Ok(StreamPart::TextDelta { - id: index.to_string(), - delta: text, - provider_metadata: None, - }); - } - if let Some(text) = delta.content { - yield Ok(StreamPart::TextDelta { id: index.to_string(), delta: text, provider_metadata: None }); - } - if let Some(citation) = delta.citation { - if let Some(BlockState::Text { citations }) = blocks.get_mut(&index) { - citations.push(citation.clone()); - } - if let Some(source) = citation_source(&citation, &citation_documents, generate_id.as_ref()) { - yield Ok(StreamPart::Source(source)); - } - } - if let Some(partial) = delta.partial_json { - // Accumulate the partial JSON fragment and emit - // a ToolInputDelta. Empty fragments (the - // leading `input_json_delta` with - // `partial_json: ""`) are skipped, matching the - // TS SDK. - if is_json_response_from_tool && matches!(blocks.get(&index), Some(BlockState::Text { .. })) && !partial.is_empty() { - yield Ok(StreamPart::TextDelta { id: index.to_string(), delta: partial, provider_metadata: None }); - continue; - } - if is_json_response_from_tool { continue; } - let delta_event: Option<(String, String)> = - match blocks.get_mut(&index) { - Some(BlockState::ToolUse { - id, - accumulated_json, - provider_tool_input_type, - first_delta, - toolset_member_name, - .. - }) if !partial.is_empty() => { - let emitted_delta = if *first_delta { - if let Some(input_type) = provider_tool_input_type { - format!( - "{{\"type\": \"{input_type}\",{}", - partial.strip_prefix('{').unwrap_or(&partial) - ) - } else { - partial - } - } else { - partial - }; - accumulated_json.push_str(&emitted_delta); - *first_delta = false; - toolset_member_name.is_none().then(|| (id.clone(), emitted_delta)) - } - _ => None, - }; - if let Some((id, delta)) = delta_event { - yield Ok(StreamPart::ToolInputDelta { - id, - delta, - provider_metadata: None, - }); - } - } - if let Some(thinking) = delta.thinking { - yield Ok(StreamPart::ReasoningDelta { - id: index.to_string(), - delta: thinking, - provider_metadata: None, - }); - } - if let Some(sig) = delta.signature - && let Some(BlockState::Thinking) = blocks.get(&index) { - yield Ok(StreamPart::ReasoningDelta { - id: index.to_string(), delta: String::new(), - provider_metadata: Some(provider_namespace(CANONICAL, json!({ "signature": sig })).expect("provider metadata must be an object")), - }); - } - } - StreamEvent::ContentBlockStop { index } => { - // Removing the block releases the borrow before any - // yield. - if let Some(state) = blocks.remove(&index) { - match state { - BlockState::Text { citations } => { - yield Ok(StreamPart::TextEnd { - id: index.to_string(), - provider_metadata: citation_metadata(&citations), - }); - } - BlockState::Thinking => { - yield Ok(StreamPart::ReasoningEnd { - id: index.to_string(), - provider_metadata: None, - }); - } - BlockState::ToolUse { - id, - name, - mut accumulated_json, - provider_executed, - dynamic, - provider_tool_name, - provider_tool_input_type, - provider_metadata, - toolset_member_name, - .. - } => { - if let Some(member_name) = toolset_member_name { - // Upstream parses with secureJsonParse and keeps - // the raw input when that fails - // (anthropic-language-model.ts:2606-2632). - let parsed = if accumulated_json.is_empty() { - Some(json!({})) - } else { - serde_json::from_str::(&accumulated_json) - .ok() - .filter(|input| !contains_forbidden_prototype(input)) - }; - if let Some(input) = parsed { - accumulated_json = toolset_member_input(&member_name, &input).to_string(); - } - yield Ok(StreamPart::ToolInputDelta { - id: id.clone(), - delta: accumulated_json.clone(), - provider_metadata: None, - }); - } - yield Ok(StreamPart::ToolInputEnd { - id: id.clone(), - provider_metadata: None, - }); - let input = - finalize_streamed_tool_input( - accumulated_json, - provider_tool_name.as_deref(), - provider_tool_input_type.as_deref(), - ); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id, - tool_name: name, - input, - provider_executed, - dynamic, - provider_metadata, - })); - } - } - } - } - StreamEvent::MessageDelta { delta, usage, context_management: edits, input_transformations: transformations } => { - if let Some(reason) = delta.stop_reason { - final_finish_reason = Some(parse_stop_reason(&reason)); - } - stop_sequence = delta.stop_sequence; - container = delta.container; - context_management = edits.or(context_management); - stop_details = delta.stop_details; - input_transformations = transformations.or(input_transformations); - safeguard_results = delta.safeguard_results.or(safeguard_results); - if let Some(u) = usage { - if let (Value::Object(raw), Ok(Value::Object(update))) = - (&mut raw_usage, serde_json::to_value(&u)) - { - if let Value::Object(effective) = &mut effective_usage { - effective.extend(update.iter().filter(|(_, value)| !value.is_null()).map(|(key, value)| (key.clone(), value.clone()))); - } - raw.extend(update); - } - if let Ok(usage) = serde_json::from_value(effective_usage.clone()) { - final_usage = super::usage::usage_from_anthropic(&usage); - final_usage.raw = raw_usage.as_object().cloned(); - } - } - } - StreamEvent::MessageStop => { - message_stopped = true; - break; + ContentBlock::Thinking { .. } => { + controller.reasoning_start(index.to_string()); + self.blocks.insert(index, BlockState::Thinking); + } + ContentBlock::ToolUse { name, .. } + if self.uses_json_response_tool && name == "json" => + { + self.is_json_response_from_tool = true; + self.blocks.insert( + index, + BlockState::Text { + citations: Vec::new(), + }, + ); + controller.text_start(index.to_string()); + } + ContentBlock::ToolUse { + id, + name, + input, + toolset_name, + caller, + } => { + let custom_name = toolset_name.as_deref().map_or_else( + || name.clone(), + |name| self.tool_names.to_custom_tool_name(name).to_string(), + ); + let initial_input = initial_tool_input(&input); + controller.tool_input_start(id.clone(), custom_name.clone()); + self.blocks.insert( + index, + BlockState::ToolUse { + id, + name: custom_name, + first_delta: initial_input.is_empty(), + accumulated_json: initial_input, + provider_executed: None, + dynamic: None, + provider_tool_name: None, + provider_tool_input_type: None, + provider_metadata: tool_call_metadata( + caller.as_ref(), + toolset_name.as_deref(), + CANONICAL, + ), + toolset_member_name: toolset_name.map(|_| name), + }, + ); + } + // Server-side tool use follows the same input + // lifecycle as client tools because some code + // execution inputs arrive entirely via deltas. + ContentBlock::ServerToolUse { + id, + name, + input, + caller, + } => { + if is_tool_search_provider_name(&name) { + self.server_tool_calls.insert(id.clone(), name.clone()); + } + let provider_name = server_tool_provider_name(&name); + let custom_name = self + .tool_names + .to_custom_tool_name(provider_name) + .to_string(); + let dynamic = (provider_name == "code_execution" + && self.tool_names.mark_code_execution_dynamic()) + .then_some(true); + let initial_input = initial_tool_input(&input); + controller.enqueue(StreamPart::ToolInputStart { + id: id.clone(), + tool_name: custom_name.clone(), + provider_executed: Some(true), + dynamic, + title: None, + provider_metadata: None, + }); + self.blocks.insert( + index, + BlockState::ToolUse { + id, + name: custom_name, + first_delta: initial_input.is_empty(), + accumulated_json: initial_input, + provider_executed: Some(true), + dynamic, + provider_tool_name: Some(provider_name.to_string()), + provider_tool_input_type: match name.as_str() { + "text_editor_code_execution" | "bash_code_execution" => Some(name), + "code_execution" => Some("programmatic-tool-call".to_string()), + _ => None, }, - StreamEvent::Error { error } => { - yield Ok(StreamPart::Error { - error: anthropic_stream_error( - &error, - &stream_error_url, - stream_request_body.clone(), - stream_response_headers.clone(), - ), - }); + provider_metadata: tool_call_caller_metadata(caller.as_ref(), CANONICAL), + toolset_member_name: None, + }, + ); + } + // MCP tool use — provider-executed + dynamic. + ContentBlock::McpToolUse { + id, + name, + input, + server_name, + } => { + self.mcp_tool_calls + .insert(id.clone(), (name.clone(), server_name.clone())); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: id.clone(), + tool_name: name.clone(), + input: input.to_string(), + provider_executed: Some(true), + dynamic: Some(true), + provider_metadata: Some( + provider_namespace( + CANONICAL, + json!({ + "type": "mcp-tool-use", + "serverName": server_name, + }), + ) + .expect("provider metadata must be an object"), + ), + })); + } + // Redacted thinking — emit as ReasoningStart. + ContentBlock::RedactedThinking { data } => { + let id = index.to_string(); + controller.enqueue(StreamPart::ReasoningStart { + id: id.clone(), + provider_metadata: Some( + provider_namespace(CANONICAL, json!({ "redactedData": data })) + .expect("provider metadata must be an object"), + ), + }); + self.blocks.insert(index, BlockState::Thinking); + } + // Server-tool result blocks arrive whole on + // `content_block_start`, so they reuse the + // non-streaming payload mapping. Upstream + // mirrors its `doGenerate` switch here too + // (anthropic-language-model.ts:1901-2178). + other => { + if let ContentBlock::WebFetchToolResult { content, .. } = &other + && let Some(document) = web_fetch_document(content) + { + self.citation_documents.push(document); + } + for part in stream_parts_for_result_block( + CANONICAL, + &other, + &self.tool_names, + &self.mcp_tool_calls, + &self.server_tool_calls, + self.ctx.generate_id.as_ref(), + ) { + controller.enqueue(part); + } + } + } + } + fn content_block_delta( + &mut self, + index: usize, + delta: DeltaBlock, + controller: &mut Controller, + ) { + if !self.uses_json_response_tool + && let Some(text) = delta.text + { + controller.text_delta(index.to_string(), text); + } + if let Some(text) = delta.content { + controller.text_delta(index.to_string(), text); + } + if let Some(citation) = delta.citation { + if let Some(BlockState::Text { citations }) = self.blocks.get_mut(&index) { + citations.push(citation.clone()); + } + if let Some(source) = citation_source( + &citation, + &self.citation_documents, + self.ctx.generate_id.as_ref(), + ) { + controller.enqueue(StreamPart::Source(source)); + } + } + if let Some(partial) = delta.partial_json { + // Accumulate the partial JSON fragment and emit + // a ToolInputDelta. Empty fragments (the + // leading `input_json_delta` with + // `partial_json: ""`) are skipped, matching the + // TS SDK. + if self.is_json_response_from_tool + && matches!(self.blocks.get(&index), Some(BlockState::Text { .. })) + && !partial.is_empty() + { + controller.text_delta(index.to_string(), partial); + return; + } + if self.is_json_response_from_tool { + return; + } + let delta_event: Option<(String, String)> = match self.blocks.get_mut(&index) { + Some(BlockState::ToolUse { + id, + accumulated_json, + provider_tool_input_type, + first_delta, + toolset_member_name, + .. + }) if !partial.is_empty() => { + let emitted_delta = if *first_delta { + if let Some(input_type) = provider_tool_input_type { + format!( + "{{\"type\": \"{input_type}\",{}", + partial.strip_prefix('{').unwrap_or(&partial) + ) + } else { + partial } - _ => {} - } + } else { + partial + }; + accumulated_json.push_str(&emitted_delta); + *first_delta = false; + toolset_member_name + .is_none() + .then(|| (id.clone(), emitted_delta)) } - Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Ok(StreamPart::Error { error }); - if !recoverable { - return; + _ => None, + }; + if let Some((id, delta)) = delta_event { + controller.tool_input_delta(id, delta); + } + } + if let Some(thinking) = delta.thinking { + controller.reasoning_delta(index.to_string(), thinking); + } + if let Some(sig) = delta.signature + && let Some(BlockState::Thinking) = self.blocks.get(&index) + { + controller.enqueue(StreamPart::ReasoningDelta { + id: index.to_string(), + delta: String::new(), + provider_metadata: Some( + provider_namespace(CANONICAL, json!({ "signature": sig })) + .expect("provider metadata must be an object"), + ), + }); + } + } + + fn content_block_stop(&mut self, index: usize, controller: &mut Controller) { + let Some(state) = self.blocks.remove(&index) else { + return; + }; + match state { + BlockState::Text { citations } => { + controller.enqueue(StreamPart::TextEnd { + id: index.to_string(), + provider_metadata: citation_metadata(&citations), + }); + } + BlockState::Thinking => { + controller.reasoning_end(index.to_string()); + } + BlockState::ToolUse { + id, + name, + mut accumulated_json, + provider_executed, + dynamic, + provider_tool_name, + provider_tool_input_type, + provider_metadata, + toolset_member_name, + .. + } => { + if let Some(member_name) = toolset_member_name { + // Upstream parses with secureJsonParse and keeps the raw + // input when that fails (anthropic-language-model.ts:2606-2632). + let parsed = if accumulated_json.is_empty() { + Some(json!({})) + } else { + serde_json::from_str::(&accumulated_json) + .ok() + .filter(|input| !contains_forbidden_prototype(input)) + }; + if let Some(input) = parsed { + accumulated_json = toolset_member_input(&member_name, &input).to_string(); } + controller.tool_input_delta(id.clone(), accumulated_json.clone()); + } + controller.tool_input_end(id.clone()); + let input = finalize_streamed_tool_input( + accumulated_json, + provider_tool_name.as_deref(), + provider_tool_input_type.as_deref(), + ); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: id, + tool_name: name, + input, + provider_executed, + dynamic, + provider_metadata, + })); + } + } + } + + fn message_delta( + &mut self, + delta: MessageDeltaBody, + usage: Option, + edits: Option, + transformations: Option, + ) { + if let Some(reason) = delta.stop_reason { + self.final_finish_reason = Some(parse_stop_reason(&reason)); + } + self.stop_sequence = delta.stop_sequence; + self.container = delta.container; + self.context_management = edits.or(self.context_management.take()); + self.stop_details = delta.stop_details; + self.input_transformations = transformations.or(self.input_transformations.take()); + self.safeguard_results = delta.safeguard_results.or(self.safeguard_results.take()); + if let Some(u) = usage { + if let (Value::Object(raw), Ok(Value::Object(update))) = + (&mut self.raw_usage, serde_json::to_value(&u)) + { + if let Value::Object(effective) = &mut self.effective_usage { + effective.extend( + update + .iter() + .filter(|(_, value)| !value.is_null()) + .map(|(key, value)| (key.clone(), value.clone())), + ); } + raw.extend(update); + } + if let Ok(usage) = serde_json::from_value(self.effective_usage.clone()) { + self.final_usage = super::usage::usage_from_anthropic(&usage); + self.final_usage.raw = self.raw_usage.as_object().cloned(); } } + } +} + +impl Transformer for AnthropicMessagesStream { + /// Each event as the AI SDK's `ParseResult`: the raw JSON (kept only when + /// raw chunks are requested) and the chunk schema's result. + type Input = (Option, Result); + type Output = StreamPart; + + fn start(&mut self, controller: &mut Controller) { + self.ctx.start(controller); + } + fn transform( + &mut self, + (raw, event): (Option, Result), + controller: &mut Controller, + ) { + // Every chunk, valid or not, is forwarded first when raw chunks are + // requested (anthropic@4.0.71 `anthropic-language-model.ts:1956-1958`); + // a chunk that is not JSON has no raw value. + if self.ctx.include_raw_chunks + && (raw.is_some() + || event + .as_ref() + .is_err_and(AiMuxError::is_recoverable_stream_error)) + { + controller.enqueue(StreamPart::Raw { + raw_value: raw.unwrap_or(Value::Null), + }); + } + let stream_event = match event { + Ok(stream_event) => stream_event, + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + controller.enqueue(StreamPart::Error { error }); + if !recoverable { + controller.terminate(); + } + return; + } + }; + match stream_event { + StreamEvent::MessageStart { message } => self.message_start(message, controller), + StreamEvent::ContentBlockStart { + index, + content_block, + } => self.content_block_start(index, content_block, controller), + StreamEvent::ContentBlockDelta { index, delta } => { + self.content_block_delta(index, delta, controller); + } + StreamEvent::ContentBlockStop { index } => self.content_block_stop(index, controller), + StreamEvent::MessageDelta { + delta, + usage, + context_management: edits, + input_transformations: transformations, + } => self.message_delta(delta, usage, edits, transformations), + StreamEvent::MessageStop => { + // The rest of the body is not read: the stream finishes now. + self.message_stopped = true; + self.finish(controller); + controller.terminate(); + } + StreamEvent::Error { error } => { + controller.enqueue(StreamPart::Error { + error: anthropic_stream_error( + &error, + &self.ctx.url, + self.ctx.request_body.clone(), + self.ctx.response_headers.clone(), + ), + }); + } + StreamEvent::Other => {} + } + } + + fn flush(mut self, controller: &mut Controller) { + self.finish(controller); + } +} + +impl AnthropicMessagesStream { + /// Emit `Finish` once `message_stop` has arrived. + fn finish(&mut self, controller: &mut Controller) { // The provider completes a message explicitly; EOF after an error // does not synthesize a finish event. - if !message_stopped { + if !self.message_stopped { return; } - if is_json_response_from_tool && let Some(reason) = &mut final_finish_reason - && reason.unified == FinishReasonUnified::ToolCalls { + if self.is_json_response_from_tool + && let Some(reason) = &mut self.final_finish_reason + && reason.unified == FinishReasonUnified::ToolCalls + { reason.unified = FinishReasonUnified::Stop; } - yield Ok(StreamPart::Finish { - finish_reason: final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }), - usage: final_usage, - provider_metadata: Some({ - let mut metadata = super::usage::result_provider_metadata( - options_name.as_str(), - &raw_usage, - stop_sequence.as_deref(), - container.as_ref(), - context_management.as_ref(), - used_custom_options_key, - ); - super::usage::extend_result_metadata(&mut metadata, stop_details.as_ref(), input_transformations.as_ref(), safeguard_results.as_ref()); - metadata - }), + let mut metadata = super::usage::result_provider_metadata( + self.options_name.as_str(), + &self.raw_usage, + self.stop_sequence.as_deref(), + self.container.as_ref(), + self.context_management.as_ref(), + self.used_custom_options_key, + ); + super::usage::extend_result_metadata( + &mut metadata, + self.stop_details.as_ref(), + self.input_transformations.as_ref(), + self.safeguard_results.as_ref(), + ); + controller.enqueue(StreamPart::Finish { + finish_reason: self.final_finish_reason.take().unwrap_or(FINISH_OTHER), + usage: std::mem::take(&mut self.final_usage), + provider_metadata: Some(metadata), }); - }; + } +} - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - response: Some(aimux_core::shared::StreamResponseInfo { - headers: Some(response_headers), - }), - }) +/// The Messages stream driven event by event. +#[cfg(test)] +mod stream_tests { + use super::*; + use crate::stream::testing::run; + + fn stream() -> AnthropicMessagesStream { + AnthropicMessagesStream { + options_name: CANONICAL.into(), + ..Default::default() + } + } + + fn event(event: Value) -> (Option, Result) { + ( + Some(event.clone()), + serde_json::from_value(event).map_err(AiMuxError::from), + ) + } + + fn message_start() -> (Option, Result) { + event(json!({ + "type": "message_start", + "message": { + "id": "msg_1", "model": "claude-x", + "usage": { "input_tokens": 3, "output_tokens": 1 } + } + })) + } + + fn text_block( + index: usize, + text: &str, + ) -> [(Option, Result); 3] { + [ + event(json!({ + "type": "content_block_start", "index": index, + "content_block": { "type": "text", "text": "" } + })), + event(json!({ + "type": "content_block_delta", "index": index, + "delta": { "type": "text_delta", "text": text } + })), + event(json!({ "type": "content_block_stop", "index": index })), + ] + } + + #[test] + fn interleaved_text_and_tool_blocks_then_message_stop_finish_the_stream() { + let mut input = vec![message_start()]; + input.push(event(json!({ + "type": "content_block_start", "index": 1, + "content_block": { "type": "tool_use", "id": "toolu_1", "name": "weather", "input": {} } + }))); + input.extend(text_block(0, "Checking")); + input.extend([ + event(json!({ + "type": "content_block_delta", "index": 1, + "delta": { "type": "input_json_delta", "partial_json": "{\"city\":" } + })), + event(json!({ + "type": "content_block_delta", "index": 1, + "delta": { "type": "input_json_delta", "partial_json": "\"Paris\"}" } + })), + event(json!({ "type": "content_block_stop", "index": 1 })), + event(json!({ + "type": "message_delta", + "delta": { "stop_reason": "tool_use" }, + "usage": { "output_tokens": 9 } + })), + event(json!({ "type": "message_stop" })), + // Not read: the stream ended at `message_stop`. + event(json!({ + "type": "content_block_start", "index": 2, + "content_block": { "type": "text", "text": "" } + })), + ]); + assert_eq!( + run(stream(), input), + [ + "stream-start", + "response-metadata msg_1", + "tool-input-start toolu_1 weather", + "text-start 0", + "text-delta 0 Checking", + "text-end 0", + "tool-input-delta toolu_1 {\"city\":", + "tool-input-delta toolu_1 \"Paris\"}", + "tool-input-end toolu_1", + "tool-call toolu_1 weather {\"city\":\"Paris\"}", + "finish ToolCalls", + ] + ); + } + + #[test] + fn an_error_event_is_an_error_part_and_the_stream_goes_on() { + let mut input = vec![message_start()]; + input.push(event(json!({ + "type": "error", + "error": { "type": "overloaded_error", "message": "Overloaded" } + }))); + input.extend(text_block(0, "late")); + input.push(event(json!({ "type": "message_stop" }))); + assert_eq!( + run(stream(), input), + [ + "stream-start", + "response-metadata msg_1", + "error API call error: HTTP 529: Overloaded", + "text-start 0", + "text-delta 0 late", + "text-end 0", + "finish Other", + ] + ); + } + + #[test] + fn flush_without_message_stop_emits_no_finish() { + let mut input = vec![message_start()]; + input.extend(text_block(0, "cut")); + assert_eq!( + run(stream(), input), + [ + "stream-start", + "response-metadata msg_1", + "text-start 0", + "text-delta 0 cut", + "text-end 0", + ] + ); + } } diff --git a/aimux-providers/src/bedrock/model.rs b/aimux-providers/src/bedrock/model.rs index c6108d2a9..b7e71cf8e 100644 --- a/aimux-providers/src/bedrock/model.rs +++ b/aimux-providers/src/bedrock/model.rs @@ -23,11 +23,14 @@ use aimux_core::types::{ use serde_json::json; -use aimux_provider_utils::HttpBody; +use aimux_provider_utils::{HttpBody, Transformer, pipe_through}; + +use crate::stream::{Controller, Emit, FINISH_ERROR, FINISH_OTHER, StreamContext}; use super::convert::{ build_request_body_checked, convert_usage, map_finish_reason, normalize_tool_call_id, }; +use super::event_stream::EventStreamMessage; use super::options; use super::types::{BedrockContentBlock, BedrockConverseResponse}; use crate::shared::EndpointConfig; @@ -200,10 +203,7 @@ impl LanguageModel for BedrockModel { .stop_reason .as_deref() .map(map_finish_reason) - .unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }); + .unwrap_or(FINISH_OTHER); if is_json_response_from_tool && finish_reason.unified == FinishReasonUnified::ToolCalls { finish_reason.unified = FinishReasonUnified::Stop; @@ -271,7 +271,6 @@ impl LanguageModel for BedrockModel { let response_headers = resp.response_headers; let body_stream = resp.value; - let include_raw_chunks = options.include_raw_chunks.unwrap_or(false); let request_id = response_headers .get("x-amzn-requestid") .cloned() @@ -282,305 +281,405 @@ impl LanguageModel for BedrockModel { .get("date") .and_then(|date| aimux_core::Timestamp::from_http_date(date)); - let model_id = self.model_id.clone(); + let stream = pipe_through( + super::event_stream::decode_stream(body_stream), + AmazonBedrockChatStream { + ctx: StreamContext::new( + options, + warnings, + error_url, + &error_body, + &response_headers, + ), + request_id, + response_timestamp, + model_id: self.model_id.clone(), + uses_json_tool, + extractor: uses_json_instruction.then(JsonObjectTextExtractor::default), + ..AmazonBedrockChatStream::default() + }, + ); - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), + response: Some(aimux_core::shared::StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: request_id, - timestamp: response_timestamp, - model_id: Some(model_id.clone()), - })); +/// The `TransformStream` of `BedrockChatLanguageModel.doStream`. +#[derive(Default)] +struct AmazonBedrockChatStream { + ctx: StreamContext, + request_id: Option, + response_timestamp: Option, + model_id: String, + uses_json_tool: bool, + extractor: Option, + is_json_response_from_tool: bool, + text_blocks: HashSet, + reasoning_blocks: HashMap>, + block_counter: usize, + /// Tool call state: block_index → (id, name, accumulated_json) + tool_blocks: HashMap, + final_usage: Usage, + final_finish_reason: Option, + /// Provider metadata accumulated for the Finish chunk from `metadata` + /// and `messageStop` events (findings #26, #27). Mirrors the TS + /// `providerMetadata` / `stopSequence` accumulation in + /// `amazon-bedrock-chat-language-model.ts` (`doStream`). + finish_meta: serde_json::Map, + stop_sequence: Option, +} - let mut messages = super::event_stream::decode_stream(body_stream); - - let mut extractor = uses_json_instruction.then(JsonObjectTextExtractor::default); - let mut is_json_response_from_tool = false; - let mut text_blocks = HashSet::new(); - let mut reasoning_blocks: HashMap> = HashMap::new(); - let mut block_counter = 0usize; - // Tool call state: block_index → (id, name, accumulated_json) - let mut tool_blocks: HashMap = HashMap::new(); - let mut final_usage: Usage = Usage::default(); - let mut final_finish_reason: Option = None; - // Provider metadata accumulated for the Finish chunk from `metadata` - // and `messageStop` events (findings #26, #27). Mirrors the TS - // `providerMetadata` / `stopSequence` accumulation in - // `amazon-bedrock-chat-language-model.ts` (`doStream`). - let mut finish_meta: serde_json::Map = - serde_json::Map::new(); - let mut stop_sequence: Option = None; - - while let Some(message) = messages.next().await { - let msg = match message { - Ok(msg) => msg, - Err(error) => { yield Err(error); return; } +impl AmazonBedrockChatStream { + fn content_block_start(&mut self, payload: &serde_json::Value, controller: &mut Controller) { + let idx = payload + .get("contentBlockIndex") + .and_then(serde_json::Value::as_u64) + .unwrap_or(self.block_counter as u64) as usize; + + // Check if this is a tool use block. + if let Some(start) = payload.get("start") { + if let Some(tool_use) = start.get("toolUse") { + let name = tool_use + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let id = tool_use + .get("toolUseId") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let id = normalize_tool_call_id( + &if id.is_empty() { + aimux_provider_utils::generate_id() + } else { + id + }, + self.model_id.contains("mistral."), + ); + if !(self.uses_json_tool && name == "json") { + controller.tool_input_start(id.clone(), name.clone()); + } + self.tool_blocks.insert(idx, (id, name, String::new())); + } else { + // Text block. + self.block_counter = idx + 1; + let id = idx.to_string(); + self.text_blocks.insert(idx); + controller.text_start(id); + } + } else { + // Default: text block. + let id = idx.to_string(); + self.text_blocks.insert(idx); + controller.text_start(id); + } + } + + fn content_block_delta(&mut self, payload: &serde_json::Value, controller: &mut Controller) { + let idx = payload + .get("contentBlockIndex") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) as usize; + + let Some(delta) = payload.get("delta") else { + return; + }; + // Text delta + if let Some(text) = delta.get("text").and_then(|v| v.as_str()) + && !text.is_empty() + { + if self.text_blocks.insert(idx) { + let id = idx.to_string(); + self.text_blocks.insert(idx); + controller.text_start(id); + } + let delta = self + .extractor + .as_mut() + .map(|extractor| extractor.process(text)) + .unwrap_or_else(|| text.to_string()); + if !delta.is_empty() { + controller.text_delta(idx.to_string(), delta); + } + } + // Tool use input delta + if let Some(partial) = delta.get("toolUse").and_then(|t| t.get("input")) + && let Some(partial_str) = partial.as_str() + && let Some((id, name, acc)) = self.tool_blocks.get_mut(&idx) + && !partial_str.is_empty() + { + acc.push_str(partial_str); + if !(self.uses_json_tool && name == "json") { + let id = id.clone(); + controller.tool_input_delta(id, partial_str.to_string()); + } + } + if let Some(rc) = delta.get("reasoningContent") { + if !["text", "signature", "data", "redactedContent"] + .iter() + .any(|key| { + rc.get(key) + .and_then(|v| v.as_str()) + .is_some_and(|s| !s.is_empty()) + }) + { + return; + } + let id = idx.to_string(); + if let std::collections::hash_map::Entry::Vacant(entry) = + self.reasoning_blocks.entry(idx) + { + entry.insert(None); + controller.reasoning_start(id.clone()); + } + if let Some(text) = rc.get("text").and_then(|v| v.as_str()) { + controller.reasoning_delta(id.clone(), text.to_string()); + } else if let Some(sig) = rc + .get("signature") + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + { + controller.enqueue(StreamPart::ReasoningDelta { + id: id.clone(), + delta: String::new(), + provider_metadata: Some(options::metadata(json!({"signature": sig}))), + }); + } else if let Some(data) = rc + .get("data") + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + { + controller.enqueue(StreamPart::ReasoningDelta { + id: id.clone(), + delta: String::new(), + provider_metadata: Some(options::metadata(json!({"redactedData": data}))), + }); + } else if let Some(data) = rc.get("redactedContent").and_then(|v| v.as_str()) { + self.reasoning_blocks + .entry(idx) + .or_default() + .get_or_insert_with(String::new) + .push_str(data); + } + } + } + + fn content_block_stop(&mut self, payload: &serde_json::Value, controller: &mut Controller) { + let idx = payload + .get("contentBlockIndex") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) as usize; + + if let Some((id, name, acc)) = self.tool_blocks.remove(&idx) { + if self.uses_json_tool && name == "json" { + self.is_json_response_from_tool = true; + let id = idx.to_string(); + controller.text_start(id.clone()); + controller.text_delta(id.clone(), acc); + controller.text_end(id); + } else { + controller.tool_input_end(id.clone()); + // Empty input normalizes to "{}" per the upstream + // provider. + let input = if acc.is_empty() { + "{}".to_string() + } else { + acc }; - let payload: serde_json::Value = match serde_json::from_str(&msg.data) { - Ok(v) => v, + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: normalize_tool_call_id(&id, self.model_id.contains("mistral.")), + tool_name: name, + input, + provider_executed: None, + dynamic: None, + provider_metadata: None, + })); + } + } else if let Some(redacted) = self.reasoning_blocks.remove(&idx) { + controller.enqueue(StreamPart::ReasoningEnd { + id: idx.to_string(), + provider_metadata: reasoning_redacted_meta(redacted), + }); + } else if self.text_blocks.remove(&idx) { + controller.text_end(idx.to_string()); + } + } + + fn message_stop(&mut self, payload: &serde_json::Value) { + if let Some(reason) = payload.get("stopReason").and_then(|v| v.as_str()) { + let mut finish_reason = map_finish_reason(reason); + if self.is_json_response_from_tool + && finish_reason.unified == FinishReasonUnified::ToolCalls + { + finish_reason.unified = FinishReasonUnified::Stop; + } + self.final_finish_reason = Some(finish_reason); + } + // #26: surface which stop sequence sentinel fired + // (additionalModelResponseFields.delta.stop_sequence), if any. + if let Some(seq) = payload + .get("additionalModelResponseFields") + .and_then(|f| f.get("delta")) + .and_then(|d| d.get("stop_sequence")) + .and_then(|v| v.as_str()) + { + self.stop_sequence = Some(seq.to_string()); + } + } + + fn metadata(&mut self, payload: &serde_json::Value, controller: &mut Controller) { + if let Some(usage) = payload.get("usage") { + let bedrock_usage: super::types::BedrockUsage = + match serde_json::from_value(usage.clone()) { + Ok(usage) => usage, Err(error) => { - final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); - yield Ok(StreamPart::Error { error: AiMuxError::from(error) }); - continue; + self.final_finish_reason = Some(FINISH_ERROR); + controller.enqueue(StreamPart::Error { + error: AiMuxError::InvalidResponseData(format!( + "{error}; usage: {usage}" + )), + }); + return; } }; - if include_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: json!({ msg.event_type.clone(): payload.clone() }) }); - } - if msg.message_type != "event" || msg.event_type.ends_with("Exception") { - final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); - let (status_code, is_retryable) = match msg.event_type.as_str() { - "internalServerException" | "InternalServerException" => (Some(500), true), - "modelStreamErrorException" | "ModelStreamErrorException" => (Some(424), true), - "serviceUnavailableException" | "ServiceUnavailableException" => (Some(503), true), - "throttlingException" | "ThrottlingException" => (Some(429), true), - "validationException" | "ValidationException" => (Some(400), false), - _ => (None, false), - }; - let message = payload.get("message").and_then(|v| v.as_str()).map(str::to_string) - .unwrap_or_else(|| format!("Amazon Bedrock stream failed with {}", msg.event_type)); - yield Ok(StreamPart::Error { error: AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { - status_code, is_retryable, provider_code: Some(msg.event_type.clone()), - data: Some(json!({msg.event_type.clone(): payload.clone()})), - ..aimux_core::ApiCallError::new(message, error_url.clone(), error_body.clone()) - })) }); - continue; - } - match msg.event_type.as_str() { - "messageStart" => { - // Nothing to emit; response metadata already sent. - } - "contentBlockStart" => { - let idx = payload - .get("contentBlockIndex") - .and_then(serde_json::Value::as_u64) - .unwrap_or(block_counter as u64) as usize; - - // Check if this is a tool use block. - if let Some(start) = payload.get("start") { - if let Some(tool_use) = start.get("toolUse") { - let name = tool_use - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let id = tool_use - .get("toolUseId") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let id = normalize_tool_call_id(&if id.is_empty() { - aimux_provider_utils::generate_id() - } else { - id - }, model_id.contains("mistral.")); - if !(uses_json_tool && name == "json") { - yield Ok(StreamPart::ToolInputStart { - id: id.clone(), - tool_name: name.clone(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - } - tool_blocks.insert(idx, (id, name, String::new())); - } else { - // Text block. - block_counter = idx + 1; - let id = idx.to_string(); - text_blocks.insert(idx); - yield Ok(StreamPart::TextStart { id, provider_metadata: None}); - } - } else { - // Default: text block. - let id = idx.to_string(); - text_blocks.insert(idx); - yield Ok(StreamPart::TextStart { id, provider_metadata: None}); - } - } - "contentBlockDelta" => { - let idx = payload - .get("contentBlockIndex") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0) as usize; - - if let Some(delta) = payload.get("delta") { - // Text delta - if let Some(text) = delta.get("text").and_then(|v| v.as_str()) - && !text.is_empty() { - if text_blocks.insert(idx) { - let id = idx.to_string(); - text_blocks.insert(idx); - yield Ok(StreamPart::TextStart { id, provider_metadata: None}); - } - let delta = extractor.as_mut().map(|extractor| extractor.process(text)).unwrap_or_else(|| text.to_string()); - if !delta.is_empty() { - yield Ok(StreamPart::TextDelta { id: idx.to_string(), delta, provider_metadata: None }); - } - } - // Tool use input delta - if let Some(partial) = - delta.get("toolUse").and_then(|t| t.get("input")) - && let Some(partial_str) = partial.as_str() - && let Some((id, name, acc)) = tool_blocks.get_mut(&idx) - && !partial_str.is_empty() { - acc.push_str(partial_str); - if !(uses_json_tool && name == "json") { - let id = id.clone(); - yield Ok(StreamPart::ToolInputDelta { - id, - delta: partial_str.to_string(), - provider_metadata: None, - }); - } - } - if let Some(rc) = delta.get("reasoningContent") { - if !["text", "signature", "data", "redactedContent"].iter().any(|key| rc.get(key).and_then(|v| v.as_str()).is_some_and(|s| !s.is_empty())) { continue; } - let id = idx.to_string(); - if let std::collections::hash_map::Entry::Vacant(entry) = reasoning_blocks.entry(idx) { - entry.insert(None); - yield Ok(StreamPart::ReasoningStart { id: id.clone(), provider_metadata: None }); - } - if let Some(text) = rc.get("text").and_then(|v| v.as_str()) { - yield Ok(StreamPart::ReasoningDelta { id: id.clone(), delta: text.to_string(), provider_metadata: None }); - } else if let Some(sig) = rc.get("signature").and_then(|v| v.as_str()).filter(|s| !s.is_empty()) { - yield Ok(StreamPart::ReasoningDelta { id: id.clone(), delta: String::new(), provider_metadata: Some(options::metadata(json!({"signature": sig}))) }); - } else if let Some(data) = rc.get("data").and_then(|v| v.as_str()).filter(|s| !s.is_empty()) { - yield Ok(StreamPart::ReasoningDelta { id: id.clone(), delta: String::new(), provider_metadata: Some(options::metadata(json!({"redactedData": data}))) }); - } else if let Some(data) = rc.get("redactedContent").and_then(|v| v.as_str()) { - reasoning_blocks.entry(idx).or_default().get_or_insert_with(String::new).push_str(data); - } - } - } - } - "contentBlockStop" => { - let idx = payload - .get("contentBlockIndex") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0) as usize; - - if let Some((id, name, acc)) = tool_blocks.remove(&idx) { - if uses_json_tool && name == "json" { - is_json_response_from_tool = true; - let id = idx.to_string(); - yield Ok(StreamPart::TextStart { id: id.clone(), provider_metadata: None }); - yield Ok(StreamPart::TextDelta { id: id.clone(), delta: acc, provider_metadata: None }); - yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); - } else { - yield Ok(StreamPart::ToolInputEnd { id: id.clone(), provider_metadata: None}); - // Empty input normalizes to "{}" per the upstream - // provider. - let input = if acc.is_empty() { - "{}".to_string() - } else { - acc - }; - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: normalize_tool_call_id(&id, model_id.contains("mistral.")), - tool_name: name, - input, - provider_executed: None, - dynamic: None, - provider_metadata: None, - })); - } - } else if let Some(redacted) = reasoning_blocks.remove(&idx) { - yield Ok(StreamPart::ReasoningEnd { id: idx.to_string(), provider_metadata: reasoning_redacted_meta(redacted) }); - } else if text_blocks.remove(&idx) { - yield Ok(StreamPart::TextEnd { id: idx.to_string(), provider_metadata: None }); - } - } - "messageStop" => { - if let Some(reason) = - payload.get("stopReason").and_then(|v| v.as_str()) - { - let mut finish_reason = map_finish_reason(reason); - if is_json_response_from_tool && finish_reason.unified == FinishReasonUnified::ToolCalls { finish_reason.unified = FinishReasonUnified::Stop; } - final_finish_reason = Some(finish_reason); - } - // #26: surface which stop sequence sentinel fired - // (additionalModelResponseFields.delta.stop_sequence), if any. - if let Some(seq) = payload - .get("additionalModelResponseFields") - .and_then(|f| f.get("delta")) - .and_then(|d| d.get("stop_sequence")) - .and_then(|v| v.as_str()) - { - stop_sequence = Some(seq.to_string()); - } - } - "metadata" => { - if let Some(usage) = payload.get("usage") { - let bedrock_usage: super::types::BedrockUsage = - match serde_json::from_value(usage.clone()) { - Ok(usage) => usage, - Err(error) => { - final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); - yield Ok(StreamPart::Error { error: AiMuxError::InvalidResponseData(format!("{error}; usage: {usage}")) }); - continue; - } - }; - final_usage = convert_usage(Some(&bedrock_usage)); - final_usage.raw = usage.as_object().cloned(); - if let Some(cache_usage) = cache_usage_metadata(Some(&bedrock_usage)) { finish_meta.insert("usage".into(), cache_usage); } - } - // #27: surface guardrails trace, performanceConfig, and - // serviceTier into the Finish chunk's provider_metadata - // (mirrors TS `doStream` metadata handling). - if let Some(trace) = payload.get("trace") { - finish_meta.insert("trace".to_string(), trace.clone()); - } - if let Some(pc) = payload.get("performanceConfig") { - finish_meta - .insert("performanceConfig".to_string(), pc.clone()); - } - if let Some(st) = payload.get("serviceTier") { - finish_meta.insert("serviceTier".to_string(), st.clone()); - } - } - _ => {} - } + self.final_usage = convert_usage(Some(&bedrock_usage)); + self.final_usage.raw = usage.as_object().cloned(); + if let Some(cache_usage) = cache_usage_metadata(Some(&bedrock_usage)) { + self.finish_meta.insert("usage".into(), cache_usage); } + } + // #27: surface guardrails trace, performanceConfig, and + // serviceTier into the Finish chunk's provider_metadata + // (mirrors TS `doStream` metadata handling). + if let Some(trace) = payload.get("trace") { + self.finish_meta.insert("trace".to_string(), trace.clone()); + } + if let Some(pc) = payload.get("performanceConfig") { + self.finish_meta + .insert("performanceConfig".to_string(), pc.clone()); + } + if let Some(st) = payload.get("serviceTier") { + self.finish_meta + .insert("serviceTier".to_string(), st.clone()); + } + } +} + +impl Transformer for AmazonBedrockChatStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut Controller) { + self.ctx.start(controller); + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: self.request_id.take(), + timestamp: self.response_timestamp.take(), + model_id: Some(self.model_id.clone()), + })); + } - // #26: merge the stop sentinel into the metadata payload, then build - // the Finish provider_metadata under `amazonBedrock` (mirrors the - // TS `doStream` flush handler). - if is_json_response_from_tool { - finish_meta.insert("isJsonResponseFromTool".into(), json!(true)); - finish_meta.insert("stopSequence".into(), serde_json::Value::Null); + fn transform( + &mut self, + message: Result, + controller: &mut Controller, + ) { + let msg = match message { + Ok(msg) => msg, + Err(error) => { + controller.error(error); + return; } - if let Some(seq) = stop_sequence { - finish_meta.insert( - "stopSequence".to_string(), - serde_json::Value::String(seq), - ); + }; + let payload: serde_json::Value = match serde_json::from_str(&msg.data) { + Ok(v) => v, + Err(error) => { + self.final_finish_reason = Some(FINISH_ERROR); + controller.enqueue(StreamPart::Error { + error: AiMuxError::from(error), + }); + return; } - let provider_metadata = if finish_meta.is_empty() { - None - } else { - let payload = serde_json::Value::Object(finish_meta); - Some(options::metadata(payload)) + }; + self.ctx.raw( + &json!({ msg.event_type.clone(): payload.clone() }), + controller, + ); + if msg.message_type != "event" || msg.event_type.ends_with("Exception") { + self.final_finish_reason = Some(FINISH_ERROR); + let (status_code, is_retryable) = match msg.event_type.as_str() { + "internalServerException" | "InternalServerException" => (Some(500), true), + "modelStreamErrorException" | "ModelStreamErrorException" => (Some(424), true), + "serviceUnavailableException" | "ServiceUnavailableException" => (Some(503), true), + "throttlingException" | "ThrottlingException" => (Some(429), true), + "validationException" | "ValidationException" => (Some(400), false), + _ => (None, false), }; - - yield Ok(StreamPart::Finish { - finish_reason: final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }), - usage: final_usage, - provider_metadata, + let message = payload + .get("message") + .and_then(|v| v.as_str()) + .map(str::to_string) + .unwrap_or_else(|| format!("Amazon Bedrock stream failed with {}", msg.event_type)); + controller.enqueue(StreamPart::Error { + error: AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { + status_code, + is_retryable, + provider_code: Some(msg.event_type.clone()), + data: Some(json!({msg.event_type.clone(): payload.clone()})), + ..aimux_core::ApiCallError::new( + message, + self.ctx.url.clone(), + self.ctx.request_body.clone(), + ) + })), }); + return; + } + match msg.event_type.as_str() { + "messageStart" => { + // Nothing to emit; response metadata already sent. + } + "contentBlockStart" => self.content_block_start(&payload, controller), + "contentBlockDelta" => self.content_block_delta(&payload, controller), + "contentBlockStop" => self.content_block_stop(&payload, controller), + "messageStop" => self.message_stop(&payload), + "metadata" => self.metadata(&payload, controller), + _ => {} + } + } + + fn flush(mut self, controller: &mut Controller) { + // #26: merge the stop sentinel into the metadata payload, then build + // the Finish provider_metadata under `amazonBedrock` (mirrors the + // TS `doStream` flush handler). + if self.is_json_response_from_tool { + self.finish_meta + .insert("isJsonResponseFromTool".into(), json!(true)); + self.finish_meta + .insert("stopSequence".into(), serde_json::Value::Null); + } + if let Some(seq) = self.stop_sequence { + self.finish_meta + .insert("stopSequence".to_string(), serde_json::Value::String(seq)); + } + let provider_metadata = if self.finish_meta.is_empty() { + None + } else { + let payload = serde_json::Value::Object(self.finish_meta); + Some(options::metadata(payload)) }; - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - response: Some(aimux_core::shared::StreamResponseInfo { - headers: Some(response_headers), - }), - }) + controller.enqueue(StreamPart::Finish { + finish_reason: self.final_finish_reason.unwrap_or(FINISH_OTHER), + usage: self.final_usage, + provider_metadata, + }); } } diff --git a/aimux-providers/src/cohere/model.rs b/aimux-providers/src/cohere/model.rs index e3d5d19a3..583e00aa8 100644 --- a/aimux-providers/src/cohere/model.rs +++ b/aimux-providers/src/cohere/model.rs @@ -6,7 +6,6 @@ use aimux_core::tool::RawToolCall; use async_trait::async_trait; -use futures::StreamExt; use serde_json::Value; use aimux_core::error::AiMuxError; @@ -14,9 +13,12 @@ use aimux_core::language_model::{LanguageModel, SupportedUrls}; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, Source, StreamResult}; use aimux_core::stream_part::StreamPart; -use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; +use aimux_core::types::{FinishReason, ResponseMetadata, Usage}; + +use aimux_provider_utils::{Transformer, pipe_through}; use crate::shared::EndpointConfig; +use crate::stream::{Controller, Emit, FINISH_ERROR, FINISH_OTHER, StreamContext}; use super::convert::{build_request_body, parse_finish_reason}; use super::types::{ChatResponse, StreamEvent, UsageResponse}; @@ -253,100 +255,17 @@ impl LanguageModel for CohereModel { let response_headers = resp.response_headers; let sse_stream = resp.value; let stream_warnings = body_result.warnings; - let include_raw_chunks = options.include_raw_chunks.unwrap_or(false); - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings: stream_warnings }); - let mut final_usage = Usage::default(); - let mut finish_reason = FinishReason { unified: FinishReasonUnified::Other, raw: None }; - let mut pending_tool_call: Option = None; - let mut is_reasoning = false; - let mut sse_iter = sse_stream; - while let Some(event) = sse_iter.next().await { - let raw = match event { - Ok(raw) => raw, - Err(error) => { - if !error.is_recoverable_stream_error() { yield Err(error); return; } - if include_raw_chunks { yield Ok(StreamPart::Raw { raw_value: Value::Null }); } - finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; - yield Ok(StreamPart::Error { error }); - continue; - } - }; - if include_raw_chunks { yield Ok(StreamPart::Raw { raw_value: raw.clone() }); } - let parsed = match StreamEvent::parse(raw) { - Ok(parsed) => parsed, - Err(error) => { - finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; - yield Ok(StreamPart::Error { error }); - continue; - } - }; - let id = parsed.index.map(|index| index.to_string()).unwrap_or_default(); - match parsed.event_type.as_str() { - "message-start" => yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { id: parsed.id, timestamp: None, model_id: None })), - "content-start" => { - let content = parsed.delta.as_ref().and_then(|delta| delta.message.as_ref()).and_then(|message| message.content.as_ref()).unwrap(); - if content["type"] == "thinking" { - is_reasoning = true; - yield Ok(StreamPart::ReasoningStart { id, provider_metadata: None }); - } else { - yield Ok(StreamPart::TextStart { id, provider_metadata: None }); - } - } - "content-delta" => { - let content = parsed.delta.as_ref().and_then(|delta| delta.message.as_ref()).and_then(|message| message.content.as_ref()).unwrap(); - if let Some(text) = content.get("text").and_then(Value::as_str) { - yield Ok(StreamPart::TextDelta { id, delta: text.into(), provider_metadata: None }); - } else { - yield Ok(StreamPart::ReasoningDelta { id, delta: content["thinking"].as_str().unwrap().into(), provider_metadata: None }); - } - } - "content-end" => { - if is_reasoning { - yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); - is_reasoning = false; - } else { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); - } - } - "tool-call-start" => { - let tool = parsed.delta.as_ref().and_then(|delta| delta.message.as_ref()).and_then(|message| message.tool_calls.as_ref()).unwrap(); - let id = tool["id"].as_str().unwrap().to_string(); - let name = tool["function"]["name"].as_str().unwrap().to_string(); - let arguments = tool["function"]["arguments"].as_str().unwrap().to_string(); - pending_tool_call = Some(PendingToolCall { id: id.clone(), name: name.clone(), arguments: arguments.clone() }); - yield Ok(StreamPart::ToolInputStart { id: id.clone(), tool_name: name, provider_executed: None, dynamic: None, title: None, provider_metadata: None }); - if !arguments.is_empty() { yield Ok(StreamPart::ToolInputDelta { id, delta: arguments, provider_metadata: None }); } - } - "tool-call-delta" => { - if let Some(tool) = &mut pending_tool_call { - let delta = parsed.delta.as_ref().and_then(|delta| delta.message.as_ref()).and_then(|message| message.tool_calls.as_ref()).unwrap()["function"]["arguments"].as_str().unwrap().to_string(); - tool.arguments.push_str(&delta); - yield Ok(StreamPart::ToolInputDelta { id: tool.id.clone(), delta, provider_metadata: None }); - } - } - "tool-call-end" => { - if let Some(tool) = pending_tool_call.take() { - yield Ok(StreamPart::ToolInputEnd { id: tool.id.clone(), provider_metadata: None }); - let text = tool.arguments.trim(); - let input = match serde_json::from_str::(if text.is_empty() { "{}" } else { text }) { - Ok(value) if !contains_prototype_key(&value) => value.to_string(), - Ok(_) => { yield Err(AiMuxError::InvalidResponseData("Object contains forbidden prototype property".into())); return; } - Err(error) => { yield Err(AiMuxError::JsonParse(error.to_string())); return; } - }; - yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: tool.id, tool_name: tool.name, input, provider_executed: None, dynamic: None, provider_metadata: None })); - } - } - "message-end" => { - let delta = parsed.delta.unwrap(); - finish_reason = parse_finish_reason(delta.finish_reason.as_ref().unwrap()); - final_usage = convert_usage(delta.usage.as_ref().unwrap()); - } - _ => {} - } - } - yield Ok(StreamPart::Finish { finish_reason, usage: final_usage, provider_metadata: None }); - }; + let stream = pipe_through( + sse_stream, + CohereChatStream { + ctx: StreamContext { + warnings: stream_warnings, + include_raw_chunks: options.include_raw_chunks == Some(true), + ..StreamContext::default() + }, + ..CohereChatStream::default() + }, + ); Ok(StreamResult { stream: Box::pin(stream), @@ -358,6 +277,178 @@ impl LanguageModel for CohereModel { } } +/// The `TransformStream` of `CohereChatLanguageModel.doStream`. +#[derive(Default)] +struct CohereChatStream { + ctx: StreamContext, + final_usage: Usage, + /// `other` until a chunk says otherwise. + finish_reason: Option, + pending_tool_call: Option, + is_reasoning: bool, +} + +impl Transformer for CohereChatStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut Controller) { + self.ctx.start(controller); + } + + fn transform(&mut self, event: Result, controller: &mut Controller) { + let raw = match event { + Ok(raw) => raw, + Err(error) => { + if !error.is_recoverable_stream_error() { + controller.error(error); + return; + } + self.ctx.raw(&Value::Null, controller); + self.finish_reason = Some(FINISH_ERROR); + controller.enqueue(StreamPart::Error { error }); + return; + } + }; + self.ctx.raw(&raw, controller); + let parsed = match StreamEvent::parse(raw) { + Ok(parsed) => parsed, + Err(error) => { + self.finish_reason = Some(FINISH_ERROR); + controller.enqueue(StreamPart::Error { error }); + return; + } + }; + let id = parsed + .index + .map(|index| index.to_string()) + .unwrap_or_default(); + match parsed.event_type.as_str() { + "message-start" => controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: parsed.id, + timestamp: None, + model_id: None, + })), + "content-start" => { + let content = parsed + .delta + .as_ref() + .and_then(|delta| delta.message.as_ref()) + .and_then(|message| message.content.as_ref()) + .unwrap(); + if content["type"] == "thinking" { + self.is_reasoning = true; + controller.reasoning_start(id); + } else { + controller.text_start(id); + } + } + "content-delta" => { + let content = parsed + .delta + .as_ref() + .and_then(|delta| delta.message.as_ref()) + .and_then(|message| message.content.as_ref()) + .unwrap(); + if let Some(text) = content.get("text").and_then(Value::as_str) { + controller.text_delta(id, text); + } else { + controller.reasoning_delta(id, content["thinking"].as_str().unwrap()); + } + } + "content-end" => { + if self.is_reasoning { + controller.reasoning_end(id); + self.is_reasoning = false; + } else { + controller.text_end(id); + } + } + "tool-call-start" => { + let tool = parsed + .delta + .as_ref() + .and_then(|delta| delta.message.as_ref()) + .and_then(|message| message.tool_calls.as_ref()) + .unwrap(); + let id = tool["id"].as_str().unwrap().to_string(); + let name = tool["function"]["name"].as_str().unwrap().to_string(); + let arguments = tool["function"]["arguments"].as_str().unwrap().to_string(); + self.pending_tool_call = Some(PendingToolCall { + id: id.clone(), + name: name.clone(), + arguments: arguments.clone(), + }); + controller.tool_input_start(id.clone(), name); + if !arguments.is_empty() { + controller.tool_input_delta(id, arguments); + } + } + "tool-call-delta" => { + if let Some(tool) = &mut self.pending_tool_call { + let delta = parsed + .delta + .as_ref() + .and_then(|delta| delta.message.as_ref()) + .and_then(|message| message.tool_calls.as_ref()) + .unwrap()["function"]["arguments"] + .as_str() + .unwrap() + .to_string(); + tool.arguments.push_str(&delta); + controller.tool_input_delta(tool.id.clone(), delta); + } + } + "tool-call-end" => { + if let Some(tool) = self.pending_tool_call.take() { + controller.tool_input_end(tool.id.clone()); + let text = tool.arguments.trim(); + let input = match serde_json::from_str::(if text.is_empty() { + "{}" + } else { + text + }) { + Ok(value) if !contains_prototype_key(&value) => value.to_string(), + Ok(_) => { + controller.error(AiMuxError::InvalidResponseData( + "Object contains forbidden prototype property".into(), + )); + return; + } + Err(error) => { + controller.error(AiMuxError::JsonParse(error.to_string())); + return; + } + }; + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: tool.id, + tool_name: tool.name, + input, + provider_executed: None, + dynamic: None, + provider_metadata: None, + })); + } + } + "message-end" => { + let delta = parsed.delta.unwrap(); + self.finish_reason = + Some(parse_finish_reason(delta.finish_reason.as_ref().unwrap())); + self.final_usage = convert_usage(delta.usage.as_ref().unwrap()); + } + _ => {} + } + } + + fn flush(self, controller: &mut Controller) { + controller.enqueue(StreamPart::Finish { + finish_reason: self.finish_reason.unwrap_or(FINISH_OTHER), + usage: self.final_usage, + provider_metadata: None, + }); + } +} + fn contains_prototype_key(value: &Value) -> bool { match value { Value::Object(object) => object.iter().any(|(key, value)| { diff --git a/aimux-providers/src/deepseek/model.rs b/aimux-providers/src/deepseek/model.rs index 69c91455d..6675d1b98 100644 --- a/aimux-providers/src/deepseek/model.rs +++ b/aimux-providers/src/deepseek/model.rs @@ -4,7 +4,6 @@ use aimux_core::tool::RawToolCall; use std::collections::{BTreeMap, HashMap}; use async_trait::async_trait; -use futures::StreamExt; use serde_json::{Map, Value, json}; use aimux_core::error::AiMuxError; @@ -13,10 +12,8 @@ use aimux_core::options::{CallOptions, ResponseFormat}; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, StreamResult}; use aimux_core::shared::{RequestInfo, ResponseInfo, StreamResponseInfo, provider_namespace}; use aimux_core::stream_part::StreamPart; -use aimux_core::types::{ - FinishReason, FinishReasonUnified, ReasoningEffort, ResponseMetadata, Warning, -}; -use aimux_provider_utils::{StreamingToolCallDelta, StreamingToolCallTracker, generate_id}; +use aimux_core::types::{FinishReason, ReasoningEffort, ResponseMetadata, Warning}; +use aimux_provider_utils::{Transformer, generate_id, pipe_through}; use super::convert::convert_to_deepseek_chat_messages; use super::finish_reason::map_deepseek_finish_reason; @@ -26,6 +23,9 @@ use super::prepare_tools::prepare_tools; use super::types::{ChatChunk, ChatResponse}; use super::usage::convert_deepseek_usage; use crate::shared::EndpointConfig; +use crate::stream::{ + ChatToolCalls, Controller, Emit, FINISH_ERROR, FINISH_OTHER, StreamContext, ToolCallDelta, +}; /// What a DeepSeek chat model needs to talk to the API (`DeepSeekChatConfig`). pub(crate) struct DeepSeekChatConfig { @@ -496,279 +496,268 @@ impl LanguageModel for DeepSeekChatLanguageModel { ) .await?; let response_headers = resp.response_headers; - let sse_stream = resp.value; - - let provider_options_name = self.provider_options_name().to_string(); - let emit_raw_chunks = options.include_raw_chunks == Some(true); - let stream_error_body = body.clone(); - let stream_response_headers = response_headers.clone(); - - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); - - let mut tool_calls = StreamingToolCallTracker::new().with_generate_id(generate_id); - let mut tool_parts = Vec::new(); - let mut finish_reason = FinishReason { unified: FinishReasonUnified::Other, raw: None }; - let mut usage: Option = None; - let mut system_fingerprint: Option = None; - let mut is_first_chunk = true; - let mut is_active_reasoning = false; - let mut is_active_text = false; - let mut response_object: Option = None; - let mut choice_index: Option = None; - let mut message_role: Option = None; - let mut tool_call_types: BTreeMap = BTreeMap::new(); - let mut content_logprobs: Vec = Vec::new(); - let mut reasoning_logprobs: Vec = Vec::new(); - - let mut events = sse_stream; - while let Some(event) = events.next().await { - let parsed = match event { - Ok(parsed) => parsed, - Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Err(error); - if !recoverable { - return; - } - continue; - } - }; - // Emit the raw chunk if requested (before anything else). - if emit_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); - } + let stream = pipe_through( + resp.value, + DeepSeekChatStream { + ctx: StreamContext::new(options, warnings, endpoint, &body, &response_headers), + provider_options_name: self.provider_options_name().to_string(), + ..DeepSeekChatStream::default() + }, + ); - // handle error chunks: - if parsed.get("error").is_some() { - finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; - yield Ok(StreamPart::Error { - error: deepseek_stream_error( - &parsed, - &endpoint, - stream_error_body.clone(), - stream_response_headers.clone(), - ), - }); - continue; - } + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(RequestInfo { body: Some(body) }), + response: Some(StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} - // handle failed chunk validation: - let chunk: ChatChunk = match serde_json::from_value(parsed) { - Ok(chunk) => chunk, - Err(error) => { - finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; - yield Ok(StreamPart::Error { error: error.into() }); - continue; - } - }; +/// The `TransformStream` of `DeepSeekChatLanguageModel.doStream`. +#[derive(Default)] +struct DeepSeekChatStream { + ctx: StreamContext, + provider_options_name: String, + tool_calls: ChatToolCalls, + /// `other` until a chunk says otherwise. + finish_reason: Option, + usage: Option, + system_fingerprint: Option, + metadata_emitted: bool, + is_active_reasoning: bool, + is_active_text: bool, + response_object: Option, + choice_index: Option, + message_role: Option, + tool_call_types: BTreeMap, + content_logprobs: Vec, + reasoning_logprobs: Vec, +} - if is_first_chunk { - is_first_chunk = false; - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: chunk.id.clone(), - timestamp: timestamp(chunk.created), - model_id: chunk.model.clone(), - })); - } +impl DeepSeekChatStream { + fn end_active_reasoning(&mut self, controller: &mut Controller) { + if self.is_active_reasoning { + controller.reasoning_end("reasoning-0"); + self.is_active_reasoning = false; + } + } +} - if let Some(chunk_usage) = chunk.usage.filter(|usage| !usage.is_null()) { - usage = Some(chunk_usage); - } +impl Transformer for DeepSeekChatStream { + type Input = Result; + type Output = StreamPart; - if chunk.object.is_some() { - response_object = chunk.object; - } + fn start(&mut self, controller: &mut Controller) { + self.ctx.start(controller); + } - // The fingerprint is repeated on stream chunks; keep the - // latest non-null value in case it changes during the response. - if chunk.system_fingerprint.is_some() { - system_fingerprint = chunk.system_fingerprint; + fn transform(&mut self, event: Result, controller: &mut Controller) { + let parsed = match event { + Ok(parsed) => parsed, + Err(error) => { + if error.is_recoverable_stream_error() { + controller.enqueue_error(error); + } else { + controller.error(error); } + return; + } + }; - let Some(choice) = chunk.choices.into_iter().next() else { - continue; - }; + // Emit the raw chunk if requested (before anything else). + self.ctx.raw(&parsed, controller); + + // handle error chunks: + if parsed.get("error").is_some() { + self.finish_reason = Some(FINISH_ERROR); + controller.enqueue(StreamPart::Error { + error: deepseek_stream_error( + &parsed, + &self.ctx.url, + self.ctx.request_body.clone(), + self.ctx.response_headers.clone(), + ), + }); + return; + } - if choice.index.is_some() { - choice_index = choice.index; - } + // handle failed chunk validation: + let chunk: ChatChunk = match serde_json::from_value(parsed) { + Ok(chunk) => chunk, + Err(error) => { + self.finish_reason = Some(FINISH_ERROR); + controller.enqueue(StreamPart::Error { + error: error.into(), + }); + return; + } + }; - if let Some(reason) = choice.finish_reason { - finish_reason = FinishReason { - unified: map_deepseek_finish_reason(Some(&reason)), - raw: Some(reason), - }; - } + if !self.metadata_emitted { + self.metadata_emitted = true; + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: chunk.id.clone(), + timestamp: timestamp(chunk.created), + model_id: chunk.model.clone(), + })); + } - if let Some(logprobs) = &choice.logprobs { - for (key, sink) in [ - ("content", &mut content_logprobs), - ("reasoning_content", &mut reasoning_logprobs), - ] { - if let Some(entries) = logprobs.get(key).and_then(Value::as_array) { - sink.extend(entries.iter().cloned()); - } - } - } + if let Some(chunk_usage) = chunk.usage.filter(|usage| !usage.is_null()) { + self.usage = Some(chunk_usage); + } - let Some(delta) = choice.delta else { - continue; - }; + if chunk.object.is_some() { + self.response_object = chunk.object; + } - if delta.role.is_some() { - message_role = delta.role; - } + // The fingerprint is repeated on stream chunks; keep the + // latest non-null value in case it changes during the response. + if chunk.system_fingerprint.is_some() { + self.system_fingerprint = chunk.system_fingerprint; + } - // enqueue reasoning before text deltas: - if let Some(reasoning) = delta - .reasoning_content - .filter(|reasoning| !reasoning.is_empty()) - { - if !is_active_reasoning { - yield Ok(StreamPart::ReasoningStart { - id: "reasoning-0".to_string(), - provider_metadata: None, - }); - is_active_reasoning = true; - } - yield Ok(StreamPart::ReasoningDelta { - id: "reasoning-0".to_string(), - delta: reasoning, - provider_metadata: None, - }); - } + let Some(choice) = chunk.choices.into_iter().next() else { + return; + }; - if let Some(text) = delta.content.filter(|text| !text.is_empty()) { - if !is_active_text { - yield Ok(StreamPart::TextStart { - id: "txt-0".to_string(), - provider_metadata: None, - }); - is_active_text = true; - } - - // end reasoning when text starts: - if is_active_reasoning { - yield Ok(StreamPart::ReasoningEnd { - id: "reasoning-0".to_string(), - provider_metadata: None, - }); - is_active_reasoning = false; - } - - yield Ok(StreamPart::TextDelta { - id: "txt-0".to_string(), - delta: text, - provider_metadata: None, - }); - } + if choice.index.is_some() { + self.choice_index = choice.index; + } + + if let Some(reason) = choice.finish_reason { + self.finish_reason = Some(FinishReason { + unified: map_deepseek_finish_reason(Some(&reason)), + raw: Some(reason), + }); + } - if let Some(deltas) = delta.tool_calls.filter(|deltas| !deltas.is_empty()) { - // end reasoning when tool calls start: - if is_active_reasoning { - yield Ok(StreamPart::ReasoningEnd { - id: "reasoning-0".to_string(), - provider_metadata: None, - }); - is_active_reasoning = false; - } - - for tool_call in &deltas { - if let Some(kind) = &tool_call.r#type { - tool_call_types.insert(tool_call.index, kind.clone()); - } - let delta = StreamingToolCallDelta { - index: Some(tool_call.index), - id: tool_call.id.as_deref(), - r#type: tool_call.r#type.as_deref(), - name: tool_call.function.name.as_deref(), - arguments: tool_call.function.arguments.as_deref(), - provider_metadata: None, - }; - if let Err(error) = tool_calls.process(delta, &mut tool_parts) { - yield Err(error.into()); - return; - } - for part in tool_parts.drain(..) { - yield Ok(part); - } - } + if let Some(logprobs) = &choice.logprobs { + for (key, sink) in [ + ("content", &mut self.content_logprobs), + ("reasoning_content", &mut self.reasoning_logprobs), + ] { + if let Some(entries) = logprobs.get(key).and_then(Value::as_array) { + sink.extend(entries.iter().cloned()); } } + } - if is_active_reasoning { - yield Ok(StreamPart::ReasoningEnd { - id: "reasoning-0".to_string(), - provider_metadata: None, - }); - } - if is_active_text { - yield Ok(StreamPart::TextEnd { - id: "txt-0".to_string(), - provider_metadata: None, - }); + let Some(delta) = choice.delta else { + return; + }; + + if delta.role.is_some() { + self.message_role = delta.role; + } + + // enqueue reasoning before text deltas: + if let Some(reasoning) = delta + .reasoning_content + .filter(|reasoning| !reasoning.is_empty()) + { + if !self.is_active_reasoning { + controller.reasoning_start("reasoning-0"); + self.is_active_reasoning = true; } - tool_calls.finish(&mut tool_parts); - for part in tool_parts.drain(..) { - yield Ok(part); + controller.reasoning_delta("reasoning-0", reasoning); + } + + if let Some(text) = delta.content.filter(|text| !text.is_empty()) { + if !self.is_active_text { + controller.text_start("txt-0"); + self.is_active_text = true; } - let mut metadata = Map::new(); - for (key, name) in [ - ("prompt_cache_hit_tokens", "promptCacheHitTokens"), - ("prompt_cache_miss_tokens", "promptCacheMissTokens"), - ] { - if let Some(value) = usage - .as_ref() - .and_then(|usage| usage.get(key)) - .filter(|value| !value.is_null()) - { - metadata.insert(name.into(), value.clone()); + // end reasoning when text starts: + self.end_active_reasoning(controller); + + controller.text_delta("txt-0", text); + } + + if let Some(deltas) = delta.tool_calls.filter(|deltas| !deltas.is_empty()) { + // end reasoning when tool calls start: + self.end_active_reasoning(controller); + + for tool_call in &deltas { + if let Some(kind) = &tool_call.r#type { + self.tool_call_types.insert(tool_call.index, kind.clone()); + } + let delta = ToolCallDelta { + index: Some(tool_call.index), + id: tool_call.id.as_deref(), + name: tool_call.function.name.as_deref(), + arguments: tool_call.function.arguments.as_deref(), + ..ToolCallDelta::default() + }; + if !self.tool_calls.delta(delta, controller) { + return; } } - if let Some(object) = response_object { - metadata.insert("responseObject".into(), json!(object)); - } - if let Some(index) = choice_index { - metadata.insert("choiceIndex".into(), json!(index)); - } - if let Some(role) = message_role { - metadata.insert("messageRole".into(), json!(role)); - } - if !tool_call_types.is_empty() { - metadata.insert( - "toolCallTypes".into(), - json!(tool_call_types.into_values().collect::>()), - ); - } - let mut logprobs = Map::new(); - if !content_logprobs.is_empty() { - logprobs.insert("content".into(), Value::Array(content_logprobs)); - } - if !reasoning_logprobs.is_empty() { - logprobs.insert("reasoning_content".into(), Value::Array(reasoning_logprobs)); - } - if !logprobs.is_empty() { - metadata.insert("logprobs".into(), Value::Object(logprobs)); - } - if let Some(system_fingerprint) = system_fingerprint { - metadata.insert("systemFingerprint".into(), json!(system_fingerprint)); - } - yield Ok(StreamPart::Finish { - finish_reason, - usage: convert_deepseek_usage(usage.as_ref()), - provider_metadata: Some(provider_namespace(&provider_options_name, Value::Object(metadata)).expect("metadata payload is an object")), - }); - }; + } + } - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(RequestInfo { body: Some(body) }), - response: Some(StreamResponseInfo { - headers: Some(response_headers), - }), - }) + fn flush(mut self, controller: &mut Controller) { + self.end_active_reasoning(controller); + if self.is_active_text { + controller.text_end("txt-0"); + } + self.tool_calls.finish(controller); + + let mut metadata = Map::new(); + for (key, name) in [ + ("prompt_cache_hit_tokens", "promptCacheHitTokens"), + ("prompt_cache_miss_tokens", "promptCacheMissTokens"), + ] { + if let Some(value) = self + .usage + .as_ref() + .and_then(|usage| usage.get(key)) + .filter(|value| !value.is_null()) + { + metadata.insert(name.into(), value.clone()); + } + } + if let Some(object) = self.response_object { + metadata.insert("responseObject".into(), json!(object)); + } + if let Some(index) = self.choice_index { + metadata.insert("choiceIndex".into(), json!(index)); + } + if let Some(role) = self.message_role { + metadata.insert("messageRole".into(), json!(role)); + } + if !self.tool_call_types.is_empty() { + metadata.insert( + "toolCallTypes".into(), + json!(self.tool_call_types.into_values().collect::>()), + ); + } + let mut logprobs = Map::new(); + if !self.content_logprobs.is_empty() { + logprobs.insert("content".into(), Value::Array(self.content_logprobs)); + } + if !self.reasoning_logprobs.is_empty() { + logprobs.insert( + "reasoning_content".into(), + Value::Array(self.reasoning_logprobs), + ); + } + if !logprobs.is_empty() { + metadata.insert("logprobs".into(), Value::Object(logprobs)); + } + if let Some(system_fingerprint) = self.system_fingerprint { + metadata.insert("systemFingerprint".into(), json!(system_fingerprint)); + } + controller.enqueue(StreamPart::Finish { + finish_reason: self.finish_reason.unwrap_or(FINISH_OTHER), + usage: convert_deepseek_usage(self.usage.as_ref()), + provider_metadata: Some( + provider_namespace(&self.provider_options_name, Value::Object(metadata)) + .expect("metadata payload is an object"), + ), + }); } } diff --git a/aimux-providers/src/google/model.rs b/aimux-providers/src/google/model.rs index 1d7ac9306..bbc08f82b 100644 --- a/aimux-providers/src/google/model.rs +++ b/aimux-providers/src/google/model.rs @@ -1,5 +1,7 @@ //! Google Gemini language model — implements `LanguageModel`. +use std::collections::HashSet; + use aimux_core::tool::RawToolCall; use aimux_core::tool::ToolResult; use async_trait::async_trait; @@ -17,6 +19,9 @@ use aimux_core::stream_part::StreamPart; use aimux_core::types::{ FinishReason, FinishReasonUnified, ProviderMetadata, ResponseMetadata, Usage, }; +use aimux_provider_utils::{Transformer, pipe_through}; + +use crate::stream::{Controller, Emit, FINISH_ERROR, FINISH_OTHER, StreamContext}; use aimux_core::language_model::SupportedUrls; @@ -25,7 +30,7 @@ use super::convert::{ parse_finish_reason, validate_call_options, }; use super::options::google_metadata; -use super::types::{Candidate, GenerateContentResponse, GoogleStreamEvent}; +use super::types::{Candidate, GenerateContentResponse, GoogleStreamEvent, StreamChunk}; use crate::shared::EndpointConfig; /// A Google Gemini language model. @@ -126,10 +131,7 @@ impl LanguageModel for GoogleModel { .finish_reason .as_deref() .map(|r| parse_finish_reason(r, has_tool_calls)) - .unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }) + .unwrap_or(FINISH_OTHER) }; let usage = data @@ -197,450 +199,498 @@ impl LanguageModel for GoogleModel { Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), first_event => first_event, }; - let stream_error_url = endpoint.clone(); - let stream_request_body = body.clone(); - let stream_error_headers = response_headers.clone(); - - let generate_id = self.generate_id.clone(); - let include_raw_chunks = options.include_raw_chunks.unwrap_or(false); - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings: tool_warnings }); - - let mut sse_stream = futures::stream::iter(first_event.into_iter()).chain(sse_stream); - let mut text_id: Option = None; - let mut reasoning_id: Option = None; - let mut block_counter = 0usize; - let mut final_usage: Usage = Usage::default(); - let mut final_finish_reason: Option = None; - let mut has_tool_calls = false; - let mut response_metadata_emitted = false; - let mut stream_errored = false; - let mut prompt_blocked = false; - - // Provider-metadata accumulators (mirrors TS `lastGroundingMetadata` / - // `lastUrlContextMetadata` + the finishReason-chunk snapshot). - let mut last_grounding_metadata: Option = None; - let mut last_url_context_metadata: Option = None; - let mut last_prompt_feedback: Option = None; - let mut last_safety_ratings: Option = None; - let mut last_finish_message: Option = None; - let mut last_usage_metadata_value: Option = None; - - // Source dedup across chunks (url sources only). - let mut emitted_source_urls: std::collections::HashSet = - std::collections::HashSet::new(); - let mut source_id = 0usize; - - // Associates code-execution results / server tool responses with - // their preceding call. - let mut last_code_execution_tool_call_id: Option = None; - let mut last_server_tool_call_id: Option = None; - - while let Some(event) = sse_stream.next().await { - if stream_errored { - break; - } - let event = match event { - Ok(raw_value) => { - if include_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: raw_value.clone() }); - } - serde_json::from_value::(raw_value).map_err(AiMuxError::from) - } - Err(error) => Err(error), - }; - match event { - Ok(GoogleStreamEvent::Chunk(chunk)) => { - - // Emit ResponseMetadata from the first chunk that has - // a responseId (matches TS behaviour). - if !response_metadata_emitted - && let Some(id) = &chunk.response_id { - response_metadata_emitted = true; - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: Some(id.clone()), - timestamp: None, - model_id: None, - })); - } + let stream = pipe_through( + futures::stream::iter(first_event).chain(sse_stream), + GoogleLanguageModelStream { + ctx: StreamContext { + generate_id: self.generate_id.clone(), + ..StreamContext::new(options, tool_warnings, endpoint, &body, &response_headers) + }, + code_execution_tool_name, + ..GoogleLanguageModelStream::default() + }, + ); - if let Some(usage) = &chunk.usage_metadata { - final_usage = convert_usage(usage); - last_usage_metadata_value = - serde_json::to_value(usage).ok(); - } - if let Some(pf) = &chunk.prompt_feedback - && !prompt_blocked { - last_prompt_feedback = Some(pf.clone()); - if let Some(reason) = confirmed_prompt_block_reason(Some(pf)) { - prompt_blocked = true; - final_finish_reason = Some(FinishReason { - unified: FinishReasonUnified::ContentFilter, raw: Some(reason), - }); - } - } + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), + response: Some(aimux_core::shared::StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} - let Some(candidates) = chunk.candidates else { - continue; - }; - let Some(candidate) = candidates.into_iter().next() else { - continue; - }; +/// The `TransformStream` of `GoogleLanguageModel.doStream`. +#[derive(Default)] +struct GoogleLanguageModelStream { + ctx: StreamContext, + code_execution_tool_name: String, + text_id: Option, + reasoning_id: Option, + block_counter: usize, + final_usage: Usage, + final_finish_reason: Option, + has_tool_calls: bool, + response_metadata_emitted: bool, + stream_errored: bool, + prompt_blocked: bool, + // Provider-metadata accumulators (mirrors TS `lastGroundingMetadata` / + // `lastUrlContextMetadata` + the finishReason-chunk snapshot). + last_grounding_metadata: Option, + last_url_context_metadata: Option, + last_prompt_feedback: Option, + last_safety_ratings: Option, + last_finish_message: Option, + last_usage_metadata_value: Option, + // Source dedup across chunks (url sources only). + emitted_source_urls: HashSet, + source_id: usize, + // Associates code-execution results / server tool responses with their + // preceding call. + last_code_execution_tool_call_id: Option, + last_server_tool_call_id: Option, +} - if let Some(gm) = &candidate.grounding_metadata { - last_grounding_metadata = Some(gm.clone()); - } - if let Some(ucm) = &candidate.url_context_metadata { - last_url_context_metadata = Some(ucm.clone()); - } +impl GoogleLanguageModelStream { + fn transform_chunk(&mut self, chunk: Box, controller: &mut Controller) { + // Emit ResponseMetadata from the first chunk that has + // a responseId (matches TS behaviour). + if !self.response_metadata_emitted + && let Some(id) = &chunk.response_id + { + self.response_metadata_emitted = true; + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: Some(id.clone()), + timestamp: None, + model_id: None, + })); + } - if let Some(sr) = &candidate.safety_ratings { - last_safety_ratings = serde_json::to_value(sr).ok(); - } - if let Some(fm) = &candidate.finish_message { - last_finish_message = Some(json!(fm)); - } - if prompt_blocked { continue; } - - // Extract url sources from this chunk's grounding metadata - // (deduplicated across chunks; document sources are not - // emitted in the stream, matching TS). - let chunk_sources = - extract_sources(candidate.grounding_metadata.as_ref(), &mut source_id); - for src in chunk_sources { - if let GenerateContent::Source(Source::Url { - url, - id: _, - title, - provider_metadata: None, - }) = src - && emitted_source_urls.insert(url.clone()) { - yield Ok(StreamPart::Source(Source::Url { - id: generate_id(), - url, - title, - provider_metadata: None, - })); - } - } + if let Some(usage) = &chunk.usage_metadata { + self.final_usage = convert_usage(usage); + self.last_usage_metadata_value = serde_json::to_value(usage).ok(); + } + if let Some(pf) = &chunk.prompt_feedback + && !self.prompt_blocked + { + self.last_prompt_feedback = Some(pf.clone()); + if let Some(reason) = confirmed_prompt_block_reason(Some(pf)) { + self.prompt_blocked = true; + self.final_finish_reason = Some(FinishReason { + unified: FinishReasonUnified::ContentFilter, + raw: Some(reason), + }); + } + } - if let Some(parts) = - candidate.content.as_ref().and_then(|c| c.parts.as_ref()) + let Some(candidates) = chunk.candidates else { + return; + }; + let Some(candidate) = candidates.into_iter().next() else { + return; + }; + + if let Some(gm) = &candidate.grounding_metadata { + self.last_grounding_metadata = Some(gm.clone()); + } + if let Some(ucm) = &candidate.url_context_metadata { + self.last_url_context_metadata = Some(ucm.clone()); + } + + if let Some(sr) = &candidate.safety_ratings { + self.last_safety_ratings = serde_json::to_value(sr).ok(); + } + if let Some(fm) = &candidate.finish_message { + self.last_finish_message = Some(json!(fm)); + } + if self.prompt_blocked { + return; + } + + // Extract url sources from this chunk's grounding metadata + // (deduplicated across chunks; document sources are not + // emitted in the stream, matching TS). + let chunk_sources = + extract_sources(candidate.grounding_metadata.as_ref(), &mut self.source_id); + for src in chunk_sources { + if let GenerateContent::Source(Source::Url { + url, + id: _, + title, + provider_metadata: None, + }) = src + && self.emitted_source_urls.insert(url.clone()) + { + controller.enqueue(StreamPart::Source(Source::Url { + id: (self.ctx.generate_id)(), + url, + title, + provider_metadata: None, + })); + } + } + + if let Some(parts) = candidate.content.as_ref().and_then(|c| c.parts.as_ref()) { + for part in parts { + // thoughtSignature → provider_metadata (upstream :778-782) + let thought_sig_meta: Option = part + .get("thoughtSignature") + .and_then(|v| v.as_str()) + .filter(|signature| !signature.is_empty()) + .map(|s| google_metadata(json!({ "thoughtSignature": s }))); + + // text part + if let Some(text) = part.get("text").and_then(|v| v.as_str()) { + if text.is_empty() { + // Empty-text part may carry a thoughtSignature + // (upstream :784-794): emit as a zero-length + // delta with the signature metadata on + // whichever block is open — text or + // reasoning (mirrors the non-streaming + // path, which attaches to the last + // content item regardless of type). + if let Some(meta) = &thought_sig_meta + && let Some(id) = &self.text_id { - for part in parts { - // thoughtSignature → provider_metadata (upstream :778-782) - let thought_sig_meta: Option = part - .get("thoughtSignature") - .and_then(|v| v.as_str()) - .filter(|signature| !signature.is_empty()) - .map(|s| google_metadata(json!({ "thoughtSignature": s }))); - - // text part - if let Some(text) = part.get("text").and_then(|v| v.as_str()) { - if text.is_empty() { - // Empty-text part may carry a thoughtSignature - // (upstream :784-794): emit as a zero-length - // delta with the signature metadata on - // whichever block is open — text or - // reasoning (mirrors the non-streaming - // path, which attaches to the last - // content item regardless of type). - if let Some(meta) = &thought_sig_meta - && let Some(id) = &text_id { - yield Ok(StreamPart::TextDelta { - id: id.clone(), - delta: String::new(), - provider_metadata: Some(meta.clone()), - }); - - } - } else { - let is_thought = part.get("thought").and_then(serde_json::Value::as_bool).unwrap_or(false); - if is_thought { - // Close any open text block before starting reasoning - if let Some(id) = text_id.take() { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); - } - // Start a reasoning block if not already active - if reasoning_id.is_none() { - let id = format!("{block_counter}"); - block_counter += 1; - reasoning_id = Some(id.clone()); - yield Ok(StreamPart::ReasoningStart { - id, - provider_metadata: thought_sig_meta.clone(), - }); - } - if let Some(id) = &reasoning_id { - yield Ok(StreamPart::ReasoningDelta { - id: id.clone(), - delta: text.to_string(), - provider_metadata: thought_sig_meta.clone(), - }); - } - } else { - // Close any open reasoning block before starting text - if let Some(id) = reasoning_id.take() { - yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); - } - if text_id.is_none() { - let id = format!("{block_counter}"); - block_counter += 1; - text_id = Some(id.clone()); - yield Ok(StreamPart::TextStart { id, provider_metadata: thought_sig_meta.clone() }); - } - if let Some(id) = &text_id { - yield Ok(StreamPart::TextDelta { - id: id.clone(), - delta: text.to_string(), - provider_metadata: thought_sig_meta.clone(), - }); - } - } - } - } else if let Some(ec) = part.get("executableCode") { - // Provider-executed code execution. - let has_code = ec - .get("code") - .and_then(|v| v.as_str()) - .map(|s| !s.is_empty()) - .unwrap_or(false); - if has_code { - let id = generate_id(); - last_code_execution_tool_call_id = Some(id.clone()); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id.clone(), - tool_name: code_execution_tool_name.clone(), - input: ec.to_string(), - provider_executed: Some(true), - dynamic: None, - provider_metadata: None, - })); - // provider-executed → does NOT set has_tool_calls - } - } else if let Some(cer) = part.get("codeExecutionResult") { - // Result corresponds to the most recent - // executableCode part. Gemini may emit - // several results for that one call, so - // retain the association until a new call. - if let Some(call_id) = - last_code_execution_tool_call_id.as_ref() - { - let outcome = - cer.get("outcome").cloned().unwrap_or(json!(null)); - let output = cer - .get("output") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - .unwrap_or_default(); - yield Ok(StreamPart::ToolResult(ToolResult { - tool_call_id: call_id.clone(), - tool_name: code_execution_tool_name.clone(), - result: json!({ "outcome": outcome, "output": output }), - is_error: None, - preliminary: None, - dynamic: None, - provider_metadata: None, - })); - } - } else if let Some(tc) = part.get("toolCall") { - // Server-side tool call (provider-executed). - let tool_type = tc - .get("toolType") - .and_then(|v| v.as_str()) - .unwrap_or(""); - let id = tc - .get("id") - .and_then(|v| v.as_str()) - .filter(|id| !id.is_empty()) - .map(std::string::ToString::to_string) - .unwrap_or_else(|| generate_id()); - last_server_tool_call_id = Some(id.clone()); - let args = tc.get("args").cloned().unwrap_or(json!({})); - let server_meta = server_tool_metadata( - &id, - tool_type, - part.get("thoughtSignature").and_then(|v| v.as_str()), - ); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id, - tool_name: format!("server:{tool_type}"), - input: args.to_string(), - provider_executed: Some(true), - dynamic: Some(true), - provider_metadata: Some(server_meta), - })); - // provider-executed → does NOT set has_tool_calls - } else if let Some(tr) = part.get("toolResponse") { - // Server-side tool response. - let tool_type = tr.get("toolType").and_then(|v| v.as_str()).unwrap_or(""); - let id = last_server_tool_call_id - .take() - .or_else(|| { - tr.get("id") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - }) - .filter(|id| !id.is_empty()) - .unwrap_or_else(|| generate_id()); - let response = - tr.get("response").cloned().unwrap_or(json!({})); - let server_meta = server_tool_metadata( - &id, - tool_type, - part.get("thoughtSignature").and_then(|v| v.as_str()), - ); - yield Ok(StreamPart::ToolResult(ToolResult { - tool_call_id: id, - tool_name: format!("server:{tool_type}"), - result: response, - is_error: None, - preliminary: None, - dynamic: None, - provider_metadata: Some(server_meta), - })); - } else if let Some(inline) = part.get("inlineData") { - // File output — upstream :847-877. - // Close any open text/reasoning block before file. - if let Some(id) = text_id.take() { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); - } - if let Some(id) = reasoning_id.take() { - yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); - } - if let (Some(data), Some(mime)) = ( - inline.get("data").and_then(|v| v.as_str()), - inline.get("mimeType").and_then(|v| v.as_str()), - ) { - let file = LanguageModelFile { - data: GeneratedFileData::Data { data: FileBytes::Base64(data.to_string()) }, - media_type: mime.to_string(), - provider_metadata: thought_sig_meta.clone(), - }; - yield Ok(if part.get("thought").and_then(serde_json::Value::as_bool).unwrap_or(false) { StreamPart::ReasoningFile(file) } else { StreamPart::File(file) }); - } - } + controller.enqueue(StreamPart::TextDelta { + id: id.clone(), + delta: String::new(), + provider_metadata: Some(meta.clone()), + }); + } + } else { + let is_thought = part + .get("thought") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + if is_thought { + // Close any open text block before starting reasoning + if let Some(id) = self.text_id.take() { + controller.text_end(id); } - for part in parts { - let thought_sig_meta = part.get("thoughtSignature").and_then(Value::as_str) - .filter(|signature| !signature.is_empty()) - .map(|signature| google_metadata(json!({"thoughtSignature": signature}))); - if let Some(fc) = part.get("functionCall") { - // Single-chunk complete function call (the - // common case for the public Gemini API). - if fc.get("name").and_then(Value::as_str).is_none() || fc.get("partialArgs").is_some_and(|value| !value.is_null()) || fc.get("willContinue").and_then(Value::as_bool) == Some(true) { continue; } - let name = fc - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or(""); - let id = fc - .get("id") - .and_then(|v| v.as_str()) - .filter(|id| !id.is_empty()) - .filter(|id| !id.is_empty()) - .map(std::string::ToString::to_string) - .unwrap_or_else(|| generate_id()); - let args = fc.get("args").filter(|value| !value.is_null()).cloned().unwrap_or(json!({})); - - yield Ok(StreamPart::ToolInputStart { - id: id.clone(), - tool_name: name.to_string(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: thought_sig_meta.clone(), - }); - let args_str = args.as_str().map(str::to_owned).unwrap_or_else(|| args.to_string()); - if fc.get("args").is_some_and(|value| !value.is_null()) { - yield Ok(StreamPart::ToolInputDelta { - id: id.clone(), - delta: args_str.clone(), - provider_metadata: thought_sig_meta.clone(), - }); - } - yield Ok(StreamPart::ToolInputEnd { id: id.clone(), provider_metadata: thought_sig_meta.clone()}); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id, - tool_name: name.to_string(), - input: args_str, - provider_executed: None, - dynamic: None, - provider_metadata: thought_sig_meta.clone(), - })); - has_tool_calls = true; - } + // Start a reasoning block if not already active + if self.reasoning_id.is_none() { + let id = format!("{}", self.block_counter); + self.block_counter += 1; + self.reasoning_id = Some(id.clone()); + controller.enqueue(StreamPart::ReasoningStart { + id, + provider_metadata: thought_sig_meta.clone(), + }); + } + if let Some(id) = &self.reasoning_id { + controller.enqueue(StreamPart::ReasoningDelta { + id: id.clone(), + delta: text.to_string(), + provider_metadata: thought_sig_meta.clone(), + }); + } + } else { + // Close any open reasoning block before starting text + if let Some(id) = self.reasoning_id.take() { + controller.reasoning_end(id); + } + if self.text_id.is_none() { + let id = format!("{}", self.block_counter); + self.block_counter += 1; + self.text_id = Some(id.clone()); + controller.enqueue(StreamPart::TextStart { + id, + provider_metadata: thought_sig_meta.clone(), + }); + } + if let Some(id) = &self.text_id { + controller.enqueue(StreamPart::TextDelta { + id: id.clone(), + delta: text.to_string(), + provider_metadata: thought_sig_meta.clone(), + }); } - } - - if let Some(reason) = candidate.finish_reason.as_deref() { - final_finish_reason = - Some(parse_finish_reason(reason, has_tool_calls)); } } - Ok(GoogleStreamEvent::Error(error)) => { - yield Ok(StreamPart::Error { - error: super::google_stream_error( - &error.error, - &stream_error_url, - stream_request_body.clone(), - stream_error_headers.clone(), - ), - }); - stream_errored = true; - break; + } else if let Some(ec) = part.get("executableCode") { + // Provider-executed code execution. + let has_code = ec + .get("code") + .and_then(|v| v.as_str()) + .map(|s| !s.is_empty()) + .unwrap_or(false); + if has_code { + let id = (self.ctx.generate_id)(); + self.last_code_execution_tool_call_id = Some(id.clone()); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: id.clone(), + tool_name: self.code_execution_tool_name.clone(), + input: ec.to_string(), + provider_executed: Some(true), + dynamic: None, + provider_metadata: None, + })); + // provider-executed → does NOT set has_tool_calls } - Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Ok(StreamPart::Error { error }); - if !recoverable { - return; - } + } else if let Some(cer) = part.get("codeExecutionResult") { + // Result corresponds to the most recent + // executableCode part. Gemini may emit + // several results for that one call, so + // retain the association until a new call. + if let Some(call_id) = self.last_code_execution_tool_call_id.as_ref() { + let outcome = cer.get("outcome").cloned().unwrap_or(json!(null)); + let output = cer + .get("output") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string) + .unwrap_or_default(); + controller.enqueue(StreamPart::ToolResult(ToolResult { + tool_call_id: call_id.clone(), + tool_name: self.code_execution_tool_name.clone(), + result: json!({ "outcome": outcome, "output": output }), + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: None, + })); + } + } else if let Some(tc) = part.get("toolCall") { + // Server-side tool call (provider-executed). + let tool_type = tc.get("toolType").and_then(|v| v.as_str()).unwrap_or(""); + let id = tc + .get("id") + .and_then(|v| v.as_str()) + .filter(|id| !id.is_empty()) + .map(std::string::ToString::to_string) + .unwrap_or_else(|| (self.ctx.generate_id)()); + self.last_server_tool_call_id = Some(id.clone()); + let args = tc.get("args").cloned().unwrap_or(json!({})); + let server_meta = server_tool_metadata( + &id, + tool_type, + part.get("thoughtSignature").and_then(|v| v.as_str()), + ); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: id, + tool_name: format!("server:{tool_type}"), + input: args.to_string(), + provider_executed: Some(true), + dynamic: Some(true), + provider_metadata: Some(server_meta), + })); + // provider-executed → does NOT set has_tool_calls + } else if let Some(tr) = part.get("toolResponse") { + // Server-side tool response. + let tool_type = tr.get("toolType").and_then(|v| v.as_str()).unwrap_or(""); + let id = self + .last_server_tool_call_id + .take() + .or_else(|| { + tr.get("id") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string) + }) + .filter(|id| !id.is_empty()) + .unwrap_or_else(|| (self.ctx.generate_id)()); + let response = tr.get("response").cloned().unwrap_or(json!({})); + let server_meta = server_tool_metadata( + &id, + tool_type, + part.get("thoughtSignature").and_then(|v| v.as_str()), + ); + controller.enqueue(StreamPart::ToolResult(ToolResult { + tool_call_id: id, + tool_name: format!("server:{tool_type}"), + result: response, + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: Some(server_meta), + })); + } else if let Some(inline) = part.get("inlineData") { + // File output — upstream :847-877. + // Close any open text/reasoning block before file. + if let Some(id) = self.text_id.take() { + controller.text_end(id); + } + if let Some(id) = self.reasoning_id.take() { + controller.reasoning_end(id); + } + if let (Some(data), Some(mime)) = ( + inline.get("data").and_then(|v| v.as_str()), + inline.get("mimeType").and_then(|v| v.as_str()), + ) { + let file = LanguageModelFile { + data: GeneratedFileData::Data { + data: FileBytes::Base64(data.to_string()), + }, + media_type: mime.to_string(), + provider_metadata: thought_sig_meta.clone(), + }; + controller.enqueue( + if part + .get("thought") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false) + { + StreamPart::ReasoningFile(file) + } else { + StreamPart::File(file) + }, + ); } } } + for part in parts { + let thought_sig_meta = part + .get("thoughtSignature") + .and_then(Value::as_str) + .filter(|signature| !signature.is_empty()) + .map(|signature| google_metadata(json!({"thoughtSignature": signature}))); + if let Some(fc) = part.get("functionCall") { + // Single-chunk complete function call (the + // common case for the public Gemini API). + if fc.get("name").and_then(Value::as_str).is_none() + || fc.get("partialArgs").is_some_and(|value| !value.is_null()) + || fc.get("willContinue").and_then(Value::as_bool) == Some(true) + { + continue; + } + let name = fc.get("name").and_then(|v| v.as_str()).unwrap_or(""); + let id = fc + .get("id") + .and_then(|v| v.as_str()) + .filter(|id| !id.is_empty()) + .map(std::string::ToString::to_string) + .unwrap_or_else(|| (self.ctx.generate_id)()); + let args = fc + .get("args") + .filter(|value| !value.is_null()) + .cloned() + .unwrap_or(json!({})); + + controller.enqueue(StreamPart::ToolInputStart { + id: id.clone(), + tool_name: name.to_string(), + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: thought_sig_meta.clone(), + }); + let args_str = args + .as_str() + .map(str::to_owned) + .unwrap_or_else(|| args.to_string()); + if fc.get("args").is_some_and(|value| !value.is_null()) { + controller.enqueue(StreamPart::ToolInputDelta { + id: id.clone(), + delta: args_str.clone(), + provider_metadata: thought_sig_meta.clone(), + }); + } + controller.enqueue(StreamPart::ToolInputEnd { + id: id.clone(), + provider_metadata: thought_sig_meta.clone(), + }); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: id, + tool_name: name.to_string(), + input: args_str, + provider_executed: None, + dynamic: None, + provider_metadata: thought_sig_meta.clone(), + })); + self.has_tool_calls = true; + } + } + } - // Close any remaining open text/reasoning segment. - if let Some(id) = text_id.take() { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None}); + if let Some(reason) = candidate.finish_reason.as_deref() { + self.final_finish_reason = Some(parse_finish_reason(reason, self.has_tool_calls)); + } + } +} + +impl Transformer for GoogleLanguageModelStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut Controller) { + self.ctx.start(controller); + } + + fn transform(&mut self, event: Result, controller: &mut Controller) { + let event = match event { + Ok(raw_value) => { + self.ctx.raw(&raw_value, controller); + serde_json::from_value::(raw_value).map_err(AiMuxError::from) } - if let Some(id) = reasoning_id.take() { - yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None}); + Err(error) => Err(error), + }; + + match event { + Ok(GoogleStreamEvent::Chunk(chunk)) => self.transform_chunk(chunk, controller), + Ok(GoogleStreamEvent::Error(error)) => { + controller.enqueue(StreamPart::Error { + error: super::google_stream_error( + &error.error, + &self.ctx.url, + self.ctx.request_body.clone(), + self.ctx.response_headers.clone(), + ), + }); + // The rest of the body is not read: the stream finishes now. + self.stream_errored = true; + self.finish(controller); + controller.terminate(); } + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + controller.enqueue(StreamPart::Error { error }); + if !recoverable { + controller.terminate(); + } + } + } + } - let provider_metadata = Some(google_metadata(serde_json::json!({ - "promptFeedback": last_prompt_feedback, - "groundingMetadata": last_grounding_metadata, - "urlContextMetadata": last_url_context_metadata, - "safetyRatings": last_safety_ratings, - "usageMetadata": last_usage_metadata_value, - "finishMessage": last_finish_message, - "serviceTier": last_usage_metadata_value.as_ref().and_then(|usage| usage.get("serviceTier")), - }))); - - yield Ok(StreamPart::Finish { - finish_reason: if stream_errored { - FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - } - } else { - final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }) - }, - usage: if stream_errored { Usage::default() } else { final_usage }, - provider_metadata, - }); - }; + fn flush(mut self, controller: &mut Controller) { + self.finish(controller); + } +} - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - response: Some(aimux_core::shared::StreamResponseInfo { - headers: Some(response_headers), - }), - }) +impl GoogleLanguageModelStream { + /// Close the open segments and emit `Finish`: at the end of the body, or + /// right after an error event. + fn finish(&mut self, controller: &mut Controller) { + // Close any remaining open text/reasoning segment. + if let Some(id) = self.text_id.take() { + controller.text_end(id); + } + if let Some(id) = self.reasoning_id.take() { + controller.reasoning_end(id); + } + + let provider_metadata = Some(google_metadata(serde_json::json!({ + "promptFeedback": self.last_prompt_feedback, + "groundingMetadata": self.last_grounding_metadata, + "urlContextMetadata": self.last_url_context_metadata, + "safetyRatings": self.last_safety_ratings, + "usageMetadata": self.last_usage_metadata_value, + "finishMessage": self.last_finish_message, + "serviceTier": self.last_usage_metadata_value.as_ref().and_then(|usage| usage.get("serviceTier")), + }))); + + controller.enqueue(StreamPart::Finish { + finish_reason: if self.stream_errored { + FINISH_ERROR + } else { + self.final_finish_reason.take().unwrap_or(FINISH_OTHER) + }, + usage: if self.stream_errored { + Usage::default() + } else { + std::mem::take(&mut self.final_usage) + }, + provider_metadata, + }); } } diff --git a/aimux-providers/src/groq/mod.rs b/aimux-providers/src/groq/mod.rs index 82658b134..54dffbbdc 100644 --- a/aimux-providers/src/groq/mod.rs +++ b/aimux-providers/src/groq/mod.rs @@ -22,6 +22,8 @@ mod convert; mod error; mod finish_reason; mod model; +#[cfg(test)] +pub(crate) use model::chat_test_stream; mod options; mod prepare_tools; mod transcription; diff --git a/aimux-providers/src/groq/model.rs b/aimux-providers/src/groq/model.rs index ff02d4999..1511fd250 100644 --- a/aimux-providers/src/groq/model.rs +++ b/aimux-providers/src/groq/model.rs @@ -2,13 +2,12 @@ //! //! Mirrors `groq-chat-language-model.ts`: Groq's own implementation of the //! chat completions endpoint. Streamed tool calls go through the shared -//! [`StreamingToolCallTracker`]. +//! [`ChatToolCalls`]. use aimux_core::tool::RawToolCall; use std::collections::HashMap; use async_trait::async_trait; -use futures::StreamExt; use serde_json::{Map, Value, json}; use aimux_core::error::AiMuxError; @@ -17,21 +16,20 @@ use aimux_core::options::{CallOptions, ResponseFormat}; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, StreamResult}; use aimux_core::shared::{RequestInfo, ResponseInfo, StreamResponseInfo}; use aimux_core::stream_part::StreamPart; -use aimux_core::types::{ - FinishReason, FinishReasonUnified, ReasoningEffort, ResponseMetadata, Warning, -}; -use aimux_provider_utils::{ - StreamingToolCallDelta, StreamingToolCallTracker, TypeValidation, generate_id, -}; +use aimux_core::types::{FinishReason, ReasoningEffort, ResponseMetadata, Warning}; +use aimux_provider_utils::{Transformer, generate_id, pipe_through}; use crate::shared::EndpointConfig; +use crate::stream::{ + ChatToolCalls, Controller, Emit, FINISH_ERROR, FINISH_OTHER, StreamContext, ToolCallDelta, +}; use super::convert::convert_to_groq_chat_messages; use super::error::{GroqErrorData, groq_failed_response_handler}; use super::finish_reason::map_groq_finish_reason; use super::options::{self, GroqLanguageModelChatOptions, parse_groq_options}; use super::prepare_tools::prepare_tools; -use super::types::{GroqChatChunk, GroqChatResponse}; +use super::types::{GroqChatChunk, GroqChatResponse, GroqUsage}; use super::usage::convert_groq_usage; /// A Groq chat language model. @@ -361,209 +359,187 @@ impl LanguageModel for GroqChatLanguageModel { ) .await?; let response_headers = resp.response_headers; - let mut sse_stream = resp.value; - let emit_raw_chunks = options.include_raw_chunks == Some(true); - let stream_error_url = endpoint; - let stream_error_body = body.clone(); - let stream_response_headers = response_headers.clone(); + let stream = pipe_through( + resp.value, + GroqChatStream { + ctx: StreamContext::new(options, warnings, endpoint, &body, &response_headers), + ..GroqChatStream::default() + }, + ); - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(RequestInfo { body: Some(body) }), + response: Some(StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} - let mut tool_calls = StreamingToolCallTracker::new() - .with_generate_id(generate_id) - .with_type_validation(TypeValidation::Required); - let mut tool_parts = Vec::new(); +/// The `TransformStream` of `GroqChatLanguageModel.doStream`. +#[derive(Default)] +struct GroqChatStream { + ctx: StreamContext, + tool_calls: ChatToolCalls, + /// `other` until a chunk says otherwise. + finish_reason: Option, + usage: Option, + metadata_emitted: bool, + is_active_text: bool, + is_active_reasoning: bool, +} - let mut finish_reason = FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }; - let mut usage = None; - let mut is_first_chunk = true; - let mut is_active_text = false; - let mut is_active_reasoning = false; - - while let Some(event) = sse_stream.next().await { - let parsed = match event { - Ok(parsed) => parsed, - // a chunk that fails to parse is reported as an error part - // (`chunk.success === false`); a transport failure ends the stream. - Err(error) => { - if !error.is_recoverable_stream_error() { - yield Err(error); - return; - } - if emit_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: Value::Null }); - } - finish_reason = FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - }; - yield Ok(StreamPart::Error { error }); - continue; - } - }; +/// A chat stream with an empty request context and tool-call ids from +/// `generate_id`, for tests that drive it chunk by chunk. +#[cfg(test)] +pub(crate) fn chat_test_stream( + generate_id: std::sync::Arc String + Send + Sync>, +) -> impl Transformer, Output = StreamPart> { + GroqChatStream { + tool_calls: ChatToolCalls::new(generate_id), + ..GroqChatStream::default() + } +} - // Emit the raw chunk if requested (before anything else). - if emit_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); - } +impl GroqChatStream { + fn end_active_reasoning(&mut self, controller: &mut Controller) { + if self.is_active_reasoning { + controller.reasoning_end("reasoning-0"); + self.is_active_reasoning = false; + } + } +} - // handle error chunks: - if parsed.get("error").is_some() { - finish_reason = FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - }; - yield Ok(StreamPart::Error { - error: create_groq_stream_error( - &parsed, - &stream_error_url, - stream_error_body.clone(), - stream_response_headers.clone(), - ), - }); - continue; - } +impl Transformer for GroqChatStream { + type Input = Result; + type Output = StreamPart; - // handle failed chunk parsing / validation: - let value: GroqChatChunk = match serde_json::from_value(parsed) { - Ok(value) => value, - Err(error) => { - finish_reason = FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - }; - yield Ok(StreamPart::Error { error: error.into() }); - continue; - } - }; + fn start(&mut self, controller: &mut Controller) { + self.ctx.start(controller); + } - if is_first_chunk { - is_first_chunk = false; - let metadata = response_metadata(value.id, value.created, value.model); - yield Ok(StreamPart::ResponseMetadata(metadata)); + fn transform(&mut self, event: Result, controller: &mut Controller) { + let parsed = match event { + Ok(parsed) => parsed, + // a chunk that fails to parse is reported as an error part + // (`chunk.success === false`); a transport failure ends the stream. + Err(error) => { + if !error.is_recoverable_stream_error() { + controller.error(error); + return; } + self.ctx.raw(&Value::Null, controller); + self.finish_reason = Some(FINISH_ERROR); + controller.enqueue(StreamPart::Error { error }); + return; + } + }; - if let Some(chunk_usage) = value.x_groq.and_then(|x_groq| x_groq.usage) { - usage = Some(chunk_usage); - } + // Emit the raw chunk if requested (before anything else). + self.ctx.raw(&parsed, controller); + + // handle error chunks: + if parsed.get("error").is_some() { + self.finish_reason = Some(FINISH_ERROR); + controller.enqueue(StreamPart::Error { + error: create_groq_stream_error( + &parsed, + &self.ctx.url, + self.ctx.request_body.clone(), + self.ctx.response_headers.clone(), + ), + }); + return; + } - let Some(choice) = value.choices.into_iter().next() else { - continue; - }; + // handle failed chunk parsing / validation: + let value: GroqChatChunk = match serde_json::from_value(parsed) { + Ok(value) => value, + Err(error) => { + self.finish_reason = Some(FINISH_ERROR); + controller.enqueue(StreamPart::Error { + error: error.into(), + }); + return; + } + }; - if let Some(reason) = choice.finish_reason { - finish_reason = FinishReason { - unified: map_groq_finish_reason(Some(&reason)), - raw: Some(reason), - }; - } + if !self.metadata_emitted { + self.metadata_emitted = true; + let metadata = response_metadata(value.id, value.created, value.model); + controller.enqueue(StreamPart::ResponseMetadata(metadata)); + } - let Some(delta) = choice.delta else { - continue; - }; + if let Some(chunk_usage) = value.x_groq.and_then(|x_groq| x_groq.usage) { + self.usage = Some(chunk_usage); + } - if let Some(reasoning) = delta.reasoning.filter(|text| !text.is_empty()) { - if !is_active_reasoning { - yield Ok(StreamPart::ReasoningStart { - id: "reasoning-0".to_string(), - provider_metadata: None, - }); - is_active_reasoning = true; - } - yield Ok(StreamPart::ReasoningDelta { - id: "reasoning-0".to_string(), - delta: reasoning, - provider_metadata: None, - }); - } + let Some(choice) = value.choices.into_iter().next() else { + return; + }; - if let Some(text) = delta.content.filter(|text| !text.is_empty()) { - // end the active reasoning block before text starts - if is_active_reasoning { - yield Ok(StreamPart::ReasoningEnd { - id: "reasoning-0".to_string(), - provider_metadata: None, - }); - is_active_reasoning = false; - } - if !is_active_text { - yield Ok(StreamPart::TextStart { - id: "txt-0".to_string(), - provider_metadata: None, - }); - is_active_text = true; - } - yield Ok(StreamPart::TextDelta { - id: "txt-0".to_string(), - delta: text, - provider_metadata: None, - }); - } + if let Some(reason) = choice.finish_reason { + self.finish_reason = Some(FinishReason { + unified: map_groq_finish_reason(Some(&reason)), + raw: Some(reason), + }); + } - if let Some(deltas) = delta.tool_calls.filter(|deltas| !deltas.is_empty()) { - // end the active reasoning block before tool calls start - if is_active_reasoning { - yield Ok(StreamPart::ReasoningEnd { - id: "reasoning-0".to_string(), - provider_metadata: None, - }); - is_active_reasoning = false; - } - for tool_call in &deltas { - let delta = StreamingToolCallDelta { - index: Some(tool_call.index), - id: tool_call.id.as_deref(), - r#type: tool_call.r#type.as_deref(), - name: tool_call.function.name.as_deref(), - arguments: tool_call.function.arguments.as_deref(), - provider_metadata: None, - }; - if let Err(error) = tool_calls.process(delta, &mut tool_parts) { - yield Err(error.into()); - return; - } - for part in tool_parts.drain(..) { - yield Ok(part); - } - } - } + let Some(delta) = choice.delta else { + return; + }; + + if let Some(reasoning) = delta.reasoning.filter(|text| !text.is_empty()) { + if !self.is_active_reasoning { + controller.reasoning_start("reasoning-0"); + self.is_active_reasoning = true; } + controller.reasoning_delta("reasoning-0", reasoning); + } - if is_active_reasoning { - yield Ok(StreamPart::ReasoningEnd { - id: "reasoning-0".to_string(), - provider_metadata: None, - }); + if let Some(text) = delta.content.filter(|text| !text.is_empty()) { + // end the active reasoning block before text starts + self.end_active_reasoning(controller); + if !self.is_active_text { + controller.text_start("txt-0"); + self.is_active_text = true; } - if is_active_text { - yield Ok(StreamPart::TextEnd { - id: "txt-0".to_string(), + controller.text_delta("txt-0", text); + } + + if let Some(deltas) = delta.tool_calls.filter(|deltas| !deltas.is_empty()) { + // end the active reasoning block before tool calls start + self.end_active_reasoning(controller); + for tool_call in &deltas { + let delta = ToolCallDelta { + index: Some(tool_call.index), + id: tool_call.id.as_deref(), + name: tool_call.function.name.as_deref(), + arguments: tool_call.function.arguments.as_deref(), + // A call must start with `type: 'function'`. + invalid_type: tool_call.r#type.as_deref() != Some("function"), provider_metadata: None, - }); - } - tool_calls.finish(&mut tool_parts); - for part in tool_parts.drain(..) { - yield Ok(part); + }; + if !self.tool_calls.delta(delta, controller) { + return; + } } + } + } - yield Ok(StreamPart::Finish { - finish_reason, - usage: convert_groq_usage(usage.as_ref()), - provider_metadata: None, - }); - }; + fn flush(mut self, controller: &mut Controller) { + self.end_active_reasoning(controller); + if self.is_active_text { + controller.text_end("txt-0"); + } + self.tool_calls.finish(controller); - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(RequestInfo { body: Some(body) }), - response: Some(StreamResponseInfo { - headers: Some(response_headers), - }), - }) + controller.enqueue(StreamPart::Finish { + finish_reason: self.finish_reason.unwrap_or(FINISH_OTHER), + usage: convert_groq_usage(self.usage.as_ref()), + provider_metadata: None, + }); } } diff --git a/aimux-providers/src/huggingface/responses.rs b/aimux-providers/src/huggingface/responses.rs index 5c638f186..d965c10bc 100644 --- a/aimux-providers/src/huggingface/responses.rs +++ b/aimux-providers/src/huggingface/responses.rs @@ -37,9 +37,12 @@ use aimux_core::tool::Tool; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage, Warning}; use aimux_provider_utils::{ - MediaTypeData, detect_media_type, get_top_level_media_type, is_full_media_type, + MediaTypeData, Transformer, detect_media_type, get_top_level_media_type, is_full_media_type, + pipe_through, }; +use crate::stream::{Controller, Emit, FINISH_OTHER, StreamContext}; + use crate::shared::EndpointConfig; fn huggingface_failed_response_handler() -> aimux_provider_utils::ResponseHandler { @@ -236,295 +239,322 @@ impl LanguageModel for HuggingFaceResponsesModel { { return Err(error); } - let stream_error_url = endpoint; - let stream_request_body = body.clone(); - let stream_response_headers = response_headers.clone(); - - let warnings = request.warnings.clone(); - let stream = async_stream::stream! { - // First part: StreamStart. - yield Ok(StreamPart::StreamStart { warnings }); - - let mut finish_reason = FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }; - let mut response_id: Option = None; - let mut usage_raw: Option = None; - let mut event_iter = futures::stream::iter(first_event.into_iter()).chain(sse_stream); - - while let Some(event) = event_iter.next().await { - match event { - Ok(parsed) => { - if let Some(error) = huggingface_stream_error( - &parsed, - &stream_error_url, - stream_request_body.clone(), - stream_response_headers.clone(), - ) { - yield Ok(StreamPart::Error { error }); - break; + let stream = pipe_through( + futures::stream::iter(first_event).chain(sse_stream), + HuggingFaceResponsesStream { + ctx: StreamContext { + warnings: request.warnings.clone(), + url: endpoint, + request_body: body.clone(), + response_headers: response_headers.clone(), + ..StreamContext::default() + }, + finish_reason: FINISH_OTHER, + response_id: None, + usage_raw: None, + finished: false, + }, + ); + + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), + response: Some(aimux_core::shared::StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} + +/// The `TransformStream` of the Hugging Face Responses `doStream`. +struct HuggingFaceResponsesStream { + /// Raw chunks are never forwarded. + ctx: StreamContext, + finish_reason: FinishReason, + response_id: Option, + usage_raw: Option, + /// Set by an error chunk: later chunks are ignored, the stream still + /// finishes. + finished: bool, +} + +impl Transformer for HuggingFaceResponsesStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut Controller) { + self.ctx.start(controller); + } + + fn transform(&mut self, event: Result, controller: &mut Controller) { + self.transform_event(event, controller); + if self.finished { + // An error event ends the response: the rest of the body is not + // read and the stream finishes now. + self.finish(controller); + controller.terminate(); + } + } + + fn flush(mut self, controller: &mut Controller) { + self.finish(controller); + } +} + +impl HuggingFaceResponsesStream { + /// Handle one parsed SSE event. + fn transform_event(&mut self, event: Result, controller: &mut Controller) { + let Self { + ref mut finish_reason, + ref mut response_id, + ref mut usage_raw, + .. + } = *self; + match event { + Ok(parsed) => { + if let Some(error) = huggingface_stream_error( + &parsed, + &self.ctx.url, + self.ctx.request_body.clone(), + self.ctx.response_headers.clone(), + ) { + controller.enqueue(StreamPart::Error { error }); + self.finished = true; + return; + } + let chunk_type = parsed.get("type").and_then(|v| v.as_str()).unwrap_or(""); + + match chunk_type { + "response.created" => { + if let Some(resp) = parsed.get("response") { + *response_id = resp + .get("id") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + let created_at = resp + .get("created_at") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0); + let model = resp + .get("model") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: response_id.clone(), + timestamp: format_timestamp(created_at), + model_id: model, + })); } - let chunk_type = parsed - .get("type") - .and_then(|v| v.as_str()) - .unwrap_or(""); + } - match chunk_type { - "response.created" => { - if let Some(resp) = parsed.get("response") { - response_id = resp - .get("id") + "response.output_item.added" => { + if let Some(item) = parsed.get("item") { + let item_type = item.get("type").and_then(|v| v.as_str()).unwrap_or(""); + match item_type { + "message" => { + let role = + item.get("role").and_then(|v| v.as_str()).unwrap_or(""); + if role == "assistant" { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::TextStart { + id: id.clone(), + provider_metadata: Some( + super::options::huggingface_metadata( + json!({ "itemId": id }), + ), + ), + }); + } + } + "function_call" => { + let call_id = item + .get("call_id") .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - let created_at = resp - .get("created_at") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0); - let model = resp - .get("model") + .unwrap_or("") + .to_string(); + let name = item + .get("name") .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: response_id.clone(), - timestamp: format_timestamp(created_at), - model_id: model, - })); + .unwrap_or("") + .to_string(); + controller.tool_input_start(call_id, name); } - } - - "response.output_item.added" => { - if let Some(item) = parsed.get("item") { - let item_type = item - .get("type") + "reasoning" => { + let id = item + .get("id") .and_then(|v| v.as_str()) - .unwrap_or(""); - match item_type { - "message" => { - let role = item - .get("role") - .and_then(|v| v.as_str()) - .unwrap_or(""); - if role == "assistant" { - let id = item - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::TextStart { - id: id.clone(), - provider_metadata: Some(super::options::huggingface_metadata(json!({ "itemId": id }))), - }); - } - } - "function_call" => { - let call_id = item - .get("call_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let name = item - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::ToolInputStart { - id: call_id, - tool_name: name, - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - } - "reasoning" => { - let id = item - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::ReasoningStart { - id: id.clone(), - provider_metadata: Some(super::options::huggingface_metadata(json!({ "itemId": id }))), - }); - } - _ => {} - } + .unwrap_or("") + .to_string(); + controller.enqueue(StreamPart::ReasoningStart { + id: id.clone(), + provider_metadata: Some( + super::options::huggingface_metadata( + json!({ "itemId": id }), + ), + ), + }); } + _ => {} } + } + } - "response.output_item.done" => { - if let Some(item) = parsed.get("item") { - let item_type = item - .get("type") - .and_then(|v| v.as_str()) - .unwrap_or(""); - match item_type { - "message" => { - let role = item - .get("role") - .and_then(|v| v.as_str()) - .unwrap_or(""); - if role == "assistant" { - let id = item - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::TextEnd { id, provider_metadata: None}); - } - } - "function_call" => { - let call_id = item - .get("call_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let name = item - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let arguments = item - .get("arguments") - .and_then(|v| v.as_str()) - .unwrap_or("{}"); - let input = arguments.to_string(); - - yield Ok(StreamPart::ToolInputEnd { - id: call_id.clone(), - provider_metadata: None, - }); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: call_id.clone(), - tool_name: name.clone(), - input, - provider_executed: None, - dynamic: None, - provider_metadata: None, - })); - - if let Some(output) = - item.get("output").and_then(|v| v.as_str()) - { - yield Ok(StreamPart::ToolResult(ToolResult { - tool_call_id: call_id, - tool_name: name.clone(), - result: Value::String(output.to_string()), - is_error: None, - preliminary: None, - dynamic: None, - provider_metadata: None, - })); - } - } - _ => {} + "response.output_item.done" => { + if let Some(item) = parsed.get("item") { + let item_type = item.get("type").and_then(|v| v.as_str()).unwrap_or(""); + match item_type { + "message" => { + let role = + item.get("role").and_then(|v| v.as_str()).unwrap_or(""); + if role == "assistant" { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.text_end(id); } } - } - - "response.output_text.delta" => { - let item_id = parsed - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let delta = parsed - .get("delta") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::TextDelta { id: item_id, delta, provider_metadata: None}); - } - - "response.reasoning_text.delta" => { - let item_id = parsed - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let delta = parsed - .get("delta") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::ReasoningDelta { - id: item_id, - delta, - provider_metadata: None, - }); - } - - "response.reasoning_text.done" => { - let item_id = parsed - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::ReasoningEnd { id: item_id, - provider_metadata: None, - }); - } - - "response.completed" => { - if let Some(resp) = parsed.get("response") { - response_id = resp - .get("id") + "function_call" => { + let call_id = item + .get("call_id") .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - - let incomplete_reason = resp - .get("incomplete_details") - .and_then(|d| d.get("reason")) - .and_then(|r| r.as_str()); - - finish_reason = FinishReason { - unified: map_finish_reason( - incomplete_reason.unwrap_or("stop"), - ), - raw: incomplete_reason.map(std::string::ToString::to_string), - }; + .unwrap_or("") + .to_string(); + let name = item + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let arguments = item + .get("arguments") + .and_then(|v| v.as_str()) + .unwrap_or("{}"); + let input = arguments.to_string(); + + controller.tool_input_end(call_id.clone()); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: call_id.clone(), + tool_name: name.clone(), + input, + provider_executed: None, + dynamic: None, + provider_metadata: None, + })); - if let Some(usage) = resp.get("usage") - && !usage.is_null() + if let Some(output) = + item.get("output").and_then(|v| v.as_str()) { - usage_raw = Some(usage.clone()); + controller.enqueue(StreamPart::ToolResult(ToolResult { + tool_call_id: call_id, + tool_name: name.clone(), + result: Value::String(output.to_string()), + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: None, + })); } } + _ => {} } + } + } + + "response.output_text.delta" => { + let item_id = parsed + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let delta = parsed + .get("delta") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.text_delta(item_id, delta); + } - _ => { - // Unknown chunk type — ignore (matches the TS - // fallback schema which silently drops unknowns). + "response.reasoning_text.delta" => { + let item_id = parsed + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let delta = parsed + .get("delta") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.reasoning_delta(item_id, delta); + } + + "response.reasoning_text.done" => { + let item_id = parsed + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.reasoning_end(item_id); + } + + "response.completed" => { + if let Some(resp) = parsed.get("response") { + *response_id = resp + .get("id") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + + let incomplete_reason = resp + .get("incomplete_details") + .and_then(|d| d.get("reason")) + .and_then(|r| r.as_str()); + + *finish_reason = FinishReason { + unified: map_finish_reason(incomplete_reason.unwrap_or("stop")), + raw: incomplete_reason.map(std::string::ToString::to_string), + }; + + if let Some(usage) = resp.get("usage") + && !usage.is_null() + { + *usage_raw = Some(usage.clone()); } } } - Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Err(error); - if !recoverable { - return; - } + + _ => { + // Unknown chunk type — ignore (matches the TS + // fallback schema which silently drops unknowns). } } } + Err(error) => { + if error.is_recoverable_stream_error() { + controller.enqueue_error(error); + } else { + controller.error(error); + } + } + } + } - let usage = convert_usage(usage_raw.as_ref()); - - yield Ok(StreamPart::Finish { - finish_reason, - usage, - provider_metadata: Some(super::options::huggingface_metadata(json!({ "responseId": response_id }))), - }); - }; + /// Emit `Finish`: at the end of the body, or right after an error + /// event. + fn finish(&mut self, controller: &mut Controller) { + let usage = convert_usage(self.usage_raw.as_ref()); - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - response: Some(aimux_core::shared::StreamResponseInfo { - headers: Some(response_headers), - }), - }) + controller.enqueue(StreamPart::Finish { + finish_reason: self.finish_reason.clone(), + usage, + provider_metadata: Some(super::options::huggingface_metadata( + json!({ "responseId": self.response_id }), + )), + }); } } diff --git a/aimux-providers/src/lib.rs b/aimux-providers/src/lib.rs index 861a4e27d..280517f4c 100644 --- a/aimux-providers/src/lib.rs +++ b/aimux-providers/src/lib.rs @@ -173,6 +173,7 @@ pub use default_providers::{ }; pub mod preset; pub(crate) mod shared; +pub(crate) mod stream; pub mod vertex; pub mod voyage; diff --git a/aimux-providers/src/mistral/mod.rs b/aimux-providers/src/mistral/mod.rs index c8b4c7ed9..49ecce861 100644 --- a/aimux-providers/src/mistral/mod.rs +++ b/aimux-providers/src/mistral/mod.rs @@ -15,6 +15,8 @@ pub mod convert; pub mod embedding; mod model; +#[cfg(test)] +pub(crate) use model::chat_test_stream; pub(crate) mod options; mod speech; mod transcription; diff --git a/aimux-providers/src/mistral/model.rs b/aimux-providers/src/mistral/model.rs index b33a61dc0..e3ad9af4d 100644 --- a/aimux-providers/src/mistral/model.rs +++ b/aimux-providers/src/mistral/model.rs @@ -5,13 +5,13 @@ //! - `content` in responses can be a string or an array of typed parts //! (text, thinking, image_url). Thinking parts are extracted as reasoning //! in streaming mode. -//! - Streamed tool calls are assembled by the shared `StreamingToolCallTracker`. +//! - Streamed tool calls are assembled by the shared `ChatToolCalls`. //! - Usage supports `num_cached_tokens` / `prompt_tokens_details.cached_tokens`. //! - Finish reasons include `model_length`. use aimux_core::tool::RawToolCall; + use async_trait::async_trait; -use futures::StreamExt; use serde_json::Value; use std::sync::Arc; @@ -20,11 +20,12 @@ use aimux_core::language_model::LanguageModel; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, StreamResult}; use aimux_core::stream_part::StreamPart; -use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; +use aimux_core::types::{FinishReason, ResponseMetadata, Usage}; -use aimux_provider_utils::{StreamingToolCallDelta, StreamingToolCallTracker}; +use aimux_provider_utils::{Transformer, pipe_through}; use crate::shared::EndpointConfig; +use crate::stream::{ChatToolCalls, Controller, Emit, FINISH_OTHER, StreamContext, ToolCallDelta}; use super::convert::{build_request_body, parse_finish_reason, request_warnings}; use super::types::{ChatCompletionResponse, StreamChunk, UsageResponse}; @@ -247,10 +248,7 @@ impl LanguageModel for MistralModel { .finish_reason .as_deref() .map(parse_finish_reason) - .unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }); + .unwrap_or(FINISH_OTHER); let usage = convert_usage(&data.usage); @@ -286,208 +284,213 @@ impl LanguageModel for MistralModel { .await?; let response_headers = resp.response_headers; - let mut sse_stream = resp.value; - - let warnings = request_warnings(options, &self.model_id); - let include_raw_chunks = options.include_raw_chunks.unwrap_or(false); - - let stream_error_url = endpoint; - let stream_error_body = body.clone(); - let stream_response_headers = response_headers.clone(); - - let generate_id_fn = self.generate_id.clone(); - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); - - let text_id = 0usize; - let mut text_started = false; - let mut reasoning_started = false; - let tool_generate_id = generate_id_fn.clone(); - let mut tool_calls = StreamingToolCallTracker::new().with_generate_id(move || tool_generate_id()); - let mut tool_parts = Vec::new(); - let mut reasoning_id: Option = None; - let mut final_usage = Usage::default(); - let mut final_finish_reason: Option = None; - let mut response_metadata_emitted = false; - - while let Some(event) = sse_stream.next().await { - - match event { - Ok(parsed) => { - if include_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); - } - if let Some(err_obj) = parsed.get("error") { - yield Ok(StreamPart::Error { - error: super::mistral_stream_error( - err_obj, - &stream_error_url, - stream_error_body.clone(), - stream_response_headers.clone(), - ), - }); - continue; - } + let stream = pipe_through( + resp.value, + MistralChatStream::new(StreamContext { + generate_id: self.generate_id.clone(), + ..StreamContext::new( + options, + request_warnings(options, &self.model_id), + endpoint, + &body, + &response_headers, + ) + }), + ); - let chunk: StreamChunk = match serde_json::from_value(parsed) { - Ok(c) => c, - Err(e) => { - yield Ok(StreamPart::Error { error: e.into() }); - continue; - } - }; - - // Emit ResponseMetadata from the first valid chunk. - if !response_metadata_emitted { - response_metadata_emitted = true; - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: chunk.id.clone(), - timestamp: chunk - .created - .and_then(|secs| aimux_core::Timestamp::from_unix_seconds(secs as f64)), - model_id: chunk.model.clone(), - })); - } + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), + response: Some(aimux_core::shared::StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} - // Update usage. - if let Some(usage) = &chunk.usage { - final_usage = convert_usage(usage); - } +/// The id of the (single) text block of a streamed response. +const TEXT_ID: &str = "0"; + +/// The `TransformStream` of the Mistral chat model's `doStream`. +#[derive(Default)] +struct MistralChatStream { + /// Its id generator names reasoning blocks and id-less tool calls. + ctx: StreamContext, + tool_calls: ChatToolCalls, + text_started: bool, + /// The open reasoning block. + reasoning_id: Option, + final_usage: Usage, + final_finish_reason: Option, + response_metadata_emitted: bool, +} - // Upstream processes only the first choice. - if let Some(choice) = chunk.choices.into_iter().next() { - // Reasoning content (from thinking parts in array content). - for reasoning_delta in choice.delta.content.as_ref().and_then(Value::as_array).into_iter().flatten() - .filter_map(extract_reasoning_part) { - if !reasoning_started { - // End any active text before starting reasoning. - if text_started { - yield Ok(StreamPart::TextEnd { - id: format!("{text_id}"), - provider_metadata: None, - }); - text_started = false; - } - let rid = generate_id_fn(); - reasoning_id = Some(rid.clone()); - reasoning_started = true; - yield Ok(StreamPart::ReasoningStart { id: rid, - provider_metadata: None, - }); - } - if let Some(ref rid) = reasoning_id { - yield Ok(StreamPart::ReasoningDelta { - id: rid.clone(), - delta: reasoning_delta, - provider_metadata: None, - }); - } - } - - // Text content. - if let Some(text_delta) = - extract_text_content(&choice.delta.content) - && !text_delta.is_empty() { - if !text_started { - // End reasoning before starting text. - if reasoning_started { - if let Some(ref rid) = reasoning_id { - yield Ok(StreamPart::ReasoningEnd { - id: rid.clone(), - provider_metadata: None, +/// A chat stream with an empty request context and tool-call ids from +/// `generate_id`, for tests that drive it chunk by chunk. +#[cfg(test)] +pub(crate) fn chat_test_stream( + generate_id: Arc String + Send + Sync>, +) -> impl Transformer, Output = StreamPart> { + MistralChatStream::new(StreamContext { + generate_id, + ..StreamContext::default() + }) +} + +impl MistralChatStream { + fn new(ctx: StreamContext) -> Self { + Self { + tool_calls: ChatToolCalls::new(ctx.generate_id.clone()), + ctx, + ..Self::default() + } + } + + /// Handle one parsed SSE `data:` payload. + fn transform_chunk(&mut self, parsed: Value, controller: &mut Controller) { + self.ctx.raw(&parsed, controller); + + if let Some(err_obj) = parsed.get("error") { + controller.enqueue(StreamPart::Error { + error: super::mistral_stream_error( + err_obj, + &self.ctx.url, + self.ctx.request_body.clone(), + self.ctx.response_headers.clone(), + ), }); - } - reasoning_started = false; - reasoning_id = None; - } - yield Ok(StreamPart::TextStart { - id: format!("{text_id}"), - provider_metadata: None, - }); - text_started = true; - } - yield Ok(StreamPart::TextDelta { - id: format!("{text_id}"), - delta: text_delta, - provider_metadata: None, - }); - } - - // Tool calls: assembled by the shared tracker. - if let Some(tool_call_deltas) = choice.delta.tool_calls { - for dtc in &tool_call_deltas { - let delta = StreamingToolCallDelta { - index: dtc.index, - id: dtc.id.as_deref(), - name: dtc.function.name.as_deref(), - arguments: dtc.function.arguments.as_deref(), - ..Default::default() - }; - // A new call without a function name is - // invalid response data, as in the AI SDK; - // the stream ends. - if let Err(error) = tool_calls.process(delta, &mut tool_parts) { - yield Err(error.into()); - return; - } - for part in tool_parts.drain(..) { - yield Ok(part); - } - } - } - - // Finish reason. - if let Some(reason) = choice.finish_reason { - final_finish_reason = Some(parse_finish_reason(&reason)); - } - } - } - Err(error) => { - if include_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: Value::Null }); - } - yield Ok(StreamPart::Error { error }); - } + return; + } + + let chunk: StreamChunk = match serde_json::from_value(parsed) { + Ok(c) => c, + Err(e) => { + controller.enqueue(StreamPart::Error { error: e.into() }); + return; + } + }; + + // Emit ResponseMetadata from the first valid chunk. + if !self.response_metadata_emitted { + self.response_metadata_emitted = true; + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: chunk.id.clone(), + timestamp: chunk + .created + .and_then(|secs| aimux_core::Timestamp::from_unix_seconds(secs as f64)), + model_id: chunk.model.clone(), + })); + } + + // Update usage. + if let Some(usage) = &chunk.usage { + self.final_usage = convert_usage(usage); + } + + // Upstream processes only the first choice. + let Some(choice) = chunk.choices.into_iter().next() else { + return; + }; + + // Reasoning content (from thinking parts in array content). + for reasoning_delta in choice + .delta + .content + .as_ref() + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(extract_reasoning_part) + { + if self.reasoning_id.is_none() { + // End any active text before starting reasoning. + if self.text_started { + controller.text_end(TEXT_ID); + self.text_started = false; } + let rid = (self.ctx.generate_id)(); + self.reasoning_id = Some(rid.clone()); + controller.reasoning_start(rid); + } + if let Some(ref rid) = self.reasoning_id { + controller.reasoning_delta(rid.clone(), reasoning_delta); } + } - // Close any remaining open segments. - if text_started { - yield Ok(StreamPart::TextEnd { - id: format!("{text_id}"), - provider_metadata: None, - }); + // Text content. + if let Some(text_delta) = extract_text_content(&choice.delta.content) + && !text_delta.is_empty() + { + if !self.text_started { + // End reasoning before starting text. + if let Some(rid) = self.reasoning_id.take() { + controller.reasoning_end(rid); + } + controller.text_start(TEXT_ID); + self.text_started = true; } - if reasoning_started - && let Some(ref rid) = reasoning_id { - yield Ok(StreamPart::ReasoningEnd { - id: rid.clone(), - provider_metadata: None, - }); + controller.text_delta(TEXT_ID, text_delta); + } + + // Tool calls: assembled by the shared tracker. + if let Some(tool_call_deltas) = choice.delta.tool_calls { + for dtc in &tool_call_deltas { + let delta = ToolCallDelta { + index: dtc.index, + id: dtc.id.as_deref(), + name: dtc.function.name.as_deref(), + arguments: dtc.function.arguments.as_deref(), + ..ToolCallDelta::default() + }; + // A new call without a function name is invalid response + // data, as in the AI SDK; the stream ends. + if !self.tool_calls.delta(delta, controller) { + return; } + } + } + + // Finish reason. + if let Some(reason) = choice.finish_reason { + self.final_finish_reason = Some(parse_finish_reason(&reason)); + } + } +} + +impl Transformer for MistralChatStream { + type Input = Result; + type Output = StreamPart; - tool_calls.finish(&mut tool_parts); - for part in tool_parts.drain(..) { - yield Ok(part); + fn start(&mut self, controller: &mut Controller) { + self.ctx.start(controller); + } + + fn transform(&mut self, event: Result, controller: &mut Controller) { + match event { + Ok(parsed) => self.transform_chunk(parsed, controller), + Err(error) => { + self.ctx.raw(&Value::Null, controller); + controller.enqueue(StreamPart::Error { error }); } + } + } - yield Ok(StreamPart::Finish { - finish_reason: final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }), - usage: final_usage, - provider_metadata: None, - }); - }; + fn flush(mut self, controller: &mut Controller) { + // Close any remaining open segments. + if self.text_started { + controller.text_end(TEXT_ID); + } + if let Some(rid) = self.reasoning_id.take() { + controller.reasoning_end(rid); + } - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - response: Some(aimux_core::shared::StreamResponseInfo { - headers: Some(response_headers), - }), - }) + self.tool_calls.finish(controller); + + controller.enqueue(StreamPart::Finish { + finish_reason: self.final_finish_reason.unwrap_or(FINISH_OTHER), + usage: self.final_usage, + provider_metadata: None, + }); } } diff --git a/aimux-providers/src/open_responses.rs b/aimux-providers/src/open_responses.rs index 236ec1f91..4f285f390 100644 --- a/aimux-providers/src/open_responses.rs +++ b/aimux-providers/src/open_responses.rs @@ -33,7 +33,11 @@ use aimux_core::types::{ use aimux_core::embedding_model::EmbeddingModel; use aimux_core::image_model::ImageModel; use aimux_core::provider::Provider; -use aimux_provider_utils::{FetchFunction, HeaderMapOpt, Resolvable, validate_base_url}; +use aimux_provider_utils::{ + FetchFunction, HeaderMapOpt, Resolvable, Transformer, pipe_through, validate_base_url, +}; + +use crate::stream::{Controller, Emit, FINISH_OTHER}; use crate::shared::{ Credential, EndpointConfig, EndpointHeaders, ProviderHeaders, TransformRequestBody, @@ -488,319 +492,317 @@ impl LanguageModel for OpenResponsesModel { )); } - let stream = async_stream::stream! { - // First part: StreamStart. - yield Ok(StreamPart::StreamStart { warnings }); + let stream = pipe_through( + futures::stream::iter(first_event).chain(sse_stream), + OpenResponsesStream { + warnings, + final_usage: Usage::default(), + has_tool_calls: false, + finish_reason: FINISH_OTHER, + is_active_reasoning: false, + tool_calls: HashMap::new(), + }, + ); - let mut final_usage = Usage::default(); - let mut has_tool_calls = false; - let mut finish_reason = FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }; - let mut is_active_reasoning = false; - - // Tool-call accumulators keyed by item_id. - let mut tool_calls: HashMap = HashMap::new(); - - let mut event_iter = - futures::stream::iter(first_event.into_iter()).chain(sse_stream); - - while let Some(event) = event_iter.next().await { - match event { - Ok(chunk) => { - let chunk_type = chunk - .get("type") - .and_then(|t| t.as_str()) - .unwrap_or(""); - - match chunk_type { - // -- Tool call / reasoning / message item added -- - "response.output_item.added" => { - if let Some(item) = chunk.get("item") { - let item_type = item - .get("type") - .and_then(|t| t.as_str()) - .unwrap_or(""); - match item_type { - "function_call" => { - let id = item - .get("id") + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), + response: Some(aimux_core::shared::StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} + +/// The `TransformStream` of `OpenResponsesLanguageModel.doStream`. +struct OpenResponsesStream { + warnings: Vec, + final_usage: Usage, + has_tool_calls: bool, + finish_reason: FinishReason, + is_active_reasoning: bool, + /// Tool-call accumulators keyed by item_id. + tool_calls: HashMap, +} + +impl Transformer for OpenResponsesStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut Controller) { + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + } + + fn transform(&mut self, event: Result, controller: &mut Controller) { + let Self { + ref mut final_usage, + ref mut has_tool_calls, + ref mut finish_reason, + ref mut is_active_reasoning, + ref mut tool_calls, + .. + } = *self; + match event { + Ok(chunk) => { + let chunk_type = chunk.get("type").and_then(|t| t.as_str()).unwrap_or(""); + + match chunk_type { + // -- Tool call / reasoning / message item added -- + "response.output_item.added" => { + if let Some(item) = chunk.get("item") { + let item_type = item.get("type").and_then(|t| t.as_str()).unwrap_or(""); + match item_type { + "function_call" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + tool_calls.insert( + id, + ToolCallAccum { + tool_name: item + .get("name") .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - tool_calls.insert( - id, - ToolCallAccum { - tool_name: item - .get("name") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - tool_call_id: item - .get("call_id") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - arguments: item - .get("arguments") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - }, - ); - } - "reasoning" => { - let id = item - .get("id") + .map(std::string::ToString::to_string), + tool_call_id: item + .get("call_id") .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::ReasoningStart { - id, - provider_metadata: None, - }); - is_active_reasoning = true; - } - "message" => { - let id = item - .get("id") + .map(std::string::ToString::to_string), + arguments: item + .get("arguments") .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::TextStart { id, provider_metadata: None}); - } - _ => {} - } + .map(std::string::ToString::to_string), + }, + ); } + "reasoning" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.reasoning_start(id); + *is_active_reasoning = true; + } + "message" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.text_start(id); + } + _ => {} } - "response.function_call_arguments.delta" => { - let item_id = chunk - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let delta = chunk - .get("delta") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - tool_calls - .entry(item_id) - .and_modify(|tc| { - let existing = tc.arguments.take().unwrap_or_default(); - tc.arguments = Some(existing + &delta); - }) - .or_insert(ToolCallAccum { - tool_name: None, - tool_call_id: None, - arguments: Some(delta), - }); - } - "response.function_call_arguments.done" => { - let item_id = chunk - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let arguments = chunk - .get("arguments") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - tool_calls - .entry(item_id) - .and_modify(|tc| { - tc.arguments = Some(arguments.clone()); - }) - .or_insert(ToolCallAccum { - tool_name: None, - tool_call_id: None, - arguments: Some(arguments), - }); - } - "response.output_item.done" => { - if let Some(item) = chunk.get("item") { - let item_type = item - .get("type") - .and_then(|t| t.as_str()) - .unwrap_or(""); - match item_type { - "function_call" => { - let id = item - .get("id") + } + } + "response.function_call_arguments.delta" => { + let item_id = chunk + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let delta = chunk + .get("delta") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + tool_calls + .entry(item_id) + .and_modify(|tc| { + let existing = tc.arguments.take().unwrap_or_default(); + tc.arguments = Some(existing + &delta); + }) + .or_insert(ToolCallAccum { + tool_name: None, + tool_call_id: None, + arguments: Some(delta), + }); + } + "response.function_call_arguments.done" => { + let item_id = chunk + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let arguments = chunk + .get("arguments") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + tool_calls + .entry(item_id) + .and_modify(|tc| { + tc.arguments = Some(arguments.clone()); + }) + .or_insert(ToolCallAccum { + tool_name: None, + tool_call_id: None, + arguments: Some(arguments), + }); + } + "response.output_item.done" => { + if let Some(item) = chunk.get("item") { + let item_type = item.get("type").and_then(|t| t.as_str()).unwrap_or(""); + match item_type { + "function_call" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let accum = tool_calls.remove(&id); + let tool_name = accum + .as_ref() + .and_then(|a| a.tool_name.clone()) + .or_else(|| { + item.get("name") .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let accum = tool_calls.remove(&id); - let tool_name = accum - .as_ref() - .and_then(|a| a.tool_name.clone()) - .or_else(|| { - item.get("name") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - }) - .unwrap_or_default(); - let tool_call_id = accum - .as_ref() - .and_then(|a| a.tool_call_id.clone()) - .or_else(|| { - item.get("call_id") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - }) - .unwrap_or_default(); - let arguments = accum - .as_ref() - .and_then(|a| a.arguments.clone()) - .or_else(|| { - item.get("arguments") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - }) - .unwrap_or_default(); - let input = arguments; - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id, - tool_name, - input, - provider_executed: None, - dynamic: None, - provider_metadata: None, - })); - has_tool_calls = true; - } - "reasoning" => { - let id = item - .get("id") + .map(std::string::ToString::to_string) + }) + .unwrap_or_default(); + let tool_call_id = accum + .as_ref() + .and_then(|a| a.tool_call_id.clone()) + .or_else(|| { + item.get("call_id") .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::ReasoningEnd { - id, - provider_metadata: None, - }); - is_active_reasoning = false; - } - "message" => { - let id = item - .get("id") + .map(std::string::ToString::to_string) + }) + .unwrap_or_default(); + let arguments = accum + .as_ref() + .and_then(|a| a.arguments.clone()) + .or_else(|| { + item.get("arguments") .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::TextEnd { id, provider_metadata: None}); - } - _ => {} - } + .map(std::string::ToString::to_string) + }) + .unwrap_or_default(); + let input = arguments; + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id, + tool_name, + input, + provider_executed: None, + dynamic: None, + provider_metadata: None, + })); + *has_tool_calls = true; } + "reasoning" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.reasoning_end(id); + *is_active_reasoning = false; + } + "message" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.text_end(id); + } + _ => {} } + } + } - // -- Reasoning text delta (LM Studio extension) -- - "response.reasoning_text.delta" => { - let id = chunk - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let delta = chunk - .get("delta") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::ReasoningDelta { - id, - delta, - provider_metadata: None, - }); - } + // -- Reasoning text delta (LM Studio extension) -- + "response.reasoning_text.delta" => { + let id = chunk + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let delta = chunk + .get("delta") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.reasoning_delta(id, delta); + } - // -- Text delta -- - "response.output_text.delta" => { - let id = chunk - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let delta = chunk - .get("delta") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::TextDelta { id, delta, provider_metadata: None}); - } + // -- Text delta -- + "response.output_text.delta" => { + let id = chunk + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let delta = chunk + .get("delta") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.text_delta(id, delta); + } - // -- Completion events -- - "response.completed" | "response.incomplete" => { - if let Some(response) = chunk.get("response") { - let reason = response - .get("incomplete_details") - .and_then(|d| d.get("reason")) - .and_then(|r| r.as_str()); - finish_reason = FinishReason { - unified: map_open_responses_finish_reason( - reason, - has_tool_calls, - ), - raw: reason.map(std::string::ToString::to_string), - }; - if let Some(usage_val) = response.get("usage") { - final_usage = extract_usage_from_value(usage_val); - } - } - } - "response.failed" => { - if let Some(response) = chunk.get("response") { - let raw = response - .get("error") - .and_then(|e| e.get("code")) - .and_then(|c| c.as_str()) - .or_else(|| { - response.get("status").and_then(|s| s.as_str()) - }); - finish_reason = FinishReason { - unified: FinishReasonUnified::Error, - raw: raw.map(std::string::ToString::to_string), - }; - if let Some(usage_val) = response.get("usage") { - final_usage = extract_usage_from_value(usage_val); - } - } - } - _ => { - // Ignore unrecognised event types. + // -- Completion events -- + "response.completed" | "response.incomplete" => { + if let Some(response) = chunk.get("response") { + let reason = response + .get("incomplete_details") + .and_then(|d| d.get("reason")) + .and_then(|r| r.as_str()); + *finish_reason = FinishReason { + unified: map_open_responses_finish_reason(reason, *has_tool_calls), + raw: reason.map(std::string::ToString::to_string), + }; + if let Some(usage_val) = response.get("usage") { + *final_usage = extract_usage_from_value(usage_val); } } } - Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Err(error); - if !recoverable { - return; + "response.failed" => { + if let Some(response) = chunk.get("response") { + let raw = response + .get("error") + .and_then(|e| e.get("code")) + .and_then(|c| c.as_str()) + .or_else(|| response.get("status").and_then(|s| s.as_str())); + *finish_reason = FinishReason { + unified: FinishReasonUnified::Error, + raw: raw.map(std::string::ToString::to_string), + }; + if let Some(usage_val) = response.get("usage") { + *final_usage = extract_usage_from_value(usage_val); + } } } + _ => { + // Ignore unrecognised event types. + } } } - - // Flush: close any dangling reasoning segment. - if is_active_reasoning { - yield Ok(StreamPart::ReasoningEnd { - id: "reasoning-0".to_string(), - provider_metadata: None, - }); + Err(error) => { + if error.is_recoverable_stream_error() { + controller.enqueue_error(error); + } else { + controller.error(error); + } } + } + } - yield Ok(StreamPart::Finish { - finish_reason, - usage: final_usage, - provider_metadata: None, - }); - }; + fn flush(self, controller: &mut Controller) { + // Close any dangling reasoning segment. + if self.is_active_reasoning { + controller.reasoning_end("reasoning-0"); + } - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - response: Some(aimux_core::shared::StreamResponseInfo { - headers: Some(response_headers), - }), - }) + controller.enqueue(StreamPart::Finish { + finish_reason: self.finish_reason, + usage: self.final_usage, + provider_metadata: None, + }); } } diff --git a/aimux-providers/src/openai/model.rs b/aimux-providers/src/openai/model.rs index e02a127d8..23527ee08 100644 --- a/aimux-providers/src/openai/model.rs +++ b/aimux-providers/src/openai/model.rs @@ -8,6 +8,7 @@ //! the conversion + streaming logic while supplying their own URL and auth. use aimux_core::tool::RawToolCall; + use async_trait::async_trait; use futures::StreamExt; use serde_json::{Value, json}; @@ -18,10 +19,12 @@ use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, Source, StreamResult}; use aimux_core::shared::provider_namespace; use aimux_core::stream_part::StreamPart; -use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; +use aimux_core::types::{FinishReason, ResponseMetadata, Usage}; + +use aimux_provider_utils::{HttpRequest, Transformer, pipe_through}; -use aimux_provider_utils::{ - HttpRequest, StreamingToolCallDelta, StreamingToolCallTracker, TypeValidation, generate_id, +use crate::stream::{ + ChatToolCalls, Controller, Emit, FINISH_ERROR, FINISH_OTHER, StreamContext, ToolCallDelta, }; use super::config::{OpenAIModelConfig, TransformRequestBody}; @@ -316,10 +319,7 @@ pub(crate) async fn execute_generate( .finish_reason .as_deref() .map(parse_finish_reason) - .unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }); + .unwrap_or(FINISH_OTHER); let usage = data.usage.as_ref().map_or_else(Usage::default, |usage| { convert_usage(usage, response_value.get("usage")) @@ -423,311 +423,268 @@ pub(crate) async fn execute_stream( )); } - // M2 (RFC-0016): capture whether raw chunks should be emitted — the - // borrowed `options` cannot be moved into the generator. - let emit_raw_chunks = options.include_raw_chunks == Some(true); - let stream_error_url = endpoint; - let stream_error_body = body.clone(); - let stream_response_headers = response_headers.clone(); - - let stream = async_stream::stream! { - // First part: StreamStart. - yield Ok(StreamPart::StreamStart { warnings }); - - let text_id = 0usize; - let mut text_started = false; - let reasoning_id = "reasoning-0".to_string(); - let mut reasoning_started = false; - let mut final_usage = Usage::default(); - let mut final_usage_raw: Option = None; - let mut final_finish_reason: Option = None; - let mut response_metadata_emitted = false; - let mut final_logprobs: Option = None; - - // Streamed tool calls, correlated by wire id, index and function name - // (the AI SDK's StreamingToolCallTracker) and finalized on flush. - // Same options as `@ai-sdk/openai`'s chat model. - let mut tool_calls = StreamingToolCallTracker::new() - .with_generate_id(generate_id) - .with_type_validation(TypeValidation::IfPresent); - let mut tool_parts = Vec::new(); - - // Process the first event (already peeked) then the rest. - let mut event_iter = - futures::stream::iter(first_event.into_iter()).chain(sse_stream); - - let mut stream_errored = false; - - while let Some(event) = event_iter.next().await { - match event { - Ok(parsed) => { - - // M2 (RFC-0016): emit the raw provider chunk for debugging - // before it is consumed below. JSON payloads only — the - // "[DONE]" sentinel is skipped by the early break above. - if emit_raw_chunks { - yield Ok(StreamPart::Raw { - raw_value: parsed.clone(), - }); - } + let transformer = OpenAIChatStream { + ctx: StreamContext::new(options, warnings, endpoint, &body, &response_headers), + ..OpenAIChatStream::default() + }; + // Process the first event (already peeked) then the rest. + let events = futures::stream::iter(first_event).chain(sse_stream); + let stream = pipe_through(events, transformer); - // Check for mid-stream error. - if let Some(err_obj) = parsed.get("error") { - yield Ok(StreamPart::Error { - error: super::openai_stream_error( - err_obj, - &stream_error_url, - stream_error_body.clone(), - stream_response_headers.clone(), - ), - }); - stream_errored = true; - continue; - } + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), + response: Some(aimux_core::shared::StreamResponseInfo { + headers: Some(response_headers), + }), + }) +} - // The chunk's top-level `usage`, taken from the raw JSON before - // the chunk is consumed below. `chunk_usage_raw` keeps the - // provider's original object for `Usage.raw` (M10). - let chunk_usage_raw: Option = parsed.get("usage").cloned(); - let chunk_usage: Option = chunk_usage_raw - .as_ref() - .and_then(|u| serde_json::from_value(u.clone()).ok()); - - // Parse as StreamChunk. - let chunk: StreamChunk = match serde_json::from_value(parsed) { - Ok(c) => c, - Err(e) => { - stream_errored = true; - yield Ok(StreamPart::Error { error: e.into() }); - continue; - } - }; +const TEXT_ID: &str = "0"; +const REASONING_ID: &str = "reasoning-0"; + +/// The `TransformStream` of the Chat Completions stream ([`execute_stream`]): +/// turns the parsed SSE events into stream parts. +#[derive(Default)] +struct OpenAIChatStream { + ctx: StreamContext, + /// Finalized on flush. + tool_calls: ChatToolCalls, + text_started: bool, + reasoning_started: bool, + final_usage: Usage, + final_usage_raw: Option, + final_finish_reason: Option, + response_metadata_emitted: bool, + final_logprobs: Option, + stream_errored: bool, +} - // Emit ResponseMetadata from the first valid chunk. - if !response_metadata_emitted - && (chunk.id.as_ref().is_some_and(|id| !id.is_empty()) || chunk.model.as_ref().is_some_and(|model| !model.is_empty()) || chunk.created.is_some_and(|created| created != 0)) - { - response_metadata_emitted = true; - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: chunk.id.clone(), - timestamp: chunk - .created - .filter(|created| *created != 0) - .and_then(|secs| aimux_core::Timestamp::from_unix_seconds(secs as f64)), - model_id: chunk.model.clone(), - })); - } +/// A chat stream with an empty request context and tool-call ids from +/// `generate_id`, for tests that drive it chunk by chunk. +#[cfg(test)] +pub(crate) fn chat_test_stream( + generate_id: std::sync::Arc String + Send + Sync>, +) -> impl Transformer, Output = StreamPart> { + OpenAIChatStream { + tool_calls: ChatToolCalls::new(generate_id), + ..OpenAIChatStream::default() + } +} - // Update usage from the chunk that carries it. - if let Some(usage) = &chunk_usage { - final_usage = convert_usage(usage, chunk_usage_raw.as_ref()); - final_usage_raw = Some(usage.clone()); - } +impl OpenAIChatStream { + fn end_reasoning(&mut self, controller: &mut Controller) { + if self.reasoning_started { + controller.reasoning_end(REASONING_ID); + self.reasoning_started = false; + } + } +} - // Process choices. - for choice in chunk.choices { - // The last chunk that carries logprobs sets them. - if let Some(content) = choice.logprobs.as_ref().and_then(logprobs_content) { - final_logprobs = Some(content); - } - - // Reasoning delta: prefer reasoning_content over reasoning. - let reasoning_delta = choice - .delta - .reasoning_content - .clone() - .or_else(|| choice.delta.reasoning.clone()); - if let Some(reasoning) = reasoning_delta - && !reasoning.is_empty() - { - if !reasoning_started { - reasoning_started = true; - yield Ok(StreamPart::ReasoningStart { - id: reasoning_id.clone(), - provider_metadata: None, - }); - } - yield Ok(StreamPart::ReasoningDelta { - id: reasoning_id.clone(), - delta: reasoning, - provider_metadata: None, - }); - } - - // Text delta. - if let Some(content) = choice.delta.content { - // End active reasoning block before text starts. - if reasoning_started { - yield Ok(StreamPart::ReasoningEnd { - id: reasoning_id.clone(), - provider_metadata: None, - }); - reasoning_started = false; - } - if !text_started { - text_started = true; - yield Ok(StreamPart::TextStart { - id: format!("{text_id}"), - provider_metadata: None, - }); - } - yield Ok(StreamPart::TextDelta { - id: format!("{text_id}"), - delta: content, - provider_metadata: None, - }); - } - - // Tool-call deltas. - if let Some(tool_call_deltas) = choice.delta.tool_calls { - // End active reasoning block before tool calls start. - if reasoning_started { - yield Ok(StreamPart::ReasoningEnd { - id: reasoning_id.clone(), - provider_metadata: None, - }); - reasoning_started = false; - } - for dtc in &tool_call_deltas { - let function = dtc.function.as_ref(); - let delta = StreamingToolCallDelta { - index: Some(dtc.index), - id: dtc.id.as_deref(), - r#type: dtc.r#type.as_deref(), - name: function.and_then(|f| f.name.as_deref()), - arguments: function.and_then(|f| f.arguments.as_deref()), - provider_metadata: None, - }; - // A malformed delta (new call without a - // function name, or a non-`function` type) - // is invalid response data, as in the AI SDK; - // the stream ends. - if let Err(error) = tool_calls.process(delta, &mut tool_parts) { - yield Err(error.into()); - return; - } - for part in tool_parts.drain(..) { - yield Ok(part); - } - } - } - - // Annotations / citations (URL citations → Source). - if let Some(annotations) = choice.delta.annotations { - for (i, ann) in annotations.iter().enumerate() { - if ann.get("type").and_then(|v| v.as_str()) - == Some("url_citation") - && let Some(uc) = ann.get("url_citation") - { - yield Ok(StreamPart::Source(Source::Url { - id: format!("annotation-{i}"), - url: uc - .get("url") - .and_then(|v| v.as_str()) - .unwrap_or_default() - .to_string(), - title: uc - .get("title") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - provider_metadata: None, - })); - } - } - } - - // Finish reason. - if let Some(reason) = choice.finish_reason { - // Close any open reasoning segment. - if reasoning_started { - yield Ok(StreamPart::ReasoningEnd { - id: reasoning_id.clone(), - provider_metadata: None, +impl Transformer for OpenAIChatStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut Controller) { + self.ctx.start(controller); + } + + fn transform(&mut self, event: Result, controller: &mut Controller) { + let parsed = match event { + Ok(parsed) => parsed, + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + self.stream_errored = true; + controller.enqueue(StreamPart::Error { error }); + if !recoverable { + controller.terminate(); + } + return; + } + }; + + // M2 (RFC-0016): emit the raw provider chunk for debugging before it + // is consumed below. + self.ctx.raw(&parsed, controller); + + // Check for mid-stream error. + if let Some(err_obj) = parsed.get("error") { + controller.enqueue(StreamPart::Error { + error: super::openai_stream_error( + err_obj, + &self.ctx.url, + self.ctx.request_body.clone(), + self.ctx.response_headers.clone(), + ), }); - reasoning_started = false; - } + self.stream_errored = true; + return; + } - stream_errored = false; - final_finish_reason = Some(parse_finish_reason(&reason)); - } - } + // The chunk's top-level `usage`, taken from the raw JSON before the + // chunk is consumed below. `chunk_usage_raw` keeps the provider's + // original object for `Usage.raw` (M10). + let chunk_usage_raw: Option = parsed.get("usage").cloned(); + let chunk_usage: Option = chunk_usage_raw + .as_ref() + .and_then(|u| serde_json::from_value(u.clone()).ok()); + + let chunk: StreamChunk = match serde_json::from_value(parsed) { + Ok(c) => c, + Err(e) => { + self.stream_errored = true; + controller.enqueue(StreamPart::Error { error: e.into() }); + return; + } + }; + + // Emit ResponseMetadata from the first valid chunk. + if !self.response_metadata_emitted + && (chunk.id.as_ref().is_some_and(|id| !id.is_empty()) + || chunk.model.as_ref().is_some_and(|model| !model.is_empty()) + || chunk.created.is_some_and(|created| created != 0)) + { + self.response_metadata_emitted = true; + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: chunk.id.clone(), + timestamp: chunk + .created + .filter(|created| *created != 0) + .and_then(|secs| aimux_core::Timestamp::from_unix_seconds(secs as f64)), + model_id: chunk.model.clone(), + })); + } + + // Update usage from the chunk that carries it. + if let Some(usage) = &chunk_usage { + self.final_usage = convert_usage(usage, chunk_usage_raw.as_ref()); + self.final_usage_raw = Some(usage.clone()); + } + + for choice in chunk.choices { + // The last chunk that carries logprobs sets them. + if let Some(content) = choice.logprobs.as_ref().and_then(logprobs_content) { + self.final_logprobs = Some(content); + } + + // Reasoning delta: prefer reasoning_content over reasoning. + let reasoning_delta = choice + .delta + .reasoning_content + .clone() + .or_else(|| choice.delta.reasoning.clone()); + if let Some(reasoning) = reasoning_delta + && !reasoning.is_empty() + { + if !self.reasoning_started { + self.reasoning_started = true; + controller.reasoning_start(REASONING_ID); } - Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - stream_errored = true; - yield Ok(StreamPart::Error { error }); - if !recoverable { + controller.reasoning_delta(REASONING_ID, reasoning); + } + + // Text delta; an active reasoning block ends before text starts. + if let Some(content) = choice.delta.content { + self.end_reasoning(controller); + if !self.text_started { + self.text_started = true; + controller.text_start(TEXT_ID); + } + controller.text_delta(TEXT_ID, content); + } + + // Tool-call deltas; an active reasoning block ends first. + if let Some(tool_call_deltas) = choice.delta.tool_calls { + self.end_reasoning(controller); + for dtc in &tool_call_deltas { + let function = dtc.function.as_ref(); + let delta = ToolCallDelta { + index: Some(dtc.index), + id: dtc.id.as_deref(), + name: function.and_then(|f| f.name.as_deref()), + arguments: function.and_then(|f| f.arguments.as_deref()), + // `type` is optional, but must be 'function' when present. + invalid_type: dtc.r#type.as_deref().is_some_and(|t| t != "function"), + provider_metadata: None, + }; + // A malformed delta (new call without a function name, or + // a non-`function` type) is invalid response data, as in + // the AI SDK; the stream ends. + if !self.tool_calls.delta(delta, controller) { return; } } } - } - // Close any remaining open reasoning segment. - if reasoning_started { - yield Ok(StreamPart::ReasoningEnd { - id: reasoning_id.clone(), - provider_metadata: None, - }); - } + // Annotations / citations (URL citations → Source). + if let Some(annotations) = choice.delta.annotations { + for (i, ann) in annotations.iter().enumerate() { + if ann.get("type").and_then(|v| v.as_str()) == Some("url_citation") + && let Some(uc) = ann.get("url_citation") + { + controller.enqueue(StreamPart::Source(Source::Url { + id: format!("annotation-{i}"), + url: uc + .get("url") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string(), + title: uc + .get("title") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string), + provider_metadata: None, + })); + } + } + } - // Close any remaining open text segment. - if text_started { - yield Ok(StreamPart::TextEnd { - id: format!("{text_id}"), - provider_metadata: None, - }); + // Finish reason; closes any open reasoning segment. + if let Some(reason) = choice.finish_reason { + self.end_reasoning(controller); + self.stream_errored = false; + self.final_finish_reason = Some(parse_finish_reason(&reason)); + } } + } - // A parsable argument buffer can still be a prefix of a longer input: - // like the AI SDK's tracker, finalize only when the stream flushes. - tool_calls.finish(&mut tool_parts); - for part in tool_parts.drain(..) { - yield Ok(part); + fn flush(mut self, controller: &mut Controller) { + // Close any remaining open reasoning and text segments. + self.end_reasoning(controller); + if self.text_started { + controller.text_end(TEXT_ID); } - // Build provider metadata for the Finish part. + self.tool_calls.finish(controller); + + // Provider metadata for the Finish part: logprobs + prediction tokens. let mut pm_openai = serde_json::json!({}); - if let Some(ref logprobs) = final_logprobs { + if let Some(ref logprobs) = self.final_logprobs { pm_openai["logprobs"] = json!(logprobs); } - // Prediction tokens from raw usage. - if let Some(ref raw_usage) = final_usage_raw - && let Some(ref details) = raw_usage.completion_tokens_details { - if let Some(apt) = details.accepted_prediction_tokens { - pm_openai["acceptedPredictionTokens"] = json!(apt); - } - if let Some(rpt) = details.rejected_prediction_tokens { - pm_openai["rejectedPredictionTokens"] = json!(rpt); - } + if let Some(ref raw_usage) = self.final_usage_raw + && let Some(ref details) = raw_usage.completion_tokens_details + { + if let Some(apt) = details.accepted_prediction_tokens { + pm_openai["acceptedPredictionTokens"] = json!(apt); } - let provider_metadata = provider_namespace("openai", pm_openai).expect("provider metadata must be an object"); - - // Final part: Finish. - yield Ok(StreamPart::Finish { - finish_reason: if stream_errored { - FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - } + if let Some(rpt) = details.rejected_prediction_tokens { + pm_openai["rejectedPredictionTokens"] = json!(rpt); + } + } + let provider_metadata = + provider_namespace("openai", pm_openai).expect("provider metadata must be an object"); + + controller.enqueue(StreamPart::Finish { + finish_reason: if self.stream_errored { + FINISH_ERROR } else { - final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }) + self.final_finish_reason.unwrap_or(FINISH_OTHER) }, - usage: final_usage, + usage: self.final_usage, provider_metadata: Some(provider_metadata), }); - }; - - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - response: Some(aimux_core::shared::StreamResponseInfo { - headers: Some(response_headers), - }), - }) + } } // ── Model listing (RFC-0027) ───────────────────────────────────────────────── @@ -782,3 +739,121 @@ pub(crate) async fn list_models_once( }) .collect()) } + +/// The Chat Completions stream driven chunk by chunk. +#[cfg(test)] +mod stream_tests { + use serde_json::json; + + use super::*; + use crate::stream::testing::run; + + fn delta(delta: Value) -> Result { + Ok(json!({ "choices": [{ "delta": delta }] })) + } + + fn tool_calls(tool_calls: Value) -> Result { + delta(json!({ "tool_calls": tool_calls })) + } + + #[test] + fn interleaved_tool_calls_stream_as_they_arrive_and_end_on_flush_in_index_order() { + let parts = run( + OpenAIChatStream::default(), + [ + tool_calls(json!([{ + "index": 1, "id": "call_b", "type": "function", + "function": { "name": "time", "arguments": "{" } + }])), + tool_calls(json!([{ + "index": 0, "id": "call_a", "type": "function", + "function": { "name": "weather", "arguments": "{\"city\":" } + }])), + tool_calls(json!([ + { "index": 1, "function": { "arguments": "}" } }, + { "index": 0, "function": { "arguments": "\"Paris\"}" } } + ])), + Ok(json!({ "choices": [{ "delta": {}, "finish_reason": "tool_calls" }] })), + ], + ); + assert_eq!( + parts, + [ + "stream-start", + "tool-input-start call_b time", + "tool-input-delta call_b {", + "tool-input-start call_a weather", + "tool-input-delta call_a {\"city\":", + "tool-input-delta call_b }", + "tool-input-delta call_a \"Paris\"}", + "tool-input-end call_a", + "tool-call call_a weather {\"city\":\"Paris\"}", + "tool-input-end call_b", + "tool-call call_b time {}", + "finish ToolCalls", + ] + ); + } + + #[test] + fn an_error_frame_is_an_error_part_and_the_stream_finishes_with_error() { + let parts = run( + OpenAIChatStream::default(), + [ + delta(json!({ "content": "Hi" })), + Ok(json!({ "error": { "message": "overloaded", "type": "server_error" } })), + ], + ); + assert_eq!(parts.len(), 6, "{parts:?}"); + assert_eq!( + parts[..3], + ["stream-start", "text-start 0", "text-delta 0 Hi"] + ); + assert!(parts[3].starts_with("error ") && parts[3].contains("overloaded")); + assert_eq!(parts[4..], ["text-end 0", "finish Error"]); + } + + #[test] + fn flush_closes_open_blocks_and_finishes_with_other_without_a_finish_reason() { + let parts = run( + OpenAIChatStream::default(), + [ + delta(json!({ "reasoning_content": "Hmm" })), + delta(json!({ "content": "Hi" })), + ], + ); + assert_eq!( + parts, + [ + "stream-start", + "reasoning-start reasoning-0", + "reasoning-delta reasoning-0 Hmm", + "reasoning-end reasoning-0", + "text-start 0", + "text-delta 0 Hi", + "text-end 0", + "finish Other", + ] + ); + } + + #[test] + fn a_tool_call_of_another_type_ends_the_stream_without_flush() { + let parts = run( + OpenAIChatStream::default(), + [ + tool_calls(json!([{ + "index": 0, "id": "call_a", "type": "custom", + "function": { "name": "weather", "arguments": "{}" } + }])), + delta(json!({ "content": "ignored" })), + ], + ); + assert_eq!(parts.len(), 2, "{parts:?}"); + assert_eq!(parts[0], "stream-start"); + assert!( + parts[1].starts_with("err ") && parts[1].contains("Expected 'function' type."), + "{parts:?}" + ); + } +} diff --git a/aimux-providers/src/openai/responses/responses_convert.rs b/aimux-providers/src/openai/responses/responses_convert.rs index 492490ffc..f9e0de143 100644 --- a/aimux-providers/src/openai/responses/responses_convert.rs +++ b/aimux-providers/src/openai/responses/responses_convert.rs @@ -17,6 +17,9 @@ use aimux_core::tool::RawToolCall; use std::collections::{HashMap, HashSet, VecDeque}; use std::pin::Pin; +use aimux_provider_utils::{Transformer, pipe_through}; + +use crate::stream::{Controller, Emit, FINISH_ERROR, FINISH_OTHER, StreamContext}; use futures::{Stream, StreamExt}; use serde_json::{Value, json}; @@ -616,685 +619,922 @@ where } } - let stream = async_stream::stream! { - // First part: StreamStart. - yield Ok(StreamPart::StreamStart { warnings }); - - // Ongoing tool calls keyed by output_index. - let mut ongoing_tool_calls: HashMap = HashMap::new(); - // Active reasoning items keyed by item_id. - let mut active_reasoning: HashMap = HashMap::new(); + let shell_provider_executed = hosted_shell(&request_body); + let stream = pipe_through( + futures::stream::iter(first_event).chain(sse_stream), + OpenAIResponsesStream { + ctx: StreamContext { + warnings, + include_raw_chunks, + url: request_url, + request_body, + response_headers, + ..StreamContext::default() + }, + provider_key, + store_flag, + tool_names, + prompt_approval_tool_call_ids, + shell_provider_executed, + ..OpenAIResponsesStream::default() + }, + ); + Ok(Box::pin(stream)) +} - let mut has_function_call = false; - let mut final_usage: Option = None; - let mut final_raw_usage: Option = None; - let mut final_service_tier: Option = None; - let mut final_reasoning_context: Option = None; - let mut final_finish_reason: Option = None; - let mut response_id: Option = None; - let mut stream_errored = false; - let mut text_states: HashMap = HashMap::new(); - let mut output_ids: HashMap = HashMap::new(); - let mut text_annotations: HashMap> = HashMap::new(); - let mut logprobs = Vec::new(); - let shell_provider_executed = hosted_shell(&request_body); - let mut provider_items: HashMap = HashMap::new(); - let mut emitted_calls = HashSet::new(); - let mut ended_inputs = HashSet::new(); - let mut diff_received = HashSet::new(); - let mut hosted_tool_search_ids = VecDeque::new(); - let mut approval_tool_call_ids: HashMap = HashMap::new(); +/// The `TransformStream` of the OpenAI Responses `doStream`. +#[derive(Default)] +struct OpenAIResponsesStream { + ctx: StreamContext, + provider_key: String, + store_flag: bool, + tool_names: HashMap, + prompt_approval_tool_call_ids: HashMap, + shell_provider_executed: bool, + /// Ongoing tool calls keyed by output_index. + ongoing_tool_calls: HashMap, + /// Active reasoning items keyed by item_id. + active_reasoning: HashMap, + has_function_call: bool, + final_usage: Option, + final_raw_usage: Option, + final_service_tier: Option, + final_reasoning_context: Option, + final_finish_reason: Option, + response_id: Option, + stream_errored: bool, + text_states: HashMap, + output_ids: HashMap, + text_annotations: HashMap>, + logprobs: Vec, + provider_items: HashMap, + emitted_calls: HashSet, + ended_inputs: HashSet, + diff_received: HashSet, + hosted_tool_search_ids: VecDeque, + approval_tool_call_ids: HashMap, +} - let mut event_iter = - futures::stream::iter(first_event.into_iter()).chain(sse_stream); +impl Transformer for OpenAIResponsesStream { + type Input = Result; + type Output = StreamPart; - while let Some(event) = event_iter.next().await { - match event { - Ok(parsed) => { - if include_raw_chunks { yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); } - let etype = parsed - .get("type") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); + fn start(&mut self, controller: &mut Controller) { + self.ctx.start(controller); + } - match etype.as_str() { - // ── response.created → ResponseMetadata ─────────────────── - "response.created" => { - if let Some(resp_obj) = parsed.get("response") { - response_id = resp_obj - .get("id") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - let model_id = resp_obj - .get("model") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - let timestamp = resp_obj - .get("created_at") - .and_then(serde_json::Value::as_u64) - .and_then(|secs| aimux_core::Timestamp::from_unix_seconds(secs as f64)); - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: response_id.clone(), - timestamp, - model_id, - })); - } - } + fn transform(&mut self, event: Result, controller: &mut Controller) { + let Self { + ref provider_key, + store_flag, + ctx: + StreamContext { + ref url, + ref request_body, + ref response_headers, + include_raw_chunks, + .. + }, + ref tool_names, + ref prompt_approval_tool_call_ids, + shell_provider_executed, + ref mut ongoing_tool_calls, + ref mut active_reasoning, + ref mut has_function_call, + ref mut final_usage, + ref mut final_raw_usage, + ref mut final_service_tier, + ref mut final_reasoning_context, + ref mut final_finish_reason, + ref mut response_id, + ref mut stream_errored, + ref mut text_states, + ref mut output_ids, + ref mut text_annotations, + ref mut logprobs, + ref mut provider_items, + ref mut emitted_calls, + ref mut ended_inputs, + ref mut diff_received, + ref mut hosted_tool_search_ids, + ref mut approval_tool_call_ids, + .. + } = *self; + match event { + Ok(parsed) => { + if include_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: parsed.clone(), + }); + } + let etype = parsed + .get("type") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); - // ── response.output_item.added ───────────────────────────── - "response.output_item.added" => { - let output_index = parsed - .get("output_index") + match etype.as_str() { + // ── response.created → ResponseMetadata ─────────────────── + "response.created" => { + if let Some(resp_obj) = parsed.get("response") { + *response_id = resp_obj + .get("id") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + let model_id = resp_obj + .get("model") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + let timestamp = resp_obj + .get("created_at") .and_then(serde_json::Value::as_u64) - .unwrap_or(0) as usize; - if let Some(item) = parsed.get("item") { - let item_type = item - .get("type") - .and_then(|v| v.as_str()) - .unwrap_or(""); - match item_type { - + .and_then(|secs| { + aimux_core::Timestamp::from_unix_seconds(secs as f64) + }); + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: response_id.clone(), + timestamp, + model_id, + })); + } + } - "message" => { - let id = item - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - output_ids.insert(output_index, id.clone()); - text_states.insert(id.clone(), item.clone()); - yield Ok(StreamPart::TextStart { id, provider_metadata: Some(provider_namespace(&provider_key, text_metadata(item, None)).expect("provider metadata must be an object"))}); - } - "function_call" => { - let call_id = item - .get("call_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let name = item - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - ongoing_tool_calls.insert( - output_index, - OngoingToolCall { - tool_call_id: call_id.clone(), - }, - ); - yield Ok(StreamPart::ToolInputStart { - id: call_id, - tool_name: name, - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - } - "custom_tool_call" => { - let call_id = item - .get("call_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let name = item - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - ongoing_tool_calls.insert( - output_index, - OngoingToolCall { - tool_call_id: call_id.clone(), - }, - ); - yield Ok(StreamPart::ToolInputStart { - id: call_id, - tool_name: name, - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - } - "reasoning" => { - let id = item - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let encrypted = item - .get("encrypted_content") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - // Carry itemId + encrypted_content in - // provider_metadata (same shape as the - // non-streaming path) so reasoning can - // be echoed back on the next turn. - let meta = reasoning_stream_metadata( - &provider_key, - &id, - encrypted.as_deref(), - ); - active_reasoning.insert( - id.clone(), - ReasoningState { - encrypted_content: encrypted, - summary_parts: HashMap::from([( - 0usize, - SummaryStatus::Active, - )]), - }, - ); - yield Ok(StreamPart::ReasoningStart { - id: format!("{id}:0"), - provider_metadata: Some(meta), - }); - } - _ => { - provider_items.insert(output_index, item.clone()); - for part in provider_tool_start(item, &tool_names) { - if let StreamPart::ToolCall(call) = &part { emitted_calls.insert(call.tool_call_id.clone()); } - if let StreamPart::ToolInputEnd { id, .. } = &part { ended_inputs.insert(id.clone()); } - yield Ok(part); + // ── response.output_item.added ───────────────────────────── + "response.output_item.added" => { + let output_index = parsed + .get("output_index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) as usize; + if let Some(item) = parsed.get("item") { + let item_type = item.get("type").and_then(|v| v.as_str()).unwrap_or(""); + match item_type { + "message" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + output_ids.insert(output_index, id.clone()); + text_states.insert(id.clone(), item.clone()); + controller.enqueue(StreamPart::TextStart { + id, + provider_metadata: Some( + provider_namespace( + provider_key, + text_metadata(item, None), + ) + .expect("provider metadata must be an object"), + ), + }); + } + "function_call" => { + let call_id = item + .get("call_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let name = item + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + ongoing_tool_calls.insert( + output_index, + OngoingToolCall { + tool_call_id: call_id.clone(), + }, + ); + controller.tool_input_start(call_id, name); + } + "custom_tool_call" => { + let call_id = item + .get("call_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let name = item + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + ongoing_tool_calls.insert( + output_index, + OngoingToolCall { + tool_call_id: call_id.clone(), + }, + ); + controller.tool_input_start(call_id, name); + } + "reasoning" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let encrypted = item + .get("encrypted_content") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + // Carry itemId + encrypted_content in + // provider_metadata (same shape as the + // non-streaming path) so reasoning can + // be echoed back on the next turn. + let meta = reasoning_stream_metadata( + provider_key, + &id, + encrypted.as_deref(), + ); + active_reasoning.insert( + id.clone(), + ReasoningState { + encrypted_content: encrypted, + summary_parts: HashMap::from([( + 0usize, + SummaryStatus::Active, + )]), + }, + ); + controller.enqueue(StreamPart::ReasoningStart { + id: format!("{id}:0"), + provider_metadata: Some(meta), + }); + } + _ => { + provider_items.insert(output_index, item.clone()); + for part in provider_tool_start(item, tool_names) { + if let StreamPart::ToolCall(call) = &part { + emitted_calls.insert(call.tool_call_id.clone()); } + if let StreamPart::ToolInputEnd { id, .. } = &part { + ended_inputs.insert(id.clone()); + } + controller.enqueue(part); } } } } + } - "response.image_generation_call.partial_image" => { - yield Ok(StreamPart::ToolResult(aimux_core::tool::ToolResult { - tool_call_id: parsed["item_id"].as_str().unwrap_or_default().to_owned(), - tool_name: tool_names.get("image_generation").map_or("image_generation", String::as_str).to_owned(), - result: json!({"result": parsed["partial_image_b64"]}), - preliminary: Some(true), is_error: None, dynamic: None, provider_metadata: None, - })); - } - "response.code_interpreter_call_code.delta" | "response.apply_patch_call_operation_diff.delta" => { - let index = parsed["output_index"].as_u64().unwrap_or(0) as usize; - if let Some(item) = provider_items.get(&index) { - let id = item.get("call_id").filter(|v| !v.is_null()).unwrap_or(&item["id"]).as_str().unwrap_or_default().to_owned(); - diff_received.insert(id.clone()); - yield Ok(StreamPart::ToolInputDelta { id, delta: escaped_delta(parsed["delta"].as_str().unwrap_or_default()), provider_metadata: None }); - } + "response.image_generation_call.partial_image" => { + controller.enqueue(StreamPart::ToolResult(aimux_core::tool::ToolResult { + tool_call_id: parsed["item_id"].as_str().unwrap_or_default().to_owned(), + tool_name: tool_names + .get("image_generation") + .map_or("image_generation", String::as_str) + .to_owned(), + result: json!({"result": parsed["partial_image_b64"]}), + preliminary: Some(true), + is_error: None, + dynamic: None, + provider_metadata: None, + })); + } + "response.code_interpreter_call_code.delta" + | "response.apply_patch_call_operation_diff.delta" => { + let index = parsed["output_index"].as_u64().unwrap_or(0) as usize; + if let Some(item) = provider_items.get(&index) { + let id = item + .get("call_id") + .filter(|v| !v.is_null()) + .unwrap_or(&item["id"]) + .as_str() + .unwrap_or_default() + .to_owned(); + diff_received.insert(id.clone()); + controller.tool_input_delta( + id, + escaped_delta(parsed["delta"].as_str().unwrap_or_default()), + ); } - "response.code_interpreter_call_code.done" | "response.apply_patch_call_operation_diff.done" => { - let index = parsed["output_index"].as_u64().unwrap_or(0) as usize; - if let Some(item) = provider_items.get(&index) { - let id = item.get("call_id").filter(|v| !v.is_null()).unwrap_or(&item["id"]).as_str().unwrap_or_default().to_owned(); - if ended_inputs.insert(id.clone()) { - let code = etype == "response.code_interpreter_call_code.done"; - if !code && !diff_received.contains(&id) { - yield Ok(StreamPart::ToolInputDelta { id: id.clone(), delta: escaped_delta(parsed["diff"].as_str().unwrap_or_default()), provider_metadata: None }); - } - yield Ok(StreamPart::ToolInputDelta { id: id.clone(), delta: if code { "\"}" } else { "\"}}" }.to_owned(), provider_metadata: None }); - yield Ok(StreamPart::ToolInputEnd { id: id.clone(), provider_metadata: None }); - if code { - let mut item = item.clone(); item["code"] = parsed["code"].clone(); - for part in super::provider_events::provider_tool_content(&item, &provider_key, &tool_names, shell_provider_executed, None) { - if let GenerateContent::ToolCall(call) = part { emitted_calls.insert(id.clone()); yield Ok(StreamPart::ToolCall(call)); } + } + "response.code_interpreter_call_code.done" + | "response.apply_patch_call_operation_diff.done" => { + let index = parsed["output_index"].as_u64().unwrap_or(0) as usize; + if let Some(item) = provider_items.get(&index) { + let id = item + .get("call_id") + .filter(|v| !v.is_null()) + .unwrap_or(&item["id"]) + .as_str() + .unwrap_or_default() + .to_owned(); + if ended_inputs.insert(id.clone()) { + let code = etype == "response.code_interpreter_call_code.done"; + if !code && !diff_received.contains(&id) { + controller.tool_input_delta( + id.clone(), + escaped_delta(parsed["diff"].as_str().unwrap_or_default()), + ); + } + controller.tool_input_delta( + id.clone(), + if code { "\"}" } else { "\"}}" }.to_owned(), + ); + controller.tool_input_end(id.clone()); + if code { + let mut item = item.clone(); + item["code"] = parsed["code"].clone(); + for part in super::provider_events::provider_tool_content( + &item, + provider_key, + tool_names, + shell_provider_executed, + None, + ) { + if let GenerateContent::ToolCall(call) = part { + emitted_calls.insert(id.clone()); + controller.enqueue(StreamPart::ToolCall(call)); } } } } } + } - // ── response.output_text.delta → TextDelta ──────────────── - "response.output_text.delta" => { - if request_body.get("top_logprobs").and_then(Value::as_u64).is_some_and(|n| n > 0) && let Some(value) = parsed.get("logprobs") { logprobs.push(value.clone()); } - let mut id = parsed - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - if let Some(index) = parsed.get("output_index").and_then(Value::as_u64) && let Some(original_id) = output_ids.get(&(index as usize)) { id = original_id.clone(); } - let delta = parsed - .get("delta") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::TextDelta { id, delta, provider_metadata: None}); + // ── response.output_text.delta → TextDelta ──────────────── + "response.output_text.delta" => { + if request_body + .get("top_logprobs") + .and_then(Value::as_u64) + .is_some_and(|n| n > 0) + && let Some(value) = parsed.get("logprobs") + { + logprobs.push(value.clone()); } + let mut id = parsed + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + if let Some(index) = parsed.get("output_index").and_then(Value::as_u64) + && let Some(original_id) = output_ids.get(&(index as usize)) + { + id = original_id.clone(); + } + let delta = parsed + .get("delta") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.text_delta(id, delta); + } - // ── function_call_arguments.delta → ToolInputDelta ──────── - "response.function_call_arguments.delta" - | "response.custom_tool_call_input.delta" => { - let output_index = parsed - .get("output_index") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0) as usize; - let delta = parsed - .get("delta") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - if let Some(tc) = ongoing_tool_calls.get(&output_index) { - yield Ok(StreamPart::ToolInputDelta { - id: tc.tool_call_id.clone(), - delta, - provider_metadata: None, + // ── function_call_arguments.delta → ToolInputDelta ──────── + "response.function_call_arguments.delta" + | "response.custom_tool_call_input.delta" => { + let output_index = parsed + .get("output_index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) as usize; + let delta = parsed + .get("delta") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + if let Some(tc) = ongoing_tool_calls.get(&output_index) { + controller.tool_input_delta(tc.tool_call_id.clone(), delta); + } + } + + // ── reasoning_summary_part.added ────────────────────────── + "response.reasoning_summary_part.added" => { + let item_id = parsed + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let summary_index = parsed + .get("summary_index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) as usize; + + if summary_index > 0 + && let Some(state) = active_reasoning.get_mut(&item_id) + { + let meta = reasoning_stream_metadata( + provider_key, + &item_id, + state.encrypted_content.as_deref(), + ); + // Conclude all 'can-conclude' parts. + let to_conclude: Vec = state + .summary_parts + .iter() + .filter(|(_, s)| **s == SummaryStatus::CanConclude) + .map(|(k, _)| *k) + .collect(); + for idx in to_conclude { + state.summary_parts.insert(idx, SummaryStatus::Concluded); + controller.enqueue(StreamPart::ReasoningEnd { + id: format!("{item_id}:{idx}"), + provider_metadata: Some(meta.clone()), }); } + state + .summary_parts + .insert(summary_index, SummaryStatus::Active); + controller.enqueue(StreamPart::ReasoningStart { + id: format!("{item_id}:{summary_index}"), + provider_metadata: Some(meta), + }); } + } - // ── reasoning_summary_part.added ────────────────────────── - "response.reasoning_summary_part.added" => { - let item_id = parsed - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let summary_index = parsed - .get("summary_index") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0) as usize; + // ── reasoning_summary_text.delta → ReasoningDelta ───────── + "response.reasoning_summary_text.delta" => { + let item_id = parsed + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let summary_index = parsed + .get("summary_index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) as usize; + let delta = parsed + .get("delta") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + controller.reasoning_delta(format!("{item_id}:{summary_index}"), delta); + } - if summary_index > 0 - && let Some(state) = active_reasoning.get_mut(&item_id) - { - let meta = reasoning_stream_metadata( - &provider_key, - &item_id, - state.encrypted_content.as_deref(), - ); - // Conclude all 'can-conclude' parts. - let to_conclude: Vec = state - .summary_parts - .iter() - .filter(|(_, s)| **s == SummaryStatus::CanConclude) - .map(|(k, _)| *k) - .collect(); - for idx in to_conclude { - state.summary_parts.insert(idx, SummaryStatus::Concluded); - yield Ok(StreamPart::ReasoningEnd { - id: format!("{item_id}:{idx}"), - provider_metadata: Some(meta.clone()), - }); - } + // ── reasoning_summary_part.done ─────────────────────────── + "response.reasoning_summary_part.done" => { + let item_id = parsed + .get("item_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let summary_index = parsed + .get("summary_index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) as usize; + if let Some(state) = active_reasoning.get_mut(&item_id) { + if store_flag { state .summary_parts - .insert(summary_index, SummaryStatus::Active); - yield Ok(StreamPart::ReasoningStart { + .insert(summary_index, SummaryStatus::Concluded); + controller.enqueue(StreamPart::ReasoningEnd { id: format!("{item_id}:{summary_index}"), - provider_metadata: Some(meta), + provider_metadata: Some(reasoning_stream_metadata( + provider_key, + &item_id, + state.encrypted_content.as_deref(), + )), }); + } else { + state + .summary_parts + .insert(summary_index, SummaryStatus::CanConclude); } } + } - // ── reasoning_summary_text.delta → ReasoningDelta ───────── - "response.reasoning_summary_text.delta" => { - let item_id = parsed - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let summary_index = parsed - .get("summary_index") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0) as usize; - let delta = parsed - .get("delta") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - yield Ok(StreamPart::ReasoningDelta { - id: format!("{item_id}:{summary_index}"), - delta, - provider_metadata: None, - }); - } - - // ── reasoning_summary_part.done ─────────────────────────── - "response.reasoning_summary_part.done" => { - let item_id = parsed - .get("item_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let summary_index = parsed - .get("summary_index") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0) as usize; - if let Some(state) = active_reasoning.get_mut(&item_id) { - if store_flag { - state - .summary_parts - .insert(summary_index, SummaryStatus::Concluded); - yield Ok(StreamPart::ReasoningEnd { - id: format!("{item_id}:{summary_index}"), - provider_metadata: Some(reasoning_stream_metadata( - &provider_key, - &item_id, - state.encrypted_content.as_deref(), - )), + // ── response.output_item.done ───────────────────────────── + "response.output_item.done" => { + let output_index = parsed + .get("output_index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) as usize; + if let Some(item) = parsed.get("item") { + let item_type = item.get("type").and_then(|v| v.as_str()).unwrap_or(""); + match item_type { + "message" => { + let mut id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + if let Some(original_id) = output_ids.remove(&output_index) { + id = original_id; + } + let mut metadata_item = item.clone(); + metadata_item["id"] = json!(id); + if metadata_item.get("phase").is_none() + && let Some(phase) = + text_states.get(&id).and_then(|v| v.get("phase")) + { + metadata_item["phase"] = phase.clone(); + } + let annotations = text_annotations.remove(&id); + controller.enqueue(StreamPart::TextEnd { + id, + provider_metadata: Some( + provider_namespace( + provider_key, + text_metadata(&metadata_item, annotations.as_ref()), + ) + .expect("provider metadata must be an object"), + ), }); - } else { - state - .summary_parts - .insert(summary_index, SummaryStatus::CanConclude); } - } - } - - // ── response.output_item.done ───────────────────────────── - "response.output_item.done" => { - let output_index = parsed - .get("output_index") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0) as usize; - if let Some(item) = parsed.get("item") { - let item_type = item - .get("type") - .and_then(|v| v.as_str()) - .unwrap_or(""); - match item_type { - "message" => { - let mut id = item - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - if let Some(original_id) = output_ids.remove(&output_index) { id = original_id; } - let mut metadata_item = item.clone(); - metadata_item["id"] = json!(id); - if metadata_item.get("phase").is_none() && let Some(phase) = text_states.get(&id).and_then(|v| v.get("phase")) { - metadata_item["phase"] = phase.clone(); + "function_call" => { + *has_function_call = true; + ongoing_tool_calls.remove(&output_index); + let call_id = item + .get("call_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let name = item + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let arguments = item + .get("arguments") + .and_then(|v| v.as_str()) + .unwrap_or("{}") + .to_string(); + controller.tool_input_end(call_id.clone()); + let input = arguments; + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: call_id, + tool_name: name, + input, + provider_executed: None, + dynamic: None, + provider_metadata: Some( + provider_namespace(provider_key, tool_metadata(item)) + .expect("provider metadata must be an object"), + ), + })); + } + "custom_tool_call" => { + *has_function_call = true; + ongoing_tool_calls.remove(&output_index); + let call_id = item + .get("call_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let name = item + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let input_str = + item.get("input").and_then(|v| v.as_str()).unwrap_or("{}"); + controller.tool_input_end(call_id.clone()); + let input = serde_json::to_string(input_str).expect( + "serializing a custom-tool input string cannot fail", + ); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: call_id, + tool_name: name, + input, + provider_executed: None, + dynamic: None, + provider_metadata: Some( + provider_namespace(provider_key, tool_metadata(item)) + .expect("provider metadata must be an object"), + ), + })); + } + "reasoning" => { + let id = item + .get("id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + if let Some(state) = active_reasoning.get_mut(&id) { + state.encrypted_content = item + .get("encrypted_content") + .and_then(Value::as_str) + .map(str::to_owned); + let meta = reasoning_stream_metadata( + provider_key, + &id, + state.encrypted_content.as_deref(), + ); + // Conclude all active / can-conclude parts. + let to_conclude: Vec = state + .summary_parts + .iter() + .filter(|(_, s)| { + **s == SummaryStatus::Active + || **s == SummaryStatus::CanConclude + }) + .map(|(k, _)| *k) + .collect(); + for idx in to_conclude { + state + .summary_parts + .insert(idx, SummaryStatus::Concluded); + controller.enqueue(StreamPart::ReasoningEnd { + id: format!("{id}:{idx}"), + provider_metadata: Some(meta.clone()), + }); } - let annotations = text_annotations.remove(&id); - yield Ok(StreamPart::TextEnd { id, provider_metadata: Some(provider_namespace(&provider_key, text_metadata(&metadata_item, annotations.as_ref())).expect("provider metadata must be an object"))}); - } - "function_call" => { - has_function_call = true; - ongoing_tool_calls.remove(&output_index); - let call_id = item - .get("call_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let name = item - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let arguments = item - .get("arguments") - .and_then(|v| v.as_str()) - .unwrap_or("{}") - .to_string(); - yield Ok(StreamPart::ToolInputEnd { - id: call_id.clone(), - provider_metadata: None, - }); - let input = arguments; - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: call_id, - tool_name: name, - input, - provider_executed: None, - dynamic: None, - provider_metadata: Some(provider_namespace(&provider_key, tool_metadata(item)).expect("provider metadata must be an object")), - })); } - "custom_tool_call" => { - has_function_call = true; - ongoing_tool_calls.remove(&output_index); - let call_id = item - .get("call_id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let name = item - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let input_str = item - .get("input") - .and_then(|v| v.as_str()) - .unwrap_or("{}"); - yield Ok(StreamPart::ToolInputEnd { - id: call_id.clone(), - provider_metadata: None, - }); - let input = - serde_json::to_string(input_str).expect( - "serializing a custom-tool input string cannot fail", - ); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: call_id, - tool_name: name, - input, - provider_executed: None, - dynamic: None, - provider_metadata: Some(provider_namespace(&provider_key, tool_metadata(item)).expect("provider metadata must be an object")), - })); - } - "reasoning" => { - let id = item - .get("id") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - if let Some(state) = active_reasoning.get_mut(&id) { - state.encrypted_content = item.get("encrypted_content").and_then(Value::as_str).map(str::to_owned); - let meta = reasoning_stream_metadata( - &provider_key, - &id, - state.encrypted_content.as_deref(), - ); - // Conclude all active / can-conclude parts. - let to_conclude: Vec = state - .summary_parts - .iter() - .filter(|(_, s)| { - **s == SummaryStatus::Active - || **s == SummaryStatus::CanConclude - }) - .map(|(k, _)| *k) - .collect(); - for idx in to_conclude { - state - .summary_parts - .insert(idx, SummaryStatus::Concluded); - yield Ok(StreamPart::ReasoningEnd { - id: format!("{id}:{idx}"), - provider_metadata: Some(meta.clone()), - }); + active_reasoning.remove(&id); + } + _ => { + let mut item = item.clone(); + if item_type == "local_shell_call" + && let Some(action) = item["action"].as_object_mut() + { + for (from, to) in [ + ("timeout_ms", "timeoutMs"), + ("working_directory", "workingDirectory"), + ] { + if let Some(value) = action.remove(from) { + action.insert(to.into(), value); } } - active_reasoning.remove(&id); } - _ => { - let mut item = item.clone(); - if item_type == "local_shell_call" && let Some(action) = item["action"].as_object_mut() { - for (from, to) in [("timeout_ms", "timeoutMs"), ("working_directory", "workingDirectory")] { - if let Some(value) = action.remove(from) { action.insert(to.into(), value); } - } - } - if item_type == "apply_patch_call" { - let id = item["call_id"].as_str().unwrap_or_default().to_owned(); - if ended_inputs.insert(id.clone()) && item["operation"]["type"] != "delete_file" { - if !diff_received.contains(&id) { - yield Ok(StreamPart::ToolInputDelta { id: id.clone(), delta: escaped_delta(item["operation"]["diff"].as_str().unwrap_or_default()), provider_metadata: None }); - } - yield Ok(StreamPart::ToolInputDelta { id: id.clone(), delta: "\"}}".into(), provider_metadata: None }); - yield Ok(StreamPart::ToolInputEnd { id, provider_metadata: None }); + if item_type == "apply_patch_call" { + let id = + item["call_id"].as_str().unwrap_or_default().to_owned(); + if ended_inputs.insert(id.clone()) + && item["operation"]["type"] != "delete_file" + { + if !diff_received.contains(&id) { + controller.tool_input_delta( + id.clone(), + escaped_delta( + item["operation"]["diff"] + .as_str() + .unwrap_or_default(), + ), + ); } + controller.tool_input_delta(id.clone(), "\"}}"); + controller.tool_input_end(id); } - if item_type == "tool_search_call" && item["execution"] == "server" { - item["call_id"] = Value::Null; - } - pair_tool_search_id(&mut item, &mut hosted_tool_search_ids); - let mcp_tool_call_id = item.get("approval_request_id").and_then(Value::as_str) - .and_then(|id| approval_tool_call_ids.get(id).or_else(|| prompt_approval_tool_call_ids.get(id))).cloned(); - for part in super::provider_events::provider_tool_content(&item, &provider_key, &tool_names, shell_provider_executed, mcp_tool_call_id.as_deref()) { - match part { - GenerateContent::ToolCall(call) => { - if item_type == "code_interpreter_call" || (item_type == "apply_patch_call" && item["status"] != "completed") { continue; } - if emitted_calls.insert(call.tool_call_id.clone()) || item_type == "mcp_call" { - has_function_call |= call.provider_executed != Some(true) && matches!(item_type, "computer_call" | "apply_patch_call"); - if matches!(item_type, "computer_call" | "apply_patch_call" | "tool_search_call") { - if item_type == "tool_search_call" && item["execution"] != "server" { - yield Ok(StreamPart::ToolInputStart { id: call.tool_call_id.clone(), tool_name: call.tool_name.clone(), provider_executed: call.provider_executed, dynamic: None, title: None, provider_metadata: None }); - } - if !ended_inputs.contains(&call.tool_call_id) { - if item_type == "computer_call" && !call.input.is_empty() { yield Ok(StreamPart::ToolInputDelta { id: call.tool_call_id.clone(), delta: call.input.clone(), provider_metadata: None }); } - yield Ok(StreamPart::ToolInputEnd { id: call.tool_call_id.clone(), provider_metadata: None }); + } + if item_type == "tool_search_call" + && item["execution"] == "server" + { + item["call_id"] = Value::Null; + } + pair_tool_search_id(&mut item, hosted_tool_search_ids); + let mcp_tool_call_id = item + .get("approval_request_id") + .and_then(Value::as_str) + .and_then(|id| { + approval_tool_call_ids + .get(id) + .or_else(|| prompt_approval_tool_call_ids.get(id)) + }) + .cloned(); + for part in super::provider_events::provider_tool_content( + &item, + provider_key, + tool_names, + shell_provider_executed, + mcp_tool_call_id.as_deref(), + ) { + match part { + GenerateContent::ToolCall(call) => { + if item_type == "code_interpreter_call" + || (item_type == "apply_patch_call" + && item["status"] != "completed") + { + continue; + } + if emitted_calls.insert(call.tool_call_id.clone()) + || item_type == "mcp_call" + { + *has_function_call |= call.provider_executed + != Some(true) + && matches!( + item_type, + "computer_call" | "apply_patch_call" + ); + if matches!( + item_type, + "computer_call" + | "apply_patch_call" + | "tool_search_call" + ) { + if item_type == "tool_search_call" + && item["execution"] != "server" + { + controller.enqueue( + StreamPart::ToolInputStart { + id: call.tool_call_id.clone(), + tool_name: call + .tool_name + .clone(), + provider_executed: call + .provider_executed, + dynamic: None, + title: None, + provider_metadata: None, + }, + ); + } + if !ended_inputs + .contains(&call.tool_call_id) + { + if item_type == "computer_call" + && !call.input.is_empty() + { + controller.enqueue( + StreamPart::ToolInputDelta { + id: call + .tool_call_id + .clone(), + delta: call.input.clone(), + provider_metadata: None, + }, + ); } + controller.enqueue( + StreamPart::ToolInputEnd { + id: call.tool_call_id.clone(), + provider_metadata: None, + }, + ); } - yield Ok(StreamPart::ToolCall(call)); } + controller.enqueue(StreamPart::ToolCall(call)); } - GenerateContent::ToolResult(result) => { yield Ok(StreamPart::ToolResult(result)); } - GenerateContent::ToolApprovalRequest(approval) => { - approval_tool_call_ids.insert(approval.approval_id.clone(), approval.tool_call_id.clone()); - yield Ok(StreamPart::ToolApprovalRequest(approval)); - } - GenerateContent::Custom { kind, provider_metadata } => { - yield Ok(StreamPart::Custom { kind, provider_metadata }); - } - _ => {} } + GenerateContent::ToolResult(result) => { + controller.enqueue(StreamPart::ToolResult(result)); + } + GenerateContent::ToolApprovalRequest(approval) => { + approval_tool_call_ids.insert( + approval.approval_id.clone(), + approval.tool_call_id.clone(), + ); + controller.enqueue( + StreamPart::ToolApprovalRequest(approval), + ); + } + GenerateContent::Custom { + kind, + provider_metadata, + } => { + controller.enqueue(StreamPart::Custom { + kind, + provider_metadata, + }); + } + _ => {} } - provider_items.remove(&output_index); } + provider_items.remove(&output_index); } } } + } - // ── response.output_text.annotation.added → Source ───────── - "response.output_text.annotation.added" => { - match annotation_source(parsed.get("annotation").unwrap_or(&Value::Null), &provider_key) { - Ok(source) => { - let id = parsed.get("item_id").and_then(Value::as_str).unwrap_or_default().to_string(); - text_annotations.entry(id).or_default().push(parsed["annotation"].clone()); - yield Ok(StreamPart::Source(source)); - } - Err(error) => { - stream_errored = true; - yield Ok(StreamPart::Error { error }); - } + // ── response.output_text.annotation.added → Source ───────── + "response.output_text.annotation.added" => { + match annotation_source( + parsed.get("annotation").unwrap_or(&Value::Null), + provider_key, + ) { + Ok(source) => { + let id = parsed + .get("item_id") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(); + text_annotations + .entry(id) + .or_default() + .push(parsed["annotation"].clone()); + controller.enqueue(StreamPart::Source(source)); + } + Err(error) => { + *stream_errored = true; + controller.enqueue(StreamPart::Error { error }); } } + } - // ── response.completed / response.incomplete → finish ──── - "response.completed" | "response.incomplete" => { - if let Some(resp_obj) = parsed.get("response") { - let reason = resp_obj - .get("incomplete_details") - .and_then(|d| d.get("reason")) - .and_then(|v| v.as_str()); - if !stream_errored { - final_finish_reason = Some(map_responses_finish_reason(reason, has_function_call)); - } - final_usage = - resp_obj.get("usage").and_then(parse_usage); - // Keep the raw wire usage object for `usage.raw` - // (RFC-0015 P0-3) and capture the finish-time - // provider metadata fields carried on the - // terminal response object. - final_raw_usage = resp_obj.get("usage").cloned(); - if let Some(st) = resp_obj - .get("service_tier") - .and_then(|v| v.as_str()) - { - final_service_tier = Some(st.to_string()); - } - if let Some(reasoning) = resp_obj.get("reasoning") - && let Some(ctx) = reasoning.get("context") - && !ctx.is_null() - { - final_reasoning_context = Some(ctx.clone()); - } + // ── response.completed / response.incomplete → finish ──── + "response.completed" | "response.incomplete" => { + if let Some(resp_obj) = parsed.get("response") { + let reason = resp_obj + .get("incomplete_details") + .and_then(|d| d.get("reason")) + .and_then(|v| v.as_str()); + if !*stream_errored { + *final_finish_reason = + Some(map_responses_finish_reason(reason, *has_function_call)); + } + *final_usage = resp_obj.get("usage").and_then(parse_usage); + // Keep the raw wire usage object for `usage.raw` + // (RFC-0015 P0-3) and capture the finish-time + // provider metadata fields carried on the + // terminal response object. + *final_raw_usage = resp_obj.get("usage").cloned(); + if let Some(st) = resp_obj.get("service_tier").and_then(|v| v.as_str()) + { + *final_service_tier = Some(st.to_string()); + } + if let Some(reasoning) = resp_obj.get("reasoning") + && let Some(ctx) = reasoning.get("context") + && !ctx.is_null() + { + *final_reasoning_context = Some(ctx.clone()); } } + } - // ── response.failed ─────────────────────────────────────── - "response.failed" => { - if let Some(resp_obj) = parsed.get("response") { - let reason = resp_obj - .get("incomplete_details") - .and_then(|d| d.get("reason")) - .and_then(|v| v.as_str()); - final_finish_reason = Some(match reason { - Some(r) => map_responses_finish_reason(Some(r), has_function_call), - None => FinishReason { - unified: FinishReasonUnified::Error, - raw: Some("error".to_string()), - }, + // ── response.failed ─────────────────────────────────────── + "response.failed" => { + if let Some(resp_obj) = parsed.get("response") { + let reason = resp_obj + .get("incomplete_details") + .and_then(|d| d.get("reason")) + .and_then(|v| v.as_str()); + *final_finish_reason = Some(match reason { + Some(r) => map_responses_finish_reason(Some(r), *has_function_call), + None => FinishReason { + unified: FinishReasonUnified::Error, + raw: Some("error".to_string()), + }, + }); + *final_usage = resp_obj.get("usage").and_then(parse_usage); + *final_raw_usage = resp_obj.get("usage").cloned(); + // Upstream's `response.failed` arm carries + // `reasoningContext` (but not `serviceTier`). + if let Some(reasoning) = resp_obj.get("reasoning") + && let Some(ctx) = reasoning.get("context") + && !ctx.is_null() + { + *final_reasoning_context = Some(ctx.clone()); + } + if !*stream_errored && resp_obj.get("error").is_some() { + *stream_errored = true; + controller.enqueue(StreamPart::Error { + error: super::super::openai_stream_error( + &parsed, + url, + request_body.clone(), + response_headers.clone(), + ), }); - final_usage = - resp_obj.get("usage").and_then(parse_usage); - final_raw_usage = resp_obj.get("usage").cloned(); - // Upstream's `response.failed` arm carries - // `reasoningContext` (but not `serviceTier`). - if let Some(reasoning) = resp_obj.get("reasoning") - && let Some(ctx) = reasoning.get("context") - && !ctx.is_null() - { - final_reasoning_context = Some(ctx.clone()); - } - if !stream_errored - && resp_obj.get("error").is_some() - { - stream_errored = true; - yield Ok(StreamPart::Error { error: super::super::openai_stream_error(&parsed, &request_url, request_body.clone(), response_headers.clone()) }); - - } } } + } - // ── error chunk ─────────────────────────────────────────── - "error" => { - stream_errored = true; - final_finish_reason = Some(FinishReason { - unified: FinishReasonUnified::Error, - raw: Some("error".to_string()), - }); - yield Ok(StreamPart::Error { error: super::super::openai_stream_error(&parsed, &request_url, request_body.clone(), response_headers.clone()) }); - } + // ── error chunk ─────────────────────────────────────────── + "error" => { + *stream_errored = true; + *final_finish_reason = Some(FinishReason { + unified: FinishReasonUnified::Error, + raw: Some("error".to_string()), + }); + controller.enqueue(StreamPart::Error { + error: super::super::openai_stream_error( + &parsed, + url, + request_body.clone(), + response_headers.clone(), + ), + }); + } - _ => { - // Unknown / unhandled chunk types (content_part.added, - // output_text.done, content_part.done, etc.) are - // ignored — the core streaming path is driven by - // output_item.added/done and the delta events. - } + _ => { + // Unknown / unhandled chunk types (content_part.added, + // output_text.done, content_part.done, etc.) are + // ignored — the core streaming path is driven by + // output_item.added/done and the delta events. } } - Err(error) => { - stream_errored = true; - final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); - yield Ok(StreamPart::Error { error }); - } + } + Err(error) => { + *stream_errored = true; + *final_finish_reason = Some(FINISH_ERROR); + controller.enqueue(StreamPart::Error { error }); } } + } + fn flush(self, controller: &mut Controller) { // Build provider metadata for the Finish part: { : { // responseId, serviceTier?, reasoningContext? } }. Mirrors the TS // flush() providerMetadata, which carries `service_tier` and // `reasoning.context` from the terminal response object. - let mut pm = json!({ "responseId": response_id.unwrap_or_default() }); - if let Some(st) = final_service_tier { + let mut pm = json!({ "responseId": self.response_id.unwrap_or_default() }); + if let Some(st) = self.final_service_tier { pm["serviceTier"] = json!(st); } - if let Some(ctx) = final_reasoning_context { + if let Some(ctx) = self.final_reasoning_context { pm["reasoningContext"] = ctx; } - if !logprobs.is_empty() { pm["logprobs"] = json!(logprobs); } - let provider_metadata = Some(provider_namespace(&provider_key, pm).expect("provider metadata must be an object")); + if !self.logprobs.is_empty() { + pm["logprobs"] = json!(self.logprobs); + } + let provider_metadata = Some( + provider_namespace(&self.provider_key, pm) + .expect("provider metadata must be an object"), + ); - yield Ok(StreamPart::Finish { - finish_reason: final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, raw: None, - }), - usage: convert_responses_usage(final_usage.as_ref(), final_raw_usage), + controller.enqueue(StreamPart::Finish { + finish_reason: self.final_finish_reason.unwrap_or(FINISH_OTHER), + usage: convert_responses_usage(self.final_usage.as_ref(), self.final_raw_usage), provider_metadata, }); - }; - Ok(Box::pin(stream)) + } } fn pair_tool_search_id(item: &mut Value, hosted_ids: &mut VecDeque) { diff --git a/aimux-providers/src/openai_compatible/chat.rs b/aimux-providers/src/openai_compatible/chat.rs index adcd06f15..00d0399a4 100644 --- a/aimux-providers/src/openai_compatible/chat.rs +++ b/aimux-providers/src/openai_compatible/chat.rs @@ -8,23 +8,25 @@ use aimux_core::tool::RawToolCall; use std::collections::{HashMap, HashSet}; use async_trait::async_trait; -use futures::StreamExt; use serde_json::{Map, Value, json}; use aimux_core::error::AiMuxError; use aimux_core::language_model::{LanguageModel, SupportedUrls}; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, ReasoningOutput, Source, StreamResult}; -use aimux_core::shared::{RequestInfo, ResponseInfo, StreamResponseInfo, provider_namespace}; -use aimux_core::stream_part::StreamPart; -use aimux_core::types::{ - FinishReason, FinishReasonUnified, InputTokenUsage, OutputTokenUsage, ResponseMetadata, Usage, +use aimux_core::shared::{ + RequestInfo, ResponseInfo, SharedProviderMetadata, StreamResponseInfo, provider_namespace, }; -use aimux_provider_utils::{StreamingToolCallDelta, StreamingToolCallTracker, generate_id}; +use aimux_core::stream_part::StreamPart; +use aimux_core::types::{FinishReason, InputTokenUsage, OutputTokenUsage, ResponseMetadata, Usage}; +use aimux_provider_utils::{Transformer, generate_id, pipe_through}; -use super::config::CompatModelConfig; +use super::config::{ChatSettings, CompatModelConfig, StreamMetadataExtractor}; use super::convert::{ChatBodySpec, RequestBodyResult, build_request_body, parse_finish_reason}; use super::types::{ChatCompletionResponse, StreamChunk, UsageResponse}; +use crate::stream::{ + ChatToolCalls, Controller, Emit, FINISH_ERROR, FINISH_OTHER, StreamContext, ToolCallDelta, +}; /// An OpenAI-compatible chat model. Holds no HTTP client: the request helpers /// of `aimux-provider-utils` use the injected or process-default transport. @@ -271,10 +273,7 @@ impl LanguageModel for OpenAICompatibleChatModel { .finish_reason .as_deref() .map(parse_finish_reason) - .unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }); + .unwrap_or(FINISH_OTHER); let chat = &self.config.chat; let usage = chat.convert_usage(raw.get("usage")); @@ -336,309 +335,325 @@ impl LanguageModel for OpenAICompatibleChatModel { ) .await?; let response_headers = resp.response_headers; - let mut sse_stream = resp.value; - - let chat = self.config.chat.clone(); - let mut metadata_extractor = chat - .metadata_extractor - .as_ref() - .map(|extractor| extractor.create_stream_extractor()); - let emit_raw_chunks = options.include_raw_chunks == Some(true); - let stream_error_url = endpoint; - let stream_error_body = body.clone(); - let stream_response_headers = response_headers.clone(); - - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); - - let text_id = "txt-0"; - let mut text_started = false; - let reasoning_id = "reasoning-0".to_string(); - let mut reasoning_started = false; - let mut final_usage_raw = None; - let mut final_usage_parsed: Option = None; - let mut final_finish_reason: Option = None; - let mut response_metadata_emitted = false; - - let mut tool_calls = StreamingToolCallTracker::new().with_generate_id(generate_id); - let mut tool_parts = Vec::new(); - let signature_metadata = |extra: Option<&Value>| { - thought_signature(extra).map(|signature| { - provider_namespace(&metadata_key, json!({ "thoughtSignature": signature })) - .expect("provider metadata must be an object") - }) - }; + let sse_stream = resp.value; + + let transformer = OpenAICompatibleChatStream { + ctx: StreamContext::new(options, warnings, endpoint, &body, &response_headers), + metadata_key, + metadata_extractor: self + .config + .chat + .metadata_extractor + .as_ref() + .map(|extractor| extractor.create_stream_extractor()), + chat: self.config.chat.clone(), + ..OpenAICompatibleChatStream::default() + }; + let stream = pipe_through(sse_stream, transformer); - // Some compatible servers send the first delta of a call without - // `function.name`; buffer by index until the name is known. - let mut pending: HashMap = HashMap::new(); - let mut forwarded: HashSet = HashSet::new(); + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(RequestInfo { body: Some(body) }), + response: Some(StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} - while let Some(event) = sse_stream.next().await { - match event { - Ok(parsed) => { - if emit_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); - } - if let Some(extractor) = metadata_extractor.as_mut() { - extractor.process_chunk(&parsed); - } - if let Some(error) = parsed.get("error") { - yield Ok(StreamPart::Error { - error: stream_error( - &metadata_key, - error, - &stream_error_url, - stream_error_body.clone(), - stream_response_headers.clone(), - ), - }); - final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); - continue; - } +const TEXT_ID: &str = "txt-0"; +const REASONING_ID: &str = "reasoning-0"; - let chunk_usage_raw: Option = match &chat.stream_usage_key { - Some(key) => parsed.get(key).and_then(|v| v.get("usage")).cloned(), - None => parsed.get("usage").cloned(), - } - .filter(|usage| !usage.is_null()); - - let chunk: StreamChunk = match serde_json::from_value(parsed) { - Ok(chunk) => chunk, - Err(e) => { - final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); - yield Ok(StreamPart::Error { error: e.into() }); - continue; - } - }; - - if !response_metadata_emitted - && (chunk.id.as_ref().is_some_and(|id| !id.is_empty()) - || chunk.model.as_ref().is_some_and(|model| !model.is_empty()) - || chunk.created.is_some_and(|created| created != 0)) - { - response_metadata_emitted = true; - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: chunk.id.clone(), - timestamp: chunk - .created - .filter(|created| *created != 0) - .and_then(|secs| aimux_core::Timestamp::from_unix_seconds(secs as f64)), - model_id: chunk.model.clone(), - })); - } +/// The `TransformStream` of [`OpenAICompatibleChatModel::do_stream`]: turns +/// the parsed SSE events into stream parts. +#[derive(Default)] +struct OpenAICompatibleChatStream { + ctx: StreamContext, + metadata_key: String, + /// The chat settings: usage accounting and the streaming usage key. + chat: ChatSettings, + /// The caller's extractor, fed every parsed chunk. + metadata_extractor: Option>, + + text_started: bool, + reasoning_started: bool, + final_usage_raw: Option, + final_usage_parsed: Option, + final_finish_reason: Option, + response_metadata_emitted: bool, + tool_calls: ChatToolCalls, + /// Some compatible servers send the first delta of a call without + /// `function.name`; buffer by index until the name is known. + pending: HashMap, + forwarded: HashSet, +} - if let Some(raw_usage) = &chunk_usage_raw { - final_usage_raw = Some(raw_usage.clone()); - final_usage_parsed = serde_json::from_value(raw_usage.clone()).ok(); - } +/// A chat stream with an empty request context and tool-call ids from +/// `generate_id`, for tests that drive it chunk by chunk. +/// `metadata_key` is the provider options key. +#[cfg(test)] +pub(crate) fn chat_test_stream( + metadata_key: &str, + generate_id: std::sync::Arc String + Send + Sync>, +) -> impl Transformer, Output = StreamPart> { + OpenAICompatibleChatStream { + metadata_key: metadata_key.to_string(), + tool_calls: ChatToolCalls::new(generate_id), + ..OpenAICompatibleChatStream::default() + } +} - for choice in chunk.choices.into_iter().take(1) { - let reasoning_delta = choice - .delta - .reasoning_content - .clone() - .or_else(|| choice.delta.reasoning.clone()); - let mut delta_content = Vec::new(); - if let Some(text) = reasoning_delta.filter(|text| !text.is_empty()) { - delta_content.push(GenerateContent::Reasoning(ReasoningOutput { - text, - provider_metadata: None, - })); - } - delta_content.extend(content_parts(choice.delta.content.as_ref())); - for part in delta_content { - match part { - GenerateContent::Reasoning(ReasoningOutput { text, .. }) => { - if text_started { - yield Ok(StreamPart::TextEnd { id: text_id.to_string(), provider_metadata: None }); - text_started = false; - } - if !reasoning_started { - reasoning_started = true; - yield Ok(StreamPart::ReasoningStart { id: reasoning_id.clone(), provider_metadata: None }); - } - yield Ok(StreamPart::ReasoningDelta { id: reasoning_id.clone(), delta: text, provider_metadata: None }); - } - GenerateContent::Text { text, .. } => { - if reasoning_started { - yield Ok(StreamPart::ReasoningEnd { id: reasoning_id.clone(), provider_metadata: None }); - reasoning_started = false; - } - if !text_started { - text_started = true; - yield Ok(StreamPart::TextStart { id: text_id.to_string(), provider_metadata: None }); - } - yield Ok(StreamPart::TextDelta { id: text_id.to_string(), delta: text, provider_metadata: None }); - } - _ => unreachable!(), - } - } +impl OpenAICompatibleChatStream { + fn signature_metadata(&self, extra: Option<&Value>) -> Option { + thought_signature(extra).map(|signature| { + provider_namespace(&self.metadata_key, json!({ "thoughtSignature": signature })) + .expect("provider metadata must be an object") + }) + } +} - if let Some(deltas) = choice.delta.tool_calls.filter(|deltas| !deltas.is_empty()) { - if reasoning_started { - yield Ok(StreamPart::ReasoningEnd { - id: reasoning_id.clone(), - provider_metadata: None, - }); - reasoning_started = false; - } - for dtc in &deltas { - let function = dtc.function.as_ref(); - let name = function - .and_then(|f| f.name.as_deref()) - .filter(|n| !n.trim().is_empty()); - let arguments = function.and_then(|f| f.arguments.as_deref()); - let buffered; - let delta = match dtc.index { - Some(index) if !forwarded.contains(&index) => { - let entry = pending.entry(index).or_default(); - if entry.id.is_none() { - entry.id.clone_from(&dtc.id); - } - if entry.extra.is_null() - && let Some(extra) = &dtc.extra_content - { - entry.extra = extra.clone(); - } - if let Some(arguments) = arguments { - entry.arguments.push_str(arguments); - } - let Some(name) = name else { continue }; - buffered = pending.remove(&index).unwrap_or_default(); - forwarded.insert(index); - StreamingToolCallDelta { - index: Some(index), - id: buffered.id.as_deref(), - name: Some(name), - arguments: Some(&buffered.arguments), - provider_metadata: signature_metadata(Some(&buffered.extra)), - ..Default::default() - } - } - _ => StreamingToolCallDelta { - index: dtc.index, - id: dtc.id.as_deref(), - name: function.and_then(|f| f.name.as_deref()), - arguments, - provider_metadata: signature_metadata(dtc.extra_content.as_ref()), - ..Default::default() - }, - }; - if let Err(error) = tool_calls.process(delta, &mut tool_parts) { - yield Err(error.into()); - return; - } - for part in tool_parts.drain(..) { - yield Ok(part); - } - } - } +impl Transformer for OpenAICompatibleChatStream { + type Input = Result; + type Output = StreamPart; - if let Some(annotations) = choice.delta.annotations { - for (i, annotation) in annotations.iter().enumerate() { - if annotation.get("type").and_then(Value::as_str) - == Some("url_citation") - && let Some(citation) = annotation.get("url_citation") - { - yield Ok(StreamPart::Source(Source::Url { - id: format!("annotation-{i}"), - url: citation - .get("url") - .and_then(Value::as_str) - .unwrap_or_default() - .to_string(), - title: citation - .get("title") - .and_then(Value::as_str) - .map(str::to_string), - provider_metadata: None, - })); - } - } - } + fn start(&mut self, controller: &mut Controller) { + self.ctx.start(controller); + } - if let Some(reason) = choice.finish_reason { - final_finish_reason = Some(parse_finish_reason(&reason)); - } + fn transform(&mut self, event: Result, controller: &mut Controller) { + let parsed = match event { + Ok(parsed) => parsed, + Err(error) => { + self.final_finish_reason = Some(FINISH_ERROR); + controller.enqueue(StreamPart::Error { error }); + return; + } + }; + self.ctx.raw(&parsed, controller); + if let Some(extractor) = self.metadata_extractor.as_mut() { + extractor.process_chunk(&parsed); + } + if let Some(error) = parsed.get("error") { + controller.enqueue(StreamPart::Error { + error: stream_error( + &self.metadata_key, + error, + &self.ctx.url, + self.ctx.request_body.clone(), + self.ctx.response_headers.clone(), + ), + }); + self.final_finish_reason = Some(FINISH_ERROR); + return; + } + + let chunk_usage_raw: Option = match &self.chat.stream_usage_key { + Some(key) => parsed.get(key).and_then(|v| v.get("usage")).cloned(), + None => parsed.get("usage").cloned(), + } + .filter(|usage| !usage.is_null()); + + let chunk: StreamChunk = match serde_json::from_value(parsed) { + Ok(chunk) => chunk, + Err(e) => { + self.final_finish_reason = Some(FINISH_ERROR); + controller.enqueue(StreamPart::Error { error: e.into() }); + return; + } + }; + + if !self.response_metadata_emitted + && (chunk.id.as_ref().is_some_and(|id| !id.is_empty()) + || chunk.model.as_ref().is_some_and(|model| !model.is_empty()) + || chunk.created.is_some_and(|created| created != 0)) + { + self.response_metadata_emitted = true; + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: chunk.id.clone(), + timestamp: chunk + .created + .filter(|created| *created != 0) + .and_then(|secs| aimux_core::Timestamp::from_unix_seconds(secs as f64)), + model_id: chunk.model.clone(), + })); + } + + if let Some(raw_usage) = &chunk_usage_raw { + self.final_usage_raw = Some(raw_usage.clone()); + self.final_usage_parsed = serde_json::from_value(raw_usage.clone()).ok(); + } + + for choice in chunk.choices.into_iter().take(1) { + let reasoning_delta = choice + .delta + .reasoning_content + .clone() + .or_else(|| choice.delta.reasoning.clone()); + let mut delta_content = Vec::new(); + if let Some(text) = reasoning_delta.filter(|text| !text.is_empty()) { + delta_content.push(GenerateContent::Reasoning(ReasoningOutput { + text, + provider_metadata: None, + })); + } + delta_content.extend(content_parts(choice.delta.content.as_ref())); + for part in delta_content { + match part { + GenerateContent::Reasoning(ReasoningOutput { text, .. }) => { + if self.text_started { + controller.text_end(TEXT_ID); + self.text_started = false; + } + if !self.reasoning_started { + self.reasoning_started = true; + controller.reasoning_start(REASONING_ID); } + controller.reasoning_delta(REASONING_ID, text); } - Err(error) => { - final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::Error, raw: None }); - yield Ok(StreamPart::Error { error }); + GenerateContent::Text { text, .. } => { + if self.reasoning_started { + controller.reasoning_end(REASONING_ID); + self.reasoning_started = false; + } + if !self.text_started { + self.text_started = true; + controller.text_start(TEXT_ID); + } + controller.text_delta(TEXT_ID, text); } + _ => unreachable!(), } } - if reasoning_started { - yield Ok(StreamPart::ReasoningEnd { - id: reasoning_id.clone(), - provider_metadata: None, - }); - } - if text_started { - yield Ok(StreamPart::TextEnd { - id: text_id.to_string(), - provider_metadata: None, - }); - } - for (index, entry) in pending { - let delta = StreamingToolCallDelta { - index: Some(index), - id: entry.id.as_deref(), - arguments: Some(&entry.arguments), - provider_metadata: signature_metadata(Some(&entry.extra)), - ..Default::default() - }; - if let Err(error) = tool_calls.process(delta, &mut tool_parts) { - yield Err(error.into()); - return; + if let Some(deltas) = choice.delta.tool_calls.filter(|deltas| !deltas.is_empty()) { + if self.reasoning_started { + controller.reasoning_end(REASONING_ID); + self.reasoning_started = false; } - for part in tool_parts.drain(..) { - yield Ok(part); + for dtc in &deltas { + let function = dtc.function.as_ref(); + let name = function + .and_then(|f| f.name.as_deref()) + .filter(|n| !n.trim().is_empty()); + let arguments = function.and_then(|f| f.arguments.as_deref()); + let buffered; + let delta = match dtc.index { + Some(index) if !self.forwarded.contains(&index) => { + let entry = self.pending.entry(index).or_default(); + if entry.id.is_none() { + entry.id.clone_from(&dtc.id); + } + if entry.extra.is_null() + && let Some(extra) = &dtc.extra_content + { + entry.extra = extra.clone(); + } + if let Some(arguments) = arguments { + entry.arguments.push_str(arguments); + } + let Some(name) = name else { continue }; + buffered = self.pending.remove(&index).unwrap_or_default(); + self.forwarded.insert(index); + ToolCallDelta { + index: Some(index), + id: buffered.id.as_deref(), + name: Some(name), + arguments: Some(&buffered.arguments), + provider_metadata: self.signature_metadata(Some(&buffered.extra)), + ..ToolCallDelta::default() + } + } + _ => ToolCallDelta { + index: dtc.index, + id: dtc.id.as_deref(), + name: function.and_then(|f| f.name.as_deref()), + arguments, + provider_metadata: self.signature_metadata(dtc.extra_content.as_ref()), + ..ToolCallDelta::default() + }, + }; + if !self.tool_calls.delta(delta, controller) { + return; + } } } - // A parsable argument buffer can still be a prefix of a longer - // input: finalize only when the stream flushes. - tool_calls.finish(&mut tool_parts); - for part in tool_parts.drain(..) { - yield Ok(part); - } - let mut metadata = aimux_core::shared::SharedProviderMetadata::new(); - metadata.insert(metadata_key.clone(), Map::new()); - if let Some(extractor) = &metadata_extractor - && let Some(extra) = extractor.build_metadata() - { - for (key, value) in extra { - metadata.insert(key, value); + if let Some(annotations) = choice.delta.annotations { + for (i, annotation) in annotations.iter().enumerate() { + if annotation.get("type").and_then(Value::as_str) == Some("url_citation") + && let Some(citation) = annotation.get("url_citation") + { + controller.enqueue(StreamPart::Source(Source::Url { + id: format!("annotation-{i}"), + url: citation + .get("url") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(), + title: citation + .get("title") + .and_then(Value::as_str) + .map(str::to_string), + provider_metadata: None, + })); + } } } - prediction_tokens(final_usage_parsed.as_ref(), metadata.entry(metadata_key).or_default()); + if let Some(reason) = choice.finish_reason { + self.final_finish_reason = Some(parse_finish_reason(&reason)); + } + } + } + + fn flush(mut self, controller: &mut Controller) { + if self.reasoning_started { + controller.reasoning_end(REASONING_ID); + } + if self.text_started { + controller.text_end(TEXT_ID); + } + for (index, entry) in std::mem::take(&mut self.pending) { + let delta = ToolCallDelta { + index: Some(index), + id: entry.id.as_deref(), + arguments: Some(&entry.arguments), + provider_metadata: self.signature_metadata(Some(&entry.extra)), + ..ToolCallDelta::default() + }; + if !self.tool_calls.delta(delta, controller) { + return; + } + } + self.tool_calls.finish(controller); - if final_finish_reason.is_none() { - yield Ok(StreamPart::Error { - error: AiMuxError::InvalidResponseData("Response stream ended without a finish reason.".to_string()), - }); + let mut metadata = aimux_core::shared::SharedProviderMetadata::new(); + metadata.insert(self.metadata_key.clone(), Map::new()); + if let Some(extractor) = &self.metadata_extractor + && let Some(extra) = extractor.build_metadata() + { + for (key, value) in extra { + metadata.insert(key, value); } - yield Ok(StreamPart::Finish { - finish_reason: final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - }), - usage: chat.convert_usage(final_usage_raw.as_ref()), - provider_metadata: Some(metadata), - }); - }; + } - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(RequestInfo { body: Some(body) }), - response: Some(StreamResponseInfo { - headers: Some(response_headers), - }), - }) + prediction_tokens( + self.final_usage_parsed.as_ref(), + metadata.entry(self.metadata_key).or_default(), + ); + + if self.final_finish_reason.is_none() { + controller.enqueue(StreamPart::Error { + error: AiMuxError::InvalidResponseData( + "Response stream ended without a finish reason.".to_string(), + ), + }); + } + controller.enqueue(StreamPart::Finish { + finish_reason: self.final_finish_reason.unwrap_or(FINISH_ERROR), + usage: self.chat.convert_usage(self.final_usage_raw.as_ref()), + provider_metadata: Some(metadata), + }); } } diff --git a/aimux-providers/src/openai_compatible/completion.rs b/aimux-providers/src/openai_compatible/completion.rs index 9cc8228c2..94bc61d34 100644 --- a/aimux-providers/src/openai_compatible/completion.rs +++ b/aimux-providers/src/openai_compatible/completion.rs @@ -14,10 +14,11 @@ use aimux_core::result::{GenerateContent, GenerateResult}; use aimux_core::shared::{RequestInfo, ResponseInfo, StreamResponseInfo, provider_namespace}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{ - FinishReason, FinishReasonUnified, InputTokenUsage, OutputTokenUsage, ResponseMetadata, Usage, - Warning, + FinishReason, InputTokenUsage, OutputTokenUsage, ResponseMetadata, Usage, Warning, }; -use aimux_provider_utils::HttpRequest; +use aimux_provider_utils::{HttpRequest, Transformer, pipe_through}; + +use crate::stream::{Controller, Emit, FINISH_ERROR, FINISH_OTHER, StreamContext}; use super::config::CompatModelConfig; use super::convert::{parse_finish_reason, to_camel_case}; @@ -422,10 +423,7 @@ impl LanguageModel for OpenAICompatibleCompletionModel { .finish_reason .as_deref() .map(parse_finish_reason) - .unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }); + .unwrap_or(FINISH_OTHER); Ok(GenerateResult { content: if self.native() || !choice.text.is_empty() { vec![GenerateContent::Text { @@ -472,9 +470,7 @@ impl LanguageModel for OpenAICompatibleCompletionModel { .await?; let mut events = response.value; let response_headers = response.response_headers; - let stream_response_headers = response_headers.clone(); let native = self.native(); - let raw_chunks = options.include_raw_chunks == Some(true); let mut buffered = Vec::new(); if native { while let Some(event) = events.next().await { @@ -508,62 +504,132 @@ impl LanguageModel for OpenAICompatibleCompletionModel { } } } - let stream_body = body.clone(); let compat_name = match &self.config { CompletionConfig::Compatible(config) => Some(config.provider_options_name().to_owned()), CompletionConfig::Native(_) => None, }; - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); - let mut events = futures::stream::iter(buffered).chain(events); - let mut started = false; - let mut finish_reason = FinishReason { unified: FinishReasonUnified::Other, raw: None }; - let mut final_usage = None; - let mut logprobs = None; - while let Some(event) = events.next().await { - let event = match event { - Ok(event) => event, - Err(error) => { finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; yield Ok(StreamPart::Error { error }); continue; } - }; - if raw_chunks { yield Ok(StreamPart::Raw { raw_value: event.clone() }); } - if let Some(error) = event.get("error") { - finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; - // @ai-sdk/openai maps the frame to a status - // (openai-completion-language-model.ts:305-313); - // @ai-sdk/openai-compatible passes it on as it is - // (openai-compatible-completion-language-model.ts:320-325). - let error = match &compat_name { - None => crate::openai::openai_stream_error(error, &error_url, stream_body.clone(), response_headers.clone()), - Some(name) => super::chat::stream_error(name, error, &error_url, stream_body.clone(), response_headers.clone()), - }; - yield Ok(StreamPart::Error { error }); - continue; - } - let parsed: CompletionResponse = match parse_response(event, native, true) { - Ok(parsed) => parsed, - Err(error) => { finish_reason = FinishReason { unified: FinishReasonUnified::Error, raw: None }; yield Ok(StreamPart::Error { error }); continue; } - }; - if !started { - started = true; - yield Ok(StreamPart::ResponseMetadata(metadata(&parsed))); - yield Ok(StreamPart::TextStart { id: "0".into(), provider_metadata: None }); - } - if parsed.usage.is_some() { final_usage = parsed.usage; } - if let Some(choice) = parsed.choices.into_iter().next() { - if let Some(reason) = choice.finish_reason { finish_reason = parse_finish_reason(&reason); } - if native && choice.logprobs.is_some() { logprobs = choice.logprobs; } - if !native || !choice.text.is_empty() { yield Ok(StreamPart::TextDelta { id: "0".into(), delta: choice.text, provider_metadata: None }); } - } - } - if started { yield Ok(StreamPart::TextEnd { id: "0".into(), provider_metadata: None }); } - yield Ok(StreamPart::Finish { finish_reason, usage: usage(final_usage, native), provider_metadata: native.then(|| provider_namespace("openai", match logprobs { Some(logprobs) => json!({ "logprobs": logprobs }), None => json!({}) }).expect("provider metadata must be an object")) }); - }; + let stream = pipe_through( + futures::stream::iter(buffered).chain(events), + OpenAICompatibleCompletionStream { + ctx: StreamContext::new(options, warnings, error_url, &body, &response_headers), + native, + compat_name, + ..OpenAICompatibleCompletionStream::default() + }, + ); Ok(StreamResult { stream: Box::pin(stream), request: Some(RequestInfo { body: Some(body) }), response: Some(StreamResponseInfo { - headers: Some(stream_response_headers), + headers: Some(response_headers), }), }) } } + +/// The `TransformStream` of `OpenAICompatibleCompletionLanguageModel.doStream` +/// (and of `OpenAICompletionLanguageModel.doStream` when native). +#[derive(Default)] +struct OpenAICompatibleCompletionStream { + ctx: StreamContext, + native: bool, + /// The provider options name of a compatible provider; `None` for the + /// native OpenAI completion model. + compat_name: Option, + started: bool, + /// `other` until a chunk says otherwise. + finish_reason: Option, + usage: Option>, + logprobs: Option, +} + +impl OpenAICompatibleCompletionStream { + fn report_error(&mut self, error: AiMuxError, controller: &mut Controller) { + self.finish_reason = Some(FINISH_ERROR); + controller.enqueue(StreamPart::Error { error }); + } +} + +impl Transformer for OpenAICompatibleCompletionStream { + type Input = Result; + type Output = StreamPart; + + fn start(&mut self, controller: &mut Controller) { + self.ctx.start(controller); + } + + fn transform(&mut self, event: Result, controller: &mut Controller) { + let event = match event { + Ok(event) => event, + Err(error) => return self.report_error(error, controller), + }; + self.ctx.raw(&event, controller); + if let Some(error) = event.get("error") { + // @ai-sdk/openai maps the frame to a status + // (openai-completion-language-model.ts:305-313); + // @ai-sdk/openai-compatible passes it on as it is + // (openai-compatible-completion-language-model.ts:320-325). + let error = match &self.compat_name { + None => crate::openai::openai_stream_error( + error, + &self.ctx.url, + self.ctx.request_body.clone(), + self.ctx.response_headers.clone(), + ), + Some(name) => super::chat::stream_error( + name, + error, + &self.ctx.url, + self.ctx.request_body.clone(), + self.ctx.response_headers.clone(), + ), + }; + return self.report_error(error, controller); + } + let parsed: CompletionResponse = match parse_response(event, self.native, true) { + Ok(parsed) => parsed, + Err(error) => return self.report_error(error, controller), + }; + if !self.started { + self.started = true; + controller.enqueue(StreamPart::ResponseMetadata(metadata(&parsed))); + controller.text_start("0"); + } + if parsed.usage.is_some() { + self.usage = parsed.usage; + } + if let Some(choice) = parsed.choices.into_iter().next() { + if let Some(reason) = choice.finish_reason { + self.finish_reason = Some(parse_finish_reason(&reason)); + } + if self.native && choice.logprobs.is_some() { + self.logprobs = choice.logprobs; + } + if !self.native || !choice.text.is_empty() { + controller.text_delta("0", choice.text); + } + } + } + + fn flush(self, controller: &mut Controller) { + if self.started { + controller.text_end("0"); + } + let native = self.native; + let logprobs = self.logprobs; + controller.enqueue(StreamPart::Finish { + finish_reason: self.finish_reason.unwrap_or(FINISH_OTHER), + usage: usage(self.usage, native), + provider_metadata: native.then(|| { + provider_namespace( + "openai", + match logprobs { + Some(logprobs) => json!({ "logprobs": logprobs }), + None => json!({}), + }, + ) + .expect("provider metadata must be an object") + }), + }); + } +} diff --git a/aimux-provider-utils/src/streaming_tool_call_tracker.rs b/aimux-providers/src/stream/chat_tool_calls.rs similarity index 57% rename from aimux-provider-utils/src/streaming_tool_call_tracker.rs rename to aimux-providers/src/stream/chat_tool_calls.rs index f05710b1e..587be6a5c 100644 --- a/aimux-provider-utils/src/streaming_tool_call_tracker.rs +++ b/aimux-providers/src/stream/chat_tool_calls.rs @@ -1,75 +1,52 @@ -//! Streaming tool call tracker for the OpenAI chat-completions wire format. +//! The `tool_calls[i]` deltas of a chat-completions stream. //! -//! Port of `@ai-sdk/provider-utils`'s `streaming-tool-call-tracker.ts`; see -//! [`StreamingToolCallTracker`]. +//! OpenAI-style chat streams carry no tool-call boundaries: a call is a +//! series of `tool_calls[i]` fragments, labelled by `index`, `id` and +//! `function.name`, that providers fill in inconsistently. [`ChatToolCalls`] +//! is the part of a chat stream that puts those fragments back together the +//! way `@ai-sdk/provider-utils@5.0.55`'s `streaming-tool-call-tracker.ts` +//! does, so every chat stream (OpenAI, OpenAI-compatible, DeepSeek, Groq, +//! Mistral) emits the same parts for the same deltas. -use thiserror::Error; +use std::sync::Arc; use aimux_core::error::AiMuxError; use aimux_core::stream_part::StreamPart; use aimux_core::tool::RawToolCall; use aimux_core::types::ProviderMetadata; -use crate::streaming_tool_call_argument_state::{ - StreamingToolCallArgumentState, starts_with_structured_value, -}; +use super::tool_call_arguments::{StreamingToolCallArgumentState, starts_with_structured_value}; +use super::{Controller, Emit, GenerateId}; /// Id used only when the generator returns a blank string /// (provider-utils@5.0.55 src/streaming-tool-call-tracker.ts:461). const FALLBACK_TOOL_CALL_ID: &str = "tool-call"; -/// One `tool_calls[i]` entry of a streaming chunk, borrowed from the -/// provider's own wire type. -/// -/// The upstream `function: { name, arguments }` object is flattened; a -/// missing `function` is the same as both fields being `None`. -#[derive(Debug, Clone, Default)] -pub struct StreamingToolCallDelta<'a> { - pub index: Option, - pub id: Option<&'a str>, - pub r#type: Option<&'a str>, - pub name: Option<&'a str>, +/// One `tool_calls[i]` entry of a chunk, borrowed from the vendor's wire +/// type. A missing `function` object is both `name` and `arguments` absent. +#[derive(Debug, Default)] +pub(crate) struct ToolCallDelta<'a> { + pub(crate) index: Option, + pub(crate) id: Option<&'a str>, + pub(crate) name: Option<&'a str>, /// `None` is the wire `null` / absent; `Some("")` is an empty fragment. - pub arguments: Option<&'a str>, - /// Provider metadata for the `ToolCall` part (the upstream - /// `extractMetadata` hook). Read only from the delta that starts a call. - pub provider_metadata: Option, -} - -/// How to validate the `type` field on the delta that starts a call. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] -pub enum TypeValidation { - /// No validation. - #[default] - None, - /// Reject a `type` that is present and not `"function"`. - IfPresent, - /// Reject a `type` that is not exactly `"function"`. - Required, -} - -/// A delta the tracker cannot accept; the upstream -/// `InvalidResponseDataError`. -#[derive(Debug, Error, PartialEq, Eq)] -pub enum TrackerError { - #[error("Expected 'function.name' to be a string.")] - MissingFunctionName, - #[error("Expected 'function' type.")] - InvalidType, -} - -impl From for AiMuxError { - fn from(error: TrackerError) -> Self { - AiMuxError::InvalidResponseData(error.to_string()) - } + pub(crate) arguments: Option<&'a str>, + /// The delta's `type` breaks the vendor's rule (`'function'` when + /// present for OpenAI, always `'function'` for Groq). Only checked when + /// the delta starts a call. + pub(crate) invalid_type: bool, + /// Metadata for the `ToolCall` part (the upstream `extractMetadata`); + /// only read from the delta that starts a call. + pub(crate) provider_metadata: Option, } +/// A call being streamed. struct ToolCall { id: String, name: String, arguments: String, argument_state: StreamingToolCallArgumentState, - /// Index of the starting delta; orders [`StreamingToolCallTracker::finish`]. + /// Index of the starting delta; orders [`ChatToolCalls::finish`]. index: Option, /// Every wire id and index seen for this call. wire_ids: Vec, @@ -77,25 +54,24 @@ struct ToolCall { provider_metadata: Option, } +/// Which call a delta belongs to. enum Resolution { Existing(usize), New, Ambiguous, } -/// Tracks the tool calls of one streamed response. +/// The tool calls of one chat stream, in arrival order. /// -/// Port of `@ai-sdk/provider-utils`'s `streaming-tool-call-tracker.ts`. It -/// accumulates `tool_calls[i]` deltas, emits [`StreamPart::ToolInputStart`] / -/// [`StreamPart::ToolInputDelta`] as they arrive, and emits -/// [`StreamPart::ToolInputEnd`] / [`StreamPart::ToolCall`] for every call in -/// [`StreamingToolCallTracker::finish`]. A call is never finalized earlier: a -/// parsable argument buffer can still be the prefix of a longer argument -/// string (ai-sdk #13137). +/// Every delta is emitted as it arrives ([`StreamPart::ToolInputStart`] / +/// [`StreamPart::ToolInputDelta`]); every call is ended only by +/// [`finish`](Self::finish) when the stream flushes, because a parsable +/// argument buffer can still be the prefix of longer arguments (ai-sdk +/// #13137). /// -/// Deltas are correlated to calls by wire `id`, `index` and function name, -/// not by `index` alone, so providers that reuse indices, repeat ids, drop -/// ids on continuations or send blank names are handled: +/// Deltas are matched to calls by wire `id`, `index` and function name, not +/// by `index` alone, so providers that reuse indices, repeat ids, drop ids on +/// continuations or send blank names are handled: /// /// | ID evidence | index/name evidence | start evidence | resolution | /// | --- | --- | --- | --- | @@ -108,68 +84,43 @@ enum Resolution { /// | absent | absent | named | new call | /// | absent | absent | unnamed | sole unfinished call, new call, or ambiguity | /// -/// An ambiguous delta is dropped. Behaviour follows -/// `@ai-sdk/provider-utils@5.0.55`, which keeps parallel calls apart when -/// stream labels are missing or reused (`hasEmptyArguments`). -/// -/// Protocols whose streams carry explicit tool-call boundaries (Anthropic, -/// Google, Bedrock, Cohere, the Responses API) do not need this. -pub struct StreamingToolCallTracker { +/// An ambiguous delta is dropped (`hasEmptyArguments` in 5.0.55 keeps +/// parallel calls apart when labels are missing or reused). +pub(crate) struct ChatToolCalls { calls: Vec, - generate_id: Box String + Send + Sync>, - type_validation: TypeValidation, + /// Names a call without a usable wire id. Blank or repeated outputs are + /// made unique. + generate_id: GenerateId, } -impl Default for StreamingToolCallTracker { +/// Ids from [`aimux_provider_utils::generate_id`]. +impl Default for ChatToolCalls { fn default() -> Self { - Self::new() + Self::new(Arc::new(aimux_provider_utils::generate_id)) } } -impl StreamingToolCallTracker { - /// A tracker without `type` validation. Generated ids come from `generate_id` - /// unless [`Self::with_generate_id`] sets another generator - /// (provider-utils@5.0.55 src/streaming-tool-call-tracker.ts:114). - #[must_use] - pub fn new() -> Self { +impl ChatToolCalls { + pub(crate) fn new(generate_id: GenerateId) -> Self { Self { calls: Vec::new(), - generate_id: Box::new(crate::generate_id::generate_id), - type_validation: TypeValidation::None, + generate_id, } } - /// Set the generator for calls without a usable wire id. Blank or - /// repeated outputs are turned into unique ids. - #[must_use] - pub fn with_generate_id(mut self, f: impl Fn() -> String + Send + Sync + 'static) -> Self { - self.generate_id = Box::new(f); - self - } - - #[must_use] - pub fn with_type_validation(mut self, validation: TypeValidation) -> Self { - self.type_validation = validation; - self - } - - /// Process one delta, appending the parts it produces to `out`. - /// - /// # Errors + /// Apply one delta, emitting the parts it produces. /// - /// A delta that starts a call without a function name, or whose `type` - /// fails validation. - pub fn process( - &mut self, - delta: StreamingToolCallDelta<'_>, - out: &mut Vec, - ) -> Result<(), TrackerError> { - let StreamingToolCallDelta { + /// A delta that starts a call with an invalid `type` or without a + /// function name is the upstream `InvalidResponseDataError`: the stream + /// ends with it and `false` is returned. + #[must_use] + pub(crate) fn delta(&mut self, delta: ToolCallDelta<'_>, controller: &mut Controller) -> bool { + let ToolCallDelta { index, id, - r#type, name: wire_name, arguments, + invalid_type, provider_metadata, } = delta; let wire_id = non_blank(id); @@ -178,23 +129,28 @@ impl StreamingToolCallTracker { let empty_arguments = arguments.is_none_or(|arguments| arguments.trim().is_empty()); let call = match self.resolve(wire_id, index, name, explicit_start, empty_arguments) { - Resolution::Ambiguous => return Ok(()), + Resolution::Ambiguous => return true, // Blank names cannot start a usable call, but some providers - // repeat a blank name on continuations; those were correlated + // repeat a blank name on continuations; those were matched // above, so only an unmatched blank-name delta is ignored. - Resolution::New if wire_name.is_some() && name.is_none() => return Ok(()), - Resolution::New => { - self.validate_type(r#type)?; - let name = name.ok_or(TrackerError::MissingFunctionName)?; - self.start( - wire_id, - index, - name, - arguments.unwrap_or(""), - provider_metadata, - out, - ) - } + Resolution::New if wire_name.is_some() && name.is_none() => return true, + Resolution::New => match (invalid_type, name) { + (false, Some(name)) => { + let arguments = arguments.unwrap_or(""); + self.start( + wire_id, + index, + name, + arguments, + provider_metadata, + controller, + ) + } + (true, _) => return invalid(controller, "Expected 'function' type."), + (false, None) => { + return invalid(controller, "Expected 'function.name' to be a string."); + } + }, Resolution::Existing(call) => { let tool_call = &mut self.calls[call]; if let Some(wire_id) = wire_id @@ -205,11 +161,7 @@ impl StreamingToolCallTracker { if let Some(arguments) = arguments { tool_call.argument_state.append(arguments); tool_call.arguments.push_str(arguments); - out.push(StreamPart::ToolInputDelta { - id: tool_call.id.clone(), - delta: arguments.to_string(), - provider_metadata: None, - }); + controller.tool_input_delta(tool_call.id.clone(), arguments.to_string()); } call } @@ -221,22 +173,18 @@ impl StreamingToolCallTracker { { indices.push(index); } - Ok(()) + true } - /// Finalize every call, appending `ToolInputEnd` and `ToolCall` parts to - /// `out`. Calls are ordered by index when every call has one, otherwise - /// by arrival. - pub fn finish(mut self, out: &mut Vec) { + /// End every call with `ToolInputEnd` and `ToolCall`, ordered by index + /// when every call has one, otherwise by arrival. + pub(crate) fn finish(mut self, controller: &mut Controller) { if self.calls.iter().all(|c| c.index.is_some()) { self.calls.sort_by_key(|c| c.index); // stable: ties keep arrival order } for call in self.calls { - out.push(StreamPart::ToolInputEnd { - id: call.id.clone(), - provider_metadata: None, - }); - out.push(StreamPart::ToolCall(RawToolCall { + controller.tool_input_end(call.id.clone()); + controller.enqueue(StreamPart::ToolCall(RawToolCall { tool_call_id: call.id, tool_name: call.name, // Raw argument text; Core parses it. @@ -248,7 +196,7 @@ impl StreamingToolCallTracker { } } - /// Apply the correlation table in the type docs. + /// Apply the matching table in the type docs. fn resolve( &self, wire_id: Option<&str>, @@ -276,11 +224,10 @@ impl StreamingToolCallTracker { let has_id = |c: &ToolCall| c.wire_ids.iter().any(|w| w == wire_id); if !self.calls.iter().any(has_id) { - // A named structured start with an unseen id identifies a new - // call. Empty arguments can also occur when a continuation - // changes its id: keep the incomplete matching call then, and - // start a new call only once the matching call's structured - // arguments are complete. + // A named structured start with an unseen id is a new call. + // Empty arguments also occur when a continuation changes its id: + // keep the incomplete matching call then, and start a new call + // only once its structured arguments are complete. if explicit_start { return Resolution::New; } @@ -344,19 +291,7 @@ impl StreamingToolCallTracker { } } - fn validate_type(&self, r#type: Option<&str>) -> Result<(), TrackerError> { - let valid = match self.type_validation { - TypeValidation::None => true, - TypeValidation::IfPresent => r#type.is_none_or(|t| t == "function"), - TypeValidation::Required => r#type == Some("function"), - }; - if valid { - Ok(()) - } else { - Err(TrackerError::InvalidType) - } - } - + /// Start a call: `ToolInputStart`, then its first argument fragment. fn start( &mut self, wire_id: Option<&str>, @@ -364,23 +299,12 @@ impl StreamingToolCallTracker { name: &str, arguments: &str, provider_metadata: Option, - out: &mut Vec, + controller: &mut Controller, ) -> usize { let id = self.unique_id(wire_id); - out.push(StreamPart::ToolInputStart { - id: id.clone(), - tool_name: name.to_string(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); + controller.tool_input_start(id.clone(), name.to_string()); if !arguments.is_empty() { - out.push(StreamPart::ToolInputDelta { - id: id.clone(), - delta: arguments.to_string(), - provider_metadata: None, - }); + controller.tool_input_delta(id.clone(), arguments.to_string()); } self.calls.push(ToolCall { id, @@ -422,6 +346,12 @@ impl StreamingToolCallTracker { } } +/// End the stream with the upstream `InvalidResponseDataError`. +fn invalid(controller: &mut Controller, message: &str) -> bool { + controller.error(AiMuxError::InvalidResponseData(message.to_string())); + false +} + fn non_blank(value: Option<&str>) -> Option<&str> { value.filter(|v| !v.trim().is_empty()) } diff --git a/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs b/aimux-providers/src/stream/chat_tool_calls_tests.rs similarity index 67% rename from aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs rename to aimux-providers/src/stream/chat_tool_calls_tests.rs index d89d9c7e5..fb6e41919 100644 --- a/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs +++ b/aimux-providers/src/stream/chat_tool_calls_tests.rs @@ -1,23 +1,40 @@ -//! Port of `streaming-tool-call-tracker.test.ts` from -//! `@ai-sdk/provider-utils`, case for case. +//! `streaming-tool-call-tracker.test.ts` from `@ai-sdk/provider-utils`, case +//! for case, and the T3 differential fixture recorded from the same tracker, +//! run through the chat streams that use [`ChatToolCalls`](super::ChatToolCalls). //! -//! Tracker parts are projected to a local `Part` enum so the assertions stay -//! byte-for-byte comparable with the upstream expectations. +//! A case drives a stream's `Transformer` synchronously: every +//! `tool_calls[i]` delta is one chunk, `flush` ends the stream, and only the +//! tool parts are kept. Which stream runs a case follows the tracker options +//! upstream tested: +//! +//! - no `type` check (the tracker default): Mistral, whose `index` is +//! optional and whose ids come from an injected generator; +//! - `'function'` when present (`'if-present'`): OpenAI; +//! - `'function'` required (`'required'`): Groq; +//! - `extractMetadata`: OpenAI-compatible, whose thought signature is the +//! metadata of the call a delta starts. use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; +use serde_json::{Value, json}; + +use aimux_core::error::AiMuxError; use aimux_core::shared::provider_namespace; use aimux_core::stream_part::StreamPart; use aimux_core::tool::RawToolCall; use aimux_core::types::ProviderMetadata; -use aimux_provider_utils::{ - StreamingToolCallDelta, StreamingToolCallTracker, TrackerError, TypeValidation, -}; -use serde_json::json; +use aimux_provider_utils::{TransformStreamController, Transformer}; + +use super::Controller; + +type GenerateId = Arc String + Send + Sync>; -/// The TS tracker's four events, projected out of [`StreamPart`] so the -/// assertions below read like the upstream suite. +const MISSING_NAME: &str = "Expected 'function.name' to be a string."; +const INVALID_TYPE: &str = "Expected 'function' type."; + +/// The tracker's four events, projected out of [`StreamPart`] so the +/// assertions read like the upstream suite. #[derive(Debug, Clone, PartialEq)] #[allow( clippy::enum_variant_names, @@ -44,8 +61,11 @@ enum Part { }, } -fn project(part: StreamPart) -> Part { - match part { +/// The tool part `part` is, if it is one. Any other part a chat stream emits +/// around tool calls (stream start, response metadata, finish) is dropped; +/// an error part means the case's chunk did not fit the vendor's schema. +fn project(part: StreamPart) -> Option { + Some(match part { StreamPart::ToolInputStart { id, tool_name, @@ -76,44 +96,124 @@ fn project(part: StreamPart) -> Part { input, provider_metadata, }, - other => panic!("tracker emitted an unexpected part: {other:?}"), + StreamPart::Error { error } => panic!("the stream rejected a chunk: {error:?}"), + part @ (StreamPart::ToolInputStart { .. } + | StreamPart::ToolInputDelta { .. } + | StreamPart::ToolInputEnd { .. } + | StreamPart::ToolCall(_)) => panic!("unexpected tool part: {part:?}"), + _ => return None, + }) +} + +/// A chat stream with its chunk type fixed, so the vendors can share one +/// harness. +trait ChatStream { + fn chunk(&mut self, chunk: Value, controller: &mut Controller); + fn end(self: Box, controller: &mut Controller); +} + +impl ChatStream for T +where + T: Transformer, Output = StreamPart>, +{ + fn chunk(&mut self, chunk: Value, controller: &mut Controller) { + self.transform(Ok(chunk), controller); + } + + fn end(self: Box, controller: &mut Controller) { + (*self).flush(controller); } } -/// Tracker plus the parts it has emitted so far (the TS `createCollector`). +/// A chat stream plus the tool parts it has emitted so far (the TS +/// `createCollector`). struct Harness { - /// `None` once flushed: `finish` consumes the tracker. - tracker: Option, + /// `None` once flushed. + stream: Option>, parts: Vec, } impl Harness { + /// No `type` check; every generated id is `tool-call`. fn new() -> Self { - Self::with(StreamingToolCallTracker::new().with_generate_id(|| "tool-call".to_string())) + Self::with_ids(|| "tool-call".to_string()) + } + + /// No `type` check, ids from `generate_id`. + fn with_ids(generate_id: impl Fn() -> String + Send + Sync + 'static) -> Self { + Self::of(crate::mistral::chat_test_stream(Arc::new(generate_id))) } - fn with(tracker: StreamingToolCallTracker) -> Self { + /// `type` must be `'function'` when present. + fn openai() -> Self { + Self::of(crate::openai::model::chat_test_stream(Arc::new(|| { + "tool-call".to_string() + }))) + } + + /// `type` must be `'function'` on the delta that starts a call. + fn groq() -> Self { + Self::of(crate::groq::chat_test_stream(Arc::new(|| { + "tool-call".to_string() + }))) + } + + /// Thought signatures become the metadata of the call; the provider's + /// metadata key is `google`. + fn compatible() -> Self { + Self::of(crate::openai_compatible::chat::chat_test_stream( + "google", + Arc::new(|| "tool-call".to_string()), + )) + } + + fn of(stream: impl ChatStream + 'static) -> Self { Self { - tracker: Some(tracker), + stream: Some(Box::new(stream)), parts: Vec::new(), } } - fn delta(&mut self, delta: StreamingToolCallDelta<'_>) -> Result<(), TrackerError> { - let mut parts = Vec::new(); - let tracker = self.tracker.as_mut().expect("delta after flush"); - let result = tracker.process(delta, &mut parts); - self.parts.extend(parts.into_iter().map(project)); - result + /// Feed one `tool_calls[i]` entry as a chunk. An `Err` is the message of + /// the `InvalidResponseData` error the stream ended with. Every delta + /// gets a fresh controller, so, like the upstream tracker, the stream + /// can be fed again after an error. + fn delta(&mut self, tool_call: Value) -> Result<(), String> { + let mut controller = TransformStreamController::new(); + let stream = self.stream.as_mut().expect("delta after flush"); + stream.chunk( + json!({ "choices": [{ "delta": { "tool_calls": [tool_call] } }] }), + &mut controller, + ); + self.collect(&mut controller) } fn flush(&mut self) { - let mut parts = Vec::new(); - self.tracker - .take() - .expect("flushed twice") - .finish(&mut parts); - self.parts.extend(parts.into_iter().map(project)); + let mut controller = TransformStreamController::new(); + let stream = self.stream.take().expect("flushed twice"); + stream.end(&mut controller); + // The OpenAI-compatible stream reports a missing finish reason as an + // error part; it is not a tool part. + let parts: Vec = controller.drain().filter_map(Result::ok).collect(); + self.parts.extend( + parts + .into_iter() + .filter(|part| !matches!(part, StreamPart::Error { .. })) + .filter_map(project), + ); + } + + fn collect(&mut self, controller: &mut Controller) -> Result<(), String> { + let mut result = Ok(()); + for item in controller.drain() { + match item { + Ok(part) => self.parts.extend(project(part)), + Err(AiMuxError::InvalidResponseData(message)) => result = Err(message), + Err(error) => panic!("unexpected stream error: {error:?}"), + } + } + assert_eq!(result.is_err(), controller.is_closed()); + result } fn clear(&mut self) { @@ -137,31 +237,37 @@ impl Harness { } } -/// A delta with an explicit `function` object (as every TS test passes one). -fn delta<'a>( +/// A `tool_calls[i]` entry with an explicit `function` object (as every TS +/// test passes one). `None` leaves a field out. +fn delta( index: Option, - id: Option<&'a str>, - ty: Option<&'a str>, - name: Option<&'a str>, - arguments: Option<&'a str>, -) -> StreamingToolCallDelta<'a> { - StreamingToolCallDelta { - index, - id, - r#type: ty, - name, - arguments, - provider_metadata: None, + id: Option<&str>, + ty: Option<&str>, + name: Option<&str>, + arguments: Option<&str>, +) -> Value { + let mut function = serde_json::Map::new(); + if let Some(name) = name { + function.insert("name".into(), json!(name)); + } + if let Some(arguments) = arguments { + function.insert("arguments".into(), json!(arguments)); + } + let mut tool_call = json!({ "function": function }); + if let Some(index) = index { + tool_call["index"] = json!(index); + } + if let Some(id) = id { + tool_call["id"] = json!(id); } + if let Some(ty) = ty { + tool_call["type"] = json!(ty); + } + tool_call } /// Full call start: index, id, `type: 'function'`, name and arguments. -fn start<'a>( - index: usize, - id: &'a str, - name: &'a str, - arguments: &'a str, -) -> StreamingToolCallDelta<'a> { +fn start(index: usize, id: &str, name: &str, arguments: &str) -> Value { delta( Some(index), Some(id), @@ -172,10 +278,16 @@ fn start<'a>( } /// Continuation carrying only arguments (plus an optional index). -fn cont(index: Option, arguments: &str) -> StreamingToolCallDelta<'_> { +fn cont(index: Option, arguments: &str) -> Value { delta(index, None, None, None, Some(arguments)) } +/// `tool_call` carrying a Gemini thought signature (`extra_content`). +fn signed(mut tool_call: Value, signature: &str) -> Value { + tool_call["extra_content"] = json!({ "google": { "thought_signature": signature } }); + tool_call +} + fn tc(id: &str, name: &str, input: &str) -> (String, String, String) { (id.into(), name.into(), input.into()) } @@ -207,8 +319,14 @@ fn tool_call(id: &str, name: &str, input: &str) -> Part { } } -/// Deterministic id generator: `prefix-1`, `prefix-2`, … and a call counter. -fn counting_generator(ids: &'static [&'static str]) -> (impl Fn() -> String, Arc) { +/// Deterministic id generator over `ids` (the last one repeats) and a call +/// counter. +fn counting_generator( + ids: &'static [&'static str], +) -> ( + impl Fn() -> String + Send + Sync + 'static, + Arc, +) { let calls = Arc::new(AtomicUsize::new(0)); let counter = Arc::clone(&calls); ( @@ -219,7 +337,6 @@ fn counting_generator(ids: &'static [&'static str]) -> (impl Fn() -> String, Arc calls, ) } - mod process_delta { use super::*; @@ -560,9 +677,7 @@ mod process_delta { #[test] fn generates_an_id_when_id_is_missing() { - let mut h = Harness::with( - StreamingToolCallTracker::new().with_generate_id(|| "generated-id".to_string()), - ); + let mut h = Harness::with_ids(|| "generated-id".to_string()); h.delta(delta( Some(0), @@ -584,11 +699,7 @@ mod process_delta { let result = h.delta(delta(Some(0), Some("call_1"), Some("function"), None, None)); - assert_eq!(result, Err(TrackerError::MissingFunctionName)); - assert_eq!( - TrackerError::MissingFunctionName.to_string(), - "Expected 'function.name' to be a string." - ); + assert_eq!(result, Err(MISSING_NAME.to_string())); } #[test] @@ -635,7 +746,7 @@ mod process_delta { #[test] fn keeps_id_less_calls_distinct_when_an_index_is_reused_and_type_is_omitted() { let (generate, _) = counting_generator(&["generated-1", "generated-2", "generated-3"]); - let mut h = Harness::with(StreamingToolCallTracker::new().with_generate_id(generate)); + let mut h = Harness::with_ids(generate); h.delta(delta( Some(0), @@ -676,9 +787,7 @@ mod process_delta { #[test] fn keeps_complete_same_name_calls_distinct_with_reused_index() { for id in [None, Some("dup")] { - let mut h = Harness::with( - StreamingToolCallTracker::new().with_generate_id(|| "generated-id".to_string()), - ); + let mut h = Harness::with_ids(|| "generated-id".to_string()); for value in ["{\"value\":1}", "{\"value\":2}"] { h.delta(delta( @@ -703,9 +812,7 @@ mod process_delta { #[test] fn keeps_a_partial_same_name_call_distinct_with_reused_index() { for id in [None, Some("dup")] { - let mut h = Harness::with( - StreamingToolCallTracker::new().with_generate_id(|| "generated-id".to_string()), - ); + let mut h = Harness::with_ids(|| "generated-id".to_string()); for value in ["{\"value\":1}", "{\"value\":"] { h.delta(delta( @@ -786,7 +893,7 @@ mod process_delta { #[test] fn generates_unique_ids_for_blank_and_repeated_ids() { let (generate, _) = counting_generator(&["generated-1", "generated-2"]); - let mut h = Harness::with(StreamingToolCallTracker::new().with_generate_id(generate)); + let mut h = Harness::with_ids(generate); h.delta(start(0, "", "read_file", "{}")).unwrap(); h.delta(start(1, "dup", "read_file", "{}")).unwrap(); @@ -805,9 +912,7 @@ mod process_delta { #[test] fn keeps_same_name_calls_with_repeated_ids_and_distinct_indices_separate() { - let mut h = Harness::with( - StreamingToolCallTracker::new().with_generate_id(|| "generated-id".to_string()), - ); + let mut h = Harness::with_ids(|| "generated-id".to_string()); h.delta(start(0, "dup", "same_tool", "{\"value\":0}")) .unwrap(); @@ -848,7 +953,7 @@ mod process_delta { #[test] fn creates_bounded_unique_ids_when_generate_id_returns_duplicates() { let (generate, calls) = counting_generator(&["generated-id"]); - let mut h = Harness::with(StreamingToolCallTracker::new().with_generate_id(generate)); + let mut h = Harness::with_ids(generate); for (index, name) in [(0, "first"), (1, "second"), (2, "third")] { h.delta(delta( @@ -873,7 +978,7 @@ mod process_delta { #[test] fn creates_usable_ids_when_generate_id_returns_blank_values() { let (generate, calls) = counting_generator(&[" "]); - let mut h = Harness::with(StreamingToolCallTracker::new().with_generate_id(generate)); + let mut h = Harness::with_ids(generate); for (index, name) in [(0, "first"), (1, "second")] { h.delta(delta( @@ -914,9 +1019,7 @@ mod process_delta { #[test] fn ignores_an_ambiguous_continuation_for_a_repeated_id() { - let mut h = Harness::with( - StreamingToolCallTracker::new().with_generate_id(|| "generated-id".to_string()), - ); + let mut h = Harness::with_ids(|| "generated-id".to_string()); h.delta(start(0, "dup", "read_file", "{\"path\":\"a\"}")) .unwrap(); @@ -943,9 +1046,7 @@ mod process_delta { #[test] fn uses_a_matching_name_and_index_for_an_id_less_continuation() { - let mut h = Harness::with( - StreamingToolCallTracker::new().with_type_validation(TypeValidation::Required), - ); + let mut h = Harness::groq(); h.delta(start(0, "call_1", "read_file", "{\"pa")).unwrap(); h.delta(delta( @@ -1066,14 +1167,14 @@ mod process_delta { // `function: {}` — a function object with neither name nor arguments. let result = h.delta(delta(Some(0), Some("call_1"), Some("function"), None, None)); - assert_eq!(result, Err(TrackerError::MissingFunctionName)); + assert_eq!(result, Err(MISSING_NAME.to_string())); } } mod type_validation { use super::*; - fn custom_type() -> StreamingToolCallDelta<'static> { + fn custom_type() -> Value { delta( Some(0), Some("call_1"), @@ -1083,30 +1184,22 @@ mod type_validation { ) } - fn no_type() -> StreamingToolCallDelta<'static> { + fn no_type() -> Value { delta(Some(0), Some("call_1"), None, Some("fn"), Some("")) } #[test] fn does_not_validate_type_with_none() { - let mut h = Harness::with( - StreamingToolCallTracker::new().with_type_validation(TypeValidation::None), - ); + let mut h = Harness::new(); assert_eq!(h.delta(custom_type()), Ok(())); } #[test] fn validates_type_when_present_with_if_present() { - let mut h = Harness::with( - StreamingToolCallTracker::new().with_type_validation(TypeValidation::IfPresent), - ); + let mut h = Harness::openai(); - assert_eq!(h.delta(custom_type()), Err(TrackerError::InvalidType)); - assert_eq!( - TrackerError::InvalidType.to_string(), - "Expected 'function' type." - ); + assert_eq!(h.delta(custom_type()), Err(INVALID_TYPE.to_string())); // A missing type is accepted. assert_eq!(h.delta(no_type()), Ok(())); @@ -1114,11 +1207,9 @@ mod type_validation { #[test] fn requires_function_type_with_required() { - let mut h = Harness::with( - StreamingToolCallTracker::new().with_type_validation(TypeValidation::Required), - ); + let mut h = Harness::groq(); - assert_eq!(h.delta(no_type()), Err(TrackerError::InvalidType)); + assert_eq!(h.delta(no_type()), Err(INVALID_TYPE.to_string())); assert_eq!(h.delta(start(0, "call_1", "fn", "")), Ok(())); } @@ -1151,19 +1242,15 @@ mod metadata { use super::*; // The upstream `extractMetadata` / `buildToolCallProviderMetadata` hooks - // become `provider_metadata` on the delta that starts a call. + // are the OpenAI-compatible stream's thought signatures here: the + // signature of the delta that starts a call is the call's metadata. #[test] fn includes_the_starting_delta_metadata_in_the_tool_call() { - let mut h = Harness::new(); + let mut h = Harness::compatible(); - h.delta(StreamingToolCallDelta { - provider_metadata: Some( - provider_namespace("google", json!({ "thoughtSignature": "sig123" })).unwrap(), - ), - ..start(0, "call_1", "fn", "{\"incomplete") - }) - .unwrap(); + h.delta(signed(start(0, "call_1", "fn", "{\"incomplete"), "sig123")) + .unwrap(); h.flush(); assert_eq!( @@ -1181,16 +1268,10 @@ mod metadata { #[test] fn ignores_metadata_on_continuations() { - let mut h = Harness::new(); + let mut h = Harness::compatible(); h.delta(start(0, "call_1", "fn", "{")).unwrap(); - h.delta(StreamingToolCallDelta { - provider_metadata: Some( - provider_namespace("custom", json!({ "key": "value" })).unwrap(), - ), - ..cont(Some(0), "}") - }) - .unwrap(); + h.delta(signed(cont(Some(0), "}"), "sig123")).unwrap(); h.flush(); assert_eq!(h.parts.last(), Some(&tool_call("call_1", "fn", "{}"))); @@ -1202,7 +1283,7 @@ mod generate_id { #[test] fn generates_a_16_character_id_when_wire_id_is_missing() { - let mut h = Harness::with(StreamingToolCallTracker::new()); + let mut h = Harness::with_ids(aimux_provider_utils::generate_id); h.delta(delta( Some(0), @@ -1233,9 +1314,7 @@ mod generate_id { #[test] fn keeps_the_wire_id_when_present() { - let mut h = Harness::with( - StreamingToolCallTracker::new().with_generate_id(|| "custom-id".to_string()), - ); + let mut h = Harness::with_ids(|| "custom-id".to_string()); h.delta(start(0, "call_1", "fn", "{\"key\": \"val")) .unwrap(); @@ -1247,3 +1326,163 @@ mod generate_id { assert_eq!(ids, vec!["call_1"]); } } + +/// The T3 differential fixture: +/// `fixtures/aisdk/provider-utils/t3-unattributable-deltas.json` holds the +/// deltas of each case plus what the `StreamingToolCallTracker` of the +/// `@ai-sdk/provider-utils` version pinned in `fixtures/aisdk/VERSIONS.json` +/// (5.0.55) emitted for them with the case's `typeValidation` +/// (`'if-present'`: openai chat; `'required'`: groq; `'none'`: +/// openai-compatible), recorded by `scripts/gen-t3-tracker-fixture.mjs`. +/// Each case asserts that a chat stream with the same `type` rule matches +/// per delta, on the error, and on flush. +mod t3_unattributable_delta_parity { + use super::*; + + const FIXTURE: &str = + include_str!("../../../fixtures/aisdk/provider-utils/t3-unattributable-deltas.json"); + + fn to_aisdk(part: &Part) -> Value { + match part { + Part::ToolInputStart { id, tool_name } => { + json!({"type": "tool-input-start", "id": id, "toolName": tool_name}) + } + Part::ToolInputDelta { id, delta } => { + json!({"type": "tool-input-delta", "id": id, "delta": delta}) + } + Part::ToolInputEnd { id } => json!({"type": "tool-input-end", "id": id}), + Part::ToolCall { + tool_call_id, + tool_name, + input, + .. + } => json!({ + "type": "tool-call", "toolCallId": tool_call_id, "toolName": tool_name, "input": input + }), + } + } + + /// The stream that applies the case's `typeValidation` to its deltas. + /// + /// The vendor schemas are stricter than the tracker: OpenAI and Groq + /// require `index`, and Groq rejects a `type` other than `'function'` + /// while parsing. A case those schemas would reject runs on a stream + /// whose rule gives the same verdict for its deltas: Mistral (no check) + /// for `'if-present'` deltas without an index, none of which has a + /// non-`'function'` type; OpenAI for `'required'` deltas with another + /// `type`, which both rules reject. + fn harness(case: &Value, generate_id: impl Fn() -> String + Send + Sync + 'static) -> Harness { + let deltas = case["deltas"].as_array().unwrap(); + let all_indexed = deltas.iter().all(|d| d.get("index").is_some()); + let other_type = deltas + .iter() + .any(|d| d.get("type").is_some_and(|t| t != "function")); + let generate_id: GenerateId = Arc::new(generate_id); + let stream: Box = match case["typeValidation"].as_str().unwrap() { + "if-present" if all_indexed => { + Box::new(crate::openai::model::chat_test_stream(generate_id)) + } + "if-present" => { + assert!(!other_type); + Box::new(crate::mistral::chat_test_stream(generate_id)) + } + "required" if other_type => { + Box::new(crate::openai::model::chat_test_stream(generate_id)) + } + "required" => Box::new(crate::groq::chat_test_stream(generate_id)), + "none" => Box::new(crate::mistral::chat_test_stream(generate_id)), + other => panic!("unknown typeValidation {other}"), + }; + Harness { + stream: Some(stream), + parts: Vec::new(), + } + } + + fn run_case(case: &Value) { + let name = case["name"].as_str().unwrap(); + // Ids are generated `gen-0`, `gen-1`, ... like the recording script. + let next_id = AtomicUsize::new(0); + let mut h = harness(case, move || { + format!("gen-{}", next_id.fetch_add(1, Ordering::SeqCst)) + }); + let mut steps = Vec::new(); + let mut error = Value::Null; + for tool_call in case["deltas"].as_array().unwrap() { + h.clear(); + let result = h.delta(tool_call.clone()); + steps.push(json!({"parts": h.parts.iter().map(to_aisdk).collect::>()})); + if let Err(message) = result { + error = json!({"name": "AI_InvalidResponseDataError", "message": message}); + break; + } + } + let flush = if error.is_null() { + h.clear(); + h.flush(); + Value::Array(h.parts.iter().map(to_aisdk).collect()) + } else { + Value::Null + }; + let actual = json!({"steps": steps, "error": error, "flush": flush}); + assert_eq!(actual, case["aisdk"], "{name} ({})", case["tracker"]); + } + + fn case(name: &str) -> Value { + let fixture: Value = serde_json::from_str(FIXTURE).unwrap(); + fixture["cases"] + .as_array() + .unwrap() + .iter() + .find(|c| c["name"] == name) + .unwrap_or_else(|| panic!("missing fixture case {name}")) + .clone() + } + + /// Every recorded case, including the typeValidation `required` (groq) + /// and `none` (openai-compatible) ones and the interleaved ones. + #[test] + fn every_fixture_case_matches_the_aisdk_tracker() { + let fixture: Value = serde_json::from_str(FIXTURE).unwrap(); + let cases = fixture["cases"].as_array().unwrap(); + assert!(!cases.is_empty()); + for case in cases { + run_case(case); + } + } + + /// streaming-tool-call-tracker.ts:144-146, :305-308. + #[test] + fn ambiguous_delta_across_open_calls_is_dropped() { + run_case(&case("ambiguous_delta_across_open_calls_is_dropped")); + } + + /// streaming-tool-call-tracker.ts:153-155. + #[test] + fn blank_name_matching_no_call_is_dropped() { + run_case(&case("blank_name_matching_no_call_is_dropped")); + } + + /// streaming-tool-call-tracker.ts:305-308, :384-389. + #[test] + fn nameless_delta_without_open_call_is_invalid_response_data() { + run_case(&case( + "nameless_delta_without_open_call_is_invalid_response_data", + )); + } + + /// streaming-tool-call-tracker.ts:375-381. + #[test] + fn non_function_type_is_invalid_response_data() { + run_case(&case("non_function_type_is_invalid_response_data")); + } + + /// The case 5.0.55 changed: an index-only continuation after the index + /// is reused goes to the only call there without complete structured + /// arguments (streaming-tool-call-tracker.ts:291-298, :345-347); 5.0.53 + /// dropped it. + #[test] + fn index_only_continuation_goes_to_the_only_incomplete_call() { + run_case(&case("index_reuse_with_new_id")); + } +} diff --git a/aimux-providers/src/stream/mod.rs b/aimux-providers/src/stream/mod.rs new file mode 100644 index 000000000..9f932cc73 --- /dev/null +++ b/aimux-providers/src/stream/mod.rs @@ -0,0 +1,261 @@ +//! What the vendor `doStream` transformers share. +//! +//! A vendor stream is a [`Transformer`](aimux_provider_utils::Transformer) +//! struct: the read-only [`StreamContext`] of its request plus the state it +//! accumulates between chunks. The OpenAI-style chat streams also share +//! [`ChatToolCalls`]. + +mod chat_tool_calls; +#[cfg(test)] +mod chat_tool_calls_tests; +mod tool_call_arguments; + +use std::collections::HashMap; +use std::sync::Arc; + +use serde_json::Value; + +use aimux_core::options::CallOptions; +use aimux_core::stream_part::StreamPart; +use aimux_core::types::{FinishReason, FinishReasonUnified, Warning}; +use aimux_provider_utils::TransformStreamController; + +pub(crate) use chat_tool_calls::{ChatToolCalls, ToolCallDelta}; + +/// The request a stream answers, fixed before the first chunk: the +/// warnings `start` hands out, whether raw chunks are forwarded, what an +/// in-stream error reports (the URL, the request body and the response +/// headers) and the model's id generator. A stream that reports none of +/// these leaves them empty. +pub(crate) struct StreamContext { + pub(crate) warnings: Vec, + /// `includeRawChunks`. + pub(crate) include_raw_chunks: bool, + pub(crate) url: String, + pub(crate) request_body: Value, + pub(crate) response_headers: HashMap, + /// The model's `generateId`; [`aimux_provider_utils::generate_id`] unless + /// the provider settings replace it. + pub(crate) generate_id: GenerateId, +} + +/// A model's id generator. +pub(crate) type GenerateId = Arc String + Send + Sync>; + +impl Default for StreamContext { + fn default() -> Self { + Self { + warnings: Vec::new(), + include_raw_chunks: false, + url: String::new(), + request_body: Value::Null, + response_headers: HashMap::new(), + generate_id: Arc::new(aimux_provider_utils::generate_id), + } + } +} + +impl StreamContext { + pub(crate) fn new( + options: &CallOptions, + warnings: Vec, + url: String, + request_body: &Value, + response_headers: &HashMap, + ) -> Self { + Self { + warnings, + include_raw_chunks: options.include_raw_chunks == Some(true), + url, + request_body: request_body.clone(), + response_headers: response_headers.clone(), + ..Self::default() + } + } + + /// The `stream-start` part every stream begins with. + pub(crate) fn start(&mut self, controller: &mut Controller) { + controller.enqueue(StreamPart::StreamStart { + warnings: std::mem::take(&mut self.warnings), + }); + } + + /// Forward `chunk` as a `raw` part when raw chunks were asked for. + pub(crate) fn raw(&self, chunk: &Value, controller: &mut Controller) { + if self.include_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: chunk.clone(), + }); + } + } +} + +/// The finish reason of a stream that ended without one. +pub(crate) const FINISH_OTHER: FinishReason = FinishReason { + unified: FinishReasonUnified::Other, + raw: None, +}; + +/// The finish reason after an error the stream reported. +pub(crate) const FINISH_ERROR: FinishReason = FinishReason { + unified: FinishReasonUnified::Error, + raw: None, +}; + +/// What a vendor stream writes its parts to. +pub(crate) type Controller = TransformStreamController; + +/// One-line forms of the parts a stream emits without provider metadata. +pub(crate) trait Emit { + fn text_start(&mut self, id: impl Into); + fn text_delta(&mut self, id: impl Into, delta: impl Into); + fn text_end(&mut self, id: impl Into); + fn reasoning_start(&mut self, id: impl Into); + fn reasoning_delta(&mut self, id: impl Into, delta: impl Into); + fn reasoning_end(&mut self, id: impl Into); + /// A client-side tool call's input starts. + fn tool_input_start(&mut self, id: impl Into, tool_name: impl Into); + fn tool_input_delta(&mut self, id: impl Into, delta: impl Into); + fn tool_input_end(&mut self, id: impl Into); +} + +impl Emit for Controller { + fn text_start(&mut self, id: impl Into) { + self.enqueue(StreamPart::TextStart { + id: id.into(), + provider_metadata: None, + }); + } + + fn text_delta(&mut self, id: impl Into, delta: impl Into) { + self.enqueue(StreamPart::TextDelta { + id: id.into(), + delta: delta.into(), + provider_metadata: None, + }); + } + + fn text_end(&mut self, id: impl Into) { + self.enqueue(StreamPart::TextEnd { + id: id.into(), + provider_metadata: None, + }); + } + + fn reasoning_start(&mut self, id: impl Into) { + self.enqueue(StreamPart::ReasoningStart { + id: id.into(), + provider_metadata: None, + }); + } + + fn reasoning_delta(&mut self, id: impl Into, delta: impl Into) { + self.enqueue(StreamPart::ReasoningDelta { + id: id.into(), + delta: delta.into(), + provider_metadata: None, + }); + } + + fn reasoning_end(&mut self, id: impl Into) { + self.enqueue(StreamPart::ReasoningEnd { + id: id.into(), + provider_metadata: None, + }); + } + + fn tool_input_start(&mut self, id: impl Into, tool_name: impl Into) { + self.enqueue(StreamPart::ToolInputStart { + id: id.into(), + tool_name: tool_name.into(), + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: None, + }); + } + + fn tool_input_delta(&mut self, id: impl Into, delta: impl Into) { + self.enqueue(StreamPart::ToolInputDelta { + id: id.into(), + delta: delta.into(), + provider_metadata: None, + }); + } + + fn tool_input_end(&mut self, id: impl Into) { + self.enqueue(StreamPart::ToolInputEnd { + id: id.into(), + provider_metadata: None, + }); + } +} + +/// Driving a stream's `Transformer` synchronously, without HTTP or SSE. +#[cfg(test)] +pub(crate) mod testing { + use aimux_core::error::AiMuxError; + use aimux_core::stream_part::StreamPart; + use aimux_provider_utils::Transformer; + + use super::Controller; + + /// What `pipe_through` would emit for `input`, one line per item: `start`, + /// `transform` per chunk until the stream is closed, then `flush` unless + /// it was. + pub(crate) fn run>( + mut stream: T, + input: impl IntoIterator, + ) -> Vec { + let mut controller = Controller::new(); + stream.start(&mut controller); + for chunk in input { + if controller.is_closed() { + break; + } + stream.transform(chunk, &mut controller); + } + if !controller.is_closed() { + stream.flush(&mut controller); + } + controller.drain().map(describe).collect() + } + + /// The kind of an item and the fields a test tells parts apart by. + pub(crate) fn describe(item: Result) -> String { + let part = match item { + Ok(part) => part, + Err(error) => return format!("err {error}"), + }; + match part { + StreamPart::StreamStart { .. } => "stream-start".to_string(), + StreamPart::ResponseMetadata(metadata) => { + format!("response-metadata {}", metadata.id.unwrap_or_default()) + } + StreamPart::TextStart { id, .. } => format!("text-start {id}"), + StreamPart::TextDelta { id, delta, .. } => format!("text-delta {id} {delta}"), + StreamPart::TextEnd { id, .. } => format!("text-end {id}"), + StreamPart::ReasoningStart { id, .. } => format!("reasoning-start {id}"), + StreamPart::ReasoningDelta { id, delta, .. } => { + format!("reasoning-delta {id} {delta}") + } + StreamPart::ReasoningEnd { id, .. } => format!("reasoning-end {id}"), + StreamPart::ToolInputStart { id, tool_name, .. } => { + format!("tool-input-start {id} {tool_name}") + } + StreamPart::ToolInputDelta { id, delta, .. } => { + format!("tool-input-delta {id} {delta}") + } + StreamPart::ToolInputEnd { id, .. } => format!("tool-input-end {id}"), + StreamPart::ToolCall(call) => format!( + "tool-call {} {} {}", + call.tool_call_id, call.tool_name, call.input + ), + StreamPart::Error { error } => format!("error {error}"), + StreamPart::Finish { finish_reason, .. } => { + format!("finish {:?}", finish_reason.unified) + } + other => format!("{other:?}"), + } + } +} diff --git a/aimux-provider-utils/src/streaming_tool_call_argument_state.rs b/aimux-providers/src/stream/tool_call_arguments.rs similarity index 93% rename from aimux-provider-utils/src/streaming_tool_call_argument_state.rs rename to aimux-providers/src/stream/tool_call_arguments.rs index adf24c4b6..7a7fe98d1 100644 --- a/aimux-provider-utils/src/streaming_tool_call_argument_state.rs +++ b/aimux-providers/src/stream/tool_call_arguments.rs @@ -7,14 +7,14 @@ /// Whether `value` (ignoring leading whitespace) starts with `{` or `[`. #[must_use] -pub fn starts_with_structured_value(value: Option<&str>) -> bool { +pub(super) fn starts_with_structured_value(value: Option<&str>) -> bool { value.is_some_and(|v| matches!(v.trim_start().as_bytes().first(), Some(b'{' | b'['))) } /// Incrementally tracks whether streamed tool-call arguments contain a /// complete structured JSON value. #[derive(Debug, Clone)] -pub struct StreamingToolCallArgumentState { +pub(super) struct StreamingToolCallArgumentState { structure: Structure, } @@ -37,7 +37,7 @@ enum Structure { impl StreamingToolCallArgumentState { /// Create a state seeded with `initial_value`. #[must_use] - pub fn new(initial_value: &str) -> Self { + pub(super) fn new(initial_value: &str) -> Self { let mut state = Self { structure: Structure::Undetermined, }; @@ -47,12 +47,12 @@ impl StreamingToolCallArgumentState { /// `true` once a top-level `{…}` / `[…]` value has been closed. #[must_use] - pub fn has_complete_structured_value(&self) -> bool { + pub(super) fn has_complete_structured_value(&self) -> bool { matches!(self.structure, Structure::Complete) } /// Feed the next argument fragment. - pub fn append(&mut self, delta: &str) { + pub(super) fn append(&mut self, delta: &str) { let mut rest = delta; if matches!(self.structure, Structure::Undetermined) { rest = rest.trim_start(); @@ -115,8 +115,7 @@ impl StreamingToolCallArgumentState { } /// Port of `streaming-tool-call-argument-state.test.ts` from -/// `@ai-sdk/provider-utils`. Upstream does not export the argument state -/// either, so it is tested here rather than through the crate root. +/// `@ai-sdk/provider-utils`. #[cfg(test)] mod tests { use super::*; diff --git a/aimux-providers/src/vertex/model.rs b/aimux-providers/src/vertex/model.rs index 1876d0b97..3a13cff54 100644 --- a/aimux-providers/src/vertex/model.rs +++ b/aimux-providers/src/vertex/model.rs @@ -4,10 +4,11 @@ //! [`crate::google::types`] response types. Only the endpoint construction and //! authentication differ from the public Gemini API provider. +use std::collections::HashSet; + use aimux_core::tool::RawToolCall; use aimux_core::tool::ToolResult; use async_trait::async_trait; -use futures::StreamExt; use serde_json::{Value, json}; use aimux_core::error::AiMuxError; @@ -20,12 +21,15 @@ use aimux_core::stream_part::StreamPart; use aimux_core::types::{ FinishReason, FinishReasonUnified, ProviderMetadata, ResponseMetadata, Usage, Warning, }; +use aimux_provider_utils::{Transformer, pipe_through}; + +use crate::stream::{Controller, Emit, FINISH_ERROR, FINISH_OTHER, StreamContext}; use crate::google::convert::{ build_vertex_request_body, code_execution_tool_name, extract_sources, parse_finish_reason, prepare_all_tools, validate_call_options_for_namespace, }; -use crate::google::types::{Candidate, GenerateContentResponse, GoogleStreamEvent}; +use crate::google::types::{Candidate, GenerateContentResponse, GoogleStreamEvent, StreamChunk}; use crate::google::utils::{GoogleJsonAccumulator, PartialArg}; use crate::google::options::{GOOGLE, Namespace}; @@ -171,7 +175,6 @@ impl LanguageModel for VertexModel { let options = &options; let code_execution_tool_name = code_execution_tool_name(options.tools.as_deref()); let (body, warnings) = vertex_request_body(&self.model_id, options, true)?; - let include_raw_chunks = options.include_raw_chunks == Some(true); let exchange = self.config.exchange(options.headers.as_ref()).await?; let body = exchange.transform_body(body); let endpoint = self.stream_endpoint(exchange.base_url()); @@ -185,467 +188,580 @@ impl LanguageModel for VertexModel { let response_headers = resp.response_headers; let sse_stream = resp.value; - let stream_error_url = endpoint; - let stream_request_body = body.clone(); - let stream_error_headers = response_headers.clone(); + let stream = pipe_through( + sse_stream, + VertexModelStream { + ctx: StreamContext::new(options, warnings, endpoint, &body, &response_headers), + code_execution_tool_name, + ..VertexModelStream::default() + }, + ); - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), + response: Some(aimux_core::shared::StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} - let mut sse_stream = sse_stream; - let mut text_id: Option = None; - let mut reasoning_id: Option = None; - let mut prompt_blocked = false; - let mut block_counter = 0usize; - let mut final_usage: Usage = Usage::default(); - let mut final_finish_reason: Option = None; - let mut has_tool_calls = false; - let mut active_tool_calls: Vec<(RawToolCall, GoogleJsonAccumulator)> = Vec::new(); - let mut response_metadata_emitted = false; - let mut stream_errored = false; +/// The `TransformStream` of `GoogleLanguageModel.doStream` as configured by +/// `@ai-sdk/google-vertex`. +#[derive(Default)] +struct VertexModelStream { + ctx: StreamContext, + code_execution_tool_name: String, + text_id: Option, + reasoning_id: Option, + block_counter: usize, + final_usage: Usage, + final_finish_reason: Option, + has_tool_calls: bool, + active_tool_calls: Vec<(RawToolCall, GoogleJsonAccumulator)>, + response_metadata_emitted: bool, + stream_errored: bool, + prompt_blocked: bool, + // Provider-metadata accumulators (mirrors TS `lastGroundingMetadata` / + // `lastUrlContextMetadata` + the finishReason-chunk snapshot). + last_grounding_metadata: Option, + last_url_context_metadata: Option, + last_prompt_feedback: Option, + last_safety_ratings: Option, + last_finish_message: Option, + last_usage_metadata_value: Option, + // Source dedup across chunks (url sources only). + emitted_source_urls: HashSet, + source_id: usize, + // Associates code-execution results / server tool responses with their + // preceding call. + last_code_execution_tool_call_id: Option, + last_server_tool_call_id: Option, +} - // Provider-metadata accumulators (mirrors TS `lastGroundingMetadata` / - // `lastUrlContextMetadata` + the finishReason-chunk snapshot). Same - // behaviour as the public Gemini API provider. - let mut last_grounding_metadata: Option = None; - let mut last_url_context_metadata: Option = None; - let mut last_prompt_feedback: Option = None; - let mut last_safety_ratings: Option = None; - let mut last_finish_message: Option = None; - let mut last_usage_metadata_value: Option = None; +impl VertexModelStream { + fn transform_chunk(&mut self, chunk: Box, controller: &mut Controller) { + if !self.response_metadata_emitted + && let Some(id) = &chunk.response_id + { + self.response_metadata_emitted = true; + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id: Some(id.clone()), + timestamp: None, + model_id: None, + })); + } - // Source dedup across chunks (url sources only). - let mut emitted_source_urls: std::collections::HashSet = - std::collections::HashSet::new(); - let mut source_id = 0usize; + if !self.prompt_blocked + && let Some(pf) = &chunk.prompt_feedback + { + self.last_prompt_feedback = Some(pf.clone()); + if let Some(reason) = confirmed_block_reason(Some(pf)) { + self.prompt_blocked = true; + self.final_finish_reason = Some(FinishReason { + unified: FinishReasonUnified::ContentFilter, + raw: Some(reason), + }); + } + } - // Associates code-execution results / server tool responses with - // their preceding call. - let mut last_code_execution_tool_call_id: Option = None; - let mut last_server_tool_call_id: Option = None; + let Some(candidates) = chunk.candidates else { + return; + }; + let Some(candidate) = candidates.into_iter().next() else { + return; + }; - while let Some(event) = sse_stream.next().await { - if stream_errored { - break; - } - let event = match event { - Ok(raw) => { - if include_raw_chunks { yield Ok(StreamPart::Raw { raw_value: raw.clone() }); } - if let Some(usage) = raw.get("usageMetadata") { - final_usage = vertex_usage(Some(usage)); - last_usage_metadata_value = Some(usage.clone()); - } - serde_json::from_value::(raw) - .map_err(|error| AiMuxError::InvalidResponseData(error.to_string())) - } - Err(error) => Err(error), - }; - match event { - Ok(GoogleStreamEvent::Chunk(chunk)) => { + if let Some(gm) = &candidate.grounding_metadata { + self.last_grounding_metadata = Some(gm.clone()); + } + if let Some(ucm) = &candidate.url_context_metadata { + self.last_url_context_metadata = Some(ucm.clone()); + } - if !response_metadata_emitted - && let Some(id) = &chunk.response_id { - response_metadata_emitted = true; - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { - id: Some(id.clone()), - timestamp: None, - model_id: None, - })); - } + if let Some(sr) = &candidate.safety_ratings { + self.last_safety_ratings = serde_json::to_value(sr).ok(); + } + if let Some(fm) = &candidate.finish_message { + self.last_finish_message = Some(json!(fm)); + } - if !prompt_blocked && let Some(pf) = &chunk.prompt_feedback { - last_prompt_feedback = Some(pf.clone()); - if let Some(reason) = confirmed_block_reason(Some(pf)) { - prompt_blocked = true; - final_finish_reason = Some(FinishReason { unified: FinishReasonUnified::ContentFilter, raw: Some(reason) }); - } - } + if self.prompt_blocked { + return; + } - let Some(candidates) = chunk.candidates else { - continue; - }; - let Some(candidate) = candidates.into_iter().next() else { - continue; - }; + // Extract url sources from this chunk's grounding metadata + // (deduplicated across chunks; document sources are not + // emitted in the stream, matching TS). + let chunk_sources = + extract_sources(candidate.grounding_metadata.as_ref(), &mut self.source_id); + for src in chunk_sources { + if let GenerateContent::Source(Source::Url { + url, + id, + title, + provider_metadata: None, + }) = src + && self.emitted_source_urls.insert(url.clone()) + { + controller.enqueue(StreamPart::Source(Source::Url { + id, + url, + title, + provider_metadata: None, + })); + } + } - if let Some(gm) = &candidate.grounding_metadata { - last_grounding_metadata = Some(gm.clone()); + if let Some(parts) = candidate.content.as_ref().and_then(|c| c.parts.as_ref()) { + for part in parts + .iter() + .filter(|part| part.get("functionCall").is_none()) + .chain( + parts + .iter() + .filter(|part| part.get("functionCall").is_some()), + ) + { + if let Some(text) = part.get("text").and_then(Value::as_str) { + let metadata = thought_metadata(part); + if text.is_empty() { + if metadata.is_some() + && let Some(id) = &self.text_id + { + controller.enqueue(StreamPart::TextDelta { + id: id.clone(), + delta: String::new(), + provider_metadata: metadata, + }); } - if let Some(ucm) = &candidate.url_context_metadata { - last_url_context_metadata = Some(ucm.clone()); + } else if part.get("thought") == Some(&Value::Bool(true)) { + if let Some(id) = self.text_id.take() { + controller.text_end(id); } - - if let Some(sr) = &candidate.safety_ratings { - last_safety_ratings = serde_json::to_value(sr).ok(); + if self.reasoning_id.is_none() { + let id = format!("{}", self.block_counter); + self.block_counter += 1; + self.reasoning_id = Some(id.clone()); + controller.enqueue(StreamPart::ReasoningStart { + id, + provider_metadata: metadata.clone(), + }); } - if let Some(fm) = &candidate.finish_message { - last_finish_message = Some(json!(fm)); + controller.enqueue(StreamPart::ReasoningDelta { + id: self.reasoning_id.clone().unwrap(), + delta: text.to_string(), + provider_metadata: metadata, + }); + } else { + if let Some(id) = self.reasoning_id.take() { + controller.reasoning_end(id); } - - if prompt_blocked { continue; } - - // Extract url sources from this chunk's grounding metadata - // (deduplicated across chunks; document sources are not - // emitted in the stream, matching TS). - let chunk_sources = - extract_sources(candidate.grounding_metadata.as_ref(), &mut source_id); - for src in chunk_sources { - if let GenerateContent::Source(Source::Url { - url, + if self.text_id.is_none() { + let id = format!("{}", self.block_counter); + self.block_counter += 1; + self.text_id = Some(id.clone()); + controller.enqueue(StreamPart::TextStart { id, - title, - provider_metadata: None, - }) = src - && emitted_source_urls.insert(url.clone()) { - yield Ok(StreamPart::Source(Source::Url { - id, - url, - title, - provider_metadata: None, - })); - } + provider_metadata: metadata.clone(), + }); } - - if let Some(parts) = - candidate.content.as_ref().and_then(|c| c.parts.as_ref()) + controller.enqueue(StreamPart::TextDelta { + id: self.text_id.clone().unwrap(), + delta: text.to_string(), + provider_metadata: metadata, + }); + } + } else if let Some(file) = inline_file(part) { + if let Some(id) = self.text_id.take() { + controller.text_end(id); + } + if let Some(id) = self.reasoning_id.take() { + controller.reasoning_end(id); + } + controller.enqueue(if part.get("thought") == Some(&Value::Bool(true)) { + StreamPart::ReasoningFile(file) + } else { + StreamPart::File(file) + }); + } else if let Some(fc) = part.get("functionCall") { + let name = fc.get("name").and_then(Value::as_str); + let args = fc.get("args").filter(|value| !value.is_null()); + let partial_args = fc.get("partialArgs").filter(|value| !value.is_null()); + let will_continue = + fc.get("willContinue").and_then(Value::as_bool) == Some(true); + let thought_signature = part + .get("thoughtSignature") + .and_then(Value::as_str) + .filter(|signature| !signature.is_empty()) + .map(str::to_owned); + let tool_metadata = thought_signature.as_deref().map(|signature| { + vertex_provider_metadata(json!({ "thoughtSignature": signature })) + }); + if partial_args.is_some() || (name.is_some() && will_continue) { + if let Some(name) = name { + let id = fc + .get("id") + .and_then(Value::as_str) + .filter(|id| !id.is_empty()) + .map(str::to_owned) + .unwrap_or_else(aimux_provider_utils::generate_id); + controller.enqueue(StreamPart::ToolInputStart { + id: id.clone(), + tool_name: name.to_owned(), + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: tool_metadata.clone(), + }); + self.active_tool_calls.push(( + RawToolCall { + tool_call_id: id, + tool_name: name.to_owned(), + input: String::new(), + provider_executed: None, + dynamic: None, + provider_metadata: tool_metadata.clone(), + }, + GoogleJsonAccumulator::new(), + )); + } + if let Some(partial_args) = partial_args.and_then(Value::as_array) + && let Some((call, accumulator)) = self.active_tool_calls.last_mut() { - for part in parts.iter().filter(|part| part.get("functionCall").is_none()) - .chain(parts.iter().filter(|part| part.get("functionCall").is_some())) { - if let Some(text) = part.get("text").and_then(Value::as_str) { - let metadata = thought_metadata(part); - if text.is_empty() { - if metadata.is_some() && let Some(id) = &text_id { - yield Ok(StreamPart::TextDelta { id: id.clone(), delta: String::new(), provider_metadata: metadata }); - } - } else if part.get("thought") == Some(&Value::Bool(true)) { - if let Some(id) = text_id.take() { yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); } - if reasoning_id.is_none() { - let id = format!("{block_counter}"); block_counter += 1; - reasoning_id = Some(id.clone()); - yield Ok(StreamPart::ReasoningStart { id, provider_metadata: metadata.clone() }); - } - yield Ok(StreamPart::ReasoningDelta { id: reasoning_id.clone().unwrap(), delta: text.to_string(), provider_metadata: metadata }); - } else { - if let Some(id) = reasoning_id.take() { yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); } - if text_id.is_none() { - let id = format!("{block_counter}"); block_counter += 1; - text_id = Some(id.clone()); - yield Ok(StreamPart::TextStart { id, provider_metadata: metadata.clone() }); - } - yield Ok(StreamPart::TextDelta { id: text_id.clone().unwrap(), delta: text.to_string(), provider_metadata: metadata }); - } - } else if let Some(file) = inline_file(part) { - if let Some(id) = text_id.take() { yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); } - if let Some(id) = reasoning_id.take() { yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); } - yield Ok(if part.get("thought") == Some(&Value::Bool(true)) { StreamPart::ReasoningFile(file) } else { StreamPart::File(file) }); - } else if let Some(fc) = part.get("functionCall") { - let name = fc.get("name").and_then(Value::as_str); - let args = fc.get("args").filter(|value| !value.is_null()); - let partial_args = fc.get("partialArgs").filter(|value| !value.is_null()); - let will_continue = fc.get("willContinue").and_then(Value::as_bool) == Some(true); - let thought_signature = part.get("thoughtSignature") - .and_then(Value::as_str).filter(|signature| !signature.is_empty()) - .map(str::to_owned); - let tool_metadata = thought_signature.as_deref().map(|signature| { - vertex_provider_metadata(json!({ "thoughtSignature": signature })) - }); - if partial_args.is_some() || (name.is_some() && will_continue) { - if let Some(name) = name { - let id = fc.get("id").and_then(Value::as_str) - .filter(|id| !id.is_empty()).map(str::to_owned) - .unwrap_or_else(aimux_provider_utils::generate_id); - yield Ok(StreamPart::ToolInputStart { - id: id.clone(), tool_name: name.to_owned(), - provider_executed: None, dynamic: None, title: None, - provider_metadata: tool_metadata.clone(), - }); - active_tool_calls.push((RawToolCall { - tool_call_id: id, tool_name: name.to_owned(), input: String::new(), - provider_executed: None, dynamic: None, - provider_metadata: tool_metadata.clone(), - }, GoogleJsonAccumulator::new())); - } - if let Some(partial_args) = partial_args.and_then(Value::as_array) - && let Some((call, accumulator)) = active_tool_calls.last_mut() { - let partial_args: Vec = partial_args.iter().map(|arg| PartialArg { - json_path: arg.get("jsonPath").and_then(Value::as_str).unwrap_or_default().to_owned(), - string_value: arg.get("stringValue").and_then(Value::as_str).map(str::to_owned), - number_value: arg.get("numberValue").and_then(Value::as_f64), - bool_value: arg.get("boolValue").and_then(Value::as_bool), - null_value: arg.get("nullValue").map(|_| ()), - will_continue: arg.get("willContinue").and_then(Value::as_bool), - }).collect(); - match accumulator.process_partial_args(&partial_args) { - Ok(result) => { - if !result.text_delta.is_empty() { - yield Ok(StreamPart::ToolInputDelta { - id: call.tool_call_id.clone(), delta: result.text_delta, - provider_metadata: tool_metadata, - }); - } - } - Err(error) => { yield Ok(StreamPart::Error { error }); return; } - } - if !will_continue && partial_args.iter().all(|arg| arg.will_continue != Some(true)) { - for event in finish_streaming_tool_call(&mut active_tool_calls) { yield Ok(event); } - has_tool_calls = true; - } - } - continue; - } - if name.is_none() && args.is_none() && partial_args.is_none() - && fc.get("willContinue").is_none_or(Value::is_null) { - if !active_tool_calls.is_empty() { - for event in finish_streaming_tool_call(&mut active_tool_calls) { yield Ok(event); } - has_tool_calls = true; - } - continue; - } - let Some(name) = name else { continue; }; - let id = fc.get("id").and_then(Value::as_str) - .filter(|id| !id.is_empty()).map(str::to_owned) - .unwrap_or_else(aimux_provider_utils::generate_id); - let args = args.cloned().unwrap_or(json!({})); - yield Ok(StreamPart::ToolInputStart { - id: id.clone(), - tool_name: name.to_string(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: tool_metadata.clone(), - }); - let args_str = args.as_str().map(str::to_string).unwrap_or_else(|| args.to_string()); - if fc.get("args").is_some_and(|value| !value.is_null()) { yield Ok(StreamPart::ToolInputDelta { - id: id.clone(), - delta: args_str.clone(), - provider_metadata: tool_metadata.clone(), - }); - } - yield Ok(StreamPart::ToolInputEnd { - id: id.clone(), - provider_metadata: tool_metadata.clone(), - }); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id, - tool_name: name.to_string(), - input: args_str, - provider_executed: None, - dynamic: None, - provider_metadata: tool_metadata, - })); - has_tool_calls = true; - } else if let Some(inline) = part.get("inlineData") { - if let Some(id) = text_id.take() { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None }); - } - if let Some(id) = reasoning_id.take() { - yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); - } - if let (Some(data), Some(media_type)) = ( - inline.get("data").and_then(Value::as_str), - inline.get("mimeType").and_then(Value::as_str), - ) { - let file = LanguageModelFile { - data: GeneratedFileData::Data { data: FileBytes::Base64(data.to_string()) }, - media_type: media_type.to_string(), - provider_metadata: thought_metadata(part), - }; - yield Ok(if part.get("thought").and_then(Value::as_bool) == Some(true) { - StreamPart::ReasoningFile(file) - } else { - StreamPart::File(file) + let partial_args: Vec = partial_args + .iter() + .map(|arg| PartialArg { + json_path: arg + .get("jsonPath") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + string_value: arg + .get("stringValue") + .and_then(Value::as_str) + .map(str::to_owned), + number_value: arg.get("numberValue").and_then(Value::as_f64), + bool_value: arg.get("boolValue").and_then(Value::as_bool), + null_value: arg.get("nullValue").map(|_| ()), + will_continue: arg.get("willContinue").and_then(Value::as_bool), + }) + .collect(); + match accumulator.process_partial_args(&partial_args) { + Ok(result) => { + if !result.text_delta.is_empty() { + controller.enqueue(StreamPart::ToolInputDelta { + id: call.tool_call_id.clone(), + delta: result.text_delta, + provider_metadata: tool_metadata, }); } - } else if let Some(ec) = part.get("executableCode") { - // Provider-executed code execution. - let has_code = ec - .get("code") - .and_then(|v| v.as_str()) - .map(|s| !s.is_empty()) - .unwrap_or(false); - if has_code { - let id = format!("call-{block_counter}"); - block_counter += 1; - last_code_execution_tool_call_id = Some(id.clone()); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id.clone(), - tool_name: code_execution_tool_name.clone(), - input: ec.to_string(), - provider_executed: Some(true), - dynamic: None, - provider_metadata: Some(vertex_server_tool_metadata( - &id, - "code_execution", - None, - )), - })); - // provider-executed → does NOT set has_tool_calls - } - } else if let Some(cer) = part.get("codeExecutionResult") { - // Result corresponds to the most recent - // executableCode part. Gemini may emit - // several results for that one call, so - // retain the association until a new call. - if let Some(call_id) = - last_code_execution_tool_call_id.as_ref() - { - let outcome = - cer.get("outcome").cloned().unwrap_or(json!(null)); - let output = cer - .get("output") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - .unwrap_or_default(); - yield Ok(StreamPart::ToolResult(ToolResult { - tool_call_id: call_id.clone(), - tool_name: code_execution_tool_name.clone(), - result: json!({ "outcome": outcome, "output": output }), - is_error: None, - preliminary: None, - dynamic: None, - provider_metadata: Some(vertex_server_tool_metadata( - call_id, - "code_execution", - None, - )), - })); - } - } else if let Some(tc) = part.get("toolCall") { - // Server-side tool call (provider-executed). - let tool_type = tc - .get("toolType") - .and_then(|v| v.as_str()) - .unwrap_or(""); - let id = tc - .get("id") - .and_then(|v| v.as_str()) - .filter(|id| !id.is_empty()) - .map(std::string::ToString::to_string) - .unwrap_or_else(aimux_provider_utils::generate_id); - block_counter += 1; - last_server_tool_call_id = Some(id.clone()); - let args = tc.get("args").cloned().unwrap_or(json!({})); - let thought_signature = part - .get("thoughtSignature") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - let server_meta = vertex_server_tool_metadata( - &id, - tool_type, - thought_signature.as_deref(), - ); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: id, - tool_name: format!("server:{tool_type}"), - input: args.to_string(), - provider_executed: Some(true), - dynamic: Some(true), - provider_metadata: Some(server_meta), - })); - // provider-executed → does NOT set has_tool_calls - } else if let Some(tr) = part.get("toolResponse") { - // Server-side tool response. - let tool_type = tr - .get("toolType") - .and_then(|v| v.as_str()) - .unwrap_or(""); - let id = last_server_tool_call_id - .take() - .or_else(|| { - tr.get("id") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - }) - .unwrap_or_else(|| format!("call-{block_counter}")); - block_counter += 1; - let response = - tr.get("response").cloned().unwrap_or(json!({})); - let server_meta = vertex_server_tool_metadata( - &id, - tool_type, - part.get("thoughtSignature").and_then(|v| v.as_str()), - ); - yield Ok(StreamPart::ToolResult(ToolResult { - tool_call_id: id, - tool_name: format!("server:{tool_type}"), - result: response, - is_error: None, - preliminary: None, - dynamic: None, - provider_metadata: Some(server_meta), - })); + } + Err(error) => { + controller.enqueue(StreamPart::Error { error }); + controller.terminate(); + return; } } + if !will_continue + && partial_args + .iter() + .all(|arg| arg.will_continue != Some(true)) + { + for event in finish_streaming_tool_call(&mut self.active_tool_calls) + { + controller.enqueue(event); + } + self.has_tool_calls = true; + } } - - if let Some(reason) = candidate.finish_reason.as_deref() { - final_finish_reason = - Some(parse_finish_reason(reason, has_tool_calls)); + continue; + } + if name.is_none() + && args.is_none() + && partial_args.is_none() + && fc.get("willContinue").is_none_or(Value::is_null) + { + if !self.active_tool_calls.is_empty() { + for event in finish_streaming_tool_call(&mut self.active_tool_calls) { + controller.enqueue(event); + } + self.has_tool_calls = true; } + continue; } - Ok(GoogleStreamEvent::Error(error)) => { - yield Ok(StreamPart::Error { - error: crate::google::google_stream_error( - &error.error, - &stream_error_url, - stream_request_body.clone(), - stream_error_headers.clone(), - ), + let Some(name) = name else { + continue; + }; + let id = fc + .get("id") + .and_then(Value::as_str) + .filter(|id| !id.is_empty()) + .map(str::to_owned) + .unwrap_or_else(aimux_provider_utils::generate_id); + let args = args.cloned().unwrap_or(json!({})); + controller.enqueue(StreamPart::ToolInputStart { + id: id.clone(), + tool_name: name.to_string(), + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: tool_metadata.clone(), + }); + let args_str = args + .as_str() + .map(str::to_string) + .unwrap_or_else(|| args.to_string()); + if fc.get("args").is_some_and(|value| !value.is_null()) { + controller.enqueue(StreamPart::ToolInputDelta { + id: id.clone(), + delta: args_str.clone(), + provider_metadata: tool_metadata.clone(), }); - stream_errored = true; - break; } - Err(error) => { - let recoverable = error.is_recoverable_stream_error(); - yield Ok(StreamPart::Error { error }); - if !recoverable { - return; - } + controller.enqueue(StreamPart::ToolInputEnd { + id: id.clone(), + provider_metadata: tool_metadata.clone(), + }); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: id, + tool_name: name.to_string(), + input: args_str, + provider_executed: None, + dynamic: None, + provider_metadata: tool_metadata, + })); + self.has_tool_calls = true; + } else if let Some(inline) = part.get("inlineData") { + if let Some(id) = self.text_id.take() { + controller.text_end(id); + } + if let Some(id) = self.reasoning_id.take() { + controller.reasoning_end(id); + } + if let (Some(data), Some(media_type)) = ( + inline.get("data").and_then(Value::as_str), + inline.get("mimeType").and_then(Value::as_str), + ) { + let file = LanguageModelFile { + data: GeneratedFileData::Data { + data: FileBytes::Base64(data.to_string()), + }, + media_type: media_type.to_string(), + provider_metadata: thought_metadata(part), + }; + controller.enqueue( + if part.get("thought").and_then(Value::as_bool) == Some(true) { + StreamPart::ReasoningFile(file) + } else { + StreamPart::File(file) + }, + ); + } + } else if let Some(ec) = part.get("executableCode") { + // Provider-executed code execution. + let has_code = ec + .get("code") + .and_then(|v| v.as_str()) + .map(|s| !s.is_empty()) + .unwrap_or(false); + if has_code { + let id = format!("call-{}", self.block_counter); + self.block_counter += 1; + self.last_code_execution_tool_call_id = Some(id.clone()); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: id.clone(), + tool_name: self.code_execution_tool_name.clone(), + input: ec.to_string(), + provider_executed: Some(true), + dynamic: None, + provider_metadata: Some(vertex_server_tool_metadata( + &id, + "code_execution", + None, + )), + })); + // provider-executed → does NOT set has_tool_calls } + } else if let Some(cer) = part.get("codeExecutionResult") { + // Result corresponds to the most recent + // executableCode part. Gemini may emit + // several results for that one call, so + // retain the association until a new call. + if let Some(call_id) = self.last_code_execution_tool_call_id.as_ref() { + let outcome = cer.get("outcome").cloned().unwrap_or(json!(null)); + let output = cer + .get("output") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string) + .unwrap_or_default(); + controller.enqueue(StreamPart::ToolResult(ToolResult { + tool_call_id: call_id.clone(), + tool_name: self.code_execution_tool_name.clone(), + result: json!({ "outcome": outcome, "output": output }), + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: Some(vertex_server_tool_metadata( + call_id, + "code_execution", + None, + )), + })); + } + } else if let Some(tc) = part.get("toolCall") { + // Server-side tool call (provider-executed). + let tool_type = tc.get("toolType").and_then(|v| v.as_str()).unwrap_or(""); + let id = tc + .get("id") + .and_then(|v| v.as_str()) + .filter(|id| !id.is_empty()) + .map(std::string::ToString::to_string) + .unwrap_or_else(aimux_provider_utils::generate_id); + self.block_counter += 1; + self.last_server_tool_call_id = Some(id.clone()); + let args = tc.get("args").cloned().unwrap_or(json!({})); + let thought_signature = part + .get("thoughtSignature") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + let server_meta = + vertex_server_tool_metadata(&id, tool_type, thought_signature.as_deref()); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: id, + tool_name: format!("server:{tool_type}"), + input: args.to_string(), + provider_executed: Some(true), + dynamic: Some(true), + provider_metadata: Some(server_meta), + })); + // provider-executed → does NOT set has_tool_calls + } else if let Some(tr) = part.get("toolResponse") { + // Server-side tool response. + let tool_type = tr.get("toolType").and_then(|v| v.as_str()).unwrap_or(""); + let id = self + .last_server_tool_call_id + .take() + .or_else(|| { + tr.get("id") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string) + }) + .unwrap_or_else(|| format!("call-{}", self.block_counter)); + self.block_counter += 1; + let response = tr.get("response").cloned().unwrap_or(json!({})); + let server_meta = vertex_server_tool_metadata( + &id, + tool_type, + part.get("thoughtSignature").and_then(|v| v.as_str()), + ); + controller.enqueue(StreamPart::ToolResult(ToolResult { + tool_call_id: id, + tool_name: format!("server:{tool_type}"), + result: response, + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: Some(server_meta), + })); } } + } + + if let Some(reason) = candidate.finish_reason.as_deref() { + self.final_finish_reason = Some(parse_finish_reason(reason, self.has_tool_calls)); + } + } +} + +impl Transformer for VertexModelStream { + type Input = Result; + type Output = StreamPart; - if let Some(id) = text_id.take() { - yield Ok(StreamPart::TextEnd { id, provider_metadata: None}); + fn start(&mut self, controller: &mut Controller) { + self.ctx.start(controller); + } + + fn transform(&mut self, event: Result, controller: &mut Controller) { + let event = match event { + Ok(raw) => { + self.ctx.raw(&raw, controller); + if let Some(usage) = raw.get("usageMetadata") { + self.final_usage = vertex_usage(Some(usage)); + self.last_usage_metadata_value = Some(usage.clone()); + } + serde_json::from_value::(raw) + .map_err(|error| AiMuxError::InvalidResponseData(error.to_string())) } + Err(error) => Err(error), + }; - if let Some(id) = reasoning_id.take() { - yield Ok(StreamPart::ReasoningEnd { id, provider_metadata: None }); + match event { + Ok(GoogleStreamEvent::Chunk(chunk)) => self.transform_chunk(chunk, controller), + Ok(GoogleStreamEvent::Error(error)) => { + controller.enqueue(StreamPart::Error { + error: crate::google::google_stream_error( + &error.error, + &self.ctx.url, + self.ctx.request_body.clone(), + self.ctx.response_headers.clone(), + ), + }); + // The rest of the body is not read: the stream finishes now. + self.stream_errored = true; + self.finish(controller); + controller.terminate(); + } + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + controller.enqueue(StreamPart::Error { error }); + if !recoverable { + controller.terminate(); + } } - let provider_metadata = Some(vertex_provider_metadata(json!({ - "promptFeedback": last_prompt_feedback, - "groundingMetadata": last_grounding_metadata, - "urlContextMetadata": last_url_context_metadata, - "safetyRatings": last_safety_ratings, - "serviceTier": last_usage_metadata_value.as_ref().and_then(|usage| usage.get("serviceTier")), - "usageMetadata": last_usage_metadata_value, - "finishMessage": last_finish_message, - }))); + } + } - yield Ok(StreamPart::Finish { - finish_reason: if stream_errored { - FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - } - } else { - final_finish_reason.unwrap_or(FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }) - }, - usage: if stream_errored { Usage::default() } else { final_usage }, - provider_metadata, - }); - }; + fn flush(mut self, controller: &mut Controller) { + self.finish(controller); + } +} - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - response: Some(aimux_core::shared::StreamResponseInfo { - headers: Some(response_headers), - }), - }) +impl VertexModelStream { + /// Close the open segments and emit `Finish`: at the end of the body, or + /// right after an error event. + fn finish(&mut self, controller: &mut Controller) { + if let Some(id) = self.text_id.take() { + controller.text_end(id); + } + + if let Some(id) = self.reasoning_id.take() { + controller.reasoning_end(id); + } + let provider_metadata = Some(vertex_provider_metadata(json!({ + "promptFeedback": self.last_prompt_feedback, + "groundingMetadata": self.last_grounding_metadata, + "urlContextMetadata": self.last_url_context_metadata, + "safetyRatings": self.last_safety_ratings, + "serviceTier": self.last_usage_metadata_value.as_ref().and_then(|usage| usage.get("serviceTier")), + "usageMetadata": self.last_usage_metadata_value, + "finishMessage": self.last_finish_message, + }))); + + controller.enqueue(StreamPart::Finish { + finish_reason: if self.stream_errored { + FINISH_ERROR + } else { + self.final_finish_reason.take().unwrap_or(FINISH_OTHER) + }, + usage: if self.stream_errored { + Usage::default() + } else { + std::mem::take(&mut self.final_usage) + }, + provider_metadata, + }); } } diff --git a/aimux-providers/src/xai/responses/mod.rs b/aimux-providers/src/xai/responses/mod.rs index 615fd587a..5c4f3700b 100644 --- a/aimux-providers/src/xai/responses/mod.rs +++ b/aimux-providers/src/xai/responses/mod.rs @@ -20,6 +20,9 @@ use aimux_core::tool::ToolResult; use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; +use aimux_provider_utils::{Transformer, pipe_through}; + +use crate::stream::{Controller, Emit, FINISH_OTHER, StreamContext}; use async_trait::async_trait; use futures::StreamExt; use serde_json::{Value, json}; @@ -533,374 +536,494 @@ impl LanguageModel for XaiResponsesModel { response_headers.clone(), )); } - let stream_error_url = endpoint; - let stream_request_body = body.clone(); - let stream_response_headers = response_headers.clone(); - let include_raw_chunks = options.include_raw_chunks == Some(true); - - let stream = async_stream::stream! { - yield Ok(StreamPart::StreamStart { warnings }); - - let mut final_usage: Option = None; - let mut final_finish_reason: FinishReason = FinishReason { - unified: FinishReasonUnified::Other, - raw: None, - }; - let mut cost_in_usd_ticks: Option = None; - let mut service_tier: Option = None; - let mut prompt_cache_key: Option = None; - let mut safety_identifier: Option = None; - let mut has_function_call = false; - let mut is_first_chunk = true; - - // Track content blocks. - let mut content_blocks: Vec = Vec::new(); - let mut seen_tool_calls: std::collections::HashSet = std::collections::HashSet::new(); - let mut active_reasoning: HashMap = HashMap::new(); - - // Track ongoing function calls by output_index. - let mut ongoing_tool_calls: HashMap = HashMap::new(); // output_index -> (tool_call_id, tool_name) - - let mut event_iter = futures::stream::iter(first_event.into_iter()).chain(sse_stream); - while let Some(event) = event_iter.next().await { - match event { - Ok(parsed) => { - if include_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: parsed.clone() }); - } - let event_type = types::event_type(&parsed); - - // ── response.created / response.in_progress ── - if event_type == "response.created" || event_type == "response.in_progress" { - if is_first_chunk { - is_first_chunk = false; - let response = parsed.get("response").cloned().unwrap_or(Value::Null); - let id = response.get("id").and_then(|v| v.as_str()).map(std::string::ToString::to_string); - let model = response.get("model").and_then(|v| v.as_str()).map(std::string::ToString::to_string); - let created_at = response.get("created_at").and_then(serde_json::Value::as_u64); - let timestamp = created_at.and_then(|c| { - aimux_core::Timestamp::from_unix_seconds(c as f64) - }); - yield Ok(StreamPart::ResponseMetadata(ResponseMetadata { id, timestamp, model_id: model })); - } - continue; - } + let stream = pipe_through( + futures::stream::iter(first_event).chain(sse_stream), + XaiResponsesStream { + ctx: StreamContext::new(options, warnings, endpoint, &body, &response_headers), + provider_tool_names, + ..XaiResponsesStream::default() + }, + ); - // ── reasoning_summary_part.added ── - if event_type == "response.reasoning_summary_part.added" { - let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); - let block_id = format!("reasoning-{item_id}"); - if !active_reasoning.contains_key(item_id) { - active_reasoning.insert(item_id.to_string(), ()); - yield Ok(StreamPart::ReasoningStart { - id: block_id, - provider_metadata: Some(crate::xai::options::xai_metadata(json!({ "itemId": item_id }))), - }); - } - continue; - } + Ok(StreamResult { + stream: Box::pin(stream), + request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), + response: Some(aimux_core::shared::StreamResponseInfo { + headers: Some(response_headers), + }), + }) + } +} - // ── reasoning_summary_text.delta ── - if event_type == "response.reasoning_summary_text.delta" { - let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); - let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); - let block_id = format!("reasoning-{item_id}"); - yield Ok(StreamPart::ReasoningDelta { - id: block_id, - delta: delta.to_string(), - provider_metadata: Some(crate::xai::options::xai_metadata(json!({ "itemId": item_id }))), - }); - continue; - } +/// The `TransformStream` of the xAI Responses `doStream`. +#[derive(Default)] +struct XaiResponsesStream { + ctx: StreamContext, + provider_tool_names: HashMap, + final_usage: Option, + /// `other` until the response says otherwise. + final_finish_reason: Option, + cost_in_usd_ticks: Option, + service_tier: Option, + prompt_cache_key: Option, + safety_identifier: Option, + has_function_call: bool, + metadata_emitted: bool, + /// Open text blocks. + content_blocks: Vec, + seen_tool_calls: std::collections::HashSet, + active_reasoning: HashMap, + /// Ongoing function calls: output_index -> (tool_call_id, tool_name). + ongoing_tool_calls: HashMap, + /// Set by an `error` event: later chunks are ignored, the stream still + /// finishes. + finished: bool, +} - // ── reasoning_summary_text.done ── - if event_type == "response.reasoning_summary_text.done" { - continue; - } +impl Transformer for XaiResponsesStream { + type Input = Result; + type Output = StreamPart; - // ── reasoning_text.delta ── - if event_type == "response.reasoning_text.delta" { - let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); - let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); - let block_id = format!("reasoning-{item_id}"); - if !active_reasoning.contains_key(item_id) { - active_reasoning.insert(item_id.to_string(), ()); - yield Ok(StreamPart::ReasoningStart { - id: block_id.clone(), - provider_metadata: Some(crate::xai::options::xai_metadata(json!({ "itemId": item_id }))), - }); - } - yield Ok(StreamPart::ReasoningDelta { - id: block_id, - delta: delta.to_string(), - provider_metadata: Some(crate::xai::options::xai_metadata(json!({ "itemId": item_id }))), - }); - continue; - } + fn start(&mut self, controller: &mut Controller) { + self.ctx.start(controller); + } - // ── reasoning_text.done ── - if event_type == "response.reasoning_text.done" { - continue; - } + fn transform(&mut self, event: Result, controller: &mut Controller) { + self.transform_event(event, controller); + if self.finished { + // An error event ends the response: the rest of the body is not + // read and the stream finishes now. + self.finish(controller); + controller.terminate(); + } + } - // ── output_text.delta ── - if event_type == "response.output_text.delta" { - let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); - let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); - let block_id = format!("text-{item_id}"); - if !content_blocks.contains(&block_id) { - content_blocks.push(block_id.clone()); - yield Ok(StreamPart::TextStart { id: block_id.clone(), provider_metadata: None}); - } - yield Ok(StreamPart::TextDelta { - id: block_id, - delta: delta.to_string(), - provider_metadata: None, - }); - continue; - } + fn flush(mut self, controller: &mut Controller) { + self.finish(controller); + } +} - // ── output_text.done ── - if event_type == "response.output_text.done" { - if let Some(annotations) = parsed.get("annotations").and_then(|v| v.as_array()) { - for ann in annotations { - if ann.get("type").and_then(|v| v.as_str()) == Some("url_citation") { - let url = ann.get("url").and_then(|v| v.as_str()).unwrap_or(""); - let title = ann.get("title").and_then(|v| v.as_str()).unwrap_or(url).to_string(); - yield Ok(StreamPart::Source(Source::Url { - id: generate_source_id(), - url: url.to_string(), - title: Some(title), - provider_metadata: None, - })); - } - } - } - continue; - } +impl XaiResponsesStream { + /// Handle one parsed SSE event. + fn transform_event(&mut self, event: Result, controller: &mut Controller) { + let Self { + ctx: + StreamContext { + include_raw_chunks, + url: ref stream_error_url, + request_body: ref stream_request_body, + response_headers: ref stream_response_headers, + .. + }, + ref provider_tool_names, + ref mut final_usage, + ref mut final_finish_reason, + ref mut cost_in_usd_ticks, + ref mut service_tier, + ref mut prompt_cache_key, + ref mut safety_identifier, + ref mut has_function_call, + ref mut metadata_emitted, + ref mut content_blocks, + ref mut seen_tool_calls, + ref mut active_reasoning, + ref mut ongoing_tool_calls, + ref mut finished, + .. + } = *self; + match event { + Ok(parsed) => { + if include_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: parsed.clone(), + }); + } + let event_type = types::event_type(&parsed); + + // ── response.created / response.in_progress ── + if event_type == "response.created" || event_type == "response.in_progress" { + if !*metadata_emitted { + *metadata_emitted = true; + let response = parsed.get("response").cloned().unwrap_or(Value::Null); + let id = response + .get("id") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + let model = response + .get("model") + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + let created_at = response + .get("created_at") + .and_then(serde_json::Value::as_u64); + let timestamp = created_at + .and_then(|c| aimux_core::Timestamp::from_unix_seconds(c as f64)); + controller.enqueue(StreamPart::ResponseMetadata(ResponseMetadata { + id, + timestamp, + model_id: model, + })); + } + return; + } + + // ── reasoning_summary_part.added ── + if event_type == "response.reasoning_summary_part.added" { + let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); + let block_id = format!("reasoning-{item_id}"); + if !active_reasoning.contains_key(item_id) { + active_reasoning.insert(item_id.to_string(), ()); + controller.enqueue(StreamPart::ReasoningStart { + id: block_id, + provider_metadata: Some(crate::xai::options::xai_metadata( + json!({ "itemId": item_id }), + )), + }); + } + return; + } - // ── output_text.annotation.added ── - if event_type == "response.output_text.annotation.added" { - let annotation = parsed.get("annotation").cloned().unwrap_or(Value::Null); - if annotation.get("type").and_then(|v| v.as_str()) == Some("url_citation") { - let url = annotation.get("url").and_then(|v| v.as_str()).unwrap_or(""); - let title = annotation.get("title").and_then(|v| v.as_str()).unwrap_or(url).to_string(); - yield Ok(StreamPart::Source(Source::Url { + // ── reasoning_summary_text.delta ── + if event_type == "response.reasoning_summary_text.delta" { + let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); + let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); + let block_id = format!("reasoning-{item_id}"); + controller.enqueue(StreamPart::ReasoningDelta { + id: block_id, + delta: delta.to_string(), + provider_metadata: Some(crate::xai::options::xai_metadata( + json!({ "itemId": item_id }), + )), + }); + return; + } + + // ── reasoning_summary_text.done ── + if event_type == "response.reasoning_summary_text.done" { + return; + } + + // ── reasoning_text.delta ── + if event_type == "response.reasoning_text.delta" { + let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); + let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); + let block_id = format!("reasoning-{item_id}"); + if !active_reasoning.contains_key(item_id) { + active_reasoning.insert(item_id.to_string(), ()); + controller.enqueue(StreamPart::ReasoningStart { + id: block_id.clone(), + provider_metadata: Some(crate::xai::options::xai_metadata( + json!({ "itemId": item_id }), + )), + }); + } + controller.enqueue(StreamPart::ReasoningDelta { + id: block_id, + delta: delta.to_string(), + provider_metadata: Some(crate::xai::options::xai_metadata( + json!({ "itemId": item_id }), + )), + }); + return; + } + + // ── reasoning_text.done ── + if event_type == "response.reasoning_text.done" { + return; + } + + // ── output_text.delta ── + if event_type == "response.output_text.delta" { + let item_id = parsed.get("item_id").and_then(|v| v.as_str()).unwrap_or(""); + let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); + let block_id = format!("text-{item_id}"); + if !content_blocks.contains(&block_id) { + content_blocks.push(block_id.clone()); + controller.text_start(block_id.clone()); + } + controller.text_delta(block_id, delta.to_string()); + return; + } + + // ── output_text.done ── + if event_type == "response.output_text.done" { + if let Some(annotations) = parsed.get("annotations").and_then(|v| v.as_array()) + { + for ann in annotations { + if ann.get("type").and_then(|v| v.as_str()) == Some("url_citation") { + let url = ann.get("url").and_then(|v| v.as_str()).unwrap_or(""); + let title = ann + .get("title") + .and_then(|v| v.as_str()) + .unwrap_or(url) + .to_string(); + controller.enqueue(StreamPart::Source(Source::Url { id: generate_source_id(), url: url.to_string(), title: Some(title), provider_metadata: None, })); } - continue; } + } + return; + } - // ── response.done / response.completed / response.incomplete ── - if event_type == "response.done" - || event_type == "response.completed" - || event_type == "response.incomplete" - { - let response = parsed.get("response").cloned().unwrap_or(Value::Null); - if let Some(usage) = response.get("usage") - && let Ok(u) = serde_json::from_value::(usage.clone()) { - let mut converted = convert_xai_responses_usage(&u); - converted.raw = usage.as_object().cloned(); - final_usage = Some(converted); - cost_in_usd_ticks = u.cost_in_usd_ticks; - } + // ── output_text.annotation.added ── + if event_type == "response.output_text.annotation.added" { + let annotation = parsed.get("annotation").cloned().unwrap_or(Value::Null); + if annotation.get("type").and_then(|v| v.as_str()) == Some("url_citation") { + let url = annotation.get("url").and_then(|v| v.as_str()).unwrap_or(""); + let title = annotation + .get("title") + .and_then(|v| v.as_str()) + .unwrap_or(url) + .to_string(); + controller.enqueue(StreamPart::Source(Source::Url { + id: generate_source_id(), + url: url.to_string(), + title: Some(title), + provider_metadata: None, + })); + } + return; + } - service_tier = response.get("service_tier").and_then(Value::as_str).map(str::to_string); - prompt_cache_key = response.get("prompt_cache_key").and_then(Value::as_str).map(str::to_string); - safety_identifier = response.get("safety_identifier").and_then(Value::as_str).map(str::to_string); - - if event_type == "response.incomplete" { - let reason = response - .get("incomplete_details") - .and_then(|d| d.get("reason")) - .and_then(|v| v.as_str()); - final_finish_reason = FinishReason { - unified: reason - .map(map_xai_responses_finish_reason) - .unwrap_or(FinishReasonUnified::Other), - raw: reason.map(std::string::ToString::to_string).or(Some("incomplete".to_string())), - }; - } else if let Some(status) = response.get("status").and_then(|v| v.as_str()) { - final_finish_reason = FinishReason { - unified: if has_function_call { - FinishReasonUnified::ToolCalls - } else { - map_xai_responses_finish_reason(status) - }, - raw: Some(status.to_string()), - }; - } - continue; - } + // ── response.done / response.completed / response.incomplete ── + if event_type == "response.done" + || event_type == "response.completed" + || event_type == "response.incomplete" + { + let response = parsed.get("response").cloned().unwrap_or(Value::Null); + if let Some(usage) = response.get("usage") + && let Ok(u) = + serde_json::from_value::(usage.clone()) + { + let mut converted = convert_xai_responses_usage(&u); + converted.raw = usage.as_object().cloned(); + *final_usage = Some(converted); + *cost_in_usd_ticks = u.cost_in_usd_ticks; + } - // ── response.failed ── - if event_type == "response.failed" { - let response = parsed.get("response").cloned().unwrap_or(Value::Null); - let reason = response - .get("incomplete_details") - .and_then(|d| d.get("reason")) - .and_then(|v| v.as_str()); - final_finish_reason = FinishReason { - unified: reason - .map(map_xai_responses_finish_reason) - .unwrap_or(FinishReasonUnified::Error), - raw: reason.map(std::string::ToString::to_string).or(Some("error".to_string())), - }; - if let Some(usage) = response.get("usage") - && let Ok(u) = serde_json::from_value::(usage.clone()) { - let mut converted = convert_xai_responses_usage(&u); - converted.raw = usage.as_object().cloned(); - final_usage = Some(converted); - } - if response.get("error").is_some_and(|error| !error.is_null()) { - yield Ok(StreamPart::Error { - error: super::xai_stream_error( - &parsed, - &stream_error_url, - stream_request_body.clone(), - stream_response_headers.clone(), - ), - }); - } - continue; - } + *service_tier = response + .get("service_tier") + .and_then(Value::as_str) + .map(str::to_string); + *prompt_cache_key = response + .get("prompt_cache_key") + .and_then(Value::as_str) + .map(str::to_string); + *safety_identifier = response + .get("safety_identifier") + .and_then(Value::as_str) + .map(str::to_string); + + if event_type == "response.incomplete" { + let reason = response + .get("incomplete_details") + .and_then(|d| d.get("reason")) + .and_then(|v| v.as_str()); + *final_finish_reason = Some(FinishReason { + unified: reason + .map(map_xai_responses_finish_reason) + .unwrap_or(FinishReasonUnified::Other), + raw: reason + .map(std::string::ToString::to_string) + .or(Some("incomplete".to_string())), + }); + } else if let Some(status) = response.get("status").and_then(|v| v.as_str()) { + *final_finish_reason = Some(FinishReason { + unified: if *has_function_call { + FinishReasonUnified::ToolCalls + } else { + map_xai_responses_finish_reason(status) + }, + raw: Some(status.to_string()), + }); + } + return; + } - // ── error event ── - if event_type == "error" { - yield Ok(StreamPart::Error { - error: super::xai_stream_error( - &parsed, - &stream_error_url, - stream_request_body.clone(), - stream_response_headers.clone(), - ), - }); - final_finish_reason = FinishReason { - unified: FinishReasonUnified::Error, - raw: Some("error".to_string()), - }; - break; - } + // ── response.failed ── + if event_type == "response.failed" { + let response = parsed.get("response").cloned().unwrap_or(Value::Null); + let reason = response + .get("incomplete_details") + .and_then(|d| d.get("reason")) + .and_then(|v| v.as_str()); + *final_finish_reason = Some(FinishReason { + unified: reason + .map(map_xai_responses_finish_reason) + .unwrap_or(FinishReasonUnified::Error), + raw: reason + .map(std::string::ToString::to_string) + .or(Some("error".to_string())), + }); + if let Some(usage) = response.get("usage") + && let Ok(u) = + serde_json::from_value::(usage.clone()) + { + let mut converted = convert_xai_responses_usage(&u); + converted.raw = usage.as_object().cloned(); + *final_usage = Some(converted); + } + if response.get("error").is_some_and(|error| !error.is_null()) { + controller.enqueue(StreamPart::Error { + error: super::xai_stream_error( + &parsed, + stream_error_url, + stream_request_body.clone(), + stream_response_headers.clone(), + ), + }); + } + return; + } - // ── custom_tool_call_input.delta / .done ── - if event_type == "response.custom_tool_call_input.delta" - || event_type == "response.custom_tool_call_input.done" - { - continue; - } + // ── error event ── + if event_type == "error" { + controller.enqueue(StreamPart::Error { + error: super::xai_stream_error( + &parsed, + stream_error_url, + stream_request_body.clone(), + stream_response_headers.clone(), + ), + }); + *final_finish_reason = Some(FinishReason { + unified: FinishReasonUnified::Error, + raw: Some("error".to_string()), + }); + *finished = true; + return; + } - // ── function_call_arguments.delta ── - if event_type == "response.function_call_arguments.delta" { - let output_index = parsed.get("output_index").and_then(serde_json::Value::as_u64).unwrap_or(0); - let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); - if let Some((tool_call_id, _)) = ongoing_tool_calls.get(&output_index) { - yield Ok(StreamPart::ToolInputDelta { - id: tool_call_id.clone(), - delta: delta.to_string(), - provider_metadata: None, + // ── custom_tool_call_input.delta / .done ── + if event_type == "response.custom_tool_call_input.delta" + || event_type == "response.custom_tool_call_input.done" + { + return; + } + + // ── function_call_arguments.delta ── + if event_type == "response.function_call_arguments.delta" { + let output_index = parsed + .get("output_index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0); + let delta = parsed.get("delta").and_then(|v| v.as_str()).unwrap_or(""); + if let Some((tool_call_id, _)) = ongoing_tool_calls.get(&output_index) { + controller.tool_input_delta(tool_call_id.clone(), delta.to_string()); + } + return; + } + + // ── function_call_arguments.done ── + if event_type == "response.function_call_arguments.done" { + return; + } + + // ── output_item.added / output_item.done ── + if event_type == "response.output_item.added" + || event_type == "response.output_item.done" + || matches!( + event_type, + "response.image_generation_call.in_progress" + | "response.image_generation_call.generating" + | "response.image_generation_call.completed" + ) + { + let item = if event_type.starts_with("response.image_generation_call.") { + json!({ "type": "image_generation_call", "id": parsed.get("item_id") }) + } else { + parsed.get("item").cloned().unwrap_or(Value::Null) + }; + let item_type = item.get("type").and_then(|v| v.as_str()).unwrap_or(""); + let output_index = parsed + .get("output_index") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0); + + // ── reasoning item ── + if item_type == "reasoning" { + if event_type == "response.output_item.done" { + let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); + let block_id = format!("reasoning-{part_id}"); + let encrypted = item + .get("encrypted_content") + .and_then(|v| v.as_str()) + .filter(|v| !v.is_empty()); + + if !active_reasoning.contains_key(part_id) { + active_reasoning.insert(part_id.to_string(), ()); + controller.enqueue(StreamPart::ReasoningStart { + id: block_id.clone(), + provider_metadata: Some(crate::xai::options::xai_metadata( + if part_id.is_empty() { + json!({}) + } else { + json!({ "itemId": part_id }) + }, + )), }); } - continue; - } - // ── function_call_arguments.done ── - if event_type == "response.function_call_arguments.done" { - continue; + let mut meta = json!({}); + if !part_id.is_empty() { + meta["itemId"] = json!(part_id); + } + if let Some(ec) = encrypted { + meta["reasoningEncryptedContent"] = json!(ec); + } + controller.enqueue(StreamPart::ReasoningEnd { + id: block_id, + provider_metadata: Some(crate::xai::options::xai_metadata(json!( + meta + ))), + }); + active_reasoning.remove(part_id); } + return; + } - // ── output_item.added / output_item.done ── - if event_type == "response.output_item.added" - || event_type == "response.output_item.done" - || matches!(event_type, - "response.image_generation_call.in_progress" - | "response.image_generation_call.generating" - | "response.image_generation_call.completed") - { - let item = if event_type.starts_with("response.image_generation_call.") { - json!({ "type": "image_generation_call", "id": parsed.get("item_id") }) - } else { - parsed.get("item").cloned().unwrap_or(Value::Null) - }; - let item_type = item.get("type").and_then(|v| v.as_str()).unwrap_or(""); - let output_index = parsed.get("output_index").and_then(serde_json::Value::as_u64).unwrap_or(0); - - // ── reasoning item ── - if item_type == "reasoning" { - if event_type == "response.output_item.done" { - let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); - let block_id = format!("reasoning-{part_id}"); - let encrypted = item.get("encrypted_content").and_then(|v| v.as_str()).filter(|v| !v.is_empty()); - - if !active_reasoning.contains_key(part_id) { - active_reasoning.insert(part_id.to_string(), ()); - yield Ok(StreamPart::ReasoningStart { - id: block_id.clone(), - provider_metadata: Some(crate::xai::options::xai_metadata(if part_id.is_empty() { json!({}) } else { json!({ "itemId": part_id }) })), - }); - } + // ── file_search_call item ── + if item_type == "file_search_call" || item_type == "image_generation_call" { + let tool_name = resolve_tool_name(item_type, None, provider_tool_names); + let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); + let input = if item_type == "image_generation_call" { + "{}" + } else { + "" + }; + + if !seen_tool_calls.contains(part_id) { + seen_tool_calls.insert(part_id.to_string()); + controller.tool_input_start(part_id.to_string(), tool_name.clone()); + controller.tool_input_delta(part_id.to_string(), input.to_string()); + controller.tool_input_end(part_id.to_string()); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: part_id.to_string(), + tool_name: tool_name.clone(), + input: input.to_string(), + provider_executed: Some(true), + dynamic: None, + provider_metadata: None, + })); + } - let mut meta = json!({}); - if !part_id.is_empty() { - meta["itemId"] = json!(part_id); - } - if let Some(ec) = encrypted { - meta["reasoningEncryptedContent"] = json!(ec); - } - yield Ok(StreamPart::ReasoningEnd { - id: block_id, - provider_metadata: Some(crate::xai::options::xai_metadata(json!(meta))), - }); - active_reasoning.remove(part_id); - } - continue; + if event_type == "response.output_item.done" { + if item_type == "image_generation_call" { + controller.enqueue(StreamPart::ToolResult( + image_generation_result(&item, tool_name), + )); + return; } - - // ── file_search_call item ── - if item_type == "file_search_call" || item_type == "image_generation_call" { - let tool_name = resolve_tool_name(item_type, None, &provider_tool_names); - let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); - let input = if item_type == "image_generation_call" { "{}" } else { "" }; - - if !seen_tool_calls.contains(part_id) { - seen_tool_calls.insert(part_id.to_string()); - yield Ok(StreamPart::ToolInputStart { - id: part_id.to_string(), - tool_name: tool_name.clone(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - yield Ok(StreamPart::ToolInputDelta { - id: part_id.to_string(), - delta: input.to_string(), - provider_metadata: None, - }); - yield Ok(StreamPart::ToolInputEnd { - id: part_id.to_string(), - provider_metadata: None, - }); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: part_id.to_string(), - tool_name: tool_name.clone(), - input: input.to_string(), - provider_executed: Some(true), - dynamic: None, - provider_metadata: None, - })); - } - - if event_type == "response.output_item.done" { - if item_type == "image_generation_call" { - yield Ok(StreamPart::ToolResult(image_generation_result(&item, tool_name))); - continue; - } - let queries = item.get("queries").and_then(|v| v.as_array()).cloned().unwrap_or_default(); - let results = item.get("results").and_then(|v| v.as_array()).map(|arr| { + let queries = item + .get("queries") + .and_then(|v| v.as_array()) + .cloned() + .unwrap_or_default(); + let results = item.get("results").and_then(|v| v.as_array()).map(|arr| { json!(arr.iter().map(|r| json!({ "fileId": r.get("file_id").cloned().unwrap_or(Value::Null), "filename": r.get("filename").cloned().unwrap_or(Value::Null), @@ -908,204 +1031,190 @@ impl LanguageModel for XaiResponsesModel { "text": r.get("text").cloned().unwrap_or(Value::Null), })).collect::>()) }).unwrap_or(Value::Null); - yield Ok(StreamPart::ToolResult(ToolResult { - tool_call_id: part_id.to_string(), - tool_name, - result: json!({ "queries": queries, "results": results }), - is_error: None, - preliminary: None, - dynamic: None, - provider_metadata: None, - })); - } - continue; - } + controller.enqueue(StreamPart::ToolResult(ToolResult { + tool_call_id: part_id.to_string(), + tool_name, + result: json!({ "queries": queries, "results": results }), + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: None, + })); + } + return; + } - // ── Server-side tool calls ── - if matches!( - item_type, - "web_search_call" | "x_search_call" | "code_interpreter_call" - | "code_execution_call" | "view_image_call" | "view_x_video_call" - | "custom_tool_call" | "mcp_call" - ) { - let item_name = item.get("name").and_then(|v| v.as_str()); - let tool_name = resolve_tool_name(item_type, item_name, &provider_tool_names); - let tool_input = get_tool_input(item_type, &item); - let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); + // ── Server-side tool calls ── + if matches!( + item_type, + "web_search_call" + | "x_search_call" + | "code_interpreter_call" + | "code_execution_call" + | "view_image_call" + | "view_x_video_call" + | "custom_tool_call" + | "mcp_call" + ) { + let item_name = item.get("name").and_then(|v| v.as_str()); + let tool_name = + resolve_tool_name(item_type, item_name, provider_tool_names); + let tool_input = get_tool_input(item_type, &item); + let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); + + let should_emit = if item_type == "custom_tool_call" { + event_type == "response.output_item.done" + } else { + !seen_tool_calls.contains(part_id) + }; + + if should_emit && !seen_tool_calls.contains(part_id) { + seen_tool_calls.insert(part_id.to_string()); + controller.tool_input_start(part_id.to_string(), tool_name.clone()); + controller.tool_input_delta(part_id.to_string(), tool_input.clone()); + controller.tool_input_end(part_id.to_string()); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: part_id.to_string(), + tool_name: tool_name.clone(), + input: tool_input, + provider_executed: Some(true), + dynamic: None, + provider_metadata: None, + })); + } - let should_emit = if item_type == "custom_tool_call" { - event_type == "response.output_item.done" + if event_type == "response.output_item.done" { + controller.enqueue(StreamPart::ToolResult(ToolResult { + tool_call_id: part_id.to_string(), + tool_name, + result: if item_type == "web_search_call" { + map_web_search_action( + item.get("action").unwrap_or(&Value::Null), + ) } else { - !seen_tool_calls.contains(part_id) - }; - - if should_emit && !seen_tool_calls.contains(part_id) { - seen_tool_calls.insert(part_id.to_string()); - yield Ok(StreamPart::ToolInputStart { - id: part_id.to_string(), - tool_name: tool_name.clone(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - yield Ok(StreamPart::ToolInputDelta { - id: part_id.to_string(), - delta: tool_input.clone(), - provider_metadata: None, - }); - yield Ok(StreamPart::ToolInputEnd { - id: part_id.to_string(), - provider_metadata: None, - }); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: part_id.to_string(), - tool_name: tool_name.clone(), - input: tool_input, - provider_executed: Some(true), - dynamic: None, - provider_metadata: None, - })); - } + json!({}) + }, + is_error: None, + preliminary: None, + dynamic: None, + provider_metadata: None, + })); + } + return; + } - if event_type == "response.output_item.done" { - yield Ok(StreamPart::ToolResult(ToolResult { - tool_call_id: part_id.to_string(), - tool_name, - result: if item_type == "web_search_call" { - map_web_search_action(item.get("action").unwrap_or(&Value::Null)) - } else { - json!({}) - }, - is_error: None, - preliminary: None, - dynamic: None, - provider_metadata: None, - })); + // ── message item ── + if item_type == "message" { + let content_parts = item + .get("content") + .and_then(|v| v.as_array()) + .cloned() + .unwrap_or_default(); + for cp in &content_parts { + if let Some(text) = cp.get("text").and_then(|v| v.as_str()) + && !text.is_empty() + { + let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); + let block_id = format!("text-{part_id}"); + if !content_blocks.contains(&block_id) { + content_blocks.push(block_id.clone()); + controller.text_start(block_id.clone()); + controller.text_delta(block_id, text.to_string()); } - continue; } - - // ── message item ── - if item_type == "message" { - let content_parts = item - .get("content") - .and_then(|v| v.as_array()) - .cloned() - .unwrap_or_default(); - for cp in &content_parts { - if let Some(text) = cp.get("text").and_then(|v| v.as_str()) - && !text.is_empty() { - let part_id = item.get("id").and_then(|v| v.as_str()).unwrap_or(""); - let block_id = format!("text-{part_id}"); - if !content_blocks.contains(&block_id) { - content_blocks.push(block_id.clone()); - yield Ok(StreamPart::TextStart { id: block_id.clone(), provider_metadata: None}); - yield Ok(StreamPart::TextDelta { - id: block_id, - delta: text.to_string(), - provider_metadata: None, - }); - } - } - if let Some(annotations) = cp.get("annotations").and_then(|v| v.as_array()) { - for ann in annotations { - if ann.get("type").and_then(|v| v.as_str()) == Some("url_citation") { - let url = ann.get("url").and_then(|v| v.as_str()).unwrap_or(""); - let title = ann.get("title").and_then(|v| v.as_str()).unwrap_or(url).to_string(); - yield Ok(StreamPart::Source(Source::Url { - id: generate_source_id(), - url: url.to_string(), - title: Some(title), - provider_metadata: None, - })); - } - } + if let Some(annotations) = + cp.get("annotations").and_then(|v| v.as_array()) + { + for ann in annotations { + if ann.get("type").and_then(|v| v.as_str()) + == Some("url_citation") + { + let url = + ann.get("url").and_then(|v| v.as_str()).unwrap_or(""); + let title = ann + .get("title") + .and_then(|v| v.as_str()) + .unwrap_or(url) + .to_string(); + controller.enqueue(StreamPart::Source(Source::Url { + id: generate_source_id(), + url: url.to_string(), + title: Some(title), + provider_metadata: None, + })); } } - continue; - } - - // ── function_call item ── - if item_type == "function_call" { - let call_id = item.get("call_id").and_then(|v| v.as_str()).unwrap_or(""); - let name = item.get("name").and_then(|v| v.as_str()).unwrap_or(""); - - if event_type == "response.output_item.added" { - ongoing_tool_calls.insert(output_index, (call_id.to_string(), name.to_string())); - yield Ok(StreamPart::ToolInputStart { - id: call_id.to_string(), - tool_name: name.to_string(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - } else { - // output_item.done - has_function_call = true; - ongoing_tool_calls.remove(&output_index); - let arguments = item.get("arguments").and_then(|v| v.as_str()).unwrap_or(""); - yield Ok(StreamPart::ToolInputEnd { - id: call_id.to_string(), - provider_metadata: None, - }); - let input = arguments.to_string(); - yield Ok(StreamPart::ToolCall(RawToolCall { - tool_call_id: call_id.to_string(), - tool_name: name.to_string(), - input, - provider_executed: None, - dynamic: None, - provider_metadata: None, - })); - } - continue; } } - - // All other event types (web_search_call.in_progress, etc.) are ignored. + return; } - Err(error) => { - if error.is_recoverable_stream_error() { - if include_raw_chunks { - yield Ok(StreamPart::Raw { raw_value: Value::Null }); - } - yield Ok(StreamPart::Error { error }); + + // ── function_call item ── + if item_type == "function_call" { + let call_id = item.get("call_id").and_then(|v| v.as_str()).unwrap_or(""); + let name = item.get("name").and_then(|v| v.as_str()).unwrap_or(""); + + if event_type == "response.output_item.added" { + ongoing_tool_calls + .insert(output_index, (call_id.to_string(), name.to_string())); + controller.tool_input_start(call_id.to_string(), name.to_string()); } else { - yield Err(error); - return; + // output_item.done + *has_function_call = true; + ongoing_tool_calls.remove(&output_index); + let arguments = + item.get("arguments").and_then(|v| v.as_str()).unwrap_or(""); + controller.tool_input_end(call_id.to_string()); + let input = arguments.to_string(); + controller.enqueue(StreamPart::ToolCall(RawToolCall { + tool_call_id: call_id.to_string(), + tool_name: name.to_string(), + input, + provider_executed: None, + dynamic: None, + provider_metadata: None, + })); } } } - } - // Close any remaining open text blocks. - for block_id in &content_blocks { - yield Ok(StreamPart::TextEnd { id: block_id.clone(), provider_metadata: None}); + // All other event types (web_search_call.in_progress, etc.) are ignored. + } + Err(error) => { + if error.is_recoverable_stream_error() { + if include_raw_chunks { + controller.enqueue(StreamPart::Raw { + raw_value: Value::Null, + }); + } + controller.enqueue(StreamPart::Error { error }); + } else { + controller.error(error); + } } + } + } - // Final part: Finish. - let provider_meta = response_provider_metadata( - cost_in_usd_ticks, - service_tier.as_deref(), - prompt_cache_key.as_deref(), - safety_identifier.as_deref(), - ); - - yield Ok(StreamPart::Finish { - finish_reason: final_finish_reason, - usage: final_usage.unwrap_or_else(zero_usage), - provider_metadata: provider_meta, - }); - }; + /// Close the open text blocks and emit `Finish`: at the end of the + /// body, or right after an error event. + fn finish(&mut self, controller: &mut Controller) { + // Close any remaining open text blocks. + for block_id in std::mem::take(&mut self.content_blocks) { + controller.text_end(block_id); + } - Ok(StreamResult { - stream: Box::pin(stream), - request: Some(aimux_core::shared::RequestInfo { body: Some(body) }), - response: Some(aimux_core::shared::StreamResponseInfo { - headers: Some(response_headers), - }), - }) + // Final part: Finish. + let provider_meta = response_provider_metadata( + self.cost_in_usd_ticks, + self.service_tier.as_deref(), + self.prompt_cache_key.as_deref(), + self.safety_identifier.as_deref(), + ); + + controller.enqueue(StreamPart::Finish { + finish_reason: self.final_finish_reason.clone().unwrap_or(FINISH_OTHER), + usage: self.final_usage.take().unwrap_or_else(zero_usage), + provider_metadata: provider_meta, + }); } } diff --git a/docs/aisdk-architecture-alignment.md b/docs/aisdk-architecture-alignment.md index 2c8ed170e..c4388520b 100644 --- a/docs/aisdk-architecture-alignment.md +++ b/docs/aisdk-architecture-alignment.md @@ -115,7 +115,7 @@ FFI / Node / Python / 5 个 C-ABI 绑定 / CLI / Web 由 descriptor + manifest | B 轨 protocol registry = L1 协议层(RFC-0032,B1 "protocol 列 + `from_resolved`",B2–B8 "17 个 LanguageModel → ~7 个协议、33 个 wrapper 退役") | 保留方向,形式调整 | "协议"对应本文 §3.3 / §4.1 的模型实现家族(OpenAI chat / Responses、Anthropic Messages、Google、compat 等);registry preset 由内嵌的 registry 在运行时解析成 descriptor 表,按名字创建并复用 compat 实现(D12),不为每一行生成 Rust 源码,wrapper 退役由此完成。RFC-0032 的 registry 列保留为 descriptor 的数据来源;`from_resolved` 式构造由 `create_xxx` 工厂替代;C2 "FFI 构造器转发 shim"不再需要(§6 一次切换) | | #174 Auth L1(registry `auth` schema + `apply_auth()`)、#175 Auth L2(Credential 解析序 + CredentialStore + TokenRefresher) | 调整 | 鉴权按 D11 在各厂商包内以 headers 闭包 / fetch 装饰器 / Rust 内建 credential resolver 实现(§3.3、§4.1),registry `auth` 列作为 preset descriptor 数据保留,`apply_auth()` 即生成 headers 闭包。CredentialStore / TokenRefresher 涉及宿主回调与有状态刷新,与 Q4 一并后置;RFC-0018 的无状态刷新边界保留(§4.1 Codex 行) | | #167 transport-level replay("mock 挂 HTTP 层跑真协议代码";`ProviderRecord` = registry 行 + protocol)、#179 replay 子命令 | 保留,机制按本文 | 即 D4 / D5 的 `replay_fetch` / `replay_web_socket`(§4.3):录制与回放共用 provider-utils 的 leaf 边界,真实 provider 编码器与解析器参与回放。"`ProviderRecord` = registry 行 + protocol"撤销:录制 schema 3 不保存可重建的 ProviderRecord,live replay 由宿主 registry 与 operation 目标引用驱动(D24、§4.4)。#179 子命令保留在 tools | -| #185 "删 `StreamingToolCallTracker`(纯删)→ `ToolInput{Raw,Parsed}`" | 调整 | tracker 不删,改为对齐 `@ai-sdk/provider-utils` 并移入 aimux-provider-utils,由 OpenAI chat 流驱动(#204);对齐基线是 provider-utils 5.0.55,含该版新增的 `hasEmptyArguments` 修复(同一流里新 id、同名、首片段参数为空的并行调用不再被并入前一个调用;`streaming-tool-call-tracker.ts:140-141, 275-285`);`ToolInput` 的 Raw / Parsed 区分由第二部分的四层消息协议承担(V4 tool-call `input` 为原文 String,core 层 `parse_tool_call` 产出解析值) | +| #185 "删 `StreamingToolCallTracker`(纯删)→ `ToolInput{Raw,Parsed}`" | 调整 | 公共 tracker 删除(RFC-0042 T2);关联逻辑对齐 `@ai-sdk/provider-utils` 的 tracker,作为 OpenAI chat 一族流内部的工具调用表(`aimux-providers/src/stream/` 的 `ChatToolCalls`,crate 内部),不再是独立类型;对齐基线是 provider-utils 5.0.55,含该版新增的 `hasEmptyArguments` 修复(同一流里新 id、同名、首片段参数为空的并行调用不再被并入前一个调用;`streaming-tool-call-tracker.ts:140-141, 275-285`);`ToolInput` 的 Raw / Parsed 区分由第二部分的四层消息协议承担(V4 tool-call `input` 为原文 String,core 层 `parse_tool_call` 产出解析值) | | A2–A5 机械清理、E1 内部清理与文档重组、#170 / #171 registry 维护、#180 / #181 小项、S1 / S2 / S3 / P 门禁 | 保留,不受影响 | 与本文无依赖;A2(#203)独立于本文推进。体积与性能门禁按 [RFC-0043](../rfc/0043-size-performance-budgets.md)(#208);迁移各层的 PR 描述列出产物改前改后的体积(2026-10-15 修订) | | 0.9 能力矩阵 | 保留,不属于本文,迁移之后做 | 排期见 [ROADMAP](../ROADMAP.md) §6.1(2026-10-13 修订) | | "瘦身不改变对外 API"(ROADMAP §3 / §6.1) | 保留 | 该约束限定体积工作;本文的 API 替换是独立的设计决定,不以瘦身为由 | @@ -1118,6 +1118,7 @@ RFC-0014 的 `tracing` span 保留为统一事件的内建日志适配器。HTTP | 尚未对齐 | 流中途重试 | 上游 `streamText` 有 `streamRetries`（不设时不重试；设为 n 时流中途的 provider error 最多重新调用模型 n 次，`onError` 也可请求重试，见 ai 7.0.127 `A/generate-text/stream-text.ts:374`、`:1948-1953`、`:2495-2665`）。aimux 待实现，现状与不设时相同：流返回后不重试（RFC-0031 §8.3；§9.1 已知差距） | | 产品选择 | 首块超时的起点 | 上游在 provider 响应流开始时才开始计 chunk 超时（`A/generate-text/stream-text.ts:2228-2243`、`:2671`）；aimux 的 provider 在 `do_stream` 内预读首个 SSE 事件，所以 `first_chunk_ms` 在 operation 开始时计时，覆盖请求建立。保留这一差异：provider 等到首个 SSE 事件才从 `do_stream` 返回，若计时推迟到响应流开始，返回 200 之后一直不发数据的流就没有时间上限（RFC-0031 §8.1、§8.3；§9.1 已知差距） | | 产品选择 | core 工具循环 | 保持单步，不引入 agent loop | +| 结构差异 | 流式工具调用累积 | 上游 `@ai-sdk/provider-utils` 把 tracker 作为独立模块导出；aimux 没有 tracker 类型，OpenAI chat 一族（openai、openai-compatible、DeepSeek、Groq、Mistral）的流 Transformer 以 crate 内部的 `ChatToolCalls` 字段保存进行中的调用，各厂商的类型规则作为数据传入（RFC-0042 T2、§4）。关联行为等于 provider-utils 5.0.55 基线：上游 tracker 的测试用例改为驱动真实的 chat 流（`aimux-providers/src/stream/chat_tool_calls_tests.rs`） | | 不采用 | workflow 模型序列化 | 不用于录制或重建 | | 不采用 | 历史版本和 deprecated 兼容 | 不实现 V2/V3 适配、旧方法转发、旧录制读取；旧 namespace 键不在此列，照上游读取并双写（§4.1） | | 上游对齐 | Decision | `DecisionModelV4` 与 `decide`；aimux 扩展进 `providerOptions` / `providerMetadata`，无能力结构；不支持决策的 provider 返回 `None`；deprecated evaluation 别名不移植（§3.1） | @@ -2760,7 +2761,7 @@ TS、Go、Java、Kotlin、Swift、Dart、Python DTO 与 codec | `GenerateTextResult.raw` 及可独立修改的 metadata 镜像 | 从规范公共结果移除；原始 provider 信息放诊断/recording | | `ResponseMessageBuilder` 同时聚合 reasoning 与回放 | 替换为 content accumulator + `to_response_messages` | | 手写跨语言外部标签编解码和消息镜像 | 全部改为生成 | -| `aimux-stream::ToolCallStreamPart` | 删除；tracker 移入 aimux-provider-utils 后直接输出规范 stream part（#204），`aimux-stream` 只做 SSE 解码，不引入任何协议类型 | +| `aimux-stream::ToolCallStreamPart` | 删除；公共 tracker 一并删除，OpenAI chat 一族的流在内部累积工具调用并直接输出规范 stream part（RFC-0042 §4），`aimux-stream` 只做 SSE 解码，不引入任何协议类型 | --- @@ -2794,5 +2795,5 @@ TS、Go、Java、Kotlin、Swift、Dart、Python DTO 与 codec **两项结论（与 §0.5 一致）：** -1. **不新建 `aimux-protocol` crate。** 无运行时依赖的协议类型由第一部分的 `aimux-provider`（V4 类型）与 `aimux-provider-utils`（用户层 `ModelMessage`）承担，按 AI SDK 的包归属放置；core、provider 实现、tracker 和绑定生成都依赖这两个 crate，不复制公共类型。 +1. **不新建 `aimux-protocol` crate。** 无运行时依赖的协议类型由第一部分的 `aimux-provider`（V4 类型）与 `aimux-provider-utils`（用户层 `ModelMessage`）承担，按 AI SDK 的包归属放置；core、provider 实现和绑定生成都依赖这两个 crate，不复制公共类型。 2. **统一 JSON 适配规范采纳：** base64 字节、ISO 日期、缺失属性省略、JSON null 保留、结构化原生错误；其余字段、判别标签和层间语义直接按本文头部锁定的 AI SDK 基线实施。 diff --git a/rfc/0042-v4-type-boundaries.md b/rfc/0042-v4-type-boundaries.md index e8e747d0f..0f5f3d09c 100644 --- a/rfc/0042-v4-type-boundaries.md +++ b/rfc/0042-v4-type-boundaries.md @@ -27,7 +27,7 @@ PR #200 第一部分及 §0.5 优先于其第二部分的未消解建议： | 编号 | 提案 | 接受的后果 | |---|---|---| | T1 | 不在公共 wire 引入 `ToolInput::Raw/Parsed`；用不同层 DTO 的 `String` 与 JSON value 区分 | #185 的具体实现方案改写，保留消除混层的目标 | -| T2 | 彻底删除通用 tracker（`aimux-stream` 公共 tracker、第二套 `ToolCallStreamPart` 及任何公共累积器）；累积状态是 provider-utils 内 OpenAI-chat 一族共用流处理的内部实现，不对绑定公开 | 接受 #204 的归属方向，不接受保留公共 tracker 或“整套累积算法纯删除即完成”；见 §4 | +| T2 | 彻底删除通用 tracker（`aimux-stream` 公共 tracker、第二套 `ToolCallStreamPart` 及任何公共累积器）；累积状态是 provider 层（`aimux-providers/src/stream/` 的 crate 内部 `ChatToolCalls`）OpenAI-chat 一族共用流处理的内部实现，是各 chat 流 Transformer 的字段，不对绑定公开 | 接受 #204 的归属方向，不接受保留公共 tracker 或“整套累积算法纯删除即完成”；见 §4 | | T3 | 关联语义等于登记基线上游 tracker（`@ai-sdk/provider-utils`，基线至少 5.0.55）的 correlation 行为：归不了属的增量直接丢弃，不报错、不终止流；只有新调用缺 name 才报 `InvalidResponseData` | 与 #204 `861f26e` 的报错终止方案不同；差分用例取自上游该文件的测试；上游修复（如 5.0.55 的并行调用合并修复，#228/#250）随基线升级 | | T4 | 非法调用保留原文的结构化错误证据贯穿所有 wire/repair 路径 | 不通过 stringify(parsed input) 恢复原文 | | T5 | D8 以 PR #200 的统一 manifest 为入口，并为自定义 serde 提供显式 codec 描述 | 禁止各语言另维护字段表、独立 schema 或第二条 TS 生成链 | @@ -108,7 +108,7 @@ core 的一个入口执行工具查找、严格 JSON 解析、原有安全校验 ### 4.1 边界 -`aimux-stream` 只做通用字节/行/SSE/NDJSON 解码，不认识工具、工具 schema 或 V4 内容。OpenAI-chat 风格的工具增量累积状态是 provider-utils 内 OpenAI-chat 一族共用流处理的内部实现，不作为公共 tracker 对外导出；不能把所有协议强行转成 OpenAI index。Anthropic block id、Responses item id 等由各协议按原生终结信号转换成 V4 事件。 +`aimux-stream` 只做通用字节/行/SSE/NDJSON 解码，不认识工具、工具 schema 或 V4 内容。OpenAI-chat 风格的工具增量累积状态是 provider 层内 OpenAI-chat 一族共用流处理的内部实现（`aimux-providers/src/stream/` 的 crate 内部 `ChatToolCalls`，作为各 chat 流 Transformer 的字段，每个请求一份），不作为公共 tracker 对外导出；上游把 tracker 放在 `@ai-sdk/provider-utils` 的独立模块，aimux 放在流内部，关联行为与基线一致（登记于 `docs/aisdk-architecture-alignment.md` §5）；不能把所有协议强行转成 OpenAI index。Anthropic block id、Responses item id 等由各协议按原生终结信号转换成 V4 事件。 累积器输出唯一的 `LanguageModelV4StreamPart`，不保留 `ToolCallStreamPart` 镜像。内部状态是每个请求的临时状态，不是跨 ABI operation session，不对绑定公开控制句柄。 diff --git a/scripts/check_provider_boundaries.sh b/scripts/check_provider_boundaries.sh index ee06250fe..df1578cd2 100755 --- a/scripts/check_provider_boundaries.sh +++ b/scripts/check_provider_boundaries.sh @@ -152,13 +152,14 @@ violation 'groq / deepseek literal in the shared OpenAI or compatible package' " # `shared/` and `openai_compatible/` never reach into a vendor package: a # vendor's endpoint, base URL, headers or usage accounting is its own data # behind a trait or hook of the shared layer. Their code (doc links aside) -# names only the shared layers, the preset registry and `runtime_decision`. +# names only the shared layers (including `stream`), the preset registry and +# `runtime_decision`. rule4c_hits="$( grep -rnE 'crate::[a-z_]+' aimux-providers/src/shared \ aimux-providers/src/openai_compatible --include='*.rs' | awk '{ code = $0; sub(/^[^:]*:[0-9]+:/, "", code) } code ~ /^[[:space:]]*\/\// { next } - { gsub(/crate::(shared|openai|openai_compatible|preset|runtime_decision)([^a-z_]|$)/, "", code) } + { gsub(/crate::(shared|stream|openai|openai_compatible|preset|runtime_decision)([^a-z_]|$)/, "", code) } code ~ /crate::[a-z_]/' || true )" violation 'shared or openai_compatible names a vendor package (implement its trait in the vendor module)' "$rule4c_hits"