From bc7cd60edd17c88f77d1e1684f1ec684f84ebd99 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 29 Sep 2026 04:54:24 +0000 Subject: [PATCH 1/8] refactor(stream)!: remove unused StreamingToolCallTracker No crate in the workspace referenced the tracker; providers accumulate streamed tool-call deltas themselves. Drops the module, its 787-line test file and the re-exported companion types. Breaking for aimux-stream's public Rust API; noted in CHANGELOG under Unreleased. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015EYDYDWYcPsDmjBjFeuDVe (cherry picked from commit fe8758f0556eb7efcffcc4b2c186b8ac5581ad4f) --- aimux-stream/src/lib.rs | 5 - .../src/streaming_tool_call_tracker.rs | 429 ---------- .../tests/streaming_tool_call_tracker_test.rs | 787 ------------------ 3 files changed, 1221 deletions(-) delete mode 100644 aimux-stream/src/streaming_tool_call_tracker.rs delete mode 100644 aimux-stream/tests/streaming_tool_call_tracker_test.rs diff --git a/aimux-stream/src/lib.rs b/aimux-stream/src/lib.rs index e023a3204..65c368daf 100644 --- a/aimux-stream/src/lib.rs +++ b/aimux-stream/src/lib.rs @@ -6,12 +6,7 @@ pub mod lines; pub mod sse; -pub mod streaming_tool_call_tracker; // Re-export the most commonly used items. pub use lines::extract_lines; pub use sse::{SseError, SseEvent, SseStream}; -pub use streaming_tool_call_tracker::{ - StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, - ToolCallStreamPart, TrackerError, TypeValidation, -}; diff --git a/aimux-stream/src/streaming_tool_call_tracker.rs b/aimux-stream/src/streaming_tool_call_tracker.rs deleted file mode 100644 index be9ad8907..000000000 --- a/aimux-stream/src/streaming_tool_call_tracker.rs +++ /dev/null @@ -1,429 +0,0 @@ -//! Streaming tool call tracker. -//! -//! Rust translation of `@ai-sdk/provider-utils`'s `StreamingToolCallTracker` -//! (`packages/provider-utils/src/streaming-tool-call-tracker.ts`). -//! -//! Tracks streaming tool call state across multiple deltas from an -//! OpenAI-compatible chat completion stream: accumulates `arguments` string -//! fragments by `index`, emits `tool-input-start` / `tool-input-delta` / -//! `tool-input-end` / `tool-call` events, and finalizes any unfinished tool -//! calls on [`StreamingToolCallTracker::flush`]. -//! -//! Like the TS original, a tool call is *never* finalized before `flush` — a -//! parsable argument buffer can still be the prefix of a longer argument -//! string, so acting on it early would use truncated inputs (ai-sdk #13137). - -use serde_json::Value; -use thiserror::Error; - -/// Default upper bound accepted for a tool call `index`. -const DEFAULT_MAX_INDEX: usize = 1024; - -/// The `function` sub-object of a streaming tool call delta. -#[derive(Debug, Clone, Default)] -pub struct StreamingToolCallFunction { - pub name: Option, - pub arguments: Option, -} - -/// A streaming tool call delta — the `tool_calls[i]` entry of an OpenAI-style -/// streaming chunk. -/// -/// Use the builder methods ([`StreamingToolCallDelta::index`], etc.) to -/// construct one ergonomically. `arguments: null` (TS) maps to `None`; an -/// empty-string `arguments: ''` maps to `Some("")`. -#[derive(Debug, Clone, Default)] -pub struct StreamingToolCallDelta { - pub index: Option, - pub id: Option, - /// The `type` field. Named `r#type` because `type` is a reserved word. - pub r#type: Option, - pub function: Option, - /// Provider-specific payload carried alongside the standard fields. Used by - /// `extract_metadata` to pull out provider metadata (e.g. a Google thought - /// signature). Defaults to [`Value::Null`]. - pub extra: Value, -} - -impl StreamingToolCallDelta { - #[must_use] - pub fn new() -> Self { - Self::default() - } - - #[must_use] - pub fn index(mut self, index: usize) -> Self { - self.index = Some(index); - self - } - - #[must_use] - pub fn id(mut self, id: impl Into) -> Self { - self.id = Some(id.into()); - self - } - - /// Set the `type` field (named `tool_type` because `type` is reserved). - #[must_use] - pub fn tool_type(mut self, t: impl Into) -> Self { - self.r#type = Some(t.into()); - self - } - - #[must_use] - pub fn function_name(mut self, name: impl Into) -> Self { - self.function.get_or_insert_with(Default::default).name = Some(name.into()); - self - } - - /// Set the `function.arguments` fragment. Pass `""` for an explicit empty - /// fragment; omit the call entirely for `None` (TS `arguments: null`). - #[must_use] - pub fn arguments(mut self, args: impl Into) -> Self { - self.function.get_or_insert_with(Default::default).arguments = Some(args.into()); - self - } - - #[must_use] - pub fn extra(mut self, extra: Value) -> Self { - self.extra = extra; - self - } -} - -/// The stream parts emitted by [`StreamingToolCallTracker`]. -/// -/// Mirrors the subset of `LanguageModelV4StreamPart` the TS tracker enqueues. -/// Note that [`ToolCallStreamPart::ToolCall::input`] is the raw accumulated -/// argument *string* — the tracker does not parse it (matching the TS -/// behavior). -#[derive(Debug, Clone, PartialEq)] -pub enum ToolCallStreamPart { - /// Start of a tool call's input streaming. - ToolInputStart { id: String, tool_name: String }, - /// A delta of tool call input (a partial argument fragment). - ToolInputDelta { id: String, delta: String }, - /// End of a tool call's input streaming. - ToolInputEnd { id: String }, - /// A complete, finalized tool call. - ToolCall { - tool_call_id: String, - tool_name: String, - input: String, - provider_metadata: Option, - }, -} - -/// How to validate the `type` field on a new tool call delta. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] -pub enum TypeValidation { - /// No validation (default). - #[default] - None, - /// Throw if `type` is present and not `"function"`. - IfPresent, - /// Throw if `type` is not exactly `"function"`. - Required, -} - -/// Errors raised while processing a tool call delta. -#[derive(Debug, Error, PartialEq)] -pub enum TrackerError { - #[error("Expected 'id' to be a string.")] - MissingId, - #[error("Expected 'function.name' to be a string.")] - MissingFunctionName, - #[error("Expected 'function' type.")] - InvalidType, - #[error("Tool call index out of range")] - IndexOutOfRange, -} - -struct TrackedToolCall { - id: String, - function_name: String, - arguments: String, - has_finished: bool, - metadata: Option, -} - -/// Extract provider metadata from a delta (the TS `extractMetadata` option). -type ExtractMetadataFn = Box Option>; -/// Build the `providerMetadata` for a finalized tool call (the TS -/// `buildToolCallProviderMetadata` option). -type BuildMetadataFn = Box) -> Option>; - -/// Tracks streaming tool call state across multiple deltas from an -/// OpenAI-compatible chat completion stream. -/// -/// Emitted [`ToolCallStreamPart`]s accumulate in an internal buffer; inspect -/// them with [`parts`](Self::parts) and reset between checks with -/// [`clear_parts`](Self::clear_parts) (the TS test's `parts.length = 0`). -/// -/// The type parameter `M` is the provider-metadata type. Use `()` (the -/// default) when no metadata handling is needed, or `serde_json::Value` (or -/// any `SharedV4ProviderMetadata`-like type) together with -/// [`with_extract_metadata`](Self::with_extract_metadata) / -/// [`with_build_provider_metadata`](Self::with_build_provider_metadata). -pub struct StreamingToolCallTracker { - tool_calls: Vec>>, - parts: Vec>, - // The TS `generateId` option: fallback for `toolCall.id ?? generateId()` - // when an incoming tool-call delta omits its id. - generate_id: Box String>, - type_validation: TypeValidation, - /// Upper bound accepted for a tool call `index`. Guards against a remote - /// index resizing `tool_calls` to a huge vector. Defaults to - /// [`DEFAULT_MAX_INDEX`] (1024). - max_index: usize, - extract_metadata: Option>, - build_provider_metadata: Option>, -} - -impl Default for StreamingToolCallTracker<()> { - fn default() -> Self { - Self::new() - } -} - -impl StreamingToolCallTracker { - /// Create a new tracker with no metadata handling and default settings. - #[must_use] - pub fn new() -> Self { - Self { - tool_calls: Vec::new(), - parts: Vec::new(), - generate_id: Box::new(|| "generated".to_string()), - type_validation: TypeValidation::None, - max_index: DEFAULT_MAX_INDEX, - extract_metadata: None, - build_provider_metadata: None, - } - } - - /// Set a custom id generator (the TS `generateId` option). - #[must_use] - pub fn with_generate_id String + 'static>(mut self, f: F) -> Self { - self.generate_id = Box::new(f); - self - } - - /// Set the `type` validation mode (the TS `typeValidation` option). - #[must_use] - pub fn with_type_validation(mut self, v: TypeValidation) -> Self { - self.type_validation = v; - self - } - - /// Set the maximum accepted tool call `index` (defaults to 1024). A delta - /// whose resolved `index` exceeds `max_index` returns - /// [`TrackerError::IndexOutOfRange`] instead of resizing `tool_calls` to - /// `index + 1` slots. - #[must_use] - pub fn with_max_index(mut self, max_index: usize) -> Self { - self.max_index = max_index; - self - } - - /// Set the metadata extractor (the TS `extractMetadata` option). Called - /// once when a new tool call is detected; the returned metadata is stored - /// on the tool call and passed to the builder at finalization. - #[must_use] - pub fn with_extract_metadata Option + 'static>( - mut self, - f: F, - ) -> Self { - self.extract_metadata = Some(Box::new(f)); - self - } - - /// Set the provider-metadata builder (the TS `buildToolCallProviderMetadata` - /// option). Receives the metadata previously extracted; if `None` is - /// returned, no `provider_metadata` is included in the `tool-call` event. - #[must_use] - pub fn with_build_provider_metadata) -> Option + 'static>( - mut self, - f: F, - ) -> Self { - self.build_provider_metadata = Some(Box::new(f)); - self - } - - /// Process a tool call delta from a streaming response chunk. Emits events - /// into the internal buffer. - /// - /// # Errors - /// - /// Returns `TrackerError::IndexOutOfRange` when the delta's tool index - /// exceeds the configured maximum. - pub fn process_delta(&mut self, delta: &StreamingToolCallDelta) -> Result<(), TrackerError> { - let index = delta.index.unwrap_or(self.tool_calls.len()); - // Guard against a remote `index` resizing `tool_calls` to a huge vector. - if index > self.max_index { - return Err(TrackerError::IndexOutOfRange); - } - let is_new = self - .tool_calls - .get(index) - .map(std::option::Option::is_none) - .unwrap_or(true); - if is_new { - self.process_new_tool_call(index, delta)?; - } else { - self.process_existing_tool_call(index, delta); - } - Ok(()) - } - - /// Finalize any unfinished tool calls. Should be called during the - /// stream's flush to ensure all tool calls are properly completed. Emits - /// `tool-input-end` and `tool-call` for each unfinished tool call. - pub fn flush(&mut self) { - for i in 0..self.tool_calls.len() { - let needs_finish = self - .tool_calls - .get(i) - .and_then(|o| o.as_ref()) - .is_some_and(|tc| !tc.has_finished); - if needs_finish { - self.finish_tool_call_at(i); - } - } - } - - /// The events emitted so far (accumulated across `process_delta`/`flush`). - #[must_use] - pub fn parts(&self) -> &[ToolCallStreamPart] { - &self.parts - } - - /// Clear the accumulated events (the TS test's `parts.length = 0`). - pub fn clear_parts(&mut self) { - self.parts.clear(); - } - - fn process_new_tool_call( - &mut self, - index: usize, - delta: &StreamingToolCallDelta, - ) -> Result<(), TrackerError> { - match self.type_validation { - TypeValidation::Required => { - if delta.r#type.as_deref() != Some("function") { - return Err(TrackerError::InvalidType); - } - } - TypeValidation::IfPresent => { - if delta.r#type.as_deref().is_some_and(|t| t != "function") { - return Err(TrackerError::InvalidType); - } - } - TypeValidation::None => {} - } - - let id = delta.id.clone().unwrap_or_else(|| (self.generate_id)()); - let function_name = delta - .function - .as_ref() - .and_then(|f| f.name.clone()) - .ok_or(TrackerError::MissingFunctionName)?; - - self.parts.push(ToolCallStreamPart::ToolInputStart { - id: id.clone(), - tool_name: function_name.clone(), - }); - - let metadata = self - .extract_metadata - .as_ref() - .and_then(|extract| extract(delta)); - - // TS: `toolCallDelta.function.arguments ?? ''`. - let arguments = delta - .function - .as_ref() - .and_then(|f| f.arguments.clone()) - .unwrap_or_default(); - - if index >= self.tool_calls.len() { - self.tool_calls.resize_with(index + 1, || None); - } - self.tool_calls[index] = Some(TrackedToolCall { - id: id.clone(), - function_name: function_name.clone(), - arguments: arguments.clone(), - has_finished: false, - metadata, - }); - - // Emit initial delta if arguments already present. - if !arguments.is_empty() { - self.parts.push(ToolCallStreamPart::ToolInputDelta { - id: id.clone(), - delta: arguments, - }); - } - - // Tool calls must not finalize before the stream ends (see #13137). - Ok(()) - } - - fn process_existing_tool_call(&mut self, index: usize, delta: &StreamingToolCallDelta) { - // TS: `toolCallDelta.function?.arguments != null`. - let new_args = match delta.function.as_ref().and_then(|f| f.arguments.as_ref()) { - Some(args) => args.clone(), - None => return, - }; - let id = { - let Some(Some(tool_call)) = self.tool_calls.get_mut(index) else { - return; - }; - if tool_call.has_finished { - return; - } - tool_call.arguments.push_str(&new_args); - tool_call.id.clone() - }; - self.parts.push(ToolCallStreamPart::ToolInputDelta { - id, - delta: new_args, - }); - } - - fn finish_tool_call_at(&mut self, index: usize) { - // Scope the mutable borrow of the tracked tool call, extracting the - // owned data we need to emit the final events. - let (id, function_name, arguments, metadata) = { - let Some(Some(tool_call)) = self.tool_calls.get_mut(index) else { - return; - }; - if tool_call.has_finished { - return; - } - ( - tool_call.id.clone(), - tool_call.function_name.clone(), - tool_call.arguments.clone(), - tool_call.metadata.take(), - ) - }; - if let Some(Some(tool_call)) = self.tool_calls.get_mut(index) { - tool_call.has_finished = true; - } - - self.parts - .push(ToolCallStreamPart::ToolInputEnd { id: id.clone() }); - - let provider_metadata = self - .build_provider_metadata - .as_ref() - .and_then(|build| build(metadata.as_ref())); - - self.parts.push(ToolCallStreamPart::ToolCall { - tool_call_id: id, - tool_name: function_name, - input: arguments, - provider_metadata, - }); - } -} diff --git a/aimux-stream/tests/streaming_tool_call_tracker_test.rs b/aimux-stream/tests/streaming_tool_call_tracker_test.rs deleted file mode 100644 index 25a9b7655..000000000 --- a/aimux-stream/tests/streaming_tool_call_tracker_test.rs +++ /dev/null @@ -1,787 +0,0 @@ -//! Rust translation of -//! `packages/provider-utils/src/streaming-tool-call-tracker.test.ts` (18 cases). -//! -//! The TS test collects emitted parts on a shared `controller.enqueue` sink and -//! resets it with `parts.length = 0` between assertions. The Rust equivalent -//! accumulates parts inside the tracker: inspect with `parts()` and reset with -//! `clear_parts()`. - -use aimux_stream::{ - StreamingToolCallDelta, StreamingToolCallTracker, ToolCallStreamPart, TrackerError, - TypeValidation, -}; -use serde_json::{Value, json}; - -/// Convenience builder alias. -fn d() -> StreamingToolCallDelta { - StreamingToolCallDelta::default() -} - -/// A tracker with no metadata handling (`M = ()`), matching the TS tests that -/// don't exercise the metadata hooks. Pinning `M = ()` here lets type -/// inference resolve `ToolCallStreamPart::ToolCall { provider_metadata: None }` -/// in the assertions. -fn new_tracker() -> StreamingToolCallTracker<()> { - StreamingToolCallTracker::<()>::new() -} - -fn tool_call_id(part: &ToolCallStreamPart) -> &str { - match part { - ToolCallStreamPart::ToolCall { tool_call_id, .. } => tool_call_id, - _ => panic!("expected tool-call part"), - } -} - -// == processDelta ========================================================== - -#[test] -fn single_tool_call_accumulated_across_multiple_deltas() { - // TS: "should handle a single tool call accumulated across multiple deltas" - let mut tracker = new_tracker(); - - // First delta: new tool call with id and name. - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("get_weather") - .arguments("{\"ci"), - ) - .unwrap(); - - assert_eq!( - tracker.parts(), - &[ - ToolCallStreamPart::ToolInputStart { - id: "call_1".into(), - tool_name: "get_weather".into(), - }, - ToolCallStreamPart::ToolInputDelta { - id: "call_1".into(), - delta: "{\"ci".into(), - }, - ] - ); - - tracker.clear_parts(); - - // Second delta: more arguments. - tracker - .process_delta(&d().index(0).arguments("ty\": \"San")) - .unwrap(); - - assert_eq!( - tracker.parts(), - &[ToolCallStreamPart::ToolInputDelta { - id: "call_1".into(), - delta: "ty\": \"San".into(), - }] - ); - - tracker.clear_parts(); - - // Third delta: completes the JSON -- must not finalize before flush, since a - // parsable buffer can still be the prefix of longer arguments. - tracker - .process_delta(&d().index(0).arguments(" Francisco\"}")) - .unwrap(); - - assert_eq!( - tracker.parts(), - &[ToolCallStreamPart::ToolInputDelta { - id: "call_1".into(), - delta: " Francisco\"}".into(), - }] - ); - - tracker.clear_parts(); - - tracker.flush(); - - assert_eq!( - tracker.parts(), - &[ - ToolCallStreamPart::ToolInputEnd { - id: "call_1".into(), - }, - ToolCallStreamPart::ToolCall { - tool_call_id: "call_1".into(), - tool_name: "get_weather".into(), - input: "{\"city\": \"San Francisco\"}".into(), - provider_metadata: None, - }, - ] - ); -} - -#[test] -fn full_tool_call_in_single_chunk() { - // TS: "should handle a full tool call in a single chunk" - let mut tracker = new_tracker(); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("get_weather") - .arguments("{\"city\": \"London\"}"), - ) - .unwrap(); - - assert_eq!( - tracker.parts(), - &[ - ToolCallStreamPart::ToolInputStart { - id: "call_1".into(), - tool_name: "get_weather".into(), - }, - ToolCallStreamPart::ToolInputDelta { - id: "call_1".into(), - delta: "{\"city\": \"London\"}".into(), - }, - ] - ); - - tracker.clear_parts(); - tracker.flush(); - - assert_eq!( - tracker.parts(), - &[ - ToolCallStreamPart::ToolInputEnd { - id: "call_1".into(), - }, - ToolCallStreamPart::ToolCall { - tool_call_id: "call_1".into(), - tool_name: "get_weather".into(), - input: "{\"city\": \"London\"}".into(), - provider_metadata: None, - }, - ] - ); -} - -#[test] -fn not_finalize_when_argument_prefix_is_parsable_json() { - // TS: "should not finalize a tool call when its argument prefix is parsable JSON" - let mut tracker = new_tracker(); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("search") - .arguments("{\"query\": \"test\"}"), - ) - .unwrap(); - - // The parsable prefix must not emit tool-input-end / tool-call. - let types: Vec<&str> = tracker - .parts() - .iter() - .map(|p| match p { - ToolCallStreamPart::ToolInputStart { .. } => "tool-input-start", - ToolCallStreamPart::ToolInputDelta { .. } => "tool-input-delta", - ToolCallStreamPart::ToolInputEnd { .. } => "tool-input-end", - ToolCallStreamPart::ToolCall { .. } => "tool-call", - }) - .collect(); - assert_eq!(types, vec!["tool-input-start", "tool-input-delta"]); - - tracker - .process_delta(&d().index(0).arguments(", \"limit\": 10}")) - .unwrap(); - - tracker.flush(); - - assert_eq!( - tracker.parts().last().unwrap(), - &ToolCallStreamPart::ToolCall { - tool_call_id: "call_1".into(), - tool_name: "search".into(), - input: "{\"query\": \"test\"}, \"limit\": 10}".into(), - provider_metadata: None, - } - ); - - let tool_call_count = tracker - .parts() - .iter() - .filter(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })) - .count(); - assert_eq!(tool_call_count, 1); -} - -#[test] -fn multiple_concurrent_tool_calls() { - // TS: "should handle multiple concurrent tool calls" - let mut tracker = new_tracker(); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("get_weather") - .arguments(""), - ) - .unwrap(); - tracker - .process_delta( - &d().index(1) - .id("call_2") - .tool_type("function") - .function_name("get_time") - .arguments(""), - ) - .unwrap(); - - assert_eq!( - tracker.parts(), - &[ - ToolCallStreamPart::ToolInputStart { - id: "call_1".into(), - tool_name: "get_weather".into(), - }, - ToolCallStreamPart::ToolInputStart { - id: "call_2".into(), - tool_name: "get_time".into(), - }, - ] - ); -} - -#[test] -fn skip_deltas_for_already_finished_tool_calls() { - // TS: "should skip deltas for already-finished tool calls" - let mut tracker = new_tracker(); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments("{}"), - ) - .unwrap(); - - // Finalize via flush. - tracker.flush(); - tracker.clear_parts(); - - // Late delta for the same tool call. - tracker - .process_delta(&d().index(0).arguments("extra")) - .unwrap(); - - assert!(tracker.parts().is_empty()); -} - -#[test] -fn skip_delta_emission_when_arguments_null() { - // TS: "should skip delta emission when arguments are null". - // `arguments: null` maps to `None` (omitted on the builder); behaviorally - // identical to a `function` with no arguments field. - let mut tracker = new_tracker(); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments(""), - ) - .unwrap(); - - tracker.clear_parts(); - - // Delta with null arguments (no `.arguments(...)` call -> None). - tracker.process_delta(&d().index(0)).unwrap(); - - assert!(tracker.parts().is_empty()); -} - -#[test] -fn use_index_fallback_when_index_not_provided() { - // TS: "should use index fallback when index is not provided" - let mut tracker = new_tracker(); - - tracker - .process_delta( - &d().id("call_1") - .tool_type("function") - .function_name("fn1") - .arguments("{}"), - ) - .unwrap(); - tracker - .process_delta( - &d().id("call_2") - .tool_type("function") - .function_name("fn2") - .arguments("{}"), - ) - .unwrap(); - - let starts: Vec<_> = tracker - .parts() - .iter() - .filter(|p| matches!(p, ToolCallStreamPart::ToolInputStart { .. })) - .cloned() - .collect(); - assert_eq!( - starts, - vec![ - ToolCallStreamPart::ToolInputStart { - id: "call_1".into(), - tool_name: "fn1".into(), - }, - ToolCallStreamPart::ToolInputStart { - id: "call_2".into(), - tool_name: "fn2".into(), - }, - ] - ); -} - -#[test] -fn generate_id_when_id_missing() { - // TS: id is missing → `toolCall.id ?? generateId()` generates a fallback id - // instead of throwing. The default generator returns "generated". - let mut tracker = new_tracker(); - tracker - .process_delta(&d().index(0).tool_type("function").function_name("fn")) - .unwrap(); - - assert_eq!( - tracker.parts(), - &[ToolCallStreamPart::ToolInputStart { - id: "generated".into(), - tool_name: "fn".into(), - }] - ); -} - -#[test] -fn throw_when_function_name_missing() { - // TS: "should throw when function.name is missing". - // `function: {}` (TS) maps to `function: None` here; both yield a missing - // name and the same error. - let mut tracker = new_tracker(); - let err = tracker - .process_delta(&d().index(0).id("call_1").tool_type("function")) - .unwrap_err(); - assert_eq!(err, TrackerError::MissingFunctionName); - assert_eq!(err.to_string(), "Expected 'function.name' to be a string."); -} - -// == typeValidation ======================================================== - -#[test] -fn no_validate_type_with_type_validation_none() { - // TS: "should not validate type with typeValidation: none" - let mut tracker = new_tracker().with_type_validation(TypeValidation::None); - // Should not throw even with a non-function type. - let result = tracker.process_delta( - &d().index(0) - .id("call_1") - .tool_type("custom") - .function_name("fn") - .arguments(""), - ); - assert!(result.is_ok()); -} - -#[test] -fn validate_type_when_present_with_type_validation_if_present() { - // TS: "should validate type when present with typeValidation: if-present" - let mut tracker = new_tracker().with_type_validation(TypeValidation::IfPresent); - - // Should throw for a non-function type. - let err = tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("custom") - .function_name("fn") - .arguments(""), - ) - .unwrap_err(); - assert_eq!(err, TrackerError::InvalidType); - assert_eq!(err.to_string(), "Expected 'function' type."); - - // Should not throw when type is null (absent). The first delta errored - // before creating any tool call, so index 0 is still new. - let result = - tracker.process_delta(&d().index(0).id("call_1").function_name("fn").arguments("")); - assert!(result.is_ok()); -} - -#[test] -fn require_function_type_with_type_validation_required() { - // TS: "should require function type with typeValidation: required" - let mut tracker = new_tracker().with_type_validation(TypeValidation::Required); - - // Should throw when type is null/undefined. - let err = tracker - .process_delta(&d().index(0).id("call_1").function_name("fn").arguments("")) - .unwrap_err(); - assert_eq!(err, TrackerError::InvalidType); - - // Should not throw for 'function' type. - let result = tracker.process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments(""), - ); - assert!(result.is_ok()); -} - -// == flush ================================================================= - -#[test] -fn finalize_unfinished_tool_calls_on_flush() { - // TS: "should finalize unfinished tool calls on flush" - let mut tracker = new_tracker(); - - // Start a tool call but don't complete it. - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments("{\"key\": \"val"), - ) - .unwrap(); - - tracker.clear_parts(); - - // Flush should finalize. - tracker.flush(); - - assert_eq!( - tracker.parts(), - &[ - ToolCallStreamPart::ToolInputEnd { - id: "call_1".into(), - }, - ToolCallStreamPart::ToolCall { - tool_call_id: "call_1".into(), - tool_name: "fn".into(), - input: "{\"key\": \"val".into(), - provider_metadata: None, - }, - ] - ); -} - -#[test] -fn not_refinalize_already_finished_tool_calls() { - // TS: "should not re-finalize already finished tool calls" - let mut tracker = new_tracker(); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments("{}"), - ) - .unwrap(); - - // First flush finalizes the tool call. - tracker.flush(); - tracker.clear_parts(); - - tracker.flush(); - - // No events should be emitted since the tool call was already finished. - assert!(tracker.parts().is_empty()); -} - -// == metadata ============================================================= - -#[test] -fn extract_and_include_provider_metadata_in_tool_call_events() { - // TS: "should extract and include provider metadata in tool-call events" - let mut tracker = StreamingToolCallTracker::::new() - .with_extract_metadata(|delta| { - delta - .extra - .get("google")? - .get("thought_signature")? - .as_str() - .map(|s| json!({ "thoughtSignature": s })) - }) - .with_build_provider_metadata(|metadata| { - metadata - .as_ref() - .and_then(|m| m.get("thoughtSignature")) - .and_then(|s| s.as_str()) - .map(|s| json!({ "google": { "thoughtSignature": s } })) - }); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments("{}") - .extra(json!({ "google": { "thought_signature": "sig123" } })), - ) - .unwrap(); - - tracker.flush(); - - let tool_call = tracker - .parts() - .iter() - .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })) - .unwrap(); - assert_eq!( - tool_call, - &ToolCallStreamPart::ToolCall { - tool_call_id: "call_1".into(), - tool_name: "fn".into(), - input: "{}".into(), - provider_metadata: Some(json!({ "google": { "thoughtSignature": "sig123" } })), - } - ); -} - -#[test] -fn include_provider_metadata_for_unfinished_tool_calls_finalized_in_flush() { - // TS: "should include provider metadata for unfinished tool calls finalized in flush" - let mut tracker = StreamingToolCallTracker::::new() - .with_extract_metadata(|_| Some(json!({ "custom": { "key": "value" } }))) - .with_build_provider_metadata(|metadata| { - metadata - .as_ref() - .map(|m| json!({ "provider": (*m).clone() })) - }); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments("{\"incomplete"), - ) - .unwrap(); - - tracker.clear_parts(); - - tracker.flush(); - - let tool_call = tracker - .parts() - .iter() - .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })) - .unwrap(); - assert_eq!( - tool_call, - &ToolCallStreamPart::ToolCall { - tool_call_id: "call_1".into(), - tool_name: "fn".into(), - input: "{\"incomplete".into(), - provider_metadata: Some(json!({ "provider": { "custom": { "key": "value" } } })), - } - ); -} - -#[test] -fn not_include_provider_metadata_when_build_returns_none() { - // TS: "should not include providerMetadata when buildToolCallProviderMetadata returns undefined" - let mut tracker = StreamingToolCallTracker::::new() - .with_extract_metadata(|_| None) - .with_build_provider_metadata(|_| None); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments("{}"), - ) - .unwrap(); - - tracker.flush(); - - let tool_call = tracker - .parts() - .iter() - .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })) - .unwrap(); - assert_eq!( - tool_call, - &ToolCallStreamPart::ToolCall { - tool_call_id: "call_1".into(), - tool_name: "fn".into(), - input: "{}".into(), - provider_metadata: None, - } - ); - // Mirrors the TS `not.toHaveProperty('providerMetadata')` assertion: the - // metadata is explicitly `None`, never `Some`. - assert!(matches!( - tool_call, - ToolCallStreamPart::ToolCall { - provider_metadata: None, - .. - } - )); -} - -// == generateId =========================================================== - -#[test] -fn use_custom_generate_id_for_tool_call_ids_when_id_missing_in_fallback() { - // TS: "should use custom generateId for tool call IDs when id is missing in fallback". - // The id is present (`call_1`), so the custom generator is NOT used; the - // original id is kept (mirrors the TS `toolCall.id ?? generateId()` path). - use std::sync::atomic::{AtomicUsize, Ordering}; - static CALLS: AtomicUsize = AtomicUsize::new(0); - - let mut tracker = StreamingToolCallTracker::<()>::new().with_generate_id(|| { - CALLS.fetch_add(1, Ordering::SeqCst); - "custom-id".to_string() - }); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments("{\"key\": \"val"), - ) - .unwrap(); - - tracker.clear_parts(); - tracker.flush(); - - let tool_call = tracker - .parts() - .iter() - .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })) - .unwrap(); - assert_eq!(tool_call_id(tool_call), "call_1"); - // The custom generator must not have been invoked. - assert_eq!(CALLS.load(Ordering::SeqCst), 0); -} - -#[test] -fn use_custom_generate_id_when_id_is_missing() { - // P2-05: when a tool-call delta omits its id, the configured `generateId` - // fallback is invoked (TS `toolCall.id ?? generateId()`). The generated id - // is used for the emitted ToolInputStart and the finalized ToolCall. - use std::sync::atomic::{AtomicUsize, Ordering}; - static CALLS: AtomicUsize = AtomicUsize::new(0); - - let mut tracker = StreamingToolCallTracker::<()>::new().with_generate_id(|| { - CALLS.fetch_add(1, Ordering::SeqCst); - "custom-id".to_string() - }); - - // No `.id(...)` — the generator must fill it in. - tracker - .process_delta( - &d().index(0) - .tool_type("function") - .function_name("fn") - .arguments("{\"key\": \"val"), - ) - .unwrap(); - - // The generator was invoked exactly once for the missing id. - assert_eq!(CALLS.load(Ordering::SeqCst), 1); - - // The emitted start part carries the generated id. - assert_eq!( - tracker.parts(), - &[ - ToolCallStreamPart::ToolInputStart { - id: "custom-id".into(), - tool_name: "fn".into(), - }, - ToolCallStreamPart::ToolInputDelta { - id: "custom-id".into(), - delta: "{\"key\": \"val".into(), - }, - ] - ); - - tracker.clear_parts(); - tracker.flush(); - - let tool_call = tracker - .parts() - .iter() - .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })) - .unwrap(); - assert_eq!(tool_call_id(tool_call), "custom-id"); - // No additional generator invocations during flush. - assert_eq!(CALLS.load(Ordering::SeqCst), 1); -} - -// == index bound (P1-10) ================================================== - -#[test] -fn reject_tool_call_index_above_max_index() { - // A remote `index` far above the cap must not resize `tool_calls` to a - // huge vector; it returns `IndexOutOfRange`. - let mut tracker = new_tracker().with_max_index(4); - - // index == max_index is accepted (boundary). - tracker - .process_delta( - &d().index(4) - .id("call_4") - .tool_type("function") - .function_name("fn") - .arguments(""), - ) - .unwrap(); - - // index == max_index + 1 is rejected. - let err = tracker - .process_delta( - &d().index(5) - .id("call_5") - .tool_type("function") - .function_name("fn") - .arguments(""), - ) - .unwrap_err(); - assert_eq!(err, TrackerError::IndexOutOfRange); - assert_eq!(err.to_string(), "Tool call index out of range"); -} - -#[test] -fn default_max_index_rejects_huge_index() { - // The default cap (1024) rejects an absurd index that would otherwise - // resize `tool_calls` to index+1 slots. - let mut tracker = new_tracker(); - let err = tracker - .process_delta( - &d().index(1_000_000) - .id("call_x") - .tool_type("function") - .function_name("fn") - .arguments(""), - ) - .unwrap_err(); - assert_eq!(err, TrackerError::IndexOutOfRange); -} From c4398af6f9af3528d26ba3ae8c4cfbd2737dc2bd Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 29 Sep 2026 05:02:49 +0000 Subject: [PATCH 2/8] feat(stream)!: align StreamingToolCallTracker with the current AI SDK Re-adds the tracker as a port of the current @ai-sdk/provider-utils streaming-tool-call-tracker.ts instead of the older index-only version removed in the previous commit. - Correlate deltas by wire id, index and function name; drop ambiguous ones - Unique tool-call ids with bounded suffixes; blank generator output falls back to "tool-call" - Ignore unmatched blank function names; retain blank-name continuations - Flush in index order only when every call has an index - New StreamingToolCallArgumentState / starts_with_structured_value port - process_delta and flush return the emitted parts instead of buffering Tests are ported case for case from the upstream suites (45 tracker cases, 9 argument-state cases). No caller in the workspace uses the tracker yet. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015EYDYDWYcPsDmjBjFeuDVe (cherry picked from commit 6035201dbec549f5d77fa39390b070e3f1552f3f) --- aimux-stream/src/lib.rs | 9 + .../src/streaming_tool_call_argument_state.rs | 116 ++ .../src/streaming_tool_call_tracker.rs | 638 ++++++++++ ...streaming_tool_call_argument_state_test.rs | 73 ++ .../tests/streaming_tool_call_tracker_test.rs | 1098 +++++++++++++++++ 5 files changed, 1934 insertions(+) create mode 100644 aimux-stream/src/streaming_tool_call_argument_state.rs create mode 100644 aimux-stream/src/streaming_tool_call_tracker.rs create mode 100644 aimux-stream/tests/streaming_tool_call_argument_state_test.rs create mode 100644 aimux-stream/tests/streaming_tool_call_tracker_test.rs diff --git a/aimux-stream/src/lib.rs b/aimux-stream/src/lib.rs index 65c368daf..d9078d8d8 100644 --- a/aimux-stream/src/lib.rs +++ b/aimux-stream/src/lib.rs @@ -6,7 +6,16 @@ pub mod lines; pub mod sse; +pub mod streaming_tool_call_argument_state; +pub mod streaming_tool_call_tracker; // Re-export the most commonly used items. pub use lines::extract_lines; pub use sse::{SseError, SseEvent, SseStream}; +pub use streaming_tool_call_argument_state::{ + StreamingToolCallArgumentState, starts_with_structured_value, +}; +pub use streaming_tool_call_tracker::{ + StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, + ToolCallStreamPart, TrackerError, TypeValidation, +}; diff --git a/aimux-stream/src/streaming_tool_call_argument_state.rs b/aimux-stream/src/streaming_tool_call_argument_state.rs new file mode 100644 index 000000000..db156e66d --- /dev/null +++ b/aimux-stream/src/streaming_tool_call_argument_state.rs @@ -0,0 +1,116 @@ +//! Incremental structure tracking for streamed tool-call arguments. +//! +//! Rust translation of `@ai-sdk/provider-utils`'s +//! `StreamingToolCallArgumentState` +//! (`packages/provider-utils/src/streaming-tool-call-argument-state.ts`). +//! +//! The state is intentionally *structural* rather than a JSON parse: a +//! currently parsable scalar can still be the prefix of a later value, but a +//! closed top-level `{…}` / `[…]` cannot be extended. + +#[derive(Debug, Clone, PartialEq, Eq)] +enum ArgumentStructure { + Undetermined, + Other, + Structured { + stack: Vec, + in_string: bool, + escaped: bool, + complete: bool, + }, +} + +/// Whether `value` (ignoring leading whitespace) starts with `{` or `[`. +#[must_use] +pub fn starts_with_structured_value(value: Option<&str>) -> bool { + value + .and_then(|v| v.trim_start().chars().next()) + .is_some_and(|c| c == '{' || c == '[') +} + +/// Incrementally tracks whether streamed tool-call arguments contain a +/// complete structured JSON value. +#[derive(Debug, Clone)] +pub struct StreamingToolCallArgumentState { + structure: ArgumentStructure, +} + +impl StreamingToolCallArgumentState { + /// Create a state seeded with `initial_value`. + #[must_use] + pub fn new(initial_value: &str) -> Self { + let mut state = Self { + structure: ArgumentStructure::Undetermined, + }; + state.append(initial_value); + state + } + + /// `true` once a top-level `{…}` / `[…]` value has been closed. + #[must_use] + pub fn has_complete_structured_value(&self) -> bool { + matches!( + self.structure, + ArgumentStructure::Structured { complete: true, .. } + ) + } + + /// Feed the next argument fragment. + pub fn append(&mut self, delta: &str) { + for character in delta.chars() { + match &mut self.structure { + ArgumentStructure::Undetermined => { + if character.is_whitespace() { + continue; + } + self.structure = if character == '{' || character == '[' { + ArgumentStructure::Structured { + stack: vec![character], + in_string: false, + escaped: false, + complete: false, + } + } else { + ArgumentStructure::Other + }; + } + ArgumentStructure::Other | ArgumentStructure::Structured { complete: true, .. } => { + } + ArgumentStructure::Structured { + stack, + in_string, + escaped, + complete, + } => { + if *in_string { + if *escaped { + *escaped = false; + } else if character == '\\' { + *escaped = true; + } else if character == '"' { + *in_string = false; + } + continue; + } + + match character { + '"' => *in_string = true, + '{' | '[' => stack.push(character), + '}' | ']' => { + let expected = if character == '}' { '{' } else { '[' }; + if stack.last() != Some(&expected) { + self.structure = ArgumentStructure::Other; + continue; + } + stack.pop(); + if stack.is_empty() { + *complete = true; + } + } + _ => {} + } + } + } + } + } +} diff --git a/aimux-stream/src/streaming_tool_call_tracker.rs b/aimux-stream/src/streaming_tool_call_tracker.rs new file mode 100644 index 000000000..9a8d7240e --- /dev/null +++ b/aimux-stream/src/streaming_tool_call_tracker.rs @@ -0,0 +1,638 @@ +//! Streaming tool call tracker. +//! +//! Rust translation of `@ai-sdk/provider-utils`'s `StreamingToolCallTracker` +//! (`packages/provider-utils/src/streaming-tool-call-tracker.ts`). +//! +//! Tracks streaming tool call state across the deltas of an OpenAI-compatible +//! chat completion stream: accumulates `arguments` fragments, emits +//! `tool-input-start` / `tool-input-delta` / `tool-input-end` / `tool-call` +//! parts, and finalizes unfinished calls on [`StreamingToolCallTracker::flush`]. +//! +//! Deltas are correlated to calls by wire `id`, `index` and function name +//! (see [`StreamingToolCallTracker`]'s resolution table), not by `index` +//! alone, so providers that reuse indices, repeat ids, drop ids on +//! continuations or send blank names are handled. +//! +//! Like the TS original, a call is *never* finalized before `flush`: a +//! parsable argument buffer can still be the prefix of a longer argument +//! string, so acting on it early would use truncated inputs (ai-sdk #13137). + +use std::collections::{HashMap, HashSet}; + +use serde_json::Value; +use thiserror::Error; + +use crate::streaming_tool_call_argument_state::{ + StreamingToolCallArgumentState, starts_with_structured_value, +}; + +/// Fallback id used when the id generator returns a blank string. +const FALLBACK_TOOL_CALL_ID: &str = "tool-call"; + +/// The `function` sub-object of a streaming tool call delta. +#[derive(Debug, Clone, Default)] +pub struct StreamingToolCallFunction { + pub name: Option, + pub arguments: Option, +} + +/// A streaming tool call delta — the `tool_calls[i]` entry of an OpenAI-style +/// streaming chunk. +/// +/// `arguments: null` (TS) maps to `None`; `arguments: ''` maps to `Some("")`. +#[derive(Debug, Clone, Default)] +pub struct StreamingToolCallDelta { + pub index: Option, + pub id: Option, + /// The `type` field. Named `r#type` because `type` is a reserved word. + pub r#type: Option, + pub function: Option, + /// Provider-specific payload carried alongside the standard fields, read + /// by the `extract_metadata` hook (e.g. a Google thought signature). + pub extra: Value, +} + +impl StreamingToolCallDelta { + #[must_use] + pub fn new() -> Self { + Self::default() + } + + #[must_use] + pub fn index(mut self, index: usize) -> Self { + self.index = Some(index); + self + } + + #[must_use] + pub fn id(mut self, id: impl Into) -> Self { + self.id = Some(id.into()); + self + } + + /// Set the `type` field (named `tool_type` because `type` is reserved). + #[must_use] + pub fn tool_type(mut self, t: impl Into) -> Self { + self.r#type = Some(t.into()); + self + } + + #[must_use] + pub fn function_name(mut self, name: impl Into) -> Self { + self.function.get_or_insert_with(Default::default).name = Some(name.into()); + self + } + + /// Set the `function.arguments` fragment. Pass `""` for an explicit empty + /// fragment; omit the call for `None` (TS `arguments: null`). + #[must_use] + pub fn arguments(mut self, args: impl Into) -> Self { + self.function.get_or_insert_with(Default::default).arguments = Some(args.into()); + self + } + + #[must_use] + pub fn extra(mut self, extra: Value) -> Self { + self.extra = extra; + self + } +} + +/// The stream parts emitted by [`StreamingToolCallTracker`]. +/// +/// Mirrors the subset of `LanguageModelV4StreamPart` the TS tracker enqueues. +/// [`ToolCallStreamPart::ToolCall::input`] is the raw accumulated argument +/// *string*; the tracker does not parse it. +#[derive(Debug, Clone, PartialEq)] +pub enum ToolCallStreamPart { + /// Start of a tool call's input streaming. + ToolInputStart { id: String, tool_name: String }, + /// A partial argument fragment. + ToolInputDelta { id: String, delta: String }, + /// End of a tool call's input streaming. + ToolInputEnd { id: String }, + /// A complete, finalized tool call. + ToolCall { + tool_call_id: String, + tool_name: String, + input: String, + provider_metadata: Option, + }, +} + +/// How to validate the `type` field on a new tool call delta. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum TypeValidation { + /// No validation (default). + #[default] + None, + /// Error if `type` is present and not `"function"`. + IfPresent, + /// Error if `type` is not exactly `"function"`. + Required, +} + +/// Errors raised while processing a tool call delta. +#[derive(Debug, Error, PartialEq, Eq)] +pub enum TrackerError { + #[error("Expected 'function.name' to be a string.")] + MissingFunctionName, + #[error("Expected 'function' type.")] + InvalidType, + #[error("Failed to create a unique tool call ID.")] + IdExhausted, +} + +struct TrackedToolCall { + id: String, + index: Option, + sequence: usize, + function_name: String, + arguments: String, + argument_state: StreamingToolCallArgumentState, + has_finished: bool, + metadata: Option, +} + +enum ToolCallResolution { + Existing(usize), + New, + Ambiguous, +} + +type GenerateIdFn = Box String + Send + Sync>; +type ExtractMetadataFn = Box Option + Send + Sync>; +type BuildMetadataFn = Box) -> Option + Send + Sync>; + +/// Tracks streaming tool call state across multiple deltas. +/// +/// [`process_delta`](Self::process_delta) and [`flush`](Self::flush) return +/// the parts to forward downstream (the TS tracker enqueues them on a +/// controller instead). +/// +/// The type parameter `M` is the provider-metadata type; use `()` (the +/// default) when no metadata handling is needed. +/// +/// # Correlation +/// +/// | ID evidence | index/name evidence | start evidence | resolution | +/// | --- | --- | --- | --- | +/// | known | matching | any | matching call, new call, or ambiguity | +/// | known | conflicting | named | new call | +/// | unseen | matching | structured start | new call | +/// | unseen | matching | continuation | matching call or ambiguity | +/// | absent | matching | any | matching call, new call, or ambiguity | +/// | absent | absent | named | new call | +/// | absent | absent | unnamed | sole unfinished call, new call, or ambiguity | +/// +/// An *ambiguous* delta is dropped. +pub struct StreamingToolCallTracker { + tool_calls: Vec>, + tool_calls_by_id: HashMap>, + tool_calls_by_index: HashMap>, + used_tool_call_ids: HashSet, + next_generated_id_suffixes: HashMap, + generate_id: GenerateIdFn, + type_validation: TypeValidation, + extract_metadata: Option>, + build_provider_metadata: Option>, +} + +impl Default for StreamingToolCallTracker { + fn default() -> Self { + Self::new() + } +} + +impl StreamingToolCallTracker { + /// Create a tracker with no metadata handling and default settings. + /// + /// The default id generator returns a blank-free constant, so ids that + /// need generating become `tool-call`, `tool-call-1`, `tool-call-2`, … + #[must_use] + pub fn new() -> Self { + Self { + tool_calls: Vec::new(), + tool_calls_by_id: HashMap::new(), + tool_calls_by_index: HashMap::new(), + used_tool_call_ids: HashSet::new(), + next_generated_id_suffixes: HashMap::new(), + generate_id: Box::new(|| FALLBACK_TOOL_CALL_ID.to_string()), + type_validation: TypeValidation::None, + extract_metadata: None, + build_provider_metadata: None, + } + } + + /// Set the id generator (the TS `generateId` option). Blank or repeated + /// outputs are turned into usable unique ids. + #[must_use] + pub fn with_generate_id String + Send + Sync + 'static>(mut self, f: F) -> Self { + self.generate_id = Box::new(f); + self + } + + /// Set the `type` validation mode (the TS `typeValidation` option). + #[must_use] + pub fn with_type_validation(mut self, v: TypeValidation) -> Self { + self.type_validation = v; + self + } + + /// Set the metadata extractor (the TS `extractMetadata` option). Called + /// once when a new tool call is detected. + #[must_use] + pub fn with_extract_metadata< + F: Fn(&StreamingToolCallDelta) -> Option + Send + Sync + 'static, + >( + mut self, + f: F, + ) -> Self { + self.extract_metadata = Some(Box::new(f)); + self + } + + /// Set the provider-metadata builder (the TS + /// `buildToolCallProviderMetadata` option). If it returns `None`, the + /// `tool-call` part carries no `provider_metadata`. + #[must_use] + pub fn with_build_provider_metadata) -> Option + Send + Sync + 'static>( + mut self, + f: F, + ) -> Self { + self.build_provider_metadata = Some(Box::new(f)); + self + } + + /// Process a tool call delta from a streaming chunk and return the parts + /// it produces (possibly none). + /// + /// # Errors + /// + /// Returns [`TrackerError::InvalidType`] when `type` validation fails, + /// [`TrackerError::MissingFunctionName`] when a new call has no + /// (non-null) function name, and [`TrackerError::IdExhausted`] if no + /// unique id can be produced. + pub fn process_delta( + &mut self, + delta: &StreamingToolCallDelta, + ) -> Result>, TrackerError> { + let wire_name = delta.function.as_ref().and_then(|f| f.name.as_deref()); + let has_blank_name = wire_name.is_some_and(|n| n.trim().is_empty()); + let wire_id = non_blank(delta.id.as_deref()); + let name = non_blank(wire_name); + let index = delta.index; + let arguments = delta.function.as_ref().and_then(|f| f.arguments.as_deref()); + + let resolution = self.resolve_tool_call( + wire_id, + index, + name, + name.is_some() && starts_with_structured_value(arguments), + ); + + let mut parts = Vec::new(); + let call = match resolution { + ToolCallResolution::Ambiguous => return Ok(parts), + ToolCallResolution::New => { + // Blank names cannot start a usable call, but some providers + // repeat a blank name on continuations. Those were correlated + // above; only an unmatched blank-name delta is ignored. + if has_blank_name { + return Ok(parts); + } + self.process_new_tool_call(delta, wire_id, index, name, &mut parts)? + } + ToolCallResolution::Existing(call) => { + if let Some(wire_id) = wire_id { + self.associate_wire_id(call, wire_id); + } + self.process_existing_tool_call(call, arguments, &mut parts); + call + } + }; + + if let Some(index) = index { + self.tool_calls_by_index + .entry(index) + .or_default() + .insert(call); + } + Ok(parts) + } + + /// Finalize any unfinished tool calls and return the closing parts. Call + /// once when the stream ends. + pub fn flush(&mut self) -> Vec> { + // Index order is only reliable when every call has an index; for + // mixed streams keep insertion order. + let mut order: Vec = (0..self.tool_calls.len()).collect(); + if self.tool_calls.iter().all(|c| c.index.is_some()) { + order.sort_by_key(|&i| (self.tool_calls[i].index, self.tool_calls[i].sequence)); + } + + let mut parts = Vec::new(); + for call in order { + if !self.tool_calls[call].has_finished { + self.finish_tool_call(call, &mut parts); + } + } + parts + } + + fn resolve_tool_call( + &self, + wire_id: Option<&str>, + index: Option, + name: Option<&str>, + has_explicit_call_start: bool, + ) -> ToolCallResolution { + let indexed = index.and_then(|i| self.tool_calls_by_index.get(&i)); + let matching_indexed = self.filter_by_name(indexed, name); + + if let Some(wire_id) = wire_id { + if let Some(with_id) = self.tool_calls_by_id.get(wire_id) { + if index.is_some() { + let matching: Vec = matching_indexed + .iter() + .copied() + .filter(|c| with_id.contains(c)) + .collect(); + let resolved = self.resolve_matching(matching, has_explicit_call_start); + if !matches!(resolved, ToolCallResolution::New) { + return resolved; + } + + // A named delta with a distinct index starts a new call + // even when its wire id and name repeat: providers may + // reuse ids across parallel calls. + if name.is_some() { + return ToolCallResolution::New; + } + + // Conflicting labels on a continuation cannot be + // resolved safely. + if indexed.is_some() { + return ToolCallResolution::Ambiguous; + } + + return self.resolve_matching(sorted(with_id), false); + } + + if let Some(name) = name { + let matching: Vec = sorted(with_id) + .into_iter() + .filter(|&c| self.tool_calls[c].function_name == name) + .collect(); + return self.resolve_matching(matching, has_explicit_call_start); + } + + return self.resolve_matching(sorted(with_id), false); + } + + if !matching_indexed.is_empty() { + // A previously unseen id plus a named structured argument + // start is stronger evidence of a distinct call than a reused + // index/name; ids may still change on plain continuations. + return if has_explicit_call_start { + ToolCallResolution::New + } else { + self.resolve_matching(matching_indexed, false) + }; + } + + return ToolCallResolution::New; + } + + if indexed.is_some() { + // Repeated names are valid on continuations; a different name at + // the same index means a new call from a provider that reuses + // indices across parallel calls. + return self.resolve_matching(matching_indexed, has_explicit_call_start); + } + + if name.is_some() { + return ToolCallResolution::New; + } + + let unfinished: Vec = (0..self.tool_calls.len()) + .filter(|&c| !self.tool_calls[c].has_finished) + .collect(); + match unfinished.len() { + 0 => ToolCallResolution::New, + 1 => ToolCallResolution::Existing(unfinished[0]), + _ => ToolCallResolution::Ambiguous, + } + } + + fn filter_by_name(&self, calls: Option<&HashSet>, name: Option<&str>) -> Vec { + calls.map_or_else(Vec::new, |calls| { + sorted(calls) + .into_iter() + .filter(|&c| name.is_none_or(|n| self.tool_calls[c].function_name == n)) + .collect() + }) + } + + fn resolve_matching( + &self, + calls: Vec, + has_explicit_call_start: bool, + ) -> ToolCallResolution { + if calls.is_empty() { + return ToolCallResolution::New; + } + + if !has_explicit_call_start { + return if calls.len() == 1 { + ToolCallResolution::Existing(calls[0]) + } else { + ToolCallResolution::Ambiguous + }; + } + + // A repeated name can occur on continuations. A fresh structured + // argument prefix signals another call only once the matching call + // has completed its own structured payload. + let continuable: Vec = calls + .into_iter() + .filter(|&c| { + !self.tool_calls[c] + .argument_state + .has_complete_structured_value() + }) + .collect(); + match continuable.len() { + 0 => ToolCallResolution::New, + 1 => ToolCallResolution::Existing(continuable[0]), + _ => ToolCallResolution::Ambiguous, + } + } + + fn process_new_tool_call( + &mut self, + delta: &StreamingToolCallDelta, + wire_id: Option<&str>, + index: Option, + name: Option<&str>, + parts: &mut Vec>, + ) -> Result { + match self.type_validation { + TypeValidation::Required => { + if delta.r#type.as_deref() != Some("function") { + return Err(TrackerError::InvalidType); + } + } + TypeValidation::IfPresent => { + if delta.r#type.as_deref().is_some_and(|t| t != "function") { + return Err(TrackerError::InvalidType); + } + } + TypeValidation::None => {} + } + + let name = name.ok_or(TrackerError::MissingFunctionName)?; + let id = self.create_tool_call_id(wire_id)?; + + parts.push(ToolCallStreamPart::ToolInputStart { + id: id.clone(), + tool_name: name.to_string(), + }); + + let metadata = self + .extract_metadata + .as_ref() + .and_then(|extract| extract(delta)); + + let initial_arguments = delta + .function + .as_ref() + .and_then(|f| f.arguments.clone()) + .unwrap_or_default(); + + let call = self.tool_calls.len(); + self.tool_calls.push(TrackedToolCall { + id: id.clone(), + index, + sequence: call, + function_name: name.to_string(), + argument_state: StreamingToolCallArgumentState::new(&initial_arguments), + arguments: initial_arguments.clone(), + has_finished: false, + metadata, + }); + if let Some(wire_id) = wire_id { + self.associate_wire_id(call, wire_id); + } + + if !initial_arguments.is_empty() { + parts.push(ToolCallStreamPart::ToolInputDelta { + id, + delta: initial_arguments, + }); + } + + // Tool calls must not finalize before the stream ends (ai-sdk + // #13137); finalization happens in `flush`. + Ok(call) + } + + fn process_existing_tool_call( + &mut self, + call: usize, + arguments: Option<&str>, + parts: &mut Vec>, + ) { + let tool_call = &mut self.tool_calls[call]; + if tool_call.has_finished { + return; + } + if let Some(arguments) = arguments { + tool_call.argument_state.append(arguments); + tool_call.arguments.push_str(arguments); + parts.push(ToolCallStreamPart::ToolInputDelta { + id: tool_call.id.clone(), + delta: arguments.to_string(), + }); + } + } + + fn associate_wire_id(&mut self, call: usize, wire_id: &str) { + self.tool_calls_by_id + .entry(wire_id.to_string()) + .or_default() + .insert(call); + } + + fn create_tool_call_id(&mut self, wire_id: Option<&str>) -> Result { + if let Some(wire_id) = wire_id + && !self.used_tool_call_ids.contains(wire_id) + { + self.used_tool_call_ids.insert(wire_id.to_string()); + return Ok(wire_id.to_string()); + } + + let generated = non_blank(Some(&(self.generate_id)())) + .unwrap_or(FALLBACK_TOOL_CALL_ID) + .to_string(); + + if !self.used_tool_call_ids.contains(&generated) { + self.used_tool_call_ids.insert(generated.clone()); + return Ok(generated); + } + + // Resume after the last suffix checked for this generated value so + // deterministic generators stay bounded without rescanning occupied + // suffixes. + let initial_suffix = self + .next_generated_id_suffixes + .get(&generated) + .copied() + .unwrap_or(1); + let maximum_suffix = initial_suffix + self.used_tool_call_ids.len(); + for suffix in initial_suffix..=maximum_suffix { + let suffixed = format!("{generated}-{suffix}"); + if !self.used_tool_call_ids.contains(&suffixed) { + self.used_tool_call_ids.insert(suffixed.clone()); + self.next_generated_id_suffixes + .insert(generated, suffix + 1); + return Ok(suffixed); + } + } + + // Unreachable by the pigeonhole principle; guards the invariant + // instead of looping without bound. + Err(TrackerError::IdExhausted) + } + + fn finish_tool_call(&mut self, call: usize, parts: &mut Vec>) { + let tool_call = &mut self.tool_calls[call]; + tool_call.has_finished = true; + + parts.push(ToolCallStreamPart::ToolInputEnd { + id: tool_call.id.clone(), + }); + + let provider_metadata = self + .build_provider_metadata + .as_ref() + .and_then(|build| build(tool_call.metadata.as_ref())); + + parts.push(ToolCallStreamPart::ToolCall { + tool_call_id: tool_call.id.clone(), + tool_name: tool_call.function_name.clone(), + input: tool_call.arguments.clone(), + provider_metadata, + }); + } +} + +fn non_blank(value: Option<&str>) -> Option<&str> { + value.filter(|v| !v.trim().is_empty()) +} + +fn sorted(calls: &HashSet) -> Vec { + let mut calls: Vec = calls.iter().copied().collect(); + calls.sort_unstable(); + calls +} diff --git a/aimux-stream/tests/streaming_tool_call_argument_state_test.rs b/aimux-stream/tests/streaming_tool_call_argument_state_test.rs new file mode 100644 index 000000000..44cf754e2 --- /dev/null +++ b/aimux-stream/tests/streaming_tool_call_argument_state_test.rs @@ -0,0 +1,73 @@ +//! Port of `streaming-tool-call-argument-state.test.ts` from +//! `@ai-sdk/provider-utils`. + +use aimux_stream::{StreamingToolCallArgumentState, starts_with_structured_value}; + +#[test] +fn starts_with_structured_value_true() { + for value in ["{}", " {", "[]", "\n["] { + assert!(starts_with_structured_value(Some(value)), "{value:?}"); + } +} + +#[test] +fn starts_with_structured_value_false() { + assert!(!starts_with_structured_value(None)); + for value in ["", " ", "1", "\"value\""] { + assert!(!starts_with_structured_value(Some(value)), "{value:?}"); + } +} + +#[test] +fn tracks_a_structured_value_across_deltas() { + let mut state = StreamingToolCallArgumentState::new(" {\"value\":"); + assert!(!state.has_complete_structured_value()); + + state.append("1}"); + assert!(state.has_complete_structured_value()); +} + +#[test] +fn tracks_nested_objects_and_arrays() { + let state = StreamingToolCallArgumentState::new("[{\"value\":{\"items\":[1,2]}}]"); + assert!(state.has_complete_structured_value()); +} + +#[test] +fn ignores_structural_characters_inside_strings() { + let state = StreamingToolCallArgumentState::new("{\"value\":\"braces: } ] { [\"}"); + assert!(state.has_complete_structured_value()); +} + +#[test] +fn handles_escaped_quotes_across_deltas() { + let mut state = StreamingToolCallArgumentState::new("{\"value\":\"escaped quote: \\\""); + assert!(!state.has_complete_structured_value()); + + state.append(" still in string\"}"); + assert!(state.has_complete_structured_value()); +} + +#[test] +fn can_begin_after_an_empty_or_whitespace_only_delta() { + let mut state = StreamingToolCallArgumentState::new(" "); + + state.append("["); + assert!(!state.has_complete_structured_value()); + + state.append("]"); + assert!(state.has_complete_structured_value()); +} + +#[test] +fn does_not_treat_scalar_arguments_as_a_complete_structured_value() { + let state = StreamingToolCallArgumentState::new("12"); + assert!(!state.has_complete_structured_value()); +} + +#[test] +fn does_not_recover_mismatched_structures_as_complete() { + let mut state = StreamingToolCallArgumentState::new("{\"value\":]"); + state.append("}"); + assert!(!state.has_complete_structured_value()); +} diff --git a/aimux-stream/tests/streaming_tool_call_tracker_test.rs b/aimux-stream/tests/streaming_tool_call_tracker_test.rs new file mode 100644 index 000000000..7c55044c4 --- /dev/null +++ b/aimux-stream/tests/streaming_tool_call_tracker_test.rs @@ -0,0 +1,1098 @@ +//! Port of `streaming-tool-call-tracker.test.ts` from +//! `@ai-sdk/provider-utils`, case for case. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use aimux_stream::{ + StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, + ToolCallStreamPart, TrackerError, TypeValidation, +}; +use serde_json::{Value, json}; + +type Part = ToolCallStreamPart; + +/// Tracker plus the parts it has emitted so far (the TS `createCollector`). +struct Harness { + tracker: StreamingToolCallTracker, + parts: Vec>, +} + +impl Harness<()> { + fn new() -> Self { + Self::with(StreamingToolCallTracker::new()) + } + + fn with(tracker: StreamingToolCallTracker<()>) -> Self { + Self::with_meta(tracker) + } +} + +impl Harness { + fn with_meta(tracker: StreamingToolCallTracker) -> Self { + Self { + tracker, + parts: Vec::new(), + } + } + + fn delta(&mut self, delta: StreamingToolCallDelta) -> Result<(), TrackerError> { + let parts = self.tracker.process_delta(&delta)?; + self.parts.extend(parts); + Ok(()) + } + + fn flush(&mut self) { + let parts = self.tracker.flush(); + self.parts.extend(parts); + } + + fn clear(&mut self) { + self.parts.clear(); + } +} + +impl Harness { + /// `(id, name, input)` of every emitted `tool-call`, in order. + fn tool_calls(&self) -> Vec<(String, String, String)> { + self.parts + .iter() + .filter_map(|part| match part { + Part::ToolCall { + tool_call_id, + tool_name, + input, + .. + } => Some((tool_call_id.clone(), tool_name.clone(), input.clone())), + _ => None, + }) + .collect() + } +} + +/// A delta with an explicit `function` object (as every TS test passes one). +fn delta( + index: Option, + id: Option<&str>, + ty: Option<&str>, + name: Option<&str>, + arguments: Option<&str>, +) -> StreamingToolCallDelta { + StreamingToolCallDelta { + index, + id: id.map(str::to_string), + r#type: ty.map(str::to_string), + function: Some(StreamingToolCallFunction { + name: name.map(str::to_string), + arguments: arguments.map(str::to_string), + }), + extra: Value::Null, + } +} + +/// Full call start: index, id, `type: 'function'`, name and arguments. +fn start(index: usize, id: &str, name: &str, arguments: &str) -> StreamingToolCallDelta { + delta( + Some(index), + Some(id), + Some("function"), + Some(name), + Some(arguments), + ) +} + +/// Continuation carrying only arguments (plus an optional index). +fn cont(index: Option, arguments: &str) -> StreamingToolCallDelta { + delta(index, None, None, None, Some(arguments)) +} + +fn tc(id: &str, name: &str, input: &str) -> (String, String, String) { + (id.into(), name.into(), input.into()) +} + +fn input_start(id: &str, tool_name: &str) -> Part { + Part::ToolInputStart { + id: id.into(), + tool_name: tool_name.into(), + } +} + +fn input_delta(id: &str, delta: &str) -> Part { + Part::ToolInputDelta { + id: id.into(), + delta: delta.into(), + } +} + +fn input_end(id: &str) -> Part { + Part::ToolInputEnd { id: id.into() } +} + +fn tool_call(id: &str, name: &str, input: &str) -> Part { + Part::ToolCall { + tool_call_id: id.into(), + tool_name: name.into(), + input: input.into(), + provider_metadata: None, + } +} + +/// Deterministic id generator: `prefix-1`, `prefix-2`, … and a call counter. +fn counting_generator(ids: &'static [&'static str]) -> (impl Fn() -> String, Arc) { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = Arc::clone(&calls); + ( + move || { + let n = counter.fetch_add(1, Ordering::SeqCst); + ids[n.min(ids.len() - 1)].to_string() + }, + calls, + ) +} + +mod process_delta { + use super::*; + + #[test] + fn single_tool_call_accumulated_across_multiple_deltas() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "get_weather", "{\"ci")).unwrap(); + assert_eq!( + h.parts, + vec![ + input_start("call_1", "get_weather"), + input_delta("call_1", "{\"ci"), + ] + ); + h.clear(); + + h.delta(cont(Some(0), "ty\": \"San")).unwrap(); + assert_eq!(h.parts, vec![input_delta("call_1", "ty\": \"San")]); + h.clear(); + + // Completing the JSON must not finalize before flush: a parsable + // buffer can still be the prefix of longer arguments. + h.delta(cont(Some(0), " Francisco\"}")).unwrap(); + assert_eq!(h.parts, vec![input_delta("call_1", " Francisco\"}")]); + h.clear(); + + h.flush(); + assert_eq!( + h.parts, + vec![ + input_end("call_1"), + tool_call("call_1", "get_weather", "{\"city\": \"San Francisco\"}"), + ] + ); + } + + #[test] + fn full_tool_call_in_a_single_chunk() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "get_weather", "{\"city\": \"London\"}")) + .unwrap(); + assert_eq!( + h.parts, + vec![ + input_start("call_1", "get_weather"), + input_delta("call_1", "{\"city\": \"London\"}"), + ] + ); + h.clear(); + + h.flush(); + assert_eq!( + h.parts, + vec![ + input_end("call_1"), + tool_call("call_1", "get_weather", "{\"city\": \"London\"}"), + ] + ); + } + + #[test] + fn does_not_finalize_when_argument_prefix_is_parsable_json() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "search", "{\"query\": \"test\"}")) + .unwrap(); + assert_eq!( + h.parts, + vec![ + input_start("call_1", "search"), + input_delta("call_1", "{\"query\": \"test\"}"), + ] + ); + + h.delta(cont(Some(0), ", \"limit\": 10}")).unwrap(); + h.flush(); + + assert_eq!( + h.parts.last(), + Some(&tool_call( + "call_1", + "search", + "{\"query\": \"test\"}, \"limit\": 10}" + )) + ); + assert_eq!(h.tool_calls().len(), 1); + } + + #[test] + fn multiple_concurrent_tool_calls() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "get_weather", "")).unwrap(); + h.delta(start(1, "call_2", "get_time", "")).unwrap(); + + assert_eq!( + h.parts, + vec![ + input_start("call_1", "get_weather"), + input_start("call_2", "get_time"), + ] + ); + } + + #[test] + fn non_zero_and_non_contiguous_indexes() { + let mut h = Harness::new(); + + h.delta(start(1, "call_1", "fn1", "{\"value\":1}")).unwrap(); + h.delta(start(3, "call_2", "fn2", "{\"value\":2}")).unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("call_1", "fn1", "{\"value\":1}"), + tc("call_2", "fn2", "{\"value\":2}"), + ] + ); + } + + #[test] + fn keeps_distinct_tool_calls_that_reuse_an_index() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "fn", "{\"value\":1}")).unwrap(); + h.delta(start(0, "call_2", "fn", "{\"value\":2}")).unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("call_1", "fn", "{\"value\":1}"), + tc("call_2", "fn", "{\"value\":2}"), + ] + ); + } + + #[test] + fn continues_latest_call_when_index_is_omitted_after_starting_at() { + for index in [None, Some(7)] { + let mut h = Harness::new(); + + h.delta(delta( + index, + Some("call_1"), + Some("function"), + Some("fn"), + Some("{\"val"), + )) + .unwrap(); + h.delta(cont(None, "ue\":1}")).unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![tc("call_1", "fn", "{\"value\":1}")], + "index {index:?}" + ); + } + } + + #[test] + fn uses_the_index_when_continuation_ids_are_empty() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "fn", "{\"val")).unwrap(); + h.delta(delta( + Some(0), + Some(""), + Some("function"), + None, + Some("ue\":1}"), + )) + .unwrap(); + h.flush(); + + assert_eq!(h.tool_calls(), vec![tc("call_1", "fn", "{\"value\":1}")]); + } + + #[test] + fn skips_deltas_for_already_finished_tool_calls() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "fn", "{}")).unwrap(); + h.flush(); + h.clear(); + + h.delta(cont(Some(0), "extra")).unwrap(); + assert_eq!(h.parts, vec![]); + } + + #[test] + fn skips_delta_emission_when_arguments_are_null() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "fn", "")).unwrap(); + h.clear(); + + h.delta(delta(Some(0), None, None, None, None)).unwrap(); + assert_eq!(h.parts, vec![]); + } + + #[test] + fn uses_index_fallback_when_index_is_not_provided() { + let mut h = Harness::new(); + + h.delta(delta( + None, + Some("call_1"), + Some("function"), + Some("fn1"), + Some("{}"), + )) + .unwrap(); + h.delta(delta( + None, + Some("call_2"), + Some("function"), + Some("fn2"), + Some("{}"), + )) + .unwrap(); + + let starts: Vec<&Part> = h + .parts + .iter() + .filter(|p| matches!(p, Part::ToolInputStart { .. })) + .collect(); + assert_eq!( + starts, + vec![&input_start("call_1", "fn1"), &input_start("call_2", "fn2")] + ); + } + + #[test] + fn generates_an_id_when_id_is_missing() { + let mut h = Harness::with( + StreamingToolCallTracker::new().with_generate_id(|| "generated-id".to_string()), + ); + + h.delta(delta( + Some(0), + None, + Some("function"), + Some("fn"), + Some("{}"), + )) + .unwrap(); + h.flush(); + + assert_eq!(h.tool_calls(), vec![tc("generated-id", "fn", "{}")]); + } + + #[test] + fn errors_when_function_name_is_missing() { + // TS covers `name: undefined` and `name: null`; both are `None`. + let mut h = Harness::new(); + + 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." + ); + } + + #[test] + fn ignores_a_blank_function_name_without_preventing_prior_calls_from_finalizing() { + for name in ["", " "] { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "valid_tool", "{\"value\":1}")) + .unwrap(); + h.delta(start(1, "call_2", name, "{\"value\":2}")).unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![tc("call_1", "valid_tool", "{\"value\":1}")], + "name {name:?}" + ); + } + } + + #[test] + fn retains_continuation_arguments_for_a_blank_name() { + // (blank name, continuation id, continuation index) + let cases = [ + ("", Some("call_1"), None), // blank name with a matching id + (" ", None, Some(0)), // whitespace name, matching index + ]; + for (name, id, index) in cases { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "read_file", "{\"pa")).unwrap(); + h.delta(delta(index, id, None, Some(name), Some("th\":\"a\"}"))) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![tc("call_1", "read_file", "{\"path\":\"a\"}")], + "name {name:?}" + ); + } + } + + #[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)); + + h.delta(delta( + Some(0), + None, + None, + Some("read_file"), + Some("{\"path\":\"p0\"}"), + )) + .unwrap(); + h.delta(delta( + Some(0), + None, + None, + Some("write_file"), + Some("{\"path\":\"p1\"}"), + )) + .unwrap(); + h.delta(delta( + Some(0), + None, + None, + Some("read_file"), + Some("{\"path\":\"p2\"}"), + )) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("generated-1", "read_file", "{\"path\":\"p0\"}"), + tc("generated-2", "write_file", "{\"path\":\"p1\"}"), + tc("generated-3", "read_file", "{\"path\":\"p2\"}"), + ] + ); + } + + #[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()), + ); + + for value in ["{\"value\":1}", "{\"value\":2}"] { + h.delta(delta( + Some(0), + id, + Some("function"), + Some("same_tool"), + Some(value), + )) + .unwrap(); + } + h.flush(); + + let calls = h.tool_calls(); + let inputs: Vec<&str> = calls.iter().map(|c| c.2.as_str()).collect(); + assert_eq!(inputs, vec!["{\"value\":1}", "{\"value\":2}"], "id {id:?}"); + let ids: std::collections::HashSet<&str> = calls.iter().map(|c| c.0.as_str()).collect(); + assert_eq!(ids.len(), 2, "id {id:?}"); + } + } + + #[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()), + ); + + for value in ["{\"value\":1}", "{\"value\":"] { + h.delta(delta( + Some(0), + id, + Some("function"), + Some("same_tool"), + Some(value), + )) + .unwrap(); + } + h.flush(); + + let calls = h.tool_calls(); + let inputs: Vec<&str> = calls.iter().map(|c| c.2.as_str()).collect(); + assert_eq!(inputs, vec!["{\"value\":1}", "{\"value\":"], "id {id:?}"); + let ids: std::collections::HashSet<&str> = calls.iter().map(|c| c.0.as_str()).collect(); + assert_eq!(ids.len(), 2, "id {id:?}"); + } + } + + #[test] + fn keeps_interleaved_same_name_calls_with_distinct_ids_and_reused_index_separate() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "same_tool", "{\"value\":")) + .unwrap(); + h.delta(start(0, "call_2", "same_tool", "{\"value\":2}")) + .unwrap(); + h.delta(delta(Some(0), Some("call_1"), None, None, Some("1}"))) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("call_1", "same_tool", "{\"value\":1}"), + tc("call_2", "same_tool", "{\"value\":2}"), + ] + ); + } + + #[test] + fn ignores_an_index_only_continuation_after_the_index_is_reused() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "first", "{\"value\":1}")) + .unwrap(); + h.delta(start(0, "call_2", "second", "{\"value\":2}")) + .unwrap(); + h.delta(cont(Some(0), "{\"unattributed\":true}")).unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("call_1", "first", "{\"value\":1}"), + tc("call_2", "second", "{\"value\":2}"), + ] + ); + } + + #[test] + fn uses_the_index_when_continuation_ids_are_blank() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "read_file", "{\"pa")).unwrap(); + h.delta(delta(Some(0), Some(" "), None, None, Some("th\":\"a\"}"))) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![tc("call_1", "read_file", "{\"path\":\"a\"}")] + ); + } + + #[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)); + + h.delta(start(0, "", "read_file", "{}")).unwrap(); + h.delta(start(1, "dup", "read_file", "{}")).unwrap(); + h.delta(start(2, "dup", "write_file", "{}")).unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("generated-1", "read_file", "{}"), + tc("dup", "read_file", "{}"), + tc("generated-2", "write_file", "{}"), + ] + ); + } + + #[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()), + ); + + h.delta(start(0, "dup", "same_tool", "{\"value\":0}")) + .unwrap(); + h.delta(start(1, "dup", "same_tool", "{\"value\":1}")) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("dup", "same_tool", "{\"value\":0}"), + tc("generated-id", "same_tool", "{\"value\":1}"), + ] + ); + } + + #[test] + fn preserves_nonblank_ids_and_function_names_exactly() { + let mut h = Harness::new(); + + h.delta(start(0, " spaced ", " same_tool ", "{\"value\":")) + .unwrap(); + h.delta(delta(Some(0), Some(" spaced "), None, None, Some("0}"))) + .unwrap(); + h.delta(start(1, "spaced", " same_tool ", "{\"value\":1}")) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc(" spaced ", " same_tool ", "{\"value\":0}"), + tc("spaced", " same_tool ", "{\"value\":1}"), + ] + ); + } + + #[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)); + + for (index, name) in [(0, "first"), (1, "second"), (2, "third")] { + h.delta(delta( + Some(index), + None, + Some("function"), + Some(name), + Some("{}"), + )) + .unwrap(); + } + h.flush(); + + assert_eq!(calls.load(Ordering::SeqCst), 3); + let ids: Vec = h.tool_calls().into_iter().map(|c| c.0).collect(); + assert_eq!( + ids, + vec!["generated-id", "generated-id-1", "generated-id-2"] + ); + } + + #[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)); + + for (index, name) in [(0, "first"), (1, "second")] { + h.delta(delta( + Some(index), + None, + Some("function"), + Some(name), + Some("{}"), + )) + .unwrap(); + } + h.flush(); + + assert_eq!(calls.load(Ordering::SeqCst), 2); + let ids: Vec = h.tool_calls().into_iter().map(|c| c.0).collect(); + assert_eq!(ids, vec!["tool-call", "tool-call-1"]); + } + + #[test] + fn ignores_unattributable_deltas_when_multiple_calls_are_active() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "read_file", "{\"path\":\"a\"}")) + .unwrap(); + h.delta(start(1, "call_2", "write_file", "{\"path\":\"b\"}")) + .unwrap(); + h.delta(cont(None, "{\"unattributed\":true}")).unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("call_1", "read_file", "{\"path\":\"a\"}"), + tc("call_2", "write_file", "{\"path\":\"b\"}"), + ] + ); + } + + #[test] + fn ignores_an_ambiguous_continuation_for_a_repeated_id() { + let mut h = Harness::with( + StreamingToolCallTracker::new().with_generate_id(|| "generated-id".to_string()), + ); + + h.delta(start(0, "dup", "read_file", "{\"path\":\"a\"}")) + .unwrap(); + h.delta(start(1, "dup", "write_file", "{\"path\":\"b\"}")) + .unwrap(); + h.delta(delta( + None, + Some("dup"), + None, + None, + Some("{\"unattributed\":true}"), + )) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("dup", "read_file", "{\"path\":\"a\"}"), + tc("generated-id", "write_file", "{\"path\":\"b\"}"), + ] + ); + } + + #[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), + ); + + h.delta(start(0, "call_1", "read_file", "{\"pa")).unwrap(); + h.delta(delta( + Some(0), + None, + None, + Some("read_file"), + Some("th\":\"a\"}"), + )) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![tc("call_1", "read_file", "{\"path\":\"a\"}")] + ); + } + + #[test] + fn uses_the_index_when_a_continuation_has_an_unexpected_id() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "read_file", "{\"pa")).unwrap(); + h.delta(delta( + Some(0), + Some("unexpected"), + None, + None, + Some("th\":\"a\"}"), + )) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![tc("call_1", "read_file", "{\"path\":\"a\"}")] + ); + } + + #[test] + fn continues_a_call_when_its_id_changes_but_index_and_name_match() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "read_file", "{\"pa")).unwrap(); + h.delta(start(0, "unexpected", "read_file", "th\":\"a\"}")) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![tc("call_1", "read_file", "{\"path\":\"a\"}")] + ); + } + + #[test] + fn continues_a_call_when_all_labels_repeat_after_a_parsable_argument_prefix() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "calculate", "1")).unwrap(); + h.delta(start(0, "call_1", "calculate", "2")).unwrap(); + h.flush(); + + assert_eq!(h.tool_calls(), vec![tc("call_1", "calculate", "12")]); + } + + #[test] + fn continues_a_structured_argument_when_repeated_labels_precede_a_nested_object() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "calculate", "{\"value\":")) + .unwrap(); + h.delta(start(0, "call_1", "calculate", "{\"nested\":true}}")) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![tc("call_1", "calculate", "{\"value\":{\"nested\":true}}")] + ); + } + + #[test] + fn emits_tool_calls_in_index_order() { + let mut h = Harness::new(); + + h.delta(start(1, "call_1", "second", "{}")).unwrap(); + h.delta(start(0, "call_0", "first", "{}")).unwrap(); + h.flush(); + + let names: Vec = h.tool_calls().into_iter().map(|c| c.1).collect(); + assert_eq!(names, vec!["first", "second"]); + } + + #[test] + fn preserves_insertion_order_when_calls_mix_present_and_omitted_indices() { + let mut h = Harness::new(); + + h.delta(delta( + None, + Some("call_without_index"), + Some("function"), + Some("first"), + Some("{}"), + )) + .unwrap(); + h.delta(start(0, "call_with_index", "second", "{}")) + .unwrap(); + h.flush(); + + let names: Vec = h.tool_calls().into_iter().map(|c| c.1).collect(); + assert_eq!(names, vec!["first", "second"]); + } + + #[test] + fn errors_when_function_name_is_missing_from_a_new_call() { + let mut h = Harness::new(); + + // `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)); + } +} + +mod type_validation { + use super::*; + + fn custom_type() -> StreamingToolCallDelta { + delta( + Some(0), + Some("call_1"), + Some("custom"), + Some("fn"), + Some(""), + ) + } + + fn no_type() -> StreamingToolCallDelta { + 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), + ); + + 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), + ); + + assert_eq!(h.delta(custom_type()), Err(TrackerError::InvalidType)); + assert_eq!( + TrackerError::InvalidType.to_string(), + "Expected 'function' type." + ); + + // A missing type is accepted. + assert_eq!(h.delta(no_type()), Ok(())); + } + + #[test] + fn requires_function_type_with_required() { + let mut h = Harness::with( + StreamingToolCallTracker::new().with_type_validation(TypeValidation::Required), + ); + + assert_eq!(h.delta(no_type()), Err(TrackerError::InvalidType)); + + assert_eq!(h.delta(start(0, "call_1", "fn", "")), Ok(())); + } +} + +mod flush { + use super::*; + + #[test] + fn finalizes_unfinished_tool_calls() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "fn", "{\"key\": \"val")) + .unwrap(); + h.clear(); + + h.flush(); + + assert_eq!( + h.parts, + vec![ + input_end("call_1"), + tool_call("call_1", "fn", "{\"key\": \"val"), + ] + ); + } + + #[test] + fn does_not_re_finalize_already_finished_tool_calls() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "fn", "{}")).unwrap(); + h.flush(); + h.clear(); + + h.flush(); + + assert_eq!(h.parts, vec![]); + } +} + +mod metadata { + use super::*; + + fn google_tracker() -> StreamingToolCallTracker { + StreamingToolCallTracker::new() + .with_extract_metadata(|delta| { + delta.extra["extra_content"]["google"]["thought_signature"] + .as_str() + .map(|sig| json!({ "thoughtSignature": sig })) + }) + .with_build_provider_metadata(|metadata| { + metadata + .and_then(|m| m.get("thoughtSignature")) + .map(|sig| json!({ "google": { "thoughtSignature": sig } })) + }) + } + + #[test] + fn extracts_and_includes_provider_metadata_in_tool_call_parts() { + let mut h = Harness::with_meta(google_tracker()); + + h.delta( + start(0, "call_1", "fn", "{}") + .extra(json!({ "extra_content": { "google": { "thought_signature": "sig123" } } })), + ) + .unwrap(); + h.flush(); + + let tool_call = h + .parts + .iter() + .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })); + assert_eq!( + tool_call, + Some(&ToolCallStreamPart::ToolCall { + tool_call_id: "call_1".into(), + tool_name: "fn".into(), + input: "{}".into(), + provider_metadata: Some(json!({ "google": { "thoughtSignature": "sig123" } })), + }) + ); + } + + #[test] + fn includes_provider_metadata_for_unfinished_tool_calls_finalized_in_flush() { + let mut h = Harness::with_meta( + StreamingToolCallTracker::new() + .with_extract_metadata(|_| Some(json!({ "custom": { "key": "value" } }))) + .with_build_provider_metadata(|metadata| { + metadata.map(|m| json!({ "provider": m })) + }), + ); + + h.delta(start(0, "call_1", "fn", "{\"incomplete")).unwrap(); + h.clear(); + + h.flush(); + + assert_eq!( + h.parts.last(), + Some(&ToolCallStreamPart::ToolCall { + tool_call_id: "call_1".into(), + tool_name: "fn".into(), + input: "{\"incomplete".into(), + provider_metadata: Some(json!({ "provider": { "custom": { "key": "value" } } })), + }) + ); + } + + #[test] + fn omits_provider_metadata_when_the_builder_returns_none() { + let mut h = Harness::with_meta( + StreamingToolCallTracker::<()>::new() + .with_extract_metadata(|_| None) + .with_build_provider_metadata(|_| None), + ); + + h.delta(start(0, "call_1", "fn", "{}")).unwrap(); + h.flush(); + + let tool_call = h + .parts + .iter() + .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })); + assert_eq!(tool_call, Some(&super::tool_call("call_1", "fn", "{}"))); + } +} + +mod generate_id { + use super::*; + + #[test] + fn keeps_the_wire_id_when_present() { + let mut h = Harness::with( + StreamingToolCallTracker::new().with_generate_id(|| "custom-id".to_string()), + ); + + h.delta(start(0, "call_1", "fn", "{\"key\": \"val")) + .unwrap(); + h.clear(); + + h.flush(); + + let ids: Vec = h.tool_calls().into_iter().map(|c| c.0).collect(); + assert_eq!(ids, vec!["call_1"]); + } +} From 49768ff9cc42bbf297c6ea9bca7933e729bab9ef Mon Sep 17 00:00:00 2001 From: cunninghamcard-bit Date: Thu, 1 Oct 2026 16:09:31 +0000 Subject: [PATCH 3/8] feat(provider-utils)!: move the tool-call tracker next to its users and wire the OpenAI stream to it The AI SDK keeps StreamingToolCallTracker in @ai-sdk/provider-utils, where every provider that speaks the OpenAI chat-completions wire format imports it. aimux had the tracker in aimux-stream (the eventsource-parser role), with its own ToolCallStreamPart event type and no callers, while openai/model.rs kept an index-only accumulator. This commit puts the pieces where the AI SDK has them. aimux-stream - StreamingToolCallTracker, StreamingToolCallArgumentState, starts_with_structured_value and ToolCallStreamPart move out; the crate is SSE decoding only. serde_json becomes a dev-dependency (test fixture); the `tool-calls` keyword is dropped. aimux-provider-utils - streaming_tool_call_tracker.rs / streaming_tool_call_argument_state.rs land here, logic unchanged. The tracker emits aimux_core::StreamPart ToolInputStart / ToolInputDelta / ToolInputEnd / ToolCall directly, so there is no second event type (RFC-0036: ToolCallStreamPart must not be a second public protocol). The generic metadata parameter becomes serde_json::Value / ProviderMetadata; TrackerError converts into AiMuxError::InvalidResponseData. The ported upstream tests (45 + 9) move with it, projecting StreamPart back to the upstream event shape. aimux-providers - openai/model.rs replaces its HashMap with the tracker: tool_calls deltas are correlated by wire id, index and function name instead of index alone; id-less continuations follow their call; indices reused across parallel calls stay distinct; ambiguous deltas are dropped; a call without a wire id gets a generated `tool-call` / `tool-call-N` id instead of an empty string; a new call without a function name ends the stream with InvalidResponseData (AI SDK behaviour) instead of starting a call with an empty name. - DeltaToolCall.index is Option (the AI SDK schema is `index: z.number().nullish()`). Every registry-backed provider goes through openai/model.rs, so this is the one plug point until S4-2 folds the mistral and xai chat copies into the same implementation. Verification: cargo test -p aimux-stream -p aimux-provider-utils; cargo test -p aimux-providers (129 suites, 3041 passed, 0 failed, including the cassette replays); cargo clippy -p aimux-stream -p aimux-provider-utils -p aimux-providers --all-targets -- -D warnings; cargo fmt --all -- --check. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WGwanvo9WLU9HWD7sRR8TS (cherry picked from commit df4d09f0353b47ad56926f1c3a623e7a123e880e) --- CONTRIBUTING.md | 4 +- README.md | 8 +- aimux-provider-utils/Cargo.toml | 1 + aimux-provider-utils/src/lib.rs | 12 +- .../src/streaming_tool_call_argument_state.rs | 0 .../src/streaming_tool_call_tracker.rs | 119 +++++++++-------- ...streaming_tool_call_argument_state_test.rs | 2 +- .../tests/streaming_tool_call_tracker_test.rs | 124 +++++++++++++----- aimux-providers/src/openai/model.rs | 109 +++++---------- aimux-providers/src/openai/types.rs | 4 +- aimux-stream/src/lib.rs | 9 -- docs/PROJECT-OVERVIEW.md | 4 +- 12 files changed, 212 insertions(+), 184 deletions(-) rename {aimux-stream => aimux-provider-utils}/src/streaming_tool_call_argument_state.rs (100%) rename {aimux-stream => aimux-provider-utils}/src/streaming_tool_call_tracker.rs (87%) rename {aimux-stream => aimux-provider-utils}/tests/streaming_tool_call_argument_state_test.rs (95%) rename {aimux-stream => aimux-provider-utils}/tests/streaming_tool_call_tracker_test.rs (91%) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index cc1b87ac3..959bcfa67 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -9,8 +9,8 @@ to set up a development environment, run the tests, and submit changes. aimux/ ├── aimux-core/ # Core abstractions: LanguageModel / Provider / Message / StreamPart ├── aimux-providers/ # provider implementations + cassettes (counts: docs/api/providers.md) -├── aimux-stream/ # SSE parsing, streamed tool-call tracking -├── aimux-provider-utils/ # HTTP utilities: retry, backoff, error parsing, API-key loading +├── aimux-stream/ # SSE parsing +├── aimux-provider-utils/ # HTTP utilities: retry, backoff, error parsing, API-key loading, streamed tool-call tracking ├── aimux-ffi/ # C ABI (opaque handle + JSON + push callback) for non-native bindings ├── bindings/ # Node, Python, Swift, Kotlin, Flutter, Go, C — share one Rust core ├── rfc/ # Design docs (RFCs) diff --git a/README.md b/README.md index f0654816a..ef7630255 100644 --- a/README.md +++ b/README.md @@ -98,8 +98,8 @@ middleware, and telemetry per request). aimux/ ├── aimux-core # Core abstractions: LanguageModel / Provider / Message / StreamPart ├── aimux-providers # Provider implementations — registry-backed + typed (docs/api/providers.md) -├── aimux-stream # SSE parsing, streamed tool-call tracking -├── aimux-provider-utils # One-exchange HTTP helpers, response handlers, API-key loading +├── aimux-stream # SSE parsing +├── aimux-provider-utils # One-exchange HTTP helpers, response handlers, API-key loading, streamed tool-call tracking ├── aimux-ffi # C ABI (opaque handles + JSON results + owned aimux_error_t *) for non-native bindings └── tools/ # aimux-cli (cache probe) · aimux-replay · aimux-web (console) ``` @@ -122,8 +122,8 @@ cargo add aimux-core aimux-providers |-------|-------------|-----------| | `aimux-core` | Core abstractions: `LanguageModel` / `Provider` / `Message` / `StreamPart` | [crates.io](https://crates.io/crates/aimux-core) | | `aimux-providers` | Provider implementations — [registry-backed + typed](docs/api/providers.md) | [crates.io](https://crates.io/crates/aimux-providers) | -| `aimux-stream` | SSE parsing, streamed tool-call tracking | [crates.io](https://crates.io/crates/aimux-stream) | -| `aimux-provider-utils` | One-exchange HTTP helpers and typed response handlers | [crates.io](https://crates.io/crates/aimux-provider-utils) | +| `aimux-stream` | SSE parsing | [crates.io](https://crates.io/crates/aimux-stream) | +| `aimux-provider-utils` | One-exchange HTTP helpers, typed response handlers, streamed tool-call tracking | [crates.io](https://crates.io/crates/aimux-provider-utils) | | `aimux-ffi` | C ABI for non-native bindings | [crates.io](https://crates.io/crates/aimux-ffi) | **Node.js**: diff --git a/aimux-provider-utils/Cargo.toml b/aimux-provider-utils/Cargo.toml index 29a6049f2..bf427a093 100644 --- a/aimux-provider-utils/Cargo.toml +++ b/aimux-provider-utils/Cargo.toml @@ -16,6 +16,7 @@ reqwest = { workspace = true } http = "1" serde = { workspace = true } serde_json = { workspace = true } +thiserror = { workspace = true } tokio = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } diff --git a/aimux-provider-utils/src/lib.rs b/aimux-provider-utils/src/lib.rs index b6d43f16b..f3b0245c0 100644 --- a/aimux-provider-utils/src/lib.rs +++ b/aimux-provider-utils/src/lib.rs @@ -3,7 +3,8 @@ //! Shared utilities for provider implementations. //! //! Provides one-exchange HTTP helpers, response handlers, API key loading, -//! header management, and URL utilities — the Rust equivalents of +//! header management, URL utilities and the streamed tool-call tracker for +//! the OpenAI chat-completions wire format — the Rust equivalents of //! `@ai-sdk/provider-utils`. Operation retry and timeout live in `aimux-core`. pub mod api_key; @@ -19,6 +20,8 @@ pub mod post_to_api; pub mod read_response_with_size_limit; pub mod response_handler; pub mod retry; +pub mod streaming_tool_call_argument_state; +pub mod streaming_tool_call_tracker; pub mod url; /// WebSocket client for realtime provider APIs (RFC-0028). Empty unless the /// `ws` feature is enabled. @@ -43,4 +46,11 @@ pub use response_handler::{ invalid_response_api_call, stream_error_api_call, }; pub use retry::RetryConfig; +pub use streaming_tool_call_argument_state::{ + StreamingToolCallArgumentState, starts_with_structured_value, +}; +pub use streaming_tool_call_tracker::{ + StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, TrackerError, + TypeValidation, +}; pub use url::{validate_base_url, without_trailing_slash, without_trailing_slash_opt}; diff --git a/aimux-stream/src/streaming_tool_call_argument_state.rs b/aimux-provider-utils/src/streaming_tool_call_argument_state.rs similarity index 100% rename from aimux-stream/src/streaming_tool_call_argument_state.rs rename to aimux-provider-utils/src/streaming_tool_call_argument_state.rs diff --git a/aimux-stream/src/streaming_tool_call_tracker.rs b/aimux-provider-utils/src/streaming_tool_call_tracker.rs similarity index 87% rename from aimux-stream/src/streaming_tool_call_tracker.rs rename to aimux-provider-utils/src/streaming_tool_call_tracker.rs index 9a8d7240e..ce9666dd8 100644 --- a/aimux-stream/src/streaming_tool_call_tracker.rs +++ b/aimux-provider-utils/src/streaming_tool_call_tracker.rs @@ -5,8 +5,9 @@ //! //! Tracks streaming tool call state across the deltas of an OpenAI-compatible //! chat completion stream: accumulates `arguments` fragments, emits -//! `tool-input-start` / `tool-input-delta` / `tool-input-end` / `tool-call` -//! parts, and finalizes unfinished calls on [`StreamingToolCallTracker::flush`]. +//! [`StreamPart::ToolInputStart`] / [`StreamPart::ToolInputDelta`] / +//! [`StreamPart::ToolInputEnd`] / [`StreamPart::ToolCall`] parts, and +//! finalizes unfinished calls on [`StreamingToolCallTracker::flush`]. //! //! Deltas are correlated to calls by wire `id`, `index` and function name //! (see [`StreamingToolCallTracker`]'s resolution table), not by `index` @@ -16,12 +17,22 @@ //! Like the TS original, a call is *never* finalized before `flush`: a //! parsable argument buffer can still be the prefix of a longer argument //! string, so acting on it early would use truncated inputs (ai-sdk #13137). +//! +//! This is a tool for the OpenAI chat-completions wire format only. Protocols +//! whose streams carry explicit tool-call boundaries (Anthropic content +//! blocks, Google complete `functionCall` parts, Bedrock content blocks, +//! Cohere `tool-call-*` events, the Responses API's output items) do not need +//! it. use std::collections::{HashMap, HashSet}; use serde_json::Value; use thiserror::Error; +use aimux_core::error::AiMuxError; +use aimux_core::stream_part::StreamPart; +use aimux_core::types::ProviderMetadata; + use crate::streaming_tool_call_argument_state::{ StreamingToolCallArgumentState, starts_with_structured_value, }; @@ -98,28 +109,6 @@ impl StreamingToolCallDelta { } } -/// The stream parts emitted by [`StreamingToolCallTracker`]. -/// -/// Mirrors the subset of `LanguageModelV4StreamPart` the TS tracker enqueues. -/// [`ToolCallStreamPart::ToolCall::input`] is the raw accumulated argument -/// *string*; the tracker does not parse it. -#[derive(Debug, Clone, PartialEq)] -pub enum ToolCallStreamPart { - /// Start of a tool call's input streaming. - ToolInputStart { id: String, tool_name: String }, - /// A partial argument fragment. - ToolInputDelta { id: String, delta: String }, - /// End of a tool call's input streaming. - ToolInputEnd { id: String }, - /// A complete, finalized tool call. - ToolCall { - tool_call_id: String, - tool_name: String, - input: String, - provider_metadata: Option, - }, -} - /// How to validate the `type` field on a new tool call delta. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum TypeValidation { @@ -133,6 +122,9 @@ pub enum TypeValidation { } /// Errors raised while processing a tool call delta. +/// +/// The TS tracker throws `InvalidResponseDataError`; these convert into +/// [`AiMuxError::InvalidResponseData`]. #[derive(Debug, Error, PartialEq, Eq)] pub enum TrackerError { #[error("Expected 'function.name' to be a string.")] @@ -143,7 +135,13 @@ pub enum TrackerError { IdExhausted, } -struct TrackedToolCall { +impl From for AiMuxError { + fn from(error: TrackerError) -> Self { + AiMuxError::InvalidResponseData(error.to_string()) + } +} + +struct TrackedToolCall { id: String, index: Option, sequence: usize, @@ -151,7 +149,7 @@ struct TrackedToolCall { arguments: String, argument_state: StreamingToolCallArgumentState, has_finished: bool, - metadata: Option, + metadata: Option, } enum ToolCallResolution { @@ -161,17 +159,14 @@ enum ToolCallResolution { } type GenerateIdFn = Box String + Send + Sync>; -type ExtractMetadataFn = Box Option + Send + Sync>; -type BuildMetadataFn = Box) -> Option + Send + Sync>; +type ExtractMetadataFn = Box Option + Send + Sync>; +type BuildMetadataFn = Box) -> Option + Send + Sync>; /// Tracks streaming tool call state across multiple deltas. /// /// [`process_delta`](Self::process_delta) and [`flush`](Self::flush) return -/// the parts to forward downstream (the TS tracker enqueues them on a -/// controller instead). -/// -/// The type parameter `M` is the provider-metadata type; use `()` (the -/// default) when no metadata handling is needed. +/// the [`StreamPart`]s to forward downstream (the TS tracker enqueues them on +/// a controller instead). /// /// # Correlation /// @@ -186,25 +181,25 @@ type BuildMetadataFn = Box) -> Option + Send + Sync>; /// | absent | absent | unnamed | sole unfinished call, new call, or ambiguity | /// /// An *ambiguous* delta is dropped. -pub struct StreamingToolCallTracker { - tool_calls: Vec>, +pub struct StreamingToolCallTracker { + tool_calls: Vec, tool_calls_by_id: HashMap>, tool_calls_by_index: HashMap>, used_tool_call_ids: HashSet, next_generated_id_suffixes: HashMap, generate_id: GenerateIdFn, type_validation: TypeValidation, - extract_metadata: Option>, - build_provider_metadata: Option>, + extract_metadata: Option, + build_provider_metadata: Option, } -impl Default for StreamingToolCallTracker { +impl Default for StreamingToolCallTracker { fn default() -> Self { Self::new() } } -impl StreamingToolCallTracker { +impl StreamingToolCallTracker { /// Create a tracker with no metadata handling and default settings. /// /// The default id generator returns a blank-free constant, so ids that @@ -240,10 +235,11 @@ impl StreamingToolCallTracker { } /// Set the metadata extractor (the TS `extractMetadata` option). Called - /// once when a new tool call is detected. + /// once when a new tool call is detected; the value is kept with the call + /// and handed to the provider-metadata builder when the call finalizes. #[must_use] pub fn with_extract_metadata< - F: Fn(&StreamingToolCallDelta) -> Option + Send + Sync + 'static, + F: Fn(&StreamingToolCallDelta) -> Option + Send + Sync + 'static, >( mut self, f: F, @@ -254,9 +250,11 @@ impl StreamingToolCallTracker { /// Set the provider-metadata builder (the TS /// `buildToolCallProviderMetadata` option). If it returns `None`, the - /// `tool-call` part carries no `provider_metadata`. + /// `ToolCall` part carries no `provider_metadata`. #[must_use] - pub fn with_build_provider_metadata) -> Option + Send + Sync + 'static>( + pub fn with_build_provider_metadata< + F: Fn(Option<&Value>) -> Option + Send + Sync + 'static, + >( mut self, f: F, ) -> Self { @@ -276,7 +274,7 @@ impl StreamingToolCallTracker { pub fn process_delta( &mut self, delta: &StreamingToolCallDelta, - ) -> Result>, TrackerError> { + ) -> Result, TrackerError> { let wire_name = delta.function.as_ref().and_then(|f| f.name.as_deref()); let has_blank_name = wire_name.is_some_and(|n| n.trim().is_empty()); let wire_id = non_blank(delta.id.as_deref()); @@ -323,7 +321,7 @@ impl StreamingToolCallTracker { /// Finalize any unfinished tool calls and return the closing parts. Call /// once when the stream ends. - pub fn flush(&mut self) -> Vec> { + pub fn flush(&mut self) -> Vec { // Index order is only reliable when every call has an index; for // mixed streams keep insertion order. let mut order: Vec = (0..self.tool_calls.len()).collect(); @@ -475,7 +473,7 @@ impl StreamingToolCallTracker { wire_id: Option<&str>, index: Option, name: Option<&str>, - parts: &mut Vec>, + parts: &mut Vec, ) -> Result { match self.type_validation { TypeValidation::Required => { @@ -494,9 +492,13 @@ impl StreamingToolCallTracker { let name = name.ok_or(TrackerError::MissingFunctionName)?; let id = self.create_tool_call_id(wire_id)?; - parts.push(ToolCallStreamPart::ToolInputStart { + parts.push(StreamPart::ToolInputStart { id: id.clone(), tool_name: name.to_string(), + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: None, }); let metadata = self @@ -526,9 +528,10 @@ impl StreamingToolCallTracker { } if !initial_arguments.is_empty() { - parts.push(ToolCallStreamPart::ToolInputDelta { + parts.push(StreamPart::ToolInputDelta { id, delta: initial_arguments, + provider_metadata: None, }); } @@ -541,7 +544,7 @@ impl StreamingToolCallTracker { &mut self, call: usize, arguments: Option<&str>, - parts: &mut Vec>, + parts: &mut Vec, ) { let tool_call = &mut self.tool_calls[call]; if tool_call.has_finished { @@ -550,9 +553,10 @@ impl StreamingToolCallTracker { if let Some(arguments) = arguments { tool_call.argument_state.append(arguments); tool_call.arguments.push_str(arguments); - parts.push(ToolCallStreamPart::ToolInputDelta { + parts.push(StreamPart::ToolInputDelta { id: tool_call.id.clone(), delta: arguments.to_string(), + provider_metadata: None, }); } } @@ -605,12 +609,13 @@ impl StreamingToolCallTracker { Err(TrackerError::IdExhausted) } - fn finish_tool_call(&mut self, call: usize, parts: &mut Vec>) { + fn finish_tool_call(&mut self, call: usize, parts: &mut Vec) { let tool_call = &mut self.tool_calls[call]; tool_call.has_finished = true; - parts.push(ToolCallStreamPart::ToolInputEnd { + parts.push(StreamPart::ToolInputEnd { id: tool_call.id.clone(), + provider_metadata: None, }); let provider_metadata = self @@ -618,10 +623,16 @@ impl StreamingToolCallTracker { .as_ref() .and_then(|build| build(tool_call.metadata.as_ref())); - parts.push(ToolCallStreamPart::ToolCall { + parts.push(StreamPart::ToolCall { tool_call_id: tool_call.id.clone(), tool_name: tool_call.function_name.clone(), - input: tool_call.arguments.clone(), + // Raw argument text; Core parses it after `stream_text`. + input: Value::String(tool_call.arguments.clone()), + provider_executed: None, + dynamic: None, + thought_signature: None, + invalid: None, + error: None, provider_metadata, }); } diff --git a/aimux-stream/tests/streaming_tool_call_argument_state_test.rs b/aimux-provider-utils/tests/streaming_tool_call_argument_state_test.rs similarity index 95% rename from aimux-stream/tests/streaming_tool_call_argument_state_test.rs rename to aimux-provider-utils/tests/streaming_tool_call_argument_state_test.rs index 44cf754e2..fc05f4268 100644 --- a/aimux-stream/tests/streaming_tool_call_argument_state_test.rs +++ b/aimux-provider-utils/tests/streaming_tool_call_argument_state_test.rs @@ -1,7 +1,7 @@ //! Port of `streaming-tool-call-argument-state.test.ts` from //! `@ai-sdk/provider-utils`. -use aimux_stream::{StreamingToolCallArgumentState, starts_with_structured_value}; +use aimux_provider_utils::{StreamingToolCallArgumentState, starts_with_structured_value}; #[test] fn starts_with_structured_value_true() { diff --git a/aimux-stream/tests/streaming_tool_call_tracker_test.rs b/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs similarity index 91% rename from aimux-stream/tests/streaming_tool_call_tracker_test.rs rename to aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs index 7c55044c4..238a0ecb9 100644 --- a/aimux-stream/tests/streaming_tool_call_tracker_test.rs +++ b/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs @@ -1,35 +1,99 @@ //! Port of `streaming-tool-call-tracker.test.ts` from //! `@ai-sdk/provider-utils`, case for case. +//! +//! Tracker parts are projected to a local `Part` enum so the assertions stay +//! byte-for-byte comparable with the upstream expectations. use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; -use aimux_stream::{ - StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, - ToolCallStreamPart, TrackerError, TypeValidation, +use aimux_core::stream_part::StreamPart; +use aimux_provider_utils::{ + StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, TrackerError, + TypeValidation, }; use serde_json::{Value, json}; -type Part = ToolCallStreamPart; +/// The TS tracker's four events, projected out of [`StreamPart`] so the +/// assertions below read like the upstream suite. +#[derive(Debug, Clone, PartialEq)] +#[allow( + clippy::enum_variant_names, + reason = "named after the upstream stream parts" +)] +enum Part { + ToolInputStart { + id: String, + tool_name: String, + }, + ToolInputDelta { + id: String, + delta: String, + }, + ToolInputEnd { + id: String, + }, + ToolCall { + tool_call_id: String, + tool_name: String, + /// The raw argument text (`StreamPart::ToolCall.input` is a + /// `Value::String` from the tracker). + input: String, + provider_metadata: Option, + }, +} + +fn project(part: StreamPart) -> Part { + match part { + StreamPart::ToolInputStart { + id, + tool_name, + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: None, + } => Part::ToolInputStart { id, tool_name }, + StreamPart::ToolInputDelta { + id, + delta, + provider_metadata: None, + } => Part::ToolInputDelta { id, delta }, + StreamPart::ToolInputEnd { + id, + provider_metadata: None, + } => Part::ToolInputEnd { id }, + StreamPart::ToolCall { + tool_call_id, + tool_name, + input: Value::String(input), + provider_executed: None, + dynamic: None, + thought_signature: None, + invalid: None, + error: None, + provider_metadata, + } => Part::ToolCall { + tool_call_id, + tool_name, + input, + provider_metadata, + }, + other => panic!("tracker emitted an unexpected part: {other:?}"), + } +} /// Tracker plus the parts it has emitted so far (the TS `createCollector`). -struct Harness { - tracker: StreamingToolCallTracker, - parts: Vec>, +struct Harness { + tracker: StreamingToolCallTracker, + parts: Vec, } -impl Harness<()> { +impl Harness { fn new() -> Self { Self::with(StreamingToolCallTracker::new()) } - fn with(tracker: StreamingToolCallTracker<()>) -> Self { - Self::with_meta(tracker) - } -} - -impl Harness { - fn with_meta(tracker: StreamingToolCallTracker) -> Self { + fn with(tracker: StreamingToolCallTracker) -> Self { Self { tracker, parts: Vec::new(), @@ -38,21 +102,19 @@ impl Harness { fn delta(&mut self, delta: StreamingToolCallDelta) -> Result<(), TrackerError> { let parts = self.tracker.process_delta(&delta)?; - self.parts.extend(parts); + self.parts.extend(parts.into_iter().map(project)); Ok(()) } fn flush(&mut self) { let parts = self.tracker.flush(); - self.parts.extend(parts); + self.parts.extend(parts.into_iter().map(project)); } fn clear(&mut self) { self.parts.clear(); } -} -impl Harness { /// `(id, name, input)` of every emitted `tool-call`, in order. fn tool_calls(&self) -> Vec<(String, String, String)> { self.parts @@ -992,7 +1054,7 @@ mod flush { mod metadata { use super::*; - fn google_tracker() -> StreamingToolCallTracker { + fn google_tracker() -> StreamingToolCallTracker { StreamingToolCallTracker::new() .with_extract_metadata(|delta| { delta.extra["extra_content"]["google"]["thought_signature"] @@ -1008,7 +1070,7 @@ mod metadata { #[test] fn extracts_and_includes_provider_metadata_in_tool_call_parts() { - let mut h = Harness::with_meta(google_tracker()); + let mut h = Harness::with(google_tracker()); h.delta( start(0, "call_1", "fn", "{}") @@ -1017,13 +1079,10 @@ mod metadata { .unwrap(); h.flush(); - let tool_call = h - .parts - .iter() - .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })); + let tool_call = h.parts.iter().find(|p| matches!(p, Part::ToolCall { .. })); assert_eq!( tool_call, - Some(&ToolCallStreamPart::ToolCall { + Some(&Part::ToolCall { tool_call_id: "call_1".into(), tool_name: "fn".into(), input: "{}".into(), @@ -1034,7 +1093,7 @@ mod metadata { #[test] fn includes_provider_metadata_for_unfinished_tool_calls_finalized_in_flush() { - let mut h = Harness::with_meta( + let mut h = Harness::with( StreamingToolCallTracker::new() .with_extract_metadata(|_| Some(json!({ "custom": { "key": "value" } }))) .with_build_provider_metadata(|metadata| { @@ -1049,7 +1108,7 @@ mod metadata { assert_eq!( h.parts.last(), - Some(&ToolCallStreamPart::ToolCall { + Some(&Part::ToolCall { tool_call_id: "call_1".into(), tool_name: "fn".into(), input: "{\"incomplete".into(), @@ -1060,8 +1119,8 @@ mod metadata { #[test] fn omits_provider_metadata_when_the_builder_returns_none() { - let mut h = Harness::with_meta( - StreamingToolCallTracker::<()>::new() + let mut h = Harness::with( + StreamingToolCallTracker::new() .with_extract_metadata(|_| None) .with_build_provider_metadata(|_| None), ); @@ -1069,10 +1128,7 @@ mod metadata { h.delta(start(0, "call_1", "fn", "{}")).unwrap(); h.flush(); - let tool_call = h - .parts - .iter() - .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })); + let tool_call = h.parts.iter().find(|p| matches!(p, Part::ToolCall { .. })); assert_eq!(tool_call, Some(&super::tool_call("call_1", "fn", "{}"))); } } diff --git a/aimux-providers/src/openai/model.rs b/aimux-providers/src/openai/model.rs index 6445d52c2..107f984ee 100644 --- a/aimux-providers/src/openai/model.rs +++ b/aimux-providers/src/openai/model.rs @@ -20,7 +20,9 @@ use aimux_core::result::{GenerateContent, GenerateResult, StreamResult}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; -use aimux_provider_utils::HttpRequest; +use aimux_provider_utils::{ + HttpRequest, StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, +}; use super::OpenAIConfig; use super::convert::{RequestBodyResult, build_request_body_with_warnings, parse_finish_reason}; @@ -176,15 +178,6 @@ fn convert_usage(usage: &UsageResponse, usage_raw: Option<&Value>) -> Usage { } } -// ── Tool-call accumulator (streaming) ──────────────────────────────────────── - -/// Accumulates a streamed tool call's id, name, and argument fragments. -struct ToolCallAccumulator { - id: String, - name: String, - arguments: String, -} - #[async_trait] impl LanguageModel for OpenAIModel { /// Provider identity for recording/routing. Uses `config.provider` (the @@ -491,9 +484,9 @@ pub async fn execute_stream( let mut response_metadata_emitted = false; let mut final_logprobs: Option = None; - // Tool-call accumulators keyed by OpenAI's `index` field. - let mut tool_calls: HashMap = HashMap::new(); - let mut tool_call_order: Vec = Vec::new(); + // Streamed tool calls, correlated by wire id, index and function name + // (the AI SDK's StreamingToolCallTracker) and finalized on flush. + let mut tool_calls = StreamingToolCallTracker::new(); // Process the first event (already peeked) then the rest. let mut event_iter = @@ -643,49 +636,30 @@ pub async fn execute_stream( reasoning_started = false; } for dtc in tool_call_deltas { - let idx = dtc.index; - let func = dtc.function.unwrap_or_default(); - - // New tool call: has id and/or name. - let is_new = !tool_calls.contains_key(&idx); - if is_new { - let id = dtc.id.unwrap_or_default(); - let name = func.name.unwrap_or_default(); - tool_calls.insert( - idx, - ToolCallAccumulator { - id: id.clone(), - name: name.clone(), - arguments: String::new(), - }, - ); - tool_call_order.push(idx); - yield Ok(StreamPart::ToolInputStart { - id, - tool_name: name, - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - } - - // Argument delta. - // For new tool calls, skip the delta when - // arguments are empty (matches TS — the - // initial `""` is not emitted). For - // continuation chunks, always emit (even - // empty, matching TS). - if let Some(args) = func.arguments - && (!is_new || !args.is_empty()) - && let Some(acc) = tool_calls.get_mut(&idx) { - acc.arguments.push_str(&args); - yield Ok(StreamPart::ToolInputDelta { - id: acc.id.clone(), - delta: args, - provider_metadata: None, - }); + let delta = StreamingToolCallDelta { + index: dtc.index, + id: dtc.id, + r#type: None, + function: dtc.function.map(|f| StreamingToolCallFunction { + name: f.name, + arguments: f.arguments, + }), + extra: Value::Null, + }; + match tool_calls.process_delta(&delta) { + Ok(parts) => { + for part in parts { + yield Ok(part); } + } + // A malformed delta (new call without a + // function name) is invalid response + // data, as in the AI SDK; the stream ends. + Err(error) => { + yield Err(error.into()); + return; + } + } } } @@ -763,27 +737,10 @@ pub async fn execute_stream( }); } - // A parsable argument buffer can still be a prefix of a longer input. - // Match AI SDK's tracker by finalizing only when the stream flushes. - for &idx in &tool_call_order { - if let Some(acc) = tool_calls.get(&idx) { - yield Ok(StreamPart::ToolInputEnd { - id: acc.id.clone(), - provider_metadata: None, - }); - let input = Value::String(acc.arguments.clone()); - yield Ok(StreamPart::ToolCall { - tool_call_id: acc.id.clone(), - tool_name: acc.name.clone(), - input, - provider_executed: None, - dynamic: None, - thought_signature: None, - invalid: None, - error: None, - provider_metadata: None, - }); - } + // 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. + for part in tool_calls.flush() { + yield Ok(part); } // Build provider metadata for the Finish part. diff --git a/aimux-providers/src/openai/types.rs b/aimux-providers/src/openai/types.rs index ea78af6ec..1a8c83400 100644 --- a/aimux-providers/src/openai/types.rs +++ b/aimux-providers/src/openai/types.rs @@ -156,8 +156,10 @@ pub struct Delta { #[derive(Debug, Deserialize)] pub struct DeltaToolCall { + /// Optional: some OpenAI-compatible providers omit it (the AI SDK schema + /// is `index: z.number().nullish()`); the tracker falls back to id/name. #[serde(default)] - pub index: usize, + pub index: Option, #[serde(default)] pub id: Option, #[serde(default)] diff --git a/aimux-stream/src/lib.rs b/aimux-stream/src/lib.rs index d9078d8d8..65c368daf 100644 --- a/aimux-stream/src/lib.rs +++ b/aimux-stream/src/lib.rs @@ -6,16 +6,7 @@ pub mod lines; pub mod sse; -pub mod streaming_tool_call_argument_state; -pub mod streaming_tool_call_tracker; // Re-export the most commonly used items. pub use lines::extract_lines; pub use sse::{SseError, SseEvent, SseStream}; -pub use streaming_tool_call_argument_state::{ - StreamingToolCallArgumentState, starts_with_structured_value, -}; -pub use streaming_tool_call_tracker::{ - StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, - ToolCallStreamPart, TrackerError, TypeValidation, -}; diff --git a/docs/PROJECT-OVERVIEW.md b/docs/PROJECT-OVERVIEW.md index 3207e2521..b29c96998 100644 --- a/docs/PROJECT-OVERVIEW.md +++ b/docs/PROJECT-OVERVIEW.md @@ -165,8 +165,8 @@ aimux/ │ ├── native protocols # standalone model + convert, handles provider-specific differences │ ├── OpenAI compatible # registry-backed: provider-registry.json + provider(name, ...) entry (RFC-0017 phase 4) │ └── modalities/search # voice / image / video / search implementations -├── aimux-stream # SSE parsing, streamed tool-call tracking -├── aimux-provider-utils # One-exchange HTTP helpers, response handlers, API-key loading +├── aimux-stream # SSE parsing +├── aimux-provider-utils # One-exchange HTTP helpers, response handlers, API-key loading, streamed tool-call tracking ├── aimux-ffi # C ABI (FFI infrastructure, shared by all bindings) └── bindings/ # 6 language bindings ├── node/ # napi-rs v3 + typed TS wrapper From 47c495b419bc87aec09e2cf41043875af5e57562 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 04:49:27 +0000 Subject: [PATCH 4/8] test(provider-utils): note the None/undefined mapping in the argument-state port Carried over from the SSE test port commit so the two split branches leave the file identical. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WGwanvo9WLU9HWD7sRR8TS (cherry picked from commit 7bc200eae6cdb57f9af4ce58428b4a48e79ad0d3) --- .../tests/streaming_tool_call_argument_state_test.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/aimux-provider-utils/tests/streaming_tool_call_argument_state_test.rs b/aimux-provider-utils/tests/streaming_tool_call_argument_state_test.rs index fc05f4268..ea6d87076 100644 --- a/aimux-provider-utils/tests/streaming_tool_call_argument_state_test.rs +++ b/aimux-provider-utils/tests/streaming_tool_call_argument_state_test.rs @@ -12,6 +12,7 @@ fn starts_with_structured_value_true() { #[test] fn starts_with_structured_value_false() { + // TS covers `undefined` and `null` separately; Rust has one `None`. assert!(!starts_with_structured_value(None)); for value in ["", " ", "1", "\"value\""] { assert!(!starts_with_structured_value(Some(value)), "{value:?}"); From 9fad0522528f86266f4ecb0e4c1dec1a0ef9fc49 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 04:49:27 +0000 Subject: [PATCH 5/8] docs(changelog): record the tracker move to aimux-provider-utils Describes only the final API: the tracker lives in aimux-provider-utils and emits StreamPart directly; the intermediate ToolCallStreamPart- returning API from earlier commits is not mentioned. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WGwanvo9WLU9HWD7sRR8TS (cherry picked from commit 264ade36cb15d6801a443cac8cb084415604210c) --- CHANGELOG.md | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8157ad3a5..b05151d85 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -41,6 +41,39 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Removed `NdjsonStream` / `NdjsonError`: nothing in the workspace used them and the AI SDK has no counterpart. `tokio` is now a dev-dependency only and the unused direct `serde` dependency is dropped. +- `StreamingToolCallTracker`, `ToolCallStreamPart` and `TrackerError` are no + longer exported from `aimux-stream`; nothing in the workspace used them and + the AI SDK keeps the tracker in `@ai-sdk/provider-utils`, next to the + provider code that drives it. See `aimux-provider-utils` below. + +**Rust (aimux-provider-utils)** + +- Gains `StreamingToolCallTracker` (plus `StreamingToolCallDelta`, + `StreamingToolCallFunction`, `TypeValidation`, `TrackerError`, + `StreamingToolCallArgumentState` and `starts_with_structured_value`), + a port of the current `@ai-sdk/provider-utils` tracker. The previous + `aimux-stream` version was a port of an older, index-only tracker. Deltas + are correlated by wire `id`, `index` and function name (ambiguous deltas + are dropped), ids are de-duplicated with bounded suffixes, blank function + names are ignored, and `flush` orders calls by index only when every call + has one. It emits `aimux_core::StreamPart` tool-input parts directly + (`ToolInputStart` / `ToolInputDelta` / `ToolInputEnd` / `ToolCall`); there + is no separate `ToolCallStreamPart` event type. `TrackerError` converts + into `AiMuxError::InvalidResponseData`. Metadata hooks are typed with + `serde_json::Value` / `ProviderMetadata`, and the builder closures must be + `Send + Sync`. + +**Rust (aimux-providers)** + +- The OpenAI chat-completions stream (`openai/model.rs`, which serves every + registry-backed provider) correlates `tool_calls` deltas with the tracker + instead of by `index` alone: deltas are matched by wire id, index and + function name; a continuation without an id follows its call; indices + reused across parallel calls stay distinct; ambiguous deltas are dropped; + a call whose delta carries no id gets a generated `tool-call` / + `tool-call-N` id instead of an empty string; a new call without a function + name ends the stream with `InvalidResponseData` (previously it started a + call with an empty name). `DeltaToolCall.index` is now `Option`. ## [0.5.0] - 2026-09-27 From 9fd2973403cd1d519b0b3862777b8bcccb848ce6 Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Fri, 9 Oct 2026 12:15:23 +0800 Subject: [PATCH 6/8] refactor(provider-utils)!: rewrite the tool-call tracker in idiomatic Rust Same behaviour as the upstream port (checked differentially against the previous implementation on ~1.1M random deltas), with Rust structure: - Correlate by linear scan over per-call wire ids and indices instead of id/index hash maps re-sorted on every delta; no per-delta allocations besides the emitted parts. - `finish(self)` consumes the tracker, so the never-true-before-flush `has_finished` state is gone and post-flush use cannot compile. - `StreamingToolCallDelta<'a>` borrows from the provider wire type and carries `provider_metadata`, replacing `StreamingToolCallFunction`, the untyped `extra` bag, the test-only builders and the two boxed metadata hooks. - `process` / `finish` append to a caller-owned buffer. - Drop the unreachable `TrackerError::IdExhausted`. - The argument state scans bytes and models completion as its own state. Co-authored-by: Amp Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee (cherry picked from commit b0fd923a3f312ee179c2e72a62e795ea44fb6cff) --- CHANGELOG.md | 30 +- aimux-provider-utils/src/lib.rs | 3 +- .../src/streaming_tool_call_argument_state.rs | 143 ++-- .../src/streaming_tool_call_tracker.rs | 747 ++++++------------ .../tests/streaming_tool_call_tracker_test.rs | 158 ++-- aimux-providers/src/openai/model.rs | 44 +- aimux-stream/src/lib.rs | 5 +- bindings/node/src/types/InputRecord.ts | 4 +- bindings/node/src/types/OutcomeRecord.ts | 6 +- 9 files changed, 414 insertions(+), 726 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b05151d85..87669b4f4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -49,19 +49,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 **Rust (aimux-provider-utils)** - Gains `StreamingToolCallTracker` (plus `StreamingToolCallDelta`, - `StreamingToolCallFunction`, `TypeValidation`, `TrackerError`, - `StreamingToolCallArgumentState` and `starts_with_structured_value`), - a port of the current `@ai-sdk/provider-utils` tracker. The previous - `aimux-stream` version was a port of an older, index-only tracker. Deltas - are correlated by wire `id`, `index` and function name (ambiguous deltas - are dropped), ids are de-duplicated with bounded suffixes, blank function - names are ignored, and `flush` orders calls by index only when every call - has one. It emits `aimux_core::StreamPart` tool-input parts directly - (`ToolInputStart` / `ToolInputDelta` / `ToolInputEnd` / `ToolCall`); there - is no separate `ToolCallStreamPart` event type. `TrackerError` converts - into `AiMuxError::InvalidResponseData`. Metadata hooks are typed with - `serde_json::Value` / `ProviderMetadata`, and the builder closures must be - `Send + Sync`. + `TypeValidation`, `TrackerError`, `StreamingToolCallArgumentState` and + `starts_with_structured_value`), a port of the current + `@ai-sdk/provider-utils` tracker. The previous `aimux-stream` version was + a port of an older, index-only tracker. Deltas are correlated by wire + `id`, `index` and function name (ambiguous deltas are dropped), ids are + de-duplicated with `-N` suffixes, blank function names are ignored, and + `finish` orders calls by index only when every call has one. It emits + `aimux_core::StreamPart` tool-input parts directly (`ToolInputStart` / + `ToolInputDelta` / `ToolInputEnd` / `ToolCall`); there is no separate + `ToolCallStreamPart` event type. `StreamingToolCallDelta` borrows from the + provider's wire type and flattens upstream's `function` object; + `process` and `finish` append to a caller-owned `Vec`, and + `finish` consumes the tracker. The upstream `extractMetadata` / + `buildToolCallProviderMetadata` hooks are replaced by + `StreamingToolCallDelta::provider_metadata`, read from the delta that + starts a call. `TrackerError` converts into + `AiMuxError::InvalidResponseData`. **Rust (aimux-providers)** diff --git a/aimux-provider-utils/src/lib.rs b/aimux-provider-utils/src/lib.rs index f3b0245c0..48e0df249 100644 --- a/aimux-provider-utils/src/lib.rs +++ b/aimux-provider-utils/src/lib.rs @@ -50,7 +50,6 @@ pub use streaming_tool_call_argument_state::{ StreamingToolCallArgumentState, starts_with_structured_value, }; pub use streaming_tool_call_tracker::{ - StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, TrackerError, - TypeValidation, + StreamingToolCallDelta, StreamingToolCallTracker, TrackerError, TypeValidation, }; pub use url::{validate_base_url, without_trailing_slash, without_trailing_slash_opt}; diff --git a/aimux-provider-utils/src/streaming_tool_call_argument_state.rs b/aimux-provider-utils/src/streaming_tool_call_argument_state.rs index db156e66d..3de33b7f1 100644 --- a/aimux-provider-utils/src/streaming_tool_call_argument_state.rs +++ b/aimux-provider-utils/src/streaming_tool_call_argument_state.rs @@ -1,38 +1,37 @@ //! Incremental structure tracking for streamed tool-call arguments. //! -//! Rust translation of `@ai-sdk/provider-utils`'s -//! `StreamingToolCallArgumentState` -//! (`packages/provider-utils/src/streaming-tool-call-argument-state.ts`). -//! +//! Port of `@ai-sdk/provider-utils`'s `streaming-tool-call-argument-state.ts`. //! The state is intentionally *structural* rather than a JSON parse: a //! currently parsable scalar can still be the prefix of a later value, but a //! closed top-level `{…}` / `[…]` cannot be extended. -#[derive(Debug, Clone, PartialEq, Eq)] -enum ArgumentStructure { - Undetermined, - Other, - Structured { - stack: Vec, - in_string: bool, - escaped: bool, - complete: bool, - }, -} - /// Whether `value` (ignoring leading whitespace) starts with `{` or `[`. #[must_use] pub fn starts_with_structured_value(value: Option<&str>) -> bool { - value - .and_then(|v| v.trim_start().chars().next()) - .is_some_and(|c| c == '{' || c == '[') + 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 { - structure: ArgumentStructure, + structure: Structure, +} + +#[derive(Debug, Clone)] +enum Structure { + /// Only whitespace so far. + Undetermined, + /// Not a structured value (a scalar, or a mismatched bracket). + Other, + /// Inside a top-level `{…}` / `[…]`; `stack` holds the open brackets. + Open { + stack: Vec, + in_string: bool, + escaped: bool, + }, + /// The top-level value has been closed. + Complete, } impl StreamingToolCallArgumentState { @@ -40,7 +39,7 @@ impl StreamingToolCallArgumentState { #[must_use] pub fn new(initial_value: &str) -> Self { let mut state = Self { - structure: ArgumentStructure::Undetermined, + structure: Structure::Undetermined, }; state.append(initial_value); state @@ -49,67 +48,67 @@ impl StreamingToolCallArgumentState { /// `true` once a top-level `{…}` / `[…]` value has been closed. #[must_use] pub fn has_complete_structured_value(&self) -> bool { - matches!( - self.structure, - ArgumentStructure::Structured { complete: true, .. } - ) + matches!(self.structure, Structure::Complete) } /// Feed the next argument fragment. pub fn append(&mut self, delta: &str) { - for character in delta.chars() { - match &mut self.structure { - ArgumentStructure::Undetermined => { - if character.is_whitespace() { - continue; - } - self.structure = if character == '{' || character == '[' { - ArgumentStructure::Structured { - stack: vec![character], - in_string: false, - escaped: false, - complete: false, - } - } else { - ArgumentStructure::Other + let mut rest = delta; + if matches!(self.structure, Structure::Undetermined) { + rest = rest.trim_start(); + match rest.as_bytes().first() { + None => return, + Some(b'{' | b'[') => { + self.structure = Structure::Open { + stack: Vec::new(), + in_string: false, + escaped: false, }; } - ArgumentStructure::Other | ArgumentStructure::Structured { complete: true, .. } => { + Some(_) => { + self.structure = Structure::Other; + return; } - ArgumentStructure::Structured { - stack, - in_string, - escaped, - complete, - } => { - if *in_string { - if *escaped { - *escaped = false; - } else if character == '\\' { - *escaped = true; - } else if character == '"' { - *in_string = false; - } - continue; - } + } + } + + let Structure::Open { + stack, + in_string, + escaped, + } = &mut self.structure + else { + return; + }; - match character { - '"' => *in_string = true, - '{' | '[' => stack.push(character), - '}' | ']' => { - let expected = if character == '}' { '{' } else { '[' }; - if stack.last() != Some(&expected) { - self.structure = ArgumentStructure::Other; - continue; - } - stack.pop(); - if stack.is_empty() { - *complete = true; - } - } - _ => {} + // Every structural character is ASCII and UTF-8 continuation bytes + // never collide with ASCII, so scanning bytes is exact. + for &byte in rest.as_bytes() { + if *in_string { + if *escaped { + *escaped = false; + } else if byte == b'\\' { + *escaped = true; + } else if byte == b'"' { + *in_string = false; + } + continue; + } + match byte { + b'"' => *in_string = true, + b'{' | b'[' => stack.push(byte), + b'}' | b']' => { + let opening = if byte == b'}' { b'{' } else { b'[' }; + if stack.pop() != Some(opening) { + self.structure = Structure::Other; + return; + } + if stack.is_empty() { + self.structure = Structure::Complete; + return; } } + _ => {} } } } diff --git a/aimux-provider-utils/src/streaming_tool_call_tracker.rs b/aimux-provider-utils/src/streaming_tool_call_tracker.rs index ce9666dd8..cb9fedd43 100644 --- a/aimux-provider-utils/src/streaming_tool_call_tracker.rs +++ b/aimux-provider-utils/src/streaming_tool_call_tracker.rs @@ -1,30 +1,31 @@ -//! Streaming tool call tracker. +//! Streaming tool call tracker for the OpenAI chat-completions wire format. //! -//! Rust translation of `@ai-sdk/provider-utils`'s `StreamingToolCallTracker` -//! (`packages/provider-utils/src/streaming-tool-call-tracker.ts`). +//! 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). //! -//! Tracks streaming tool call state across the deltas of an OpenAI-compatible -//! chat completion stream: accumulates `arguments` fragments, emits -//! [`StreamPart::ToolInputStart`] / [`StreamPart::ToolInputDelta`] / -//! [`StreamPart::ToolInputEnd`] / [`StreamPart::ToolCall`] parts, and -//! finalizes unfinished calls on [`StreamingToolCallTracker::flush`]. +//! 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 correlated to calls by wire `id`, `index` and function name -//! (see [`StreamingToolCallTracker`]'s resolution table), 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 | +//! | --- | --- | --- | --- | +//! | known | matching | any | matching call, new call, or ambiguity | +//! | known | conflicting | named | new call | +//! | unseen | matching | structured start | new call | +//! | unseen | matching | continuation | matching call or ambiguity | +//! | absent | matching | any | matching call, new call, or ambiguity | +//! | absent | absent | named | new call | +//! | absent | absent | unnamed | sole call, new call, or ambiguity | //! -//! Like the TS original, a call is *never* finalized before `flush`: a -//! parsable argument buffer can still be the prefix of a longer argument -//! string, so acting on it early would use truncated inputs (ai-sdk #13137). +//! An ambiguous delta is dropped. //! -//! This is a tool for the OpenAI chat-completions wire format only. Protocols -//! whose streams carry explicit tool-call boundaries (Anthropic content -//! blocks, Google complete `functionCall` parts, Bedrock content blocks, -//! Cohere `tool-call-*` events, the Responses API's output items) do not need -//! it. - -use std::collections::{HashMap, HashSet}; +//! Protocols whose streams carry explicit tool-call boundaries (Anthropic, +//! Google, Bedrock, Cohere, the Responses API) do not need this. use serde_json::Value; use thiserror::Error; @@ -37,102 +38,48 @@ use crate::streaming_tool_call_argument_state::{ StreamingToolCallArgumentState, starts_with_structured_value, }; -/// Fallback id used when the id generator returns a blank string. +/// Id used when the id generator returns a blank string; also the output of +/// the default generator. const FALLBACK_TOOL_CALL_ID: &str = "tool-call"; -/// The `function` sub-object of a streaming tool call delta. -#[derive(Debug, Clone, Default)] -pub struct StreamingToolCallFunction { - pub name: Option, - pub arguments: Option, -} - -/// A streaming tool call delta — the `tool_calls[i]` entry of an OpenAI-style -/// streaming chunk. +/// One `tool_calls[i]` entry of a streaming chunk, borrowed from the +/// provider's own wire type. /// -/// `arguments: null` (TS) maps to `None`; `arguments: ''` maps to `Some("")`. +/// 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 { +pub struct StreamingToolCallDelta<'a> { pub index: Option, - pub id: Option, - /// The `type` field. Named `r#type` because `type` is a reserved word. - pub r#type: Option, - pub function: Option, - /// Provider-specific payload carried alongside the standard fields, read - /// by the `extract_metadata` hook (e.g. a Google thought signature). - pub extra: Value, -} - -impl StreamingToolCallDelta { - #[must_use] - pub fn new() -> Self { - Self::default() - } - - #[must_use] - pub fn index(mut self, index: usize) -> Self { - self.index = Some(index); - self - } - - #[must_use] - pub fn id(mut self, id: impl Into) -> Self { - self.id = Some(id.into()); - self - } - - /// Set the `type` field (named `tool_type` because `type` is reserved). - #[must_use] - pub fn tool_type(mut self, t: impl Into) -> Self { - self.r#type = Some(t.into()); - self - } - - #[must_use] - pub fn function_name(mut self, name: impl Into) -> Self { - self.function.get_or_insert_with(Default::default).name = Some(name.into()); - self - } - - /// Set the `function.arguments` fragment. Pass `""` for an explicit empty - /// fragment; omit the call for `None` (TS `arguments: null`). - #[must_use] - pub fn arguments(mut self, args: impl Into) -> Self { - self.function.get_or_insert_with(Default::default).arguments = Some(args.into()); - self - } - - #[must_use] - pub fn extra(mut self, extra: Value) -> Self { - self.extra = extra; - self - } + pub id: Option<&'a str>, + pub r#type: Option<&'a str>, + pub 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 a new tool call delta. +/// 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). + /// No validation. #[default] None, - /// Error if `type` is present and not `"function"`. + /// Reject a `type` that is present and not `"function"`. IfPresent, - /// Error if `type` is not exactly `"function"`. + /// Reject a `type` that is not exactly `"function"`. Required, } -/// Errors raised while processing a tool call delta. -/// -/// The TS tracker throws `InvalidResponseDataError`; these convert into -/// [`AiMuxError::InvalidResponseData`]. +/// 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, - #[error("Failed to create a unique tool call ID.")] - IdExhausted, } impl From for AiMuxError { @@ -141,56 +88,30 @@ impl From for AiMuxError { } } -struct TrackedToolCall { +struct ToolCall { id: String, - index: Option, - sequence: usize, - function_name: String, + name: String, arguments: String, argument_state: StreamingToolCallArgumentState, - has_finished: bool, - metadata: Option, + /// Index of the starting delta; orders [`StreamingToolCallTracker::finish`]. + index: Option, + /// Every wire id and index seen for this call. + wire_ids: Vec, + indices: Vec, + provider_metadata: Option, } -enum ToolCallResolution { +enum Resolution { Existing(usize), New, Ambiguous, } -type GenerateIdFn = Box String + Send + Sync>; -type ExtractMetadataFn = Box Option + Send + Sync>; -type BuildMetadataFn = Box) -> Option + Send + Sync>; - -/// Tracks streaming tool call state across multiple deltas. -/// -/// [`process_delta`](Self::process_delta) and [`flush`](Self::flush) return -/// the [`StreamPart`]s to forward downstream (the TS tracker enqueues them on -/// a controller instead). -/// -/// # Correlation -/// -/// | ID evidence | index/name evidence | start evidence | resolution | -/// | --- | --- | --- | --- | -/// | known | matching | any | matching call, new call, or ambiguity | -/// | known | conflicting | named | new call | -/// | unseen | matching | structured start | new call | -/// | unseen | matching | continuation | matching call or ambiguity | -/// | absent | matching | any | matching call, new call, or ambiguity | -/// | absent | absent | named | new call | -/// | absent | absent | unnamed | sole unfinished call, new call, or ambiguity | -/// -/// An *ambiguous* delta is dropped. +/// Tracks the tool calls of one streamed response. See the module docs. pub struct StreamingToolCallTracker { - tool_calls: Vec, - tool_calls_by_id: HashMap>, - tool_calls_by_index: HashMap>, - used_tool_call_ids: HashSet, - next_generated_id_suffixes: HashMap, - generate_id: GenerateIdFn, + calls: Vec, + generate_id: Box String + Send + Sync>, type_validation: TypeValidation, - extract_metadata: Option, - build_provider_metadata: Option, } impl Default for StreamingToolCallTracker { @@ -200,299 +121,226 @@ impl Default for StreamingToolCallTracker { } impl StreamingToolCallTracker { - /// Create a tracker with no metadata handling and default settings. - /// - /// The default id generator returns a blank-free constant, so ids that - /// need generating become `tool-call`, `tool-call-1`, `tool-call-2`, … + /// A tracker without `type` validation whose generated ids are + /// `tool-call`, `tool-call-1`, `tool-call-2`, … #[must_use] pub fn new() -> Self { Self { - tool_calls: Vec::new(), - tool_calls_by_id: HashMap::new(), - tool_calls_by_index: HashMap::new(), - used_tool_call_ids: HashSet::new(), - next_generated_id_suffixes: HashMap::new(), + calls: Vec::new(), generate_id: Box::new(|| FALLBACK_TOOL_CALL_ID.to_string()), type_validation: TypeValidation::None, - extract_metadata: None, - build_provider_metadata: None, } } - /// Set the id generator (the TS `generateId` option). Blank or repeated - /// outputs are turned into usable unique ids. + /// 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 String + Send + Sync + 'static>(mut self, f: F) -> Self { + pub fn with_generate_id(mut self, f: impl Fn() -> String + Send + Sync + 'static) -> Self { self.generate_id = Box::new(f); self } - /// Set the `type` validation mode (the TS `typeValidation` option). - #[must_use] - pub fn with_type_validation(mut self, v: TypeValidation) -> Self { - self.type_validation = v; - self - } - - /// Set the metadata extractor (the TS `extractMetadata` option). Called - /// once when a new tool call is detected; the value is kept with the call - /// and handed to the provider-metadata builder when the call finalizes. - #[must_use] - pub fn with_extract_metadata< - F: Fn(&StreamingToolCallDelta) -> Option + Send + Sync + 'static, - >( - mut self, - f: F, - ) -> Self { - self.extract_metadata = Some(Box::new(f)); - self - } - - /// Set the provider-metadata builder (the TS - /// `buildToolCallProviderMetadata` option). If it returns `None`, the - /// `ToolCall` part carries no `provider_metadata`. #[must_use] - pub fn with_build_provider_metadata< - F: Fn(Option<&Value>) -> Option + Send + Sync + 'static, - >( - mut self, - f: F, - ) -> Self { - self.build_provider_metadata = Some(Box::new(f)); + pub fn with_type_validation(mut self, validation: TypeValidation) -> Self { + self.type_validation = validation; self } - /// Process a tool call delta from a streaming chunk and return the parts - /// it produces (possibly none). + /// Process one delta, appending the parts it produces to `out`. /// /// # Errors /// - /// Returns [`TrackerError::InvalidType`] when `type` validation fails, - /// [`TrackerError::MissingFunctionName`] when a new call has no - /// (non-null) function name, and [`TrackerError::IdExhausted`] if no - /// unique id can be produced. - pub fn process_delta( + /// A delta that starts a call without a function name, or whose `type` + /// fails validation. + pub fn process( &mut self, - delta: &StreamingToolCallDelta, - ) -> Result, TrackerError> { - let wire_name = delta.function.as_ref().and_then(|f| f.name.as_deref()); - let has_blank_name = wire_name.is_some_and(|n| n.trim().is_empty()); - let wire_id = non_blank(delta.id.as_deref()); - let name = non_blank(wire_name); - let index = delta.index; - let arguments = delta.function.as_ref().and_then(|f| f.arguments.as_deref()); - - let resolution = self.resolve_tool_call( - wire_id, + delta: StreamingToolCallDelta<'_>, + out: &mut Vec, + ) -> Result<(), TrackerError> { + let StreamingToolCallDelta { index, - name, - name.is_some() && starts_with_structured_value(arguments), - ); - - let mut parts = Vec::new(); - let call = match resolution { - ToolCallResolution::Ambiguous => return Ok(parts), - ToolCallResolution::New => { - // Blank names cannot start a usable call, but some providers - // repeat a blank name on continuations. Those were correlated - // above; only an unmatched blank-name delta is ignored. - if has_blank_name { - return Ok(parts); - } - self.process_new_tool_call(delta, wire_id, index, name, &mut parts)? + id, + r#type, + name: wire_name, + arguments, + provider_metadata, + } = delta; + let wire_id = non_blank(id); + let name = non_blank(wire_name); + let explicit_start = name.is_some() && starts_with_structured_value(arguments); + + let call = match self.resolve(wire_id, index, name, explicit_start) { + Resolution::Ambiguous => return Ok(()), + // Blank names cannot start a usable call, but some providers + // repeat a blank name on continuations; those were correlated + // 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, + ) } - ToolCallResolution::Existing(call) => { - if let Some(wire_id) = wire_id { - self.associate_wire_id(call, wire_id); + Resolution::Existing(call) => { + let tool_call = &mut self.calls[call]; + if let Some(wire_id) = wire_id + && !tool_call.wire_ids.iter().any(|w| w == wire_id) + { + tool_call.wire_ids.push(wire_id.to_string()); + } + 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, + }); } - self.process_existing_tool_call(call, arguments, &mut parts); call } }; - if let Some(index) = index { - self.tool_calls_by_index - .entry(index) - .or_default() - .insert(call); + let indices = &mut self.calls[call].indices; + if let Some(index) = index + && !indices.contains(&index) + { + indices.push(index); } - Ok(parts) + Ok(()) } - /// Finalize any unfinished tool calls and return the closing parts. Call - /// once when the stream ends. - pub fn flush(&mut self) -> Vec { - // Index order is only reliable when every call has an index; for - // mixed streams keep insertion order. - let mut order: Vec = (0..self.tool_calls.len()).collect(); - if self.tool_calls.iter().all(|c| c.index.is_some()) { - order.sort_by_key(|&i| (self.tool_calls[i].index, self.tool_calls[i].sequence)); + /// 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) { + if self.calls.iter().all(|c| c.index.is_some()) { + self.calls.sort_by_key(|c| c.index); // stable: ties keep arrival order } - - let mut parts = Vec::new(); - for call in order { - if !self.tool_calls[call].has_finished { - self.finish_tool_call(call, &mut parts); - } + for call in self.calls { + out.push(StreamPart::ToolInputEnd { + id: call.id.clone(), + provider_metadata: None, + }); + out.push(StreamPart::ToolCall { + tool_call_id: call.id, + tool_name: call.name, + // Raw argument text; Core parses it. + input: Value::String(call.arguments), + provider_executed: None, + dynamic: None, + thought_signature: None, + invalid: None, + error: None, + provider_metadata: call.provider_metadata, + }); } - parts } - fn resolve_tool_call( + /// Apply the correlation table in the module docs. + fn resolve( &self, wire_id: Option<&str>, index: Option, name: Option<&str>, - has_explicit_call_start: bool, - ) -> ToolCallResolution { - let indexed = index.and_then(|i| self.tool_calls_by_index.get(&i)); - let matching_indexed = self.filter_by_name(indexed, name); - - if let Some(wire_id) = wire_id { - if let Some(with_id) = self.tool_calls_by_id.get(wire_id) { - if index.is_some() { - let matching: Vec = matching_indexed - .iter() - .copied() - .filter(|c| with_id.contains(c)) - .collect(); - let resolved = self.resolve_matching(matching, has_explicit_call_start); - if !matches!(resolved, ToolCallResolution::New) { - return resolved; - } - - // A named delta with a distinct index starts a new call - // even when its wire id and name repeat: providers may - // reuse ids across parallel calls. - if name.is_some() { - return ToolCallResolution::New; - } - - // Conflicting labels on a continuation cannot be - // resolved safely. - if indexed.is_some() { - return ToolCallResolution::Ambiguous; - } - - return self.resolve_matching(sorted(with_id), false); - } - - if let Some(name) = name { - let matching: Vec = sorted(with_id) - .into_iter() - .filter(|&c| self.tool_calls[c].function_name == name) - .collect(); - return self.resolve_matching(matching, has_explicit_call_start); - } - - return self.resolve_matching(sorted(with_id), false); + explicit_start: bool, + ) -> Resolution { + let at_index = |c: &ToolCall| index.is_some_and(|i| c.indices.contains(&i)); + let named = |c: &ToolCall| name.is_none_or(|n| c.name == n); + let index_known = self.calls.iter().any(at_index); + + let Some(wire_id) = wire_id else { + if index_known { + // Repeated names are valid on continuations; a different name + // at the same index is a new call from a provider that reuses + // indices across parallel calls. + return self.matching(|c| at_index(c) && named(c), explicit_start); } - - if !matching_indexed.is_empty() { - // A previously unseen id plus a named structured argument - // start is stronger evidence of a distinct call than a reused - // index/name; ids may still change on plain continuations. - return if has_explicit_call_start { - ToolCallResolution::New - } else { - self.resolve_matching(matching_indexed, false) - }; + if name.is_some() { + return Resolution::New; } + return self.matching(|_| true, false); + }; - return ToolCallResolution::New; + let has_id = |c: &ToolCall| c.wire_ids.iter().any(|w| w == wire_id); + if !self.calls.iter().any(has_id) { + // A previously unseen id plus a named structured start is + // stronger evidence of a distinct call than a reused index/name; + // ids may still change on plain continuations. + if explicit_start { + return Resolution::New; + } + return self.matching(|c| at_index(c) && named(c), false); } - if indexed.is_some() { - // Repeated names are valid on continuations; a different name at - // the same index means a new call from a provider that reuses - // indices across parallel calls. - return self.resolve_matching(matching_indexed, has_explicit_call_start); + if index.is_none() { + return match name { + Some(name) => self.matching(|c| has_id(c) && c.name == name, explicit_start), + None => self.matching(has_id, false), + }; } + let resolution = self.matching(|c| has_id(c) && at_index(c) && named(c), explicit_start); + if !matches!(resolution, Resolution::New) { + return resolution; + } if name.is_some() { - return ToolCallResolution::New; + // A named delta at a distinct index starts a new call even when + // its id and name repeat: providers may reuse ids across + // parallel calls. + return Resolution::New; } - - let unfinished: Vec = (0..self.tool_calls.len()) - .filter(|&c| !self.tool_calls[c].has_finished) - .collect(); - match unfinished.len() { - 0 => ToolCallResolution::New, - 1 => ToolCallResolution::Existing(unfinished[0]), - _ => ToolCallResolution::Ambiguous, + if index_known { + // Conflicting labels on a continuation cannot be resolved safely. + return Resolution::Ambiguous; } + self.matching(has_id, false) } - fn filter_by_name(&self, calls: Option<&HashSet>, name: Option<&str>) -> Vec { - calls.map_or_else(Vec::new, |calls| { - sorted(calls) - .into_iter() - .filter(|&c| name.is_none_or(|n| self.tool_calls[c].function_name == n)) - .collect() - }) - } - - fn resolve_matching( - &self, - calls: Vec, - has_explicit_call_start: bool, - ) -> ToolCallResolution { - if calls.is_empty() { - return ToolCallResolution::New; - } - - if !has_explicit_call_start { - return if calls.len() == 1 { - ToolCallResolution::Existing(calls[0]) - } else { - ToolCallResolution::Ambiguous - }; + /// Resolve to the single call matching `predicate`. With an explicit + /// structured start, a call whose structured arguments are already + /// complete cannot continue, so it is not a candidate. + fn matching(&self, predicate: impl Fn(&ToolCall) -> bool, explicit_start: bool) -> Resolution { + let mut candidates = self.calls.iter().enumerate().filter(|(_, c)| { + predicate(c) && !(explicit_start && c.argument_state.has_complete_structured_value()) + }); + match (candidates.next(), candidates.next()) { + (None, _) => Resolution::New, + (Some((call, _)), None) => Resolution::Existing(call), + (Some(_), Some(_)) => Resolution::Ambiguous, } + } - // A repeated name can occur on continuations. A fresh structured - // argument prefix signals another call only once the matching call - // has completed its own structured payload. - let continuable: Vec = calls - .into_iter() - .filter(|&c| { - !self.tool_calls[c] - .argument_state - .has_complete_structured_value() - }) - .collect(); - match continuable.len() { - 0 => ToolCallResolution::New, - 1 => ToolCallResolution::Existing(continuable[0]), - _ => ToolCallResolution::Ambiguous, + 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) } } - fn process_new_tool_call( + fn start( &mut self, - delta: &StreamingToolCallDelta, wire_id: Option<&str>, index: Option, - name: Option<&str>, - parts: &mut Vec, - ) -> Result { - match self.type_validation { - TypeValidation::Required => { - if delta.r#type.as_deref() != Some("function") { - return Err(TrackerError::InvalidType); - } - } - TypeValidation::IfPresent => { - if delta.r#type.as_deref().is_some_and(|t| t != "function") { - return Err(TrackerError::InvalidType); - } - } - TypeValidation::None => {} - } - - let name = name.ok_or(TrackerError::MissingFunctionName)?; - let id = self.create_tool_call_id(wire_id)?; - - parts.push(StreamPart::ToolInputStart { + name: &str, + arguments: &str, + provider_metadata: Option, + out: &mut Vec, + ) -> usize { + let id = self.unique_id(wire_id); + out.push(StreamPart::ToolInputStart { id: id.clone(), tool_name: name.to_string(), provider_executed: None, @@ -500,150 +348,49 @@ impl StreamingToolCallTracker { title: None, provider_metadata: None, }); - - let metadata = self - .extract_metadata - .as_ref() - .and_then(|extract| extract(delta)); - - let initial_arguments = delta - .function - .as_ref() - .and_then(|f| f.arguments.clone()) - .unwrap_or_default(); - - let call = self.tool_calls.len(); - self.tool_calls.push(TrackedToolCall { - id: id.clone(), - index, - sequence: call, - function_name: name.to_string(), - argument_state: StreamingToolCallArgumentState::new(&initial_arguments), - arguments: initial_arguments.clone(), - has_finished: false, - metadata, - }); - if let Some(wire_id) = wire_id { - self.associate_wire_id(call, wire_id); - } - - if !initial_arguments.is_empty() { - parts.push(StreamPart::ToolInputDelta { - id, - delta: initial_arguments, - provider_metadata: None, - }); - } - - // Tool calls must not finalize before the stream ends (ai-sdk - // #13137); finalization happens in `flush`. - Ok(call) - } - - fn process_existing_tool_call( - &mut self, - call: usize, - arguments: Option<&str>, - parts: &mut Vec, - ) { - let tool_call = &mut self.tool_calls[call]; - if tool_call.has_finished { - return; - } - if let Some(arguments) = arguments { - tool_call.argument_state.append(arguments); - tool_call.arguments.push_str(arguments); - parts.push(StreamPart::ToolInputDelta { - id: tool_call.id.clone(), + if !arguments.is_empty() { + out.push(StreamPart::ToolInputDelta { + id: id.clone(), delta: arguments.to_string(), provider_metadata: None, }); } + self.calls.push(ToolCall { + id, + name: name.to_string(), + arguments: arguments.to_string(), + argument_state: StreamingToolCallArgumentState::new(arguments), + index, + wire_ids: wire_id.map(str::to_string).into_iter().collect(), + indices: Vec::new(), + provider_metadata, + }); + self.calls.len() - 1 } - fn associate_wire_id(&mut self, call: usize, wire_id: &str) { - self.tool_calls_by_id - .entry(wire_id.to_string()) - .or_default() - .insert(call); - } - - fn create_tool_call_id(&mut self, wire_id: Option<&str>) -> Result { + /// The wire id if unused, else the generated id, else the generated id + /// with the first free `-N` suffix. + fn unique_id(&self, wire_id: Option<&str>) -> String { + let used = |id: &str| self.calls.iter().any(|c| c.id == id); if let Some(wire_id) = wire_id - && !self.used_tool_call_ids.contains(wire_id) + && !used(wire_id) { - self.used_tool_call_ids.insert(wire_id.to_string()); - return Ok(wire_id.to_string()); - } - - let generated = non_blank(Some(&(self.generate_id)())) - .unwrap_or(FALLBACK_TOOL_CALL_ID) - .to_string(); - - if !self.used_tool_call_ids.contains(&generated) { - self.used_tool_call_ids.insert(generated.clone()); - return Ok(generated); + return wire_id.to_string(); } - - // Resume after the last suffix checked for this generated value so - // deterministic generators stay bounded without rescanning occupied - // suffixes. - let initial_suffix = self - .next_generated_id_suffixes - .get(&generated) - .copied() - .unwrap_or(1); - let maximum_suffix = initial_suffix + self.used_tool_call_ids.len(); - for suffix in initial_suffix..=maximum_suffix { - let suffixed = format!("{generated}-{suffix}"); - if !self.used_tool_call_ids.contains(&suffixed) { - self.used_tool_call_ids.insert(suffixed.clone()); - self.next_generated_id_suffixes - .insert(generated, suffix + 1); - return Ok(suffixed); - } + let generated = (self.generate_id)(); + let generated = non_blank(Some(&generated)).unwrap_or(FALLBACK_TOOL_CALL_ID); + if !used(generated) { + return generated.to_string(); } - - // Unreachable by the pigeonhole principle; guards the invariant - // instead of looping without bound. - Err(TrackerError::IdExhausted) - } - - fn finish_tool_call(&mut self, call: usize, parts: &mut Vec) { - let tool_call = &mut self.tool_calls[call]; - tool_call.has_finished = true; - - parts.push(StreamPart::ToolInputEnd { - id: tool_call.id.clone(), - provider_metadata: None, - }); - - let provider_metadata = self - .build_provider_metadata - .as_ref() - .and_then(|build| build(tool_call.metadata.as_ref())); - - parts.push(StreamPart::ToolCall { - tool_call_id: tool_call.id.clone(), - tool_name: tool_call.function_name.clone(), - // Raw argument text; Core parses it after `stream_text`. - input: Value::String(tool_call.arguments.clone()), - provider_executed: None, - dynamic: None, - thought_signature: None, - invalid: None, - error: None, - provider_metadata, - }); + // Ids are never released, so the first free suffix equals what the + // upstream resume-from-last-suffix search returns. + (1..) + .map(|suffix| format!("{generated}-{suffix}")) + .find(|id| !used(id)) + .expect("finitely many ids are in use") } } fn non_blank(value: Option<&str>) -> Option<&str> { value.filter(|v| !v.trim().is_empty()) } - -fn sorted(calls: &HashSet) -> Vec { - let mut calls: Vec = calls.iter().copied().collect(); - calls.sort_unstable(); - calls -} diff --git a/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs b/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs index 238a0ecb9..82b6b2726 100644 --- a/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs +++ b/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs @@ -9,8 +9,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use aimux_core::stream_part::StreamPart; use aimux_provider_utils::{ - StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, TrackerError, - TypeValidation, + StreamingToolCallDelta, StreamingToolCallTracker, TrackerError, TypeValidation, }; use serde_json::{Value, json}; @@ -84,7 +83,8 @@ fn project(part: StreamPart) -> Part { /// Tracker plus the parts it has emitted so far (the TS `createCollector`). struct Harness { - tracker: StreamingToolCallTracker, + /// `None` once flushed: `finish` consumes the tracker. + tracker: Option, parts: Vec, } @@ -95,19 +95,25 @@ impl Harness { fn with(tracker: StreamingToolCallTracker) -> Self { Self { - tracker, + tracker: Some(tracker), parts: Vec::new(), } } - fn delta(&mut self, delta: StreamingToolCallDelta) -> Result<(), TrackerError> { - let parts = self.tracker.process_delta(&delta)?; + 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)); - Ok(()) + result } fn flush(&mut self) { - let parts = self.tracker.flush(); + let mut parts = Vec::new(); + self.tracker + .take() + .expect("flushed twice") + .finish(&mut parts); self.parts.extend(parts.into_iter().map(project)); } @@ -133,27 +139,30 @@ impl Harness { } /// A delta with an explicit `function` object (as every TS test passes one). -fn delta( +fn delta<'a>( index: Option, - id: Option<&str>, - ty: Option<&str>, - name: Option<&str>, - arguments: Option<&str>, -) -> StreamingToolCallDelta { + id: Option<&'a str>, + ty: Option<&'a str>, + name: Option<&'a str>, + arguments: Option<&'a str>, +) -> StreamingToolCallDelta<'a> { StreamingToolCallDelta { index, - id: id.map(str::to_string), - r#type: ty.map(str::to_string), - function: Some(StreamingToolCallFunction { - name: name.map(str::to_string), - arguments: arguments.map(str::to_string), - }), - extra: Value::Null, + id, + r#type: ty, + name, + arguments, + provider_metadata: None, } } /// Full call start: index, id, `type: 'function'`, name and arguments. -fn start(index: usize, id: &str, name: &str, arguments: &str) -> StreamingToolCallDelta { +fn start<'a>( + index: usize, + id: &'a str, + name: &'a str, + arguments: &'a str, +) -> StreamingToolCallDelta<'a> { delta( Some(index), Some(id), @@ -164,7 +173,7 @@ fn start(index: usize, id: &str, name: &str, arguments: &str) -> StreamingToolCa } /// Continuation carrying only arguments (plus an optional index). -fn cont(index: Option, arguments: &str) -> StreamingToolCallDelta { +fn cont(index: Option, arguments: &str) -> StreamingToolCallDelta<'_> { delta(index, None, None, None, Some(arguments)) } @@ -394,18 +403,6 @@ mod process_delta { assert_eq!(h.tool_calls(), vec![tc("call_1", "fn", "{\"value\":1}")]); } - #[test] - fn skips_deltas_for_already_finished_tool_calls() { - let mut h = Harness::new(); - - h.delta(start(0, "call_1", "fn", "{}")).unwrap(); - h.flush(); - h.clear(); - - h.delta(cont(Some(0), "extra")).unwrap(); - assert_eq!(h.parts, vec![]); - } - #[test] fn skips_delta_emission_when_arguments_are_null() { let mut h = Harness::new(); @@ -964,7 +961,7 @@ mod process_delta { mod type_validation { use super::*; - fn custom_type() -> StreamingToolCallDelta { + fn custom_type() -> StreamingToolCallDelta<'static> { delta( Some(0), Some("call_1"), @@ -974,7 +971,7 @@ mod type_validation { ) } - fn no_type() -> StreamingToolCallDelta { + fn no_type() -> StreamingToolCallDelta<'static> { delta(Some(0), Some("call_1"), None, Some("fn"), Some("")) } @@ -1036,100 +1033,49 @@ mod flush { ] ); } - - #[test] - fn does_not_re_finalize_already_finished_tool_calls() { - let mut h = Harness::new(); - - h.delta(start(0, "call_1", "fn", "{}")).unwrap(); - h.flush(); - h.clear(); - - h.flush(); - - assert_eq!(h.parts, vec![]); - } } mod metadata { use super::*; - fn google_tracker() -> StreamingToolCallTracker { - StreamingToolCallTracker::new() - .with_extract_metadata(|delta| { - delta.extra["extra_content"]["google"]["thought_signature"] - .as_str() - .map(|sig| json!({ "thoughtSignature": sig })) - }) - .with_build_provider_metadata(|metadata| { - metadata - .and_then(|m| m.get("thoughtSignature")) - .map(|sig| json!({ "google": { "thoughtSignature": sig } })) - }) - } + // The upstream `extractMetadata` / `buildToolCallProviderMetadata` hooks + // become `provider_metadata` on the delta that starts a call. #[test] - fn extracts_and_includes_provider_metadata_in_tool_call_parts() { - let mut h = Harness::with(google_tracker()); + fn includes_the_starting_delta_metadata_in_the_tool_call() { + let mut h = Harness::new(); - h.delta( - start(0, "call_1", "fn", "{}") - .extra(json!({ "extra_content": { "google": { "thought_signature": "sig123" } } })), - ) + h.delta(StreamingToolCallDelta { + provider_metadata: Some(json!({ "google": { "thoughtSignature": "sig123" } })), + ..start(0, "call_1", "fn", "{\"incomplete") + }) .unwrap(); h.flush(); - let tool_call = h.parts.iter().find(|p| matches!(p, Part::ToolCall { .. })); - assert_eq!( - tool_call, - Some(&Part::ToolCall { - tool_call_id: "call_1".into(), - tool_name: "fn".into(), - input: "{}".into(), - provider_metadata: Some(json!({ "google": { "thoughtSignature": "sig123" } })), - }) - ); - } - - #[test] - fn includes_provider_metadata_for_unfinished_tool_calls_finalized_in_flush() { - let mut h = Harness::with( - StreamingToolCallTracker::new() - .with_extract_metadata(|_| Some(json!({ "custom": { "key": "value" } }))) - .with_build_provider_metadata(|metadata| { - metadata.map(|m| json!({ "provider": m })) - }), - ); - - h.delta(start(0, "call_1", "fn", "{\"incomplete")).unwrap(); - h.clear(); - - h.flush(); - assert_eq!( h.parts.last(), Some(&Part::ToolCall { tool_call_id: "call_1".into(), tool_name: "fn".into(), input: "{\"incomplete".into(), - provider_metadata: Some(json!({ "provider": { "custom": { "key": "value" } } })), + provider_metadata: Some(json!({ "google": { "thoughtSignature": "sig123" } })), }) ); } #[test] - fn omits_provider_metadata_when_the_builder_returns_none() { - let mut h = Harness::with( - StreamingToolCallTracker::new() - .with_extract_metadata(|_| None) - .with_build_provider_metadata(|_| None), - ); + fn ignores_metadata_on_continuations() { + let mut h = Harness::new(); - h.delta(start(0, "call_1", "fn", "{}")).unwrap(); + h.delta(start(0, "call_1", "fn", "{")).unwrap(); + h.delta(StreamingToolCallDelta { + provider_metadata: Some(json!({ "custom": { "key": "value" } })), + ..cont(Some(0), "}") + }) + .unwrap(); h.flush(); - let tool_call = h.parts.iter().find(|p| matches!(p, Part::ToolCall { .. })); - assert_eq!(tool_call, Some(&super::tool_call("call_1", "fn", "{}"))); + assert_eq!(h.parts.last(), Some(&tool_call("call_1", "fn", "{}"))); } } diff --git a/aimux-providers/src/openai/model.rs b/aimux-providers/src/openai/model.rs index 107f984ee..b4267c866 100644 --- a/aimux-providers/src/openai/model.rs +++ b/aimux-providers/src/openai/model.rs @@ -20,9 +20,7 @@ use aimux_core::result::{GenerateContent, GenerateResult, StreamResult}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; -use aimux_provider_utils::{ - HttpRequest, StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, -}; +use aimux_provider_utils::{HttpRequest, StreamingToolCallDelta, StreamingToolCallTracker}; use super::OpenAIConfig; use super::convert::{RequestBodyResult, build_request_body_with_warnings, parse_finish_reason}; @@ -487,6 +485,7 @@ pub async fn execute_stream( // Streamed tool calls, correlated by wire id, index and function name // (the AI SDK's StreamingToolCallTracker) and finalized on flush. let mut tool_calls = StreamingToolCallTracker::new(); + let mut tool_parts = Vec::new(); // Process the first event (already peeked) then the rest. let mut event_iter = @@ -635,30 +634,24 @@ pub async fn execute_stream( }); reasoning_started = false; } - for dtc in tool_call_deltas { + for dtc in &tool_call_deltas { + let function = dtc.function.as_ref(); let delta = StreamingToolCallDelta { index: dtc.index, - id: dtc.id, - r#type: None, - function: dtc.function.map(|f| StreamingToolCallFunction { - name: f.name, - arguments: f.arguments, - }), - extra: Value::Null, + id: dtc.id.as_deref(), + name: function.and_then(|f| f.name.as_deref()), + arguments: function.and_then(|f| f.arguments.as_deref()), + ..Default::default() }; - match tool_calls.process_delta(&delta) { - Ok(parts) => { - for part in parts { - yield Ok(part); - } - } - // A malformed delta (new call without a - // function name) is invalid response - // data, as in the AI SDK; the stream ends. - Err(error) => { - yield Err(error.into()); - return; - } + // A malformed delta (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); } } } @@ -739,7 +732,8 @@ pub async fn execute_stream( // 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. - for part in tool_calls.flush() { + tool_calls.finish(&mut tool_parts); + for part in tool_parts.drain(..) { yield Ok(part); } diff --git a/aimux-stream/src/lib.rs b/aimux-stream/src/lib.rs index 65c368daf..ae249daf6 100644 --- a/aimux-stream/src/lib.rs +++ b/aimux-stream/src/lib.rs @@ -1,8 +1,7 @@ //! # aimux-stream //! -//! Low-level streaming primitives for SSE (Server-Sent Events) parsing and -//! streamed tool-call tracking, -//! used by provider implementations to decode model API response streams. +//! Low-level SSE (Server-Sent Events) parsing, used by provider +//! implementations to decode model API response streams. pub mod lines; pub mod sse; diff --git a/bindings/node/src/types/InputRecord.ts b/bindings/node/src/types/InputRecord.ts index c6e44df47..e187e6e3d 100644 --- a/bindings/node/src/types/InputRecord.ts +++ b/bindings/node/src/types/InputRecord.ts @@ -6,11 +6,11 @@ import type { JsonValue } from "./serde_json/JsonValue"; /** * ① 输入侧:完整调用参数,足以重建 generate_text 调用。 */ -export type InputRecord = { operation?: RecordingOperation, +export type InputRecord = { operation?: RecordingOperation, /** * 完整 prompt(消息数组,含 ContentPart::Image 等多模态)。 */ -prompt: Array, +prompt: Array, /** * 序列化的 CallOptions(abort_signal/call_id 已 serde skip); * headers/provider_options/body_overrides 已递归脱敏。 diff --git a/bindings/node/src/types/OutcomeRecord.ts b/bindings/node/src/types/OutcomeRecord.ts index 38777b88a..fb13110d6 100644 --- a/bindings/node/src/types/OutcomeRecord.ts +++ b/bindings/node/src/types/OutcomeRecord.ts @@ -5,16 +5,16 @@ import type { JsonValue } from "./serde_json/JsonValue"; /** * 最终结果摘要。 */ -export type OutcomeRecord = { +export type OutcomeRecord = { /** * Normalized decision result, independent of a provider's wire format. */ -decision_result?: JsonValue | null, status: OutcomeStatus, finish_reason: string | null, error: string | null, +decision_result?: JsonValue | null, status: OutcomeStatus, finish_reason: string | null, error: string | null, /** * Lossless structured domain error. In particular, `RetryError.errors` * keeps the complete attempt history rather than only its display text. */ -error_value: JsonValue | null, +error_value: JsonValue | null, /** * 序列化的 Usage。 */ From abafcd984eaf2c0283c69e25c2a204138d59b7b1 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 11 Oct 2026 16:37:43 +0000 Subject: [PATCH 7/8] chore(stream): drop the tool-call tracker from the crate's dependencies and docs With the tracker in aimux-provider-utils, aimux-stream only decodes SSE: serde_json is needed by its tests alone, and the tool-calls keyword no longer applies. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01A3WxDBWBA4ituJjXsHaEAa --- aimux-stream/Cargo.toml | 4 ++-- aimux-stream/src/lib.rs | 6 ++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/aimux-stream/Cargo.toml b/aimux-stream/Cargo.toml index 85b6bb205..4f167f016 100644 --- a/aimux-stream/Cargo.toml +++ b/aimux-stream/Cargo.toml @@ -4,7 +4,7 @@ version.workspace = true edition.workspace = true license.workspace = true description = "Streaming primitives for aimux" -keywords = ["llm", "streaming", "sse", "tool-calls"] +keywords = ["llm", "streaming", "sse"] categories = ["asynchronous", "parser-implementations"] documentation = "https://docs.rs/aimux-stream" @@ -12,11 +12,11 @@ documentation = "https://docs.rs/aimux-stream" futures = { workspace = true } pin-project-lite = { workspace = true } bytes = { workspace = true } -serde_json = { workspace = true } thiserror = { workspace = true } sse-stream = { version = "0.3.0", default-features = false, features = ["memchr"] } [dev-dependencies] +serde_json = { workspace = true } sha2 = "0.10" tokio = { workspace = true } diff --git a/aimux-stream/src/lib.rs b/aimux-stream/src/lib.rs index ae249daf6..3edc0c1ab 100644 --- a/aimux-stream/src/lib.rs +++ b/aimux-stream/src/lib.rs @@ -1,7 +1,9 @@ //! # aimux-stream //! -//! Low-level SSE (Server-Sent Events) parsing, used by provider -//! implementations to decode model API response streams. +//! Low-level SSE (Server-Sent Events) decoding, used by +//! `aimux-provider-utils` to decode model API response streams. The crate +//! plays the role `eventsource-parser` plays for the AI SDK: it knows nothing +//! about providers, models or stream parts. pub mod lines; pub mod sse; From d02a3a8059629aa15b1f909f0fe1b228eadea2c2 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 12 Oct 2026 09:32:42 +0000 Subject: [PATCH 8/8] refactor(core)!: tool types follow the AI SDK V4 tool types One tool struct per layer, field for field with the pinned upstream types: - CallOptions.tool_choice is optional. Vendors omit an unset choice, as their upstream packages do; generate_text and stream_text default it to auto when tools are given, as upstream prepare-tool-choice does. - FunctionTool.input_examples is a list of { input: object } and ProviderTool.args is a JSON object. - Tool calls carry no separate thought signature. The value travels in provider metadata on output and provider options on input, under the vendor namespace, where each upstream vendor package reads and writes it. - Call sites pass the tool structs instead of their field lists. The fifteen tool-choice tests ported from upstream (OpenAI, Groq, Mistral, Bedrock) are kept and adapted; hand-written tests that pinned the old shapes are deleted. CHANGELOG and docs/api/gaps.md record the wire-format changes. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01A3WxDBWBA4ituJjXsHaEAa --- CHANGELOG.md | 15 + aimux-core/src/content.rs | 6 - aimux-core/src/generate.rs | 147 +-- aimux-core/src/lib.rs | 12 +- aimux-core/src/moa.rs | 702 +------------ aimux-core/src/openai_output.rs | 954 +----------------- aimux-core/src/options.rs | 6 +- aimux-core/src/parse_tool_call.rs | 62 +- aimux-core/src/recording.rs | 6 +- aimux-core/src/replay.rs | 79 +- aimux-core/src/response_messages.rs | 29 +- aimux-core/src/result.rs | 112 +- aimux-core/src/stream_part.rs | 94 +- aimux-core/src/tool.rs | 51 +- aimux-core/src/trace/layer.rs | 2 +- aimux-core/tests/tool_input_test.rs | 25 +- aimux-ffi/src/lib.rs | 4 +- .../src/streaming_tool_call_tracker.rs | 11 +- .../tests/streaming_tool_call_tracker_test.rs | 13 +- aimux-providers/src/anthropic/convert.rs | 2 +- .../src/anthropic/prepare_tools.rs | 13 +- aimux-providers/src/anthropic/stream.rs | 82 +- .../src/anthropic/tool_name_mapping.rs | 91 +- aimux-providers/src/bedrock/convert.rs | 19 +- aimux-providers/src/bedrock/model.rs | 17 +- aimux-providers/src/cohere/convert.rs | 2 +- aimux-providers/src/cohere/model.rs | 15 +- aimux-providers/src/google/convert.rs | 43 +- aimux-providers/src/google/model.rs | 86 +- aimux-providers/src/huggingface/responses.rs | 50 +- aimux-providers/src/mistral/convert.rs | 2 +- aimux-providers/src/mistral/model.rs | 15 +- aimux-providers/src/open_responses.rs | 28 +- aimux-providers/src/openai/convert.rs | 4 +- aimux-providers/src/openai/model.rs | 6 +- .../src/openai/responses/convert.rs | 2 +- .../src/openai/responses/responses_convert.rs | 32 +- aimux-providers/src/vertex/anthropic_model.rs | 22 +- aimux-providers/src/vertex/model.rs | 60 +- aimux-providers/src/xai/convert.rs | 2 +- aimux-providers/src/xai/model.rs | 24 +- aimux-providers/src/xai/responses/convert.rs | 2 +- aimux-providers/src/xai/responses/mod.rs | 56 +- aimux-providers/tests/alibaba_test.rs | 8 +- aimux-providers/tests/bedrock_convert_test.rs | 281 +----- aimux-providers/tests/codex_test.rs | 69 +- aimux-providers/tests/conformance_test.rs | 6 +- .../tests/data_loss_regression_test.rs | 28 +- aimux-providers/tests/deepseek_chat_test.rs | 19 +- aimux-providers/tests/e2e_test.rs | 6 +- aimux-providers/tests/groq_test.rs | 196 +--- .../tests/huggingface_responses_test.rs | 161 +-- aimux-providers/tests/open_responses_test.rs | 225 +---- .../tests/openai_compatible_test.rs | 17 +- aimux-providers/tests/openrouter_test.rs | 45 +- bindings/flutter/lib/types.dart | 34 +- bindings/flutter/lib/types.g.dart | 6 +- .../test/tool_call_repair_typed_test.dart | 2 - bindings/go/repair.go | 1 - bindings/go/types.go | 7 +- .../main/java/ai/arcships/aimux/Types.java | 116 +-- .../ai/arcships/aimux/ToolCallRepair.kt | 1 - .../main/kotlin/ai/arcships/aimux/Types.kt | 15 +- bindings/node/src/lib.rs | 40 +- bindings/node/src/multimodal.rs | 164 ++- bindings/node/src/types.ts | 1 + bindings/node/src/types/CallOptions.ts | 2 +- bindings/node/src/types/ContentPart.ts | 6 - bindings/node/src/types/FunctionTool.ts | 3 +- .../src/types/FunctionToolInputExample.ts | 7 + bindings/node/src/types/GenerateContent.ts | 51 +- bindings/node/src/types/ProviderTool.ts | 2 +- bindings/node/src/types/RawToolCall.ts | 2 +- bindings/node/src/types/StreamPart.ts | 51 +- bindings/node/src/types/ToolCall.ts | 7 - bindings/node/src/types/ToolResult.ts | 17 +- bindings/python/python/aimux/wrapper.py | 14 +- bindings/python/src/lib.rs | 15 +- bindings/python/src/multimodal.rs | 152 ++- .../swift/Sources/Aimux/ToolCallRepair.swift | 6 +- bindings/swift/Sources/Aimux/Types.swift | 49 +- .../swift/Tests/AimuxTests/WrapperTests.swift | 4 +- docs/api/gaps.md | 15 +- tools/aimux-web/src/wire.rs | 1 - 84 files changed, 921 insertions(+), 3936 deletions(-) create mode 100644 bindings/node/src/types/FunctionToolInputExample.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index 87669b4f4..8d57c7edd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Breaking +- Tool-call and tool-result types are unified, one per protocol layer + (`aimux-core`). A provider emits `tool::RawToolCall` (`input` is the raw + argument text) in `GenerateContent::ToolCall(..)` and + `StreamPart::ToolCall(..)`, and `tool::ToolResult` in the `ToolResult` + variants; these variants are newtype variants now. `StreamPart` is generic + over the tool-call type: a provider's `do_stream` yields `StreamPart` + (= `StreamPart`), `stream_text` and everything after it yield + `TextStreamPart` (= `StreamPart`, parsed input with `invalid` / + `error`). `RawToolCall` moved from `parse_tool_call` to `tool`. Wire format: + enum variants keep their JSON shape; a provider-layer `StreamPart::ToolCall` + has a string `input` and no `invalid` / `error`; a + `GenerateContent::ToolCall` with a non-string `input` no longer + deserializes; the standalone `tool::ToolResult` gains a required + `tool_name` and optional `dynamic` / `provider_metadata`. + - Removed the generated `ProviderName` type in every binding (Rust enum, TS const object, Go/Java/Kotlin consts, Swift enum, Dart consts, Python `Literal`) and `scripts/gen_provider_names.py`. Provider names are plain diff --git a/aimux-core/src/content.rs b/aimux-core/src/content.rs index e6dfe464a..8e97e5948 100644 --- a/aimux-core/src/content.rs +++ b/aimux-core/src/content.rs @@ -104,11 +104,6 @@ pub enum ContentPart { /// need it to replay server tool calls on a later turn. #[serde(default, skip_serializing_if = "Option::is_none")] provider_executed: Option, - /// Provider-assigned thought signature (e.g. Google Gemini - /// `thoughtSignature`). Must be echoed back verbatim on the follow-up - /// turn when the tool result is sent. - #[serde(default, skip_serializing_if = "Option::is_none")] - thought_signature: Option, /// Provider-specific options for this part (e.g. /// `anthropic.cacheControl`). #[serde(default, skip_serializing_if = "Option::is_none")] @@ -167,7 +162,6 @@ impl ContentPart { tool_name: tool_name.into(), input, provider_executed: None, - thought_signature: None, provider_options: None, } } diff --git a/aimux-core/src/generate.rs b/aimux-core/src/generate.rs index 274939bb7..08b57761c 100644 --- a/aimux-core/src/generate.rs +++ b/aimux-core/src/generate.rs @@ -21,24 +21,24 @@ use crate::language_model::LanguageModel; use crate::language_model_message::convert_to_language_model_prompt; use crate::message::{ModelMessage, ModelPrompt}; use crate::options::{CallOptions, ResponseFormat, ToolChoice}; -use crate::parse_tool_call::{RawToolCall, ToolCallRepair, parse_tool_call}; +use crate::parse_tool_call::{ToolCallRepair, parse_tool_call}; use crate::result::{ FilePart, GenerateContent, GenerateResult, ReasoningPart, SourcePart, StreamResult, StreamTextResultAggregated, }; -use crate::stream_part::StreamPart; +use crate::stream_part::{StreamPart, TextStreamPart}; use crate::tool::Tool; use crate::types::{FinishReason, ReasoningEffort, Usage, Warning}; use crate::{AbortSignal, retry, timeout}; /// Matches the AI SDK's `isOutputChunk`: only chunks containing model output /// start or reset the first/chunk output timers. -fn is_output_chunk(part: &StreamPart) -> bool { +fn is_output_chunk(part: &StreamPart) -> bool { match part { StreamPart::TextDelta { delta, .. } | StreamPart::ReasoningDelta { delta, .. } | StreamPart::ToolInputDelta { delta, .. } => !delta.is_empty(), - StreamPart::ToolCall { .. } | StreamPart::File { .. } => true, + StreamPart::ToolCall(_) | StreamPart::File { .. } => true, _ => false, } } @@ -107,6 +107,12 @@ impl GenerateTextOptions { self, prompt: crate::language_model_message::LanguageModelPrompt, ) -> CallOptions { + let tool_choice = self.tool_choice.or_else(|| { + self.tools + .as_ref() + .filter(|tools| !tools.is_empty()) + .map(|_| ToolChoice::Auto) + }); CallOptions { prompt, max_output_tokens: self.max_output_tokens, @@ -119,7 +125,7 @@ impl GenerateTextOptions { response_format: self.response_format, seed: self.seed, tools: self.tools, - tool_choice: self.tool_choice.unwrap_or_default(), + tool_choice, headers: self.headers, provider_options: self.provider_options, reasoning: self.reasoning, @@ -221,8 +227,8 @@ pub struct GenerateObjectResult { /// Result of `stream_text` (user-facing). pub struct StreamTextResult { - /// The stream of `StreamPart` items. - pub stream: Pin> + Send>>, + /// The stream of `TextStreamPart` items. + pub stream: Pin> + Send>>, /// The request body that was sent (for debugging / cache probing, RFC-0015). pub request_body: Option, /// Response headers. @@ -373,49 +379,11 @@ impl StreamTextResult { StreamPart::ReasoningEnd { provider_metadata, .. } => rm.reasoning_end(provider_metadata), - StreamPart::ToolCall { - tool_call_id, - tool_name, - input, - provider_executed, - dynamic, - thought_signature, - provider_metadata, - invalid, - error, - .. - } => { - let call = crate::tool::ToolCall { - tool_call_id, - tool_name, - input, - provider_executed, - dynamic, - thought_signature, - provider_metadata, - invalid, - error, - }; + StreamPart::ToolCall(call) => { rm.tool_call(&call); tool_calls.push(call); } - StreamPart::ToolResult { - tool_call_id, - tool_name, - result, - is_error, - preliminary, - dynamic, - provider_metadata, - } => rm.tool_result( - tool_call_id, - tool_name, - result, - is_error, - preliminary, - dynamic, - provider_metadata, - ), + StreamPart::ToolResult(result) => rm.tool_result(result), StreamPart::Source { id, source_type, @@ -656,26 +624,9 @@ pub async fn generate_text( text.push_str(t); rm.text(t, provider_metadata.as_ref()); } - GenerateContent::ToolCall { - tool_call_id, - tool_name, - input, - provider_executed, - dynamic, - thought_signature, - provider_metadata, - .. - } => { + GenerateContent::ToolCall(call) => { let parsed = parse_tool_call( - RawToolCall { - tool_call_id: tool_call_id.clone(), - tool_name: tool_name.clone(), - input: input.clone(), - provider_executed: *provider_executed, - dynamic: *dynamic, - thought_signature: thought_signature.clone(), - provider_metadata: provider_metadata.clone(), - }, + call.clone(), tools.as_deref(), repair_tool_call.as_ref(), &messages, @@ -715,25 +666,7 @@ pub async fn generate_text( } // Provider-executed results stay in the assistant message so the // provider can replay its own server-tool transcript next turn. - GenerateContent::ToolResult { - tool_call_id, - tool_name, - result, - is_error, - preliminary, - dynamic, - provider_metadata, - } => { - rm.tool_result( - tool_call_id.clone(), - tool_name.clone(), - result.clone(), - *is_error, - *preliminary, - *dynamic, - provider_metadata.clone(), - ); - } + GenerateContent::ToolResult(result) => rm.tool_result(result.clone()), } } @@ -1097,48 +1030,22 @@ pub async fn stream_text( }) }; let mut stream = stream; - let stream: Pin> + Send>> = + let stream: Pin> + Send>> = Box::pin(async_stream::stream! { while let Some(item) = stream.next().await { match item { - Ok(StreamPart::ToolCall { - tool_call_id, - tool_name, - input, - provider_executed, - dynamic, - thought_signature, - provider_metadata, - .. - }) => { + Ok(StreamPart::ToolCall(raw)) => { let parsed = parse_tool_call( - RawToolCall { - tool_call_id, - tool_name, - input: crate::parse_tool_call::raw_tool_input(&input), - provider_executed, - dynamic, - thought_signature, - provider_metadata, - }, + raw, tools.as_deref(), repair_tool_call.as_ref(), &messages, operation_instructions.as_deref(), ).await; - yield Ok(StreamPart::ToolCall { - tool_call_id: parsed.tool_call_id, - tool_name: parsed.tool_name, - input: parsed.input, - provider_executed: parsed.provider_executed, - dynamic: parsed.dynamic, - thought_signature: parsed.thought_signature, - invalid: parsed.invalid, - error: parsed.error, - provider_metadata: parsed.provider_metadata, - }); + yield Ok(StreamPart::ToolCall(parsed)); } - item => yield item, + Ok(part) => yield Ok(part.map_tool_call(|_| unreachable!("matched above"))), + Err(error) => yield Err(error), } } }); @@ -1749,7 +1656,7 @@ mod stream_aggregation_tests { use super::*; use crate::types::FinishReasonUnified; - fn result_from(parts: Vec>) -> StreamTextResult { + fn result_from(parts: Vec>) -> StreamTextResult { StreamTextResult { stream: Box::pin(futures::stream::iter(parts)), request_body: None, @@ -1757,7 +1664,7 @@ mod stream_aggregation_tests { } } - fn delta(text: &str) -> StreamPart { + fn delta(text: &str) -> TextStreamPart { StreamPart::TextDelta { id: "text-1".into(), delta: text.into(), @@ -1765,7 +1672,7 @@ mod stream_aggregation_tests { } } - fn finish() -> StreamPart { + fn finish() -> TextStreamPart { StreamPart::Finish { finish_reason: FinishReason { unified: FinishReasonUnified::Stop, diff --git a/aimux-core/src/lib.rs b/aimux-core/src/lib.rs index 4781f40bf..27c5c84b9 100644 --- a/aimux-core/src/lib.rs +++ b/aimux-core/src/lib.rs @@ -96,9 +96,8 @@ pub mod prelude { }; pub use crate::options::{CallOptions, ResponseFormat, ToolChoice}; pub use crate::parse_tool_call::{ - RawToolCall, ToolCallRepair, ToolCallRepairContext, ToolCallRepairReply, - apply_tool_call_repair, apply_tool_call_repair_to_result, tool_call_repair_context, - tool_call_repair_inputs, + ToolCallRepair, ToolCallRepairContext, ToolCallRepairReply, apply_tool_call_repair, + apply_tool_call_repair_to_result, tool_call_repair_context, tool_call_repair_inputs, }; pub use crate::provider::Provider; pub use crate::reranking_model::{ @@ -120,8 +119,11 @@ pub mod prelude { SharedProviderOptions, SharedProviderReference, Size, }; pub use crate::speech_model::{SpeechCallOptions, SpeechModel, SpeechResult, generate_speech}; - pub use crate::stream_part::StreamPart; - pub use crate::tool::{FunctionTool, ProviderTool, Tool, ToolCall, ToolResult}; + pub use crate::stream_part::{StreamPart, TextStreamPart}; + pub use crate::tool::{ + FunctionTool, FunctionToolInputExample, ProviderTool, RawToolCall, Tool, ToolCall, + ToolResult, + }; pub use crate::transcription_model::{ AudioChunk, InputAudioFormat, TranscriptionCallOptions, TranscriptionModel, TranscriptionResult, TranscriptionStreamOptions, TranscriptionStreamPart, diff --git a/aimux-core/src/moa.rs b/aimux-core/src/moa.rs index 71979d609..5111ced50 100644 --- a/aimux-core/src/moa.rs +++ b/aimux-core/src/moa.rs @@ -16,7 +16,7 @@ use serde::Deserialize; use crate::composite::{ChildModel, add_usage, build_aggregator_prompt, extract_text}; use crate::error::AiMuxError; use crate::language_model::LanguageModel; -use crate::options::{CallOptions, ToolChoice}; +use crate::options::CallOptions; use crate::result::{GenerateResult, StreamResult}; use crate::retry; use crate::stream_part::StreamPart; @@ -98,12 +98,12 @@ impl MoaModel { /// Build reference call options from the user's options. When /// `strip_reference_tools` is set, `tools` is cleared and `tool_choice` is - /// reset to `Auto` so references don't carry tool schemas. + /// cleared so references don't carry tool schemas. fn reference_options(&self, options: &CallOptions) -> CallOptions { let mut o = options.clone(); if self.config.strip_reference_tools { o.tools = None; - o.tool_choice = ToolChoice::Auto; + o.tool_choice = None; } o } @@ -306,699 +306,3 @@ impl LanguageModel for MoaModel { }) } } - -#[cfg(test)] -mod tests { - use super::*; - use crate::content::ContentPart; - use crate::language_model_message::LanguageModelPromptMessage; - use crate::message::Role; - use crate::result::GenerateContent; - use crate::stream_part::StreamPart; - use crate::types::{FinishReason, FinishReasonUnified, TokenUsage}; - use async_trait::async_trait; - use futures::stream; - use std::sync::Arc; - use std::sync::atomic::{AtomicUsize, Ordering}; - - /// A mock child that returns fixed text with fixed usage. `fail` forces an - /// error. Used as both reference and aggregator. - struct MockChild { - name: &'static str, - text: String, - fail: bool, - usage: Usage, - retry_failures: usize, - calls: Arc, - } - - #[async_trait] - impl LanguageModel for MockChild { - fn provider(&self) -> &str { - "mock" - } - fn model_id(&self) -> &str { - self.name - } - async fn do_generate(&self, _options: &CallOptions) -> Result { - let attempt = self.calls.fetch_add(1, Ordering::SeqCst) + 1; - if self.fail { - return Err(AiMuxError::Other(format!("{} failed", self.name))); - } - if attempt <= self.retry_failures { - return Err(retryable_error(self.name)); - } - Ok(GenerateResult { - content: vec![GenerateContent::Text { - text: self.text.clone(), - provider_metadata: None, - }], - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }, - usage: self.usage.clone(), - warnings: vec![], - provider_metadata: None, - response: crate::types::ResponseMetadata { - model_id: Some(self.name.into()), - ..Default::default() - }, - request_body: None, - response_headers: None, - }) - } - async fn do_stream(&self, _options: &CallOptions) -> Result { - let attempt = self.calls.fetch_add(1, Ordering::SeqCst) + 1; - if self.fail { - return Err(AiMuxError::Other(format!("{} failed", self.name))); - } - if attempt <= self.retry_failures { - return Err(retryable_error(self.name)); - } - let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::TextDelta { - id: "t1".into(), - delta: self.text.clone(), - provider_metadata: None, - }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }, - usage: self.usage.clone(), - provider_metadata: None, - }), - ]; - Ok(StreamResult { - stream: Box::pin(stream::iter(parts)), - request_body: None, - response_headers: None, - }) - } - } - - fn retryable_error(name: &str) -> AiMuxError { - AiMuxError::ApiCall(Box::new(crate::ApiCallError { - status_code: Some(503), - response_headers: Some(std::collections::HashMap::from([( - "retry-after-ms".into(), - "0".into(), - )])), - is_retryable: true, - ..crate::ApiCallError::new( - format!("{name} retryable failure"), - "https://example.test", - serde_json::json!({}), - ) - })) - } - - fn retry_child( - name: &'static str, - text: &str, - retry_failures: usize, - ) -> (ChildModel, Arc) { - let calls = Arc::new(AtomicUsize::new(0)); - ( - Arc::new(MockChild { - name, - text: text.into(), - fail: false, - usage: Usage::default(), - retry_failures, - calls: calls.clone(), - }), - calls, - ) - } - - fn mk(name: &'static str, text: &str, fail: bool, total: u32) -> ChildModel { - Arc::new(MockChild { - name, - text: text.into(), - fail, - usage: Usage { - input_tokens: TokenUsage { - total: Some(total), - ..Default::default() - }, - output_tokens: TokenUsage { - total: Some(total), - ..Default::default() - }, - raw: None, - }, - retry_failures: 0, - calls: Arc::new(AtomicUsize::new(0)), - }) - } - - fn opts_with_prompt() -> CallOptions { - CallOptions { - prompt: vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("What is Rust?")], - provider_options: None, - }], - ..Default::default() - } - } - - #[tokio::test] - async fn generate_fans_out_and_aggregates() { - let refs = vec![ - mk("ref-a", "A says hi", false, 10), - mk("ref-b", "B says hi", false, 20), - ]; - let agg = mk("aggregator", "aggregated answer", false, 5); - let moa = MoaModel::new(refs, agg, MoaConfig::default()); - - let r = moa.do_generate(&opts_with_prompt()).await.unwrap(); - // Aggregator text wins as the content. - assert_eq!(extract_text(&r.content), "aggregated answer"); - // usage = references(10 + 20) + aggregator(5) = 35 on each side. - assert_eq!(r.usage.input_tokens.total, Some(35)); - assert_eq!(r.usage.output_tokens.total, Some(35)); - } - - #[tokio::test] - async fn best_effort_drops_failed_reference_and_warns() { - let refs = vec![mk("ref-a", "A ok", false, 10), mk("ref-b", "B", true, 0)]; - let agg = mk("aggregator", "final", false, 1); - let moa = MoaModel::new(refs, agg, MoaConfig::default()); - - let r = moa.do_generate(&opts_with_prompt()).await.unwrap(); - assert_eq!(extract_text(&r.content), "final"); - // Failed reference still contributed its (zero) usage; only ref-a (10) + agg (1). - assert_eq!(r.usage.input_tokens.total, Some(11)); - assert!(r.warnings.iter().any(|w| matches!( - w, - Warning::Other { message } if message.contains("reference 1 failed") - ))); - } - - #[tokio::test] - async fn fail_fast_propagates_first_reference_error() { - let refs = vec![mk("ref-a", "A", true, 0), mk("ref-b", "B ok", false, 10)]; - let agg = mk("aggregator", "final", false, 1); - let moa = MoaModel::new( - refs, - agg, - MoaConfig { - fail_mode: MoaFailMode::FailFast, - ..Default::default() - }, - ); - let err = moa.do_generate(&opts_with_prompt()).await.unwrap_err(); - assert!(err.to_string().contains("ref-a failed")); - } - - #[tokio::test] - async fn all_references_fail_is_error() { - let refs = vec![mk("ref-a", "A", true, 0), mk("ref-b", "B", true, 0)]; - let agg = mk("aggregator", "final", false, 1); - let moa = MoaModel::new(refs, agg, MoaConfig::default()); - let err = moa.do_generate(&opts_with_prompt()).await.unwrap_err(); - assert!(err.to_string().contains("all reference models failed")); - } - - #[tokio::test] - async fn no_references_degrades_to_aggregator() { - // With 0 references, aggregator runs the original prompt directly. - let agg = mk("aggregator", "solo answer", false, 3); - let moa = MoaModel::new(vec![], agg, MoaConfig::default()); - let r = moa.do_generate(&opts_with_prompt()).await.unwrap(); - assert_eq!(extract_text(&r.content), "solo answer"); - assert_eq!(r.usage.input_tokens.total, Some(3)); - } - - #[tokio::test] - async fn stream_runs_references_then_streams_aggregator() { - let refs = vec![mk("ref-a", "A", false, 7), mk("ref-b", "B", false, 8)]; - let agg = mk("aggregator", "streamed", false, 2); - let moa = MoaModel::new(refs, agg, MoaConfig::default()); - - let result = moa.do_stream(&opts_with_prompt()).await.unwrap(); - let mut s = result.stream; - let mut saw_start = false; - let mut text = String::new(); - let mut finish_usage: Option = None; - while let Some(part) = s.next().await { - match part.unwrap() { - StreamPart::StreamStart { .. } => saw_start = true, - StreamPart::TextDelta { delta, .. } => text.push_str(&delta), - StreamPart::Finish { usage, .. } => finish_usage = Some(usage), - _ => {} - } - } - assert!(saw_start); - assert_eq!(text, "streamed"); - // references(7 + 8) + aggregator(2) = 17 on each side. - assert_eq!(finish_usage.unwrap().input_tokens.total, Some(17)); - } - - #[tokio::test] - async fn retries_references_and_aggregator_independently_without_replaying_fanout() { - let (ref_a, ref_a_calls) = retry_child("ref-a", "A", 1); - let (ref_b, ref_b_calls) = retry_child("ref-b", "B", 0); - let (aggregator, aggregator_calls) = retry_child("aggregator", "final", 1); - let moa = MoaModel::new(vec![ref_a, ref_b], aggregator, MoaConfig::default()); - assert_eq!(moa.retry_config().max_retries, 0); - - let result = crate::generate::generate_text( - &moa, - "hello", - crate::generate::GenerateTextOptions { - max_retries: Some(2), - ..Default::default() - }, - ) - .await - .unwrap(); - - assert_eq!(result.text, "final"); - assert_eq!(ref_a_calls.load(Ordering::SeqCst), 2); - assert_eq!(ref_b_calls.load(Ordering::SeqCst), 1); - assert_eq!(aggregator_calls.load(Ordering::SeqCst), 2); - } - - #[tokio::test] - async fn retries_aggregator_stream_setup_without_replaying_references() { - let (reference, reference_calls) = retry_child("reference", "A", 0); - let (aggregator, aggregator_stream_calls) = retry_child("aggregator", "final", 1); - let moa = MoaModel::new(vec![reference], aggregator, MoaConfig::default()); - - let _result = crate::generate::stream_text( - &moa, - "hello", - crate::generate::GenerateTextOptions { - max_retries: Some(2), - ..Default::default() - }, - ) - .await - .unwrap(); - - assert_eq!(reference_calls.load(Ordering::SeqCst), 1); - assert_eq!(aggregator_stream_calls.load(Ordering::SeqCst), 2); - } - - #[tokio::test] - async fn stream_best_effort_emits_drop_warning_in_start() { - let refs = vec![mk("ref-a", "A", false, 5), mk("ref-b", "B", true, 0)]; - let agg = mk("aggregator", "out", false, 1); - let moa = MoaModel::new(refs, agg, MoaConfig::default()); - - let result = moa.do_stream(&opts_with_prompt()).await.unwrap(); - let mut s = result.stream; - // First item is our StreamStart, which should carry the drop warning. - match s.next().await.unwrap().unwrap() { - StreamPart::StreamStart { warnings } => { - assert!(warnings - .iter() - .any(|w| matches!(w, Warning::Other { message } if message.contains("reference 1 failed")))); - } - other => panic!("expected StreamStart, got {other:?}"), - } - } - - #[tokio::test] - async fn strip_reference_tools_default_clears_tools() { - // Verify the helper clears tools/tool_choice when configured (default). - let agg = mk("agg", "x", false, 0); - let moa = MoaModel::new(vec![], agg, MoaConfig::default()); - let opts = CallOptions { - prompt: vec![], - tools: Some(vec![]), - tool_choice: ToolChoice::Required, - ..Default::default() - }; - let ref_opts = moa.reference_options(&opts); - assert!(ref_opts.tools.is_none()); - assert_eq!(ref_opts.tool_choice, ToolChoice::Auto); - // The original options are untouched (clone, not mutate). - assert_eq!(opts.tool_choice, ToolChoice::Required); - } - - #[tokio::test] - async fn no_references_does_not_inject_empty_heading() { - // S4 regression guard: with 0 references the aggregator prompt is the - // original prompt verbatim — no "# Reference model responses" heading. - let prompt = vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("hello")], - provider_options: None, - }]; - let built = build_aggregator_prompt(&prompt, None, &[]); - // Still just the single original message; nothing appended. - assert_eq!(built.len(), 1); - // Heading must NOT appear. - assert!(!built[0].content.iter().any(|p| matches!( - p, - ContentPart::Text { text, .. } if text.contains("Reference model responses") - ))); - } - - /// A mock aggregator whose stream yields TextDelta then a transport Err, - /// then (misbehaving) more parts. Guards that MoA relays the Err and stops. - struct MidStreamErrorAggregator; - - #[async_trait] - impl LanguageModel for MidStreamErrorAggregator { - fn provider(&self) -> &str { - "mock" - } - fn model_id(&self) -> &str { - "agg-err" - } - async fn do_generate(&self, _options: &CallOptions) -> Result { - unimplemented!() - } - async fn do_stream(&self, _options: &CallOptions) -> Result { - let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::TextDelta { - id: "t1".into(), - delta: "partial".into(), - provider_metadata: None, - }), - // A transport error mid-stream. - Err(AiMuxError::Other("aggregator transport blew up".into())), - // A misbehaving aggregator that keeps going after the Err. - Ok(StreamPart::TextDelta { - id: "t2".into(), - delta: "should-not-see".into(), - provider_metadata: None, - }), - ]; - Ok(StreamResult { - stream: Box::pin(stream::iter(parts)), - request_body: None, - response_headers: None, - }) - } - } - - #[tokio::test] - async fn stream_relays_aggregator_transport_error_and_stops() { - // N3 regression guard: a mid-stream transport Err from the aggregator - // is relayed and the stream terminates — no subsequent parts forwarded. - let refs = vec![mk("ref-a", "A", false, 5)]; - let agg = Arc::new(MidStreamErrorAggregator); - let moa = MoaModel::new(refs, agg, MoaConfig::default()); - let result = moa.do_stream(&opts_with_prompt()).await.unwrap(); - let parts: Vec<_> = result.stream.collect().await; - let mut saw_error = false; - let mut text = String::new(); - for part in parts { - match part { - Ok(StreamPart::StreamStart { .. }) => {} - Ok(StreamPart::TextDelta { delta, .. }) => text.push_str(&delta), - Err(e) => { - saw_error = true; - assert!(e.to_string().contains("blew up"), "got: {e}"); - } - _ => {} - } - } - assert!( - saw_error, - "expected the aggregator transport error to be relayed" - ); - assert!( - !text.contains("should-not-see"), - "parts after the terminal Err must not be forwarded; got: {text}" - ); - } - - struct RecoverableFrameErrorAggregator; - - #[async_trait] - impl LanguageModel for RecoverableFrameErrorAggregator { - fn provider(&self) -> &str { - "mock" - } - - fn model_id(&self) -> &str { - "agg-recoverable-error" - } - - async fn do_generate(&self, _options: &CallOptions) -> Result { - unimplemented!() - } - - async fn do_stream(&self, _options: &CallOptions) -> Result { - Ok(StreamResult { - stream: Box::pin(stream::iter([ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Err(AiMuxError::JsonParse("malformed SSE data".into())), - Ok(StreamPart::TextDelta { - id: "t1".into(), - delta: "after-error".into(), - provider_metadata: None, - }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: Some("stop".into()), - }, - usage: Usage::default(), - provider_metadata: None, - }), - ])), - request_body: None, - response_headers: None, - }) - } - } - - #[tokio::test] - async fn stream_continues_after_recoverable_aggregator_frame_error() { - let refs = vec![mk("ref-a", "A", false, 5)]; - let moa = MoaModel::new( - refs, - Arc::new(RecoverableFrameErrorAggregator), - MoaConfig::default(), - ); - let result = moa.do_stream(&opts_with_prompt()).await.unwrap(); - let outcomes: Vec<_> = result.stream.collect().await; - let error_index = outcomes - .iter() - .position(|outcome| matches!(outcome, Err(AiMuxError::JsonParse(_)))) - .expect("parse error should be relayed"); - - assert!(outcomes[error_index + 1..].iter().any(|outcome| matches!( - outcome, - Ok(StreamPart::TextDelta { delta, .. }) if delta == "after-error" - ))); - assert!(outcomes.iter().any(|outcome| matches!( - outcome, - Ok(StreamPart::Finish { finish_reason, .. }) - if finish_reason.unified == FinishReasonUnified::Stop - ))); - } - - /// Like `mk` but with full per-field usage (for G4 accumulation testing). - fn mk_full( - name: &'static str, - text: &str, - total: u32, - cache_read: u32, - cache_write: u32, - text_tokens: u32, - reasoning_tokens: u32, - ) -> ChildModel { - Arc::new(MockChild { - name, - text: text.into(), - fail: false, - usage: Usage { - input_tokens: TokenUsage { - total: Some(total), - cache_read: Some(cache_read), - cache_write: Some(cache_write), - ..Default::default() - }, - output_tokens: TokenUsage { - total: Some(total), - text: Some(text_tokens), - reasoning: Some(reasoning_tokens), - ..Default::default() - }, - raw: None, - }, - retry_failures: 0, - calls: Arc::new(AtomicUsize::new(0)), - }) - } - - #[tokio::test] - async fn generate_accumulates_per_field_usage() { - // G4: add_usage must sum cache_read/cache_write/text/reasoning, not - // just total. ref-a: cache_read=5,cache_write=2,text=3,reasoning=1; - // ref-b: cache_read=7,cache_write=4,text=6,reasoning=2; agg totals=1. - let refs = vec![ - mk_full("ref-a", "A", 10, 5, 2, 3, 1), - mk_full("ref-b", "B", 20, 7, 4, 6, 2), - ]; - let agg = mk_full("aggregator", "out", 1, 0, 0, 1, 0); - let moa = MoaModel::new(refs, agg, MoaConfig::default()); - let r = moa.do_generate(&opts_with_prompt()).await.unwrap(); - // input: cache_read 5+7+0=12, cache_write 2+4+0=6 - assert_eq!(r.usage.input_tokens.cache_read, Some(12)); - assert_eq!(r.usage.input_tokens.cache_write, Some(6)); - // output: text 3+6+1=10, reasoning 1+2+0=3 - assert_eq!(r.usage.output_tokens.text, Some(10)); - assert_eq!(r.usage.output_tokens.reasoning, Some(3)); - } - - #[tokio::test] - async fn strip_reference_tools_false_preserves_tools() { - // G5: strip_reference_tools=false keeps tools/tool_choice on references. - let agg = mk("agg", "x", false, 0); - let moa = MoaModel::new( - vec![], - agg, - MoaConfig { - strip_reference_tools: false, - ..Default::default() - }, - ); - let opts = CallOptions { - prompt: vec![], - tools: Some(vec![]), - tool_choice: ToolChoice::Required, - ..Default::default() - }; - let ref_opts = moa.reference_options(&opts); - assert_eq!(ref_opts.tools.as_ref().map(std::vec::Vec::len), Some(0)); - assert_eq!(ref_opts.tool_choice, ToolChoice::Required); - } - - #[tokio::test] - async fn aggregator_failure_on_generate_propagates() { - // G11: if the aggregator's do_generate fails, the error surfaces - // (references already ran successfully). - let refs = vec![mk("ref-a", "A", false, 5)]; - let agg = mk("aggregator", "out", true, 1); - let moa = MoaModel::new(refs, agg, MoaConfig::default()); - let err = moa.do_generate(&opts_with_prompt()).await.unwrap_err(); - assert!(err.to_string().contains("aggregator failed")); - } - - /// A mock aggregator whose stream emits non-text parts (ToolCall, Source, - /// Reasoning) so we can verify they're forwarded by MoA's do_stream. - struct RichStreamAggregator; - - #[async_trait] - impl LanguageModel for RichStreamAggregator { - fn provider(&self) -> &str { - "mock" - } - fn model_id(&self) -> &str { - "rich-agg" - } - async fn do_generate(&self, _: &CallOptions) -> Result { - unimplemented!() - } - async fn do_stream(&self, _: &CallOptions) -> Result { - let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::TextDelta { - id: "t1".into(), - delta: "answer".into(), - provider_metadata: None, - }), - Ok(StreamPart::ToolCall { - tool_call_id: "tc1".into(), - tool_name: "search".into(), - input: serde_json::json!({"q": "rust"}), - provider_executed: None, - dynamic: None, - thought_signature: None, - invalid: None, - error: None, - provider_metadata: None, - }), - Ok(StreamPart::Source { - id: "s1".into(), - source_type: "url".into(), - url: Some("https://example.com".into()), - title: Some("Ex".into()), - provider_metadata: None, - }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }, - usage: Usage::default(), - provider_metadata: None, - }), - ]; - Ok(StreamResult { - stream: Box::pin(stream::iter(parts)), - request_body: None, - response_headers: None, - }) - } - } - - #[tokio::test] - async fn stream_forwards_non_text_parts() { - // G10: ToolCall/Source deltas from the aggregator must pass through - // (not just TextDelta). - let refs = vec![mk("ref-a", "A", false, 5)]; - let agg = Arc::new(RichStreamAggregator); - let moa = MoaModel::new(refs, agg, MoaConfig::default()); - let result = moa.do_stream(&opts_with_prompt()).await.unwrap(); - let parts: Vec<_> = result.stream.collect().await; - let mut saw_tool_call = false; - let mut saw_source = false; - for part in parts { - match part.unwrap() { - StreamPart::ToolCall { tool_name, .. } => { - saw_tool_call = true; - assert_eq!(tool_name, "search"); - } - StreamPart::Source { id, .. } => { - saw_source = true; - assert_eq!(id, "s1"); - } - _ => {} - } - } - assert!(saw_tool_call, "aggregator ToolCall must be forwarded"); - assert!(saw_source, "aggregator Source must be forwarded"); - } - - #[tokio::test] - async fn nested_router_as_moa_reference() { - // G8: composability — a RouterModel can be used as a MoaModel reference - // (both are Arc). RFC-0022 §6.2 explicitly says this - // is supported and "needs testing". The router's fallback should run - // inside the MoA fanout, and its output feeds the aggregator. - let router = crate::router::RouterModel::new( - vec![ - mk("router-primary", "router picked primary", true, 0), - mk("router-backup", "router fell back", false, 4), - ], - Box::new(crate::router::RuleRouter), - crate::router::FallbackPolicy::OnError, - crate::router::RouterConfig::default(), - ); - let router_child: ChildModel = Arc::new(router); - let plain_ref = mk("plain-ref", "plain reference", false, 3); - let agg = mk("aggregator", "aggregated", false, 1); - let moa = MoaModel::new(vec![router_child, plain_ref], agg, MoaConfig::default()); - - let r = moa.do_generate(&opts_with_prompt()).await.unwrap(); - // Aggregator output wins as content. - assert_eq!(extract_text(&r.content), "aggregated"); - // Usage: router(fell back to backup=4) + plain_ref(3) + aggregator(1) = 8. - assert_eq!(r.usage.input_tokens.total, Some(8)); - } -} diff --git a/aimux-core/src/openai_output.rs b/aimux-core/src/openai_output.rs index 1eb531e1f..12ffb37bf 100644 --- a/aimux-core/src/openai_output.rs +++ b/aimux-core/src/openai_output.rs @@ -38,7 +38,8 @@ use crate::parse_tool_call::raw_tool_call_text; use crate::result::GenerateContent; use crate::result::GenerateResult; use crate::shared::FileData; -use crate::stream_part::StreamPart; +use crate::stream_part::{StreamPart, TextStreamPart}; +use crate::tool::{ToolCall, ToolResult}; use crate::types::{FinishReason, FinishReasonUnified, Usage}; // ───────────────────────────────────────────────────────────────────────────── @@ -240,25 +241,20 @@ pub fn to_chat_completion(result: &GenerateResult, model: &str) -> ChatCompletio GenerateContent::Reasoning { text, .. } => { reasoning_text.push_str(text); } - GenerateContent::ToolCall { - tool_call_id, - tool_name, - input, - .. - } => { + GenerateContent::ToolCall(call) => { // The provider's raw argument text passes through verbatim; // OpenAI's wire format requires a JSON object even when the // model emitted no arguments at all. - let arguments = if input.trim().is_empty() { + let arguments = if call.input.trim().is_empty() { "{}".to_string() } else { - input.clone() + call.input.clone() }; tool_calls.push(ChatCompletionToolCall { - id: tool_call_id.clone(), + id: call.tool_call_id.clone(), tool_type: "function".to_string(), function: ChatCompletionFunction { - name: tool_name.clone(), + name: call.tool_name.clone(), arguments, }, }); @@ -291,12 +287,12 @@ pub fn to_chat_completion(result: &GenerateResult, model: &str) -> ChatCompletio } content_text.push_str(&file_data_to_text(data, media_type)); } - GenerateContent::ToolResult { + GenerateContent::ToolResult(ToolResult { tool_name, result, is_error, .. - } => { + }) => { // Degraded mapping: provider-executed tool result as text. if !content_text.is_empty() { content_text.push('\n'); @@ -409,7 +405,7 @@ impl std::fmt::Debug for ChatCompletionStream { /// SDK; they are never held back. See `stream_text_as_openai`. #[must_use] pub fn to_chat_completion_stream( - stream: Pin> + Send>>, + stream: Pin> + Send>>, model: &str, options: OpenAiStreamOptions, ) -> ChatCompletionStream { @@ -542,7 +538,7 @@ impl StreamState { /// Process a single StreamPart, returning zero or more output chunks. fn process_part( &mut self, - part: &StreamPart, + part: &TextStreamPart, include_reasoning: bool, include_usage: bool, ) -> Vec { @@ -702,14 +698,14 @@ impl StreamState { } StreamPart::ToolInputEnd { .. } => {} - StreamPart::ToolCall { + StreamPart::ToolCall(ToolCall { tool_call_id, tool_name, input, invalid, error, .. - } => { + }) => { // Complete tool call (e.g. from non-streaming-style providers). // A provider may carry all input on its start frame and emit no // deltas. In that case the final call is the first point where @@ -791,9 +787,9 @@ impl StreamState { } } - StreamPart::ToolResult { + StreamPart::ToolResult(ToolResult { tool_name, result, .. - } => { + }) => { // Degraded: provider-executed tool result as content. if let Some(c) = self.ensure_started() { chunks.push(c); @@ -1061,923 +1057,3 @@ fn random_id() -> String { let ts = now_unix(); format!("{ts:012x}{count:012x}") } - -// ───────────────────────────────────────────────────────────────────────────── -// Tests -// ───────────────────────────────────────────────────────────────────────────── - -#[cfg(test)] -mod tests { - use super::*; - use crate::result::GenerateResult; - use crate::types::TokenUsage; - use crate::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; - use futures::StreamExt; - use serde_json::json; - - fn make_result(content: Vec) -> GenerateResult { - GenerateResult { - content, - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }, - usage: Usage::default(), - warnings: vec![], - provider_metadata: None, - response: ResponseMetadata { - id: Some("chatcmpl-test123".to_string()), - timestamp: None, - model_id: Some("gpt-4o".to_string()), - }, - request_body: None, - response_headers: None, - } - } - - #[test] - fn test_text_only_non_streaming() { - let result = make_result(vec![GenerateContent::Text { - text: "Hello world".to_string(), - provider_metadata: None, - }]); - let completion = to_chat_completion(&result, "gpt-4o"); - - assert_eq!(completion.object, "chat.completion"); - assert_eq!(completion.id, "chatcmpl-test123"); - assert_eq!(completion.model, "gpt-4o"); - assert_eq!(completion.choices.len(), 1); - assert_eq!(completion.choices[0].index, 0); - assert_eq!( - completion.choices[0].message.content.as_deref(), - Some("Hello world") - ); - assert_eq!(completion.choices[0].message.role, "assistant"); - assert_eq!(completion.choices[0].finish_reason.as_deref(), Some("stop")); - assert!(completion.choices[0].message.tool_calls.is_none()); - assert!(completion.choices[0].message.reasoning_content.is_none()); - } - - #[test] - fn test_tool_call_non_streaming() { - let result = make_result(vec![ - GenerateContent::Text { - text: "Let me check.".to_string(), - provider_metadata: None, - }, - GenerateContent::ToolCall { - tool_call_id: "call_abc".to_string(), - tool_name: "get_weather".to_string(), - input: r#"{"city":"Tokyo"}"#.to_string(), - provider_executed: None, - dynamic: None, - thought_signature: None, - provider_metadata: None, - }, - ]); - // Set finish_reason to ToolCalls. - let mut result = result; - result.finish_reason = FinishReason { - unified: FinishReasonUnified::ToolCalls, - raw: None, - }; - - let completion = to_chat_completion(&result, "gpt-4o"); - - let msg = &completion.choices[0].message; - assert_eq!(msg.content.as_deref(), Some("Let me check.")); - let tool_calls = msg.tool_calls.as_ref().unwrap(); - assert_eq!(tool_calls.len(), 1); - assert_eq!(tool_calls[0].id, "call_abc"); - assert_eq!(tool_calls[0].tool_type, "function"); - assert_eq!(tool_calls[0].function.name, "get_weather"); - assert_eq!(tool_calls[0].function.arguments, r#"{"city":"Tokyo"}"#); - assert_eq!( - completion.choices[0].finish_reason.as_deref(), - Some("tool_calls") - ); - } - - #[test] - fn test_raw_tool_call_arguments_are_not_double_encoded() { - let result = make_result(vec![GenerateContent::ToolCall { - tool_call_id: "call_raw".to_string(), - tool_name: "get_weather".to_string(), - input: r#"{"city":"Tokyo"}"#.to_string(), - provider_executed: None, - dynamic: None, - thought_signature: None, - provider_metadata: None, - }]); - - let completion = to_chat_completion(&result, "gpt-4o"); - assert_eq!( - completion.choices[0].message.tool_calls.as_ref().unwrap()[0] - .function - .arguments, - r#"{"city":"Tokyo"}"# - ); - } - - #[test] - fn blank_invalid_tool_arguments_render_as_an_empty_object() { - // Both OpenAI renderers: blank raw text has no wire representation - // other than `{}`. - let error = AiMuxError::InvalidToolInput { - tool_name: "get_weather".to_string(), - tool_input: String::new(), - cause: "Type validation failed: missing 'city'".to_string(), - }; - let arguments = parsed_tool_call_arguments(&json!({}), Some(true), Some(&error)); - assert_eq!(arguments, "{}"); - - let result = make_result(vec![GenerateContent::ToolCall { - tool_call_id: "call_blank".to_string(), - tool_name: "get_weather".to_string(), - input: " ".to_string(), - provider_executed: None, - dynamic: None, - thought_signature: None, - provider_metadata: None, - }]); - assert_eq!( - to_chat_completion(&result, "gpt-4o").choices[0] - .message - .tool_calls - .as_ref() - .unwrap()[0] - .function - .arguments, - "{}" - ); - } - - #[test] - fn parsed_tool_call_arguments_null_is_not_rewritten_to_empty_object() { - // A valid call whose parsed input happens to be `null` (no schema - // constraint rejected it) must round-trip as `null`, not `{}`. - let arguments = parsed_tool_call_arguments(&Value::Null, None, None); - assert_eq!(arguments, "null"); - } - - /// Regression coverage for the P1 finding on PR #165: `input: Value` - /// alone cannot distinguish a validly parsed JSON string from malformed - /// text wrapped as a fallback string, so an invalid call's arguments are - /// recovered from its typed error instead. Driven end to end so the - /// error/`input` pairing `parse_tool_call` actually produces is pinned - /// too, not just the renderer. - #[tokio::test] - async fn invalid_tool_calls_stream_their_raw_arguments_verbatim() { - let failing_repair = crate::parse_tool_call::ToolCallRepair::new(|_context| async { - Err(AiMuxError::Other("repair model failed".to_string())) - }); - let cases: [(&str, &str, Option<&crate::parse_tool_call::ToolCallRepair>); 3] = [ - // Valid JSON string rejected by the schema: the quotes must - // survive, or the arguments field stops being JSON at all. - ("get_weather", r#""hello""#, None), - // Unknown tool *and* malformed text: `NoSuchTool` carries no raw - // text of its own, so it comes off the unparsable `input`. - ("unknown_tool", r#"{"a":"#, None), - // A failed repair callback wraps the pre-repair error. - ("get_weather", r#"{"a":"#, Some(&failing_repair)), - ]; - let tools = [crate::tool::Tool::Function(crate::tool::FunctionTool::new( - "get_weather", - json!({"type": "object", "properties": {}, "additionalProperties": false}), - ))]; - - for (tool_name, raw_input, repair) in cases { - let parsed = crate::parse_tool_call::parse_tool_call( - crate::parse_tool_call::RawToolCall { - tool_call_id: "call_1".to_string(), - tool_name: tool_name.to_string(), - input: raw_input.to_string(), - provider_executed: None, - dynamic: None, - thought_signature: None, - provider_metadata: None, - }, - Some(&tools), - repair, - &[], - None, - ) - .await; - assert_eq!(parsed.invalid, Some(true), "{raw_input}"); - - let parts: Vec> = vec![Ok(StreamPart::ToolCall { - tool_call_id: parsed.tool_call_id, - tool_name: parsed.tool_name, - input: parsed.input, - provider_executed: parsed.provider_executed, - dynamic: parsed.dynamic, - thought_signature: parsed.thought_signature, - invalid: parsed.invalid, - error: parsed.error, - provider_metadata: parsed.provider_metadata, - })]; - let chunks = collect_stream(to_chat_completion_stream( - Box::pin(futures::stream::iter(parts)), - "gpt-4o", - OpenAiStreamOptions::default(), - )) - .await; - - let arguments = chunks - .iter() - .find_map(|c| { - c.choices - .first()? - .delta - .tool_calls - .as_ref()? - .first()? - .function - .arguments - .clone() - }) - .expect("tool call arguments chunk not found"); - assert_eq!(arguments, raw_input, "tool {tool_name}"); - } - } - - #[test] - fn test_tool_call_null_content() { - // Tool call with no text → content should be null. - let result = make_result(vec![GenerateContent::ToolCall { - tool_call_id: "call_abc".to_string(), - tool_name: "get_weather".to_string(), - input: r#"{"city":"Tokyo"}"#.to_string(), - provider_executed: None, - dynamic: None, - thought_signature: None, - provider_metadata: None, - }]); - let completion = to_chat_completion(&result, "gpt-4o"); - - assert_eq!(completion.choices[0].message.content, None); - assert!(completion.choices[0].message.tool_calls.is_some()); - } - - #[test] - fn test_reasoning_non_streaming() { - let result = make_result(vec![ - GenerateContent::Reasoning { - text: "Thinking...".to_string(), - provider_metadata: None, - }, - GenerateContent::Text { - text: "Answer".to_string(), - provider_metadata: None, - }, - ]); - let completion = to_chat_completion(&result, "gpt-4o"); - - assert_eq!( - completion.choices[0].message.reasoning_content.as_deref(), - Some("Thinking...") - ); - assert_eq!( - completion.choices[0].message.content.as_deref(), - Some("Answer") - ); - } - - #[test] - fn test_usage_round_trip() { - let usage = Usage { - input_tokens: TokenUsage { - total: Some(100), - cache_read: Some(30), - cache_write: Some(10), - ..Default::default() - }, - output_tokens: TokenUsage { - total: Some(50), - reasoning: Some(20), - ..Default::default() - }, - raw: None, - }; - let mut result = make_result(vec![]); - result.usage = usage; - - let completion = to_chat_completion(&result, "gpt-4o"); - let u = &completion.usage; - - assert_eq!(u.prompt_tokens, 100); - assert_eq!(u.completion_tokens, 50); - assert_eq!(u.total_tokens, 150); - assert_eq!(u.prompt_tokens_details.as_ref().unwrap().cached_tokens, 30); - assert_eq!( - u.prompt_tokens_details.as_ref().unwrap().cache_write_tokens, - Some(10) - ); - assert_eq!( - u.completion_tokens_details - .as_ref() - .unwrap() - .reasoning_tokens, - Some(20) - ); - } - - #[test] - fn test_finish_reason_mapping() { - let cases = [ - (FinishReasonUnified::Stop, "stop"), - (FinishReasonUnified::Length, "length"), - (FinishReasonUnified::ContentFilter, "content_filter"), - (FinishReasonUnified::ToolCalls, "tool_calls"), - (FinishReasonUnified::Error, "stop"), - ]; - for (unified, expected) in cases { - let fr = FinishReason { unified, raw: None }; - assert_eq!(finish_reason_to_openai(&fr).as_deref(), Some(expected)); - } - } - - // ── Streaming tests ── - - async fn collect_stream(stream: ChatCompletionStream) -> Vec { - let mut chunks = Vec::new(); - let mut s = stream.stream; - while let Some(result) = s.next().await { - chunks.push(result.unwrap()); - } - chunks - } - - #[tokio::test] - async fn test_stream_text_only() { - let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::TextStart { - id: "0".to_string(), - provider_metadata: None, - }), - Ok(StreamPart::TextDelta { - id: "0".to_string(), - delta: "Hello ".to_string(), - provider_metadata: None, - }), - Ok(StreamPart::TextDelta { - id: "0".to_string(), - delta: "world".to_string(), - provider_metadata: None, - }), - Ok(StreamPart::TextEnd { - id: "0".to_string(), - provider_metadata: None, - }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }, - usage: Usage::default(), - provider_metadata: None, - }), - ]; - - let input_stream = Box::pin(futures::stream::iter(parts)); - let result = - to_chat_completion_stream(input_stream, "gpt-4o", OpenAiStreamOptions::default()); - let chunks = collect_stream(result).await; - - // First chunk: role frame. - assert_eq!( - chunks[0].choices[0].delta.role.as_deref(), - Some("assistant") - ); - assert_eq!(chunks[0].choices[0].delta.content.as_deref(), Some("")); - - // Content deltas. - assert_eq!( - chunks[1].choices[0].delta.content.as_deref(), - Some("Hello ") - ); - assert_eq!(chunks[2].choices[0].delta.content.as_deref(), Some("world")); - - // Final chunk: finish_reason. - let last = chunks.last().unwrap(); - assert_eq!(last.choices[0].finish_reason.as_deref(), Some("stop")); - assert!(last.choices[0].delta.content.is_none()); - } - - #[tokio::test] - async fn test_stream_tool_calls() { - let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::ToolInputStart { - id: "call_1".to_string(), - tool_name: "get_weather".to_string(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }), - Ok(StreamPart::ToolInputDelta { - id: "call_1".to_string(), - delta: "{\"city\":".to_string(), - provider_metadata: None, - }), - Ok(StreamPart::ToolInputDelta { - id: "call_1".to_string(), - delta: "\"Tokyo\"}".to_string(), - provider_metadata: None, - }), - Ok(StreamPart::ToolInputEnd { - id: "call_1".to_string(), - provider_metadata: None, - }), - Ok(StreamPart::ToolCall { - tool_call_id: "call_1".to_string(), - tool_name: "get_weather".to_string(), - input: json!({ "city": "Tokyo" }), - provider_executed: None, - dynamic: None, - thought_signature: None, - invalid: None, - error: None, - provider_metadata: None, - }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::ToolCalls, - raw: None, - }, - usage: Usage::default(), - provider_metadata: None, - }), - ]; - - let input_stream = Box::pin(futures::stream::iter(parts)); - let result = - to_chat_completion_stream(input_stream, "gpt-4o", OpenAiStreamOptions::default()); - let chunks = collect_stream(result).await; - - // Find the tool_call opening chunk (with id + name). - let open_chunk = chunks - .iter() - .find(|c| { - c.choices - .first() - .and_then(|ch| ch.delta.tool_calls.as_ref()) - .and_then(|tcs| tcs.first()) - .and_then(|tc| tc.id.as_deref()) - == Some("call_1") - }) - .expect("tool call opening chunk not found"); - - let tc = open_chunk.choices[0] - .delta - .tool_calls - .as_ref() - .unwrap() - .first() - .unwrap(); - assert_eq!(tc.index, 0); - assert_eq!(tc.id.as_deref(), Some("call_1")); - assert_eq!(tc.tool_type.as_deref(), Some("function")); - assert_eq!(tc.function.name.as_deref(), Some("get_weather")); - - // Find argument delta chunks. - let arg_chunks: Vec<_> = chunks - .iter() - .filter(|c| { - c.choices - .first() - .and_then(|ch| ch.delta.tool_calls.as_ref()) - .and_then(|tcs| tcs.first()) - .is_some_and(|tc| tc.id.is_none() && tc.function.arguments.is_some()) - }) - .collect(); - assert_eq!(arg_chunks.len(), 2); - assert_eq!( - arg_chunks[0].choices[0].delta.tool_calls.as_ref().unwrap()[0] - .function - .arguments - .as_deref(), - Some("{\"city\":") - ); - - // Final chunk. - let last = chunks.last().unwrap(); - assert_eq!(last.choices[0].finish_reason.as_deref(), Some("tool_calls")); - } - - #[tokio::test] - async fn test_stream_complete_call_backfills_arguments_absent_from_deltas() { - let parts: Vec> = vec![ - Ok(StreamPart::ToolInputStart { - id: "call_1".to_string(), - tool_name: "fetch".to_string(), - provider_executed: Some(true), - dynamic: None, - title: None, - provider_metadata: None, - }), - Ok(StreamPart::ToolInputEnd { - id: "call_1".to_string(), - provider_metadata: None, - }), - Ok(StreamPart::ToolCall { - tool_call_id: "call_1".to_string(), - tool_name: "fetch".to_string(), - input: json!({ "url": "https://example.com" }), - provider_executed: Some(true), - dynamic: None, - thought_signature: None, - invalid: None, - error: None, - provider_metadata: None, - }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }, - usage: Usage::default(), - provider_metadata: None, - }), - ]; - - let result = to_chat_completion_stream( - Box::pin(futures::stream::iter(parts)), - "claude", - OpenAiStreamOptions::default(), - ); - let chunks = collect_stream(result).await; - let arguments = chunks - .iter() - .filter_map(|chunk| { - chunk - .choices - .first()? - .delta - .tool_calls - .as_ref()? - .first()? - .function - .arguments - .as_deref() - }) - .collect::(); - - assert_eq!(arguments, r#"{"url":"https://example.com"}"#); - } - - #[tokio::test] - async fn test_stream_multiple_tool_calls_index() { - let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::ToolInputStart { - id: "call_a".to_string(), - tool_name: "tool_a".to_string(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }), - Ok(StreamPart::ToolInputStart { - id: "call_b".to_string(), - tool_name: "tool_b".to_string(), - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::ToolCalls, - raw: None, - }, - usage: Usage::default(), - provider_metadata: None, - }), - ]; - - let input_stream = Box::pin(futures::stream::iter(parts)); - let result = - to_chat_completion_stream(input_stream, "gpt-4o", OpenAiStreamOptions::default()); - let chunks = collect_stream(result).await; - - // Find both tool call opening chunks. - let open_chunks: Vec<_> = chunks - .iter() - .filter(|c| { - c.choices - .first() - .and_then(|ch| ch.delta.tool_calls.as_ref()) - .and_then(|tcs| tcs.first()) - .and_then(|tc| tc.id.as_deref()) - .is_some() - }) - .collect(); - assert_eq!(open_chunks.len(), 2); - - let tc0 = &open_chunks[0].choices[0].delta.tool_calls.as_ref().unwrap()[0]; - let tc1 = &open_chunks[1].choices[0].delta.tool_calls.as_ref().unwrap()[0]; - assert_eq!(tc0.index, 0); - assert_eq!(tc0.id.as_deref(), Some("call_a")); - assert_eq!(tc1.index, 1); - assert_eq!(tc1.id.as_deref(), Some("call_b")); - } - - #[tokio::test] - async fn test_stream_reasoning() { - let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::ReasoningStart { - id: "reasoning-0".to_string(), - provider_metadata: None, - }), - Ok(StreamPart::ReasoningDelta { - id: "reasoning-0".to_string(), - delta: "Hmm...".to_string(), - provider_metadata: None, - }), - Ok(StreamPart::ReasoningEnd { - id: "reasoning-0".to_string(), - provider_metadata: None, - }), - Ok(StreamPart::TextDelta { - id: "0".to_string(), - delta: "Answer".to_string(), - provider_metadata: None, - }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }, - usage: Usage::default(), - provider_metadata: None, - }), - ]; - - let input_stream = Box::pin(futures::stream::iter(parts)); - let result = - to_chat_completion_stream(input_stream, "gpt-4o", OpenAiStreamOptions::default()); - let chunks = collect_stream(result).await; - - // Find reasoning chunk. - let reasoning_chunk = chunks - .iter() - .find(|c| c.choices[0].delta.reasoning_content.is_some()) - .expect("reasoning chunk not found"); - assert_eq!( - reasoning_chunk.choices[0] - .delta - .reasoning_content - .as_deref(), - Some("Hmm...") - ); - - // Find text chunk. - let text_chunk = chunks - .iter() - .find(|c| { - c.choices[0] - .delta - .content - .as_deref() - .is_some_and(|s| s == "Answer") - }) - .expect("text chunk not found"); - assert_eq!( - text_chunk.choices[0].delta.content.as_deref(), - Some("Answer") - ); - } - - #[tokio::test] - async fn test_stream_usage_in_final_chunk() { - let usage = Usage { - input_tokens: TokenUsage { - total: Some(42), - ..Default::default() - }, - output_tokens: TokenUsage { - total: Some(10), - ..Default::default() - }, - raw: None, - }; - - let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::TextDelta { - id: "0".to_string(), - delta: "Hi".to_string(), - provider_metadata: None, - }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }, - usage, - provider_metadata: None, - }), - ]; - - let input_stream = Box::pin(futures::stream::iter(parts)); - let result = to_chat_completion_stream( - input_stream, - "gpt-4o", - OpenAiStreamOptions { - include_usage: true, - include_reasoning: true, - }, - ); - let chunks = collect_stream(result).await; - - let last = chunks.last().unwrap(); - assert_eq!(last.usage.as_ref().unwrap().prompt_tokens, 42); - assert_eq!(last.usage.as_ref().unwrap().completion_tokens, 10); - assert_eq!(last.usage.as_ref().unwrap().total_tokens, 52); - } - - #[tokio::test] - async fn test_stream_no_usage_when_disabled() { - let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }, - usage: Usage { - input_tokens: TokenUsage { - total: Some(42), - ..Default::default() - }, - output_tokens: TokenUsage::default(), - raw: None, - }, - provider_metadata: None, - }), - ]; - - let input_stream = Box::pin(futures::stream::iter(parts)); - let result = to_chat_completion_stream( - input_stream, - "gpt-4o", - OpenAiStreamOptions { - include_usage: false, - include_reasoning: true, - }, - ); - let chunks = collect_stream(result).await; - - let last = chunks.last().unwrap(); - assert!(last.usage.is_none()); - } - - #[tokio::test] - async fn test_stream_error() { - let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::TextDelta { - id: "0".to_string(), - delta: "partial".to_string(), - provider_metadata: None, - }), - Ok(StreamPart::Error { - error: AiMuxError::ApiCall(Box::new(crate::error::ApiCallError::new( - "something went wrong", - "https://example.test/v1", - serde_json::json!({}), - ))), - }), - ]; - - let input_stream = Box::pin(futures::stream::iter(parts)); - let result = - to_chat_completion_stream(input_stream, "gpt-4o", OpenAiStreamOptions::default()); - let chunks = collect_stream(result).await; - - // Error should appear as a content delta. - let error_chunk = chunks - .iter() - .find(|c| { - c.choices[0] - .delta - .content - .as_deref() - .is_some_and(|s| s.contains("[error]")) - }) - .expect("error chunk not found"); - - assert!( - error_chunk.choices[0] - .delta - .content - .as_deref() - .unwrap() - .contains("something went wrong") - ); - - // Final chunk should still have finish_reason. - let last = chunks.last().unwrap(); - assert_eq!(last.choices[0].finish_reason.as_deref(), Some("stop")); - } - - #[tokio::test] - async fn recoverable_frame_error_does_not_truncate_chat_completion_stream() { - let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Err(AiMuxError::JsonParse("malformed SSE data".into())), - Ok(StreamPart::TextDelta { - id: "0".to_string(), - delta: "after-error".to_string(), - provider_metadata: None, - }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: Some("stop".into()), - }, - usage: Usage::default(), - provider_metadata: None, - }), - ]; - - let result = to_chat_completion_stream( - Box::pin(futures::stream::iter(parts)), - "gpt-4o", - OpenAiStreamOptions::default(), - ); - let outcomes: Vec<_> = result.stream.collect().await; - let error_index = outcomes - .iter() - .position(|outcome| matches!(outcome, Err(AiMuxError::JsonParse(_)))) - .expect("parse error should be preserved as an error item"); - - assert!(outcomes[error_index + 1..].iter().any(|outcome| matches!( - outcome, - Ok(chunk) - if chunk.choices.first().and_then(|choice| choice.delta.content.as_deref()) - == Some("after-error") - ))); - assert!(outcomes.iter().any(|outcome| matches!( - outcome, - Ok(chunk) - if chunk.choices.first().and_then(|choice| choice.finish_reason.as_deref()) - == Some("stop") - ))); - } - - #[tokio::test] - async fn test_stream_completes_without_finish() { - // Stream ends without Finish — should still emit a final chunk. - let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::TextDelta { - id: "0".to_string(), - delta: "Hello".to_string(), - provider_metadata: None, - }), - ]; - - let input_stream = Box::pin(futures::stream::iter(parts)); - let result = - to_chat_completion_stream(input_stream, "gpt-4o", OpenAiStreamOptions::default()); - let chunks = collect_stream(result).await; - - let last = chunks.last().unwrap(); - assert_eq!(last.choices[0].finish_reason.as_deref(), Some("stop")); - } - - #[test] - fn test_sse_encoding() { - let chunk = ChatCompletionChunk { - id: "chatcmpl-test".to_string(), - object: "chat.completion.chunk".to_string(), - created: 1234567890, - model: "gpt-4o".to_string(), - choices: vec![ChatCompletionChunkChoice { - index: 0, - delta: ChatCompletionDelta { - content: Some("Hi".to_string()), - ..Default::default() - }, - finish_reason: None, - logprobs: None, - }], - usage: None, - }; - - let sse = encode_chunk_sse(&chunk); - assert!(sse.starts_with("data: {")); - assert!(sse.ends_with("\n\n")); - assert!(sse.contains("\"content\":\"Hi\"")); - } - - #[test] - fn test_done_frame() { - assert_eq!(DONE_FRAME, "data: [DONE]\n\n"); - } -} diff --git a/aimux-core/src/options.rs b/aimux-core/src/options.rs index 3d2891396..d7a4010ac 100644 --- a/aimux-core/src/options.rs +++ b/aimux-core/src/options.rs @@ -10,7 +10,7 @@ use ts_rs::TS; use crate::AbortSignal; use crate::language_model_message::LanguageModelPrompt; -pub use crate::tool::{FunctionTool, ProviderTool, Tool, ToolChoice}; +pub use crate::tool::{FunctionTool, FunctionToolInputExample, ProviderTool, Tool, ToolChoice}; use crate::types::ReasoningEffort; /// How the model should format its response. @@ -102,7 +102,7 @@ pub struct CallOptions { pub tools: Option>, /// How the model should choose tools. - pub tool_choice: ToolChoice, + pub tool_choice: Option, /// Extra HTTP headers. pub headers: Option>, @@ -203,7 +203,7 @@ impl CallOptions { response_format: None, seed: None, tools: None, - tool_choice: ToolChoice::default(), + tool_choice: None, headers: None, provider_options: None, reasoning: None, diff --git a/aimux-core/src/parse_tool_call.rs b/aimux-core/src/parse_tool_call.rs index 16a4ca6d8..053245b6e 100644 --- a/aimux-core/src/parse_tool_call.rs +++ b/aimux-core/src/parse_tool_call.rs @@ -14,28 +14,7 @@ use serde_json::Value; use ts_rs::TS; use crate::error::AiMuxError; -use crate::tool::{Tool, ToolCall}; -use crate::types::ProviderMetadata; - -/// Provider-facing tool call before Core parses and validates its input. -/// -/// The wire shape matches [`ToolCall`] field for field, except that `input` -/// is the provider's raw argument *text* rather than a parsed value. -#[derive(Debug, Clone, Serialize, Deserialize, TS)] -#[ts(export)] -pub struct RawToolCall { - pub tool_call_id: String, - pub tool_name: String, - pub input: String, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub provider_executed: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub dynamic: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub thought_signature: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub provider_metadata: Option, -} +use crate::tool::{RawToolCall, Tool, ToolCall}; /// Context supplied to a one-shot tool-call repair callback. #[derive(Debug, Clone)] @@ -134,13 +113,13 @@ pub async fn parse_tool_call( }) }; return match parsed { - Ok(input) => valid_tool_call(tool_call, input, Some(true)), + Ok(input) => parsed_tool_call(tool_call, input, Some(true)), Err(error) => invalid_tool_call(tool_call, error), }; }; match parse_and_validate_tool_call(&tool_call, tools) { - Ok((input, dynamic)) => valid_tool_call(tool_call, input, dynamic), + Ok((input, dynamic)) => parsed_tool_call(tool_call, input, dynamic), Err(original_error) => { let Some(repair_tool_call) = repair_tool_call else { return invalid_tool_call(tool_call, original_error); @@ -181,7 +160,7 @@ fn apply_repair_outcome( ) -> ToolCall { let cause = match outcome { RepairOutcome::Repaired(repaired) => match parse_and_validate_tool_call(&repaired, tools) { - Ok((input, dynamic)) => return valid_tool_call(repaired, input, dynamic), + Ok((input, dynamic)) => return parsed_tool_call(repaired, input, dynamic), Err(repaired_error) => repaired_error, }, RepairOutcome::Unchanged => return invalid_tool_call(tool_call, original_error), @@ -298,24 +277,13 @@ fn contains_forbidden_prototype(value: &Value) -> bool { } } -/// Recover the provider's raw argument text: string inputs pass through -/// verbatim (possibly malformed JSON awaiting parse/repair); anything already -/// structured re-serializes. -pub(crate) fn raw_tool_input(input: &Value) -> String { - match input { - Value::String(input) => input.clone(), - input => serde_json::to_string(input).expect("serializing serde_json::Value cannot fail"), - } -} - -fn valid_tool_call(tool_call: RawToolCall, input: Value, dynamic: Option) -> ToolCall { +fn parsed_tool_call(tool_call: RawToolCall, input: Value, dynamic: Option) -> ToolCall { ToolCall { tool_call_id: tool_call.tool_call_id, tool_name: tool_call.tool_name, input, provider_executed: tool_call.provider_executed, dynamic, - thought_signature: tool_call.thought_signature, provider_metadata: tool_call.provider_metadata, invalid: None, error: None, @@ -332,15 +300,9 @@ fn invalid_tool_call(tool_call: RawToolCall, error: AiMuxError) -> ToolCall { let input = serde_json::from_str(&tool_call.input) .unwrap_or_else(|_| Value::String(tool_call.input.clone())); ToolCall { - tool_call_id: tool_call.tool_call_id, - tool_name: tool_call.tool_name, - input, - provider_executed: tool_call.provider_executed, - dynamic: Some(true), - thought_signature: tool_call.thought_signature, - provider_metadata: tool_call.provider_metadata, invalid: Some(true), error: Some(error), + ..parsed_tool_call(tool_call, input, Some(true)) } } @@ -400,8 +362,8 @@ pub(crate) fn raw_tool_call_text(error: &AiMuxError) -> Option { /// and that loses the JSON string quoting of a bare literal (`"Tokyo"` would /// come back as `Tokyo`), duplicate keys, number spelling and whitespace, all /// of which matter when a host re-submits the text unchanged under a new tool -/// name. Only errors from before the text was recorded fall back to -/// [`raw_tool_input`]. +/// name. Only errors from before the text was recorded fall back to the +/// compact JSON of `input`. /// /// `dynamic` is not recovered: an invalid call always reports `Some(true)`. /// That only matters for re-deriving the original error, which is read off the @@ -423,10 +385,14 @@ fn raw_from_invalid(tool_call: &ToolCall) -> Result<(RawToolCall, AiMuxError), A RawToolCall { tool_call_id: tool_call.tool_call_id.clone(), tool_name: tool_call.tool_name.clone(), - input: raw_tool_call_text(&error).unwrap_or_else(|| raw_tool_input(&tool_call.input)), + // An invalid call keeps unparsable argument text as a JSON string + // (see `invalid_tool_call`); that text is the raw input. + input: raw_tool_call_text(&error).unwrap_or_else(|| match &tool_call.input { + Value::String(text) => text.clone(), + input => input.to_string(), + }), provider_executed: tool_call.provider_executed, dynamic: tool_call.dynamic, - thought_signature: tool_call.thought_signature.clone(), provider_metadata: tool_call.provider_metadata.clone(), }, error, diff --git a/aimux-core/src/recording.rs b/aimux-core/src/recording.rs index 65e5fce7d..22dff9ba9 100644 --- a/aimux-core/src/recording.rs +++ b/aimux-core/src/recording.rs @@ -1626,11 +1626,11 @@ impl RecordingOutcomeStream { } } -impl futures::Stream for RecordingOutcomeStream +impl futures::Stream for RecordingOutcomeStream where - S: futures::Stream> + Unpin, + S: futures::Stream, crate::AiMuxError>> + Unpin, { - type Item = Result; + type Item = Result, crate::AiMuxError>; fn poll_next( mut self: std::pin::Pin<&mut Self>, diff --git a/aimux-core/src/replay.rs b/aimux-core/src/replay.rs index 462965061..6246fa03d 100644 --- a/aimux-core/src/replay.rs +++ b/aimux-core/src/replay.rs @@ -20,10 +20,11 @@ use crate::generate::{GenerateTextOptions, GenerateTextResult, generate_text}; use crate::language_model::LanguageModel; use crate::language_model_message::LanguageModelPrompt; use crate::message::{MessageContent, ModelMessage, ModelPrompt}; -use crate::options::{CallOptions, ToolChoice}; +use crate::options::CallOptions; use crate::recording::Recording; use crate::result::{GenerateContent, GenerateResult, StreamResult}; use crate::stream_part::StreamPart; +use crate::tool::RawToolCall; use crate::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; pub use crate::decision_replay::{MockDecisionReplayModel, replay_decision_with_model}; @@ -111,8 +112,7 @@ fn canonical_call_key(opts: &CallOptions) -> serde_json::Value { }) } -/// 从录制侧提取同样的规范键(向后兼容:缺失 Option 字段按 null,缺失 -/// `tool_choice` 按 `ToolChoice::default()`=Auto,与 `CallOptions` 缺省一致)。 +/// 从录制侧提取同样的规范键(缺失 Option 字段按 null)。 fn canonical_recording_key(rec: &Recording) -> serde_json::Value { let o = &rec.input.options; serde_json::json!({ @@ -122,10 +122,7 @@ fn canonical_recording_key(rec: &Recording) -> serde_json::Value { "seed": o.get("seed").cloned().unwrap_or_default(), "response_format": o.get("response_format").cloned().unwrap_or_default(), "tools": o.get("tools").cloned().unwrap_or_default(), - "tool_choice": o - .get("tool_choice") - .cloned() - .unwrap_or_else(|| serde_json::to_value(ToolChoice::default()).unwrap_or_default()), + "tool_choice": o.get("tool_choice").cloned().unwrap_or_default(), "headers": o.get("headers").cloned().unwrap_or_default(), "provider_options": o.get("provider_options").cloned().unwrap_or_default(), "body_overrides": o.get("body_overrides").cloned().unwrap_or_default(), @@ -601,15 +598,14 @@ fn rebuild_generate_result(rec: &Recording) -> Result Result Result { let mut final_usage = Usage::default(); let mut final_finish: Option = None; // tool_call 按 OpenAI `index` 稳定累积;`tool_order` 保插入顺序(确定性 emit)。 - let mut tool_calls: std::collections::HashMap = + let mut tool_calls: std::collections::HashMap = std::collections::HashMap::new(); let mut tool_order: Vec = Vec::new(); let mut saw_openai = false; @@ -791,10 +780,13 @@ fn rebuild_stream_result(rec: &Recording) -> Result { .to_string(); tool_calls.insert( idx, - ToolCallAccumulator { - id: id.clone(), - name: name.clone(), - arguments: String::new(), + RawToolCall { + tool_call_id: id.clone(), + tool_name: name.clone(), + input: String::new(), + provider_executed: None, + dynamic: None, + provider_metadata: None, }, ); tool_order.push(idx); @@ -813,9 +805,9 @@ fn rebuild_stream_result(rec: &Recording) -> Result { && (!is_new || !args.is_empty()) && let Some(acc) = tool_calls.get_mut(&idx) { - acc.arguments.push_str(args); + acc.input.push_str(args); parts.push(Ok(StreamPart::ToolInputDelta { - id: acc.id.clone(), + id: acc.tool_call_id.clone(), delta: args.to_string(), provider_metadata: None, })); @@ -861,22 +853,12 @@ fn rebuild_stream_result(rec: &Recording) -> Result { })); } for &i in &tool_order { - if let Some(acc) = tool_calls.get(&i) { + if let Some(acc) = tool_calls.remove(&i) { parts.push(Ok(StreamPart::ToolInputEnd { - id: acc.id.clone(), - provider_metadata: None, - })); - parts.push(Ok(StreamPart::ToolCall { - tool_call_id: acc.id.clone(), - tool_name: acc.name.clone(), - input: serde_json::Value::String(acc.arguments.clone()), - provider_executed: None, - dynamic: None, - thought_signature: None, - invalid: None, - error: None, + id: acc.tool_call_id.clone(), provider_metadata: None, })); + parts.push(Ok(StreamPart::ToolCall(acc))); } } @@ -1003,11 +985,7 @@ fn generate_options_from_call_options(o: CallOptions) -> GenerateTextOptions { response_format: o.response_format, seed: o.seed, tools: o.tools, - tool_choice: if o.tool_choice == crate::tool::ToolChoice::Auto { - None - } else { - Some(o.tool_choice) - }, + tool_choice: o.tool_choice, headers: o.headers, provider_options: o.provider_options, reasoning: o.reasoning, @@ -1626,18 +1604,18 @@ mod tests { .unwrap() }); assert_eq!(result.content.len(), 1); - let Some(GenerateContent::ToolCall { + let Some(GenerateContent::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. - }) = result.content.first() + })) = result.content.first() else { panic!("expected ToolCall, got {:?}", result.content); }; assert_eq!(tool_call_id, "call_abc"); assert_eq!(tool_name, "get_weather"); - assert_eq!(input, &serde_json::json!(r#"{"city":"SF"}"#)); + assert_eq!(input, r#"{"city":"SF"}"#); assert_eq!(result.finish_reason.unified, FinishReasonUnified::ToolCalls); } @@ -1671,10 +1649,11 @@ mod tests { .await .unwrap() }); - let Some(GenerateContent::ToolCall { input, .. }) = result.content.first() else { + let Some(GenerateContent::ToolCall(RawToolCall { input, .. })) = result.content.first() + else { panic!("expected ToolCall"); }; - assert_eq!(input, &serde_json::json!("not-json{")); + assert_eq!(input, "not-json{"); } #[test] @@ -1739,19 +1718,19 @@ mod tests { let calls: Vec<_> = parts .iter() .filter_map(|p| match p { - StreamPart::ToolCall { + StreamPart::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. - } => Some((tool_call_id.clone(), tool_name.clone(), input.clone())), + }) => Some((tool_call_id.clone(), tool_name.clone(), input.clone())), _ => None, }) .collect(); assert_eq!(calls.len(), 1); assert_eq!(calls[0].0, "call_1"); assert_eq!(calls[0].1, "get_weather"); - assert_eq!(calls[0].2, serde_json::json!(r#"{"city":"SF"}"#)); + assert_eq!(calls[0].2, r#"{"city":"SF"}"#); // finish_reason tool_calls。 assert!(parts.iter().any(|p| matches!( p, diff --git a/aimux-core/src/response_messages.rs b/aimux-core/src/response_messages.rs index 59c3e63ad..5c290702e 100644 --- a/aimux-core/src/response_messages.rs +++ b/aimux-core/src/response_messages.rs @@ -145,38 +145,27 @@ impl ResponseMessageBuilder { tool_name: call.tool_name.clone(), input: response_tool_call_input(&call.input, call.invalid), provider_executed: call.provider_executed, - thought_signature: call.thought_signature.clone(), provider_options: call.provider_metadata.clone(), }); } - #[allow(clippy::too_many_arguments)] - pub fn tool_result( - &mut self, - tool_call_id: String, - tool_name: String, - result: Value, - is_error: Option, - preliminary: Option, - dynamic: Option, - provider_options: Option, - ) { + pub fn tool_result(&mut self, result: crate::tool::ToolResult) { // Preliminary server-tool results are transient stream updates. The // provider contract requires a later final result, and only that // final value belongs in the replay transcript for the next turn. - if preliminary == Some(true) { + if result.preliminary == Some(true) { return; } self.flush_text(); self.flush_reasoning(); self.parts.push(ContentPart::ToolResult { - tool_call_id, - tool_name: Some(tool_name), - result, - is_error, - preliminary, - dynamic, - provider_options, + tool_call_id: result.tool_call_id, + tool_name: Some(result.tool_name), + result: result.result, + is_error: result.is_error, + preliminary: result.preliminary, + dynamic: result.dynamic, + provider_options: result.provider_metadata, }); } diff --git a/aimux-core/src/result.rs b/aimux-core/src/result.rs index c9d451aa0..c3c9811df 100644 --- a/aimux-core/src/result.rs +++ b/aimux-core/src/result.rs @@ -10,30 +10,9 @@ use ts_rs::TS; use crate::error::AiMuxError; use crate::shared::FileData; use crate::stream_part::StreamPart; +use crate::tool::{RawToolCall, ToolResult}; use crate::types::{FinishReason, ProviderMetadata, ResponseMetadata, Usage, Warning}; -use serde_json::Value; - -/// Compatibility deserializer for `GenerateContent::ToolCall.input`. -/// -/// The field used to be a `serde_json::Value` (the already-parsed arguments, -/// since providers parsed their own input) and became a `String` (the raw, -/// unparsed provider text) in the tool-input-parse-repair refactor, so a -/// `GenerateResult` persisted or recorded before it carries an -/// object/array/number/bool/null here. Accept both: a JSON string passes -/// through unchanged, and any other JSON value is re-serialized to its -/// compact JSON text so the field keeps meaning "the raw text a -/// schema-validating parse would run against". -fn deserialize_tool_call_input<'de, D>(deserializer: D) -> Result -where - D: serde::Deserializer<'de>, -{ - Ok(match Value::deserialize(deserializer)? { - Value::String(raw) => raw, - legacy => legacy.to_string(), - }) -} - /// A content item in the generation result. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] #[ts(export)] @@ -44,37 +23,9 @@ pub enum GenerateContent { #[serde(default, skip_serializing_if = "Option::is_none")] provider_metadata: Option, }, - /// A tool call requested by the model. - ToolCall { - tool_call_id: String, - tool_name: String, - /// The model's raw argument text, exactly as the provider delivered - /// it (possibly malformed). Providers never parse it — `generate_text` - /// owns parsing, schema validation, and repair. - /// - /// Always serializes as a JSON string. Deserializes a JSON string - /// (the current wire shape) unchanged, and also re-serializes the - /// pre-refactor shape — an already-parsed JSON value — to its compact - /// JSON text, so a `GenerateResult` persisted before this field - /// became a `String` keeps loading. See docs/api/gaps.md §9. - #[serde(deserialize_with = "deserialize_tool_call_input")] - input: String, - /// Whether the tool call will be executed by the provider. - /// If false/unset, the tool call is executed by the client. - #[serde(default, skip_serializing_if = "Option::is_none")] - provider_executed: Option, - /// Whether the tool is dynamic (defined at runtime, e.g. MCP tools). - #[serde(default, skip_serializing_if = "Option::is_none")] - dynamic: Option, - /// Provider-assigned thought signature (e.g. Google Gemini - /// `thoughtSignature`). Must be echoed back verbatim on the follow-up - /// turn when the tool result is sent. - #[serde(default, skip_serializing_if = "Option::is_none")] - thought_signature: Option, - /// Additional provider-specific metadata for the tool call. - #[serde(default, skip_serializing_if = "Option::is_none")] - provider_metadata: Option, - }, + /// A tool call requested by the model. `input` is the provider's raw + /// argument text; `generate_text` owns parsing, validation, and repair. + ToolCall(RawToolCall), /// A source / citation (e.g. URL citation from search-preview models). Source { id: String, @@ -106,23 +57,7 @@ pub enum GenerateContent { /// A tool result from a provider-executed tool (e.g. xAI file_search, /// web_search). Emitted alongside the preceding `ToolCall` when the /// provider executes the tool server-side. - ToolResult { - tool_call_id: String, - tool_name: String, - result: Value, - /// Whether the result is an error or error message. - #[serde(default, skip_serializing_if = "Option::is_none")] - is_error: Option, - /// Whether the result is preliminary (replaces prior, e.g. image previews). - #[serde(default, skip_serializing_if = "Option::is_none")] - preliminary: Option, - /// Whether the tool is dynamic (defined at runtime, e.g. MCP tools). - #[serde(default, skip_serializing_if = "Option::is_none")] - dynamic: Option, - /// Additional provider-specific metadata for the tool result. - #[serde(default, skip_serializing_if = "Option::is_none")] - provider_metadata: Option, - }, + ToolResult(ToolResult), } // ── M7: top-level aggregation types ───────────────────────────────────────── @@ -249,40 +184,3 @@ impl std::fmt::Debug for StreamResult { .finish() } } - -#[cfg(test)] -mod tests { - use super::*; - - /// Regression test for PR #165 review finding: `GenerateContent::ToolCall` - /// went from carrying `input: Value` (the parsed arguments) to - /// `input: String` (the raw text), so a `GenerateResult` persisted before - /// the refactor carries an object here and must still load rather than - /// fail closed. Serialization stays one-way — always a plain string. - #[test] - fn tool_call_input_loads_the_legacy_object_shape_and_the_current_string_shape() { - let wire = |input: serde_json::Value| { - serde_json::json!({ - "ToolCall": { - "tool_call_id": "call_1", - "tool_name": "get_weather", - "input": input, - } - }) - }; - let from_legacy: GenerateContent = - serde_json::from_value(wire(serde_json::json!({ "city": "Tokyo" }))).unwrap(); - let from_current: GenerateContent = - serde_json::from_value(wire(serde_json::json!(r#"{"city":"Tokyo"}"#))).unwrap(); - assert_eq!(from_legacy, from_current); - - let GenerateContent::ToolCall { input, .. } = &from_legacy else { - panic!("expected ToolCall, got {from_legacy:?}"); - }; - assert_eq!(input, r#"{"city":"Tokyo"}"#); - assert_eq!( - serde_json::to_value(&from_legacy).unwrap(), - wire(serde_json::json!(r#"{"city":"Tokyo"}"#)) - ); - } -} diff --git a/aimux-core/src/stream_part.rs b/aimux-core/src/stream_part.rs index 6a083920c..5e8888d6e 100644 --- a/aimux-core/src/stream_part.rs +++ b/aimux-core/src/stream_part.rs @@ -9,12 +9,17 @@ use ts_rs::TS; use crate::error::AiMuxError; use crate::shared::FileData; +use crate::tool::{RawToolCall, ToolCall, ToolResult}; use crate::types::{FinishReason, ProviderMetadata, Usage, Warning}; /// A single chunk in the stream returned by `do_stream`. +/// +/// Generic over the tool-call payload: providers emit the default +/// [`StreamPart`] (raw argument text); `stream_text` emits [`TextStreamPart`] +/// (parsed and validated [`ToolCall`]). #[derive(Debug, Clone, Serialize, Deserialize, TS)] -#[ts(export)] -pub enum StreamPart { +#[ts(export, concrete(C = ToolCall))] +pub enum StreamPart { // ── P0: text ── /// Start of a text segment. TextStart { @@ -79,50 +84,12 @@ pub enum StreamPart { provider_metadata: Option, }, /// A complete tool call (alternative to the start/delta/end flow). - ToolCall { - tool_call_id: String, - tool_name: String, - /// Serialized argument text in a `Value::String` from `do_stream`; - /// parsed input after `stream_text`. - input: Value, - /// Whether the tool call will be executed by the provider. - #[serde(default, skip_serializing_if = "Option::is_none")] - provider_executed: Option, - /// Whether the tool is dynamic (defined at runtime, e.g. MCP tools). - #[serde(default, skip_serializing_if = "Option::is_none")] - dynamic: Option, - /// Provider-assigned thought signature (e.g. Google Gemini - /// `thoughtSignature`). Must be echoed back verbatim on the follow-up - /// turn when the tool result is sent. - #[serde(default, skip_serializing_if = "Option::is_none")] - thought_signature: Option, - /// Set by Core when the call remains invalid after optional repair. - #[serde(default, skip_serializing_if = "Option::is_none")] - invalid: Option, - /// Typed lookup, parsing, schema, or repair failure. - #[serde(default, skip_serializing_if = "Option::is_none")] - error: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - provider_metadata: Option, - }, + /// + /// The provider layer emits [`RawToolCall`] (`input` is the raw argument + /// text); `stream_text` re-emits the parsed [`ToolCall`] ([`TextStreamPart`]). + ToolCall(C), /// A tool result (provider-executed tools). - ToolResult { - tool_call_id: String, - tool_name: String, - result: Value, - /// Whether the result is an error or error message. - #[serde(default, skip_serializing_if = "Option::is_none")] - is_error: Option, - /// Whether the result is preliminary (replaces prior, e.g. image previews). - #[serde(default, skip_serializing_if = "Option::is_none")] - preliminary: Option, - /// Whether the tool is dynamic (defined at runtime, e.g. MCP tools). - #[serde(default, skip_serializing_if = "Option::is_none")] - dynamic: Option, - /// Additional provider-specific metadata for the tool result. - #[serde(default, skip_serializing_if = "Option::is_none")] - provider_metadata: Option, - }, + ToolResult(ToolResult), // ── P2: file ── /// A file generated by the model (e.g. an image or document). @@ -178,3 +145,40 @@ pub enum StreamPart { /// A raw chunk from the provider (for debugging, when `include_raw_chunks` is set). Raw { raw_value: Value }, } + +/// The stream `stream_text` and later stages produce: tool calls are parsed. +pub type TextStreamPart = StreamPart; + +impl StreamPart { + /// Replace the tool-call payload, leaving every other variant untouched. + pub fn map_tool_call(self, f: impl FnOnce(C) -> D) -> StreamPart { + // The generic parameter forces every variant to be rebuilt. + macro_rules! carry { + ($($variant:ident { $($field:ident),* })*) => { + match self { + StreamPart::ToolCall(call) => StreamPart::ToolCall(f(call)), + StreamPart::ToolResult(result) => StreamPart::ToolResult(result), + $(StreamPart::$variant { $($field),* } => StreamPart::$variant { $($field),* },)* + } + }; + } + carry! { + TextStart { id, provider_metadata } + TextDelta { id, delta, provider_metadata } + TextEnd { id, provider_metadata } + StreamStart { warnings } + Finish { finish_reason, usage, provider_metadata } + Error { error } + ToolInputStart { id, tool_name, provider_executed, dynamic, title, provider_metadata } + ToolInputDelta { id, delta, provider_metadata } + ToolInputEnd { id, provider_metadata } + File { data, media_type, provider_metadata } + ReasoningStart { id, provider_metadata } + ReasoningDelta { id, delta, provider_metadata } + ReasoningEnd { id, provider_metadata } + ResponseMetadata { id, timestamp, model_id } + Source { id, source_type, url, title, provider_metadata } + Raw { raw_value } + } + } +} diff --git a/aimux-core/src/tool.rs b/aimux-core/src/tool.rs index 7c1176da1..421c26500 100644 --- a/aimux-core/src/tool.rs +++ b/aimux-core/src/tool.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use serde::{Deserialize, Serialize}; -use serde_json::Value; +use serde_json::{Map, Value}; use ts_rs::TS; use crate::error::AiMuxError; @@ -31,7 +31,14 @@ pub struct FunctionTool { /// Example inputs for the tool (V4 `inputExamples`), used by some providers /// to emit `input_examples` in the request body. #[serde(default, skip_serializing_if = "Option::is_none")] - pub input_examples: Option>, + pub input_examples: Option>, +} + +/// An example input for a function tool. +#[derive(Debug, Clone, Serialize, Deserialize, TS)] +#[ts(export)] +pub struct FunctionToolInputExample { + pub input: Map, } impl FunctionTool { @@ -79,7 +86,7 @@ pub struct ProviderTool { pub name: String, /// The arguments for configuring the tool. Must match the expected /// arguments defined by the provider for this tool. - pub args: Value, + pub args: Map, } /// A tool that can be either a user-defined function tool or a @@ -101,6 +108,24 @@ impl From for Tool { } } +/// Provider-facing tool call before Core parses and validates its input. +/// +/// The wire shape matches [`ToolCall`] field for field, except that `input` +/// is the provider's raw argument *text* rather than a parsed value. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] +#[ts(export)] +pub struct RawToolCall { + pub tool_call_id: String, + pub tool_name: String, + pub input: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider_executed: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub dynamic: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider_metadata: Option, +} + /// A tool call requested by the model. #[derive(Debug, Clone, Serialize, Deserialize, TS)] #[ts(export)] @@ -118,12 +143,6 @@ pub struct ToolCall { /// Whether the tool is dynamic (defined at runtime, e.g. MCP tools). #[serde(default, skip_serializing_if = "Option::is_none")] pub dynamic: Option, - /// Provider-assigned thought signature (e.g. Google Gemini - /// `thoughtSignature`). Must be echoed back verbatim on the follow-up turn - /// when the tool result is sent; thinking models reject the request - /// otherwise. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub thought_signature: Option, /// Additional provider-specific metadata associated with this call. #[serde(default, skip_serializing_if = "Option::is_none")] pub provider_metadata: Option, @@ -136,12 +155,14 @@ pub struct ToolCall { pub error: Option, } -/// The result of executing a tool call, to be sent back to the model. -#[derive(Debug, Clone, Serialize, Deserialize, TS)] +/// A tool result produced by a provider-executed tool (e.g. xAI `file_search`, +/// web search), carried next to the [`RawToolCall`] it answers. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] #[ts(export)] pub struct ToolResult { - /// Must match the corresponding `ToolCall::tool_call_id`. + /// Must match the corresponding call's `tool_call_id`. pub tool_call_id: String, + pub tool_name: String, /// The tool's output (usually a JSON-serializable value or plain text). pub result: Value, /// Whether the result is an error or error message. @@ -150,6 +171,12 @@ pub struct ToolResult { /// Whether the result is preliminary (replaces prior, e.g. image previews). #[serde(default, skip_serializing_if = "Option::is_none")] pub preliminary: Option, + /// Whether the tool is dynamic (defined at runtime, e.g. MCP tools). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub dynamic: Option, + /// Additional provider-specific metadata for the tool result. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider_metadata: Option, } /// How the model should choose tools. diff --git a/aimux-core/src/trace/layer.rs b/aimux-core/src/trace/layer.rs index a5b3b6e37..2ceec13a4 100644 --- a/aimux-core/src/trace/layer.rs +++ b/aimux-core/src/trace/layer.rs @@ -573,7 +573,7 @@ impl LanguageModel for TraceLayer { &item, Ok(StreamPart::TextDelta { .. }) | Ok(StreamPart::ReasoningDelta { .. }) - | Ok(StreamPart::ToolCall { .. }) + | Ok(StreamPart::ToolCall(_)) ); if is_model_output { ttft_obs.store(started.elapsed().as_millis() as u64, Ordering::Relaxed); diff --git a/aimux-core/tests/tool_input_test.rs b/aimux-core/tests/tool_input_test.rs index cbfcec813..655d94d3c 100644 --- a/aimux-core/tests/tool_input_test.rs +++ b/aimux-core/tests/tool_input_test.rs @@ -8,10 +8,10 @@ use aimux_core::generate::{ use aimux_core::language_model::LanguageModel; use aimux_core::openai_output::OpenAiStreamOptions; use aimux_core::options::CallOptions; -use aimux_core::parse_tool_call::{RawToolCall, ToolCallRepair, parse_tool_call}; +use aimux_core::parse_tool_call::{ToolCallRepair, parse_tool_call}; use aimux_core::result::{GenerateContent, GenerateResult, StreamResult}; use aimux_core::stream_part::StreamPart; -use aimux_core::tool::{FunctionTool, Tool}; +use aimux_core::tool::{FunctionTool, RawToolCall, Tool, ToolCall}; use aimux_core::types::{FinishReason, FinishReasonUnified, Usage}; use async_trait::async_trait; use futures::StreamExt; @@ -28,7 +28,6 @@ fn raw(name: &str, input: &str) -> RawToolCall { input: input.into(), provider_executed: None, dynamic: None, - thought_signature: None, provider_metadata: None, } } @@ -345,15 +344,14 @@ impl LanguageModel for RawToolModel { async fn do_generate(&self, _options: &CallOptions) -> Result { Ok(GenerateResult { - content: vec![GenerateContent::ToolCall { + content: vec![GenerateContent::ToolCall(RawToolCall { tool_call_id: "call-1".into(), tool_name: self.tool_name.into(), input: self.input.to_string(), provider_executed: None, dynamic: None, - thought_signature: None, provider_metadata: None, - }], + })], finish_reason: FinishReason { unified: FinishReasonUnified::ToolCalls, raw: Some("tool_calls".into()), @@ -398,17 +396,14 @@ impl LanguageModel for RawToolModel { ]); } parts.extend([ - Ok(StreamPart::ToolCall { + Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: "call-1".into(), tool_name: self.tool_name.into(), - input: json!(self.input), + input: self.input.to_string(), provider_executed: None, dynamic: None, - thought_signature: None, - invalid: None, - error: None, provider_metadata: None, - }), + })), Ok(StreamPart::Finish { finish_reason: FinishReason { unified: FinishReasonUnified::ToolCalls, @@ -444,8 +439,8 @@ async fn generate_text_parses_provider_raw_input_at_the_core_boundary() { assert_eq!(result.tool_calls[0].invalid, None); assert!(matches!( &result.raw.content[0], - GenerateContent::ToolCall { input, .. } - if input == &json!(r#"{"city":"Singapore"}"#) + GenerateContent::ToolCall(RawToolCall { input, .. }) + if input == r#"{"city":"Singapore"}"# )); } @@ -464,7 +459,7 @@ async fn stream_text_parses_provider_raw_input_at_the_core_boundary() { .unwrap(); while let Some(part) = result.stream.next().await { - if let StreamPart::ToolCall { input, invalid, .. } = part.unwrap() { + if let StreamPart::ToolCall(ToolCall { input, invalid, .. }) = part.unwrap() { assert_eq!(input, json!({ "city": "Singapore" })); assert_eq!(invalid, None); return; diff --git a/aimux-ffi/src/lib.rs b/aimux-ffi/src/lib.rs index e3e596c81..f91a2be3d 100644 --- a/aimux-ffi/src/lib.rs +++ b/aimux-ffi/src/lib.rs @@ -2102,7 +2102,9 @@ fn stream_text_with_signal( let cstr = match item { Ok(part) => stream_part_cstring(&part)?, Err(e) if e.is_recoverable_stream_error() => { - stream_part_cstring(&aimux_core::stream_part::StreamPart::Error { error: e })? + stream_part_cstring(&aimux_core::stream_part::TextStreamPart::Error { + error: e, + })? } Err(e) => return Err(e.into()), }; diff --git a/aimux-provider-utils/src/streaming_tool_call_tracker.rs b/aimux-provider-utils/src/streaming_tool_call_tracker.rs index cb9fedd43..fb862a71a 100644 --- a/aimux-provider-utils/src/streaming_tool_call_tracker.rs +++ b/aimux-provider-utils/src/streaming_tool_call_tracker.rs @@ -27,11 +27,11 @@ //! Protocols whose streams carry explicit tool-call boundaries (Anthropic, //! Google, Bedrock, Cohere, the Responses API) do not need this. -use serde_json::Value; use thiserror::Error; 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::{ @@ -228,18 +228,15 @@ impl StreamingToolCallTracker { id: call.id.clone(), provider_metadata: None, }); - out.push(StreamPart::ToolCall { + out.push(StreamPart::ToolCall(RawToolCall { tool_call_id: call.id, tool_name: call.name, // Raw argument text; Core parses it. - input: Value::String(call.arguments), + input: call.arguments, provider_executed: None, dynamic: None, - thought_signature: None, - invalid: None, - error: None, provider_metadata: call.provider_metadata, - }); + })); } } diff --git a/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs b/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs index 82b6b2726..a1ce8b16e 100644 --- a/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs +++ b/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use aimux_core::stream_part::StreamPart; +use aimux_core::tool::RawToolCall; use aimux_provider_utils::{ StreamingToolCallDelta, StreamingToolCallTracker, TrackerError, TypeValidation, }; @@ -35,8 +36,7 @@ enum Part { ToolCall { tool_call_id: String, tool_name: String, - /// The raw argument text (`StreamPart::ToolCall.input` is a - /// `Value::String` from the tracker). + /// The raw argument text. input: String, provider_metadata: Option, }, @@ -61,17 +61,14 @@ fn project(part: StreamPart) -> Part { id, provider_metadata: None, } => Part::ToolInputEnd { id }, - StreamPart::ToolCall { + StreamPart::ToolCall(RawToolCall { tool_call_id, tool_name, - input: Value::String(input), + input, provider_executed: None, dynamic: None, - thought_signature: None, - invalid: None, - error: None, provider_metadata, - } => Part::ToolCall { + }) => Part::ToolCall { tool_call_id, tool_name, input, diff --git a/aimux-providers/src/anthropic/convert.rs b/aimux-providers/src/anthropic/convert.rs index c6e98ca52..670501354 100644 --- a/aimux-providers/src/anthropic/convert.rs +++ b/aimux-providers/src/anthropic/convert.rs @@ -1919,7 +1919,7 @@ fn apply_anthropic_tools( } else { None }, - Some(&options.tool_choice), + options.tool_choice.as_ref(), disable_parallel_tool_use, caps.supports_structured_output, caps.supports_structured_output, diff --git a/aimux-providers/src/anthropic/prepare_tools.rs b/aimux-providers/src/anthropic/prepare_tools.rs index 2d42c8887..c7b4e606a 100644 --- a/aimux-providers/src/anthropic/prepare_tools.rs +++ b/aimux-providers/src/anthropic/prepare_tools.rs @@ -13,7 +13,7 @@ use std::collections::BTreeSet; use aimux_core::tool::FunctionTool; use aimux_core::tool::ToolChoice; use aimux_core::types::Warning; -use serde_json::{Value, json}; +use serde_json::{Map, Value, json}; /// Result of [`prepare_tools`] / [`prepare_tools_with_provider`]. #[derive(Debug, Clone)] @@ -121,13 +121,13 @@ pub enum AnthropicTool { Provider { id: String, name: String, - args: Value, + args: Map, }, } /// Read a camelCase arg from the `args` object, returning `Value::Null` when /// absent (mirroring the TS `undefined`-when-absent field behaviour). -fn arg(args: &Value, key: &str) -> Value { +fn arg(args: &Map, key: &str) -> Value { args.get(key).cloned().unwrap_or(Value::Null) } @@ -302,10 +302,7 @@ fn prepare_function_tool( def["allowed_callers"] = ac.clone(); } if let Some(ref examples) = tool.input_examples { - let inputs: Vec = examples - .iter() - .map(|e| e.get("input").cloned().unwrap_or(e.clone())) - .collect(); + let inputs: Vec<&Map> = examples.iter().map(|e| &e.input).collect(); def["input_examples"] = json!(inputs); } @@ -323,7 +320,7 @@ fn prepare_function_tool( /// `None` (so the caller emits an "unsupported" warning) when the id is unknown. pub(crate) fn prepare_provider_tool( id: &str, - args: &Value, + args: &Map, betas: &mut BTreeSet, ) -> Option { Some(match id { diff --git a/aimux-providers/src/anthropic/stream.rs b/aimux-providers/src/anthropic/stream.rs index 604222d7a..4f5e8dc71 100644 --- a/aimux-providers/src/anthropic/stream.rs +++ b/aimux-providers/src/anthropic/stream.rs @@ -15,6 +15,8 @@ //! - `parse_anthropic_content` — shared content-block → `GenerateContent` //! mapping used by the non-streaming path. +use aimux_core::tool::RawToolCall; +use aimux_core::tool::ToolResult; use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; @@ -399,7 +401,7 @@ pub(crate) fn stream_parts_for_result_block( ) -> Vec { let tool_result = |tool_name: String, (result, is_error): (Value, Option), tool_use_id: &str| { - StreamPart::ToolResult { + StreamPart::ToolResult(ToolResult { tool_call_id: tool_use_id.to_string(), tool_name, result, @@ -407,7 +409,7 @@ pub(crate) fn stream_parts_for_result_block( preliminary: None, dynamic: None, provider_metadata: None, - } + }) }; match block { @@ -417,7 +419,7 @@ pub(crate) fn stream_parts_for_result_block( } => { let name = names.to_custom_tool_name("web_search").to_string(); let Some(results) = content.as_array() else { - return vec![StreamPart::ToolResult { + return vec![StreamPart::ToolResult(ToolResult { tool_call_id: tool_use_id.clone(), tool_name: name, result: json!({ @@ -428,7 +430,7 @@ pub(crate) fn stream_parts_for_result_block( preliminary: None, dynamic: None, provider_metadata: None, - }]; + })]; }; // The tool result, then one Source per hit — the same pair the // non-streaming path emits. @@ -508,7 +510,7 @@ pub(crate) fn stream_parts_for_result_block( is_error, } => { let call = mcp_tool_calls.get(tool_use_id); - vec![StreamPart::ToolResult { + vec![StreamPart::ToolResult(ToolResult { tool_call_id: tool_use_id.clone(), tool_name: call.map(|(name, _)| name.clone()).unwrap_or_default(), result: content.clone(), @@ -520,7 +522,7 @@ pub(crate) fn stream_parts_for_result_block( "anthropic": { "type": "mcp-tool-use", "serverName": server } }) }), - }] + })] } _ => Vec::new(), } @@ -575,15 +577,14 @@ pub(crate) fn parse_anthropic_content( input, caller, } => { - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: id.clone(), tool_name: names.to_custom_tool_name(name).to_string(), input: input.to_string(), provider_executed: None, dynamic: None, - thought_signature: None, provider_metadata: tool_call_caller_metadata(caller.as_ref()), - }); + })); } ContentBlock::Thinking { thinking, @@ -600,7 +601,7 @@ pub(crate) fn parse_anthropic_content( // calls so they round-trip on follow-up turns. ContentBlock::ServerToolUse { id, name, input } => { let provider_name = server_tool_provider_name(name); - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: id.clone(), tool_name: names.to_custom_tool_name(provider_name).to_string(), input: normalized_server_tool_input(name, input).to_string(), @@ -608,9 +609,8 @@ pub(crate) fn parse_anthropic_content( dynamic: (provider_name == "code_execution" && names.mark_code_execution_dynamic()) .then_some(true), - thought_signature: None, provider_metadata: None, - }); + })); } // MCP tool use — provider-executed + dynamic (upstream :1166-1182). ContentBlock::McpToolUse { @@ -619,17 +619,16 @@ pub(crate) fn parse_anthropic_content( input, server_name, } => { - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: id.clone(), tool_name: name.clone(), input: input.to_string(), provider_executed: Some(true), dynamic: Some(true), - thought_signature: None, provider_metadata: Some(json!({ "anthropic": { "type": "mcp-tool-use", "serverName": server_name } })), - }); + })); } // Redacted thinking — upstream emits as reasoning with redactedData ContentBlock::RedactedThinking { data } => { @@ -653,7 +652,7 @@ pub(crate) fn parse_anthropic_content( // object (upstream branches on `Array.isArray`). match payload.as_array() { Some(results) => { - content.push(GenerateContent::ToolResult { + content.push(GenerateContent::ToolResult(ToolResult { tool_call_id: tool_use_id.clone(), tool_name: names.to_custom_tool_name("web_search").to_string(), result: Value::Array( @@ -663,7 +662,7 @@ pub(crate) fn parse_anthropic_content( preliminary: None, dynamic: None, provider_metadata: None, - }); + })); // Each result also becomes a `Source` — that is how the // URLs and titles reach `result.sources`. for result in results { @@ -681,7 +680,7 @@ pub(crate) fn parse_anthropic_content( }); } } - None => content.push(GenerateContent::ToolResult { + None => content.push(GenerateContent::ToolResult(ToolResult { tool_call_id: tool_use_id.clone(), tool_name: names.to_custom_tool_name("web_search").to_string(), result: json!({ @@ -693,7 +692,7 @@ pub(crate) fn parse_anthropic_content( preliminary: None, dynamic: None, provider_metadata: None, - }), + })), } } ContentBlock::WebFetchToolResult { @@ -701,7 +700,7 @@ pub(crate) fn parse_anthropic_content( content: payload, } => { let (result, is_error) = map_web_fetch_result(payload); - content.push(GenerateContent::ToolResult { + content.push(GenerateContent::ToolResult(ToolResult { tool_call_id: tool_use_id.clone(), tool_name: names.to_custom_tool_name("web_fetch").to_string(), result, @@ -709,14 +708,14 @@ pub(crate) fn parse_anthropic_content( preliminary: None, dynamic: None, provider_metadata: None, - }); + })); } ContentBlock::CodeExecutionToolResult { tool_use_id, content: payload, } => { let (result, is_error) = map_code_execution_result(payload); - content.push(GenerateContent::ToolResult { + content.push(GenerateContent::ToolResult(ToolResult { tool_call_id: tool_use_id.clone(), tool_name: names.to_custom_tool_name("code_execution").to_string(), result, @@ -724,7 +723,7 @@ pub(crate) fn parse_anthropic_content( preliminary: None, dynamic: None, provider_metadata: None, - }); + })); } // Upstream shares one arm for these two and passes `content` // through unmapped (anthropic-language-model.ts:1323-1334). @@ -736,7 +735,7 @@ pub(crate) fn parse_anthropic_content( tool_use_id, content: payload, } => { - content.push(GenerateContent::ToolResult { + content.push(GenerateContent::ToolResult(ToolResult { tool_call_id: tool_use_id.clone(), tool_name: names.to_custom_tool_name("code_execution").to_string(), result: payload.clone(), @@ -744,7 +743,7 @@ pub(crate) fn parse_anthropic_content( preliminary: None, dynamic: None, provider_metadata: None, - }); + })); } ContentBlock::ToolSearchToolResult { tool_use_id, @@ -755,7 +754,7 @@ pub(crate) fn parse_anthropic_content( names, server_tool_calls.get(tool_use_id.as_str()).copied(), ); - content.push(GenerateContent::ToolResult { + content.push(GenerateContent::ToolResult(ToolResult { tool_call_id: tool_use_id.clone(), tool_name: names.to_custom_tool_name(provider_name).to_string(), result, @@ -763,14 +762,14 @@ pub(crate) fn parse_anthropic_content( preliminary: None, dynamic: None, provider_metadata: None, - }); + })); } ContentBlock::AdvisorToolResult { tool_use_id, content: payload, } => { let (result, is_error) = map_advisor_result(payload); - content.push(GenerateContent::ToolResult { + content.push(GenerateContent::ToolResult(ToolResult { tool_call_id: tool_use_id.clone(), tool_name: names.to_custom_tool_name("advisor").to_string(), result, @@ -778,7 +777,7 @@ pub(crate) fn parse_anthropic_content( preliminary: None, dynamic: None, provider_metadata: None, - }); + })); } // MCP tool result — dynamic, and it inherits the name and metadata // of the `mcp_tool_use` block it answers (upstream :1184-1194). @@ -788,7 +787,7 @@ pub(crate) fn parse_anthropic_content( is_error, } => { let call = mcp_tool_calls.get(tool_use_id.as_str()); - content.push(GenerateContent::ToolResult { + content.push(GenerateContent::ToolResult(ToolResult { tool_call_id: tool_use_id.clone(), tool_name: call .map(|(name, _)| (*name).to_string()) @@ -802,7 +801,7 @@ pub(crate) fn parse_anthropic_content( "anthropic": { "type": "mcp-tool-use", "serverName": server } }) }), - }); + })); } _ => {} } @@ -1088,22 +1087,19 @@ pub(crate) async fn anthropic_stream_core( ContentBlock::McpToolUse { id, name, input, server_name } => { mcp_tool_calls .insert(id.clone(), (name.clone(), server_name.clone())); - yield Ok(StreamPart::ToolCall { + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: id.clone(), tool_name: name.clone(), - input: Value::String(input.to_string()), + input: input.to_string(), provider_executed: Some(true), dynamic: Some(true), - thought_signature: None, - invalid: None, - error: None, provider_metadata: Some(json!({ "anthropic": { "type": "mcp-tool-use", "serverName": server_name, } })), - }); + })); } // Redacted thinking — emit as ReasoningStart. ContentBlock::RedactedThinking { data } => { @@ -1299,23 +1295,19 @@ pub(crate) async fn anthropic_stream_core( id: id.clone(), provider_metadata: None, }); - let input = - Value::String(finalize_streamed_tool_input( + let input = finalize_streamed_tool_input( accumulated_json, provider_tool_name.as_deref(), provider_tool_input_type.as_deref(), - )); - yield Ok(StreamPart::ToolCall { + ); + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: id, tool_name: name, input, provider_executed, dynamic, - thought_signature: None, - invalid: None, - error: None, provider_metadata, - }); + })); } } } diff --git a/aimux-providers/src/anthropic/tool_name_mapping.rs b/aimux-providers/src/anthropic/tool_name_mapping.rs index adb611521..9bf67eaa0 100644 --- a/aimux-providers/src/anthropic/tool_name_mapping.rs +++ b/aimux-providers/src/anthropic/tool_name_mapping.rs @@ -115,26 +115,9 @@ impl ToolNameMapping { #[cfg(test)] mod tests { use super::*; - use aimux_core::tool::{FunctionTool, ProviderTool}; + use aimux_core::tool::FunctionTool; use serde_json::json; - fn provider_tool(id: &str, name: &str) -> Tool { - Tool::Provider(ProviderTool { - id: id.to_string(), - name: name.to_string(), - args: json!({}), - }) - } - - #[test] - fn maps_renamed_provider_tool_both_ways() { - let tools = vec![provider_tool("anthropic.web_search_20250305", "mySearch")]; - let mapping = ToolNameMapping::new(Some(&tools)); - - assert_eq!(mapping.to_custom_tool_name("web_search"), "mySearch"); - assert_eq!(mapping.to_provider_tool_name("mySearch"), "web_search"); - } - #[test] fn unmapped_names_pass_through() { let mapping = ToolNameMapping::new(None); @@ -149,76 +132,4 @@ mod tests { // A function tool named `web_search` must not hijack the provider name. assert_eq!(mapping.to_custom_tool_name("web_search"), "web_search"); } - - #[test] - fn all_code_execution_versions_share_one_provider_name() { - for id in [ - "anthropic.code_execution_20250522", - "anthropic.code_execution_20250825", - "anthropic.code_execution_20260120", - ] { - let tools = vec![provider_tool(id, "runCode")]; - let mapping = ToolNameMapping::new(Some(&tools)); - assert_eq!( - mapping.to_custom_tool_name("code_execution"), - "runCode", - "{id} must map to the code_execution provider name" - ); - } - } - - #[test] - fn marks_implicit_code_execution_dynamic_only_for_2026_web_tools() { - let web_only = vec![provider_tool("anthropic.web_search_20260209", "search")]; - assert!(ToolNameMapping::new(Some(&web_only)).mark_code_execution_dynamic()); - - let web_and_code = vec![ - provider_tool("anthropic.web_search_20260209", "search"), - provider_tool("anthropic.code_execution_20260120", "runCode"), - ]; - assert!(!ToolNameMapping::new(Some(&web_and_code)).mark_code_execution_dynamic()); - - let old_web = vec![provider_tool("anthropic.web_search_20250305", "search")]; - assert!(!ToolNameMapping::new(Some(&old_web)).mark_code_execution_dynamic()); - } - - /// Every provider tool id this crate can build a request for must also be - /// mappable, and to the same name the request body uses. The two tables are - /// written out separately, so without this they drift silently: a renamed - /// tool would round-trip under the wrong name. - #[test] - fn every_known_provider_tool_id_maps_to_its_request_name() { - use crate::anthropic::prepare_tools::prepare_provider_tool; - use std::collections::BTreeSet; - - // Ids are recovered from the source rather than duplicated here, so a - // newly supported tool cannot be added without this test seeing it. - let src = include_str!("prepare_tools.rs"); - let ids: BTreeSet<&str> = src - .match_indices("\"anthropic.") - .filter_map(|(i, _)| { - let rest = &src[i + 1..]; - rest.find('"').map(|end| &rest[..end]) - }) - .collect(); - assert!( - ids.len() > 10, - "expected to recover the id table, got {ids:?}" - ); - - for id in ids { - let mut betas = BTreeSet::new(); - let Some(def) = prepare_provider_tool(id, &json!({}), &mut betas) else { - continue; - }; - let request_name = def["name"].as_str().expect("provider tool has a name"); - let tools = vec![provider_tool(id, "myCustomName")]; - let mapping = ToolNameMapping::new(Some(&tools)); - assert_eq!( - mapping.to_custom_tool_name(request_name), - "myCustomName", - "{id} is sent as `{request_name}` but is not in the name mapping" - ); - } - } } diff --git a/aimux-providers/src/bedrock/convert.rs b/aimux-providers/src/bedrock/convert.rs index 9b7d8dc5c..2b5842dc4 100644 --- a/aimux-providers/src/bedrock/convert.rs +++ b/aimux-providers/src/bedrock/convert.rs @@ -453,7 +453,7 @@ pub fn supports_strict_tools(model_id: &str) -> bool { #[must_use] pub fn prepare_tools( tools: &Option>, - tool_choice: &ToolChoice, + tool_choice: Option<&ToolChoice>, model_id: &str, ) -> Value { let non_empty = tools.as_ref().filter(|t| !t.is_empty()); @@ -462,13 +462,15 @@ pub fn prepare_tools( }; // `toolChoice: none` clears tools entirely (matches TS). - if matches!(tool_choice, ToolChoice::None) { + if matches!(tool_choice, Some(ToolChoice::None)) { return json!({}); } // `toolChoice: tool` filters function tools to the named one. let filtered: Vec<&FunctionTool> = match tool_choice { - ToolChoice::Tool { tool_name } => tools.iter().filter(|t| &t.name == tool_name).collect(), + Some(ToolChoice::Tool { tool_name }) => { + tools.iter().filter(|t| &t.name == tool_name).collect() + } _ => tools.iter().collect(), }; @@ -498,10 +500,11 @@ pub fn prepare_tools( } let tool_choice_val = match tool_choice { - ToolChoice::Auto => json!({ "auto": {} }), - ToolChoice::Required => json!({ "any": {} }), - ToolChoice::Tool { tool_name } => json!({ "tool": { "name": tool_name } }), - ToolChoice::None => unreachable!(), + Some(ToolChoice::Auto) => json!({ "auto": {} }), + Some(ToolChoice::Required) => json!({ "any": {} }), + Some(ToolChoice::Tool { tool_name }) => json!({ "tool": { "name": tool_name } }), + Some(ToolChoice::None) => unreachable!(), + None => return json!({ "tools": tool_specs }), }; json!({ "tools": tool_specs, "toolChoice": tool_choice_val }) @@ -603,7 +606,7 @@ pub fn build_request_body(model_id: &str, options: &CallOptions) -> Value { }) .collect() }); - let tool_config = prepare_tools(&function_tools, &options.tool_choice, model_id); + let tool_config = prepare_tools(&function_tools, options.tool_choice.as_ref(), model_id); if tool_config .as_object() .map(|o| !o.is_empty()) diff --git a/aimux-providers/src/bedrock/model.rs b/aimux-providers/src/bedrock/model.rs index b1c0c96ec..8d515761c 100644 --- a/aimux-providers/src/bedrock/model.rs +++ b/aimux-providers/src/bedrock/model.rs @@ -5,6 +5,7 @@ //! across all Bedrock-backed models (Anthropic Claude, Meta Llama, Mistral, //! etc.). +use aimux_core::tool::RawToolCall; use std::collections::HashMap; use async_trait::async_trait; @@ -480,22 +481,19 @@ impl LanguageModel for BedrockModel { yield Ok(StreamPart::ToolInputEnd { id: id.clone(), provider_metadata: None}); // Empty input normalizes to "{}" per the upstream // provider. - let input = serde_json::Value::String(if acc.is_empty() { + let input = if acc.is_empty() { "{}".to_string() } else { acc - }); - yield Ok(StreamPart::ToolCall { + }; + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: id, tool_name: name, input, provider_executed: None, dynamic: None, - thought_signature: None, - invalid: None, - error: None, provider_metadata: None, - }); + })); } else if reasoning_id.is_some() { let id = idx.to_string(); if reasoning_id.as_deref() == Some(id.as_str()) { @@ -649,15 +647,14 @@ fn extract_content(block: &BedrockContentBlock, content: &mut Vec GooglePrompt { convert_to_google_messages_for_namespace(prompt, ProviderMetadataNamespace::Google) @@ -261,7 +260,6 @@ fn convert_assistant_parts( tool_call_id, tool_name, input, - thought_signature, provider_options, .. } => { @@ -272,11 +270,9 @@ fn convert_assistant_parts( let server_tool_type = google_options .and_then(|options| options.get("serverToolType")) .and_then(|value| value.as_str()); - let signature = thought_signature.as_deref().or_else(|| { - google_options - .and_then(|options| options.get("thoughtSignature")) - .and_then(|value| value.as_str()) - }); + let signature = google_options + .and_then(|options| options.get("thoughtSignature")) + .and_then(|value| value.as_str()); let mut part_value = if let (Some(server_id), Some(server_type)) = (server_tool_call_id, server_tool_type) @@ -421,7 +417,10 @@ pub struct PreparedTools { /// tools (`google_search`, `code_execution`, …) are out of scope for the Rust /// port — only `FunctionTool`s are supported. #[must_use] -pub fn prepare_tools(tools: &Option>, tool_choice: &ToolChoice) -> PreparedTools { +pub fn prepare_tools( + tools: &Option>, + tool_choice: Option<&ToolChoice>, +) -> PreparedTools { // Coerce empty arrays to None (matches TS `tools?.length ? tools : undefined`). let non_empty = tools.as_ref().filter(|&t| !t.is_empty()); @@ -454,22 +453,23 @@ pub fn prepare_tools(tools: &Option>, tool_choice: &ToolChoice ]); let tool_config = match tool_choice { - ToolChoice::Auto => { + None => has_strict.then(|| json!({ "functionCallingConfig": { "mode": "VALIDATED" } })), + Some(ToolChoice::Auto) => { if has_strict { Some(json!({ "functionCallingConfig": { "mode": "VALIDATED" } })) } else { Some(json!({ "functionCallingConfig": { "mode": "AUTO" } })) } } - ToolChoice::None => Some(json!({ "functionCallingConfig": { "mode": "NONE" } })), - ToolChoice::Required => { + Some(ToolChoice::None) => Some(json!({ "functionCallingConfig": { "mode": "NONE" } })), + Some(ToolChoice::Required) => { if has_strict { Some(json!({ "functionCallingConfig": { "mode": "VALIDATED" } })) } else { Some(json!({ "functionCallingConfig": { "mode": "ANY" } })) } } - ToolChoice::Tool { tool_name } => { + Some(ToolChoice::Tool { tool_name }) => { if has_strict { Some(json!({ "functionCallingConfig": { @@ -612,7 +612,7 @@ pub struct PreparedToolsWithWarnings { #[must_use] pub fn prepare_all_tools( tools: &Option>, - tool_choice: &ToolChoice, + tool_choice: Option<&ToolChoice>, model_id: &str, ) -> PreparedToolsWithWarnings { let caps = get_google_model_capabilities(model_id); @@ -656,23 +656,23 @@ pub fn prepare_all_tools( let mut combined_config = Map::new(); let fc = match tool_choice { - ToolChoice::None => { + Some(ToolChoice::None) => { let mut m = Map::new(); m.insert("mode".to_string(), json!("NONE")); m } - ToolChoice::Required => { + Some(ToolChoice::Required) => { let mut m = Map::new(); m.insert("mode".to_string(), json!("ANY")); m } - ToolChoice::Tool { tool_name } => { + Some(ToolChoice::Tool { tool_name }) => { let mut m = Map::new(); m.insert("mode".to_string(), json!("ANY")); m.insert("allowedFunctionNames".to_string(), json!([tool_name])); m } - ToolChoice::Auto => { + None | Some(ToolChoice::Auto) => { let mut m = Map::new(); m.insert("mode".to_string(), json!("VALIDATED")); m @@ -1099,7 +1099,7 @@ fn build_request_body_with_warnings_for_namespace( ); } - let prepared = prepare_all_tools(&options.tools, &options.tool_choice, model_id); + let prepared = prepare_all_tools(&options.tools, options.tool_choice.as_ref(), model_id); if let Some(tools) = prepared.tools { body.insert("tools".to_string(), Value::Array(tools)); } @@ -1322,7 +1322,6 @@ mod tests { tool_name: "runCode".to_string(), input: json!({ "language": "PYTHON", "code": "print(2)" }), provider_executed: Some(true), - thought_signature: None, provider_options: Some(code_metadata("googleVertex", call_id)), }, ContentPart::ToolResult { @@ -1364,7 +1363,6 @@ mod tests { tool_name: "weather".to_string(), input: json!({}), provider_executed: None, - thought_signature: None, provider_options: Some(json!({ "googleVertex": { "thoughtSignature": "google-vertex" }, "vertex": { "thoughtSignature": "vertex" }, @@ -1392,7 +1390,6 @@ mod tests { tool_name: "weather".to_string(), input: json!({}), provider_executed: None, - thought_signature: None, provider_options: Some(json!({ "google": { "thoughtSignature": "gateway-signature" }, })), diff --git a/aimux-providers/src/google/model.rs b/aimux-providers/src/google/model.rs index aaa923a58..8100748c4 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 aimux_core::tool::RawToolCall; +use aimux_core::tool::ToolResult; use std::collections::HashMap; use async_trait::async_trait; @@ -412,10 +414,6 @@ impl LanguageModel for GoogleModel { .unwrap_or_else(|| format!("call-{block_counter}")); block_counter += 1; let args = fc.get("args").cloned().unwrap_or(json!({})); - let thought_signature = part - .get("thoughtSignature") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); yield Ok(StreamPart::ToolInputStart { id: id.clone(), @@ -423,26 +421,23 @@ impl LanguageModel for GoogleModel { provider_executed: None, dynamic: None, title: None, - provider_metadata: None, + provider_metadata: thought_sig_meta.clone(), }); let args_str = args.to_string(); yield Ok(StreamPart::ToolInputDelta { id: id.clone(), delta: args_str, - provider_metadata: None, + provider_metadata: thought_sig_meta.clone(), }); - yield Ok(StreamPart::ToolInputEnd { id: id.clone(), provider_metadata: None}); - yield Ok(StreamPart::ToolCall { + 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: Value::String(args.to_string()), + input: args.to_string(), provider_executed: None, dynamic: None, - thought_signature, - invalid: None, - error: None, provider_metadata: thought_sig_meta.clone(), - }); + })); has_tool_calls = true; } else if let Some(ec) = part.get("executableCode") { // Provider-executed code execution. @@ -455,21 +450,18 @@ impl LanguageModel for GoogleModel { let id = format!("call-{block_counter}"); block_counter += 1; last_code_execution_tool_call_id = Some(id.clone()); - yield Ok(StreamPart::ToolCall { + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: id.clone(), tool_name: code_execution_tool_name.clone(), - input: Value::String(ec.to_string()), + input: ec.to_string(), provider_executed: Some(true), dynamic: None, - thought_signature: None, - invalid: None, - error: None, provider_metadata: Some(server_tool_metadata( &id, "code_execution", None, )), - }); + })); // provider-executed → does NOT set has_tool_calls } } else if let Some(cer) = part.get("codeExecutionResult") { @@ -487,7 +479,7 @@ impl LanguageModel for GoogleModel { .and_then(|v| v.as_str()) .map(std::string::ToString::to_string) .unwrap_or_default(); - yield Ok(StreamPart::ToolResult { + yield Ok(StreamPart::ToolResult(ToolResult { tool_call_id: call_id.clone(), tool_name: code_execution_tool_name.clone(), result: json!({ "outcome": outcome, "output": output }), @@ -499,7 +491,7 @@ impl LanguageModel for GoogleModel { "code_execution", None, )), - }); + })); } } else if let Some(tc) = part.get("toolCall") { // Server-side tool call (provider-executed). @@ -521,17 +513,14 @@ impl LanguageModel for GoogleModel { if let Some(s) = part.get("thoughtSignature").and_then(|v| v.as_str()) { server_meta["google"]["thoughtSignature"] = json!(s); } - yield Ok(StreamPart::ToolCall { + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: id, tool_name: format!("server:{tool_type}"), - input: Value::String(args.to_string()), + input: args.to_string(), provider_executed: Some(true), dynamic: Some(true), - thought_signature: None, - invalid: None, - error: None, 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. @@ -553,7 +542,7 @@ impl LanguageModel for GoogleModel { if let Some(s) = part.get("thoughtSignature").and_then(|v| v.as_str()) { server_meta["google"]["thoughtSignature"] = json!(s); } - yield Ok(StreamPart::ToolResult { + yield Ok(StreamPart::ToolResult(ToolResult { tool_call_id: id, tool_name: format!("server:{tool_type}"), result: response, @@ -561,7 +550,7 @@ impl LanguageModel for GoogleModel { 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. @@ -735,15 +724,14 @@ fn extract_content_from_candidate( if has_code { let id = format!("call-{}", content.len()); last_code_execution_tool_call_id = Some(id.clone()); - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: id.clone(), tool_name: code_execution_tool_name.to_string(), input: ec.to_string(), provider_executed: Some(true), dynamic: None, - thought_signature: None, provider_metadata: Some(server_tool_metadata(&id, "code_execution", None)), - }); + })); } } else if let Some(cer) = part.get("codeExecutionResult") { // One executableCode may be followed by multiple results. @@ -754,7 +742,7 @@ fn extract_content_from_candidate( .and_then(|v| v.as_str()) .map(std::string::ToString::to_string) .unwrap_or_default(); - content.push(GenerateContent::ToolResult { + content.push(GenerateContent::ToolResult(ToolResult { tool_call_id: call_id.clone(), tool_name: code_execution_tool_name.to_string(), result: json!({ "outcome": outcome, "output": output }), @@ -766,7 +754,7 @@ fn extract_content_from_candidate( "code_execution", None, )), - }); + })); } } else if let Some(text) = part.get("text").and_then(|v| v.as_str()) { if text.is_empty() { @@ -804,19 +792,14 @@ fn extract_content_from_candidate( .unwrap_or("") .to_string(); let input = fc.get("args").cloned().unwrap_or(json!({})); - let thought_signature = part - .get("thoughtSignature") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string); - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: id, tool_name: name, input: input.to_string(), provider_executed: None, dynamic: None, - thought_signature, provider_metadata: thought_sig_meta.clone(), - }); + })); has_tool_calls = true; } else if let Some(inline) = part.get("inlineData") { // File output (e.g. gemini-2.5-flash-image) — upstream :478-490. @@ -844,25 +827,20 @@ fn extract_content_from_candidate( .to_string(); last_server_tool_call_id = Some(id.clone()); let input = 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 mut server_meta = json!({ "google": { "serverToolCallId": id, "serverToolType": tool_type } }); if let Some(s) = part.get("thoughtSignature").and_then(|v| v.as_str()) { server_meta["google"]["thoughtSignature"] = json!(s); } - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: id, tool_name: format!("server:{tool_type}"), input: input.to_string(), provider_executed: Some(true), dynamic: Some(true), - thought_signature, 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 (upstream :512-533). @@ -882,7 +860,7 @@ fn extract_content_from_candidate( if let Some(s) = part.get("thoughtSignature").and_then(|v| v.as_str()) { server_meta["google"]["thoughtSignature"] = json!(s); } - content.push(GenerateContent::ToolResult { + content.push(GenerateContent::ToolResult(ToolResult { tool_call_id: id, tool_name: format!("server:{tool_type}"), result: response, @@ -890,7 +868,7 @@ fn extract_content_from_candidate( preliminary: None, dynamic: None, provider_metadata: Some(server_meta), - }); + })); last_server_tool_call_id = None; } } @@ -913,18 +891,18 @@ fn set_provider_metadata(item: &mut GenerateContent, meta: ProviderMetadata) { | GenerateContent::Reasoning { provider_metadata, .. } - | GenerateContent::ToolCall { + | GenerateContent::ToolCall(RawToolCall { provider_metadata, .. - } + }) | GenerateContent::File { provider_metadata, .. } | GenerateContent::Source { provider_metadata, .. } - | GenerateContent::ToolResult { + | GenerateContent::ToolResult(ToolResult { provider_metadata, .. - } => { + }) => { *provider_metadata = Some(meta); } } diff --git a/aimux-providers/src/huggingface/responses.rs b/aimux-providers/src/huggingface/responses.rs index 2c8ccd8c9..3ff7391f6 100644 --- a/aimux-providers/src/huggingface/responses.rs +++ b/aimux-providers/src/huggingface/responses.rs @@ -14,6 +14,8 @@ //! - `map-huggingface-responses-finish-reason.ts` — finish-reason mapping //! - `huggingface-responses-language-model-options.ts` — provider options +use aimux_core::tool::RawToolCall; +use aimux_core::tool::ToolResult; use std::collections::HashMap; use async_trait::async_trait; @@ -419,28 +421,25 @@ impl LanguageModel for HuggingFaceResponsesModel { .get("arguments") .and_then(|v| v.as_str()) .unwrap_or("{}"); - let input = Value::String(arguments.to_string()); + let input = arguments.to_string(); yield Ok(StreamPart::ToolInputEnd { id: call_id.clone(), provider_metadata: None, }); - yield Ok(StreamPart::ToolCall { + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: call_id.clone(), tool_name: name.clone(), input, provider_executed: None, dynamic: None, - thought_signature: None, - invalid: None, - error: None, provider_metadata: None, - }); + })); if let Some(output) = item.get("output").and_then(|v| v.as_str()) { - yield Ok(StreamPart::ToolResult { + yield Ok(StreamPart::ToolResult(ToolResult { tool_call_id: call_id, tool_name: name.clone(), result: Value::String(output.to_string()), @@ -448,7 +447,7 @@ impl LanguageModel for HuggingFaceResponsesModel { preliminary: None, dynamic: None, provider_metadata: None, - }); + })); } } _ => {} @@ -643,7 +642,7 @@ pub fn build_request_body_with_warnings( // Prepare tools. let (tools, tool_choice, tool_warnings) = - prepare_responses_tools(&options.tools, &options.tool_choice); + prepare_responses_tools(&options.tools, options.tool_choice.as_ref()); warnings.extend(tool_warnings); // Assemble the body. Key insertion order follows the TS `baseArgs` object @@ -1018,7 +1017,7 @@ fn detect_media_type_from_base64(data: &str, top_level: &str) -> Option #[must_use] pub fn prepare_responses_tools( tools: &Option>, - tool_choice: &ToolChoice, + tool_choice: Option<&ToolChoice>, ) -> (Option>, Option, Vec) { let tools = match tools { Some(t) if !t.is_empty() => t, @@ -1050,10 +1049,10 @@ pub fn prepare_responses_tools( } let mapped_tool_choice = match tool_choice { - ToolChoice::Auto => Some(json!("auto")), - ToolChoice::Required => Some(json!("required")), - ToolChoice::None => None, // not supported, ignore - ToolChoice::Tool { tool_name } => Some(json!({ + Some(ToolChoice::Auto) => Some(json!("auto")), + Some(ToolChoice::Required) => Some(json!("required")), + None | Some(ToolChoice::None) => None, // not supported, ignore + Some(ToolChoice::Tool { tool_name }) => Some(json!({ "type": "function", "function": { "name": tool_name } })), @@ -1140,15 +1139,14 @@ fn build_generate_content(response: &Value) -> Result, AiMu .and_then(|v| v.as_str()) .unwrap_or("{}"); let input = arguments.to_string(); - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: call_id.to_string(), tool_name: name.to_string(), input, provider_executed: None, dynamic: None, - thought_signature: None, provider_metadata: None, - }); + })); } "mcp_call" => { @@ -1159,18 +1157,17 @@ fn build_generate_content(response: &Value) -> Result, AiMu .and_then(|v| v.as_str()) .unwrap_or("{}"); let input = arguments.to_string(); - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: id.to_string(), tool_name: name.to_string(), input, provider_executed: Some(true), dynamic: Some(true), - thought_signature: None, provider_metadata: None, - }); + })); if let Some(output) = part.get("output").filter(|v| !v.is_null()) { - content.push(GenerateContent::ToolResult { + content.push(GenerateContent::ToolResult(ToolResult { tool_call_id: id.to_string(), tool_name: name.to_string(), result: output.clone(), @@ -1178,25 +1175,24 @@ fn build_generate_content(response: &Value) -> Result, AiMu preliminary: None, dynamic: Some(true), provider_metadata: None, - }); + })); } } "mcp_list_tools" => { let id = part.get("id").and_then(|v| v.as_str()).unwrap_or(""); let server_label = part.get("server_label").cloned().unwrap_or(Value::Null); - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: id.to_string(), tool_name: "list_tools".to_string(), input: json!({ "server_label": server_label }).to_string(), provider_executed: Some(true), dynamic: Some(true), - thought_signature: None, provider_metadata: None, - }); + })); if let Some(tools) = part.get("tools").filter(|v| !v.is_null()) { - content.push(GenerateContent::ToolResult { + content.push(GenerateContent::ToolResult(ToolResult { tool_call_id: id.to_string(), tool_name: "list_tools".to_string(), result: json!({ "tools": tools }), @@ -1204,7 +1200,7 @@ fn build_generate_content(response: &Value) -> Result, AiMu preliminary: None, dynamic: Some(true), provider_metadata: None, - }); + })); } } diff --git a/aimux-providers/src/mistral/convert.rs b/aimux-providers/src/mistral/convert.rs index 88d737c86..a224a99c9 100644 --- a/aimux-providers/src/mistral/convert.rs +++ b/aimux-providers/src/mistral/convert.rs @@ -356,7 +356,7 @@ pub fn build_request_body(model_id: &str, options: &CallOptions, stream: bool) - }) .collect() }); - let prepared = prepare_tools(&function_tools, Some(&options.tool_choice)); + let prepared = prepare_tools(&function_tools, options.tool_choice.as_ref()); if let Some(tools) = prepared.tools { body["tools"] = json!(tools); if let Some(tc) = prepared.tool_choice { diff --git a/aimux-providers/src/mistral/model.rs b/aimux-providers/src/mistral/model.rs index 52d63e874..ddcc02018 100644 --- a/aimux-providers/src/mistral/model.rs +++ b/aimux-providers/src/mistral/model.rs @@ -10,6 +10,7 @@ //! - Usage supports `num_cached_tokens` / `prompt_tokens_details.cached_tokens`. //! - Finish reasons include `model_length`. +use aimux_core::tool::RawToolCall; use std::collections::HashMap; use async_trait::async_trait; @@ -275,15 +276,14 @@ impl LanguageModel for MistralModel { if let Some(tool_calls) = choice.message.tool_calls { for tc in tool_calls { let input = tc.function.arguments; - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: tc.id, tool_name: tc.function.name, input, provider_executed: None, dynamic: None, - thought_signature: None, provider_metadata: None, - }); + })); } } @@ -520,18 +520,15 @@ impl LanguageModel for MistralModel { provider_metadata: None, }); - let input = Value::String(args); - yield Ok(StreamPart::ToolCall { + let input = args; + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: tool_id, tool_name, input, provider_executed: None, dynamic: None, - thought_signature: None, - invalid: None, - error: None, provider_metadata: None, - }); + })); } } diff --git a/aimux-providers/src/open_responses.rs b/aimux-providers/src/open_responses.rs index 340237e52..27cc52115 100644 --- a/aimux-providers/src/open_responses.rs +++ b/aimux-providers/src/open_responses.rs @@ -6,6 +6,7 @@ //! //! Translation of `reference/ai/packages/open-responses/src/responses/`. +use aimux_core::tool::RawToolCall; use std::collections::HashMap; use std::sync::Arc; @@ -396,15 +397,14 @@ impl LanguageModel for OpenResponsesModel { .and_then(|a| a.as_str()) .unwrap_or("{}"); let input = arguments.to_string(); - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: call_id, tool_name: name, input, provider_executed: None, dynamic: None, - thought_signature: None, provider_metadata: None, - }); + })); } _ => {} } @@ -681,18 +681,15 @@ impl LanguageModel for OpenResponsesModel { .map(std::string::ToString::to_string) }) .unwrap_or_default(); - let input = Value::String(arguments); - yield Ok(StreamPart::ToolCall { + let input = arguments; + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id, tool_name, input, provider_executed: None, dynamic: None, - thought_signature: None, - invalid: None, - error: None, provider_metadata: None, - }); + })); has_tool_calls = true; } "reasoning" => { @@ -931,15 +928,12 @@ fn build_request_body( }) .unwrap_or_default(); - // Convert tool choice. - // Only emit when not the default (Auto) — Rust's ToolChoice::Auto is the - // default and indistinguishable from "not set", matching the TS behavior - // where undefined toolChoice is omitted from the body. let converted_tool_choice: Option = match &options.tool_choice { - ToolChoice::Auto => None, - ToolChoice::None => Some(json!("none")), - ToolChoice::Required => Some(json!("required")), - ToolChoice::Tool { tool_name } => Some(json!({ + None => None, + Some(ToolChoice::Auto) => Some(json!("auto")), + Some(ToolChoice::None) => Some(json!("none")), + Some(ToolChoice::Required) => Some(json!("required")), + Some(ToolChoice::Tool { tool_name }) => Some(json!({ "type": "function", "name": tool_name, })), diff --git a/aimux-providers/src/openai/convert.rs b/aimux-providers/src/openai/convert.rs index 4c018dba1..aab1998e6 100644 --- a/aimux-providers/src/openai/convert.rs +++ b/aimux-providers/src/openai/convert.rs @@ -1360,11 +1360,11 @@ fn apply_tools( prepare_tools_groq( &function_tools, options.tools.as_ref(), - Some(&options.tool_choice), + options.tool_choice.as_ref(), model_id, ) } else { - prepare_tools(&function_tools, Some(&options.tool_choice)) + prepare_tools(&function_tools, options.tool_choice.as_ref()) }; if let Some(tools) = prepared.tools { diff --git a/aimux-providers/src/openai/model.rs b/aimux-providers/src/openai/model.rs index b4267c866..87af24559 100644 --- a/aimux-providers/src/openai/model.rs +++ b/aimux-providers/src/openai/model.rs @@ -7,6 +7,7 @@ //! speak the OpenAI chat-completions wire format (notably Azure OpenAI) reuse //! the conversion + streaming logic while supplying their own URL and auth. +use aimux_core::tool::RawToolCall; use std::collections::HashMap; use async_trait::async_trait; @@ -317,15 +318,14 @@ pub async fn execute_generate( if let Some(tool_calls) = choice.message.tool_calls { for tc in tool_calls { let input = tc.function.arguments; - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: tc.id, tool_name: tc.function.name, input, provider_executed: None, dynamic: None, - thought_signature: None, provider_metadata: None, - }); + })); } } // Parse annotations (URL citations) → Source content items. diff --git a/aimux-providers/src/openai/responses/convert.rs b/aimux-providers/src/openai/responses/convert.rs index fa7c99f7d..4f0aad772 100644 --- a/aimux-providers/src/openai/responses/convert.rs +++ b/aimux-providers/src/openai/responses/convert.rs @@ -922,7 +922,7 @@ pub fn build_responses_request_body( ); // -- Tools -- - let prepared = prepare_responses_tools(&options.tools, Some(&options.tool_choice)); + let prepared = prepare_responses_tools(&options.tools, options.tool_choice.as_ref()); if let Some(tools) = prepared.tools { body["tools"] = json!(tools); if let Some(tc) = prepared.tool_choice { diff --git a/aimux-providers/src/openai/responses/responses_convert.rs b/aimux-providers/src/openai/responses/responses_convert.rs index ea51c2478..cbd5dcb17 100644 --- a/aimux-providers/src/openai/responses/responses_convert.rs +++ b/aimux-providers/src/openai/responses/responses_convert.rs @@ -13,6 +13,7 @@ //! genuinely different streaming loops are **not** force-merged into one //! function — only the shared framework is extracted. +use aimux_core::tool::RawToolCall; use std::collections::HashMap; use std::pin::Pin; @@ -190,19 +191,18 @@ pub fn build_responses_generate_result( .unwrap_or("{}") .to_string(); let input = arguments; - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: call_id, tool_name: name, input, provider_executed: None, dynamic: None, - thought_signature: None, provider_metadata: Some(json!({ (provider_key.clone()): { "itemId": part.get("id").cloned().unwrap_or(Value::Null), } })), - }); + })); } Some("custom_tool_call") => { has_function_call = true; @@ -219,19 +219,18 @@ pub fn build_responses_generate_result( let input_str = part.get("input").and_then(|v| v.as_str()).unwrap_or("{}"); let input = serde_json::to_string(input_str) .expect("serializing a custom-tool input string cannot fail"); - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: call_id, tool_name: name, input, provider_executed: None, dynamic: None, - thought_signature: None, provider_metadata: Some(json!({ (provider_key.clone()): { "itemId": part.get("id").cloned().unwrap_or(Value::Null), } })), - }); + })); } Some("reasoning") => { let summary = part.get("summary").and_then(|v| v.as_array()); @@ -793,18 +792,15 @@ where id: call_id.clone(), provider_metadata: None, }); - let input = Value::String(arguments); - yield Ok(StreamPart::ToolCall { + let input = arguments; + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: call_id, tool_name: name, input, provider_executed: None, dynamic: None, - thought_signature: None, - invalid: None, - error: None, provider_metadata: None, - }); + })); } "custom_tool_call" => { has_function_call = true; @@ -827,22 +823,18 @@ where id: call_id.clone(), provider_metadata: None, }); - let input = Value::String( + let input = serde_json::to_string(input_str).expect( "serializing a custom-tool input string cannot fail", - ), - ); - yield Ok(StreamPart::ToolCall { + ); + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: call_id, tool_name: name, input, provider_executed: None, dynamic: None, - thought_signature: None, - invalid: None, - error: None, provider_metadata: None, - }); + })); } "reasoning" => { let id = item diff --git a/aimux-providers/src/vertex/anthropic_model.rs b/aimux-providers/src/vertex/anthropic_model.rs index f707441e8..4856cf9dd 100644 --- a/aimux-providers/src/vertex/anthropic_model.rs +++ b/aimux-providers/src/vertex/anthropic_model.rs @@ -14,6 +14,7 @@ //! //! Reference: +use aimux_core::tool::RawToolCall; use std::collections::HashMap; use async_trait::async_trait; @@ -386,22 +387,19 @@ impl LanguageModel for VertexAnthropicModel { ContentBlock::McpToolUse { id, name, input, server_name } => { mcp_tool_calls .insert(id.clone(), (name.clone(), server_name.clone())); - yield Ok(StreamPart::ToolCall { + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: id.clone(), tool_name: name.clone(), - input: Value::String(input.to_string()), + input: input.to_string(), provider_executed: Some(true), dynamic: Some(true), - thought_signature: None, - invalid: None, - error: None, provider_metadata: Some(json!({ "anthropic": { "type": "mcp-tool-use", "serverName": server_name, } })), - }); + })); } ContentBlock::RedactedThinking { data } => { yield Ok(StreamPart::ReasoningStart { @@ -535,23 +533,19 @@ impl LanguageModel for VertexAnthropicModel { .. } => { yield Ok(StreamPart::ToolInputEnd { id: id.clone(), provider_metadata: None}); - let input = - Value::String(finalize_streamed_tool_input( + let input = finalize_streamed_tool_input( accumulated_json, provider_tool_name.as_deref(), provider_tool_input_type.as_deref(), - )); - yield Ok(StreamPart::ToolCall { + ); + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: id, tool_name: name, input, provider_executed, dynamic, - thought_signature: None, - invalid: None, - error: None, provider_metadata, - }); + })); } } } diff --git a/aimux-providers/src/vertex/model.rs b/aimux-providers/src/vertex/model.rs index e151adfdb..4b2366946 100644 --- a/aimux-providers/src/vertex/model.rs +++ b/aimux-providers/src/vertex/model.rs @@ -4,6 +4,8 @@ //! [`crate::google::types`] response types. Only the endpoint construction and //! authentication differ from the public Gemini API provider. +use aimux_core::tool::RawToolCall; +use aimux_core::tool::ToolResult; use std::collections::HashMap; use async_trait::async_trait; @@ -433,17 +435,14 @@ impl LanguageModel for VertexModel { id: id.clone(), provider_metadata: tool_metadata.clone(), }); - yield Ok(StreamPart::ToolCall { + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: id, tool_name: name.to_string(), - input: Value::String(args.to_string()), + input: args.to_string(), provider_executed: None, dynamic: None, - thought_signature, - invalid: None, - error: None, provider_metadata: tool_metadata, - }); + })); has_tool_calls = true; } else if let Some(ec) = part.get("executableCode") { // Provider-executed code execution. @@ -456,21 +455,18 @@ impl LanguageModel for VertexModel { let id = format!("call-{block_counter}"); block_counter += 1; last_code_execution_tool_call_id = Some(id.clone()); - yield Ok(StreamPart::ToolCall { + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: id.clone(), tool_name: code_execution_tool_name.clone(), - input: Value::String(ec.to_string()), + input: ec.to_string(), provider_executed: Some(true), dynamic: None, - thought_signature: None, - invalid: None, - error: 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") { @@ -488,7 +484,7 @@ impl LanguageModel for VertexModel { .and_then(|v| v.as_str()) .map(std::string::ToString::to_string) .unwrap_or_default(); - yield Ok(StreamPart::ToolResult { + yield Ok(StreamPart::ToolResult(ToolResult { tool_call_id: call_id.clone(), tool_name: code_execution_tool_name.clone(), result: json!({ "outcome": outcome, "output": output }), @@ -500,7 +496,7 @@ impl LanguageModel for VertexModel { "code_execution", None, )), - }); + })); } } else if let Some(tc) = part.get("toolCall") { // Server-side tool call (provider-executed). @@ -525,17 +521,14 @@ impl LanguageModel for VertexModel { tool_type, thought_signature.as_deref(), ); - yield Ok(StreamPart::ToolCall { + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: id, tool_name: format!("server:{tool_type}"), - input: Value::String(args.to_string()), + input: args.to_string(), provider_executed: Some(true), dynamic: Some(true), - thought_signature, - invalid: None, - error: None, 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. @@ -559,7 +552,7 @@ impl LanguageModel for VertexModel { tool_type, part.get("thoughtSignature").and_then(|v| v.as_str()), ); - yield Ok(StreamPart::ToolResult { + yield Ok(StreamPart::ToolResult(ToolResult { tool_call_id: id, tool_name: format!("server:{tool_type}"), result: response, @@ -567,7 +560,7 @@ impl LanguageModel for VertexModel { preliminary: None, dynamic: None, provider_metadata: Some(server_meta), - }); + })); } } } @@ -695,19 +688,18 @@ fn extract_content_from_candidate( if has_code { let id = format!("call-{}", content.len()); last_code_execution_tool_call_id = Some(id.clone()); - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: id.clone(), tool_name: code_execution_tool_name.to_string(), input: ec.to_string(), provider_executed: Some(true), dynamic: None, - thought_signature: None, provider_metadata: Some(vertex_server_tool_metadata( &id, "code_execution", None, )), - }); + })); } } else if let Some(cer) = part.get("codeExecutionResult") { // One executableCode may be followed by multiple results. @@ -717,7 +709,7 @@ fn extract_content_from_candidate( .get("output") .and_then(|v| v.as_str()) .unwrap_or_default(); - content.push(GenerateContent::ToolResult { + content.push(GenerateContent::ToolResult(ToolResult { tool_call_id: call_id.clone(), tool_name: code_execution_tool_name.to_string(), result: json!({ "outcome": outcome, "output": output }), @@ -729,7 +721,7 @@ fn extract_content_from_candidate( "code_execution", None, )), - }); + })); } } else if let Some(text) = part.get("text").and_then(|v| v.as_str()) { if !text.is_empty() { @@ -763,15 +755,14 @@ fn extract_content_from_candidate( let provider_metadata = thought_signature.as_deref().map(|signature| { vertex_provider_metadata(json!({ "thoughtSignature": signature })) }); - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: id, tool_name: name, input: input.to_string(), provider_executed: None, dynamic: None, - thought_signature, provider_metadata, - }); + })); has_tool_calls = true; } else if let Some(tc) = part.get("toolCall") { let tool_type = tc.get("toolType").and_then(|v| v.as_str()).unwrap_or(""); @@ -788,15 +779,14 @@ fn extract_content_from_candidate( .map(std::string::ToString::to_string); let server_meta = vertex_server_tool_metadata(&id, tool_type, thought_signature.as_deref()); - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: id, tool_name: format!("server:{tool_type}"), input: input.to_string(), provider_executed: Some(true), dynamic: Some(true), - thought_signature, provider_metadata: Some(server_meta), - }); + })); } else if let Some(tr) = part.get("toolResponse") { let tool_type = tr.get("toolType").and_then(|v| v.as_str()).unwrap_or(""); let id = last_server_tool_call_id @@ -813,7 +803,7 @@ fn extract_content_from_candidate( tool_type, part.get("thoughtSignature").and_then(|v| v.as_str()), ); - content.push(GenerateContent::ToolResult { + content.push(GenerateContent::ToolResult(ToolResult { tool_call_id: id, tool_name: format!("server:{tool_type}"), result: response, @@ -821,7 +811,7 @@ fn extract_content_from_candidate( preliminary: None, dynamic: None, provider_metadata: Some(server_meta), - }); + })); } } } diff --git a/aimux-providers/src/xai/convert.rs b/aimux-providers/src/xai/convert.rs index 3662539d6..bfd304ed9 100644 --- a/aimux-providers/src/xai/convert.rs +++ b/aimux-providers/src/xai/convert.rs @@ -492,7 +492,7 @@ pub fn build_request_body_with_warnings( let (messages, message_warnings) = convert_to_xai_messages(&options.prompt)?; warnings.extend(message_warnings); - let prepared = prepare_tools(&options.tools, Some(&options.tool_choice)); + let prepared = prepare_tools(&options.tools, options.tool_choice.as_ref()); for tw in &prepared.tool_warnings { warnings.push(tw.clone()); } diff --git a/aimux-providers/src/xai/model.rs b/aimux-providers/src/xai/model.rs index 99cce1322..a2948521d 100644 --- a/aimux-providers/src/xai/model.rs +++ b/aimux-providers/src/xai/model.rs @@ -6,6 +6,7 @@ //! cached tokens, reasoning-effort model gating, 200-status errors) to warrant //! its own model implementation rather than reusing `OpenAIModel`. +use aimux_core::tool::RawToolCall; use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; @@ -171,15 +172,14 @@ impl LanguageModel for XaiModel { if let Some(tool_calls) = choice.message.tool_calls { for tc in tool_calls { let input = tc.function.arguments; - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: tc.id, tool_name: tc.function.name, input, provider_executed: None, dynamic: None, - thought_signature: None, provider_metadata: None, - }); + })); } } @@ -503,18 +503,15 @@ impl LanguageModel for XaiModel { id: id.clone(), provider_metadata: None, }); - let input = Value::String(args.clone()); - yield Ok(StreamPart::ToolCall { + let input = args.clone(); + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: id.clone(), tool_name: name.clone(), input, provider_executed: None, dynamic: None, - thought_signature: None, - invalid: None, - error: None, provider_metadata: None, - }); + })); } } tool_calls.clear(); @@ -557,18 +554,15 @@ impl LanguageModel for XaiModel { id: id.clone(), provider_metadata: None, }); - let input = Value::String(args.clone()); - yield Ok(StreamPart::ToolCall { + let input = args.clone(); + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: id.clone(), tool_name: name.clone(), input, provider_executed: None, dynamic: None, - thought_signature: None, - invalid: None, - error: None, provider_metadata: None, - }); + })); } } diff --git a/aimux-providers/src/xai/responses/convert.rs b/aimux-providers/src/xai/responses/convert.rs index 3226bf6e5..48c2a54b2 100644 --- a/aimux-providers/src/xai/responses/convert.rs +++ b/aimux-providers/src/xai/responses/convert.rs @@ -599,7 +599,7 @@ pub fn build_responses_request_body( let (input, input_warnings) = convert_to_xai_responses_input(&options.prompt)?; warnings.extend(input_warnings); - let prepared = prepare_responses_tools(&options.tools, Some(&options.tool_choice)); + let prepared = prepare_responses_tools(&options.tools, options.tool_choice.as_ref()); for tw in &prepared.tool_warnings { warnings.push(tw.clone()); } diff --git a/aimux-providers/src/xai/responses/mod.rs b/aimux-providers/src/xai/responses/mod.rs index e8beb3682..1076ff719 100644 --- a/aimux-providers/src/xai/responses/mod.rs +++ b/aimux-providers/src/xai/responses/mod.rs @@ -15,6 +15,8 @@ pub mod convert; pub mod types; +use aimux_core::tool::RawToolCall; +use aimux_core::tool::ToolResult; use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; @@ -132,15 +134,14 @@ impl LanguageModel for XaiResponsesModel { let tool_name = resolve_tool_name(part_type, None, &provider_tool_names); let part_id = part.get("id").and_then(|v| v.as_str()).unwrap_or(""); - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: part_id.to_string(), tool_name: tool_name.clone(), input: String::new(), provider_executed: Some(true), dynamic: None, - thought_signature: None, provider_metadata: None, - }); + })); let queries = part .get("queries") @@ -160,7 +161,7 @@ impl LanguageModel for XaiResponsesModel { }) .unwrap_or(Value::Null); - content.push(GenerateContent::ToolResult { + content.push(GenerateContent::ToolResult(ToolResult { tool_call_id: part_id.to_string(), tool_name, result: json!({ "queries": queries, "results": results }), @@ -168,7 +169,7 @@ impl LanguageModel for XaiResponsesModel { preliminary: None, dynamic: None, provider_metadata: None, - }); + })); continue; } @@ -189,15 +190,14 @@ impl LanguageModel for XaiResponsesModel { let tool_input = get_tool_input(part_type, part); let part_id = part.get("id").and_then(|v| v.as_str()).unwrap_or(""); - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: part_id.to_string(), tool_name, input: tool_input, provider_executed: Some(true), dynamic: None, - thought_signature: None, provider_metadata: None, - }); + })); continue; } @@ -246,15 +246,14 @@ impl LanguageModel for XaiResponsesModel { let name = part.get("name").and_then(|v| v.as_str()).unwrap_or(""); let arguments = part.get("arguments").and_then(|v| v.as_str()).unwrap_or(""); let input = arguments.to_string(); - content.push(GenerateContent::ToolCall { + content.push(GenerateContent::ToolCall(RawToolCall { tool_call_id: call_id.to_string(), tool_name: name.to_string(), input, provider_executed: None, dynamic: None, - thought_signature: None, provider_metadata: None, - }); + })); } "reasoning" => { let summary = part.get("summary").and_then(|v| v.as_array()); @@ -709,17 +708,14 @@ impl LanguageModel for XaiResponsesModel { id: part_id.to_string(), provider_metadata: None, }); - yield Ok(StreamPart::ToolCall { + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: part_id.to_string(), tool_name: tool_name.clone(), - input: Value::String(String::new()), + input: String::new(), provider_executed: Some(true), dynamic: None, - thought_signature: None, - invalid: None, - error: None, provider_metadata: None, - }); + })); } if event_type == "response.output_item.done" { @@ -732,7 +728,7 @@ impl LanguageModel for XaiResponsesModel { "text": r.get("text").cloned().unwrap_or(Value::Null), })).collect::>()) }).unwrap_or(Value::Null); - yield Ok(StreamPart::ToolResult { + yield Ok(StreamPart::ToolResult(ToolResult { tool_call_id: part_id.to_string(), tool_name, result: json!({ "queries": queries, "results": results }), @@ -740,7 +736,7 @@ impl LanguageModel for XaiResponsesModel { preliminary: None, dynamic: None, provider_metadata: None, - }); + })); } continue; } @@ -782,21 +778,18 @@ impl LanguageModel for XaiResponsesModel { id: part_id.to_string(), provider_metadata: None, }); - yield Ok(StreamPart::ToolCall { + yield Ok(StreamPart::ToolCall(RawToolCall { tool_call_id: part_id.to_string(), tool_name: tool_name.clone(), - input: Value::String(tool_input), + input: tool_input, provider_executed: Some(true), dynamic: None, - thought_signature: None, - invalid: None, - error: None, provider_metadata: None, - }); + })); } if event_type == "response.output_item.done" { - yield Ok(StreamPart::ToolResult { + yield Ok(StreamPart::ToolResult(ToolResult { tool_call_id: part_id.to_string(), tool_name, result: json!({}), @@ -804,7 +797,7 @@ impl LanguageModel for XaiResponsesModel { preliminary: None, dynamic: None, provider_metadata: None, - }); + })); } continue; } @@ -874,18 +867,15 @@ impl LanguageModel for XaiResponsesModel { id: call_id.to_string(), provider_metadata: None, }); - let input = Value::String(arguments.to_string()); - yield Ok(StreamPart::ToolCall { + 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, - thought_signature: None, - invalid: None, - error: None, provider_metadata: None, - }); + })); } continue; } diff --git a/aimux-providers/tests/alibaba_test.rs b/aimux-providers/tests/alibaba_test.rs index a4323ca6a..db2c913ff 100644 --- a/aimux-providers/tests/alibaba_test.rs +++ b/aimux-providers/tests/alibaba_test.rs @@ -24,6 +24,7 @@ //! DashScope's OpenAI-compatible mode. Those TS-specific behaviours are //! therefore not asserted here. +use aimux_core::tool::RawToolCall; use futures::StreamExt; use serde_json::{Value, json}; use serial_test::serial; @@ -305,7 +306,6 @@ async fn converts_tool_call_and_tool_result_messages() { tool_name: "get_weather".to_string(), input: json!({"location": "SF"}), provider_executed: None, - thought_signature: None, provider_options: None, }], ..Default::default() @@ -524,15 +524,15 @@ async fn do_generate_extracts_tool_call() { assert_eq!(result.content.len(), 1); match &result.content[0] { - GenerateContent::ToolCall { + GenerateContent::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. - } => { + }) => { assert_eq!(tool_call_id, "call_abc"); assert_eq!(tool_name, "get-weather"); - assert_eq!(input, &Value::String(r#"{"city":"SF"}"#.into())); + assert_eq!(input, r#"{"city":"SF"}"#); } other => panic!("expected ToolCall, got {other:?}"), } diff --git a/aimux-providers/tests/bedrock_convert_test.rs b/aimux-providers/tests/bedrock_convert_test.rs index fc7945103..092f5680a 100644 --- a/aimux-providers/tests/bedrock_convert_test.rs +++ b/aimux-providers/tests/bedrock_convert_test.rs @@ -1,4 +1,4 @@ -//! Pure-function tests for the Amazon Bedrock provider's conversion layer. +//! Pure-function tests for the Amazon Bedrock provider's conversion layer. //! //! Translated from the TS test files: //! - `convert-to-amazon-bedrock-chat-messages.test.ts` (64 cases) @@ -36,7 +36,7 @@ use aimux_core::options::ToolChoice; use aimux_core::tool::FunctionTool; use aimux_providers::bedrock::convert::{ - BedrockUsage, convert_prompt_to_bedrock, convert_usage, prepare_tools, supports_strict_tools, + BedrockUsage, convert_prompt_to_bedrock, convert_usage, prepare_tools, }; // ── prompt builders ───────────────────────────────────────────────────────── @@ -110,9 +110,6 @@ fn tool_result(id: &str, output: Value) -> ContentPart { ContentPart::tool_result(id.to_string(), output) } -const ANTHROPIC_MODEL: &str = "anthropic.claude-sonnet-4-5-20250929-v1:0"; -const NON_ANTHROPIC_MODEL: &str = "meta.llama3-70b-instruct-v1:0"; - // ════════════════════════════════════════════════════════════════════════════ // convert-to-amazon-bedrock-chat-messages // ════════════════════════════════════════════════════════════════════════════ @@ -950,6 +947,8 @@ fn media_type_route_to_document_text_plain() { // amazon-bedrock-prepare-tools // ════════════════════════════════════════════════════════════════════════════ +const NON_ANTHROPIC_MODEL: &str = "meta.llama3-70b-instruct-v1:0"; + fn func_tool(name: &str, description: Option<&str>, input_schema: Value) -> FunctionTool { FunctionTool { name: name.to_string(), @@ -961,108 +960,11 @@ fn func_tool(name: &str, description: Option<&str>, input_schema: Value) -> Func } } -fn func_tool_strict( - name: &str, - description: Option<&str>, - input_schema: Value, - strict: Option, -) -> FunctionTool { - FunctionTool { - name: name.to_string(), - description: description.map(std::string::ToString::to_string), - input_schema, - strict, - provider_options: None, - input_examples: None, - } -} - -/// TS: "should return empty toolConfig when tools are undefined" -#[test] -fn prepare_tools_empty_when_undefined() { - let config = prepare_tools(&None, &ToolChoice::Auto, ANTHROPIC_MODEL); - assert_eq!(config, json!({})); -} - -/// TS: "should return empty toolConfig when tools are empty" -#[test] -fn prepare_tools_empty_when_empty() { - let config = prepare_tools(&Some(vec![]), &ToolChoice::Auto, ANTHROPIC_MODEL); - assert_eq!(config, json!({})); -} - -/// TS: "should correctly prepare function tools" -#[test] -fn prepare_tools_function_tools() { - let tools = Some(vec![func_tool( - "testFunction", - Some("A test function"), - json!({ "type": "object", "properties": {} }), - )]); - let config = prepare_tools(&tools, &ToolChoice::Auto, ANTHROPIC_MODEL); - assert_eq!( - config["tools"], - json!([{ - "toolSpec": { - "name": "testFunction", - "description": "A test function", - "inputSchema": { "json": { "type": "object", "properties": {} } } - } - }]) - ); -} - -/// TS: "should exclude description when it is empty string" -#[test] -fn prepare_tools_exclude_description_empty() { - let tools = Some(vec![func_tool( - "testFunction", - Some(""), - json!({ "type": "object", "properties": {} }), - )]); - let config = prepare_tools(&tools, &ToolChoice::Auto, ANTHROPIC_MODEL); - let spec = &config["tools"][0]["toolSpec"]; - assert!(spec.get("description").is_none()); -} - -/// TS: "should exclude description when it is whitespace-only" -#[test] -fn prepare_tools_exclude_description_whitespace() { - let tools = Some(vec![func_tool( - "testFunction", - Some(" "), - json!({ "type": "object", "properties": {} }), - )]); - let config = prepare_tools(&tools, &ToolChoice::Auto, ANTHROPIC_MODEL); - let spec = &config["tools"][0]["toolSpec"]; - assert!(spec.get("description").is_none()); -} - -/// TS: "should include description when it has content" -#[test] -fn prepare_tools_include_description() { - let tools = Some(vec![func_tool( - "testFunction", - Some("Valid description"), - json!({ "type": "object", "properties": {} }), - )]); - let config = prepare_tools(&tools, &ToolChoice::Auto, ANTHROPIC_MODEL); - assert_eq!( - config["tools"][0]["toolSpec"]["description"], - json!("Valid description") - ); -} - -// SKIPPED (TS: "should warn for provider-defined tools on non-anthropic models", -// "should warn and filter out web_search_20250305 tool", "should return empty -// toolConfig when all tools are filtered out"): the Rust FunctionTool has no -// type/id, so provider-defined tools are not modelled. - /// TS: "should handle tool choice 'auto'" #[test] fn prepare_tools_tool_choice_auto() { let tools = Some(vec![func_tool("testFunction", Some("Test"), json!({}))]); - let config = prepare_tools(&tools, &ToolChoice::Auto, NON_ANTHROPIC_MODEL); + let config = prepare_tools(&tools, Some(&ToolChoice::Auto), NON_ANTHROPIC_MODEL); assert_eq!(config["toolChoice"], json!({ "auto": {} })); } @@ -1070,7 +972,7 @@ fn prepare_tools_tool_choice_auto() { #[test] fn prepare_tools_tool_choice_required() { let tools = Some(vec![func_tool("testFunction", Some("Test"), json!({}))]); - let config = prepare_tools(&tools, &ToolChoice::Required, NON_ANTHROPIC_MODEL); + let config = prepare_tools(&tools, Some(&ToolChoice::Required), NON_ANTHROPIC_MODEL); assert_eq!(config["toolChoice"], json!({ "any": {} })); } @@ -1078,7 +980,7 @@ fn prepare_tools_tool_choice_required() { #[test] fn prepare_tools_tool_choice_none_clears() { let tools = Some(vec![func_tool("testFunction", Some("Test"), json!({}))]); - let config = prepare_tools(&tools, &ToolChoice::None, NON_ANTHROPIC_MODEL); + let config = prepare_tools(&tools, Some(&ToolChoice::None), NON_ANTHROPIC_MODEL); assert_eq!(config, json!({})); } @@ -1088,9 +990,9 @@ fn prepare_tools_tool_choice_tool() { let tools = Some(vec![func_tool("testFunction", Some("Test"), json!({}))]); let config = prepare_tools( &tools, - &ToolChoice::Tool { + Some(&ToolChoice::Tool { tool_name: "testFunction".to_string(), - }, + }), NON_ANTHROPIC_MODEL, ); assert_eq!( @@ -1112,176 +1014,15 @@ fn prepare_tools_tool_choice_filters_to_named() { ]); let config = prepare_tools( &tools, - &ToolChoice::Tool { + Some(&ToolChoice::Tool { tool_name: "getWeather".to_string(), - }, + }), NON_ANTHROPIC_MODEL, ); assert_eq!(config["tools"].as_array().unwrap().len(), 1); assert_eq!(config["tools"][0]["toolSpec"]["name"], json!("getWeather")); } -/// TS: "should pass through strict mode when strict is true" -#[test] -fn prepare_tools_strict_true() { - let tools = Some(vec![func_tool_strict( - "testFunction", - Some("A test function"), - json!({ "type": "object", "properties": {} }), - Some(true), - )]); - let config = prepare_tools(&tools, &ToolChoice::Auto, ANTHROPIC_MODEL); - assert_eq!( - config["tools"], - json!([{ - "toolSpec": { - "name": "testFunction", - "description": "A test function", - "strict": true, - "inputSchema": { "json": { "type": "object", "properties": {} } } - } - }]) - ); -} - -/// TS: "should pass through strict mode when strict is false" -#[test] -fn prepare_tools_strict_false() { - let tools = Some(vec![func_tool_strict( - "testFunction", - Some("A test function"), - json!({ "type": "object", "properties": {} }), - Some(false), - )]); - let config = prepare_tools(&tools, &ToolChoice::Auto, ANTHROPIC_MODEL); - assert_eq!( - config["tools"], - json!([{ - "toolSpec": { - "name": "testFunction", - "description": "A test function", - "strict": false, - "inputSchema": { "json": { "type": "object", "properties": {} } } - } - }]) - ); -} - -/// TS: "should not include strict when strict is undefined" -#[test] -fn prepare_tools_strict_undefined() { - let tools = Some(vec![func_tool_strict( - "testFunction", - Some("A test function"), - json!({ "type": "object", "properties": {} }), - None, - )]); - let config = prepare_tools(&tools, &ToolChoice::Auto, ANTHROPIC_MODEL); - let spec = &config["tools"][0]["toolSpec"]; - assert!(spec.get("strict").is_none()); -} - -/// TS: "should pass through strict mode for multiple tools with different strict settings" -#[test] -fn prepare_tools_strict_multiple() { - let tools = Some(vec![ - func_tool_strict( - "strictTool", - Some("A strict tool"), - json!({ "type": "object", "properties": {} }), - Some(true), - ), - func_tool_strict( - "nonStrictTool", - Some("A non-strict tool"), - json!({ "type": "object", "properties": {} }), - Some(false), - ), - func_tool_strict( - "defaultTool", - Some("A tool without strict setting"), - json!({ "type": "object", "properties": {} }), - None, - ), - ]); - let config = prepare_tools(&tools, &ToolChoice::Auto, ANTHROPIC_MODEL); - let tools_arr = config["tools"].as_array().unwrap(); - assert_eq!(tools_arr.len(), 3); - assert_eq!(tools_arr[0]["toolSpec"]["strict"], json!(true)); - assert_eq!(tools_arr[1]["toolSpec"]["strict"], json!(false)); - assert!(tools_arr[2]["toolSpec"].get("strict").is_none()); -} - -/// TS: it.each "should omit strict for %s" (claude-opus-5, claude-sonnet-5, -/// claude-fable-5). These models reject newer schema fields, so `strict` is -/// omitted even when set. -#[test] -fn prepare_tools_omit_strict_for_opus_5() { - assert!(!supports_strict_tools("us.anthropic.claude-opus-5")); - let tools = Some(vec![func_tool_strict( - "testFunction", - Some("A test function"), - json!({ "type": "object", "properties": {} }), - Some(true), - )]); - let config = prepare_tools(&tools, &ToolChoice::Auto, "us.anthropic.claude-opus-5"); - assert!(config["tools"][0]["toolSpec"].get("strict").is_none()); -} - -#[test] -fn prepare_tools_omit_strict_for_sonnet_5() { - assert!(!supports_strict_tools("anthropic.claude-sonnet-5")); - let tools = Some(vec![func_tool_strict( - "testFunction", - Some("A test function"), - json!({ "type": "object", "properties": {} }), - Some(true), - )]); - let config = prepare_tools(&tools, &ToolChoice::Auto, "anthropic.claude-sonnet-5"); - assert!(config["tools"][0]["toolSpec"].get("strict").is_none()); -} - -#[test] -fn prepare_tools_omit_strict_for_fable_5() { - assert!(!supports_strict_tools("eu.anthropic.claude-fable-5")); - let tools = Some(vec![func_tool_strict( - "testFunction", - Some("A test function"), - json!({ "type": "object", "properties": {} }), - Some(true), - )]); - let config = prepare_tools(&tools, &ToolChoice::Auto, "eu.anthropic.claude-fable-5"); - assert!(config["tools"][0]["toolSpec"].get("strict").is_none()); -} - -/// TS: "should omit strict for claude-opus-4-7" -#[test] -fn prepare_tools_omit_strict_for_opus_4_7() { - assert!(!supports_strict_tools("us.anthropic.claude-opus-4-7")); - let tools = Some(vec![func_tool_strict( - "testFunction", - Some("A test function"), - json!({ "type": "object", "properties": {} }), - Some(true), - )]); - let config = prepare_tools(&tools, &ToolChoice::Auto, "us.anthropic.claude-opus-4-7"); - assert!(config["tools"][0]["toolSpec"].get("strict").is_none()); -} - -/// TS: "should omit strict for claude-opus-4-8" -#[test] -fn prepare_tools_omit_strict_for_opus_4_8() { - assert!(!supports_strict_tools("anthropic.claude-opus-4-8")); - let tools = Some(vec![func_tool_strict( - "testFunction", - Some("A test function"), - json!({ "type": "object", "properties": {} }), - Some(true), - )]); - let config = prepare_tools(&tools, &ToolChoice::Auto, "anthropic.claude-opus-4-8"); - assert!(config["tools"][0]["toolSpec"].get("strict").is_none()); -} - // ════════════════════════════════════════════════════════════════════════════ // convert-amazon-bedrock-usage // ════════════════════════════════════════════════════════════════════════════ diff --git a/aimux-providers/tests/codex_test.rs b/aimux-providers/tests/codex_test.rs index c0974477f..cd957ac24 100644 --- a/aimux-providers/tests/codex_test.rs +++ b/aimux-providers/tests/codex_test.rs @@ -18,11 +18,9 @@ use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; use aimux_core::message::Role; -use aimux_core::options::{CallOptions, Tool, ToolChoice}; +use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, StreamResult}; use aimux_core::stream_part::StreamPart; -use aimux_core::tool::FunctionTool; -use aimux_core::types::FinishReasonUnified; use aimux_providers::{CODEX_OAUTH_TOKEN_URL, CodexConfig, CodexProvider, codex_refresh_at}; @@ -40,22 +38,6 @@ fn default_options(prompt: LanguageModelPrompt) -> CallOptions { CallOptions::new(prompt) } -fn weather_tool() -> FunctionTool { - FunctionTool { - name: "weather".to_string(), - description: None, - input_schema: json!({ - "type": "object", - "properties": { "location": { "type": "string" } }, - "required": ["location"], - "additionalProperties": false, - }), - strict: None, - provider_options: None, - input_examples: None, - } -} - fn sse_event(json_str: &str) -> String { format!("data: {json_str}\n\n") } @@ -161,17 +143,6 @@ const TEXT_STREAM_EVENTS: &[&str] = &[ r#"{"type":"response.completed","response":{"id":"resp_1","created_at":1741269019,"model":"gpt-5.2-codex","incomplete_details":null,"usage":{"input_tokens":543,"input_tokens_details":{"cached_tokens":234},"output_tokens":478,"output_tokens_details":{"reasoning_tokens":123}}}}"#, ]; -/// The tool-call streaming event sequence. -const TOOL_STREAM_EVENTS: &[&str] = &[ - r#"{"type":"response.created","response":{"id":"resp_tc","created_at":1741362087,"model":"gpt-5.2-codex"}}"#, - r#"{"type":"response.output_item.added","output_index":0,"item":{"type":"function_call","id":"fc_1","call_id":"call_added","name":"weather","arguments":"","status":"completed"}}"#, - r#"{"type":"response.function_call_arguments.delta","item_id":"fc_1","output_index":0,"delta":"{\"location\":"}"#, - r#"{"type":"response.function_call_arguments.delta","item_id":"fc_1","output_index":0,"delta":"\"Rome\"}"}"#, - r#"{"type":"response.function_call_arguments.done","item_id":"fc_1","output_index":0,"arguments":"{\"location\":\"Rome\"}"}"#, - r#"{"type":"response.output_item.done","output_index":0,"item":{"type":"function_call","id":"fc_1","call_id":"call_done","name":"weather","arguments":"{\"location\":\"Rome\"}","status":"completed"}}"#, - r#"{"type":"response.completed","response":{"id":"resp_tc","created_at":1741362087,"model":"gpt-5.2-codex","incomplete_details":null,"usage":{"input_tokens":0,"input_tokens_details":{"cached_tokens":0},"output_tokens":0,"output_tokens_details":{"reasoning_tokens":0}}}}"#, -]; - /// A `response.completed` event whose nested response object carries the full /// output —what the subscription channel emits at the end of a stream. fn subscription_completed_event() -> String { @@ -249,44 +220,6 @@ async fn api_key_streams_text() { } } -#[tokio::test] -async fn api_key_streams_tool_calls() { - let server = MockServer::start().await; - mock_sse_response(&server, &sse_body(TOOL_STREAM_EVENTS)).await; - - let config = CodexConfig::new("test-key").with_base_url(server.uri()); - let model = CodexProvider::new(config).model("gpt-5.2-codex"); - - let options = CallOptions { - tools: Some(vec![Tool::from(weather_tool())]), - tool_choice: ToolChoice::Auto, - ..CallOptions::new(test_prompt()) - }; - - let result = model.do_stream(&options).await.expect("should succeed"); - let parts = collect_stream(result).await; - - match &parts[6] { - StreamPart::ToolCall { - tool_call_id, - tool_name, - input, - .. - } => { - assert_eq!(tool_call_id, "call_done"); - assert_eq!(tool_name, "weather"); - assert_eq!(input, &Value::String(r#"{"location":"Rome"}"#.into())); - } - other => panic!("expected ToolCall, got {other:?}"), - } - match &parts[7] { - StreamPart::Finish { finish_reason, .. } => { - assert_eq!(finish_reason.unified, FinishReasonUnified::ToolCalls); - } - other => panic!("expected Finish, got {other:?}"), - } -} - // Path B: subscription mode #[tokio::test] diff --git a/aimux-providers/tests/conformance_test.rs b/aimux-providers/tests/conformance_test.rs index 098f2fc04..87388af58 100644 --- a/aimux-providers/tests/conformance_test.rs +++ b/aimux-providers/tests/conformance_test.rs @@ -67,7 +67,7 @@ fn has_text(content: &[GenerateContent]) -> bool { fn has_tool_call(content: &[GenerateContent]) -> bool { content .iter() - .any(|c| matches!(c, GenerateContent::ToolCall { .. })) + .any(|c| matches!(c, GenerateContent::ToolCall(_))) } fn has_reasoning(content: &[GenerateContent]) -> bool { @@ -89,9 +89,7 @@ fn has_finish(parts: &[StreamPart]) -> bool { #[allow(dead_code)] fn has_tool_call_part(parts: &[StreamPart]) -> bool { - parts - .iter() - .any(|p| matches!(p, StreamPart::ToolCall { .. })) + parts.iter().any(|p| matches!(p, StreamPart::ToolCall(_))) } /// True when an error signals the request never matched a cassette — i.e. a diff --git a/aimux-providers/tests/data_loss_regression_test.rs b/aimux-providers/tests/data_loss_regression_test.rs index 3810d3c05..458c2a4d4 100644 --- a/aimux-providers/tests/data_loss_regression_test.rs +++ b/aimux-providers/tests/data_loss_regression_test.rs @@ -22,6 +22,7 @@ //! against re-reads of the cassette (a re-read would pass even if both the //! producer and the assertion regressed together). +use aimux_core::tool::{RawToolCall, ToolResult}; use std::collections::BTreeMap; use std::fs; use std::path::PathBuf; @@ -133,6 +134,7 @@ fn opts() -> CallOptions { fn opts_with_tools(tools: Vec) -> CallOptions { let mut o = opts(); o.tools = Some(tools); + o.tool_choice = Some(aimux_core::tool::ToolChoice::Auto); o } @@ -140,7 +142,7 @@ fn provider_tool(id: &str, name: &str) -> Tool { Tool::Provider(ProviderTool { id: id.to_string(), name: name.to_string(), - args: json!({}), + args: serde_json::Map::new(), }) } @@ -196,14 +198,13 @@ fn reasonings(content: &[GenerateContent]) -> Vec<(&str, Option<&Value>)> { .collect() } -/// `(tool_call_id, tool_name, input, provider_executed, dynamic, thought_signature, metadata)` +/// `(tool_call_id, tool_name, input, provider_executed, dynamic, metadata)` type ToolCallView<'a> = ( &'a str, &'a str, Value, Option, Option, - Option<&'a str>, Option<&'a Value>, ); @@ -211,15 +212,14 @@ fn tool_calls(content: &[GenerateContent]) -> Vec> { content .iter() .filter_map(|c| match c { - GenerateContent::ToolCall { + GenerateContent::ToolCall(RawToolCall { tool_call_id, tool_name, input, provider_executed, dynamic, - thought_signature, provider_metadata, - } => { + }) => { // Provider results intentionally carry the exact wire string; // parse only in this assertion helper so nested data-loss // checks remain readable without weakening that boundary. @@ -231,7 +231,6 @@ fn tool_calls(content: &[GenerateContent]) -> Vec> { parsed_input, *provider_executed, *dynamic, - thought_signature.as_deref(), provider_metadata.as_ref(), )) } @@ -254,7 +253,7 @@ fn tool_results(content: &[GenerateContent]) -> Vec> { content .iter() .filter_map(|c| match c { - GenerateContent::ToolResult { + GenerateContent::ToolResult(ToolResult { tool_call_id, tool_name, result, @@ -262,7 +261,7 @@ fn tool_results(content: &[GenerateContent]) -> Vec> { dynamic, provider_metadata, .. - } => Some(( + }) => Some(( tool_call_id.as_str(), tool_name.as_str(), result, @@ -546,7 +545,7 @@ async fn finding_5_gemini_thought_signature_reaches_provider_metadata() { } /// `thoughtSignature` on a `functionCall` part must reach -/// `ToolCall.thought_signature` — that is the field the follow-up turn echoes. +/// `providerMetadata.google.thoughtSignature` for use in follow-up provider options. #[tokio::test] async fn finding_5_gemini_function_call_thought_signature_round_trips() { let c = cassette("gemini/two_arg_rewrites_chain_blocking_0.json"); @@ -564,12 +563,7 @@ async fn finding_5_gemini_function_call_thought_signature_round_trips() { assert_eq!(calls[0].1, "add"); assert_eq!(calls[0].2, json!({ "x": 1, "y": 1 })); assert_eq!( - calls[0].5, - Some("signature_REDACTED_1"), - "the functionCall thoughtSignature must land on ToolCall.thought_signature" - ); - assert_eq!( - calls[0].6.expect("providerMetadata")["google"]["thoughtSignature"], + calls[0].5.expect("providerMetadata")["google"]["thoughtSignature"], json!("signature_REDACTED_1") ); } @@ -1033,7 +1027,7 @@ async fn finding_2_anthropic_mcp_tool_use_and_result_are_dynamic_and_named() { // mcp_tool_use → provider-executed + dynamic + serverName metadata. let calls = tool_calls(&result.content); assert_eq!(calls.len(), 1); - let (id, name, ref input, provider_executed, dynamic, _, meta) = calls[0]; + let (id, name, ref input, provider_executed, dynamic, meta) = calls[0]; assert_eq!(id, "mcptoolu_01SAss3KEwASziHZoMR6HcZU"); assert_eq!(name, "ask_question"); assert_eq!(input["repoName"], json!("pydantic/pydantic-ai")); diff --git a/aimux-providers/tests/deepseek_chat_test.rs b/aimux-providers/tests/deepseek_chat_test.rs index 36729ff0b..79c399d9c 100644 --- a/aimux-providers/tests/deepseek_chat_test.rs +++ b/aimux-providers/tests/deepseek_chat_test.rs @@ -26,6 +26,7 @@ //! overridden with the mock server's root URI (no `/v1` suffix), so the //! resulting request path is `/chat/completions`. +use aimux_core::tool::RawToolCall; use std::collections::HashMap; use futures::StreamExt; @@ -409,18 +410,15 @@ async fn should_extract_tool_call_content() { // The tool-call fixture has empty content and one tool call. assert_eq!(result.content.len(), 1); match &result.content[0] { - GenerateContent::ToolCall { + GenerateContent::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. - } => { + }) => { assert_eq!(tool_call_id, "call_00_9V0vrf86Pc9aelHCJMZqnJBo"); assert_eq!(tool_name, "weather"); - assert_eq!( - input, - &Value::String(r#"{"location": "San Francisco"}"#.into()) - ); + assert_eq!(input, r#"{"location": "San Francisco"}"#); } other => panic!("expected ToolCall, got {other:?}"), } @@ -745,22 +743,19 @@ async fn should_stream_tool_call() { let tool_calls: Vec<_> = parts .iter() .filter_map(|p| match p { - StreamPart::ToolCall { + StreamPart::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. - } => Some((tool_call_id, tool_name, input)), + }) => Some((tool_call_id, tool_name, input)), _ => None, }) .collect(); assert_eq!(tool_calls.len(), 1); assert_eq!(tool_calls[0].0, "call_1"); assert_eq!(tool_calls[0].1, "weather"); - assert_eq!( - tool_calls[0].2, - &Value::String(r#"{"location": "San Francisco"}"#.into()) - ); + assert_eq!(tool_calls[0].2, r#"{"location": "San Francisco"}"#); // The last part should be Finish with ToolCalls. match parts.last() { diff --git a/aimux-providers/tests/e2e_test.rs b/aimux-providers/tests/e2e_test.rs index 3f9e6e566..4cd5a9746 100644 --- a/aimux-providers/tests/e2e_test.rs +++ b/aimux-providers/tests/e2e_test.rs @@ -840,7 +840,7 @@ async fn e2e_openai_stream_tool_calls() { let has_tool_part = parts.iter().any(|p| { matches!( p, - StreamPart::ToolCall { .. } + StreamPart::ToolCall(_) | StreamPart::ToolInputStart { .. } | StreamPart::ToolInputDelta { .. } | StreamPart::ToolInputEnd { .. } @@ -857,9 +857,9 @@ async fn e2e_openai_stream_tool_calls() { // If a complete ToolCall part exists, verify its fields let tool_call = parts.iter().find_map(|p| match p { - StreamPart::ToolCall { + StreamPart::ToolCall(aimux_core::tool::ToolCall { tool_name, input, .. - } => Some((tool_name, input)), + }) => Some((tool_name, input)), _ => None, }); if let Some((tool_name, input)) = tool_call { diff --git a/aimux-providers/tests/groq_test.rs b/aimux-providers/tests/groq_test.rs index 276f3294b..3cc114448 100644 --- a/aimux-providers/tests/groq_test.rs +++ b/aimux-providers/tests/groq_test.rs @@ -18,6 +18,7 @@ //! browser_search tool, stream_options, etc.). Reasoning effort is a direct //! passthrough — no vendor normalization. +use aimux_core::tool::RawToolCall; use futures::StreamExt; use serde_json::{Value, json}; use wiremock::matchers::{method, path}; @@ -702,34 +703,6 @@ mod prepare_tools { assert_eq!(body["tools"][0]["function"]["parameters"]["type"], "object"); } - /// TS: "should add warnings for unsupported provider-defined tools" - #[tokio::test] - async fn unsupported_provider_tool_warning() { - let server = MockServer::start().await; - mock_json(&server, text_completion_body()).await; - - let model = make_provider(&server); - - let tool = Tool::Provider(aimux_core::tool::ProviderTool { - id: "some.unsupported_tool".to_string(), - name: "unsupported_tool".to_string(), - args: json!({}), - }); - let options = CallOptions { - tools: Some(vec![tool]), - ..default_options(test_prompt()) - }; - let result = model.do_generate(&options).await.unwrap(); - - // The unsupported tool should produce a warning. - assert!(result.warnings.iter().any(|w| match w { - aimux_core::types::Warning::Unsupported { feature, .. } => { - feature.contains("some.unsupported_tool") - } - _ => false, - })); - } - /// TS: "should pass through strict mode when strict is true" #[tokio::test] async fn strict_mode_true() { @@ -813,15 +786,13 @@ mod prepare_tools { let tool = FunctionTool::new("testFunction", json!({})).with_description("Test"); let options = CallOptions { tools: Some(vec![Tool::from(tool)]), - tool_choice: ToolChoice::Auto, + tool_choice: Some(ToolChoice::Auto), ..default_options(test_prompt()) }; model.do_generate(&options).await.unwrap(); let body = first_request_body(&server).await; - // tool_choice is only sent when not Auto (default) - // Actually, the OpenAI converter sends tool_choice for Auto too - assert!(body.get("tool_choice").is_some() || body.get("tool_choice").is_none()); + assert_eq!(body["tool_choice"], "auto"); } /// TS: "should handle tool choice 'required'" @@ -835,7 +806,7 @@ mod prepare_tools { let tool = FunctionTool::new("testFunction", json!({})).with_description("Test"); let options = CallOptions { tools: Some(vec![Tool::from(tool)]), - tool_choice: ToolChoice::Required, + tool_choice: Some(ToolChoice::Required), ..default_options(test_prompt()) }; model.do_generate(&options).await.unwrap(); @@ -855,7 +826,7 @@ mod prepare_tools { let tool = FunctionTool::new("testFunction", json!({})).with_description("Test"); let options = CallOptions { tools: Some(vec![Tool::from(tool)]), - tool_choice: ToolChoice::None, + tool_choice: Some(ToolChoice::None), ..default_options(test_prompt()) }; model.do_generate(&options).await.unwrap(); @@ -875,9 +846,9 @@ mod prepare_tools { let tool = FunctionTool::new("testFunction", json!({})).with_description("Test"); let options = CallOptions { tools: Some(vec![Tool::from(tool)]), - tool_choice: ToolChoice::Tool { + tool_choice: Some(ToolChoice::Tool { tool_name: "testFunction".to_string(), - }, + }), ..default_options(test_prompt()) }; model.do_generate(&options).await.unwrap(); @@ -886,138 +857,6 @@ mod prepare_tools { assert_eq!(body["tool_choice"]["type"], "function"); assert_eq!(body["tool_choice"]["function"]["name"], "testFunction"); } - - /// TS: "should handle browser search tool with supported model" - #[tokio::test] - async fn browser_search_supported_model() { - let server = MockServer::start().await; - mock_json(&server, text_completion_body()).await; - - let model = provider( - "groq", - Some("test-api-key".to_string()), - "openai/gpt-oss-120b", - Some(ProviderOptions { - base_url: Some(server.uri()), - ..Default::default() - }), - ) - .expect("groq provider should build"); - - let tool = Tool::Provider(aimux_core::tool::ProviderTool { - id: "groq.browser_search".to_string(), - name: "browser_search".to_string(), - args: json!({}), - }); - let options = CallOptions { - tools: Some(vec![tool]), - ..default_options(test_prompt()) - }; - model.do_generate(&options).await.unwrap(); - - let body = first_request_body(&server).await; - assert_eq!(body["tools"][0]["type"], "browser_search"); - } - - /// TS: "should warn when browser search is used with unsupported model" - #[tokio::test] - async fn browser_search_unsupported_model() { - let server = MockServer::start().await; - mock_json(&server, text_completion_body()).await; - - let model = make_provider(&server); - - let tool = Tool::Provider(aimux_core::tool::ProviderTool { - id: "groq.browser_search".to_string(), - name: "browser_search".to_string(), - args: json!({}), - }); - let options = CallOptions { - tools: Some(vec![tool]), - ..default_options(test_prompt()) - }; - let result = model.do_generate(&options).await.unwrap(); - - assert!(result.warnings.iter().any(|w| match w { - aimux_core::types::Warning::Unsupported { - feature, details, .. - } => { - feature.contains("groq.browser_search") - && details.as_ref().is_some_and(|d| d.contains("gemma2-9b-it")) - } - _ => false, - })); - } - - /// TS: "should handle mixed tools with model validation" - #[tokio::test] - async fn mixed_tools_with_browser_search() { - let server = MockServer::start().await; - mock_json(&server, text_completion_body()).await; - - let model = provider( - "groq", - Some("test-api-key".to_string()), - "openai/gpt-oss-20b", - Some(ProviderOptions { - base_url: Some(server.uri()), - ..Default::default() - }), - ) - .expect("groq provider should build"); - - let func_tool = FunctionTool::new("test-tool", json!({"type":"object","properties":{}})) - .with_description("A test tool"); - let browser_tool = Tool::Provider(aimux_core::tool::ProviderTool { - id: "groq.browser_search".to_string(), - name: "browser_search".to_string(), - args: json!({}), - }); - let options = CallOptions { - tools: Some(vec![Tool::from(func_tool), browser_tool]), - ..default_options(test_prompt()) - }; - model.do_generate(&options).await.unwrap(); - - let body = first_request_body(&server).await; - assert_eq!(body["tools"][0]["type"], "function"); - assert_eq!(body["tools"][0]["function"]["name"], "test-tool"); - assert_eq!(body["tools"][1]["type"], "browser_search"); - } - - /// TS: "should validate all browser search supported models" - #[tokio::test] - async fn browser_search_all_supported_models() { - for model_id in &["openai/gpt-oss-20b", "openai/gpt-oss-120b"] { - let server = MockServer::start().await; - mock_json(&server, text_completion_body()).await; - - let model = provider( - "groq", - Some("test-api-key".to_string()), - model_id, - Some(ProviderOptions { - base_url: Some(server.uri()), - ..Default::default() - }), - ) - .expect("groq provider should build"); - - let tool = Tool::Provider(aimux_core::tool::ProviderTool { - id: "groq.browser_search".to_string(), - name: "browser_search".to_string(), - args: json!({}), - }); - let options = CallOptions { - tools: Some(vec![tool]), - ..default_options(test_prompt()) - }; - model.do_generate(&options).await.unwrap(); - - let body = first_request_body(&server).await; - assert_eq!(body["tools"][0]["type"], "browser_search"); - } - } } // ════════════════════════════════════════════════════════════════════════════ @@ -1081,11 +920,11 @@ mod do_generate { .unwrap(); let tool_call = result.content.iter().find_map(|c| match c { - GenerateContent::ToolCall { + GenerateContent::ToolCall(RawToolCall { tool_call_id, tool_name, .. - } => Some((tool_call_id.clone(), tool_name.clone())), + }) => Some((tool_call_id.clone(), tool_name.clone())), _ => None, }); let (id, name) = tool_call.expect("should have tool call"); @@ -1435,9 +1274,9 @@ mod do_generate { ); let options = CallOptions { tools: Some(vec![Tool::from(tool)]), - tool_choice: ToolChoice::Tool { + tool_choice: Some(ToolChoice::Tool { tool_name: "test-tool".to_string(), - }, + }), ..default_options(test_prompt()) }; model.do_generate(&options).await.unwrap(); @@ -1725,18 +1564,18 @@ mod do_stream { // Should have a ToolCall part let tool_call = parts.iter().find_map(|p| match p { - StreamPart::ToolCall { + StreamPart::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. - } => Some((tool_call_id.clone(), tool_name.clone(), input.clone())), + }) => Some((tool_call_id.clone(), tool_name.clone(), input.clone())), _ => None, }); let (id, name, input) = tool_call.expect("should have tool call"); assert_eq!(id, "call_abc"); assert_eq!(name, "test-tool"); - assert_eq!(input, Value::String(r#"{"value":"Sparkle Day"}"#.into())); + assert_eq!(input, r#"{"value":"Sparkle Day"}"#); // Should have tool-calls finish reason let finish = parts.iter().find_map(|p| match p { @@ -1790,13 +1629,10 @@ mod do_stream { let parts = collect_stream(result).await; let tool_call = parts.iter().find_map(|p| match p { - StreamPart::ToolCall { input, .. } => Some(input.clone()), + StreamPart::ToolCall(RawToolCall { input, .. }) => Some(input.clone()), _ => None, }); - assert_eq!( - tool_call.unwrap(), - Value::String(r#"{"value":"Sparkle Day"}"#.into()) - ); + assert_eq!(tool_call.unwrap(), r#"{"value":"Sparkle Day"}"#); } /// TS: "should stream usage from x_groq.usage" diff --git a/aimux-providers/tests/huggingface_responses_test.rs b/aimux-providers/tests/huggingface_responses_test.rs index 52fe83bda..404e4ae21 100644 --- a/aimux-providers/tests/huggingface_responses_test.rs +++ b/aimux-providers/tests/huggingface_responses_test.rs @@ -16,6 +16,7 @@ //! overridden with the mock server's root URI (no `/v1` suffix), so the //! resulting request path is `/responses`. +use aimux_core::tool::{RawToolCall, ToolResult}; use std::collections::HashMap; use futures::StreamExt; @@ -29,10 +30,9 @@ use aimux_core::generate::{GenerateTextOptions, generate_text}; use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; use aimux_core::message::Role; -use aimux_core::options::{CallOptions, ResponseFormat, Tool, ToolChoice}; +use aimux_core::options::{CallOptions, ResponseFormat}; use aimux_core::result::GenerateContent; use aimux_core::stream_part::StreamPart; -use aimux_core::tool::FunctionTool; use aimux_core::types::FinishReasonUnified; use aimux_providers::huggingface::responses::convert_to_huggingface_responses_messages; @@ -536,27 +536,24 @@ async fn should_handle_mcp_tools_with_annotations() { // emits it as a separate content item too). assert_eq!(result.content.len(), 5); match &result.content[0] { - GenerateContent::ToolCall { + GenerateContent::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. - } => { + }) => { assert_eq!(tool_call_id, "mcp_search_test"); assert_eq!(tool_name, "search"); - assert_eq!( - input, - &Value::String(r#"{"query": "San Francisco tech events"}"#.into()) - ); + assert_eq!(input, r#"{"query": "San Francisco tech events"}"#); } other => panic!("expected ToolCall at [0], got {other:?}"), } match &result.content[1] { - GenerateContent::ToolResult { + GenerateContent::ToolResult(ToolResult { tool_call_id, tool_name, .. - } => { + }) => { assert_eq!(tool_call_id, "mcp_search_test"); assert_eq!(tool_name, "search"); } @@ -1110,15 +1107,15 @@ async fn should_handle_function_call_tool_responses() { // output, so there is no paired ToolResult. assert_eq!(result.content.len(), 2); match &result.content[0] { - GenerateContent::ToolCall { + GenerateContent::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. - } => { + }) => { assert_eq!(tool_call_id, "call_123"); assert_eq!(tool_name, "getWeather"); - assert_eq!(input, &Value::String(r#"{"location": "New York"}"#.into())); + assert_eq!(input, r#"{"location": "New York"}"#); } other => panic!("expected ToolCall at [0], got {other:?}"), } @@ -1192,27 +1189,24 @@ async fn should_stream_tool_calls() { other => panic!("expected ToolInputEnd, got {other:?}"), } match &parts[4] { - StreamPart::ToolCall { + StreamPart::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. - } => { + }) => { assert_eq!(tool_call_id, "call_456"); assert_eq!(tool_name, "calculator"); - assert_eq!( - input, - &Value::String(r#"{"operation": "add", "a": 5, "b": 3}"#.into()) - ); + assert_eq!(input, r#"{"operation": "add", "a": 5, "b": 3}"#); } other => panic!("expected ToolCall, got {other:?}"), } match &parts[5] { - StreamPart::ToolResult { + StreamPart::ToolResult(ToolResult { tool_call_id, result, .. - } => { + }) => { assert_eq!(tool_call_id, "call_456"); assert_eq!(result, &json!("8")); } @@ -1640,131 +1634,6 @@ async fn should_send_provider_specific_options() { // Tool preparation // ════════════════════════════════════════════════════════════════════════════ -/// TS: tool preparation › "should prepare tools correctly" -#[tokio::test] -async fn should_prepare_tools_correctly() { - let server = MockServer::start().await; - mock_json( - &server, - json!({ - "id": "resp_tools", - "model": "deepseek-ai/DeepSeek-V3-0324", - "object": "response", - "created_at": 1741257730, - "status": "completed", - "error": null, - "instructions": null, - "max_output_tokens": null, - "metadata": null, - "tool_choice": "auto", - "tools": [], - "temperature": 1.0, - "top_p": 1.0, - "incomplete_details": null, - "usage": null, - "output": [], - "output_text": "Test" - }), - ) - .await; - - let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); - - let mut options = default_options(test_prompt()); - options.tools = Some(vec![Tool::from( - FunctionTool::new( - "getWeather", - json!({ - "type": "object", - "properties": { "location": { "type": "string" } }, - "required": ["location"] - }), - ) - .with_description("Get weather information"), - )]); - options.tool_choice = ToolChoice::Tool { - tool_name: "getWeather".to_string(), - }; - - let result = model.do_generate(&options).await.expect("should succeed"); - let body = result.request_body.expect("body"); - - assert_eq!( - body["tools"], - json!([{ - "type": "function", - "name": "getWeather", - "description": "Get weather information", - "parameters": { - "type": "object", - "properties": { "location": { "type": "string" } }, - "required": ["location"] - } - }]) - ); - assert_eq!( - body["tool_choice"], - json!({ "type": "function", "function": { "name": "getWeather" } }) - ); -} - -/// TS: tool preparation › "should handle auto and required tool choices" -#[tokio::test] -async fn should_handle_auto_and_required_tool_choices() { - let server = MockServer::start().await; - mock_json( - &server, - json!({ - "id": "resp_tools", - "model": "deepseek-ai/DeepSeek-V3-0324", - "object": "response", - "created_at": 1741257730, - "status": "completed", - "error": null, - "instructions": null, - "max_output_tokens": null, - "metadata": null, - "tool_choice": "auto", - "tools": [], - "temperature": 1.0, - "top_p": 1.0, - "incomplete_details": null, - "usage": null, - "output": [], - "output_text": "Test" - }), - ) - .await; - - let provider = make_provider(&server); - let model = provider.responses_model("deepseek-ai/DeepSeek-V3-0324"); - - // Test auto - let mut options = default_options(test_prompt()); - options.tools = Some(vec![Tool::from(FunctionTool::new( - "test", - json!({ "type": "object" }), - ))]); - options.tool_choice = ToolChoice::Auto; - - let result = model.do_generate(&options).await.expect("should succeed"); - let body = result.request_body.expect("body"); - assert_eq!(body["tool_choice"], json!("auto")); - - // Test required - let mut options = default_options(test_prompt()); - options.tools = Some(vec![Tool::from(FunctionTool::new( - "test", - json!({ "type": "object" }), - ))]); - options.tool_choice = ToolChoice::Required; - - let result = model.do_generate(&options).await.expect("should succeed"); - let body = result.request_body.expect("body"); - assert_eq!(body["tool_choice"], json!("required")); -} - // ════════════════════════════════════════════════════════════════════════════ // Top-level-only media type resolution // ════════════════════════════════════════════════════════════════════════════ diff --git a/aimux-providers/tests/open_responses_test.rs b/aimux-providers/tests/open_responses_test.rs index f2821d69a..e72ffdd47 100644 --- a/aimux-providers/tests/open_responses_test.rs +++ b/aimux-providers/tests/open_responses_test.rs @@ -9,6 +9,7 @@ //! `createTestServer`. Each test starts its own `MockServer` so parallel //! `#[tokio::test]` runs do not collide. +use aimux_core::tool::RawToolCall; use std::collections::HashMap; use futures::StreamExt; @@ -20,7 +21,7 @@ use aimux_core::content::ContentPart; use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; use aimux_core::message::Role; -use aimux_core::options::{CallOptions, ResponseFormat, Tool, ToolChoice}; +use aimux_core::options::{CallOptions, ResponseFormat, Tool}; use aimux_core::result::GenerateContent; use aimux_core::stream_part::StreamPart; use aimux_core::tool::FunctionTool; @@ -165,40 +166,6 @@ fn lmstudio_basic_json() -> Value { }) } -/// The `lmstudio-tool-call.1.json` fixture. -fn lmstudio_tool_call_json() -> Value { - json!({ - "id": "resp_930de53bd4b5933673481fa630f3dc5f58027a2c67598a2a", - "object": "response", - "created_at": 1769005553, - "status": "completed", - "incomplete_details": null, - "model": "mistralai/ministral-3-14b-reasoning", - "output": [ - { - "id": "fc_ru0kcno9erlzp8573yub", - "call_id": "call_2866856768160095", - "type": "function_call", - "name": "weather", - "arguments": "{\"location\":\"San Francisco\"}", - "status": "completed" - } - ], - "error": null, - "usage": { - "input_tokens": 1189, - "output_tokens": 11, - "total_tokens": 1200, - "input_tokens_details": { - "cached_tokens": 891 - }, - "output_tokens_details": { - "reasoning_tokens": 0 - } - } - }) -} - /// The `openai-pdf-input-file.1.json` fixture. fn openai_pdf_json() -> Value { json!({ @@ -1412,190 +1379,8 @@ mod do_generate_tests { // -- Tool call parsing tests -- - #[tokio::test] - async fn parse_tool_call_from_response() { - let server = MockServer::start().await; - mock_json(&server, lmstudio_tool_call_json()).await; - - let model = make_model(&server, "gemma-7b-it"); - let options = CallOptions { - tools: Some(vec![Tool::from( - FunctionTool::new( - "weather", - json!({ - "type": "object", - "properties": { - "location": {"type": "string", "description": "The location to get the weather for"} - }, - "required": ["location"] - }), - ) - .with_description("Get the weather in a location"), - )]), - tool_choice: ToolChoice::Required, - ..CallOptions::new(test_prompt()) - }; - let result = model.do_generate(&options).await.unwrap(); - - assert_eq!(result.content.len(), 1); - match &result.content[0] { - GenerateContent::ToolCall { - tool_call_id, - tool_name, - input, - .. - } => { - assert_eq!(tool_call_id, "call_2866856768160095"); - assert_eq!(tool_name, "weather"); - assert_eq!( - input, - &Value::String(r#"{"location":"San Francisco"}"#.into()) - ); - } - other => panic!("expected ToolCall, got {other:?}"), - } - } - - #[tokio::test] - async fn tool_call_finish_reason() { - let server = MockServer::start().await; - mock_json(&server, lmstudio_tool_call_json()).await; - - let model = make_model(&server, "gemma-7b-it"); - let options = CallOptions { - tools: Some(vec![Tool::from( - FunctionTool::new("weather", json!({"type": "object"})) - .with_description("Get the weather in a location"), - )]), - tool_choice: ToolChoice::Required, - ..CallOptions::new(test_prompt()) - }; - let result = model.do_generate(&options).await.unwrap(); - - assert_eq!(result.finish_reason.unified, FinishReasonUnified::ToolCalls); - assert_eq!(result.finish_reason.raw, None); - } - - #[tokio::test] - async fn tool_call_usage() { - let server = MockServer::start().await; - mock_json(&server, lmstudio_tool_call_json()).await; - - let model = make_model(&server, "gemma-7b-it"); - let options = CallOptions { - tools: Some(vec![Tool::from( - FunctionTool::new("weather", json!({"type": "object"})) - .with_description("Get the weather in a location"), - )]), - tool_choice: ToolChoice::Required, - ..CallOptions::new(test_prompt()) - }; - let result = model.do_generate(&options).await.unwrap(); - - assert_eq!(result.usage.input_tokens.total, Some(1189)); - assert_eq!(result.usage.input_tokens.cache_read, Some(891)); - assert_eq!(result.usage.input_tokens.no_cache, Some(298)); - assert_eq!(result.usage.output_tokens.total, Some(11)); - assert_eq!(result.usage.output_tokens.reasoning, Some(0)); - assert_eq!(result.usage.output_tokens.text, Some(11)); - } - // -- Tool choice tests -- - fn test_tool() -> Tool { - Tool::from( - FunctionTool::new( - "get_weather", - json!({ - "type": "object", - "properties": {"location": {"type": "string"}}, - "required": ["location"] - }), - ) - .with_description("Get the current weather"), - ) - } - - #[tokio::test] - async fn tool_choice_auto() { - let server = MockServer::start().await; - mock_json(&server, lmstudio_basic_json()).await; - - let model = make_model(&server, "gemma-7b-it"); - let options = CallOptions { - tools: Some(vec![test_tool()]), - tool_choice: ToolChoice::Auto, - ..CallOptions::new(test_prompt()) - }; - model.do_generate(&options).await.unwrap(); - - let request = &server.received_requests().await.unwrap()[0]; - let body: Value = serde_json::from_slice(&request.body).unwrap(); - // In Rust, ToolChoice::Auto is the default and is not emitted (matches - // TS behavior where undefined toolChoice is omitted). Tools are still present. - assert!(body.get("tool_choice").is_none()); - assert!(body.get("tools").is_some()); - } - - #[tokio::test] - async fn tool_choice_none() { - let server = MockServer::start().await; - mock_json(&server, lmstudio_basic_json()).await; - - let model = make_model(&server, "gemma-7b-it"); - let options = CallOptions { - tools: Some(vec![test_tool()]), - tool_choice: ToolChoice::None, - ..CallOptions::new(test_prompt()) - }; - model.do_generate(&options).await.unwrap(); - - let request = &server.received_requests().await.unwrap()[0]; - let body: Value = serde_json::from_slice(&request.body).unwrap(); - assert_eq!(body["tool_choice"], json!("none")); - } - - #[tokio::test] - async fn tool_choice_required() { - let server = MockServer::start().await; - mock_json(&server, lmstudio_basic_json()).await; - - let model = make_model(&server, "gemma-7b-it"); - let options = CallOptions { - tools: Some(vec![test_tool()]), - tool_choice: ToolChoice::Required, - ..CallOptions::new(test_prompt()) - }; - model.do_generate(&options).await.unwrap(); - - let request = &server.received_requests().await.unwrap()[0]; - let body: Value = serde_json::from_slice(&request.body).unwrap(); - assert_eq!(body["tool_choice"], json!("required")); - } - - #[tokio::test] - async fn tool_choice_specific_tool() { - let server = MockServer::start().await; - mock_json(&server, lmstudio_basic_json()).await; - - let model = make_model(&server, "gemma-7b-it"); - let options = CallOptions { - tools: Some(vec![test_tool()]), - tool_choice: ToolChoice::Tool { - tool_name: "get_weather".to_string(), - }, - ..CallOptions::new(test_prompt()) - }; - model.do_generate(&options).await.unwrap(); - - let request = &server.received_requests().await.unwrap()[0]; - let body: Value = serde_json::from_slice(&request.body).unwrap(); - assert_eq!( - body["tool_choice"], - json!({"type": "function", "name": "get_weather"}) - ); - } - // -- System messages tests -- #[tokio::test] @@ -2009,19 +1794,19 @@ mod do_stream_tests { // ToolCall let tool_call = parts.iter().find_map(|p| match p { - StreamPart::ToolCall { + StreamPart::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. - } => Some((tool_call_id, tool_name, input)), + }) => Some((tool_call_id, tool_name, input)), _ => None, }); assert!(tool_call.is_some(), "should have ToolCall"); let (tc_id, tc_name, tc_input) = tool_call.unwrap(); assert_eq!(tc_id, "call_1"); assert_eq!(tc_name, "weather"); - assert_eq!(tc_input, &Value::String(r#"{"location":"SF"}"#.into())); + assert_eq!(tc_input, r#"{"location":"SF"}"#); // Finish with tool-calls reason let finish = parts.iter().find_map(|p| match p { diff --git a/aimux-providers/tests/openai_compatible_test.rs b/aimux-providers/tests/openai_compatible_test.rs index 8e49e65c1..e554f4287 100644 --- a/aimux-providers/tests/openai_compatible_test.rs +++ b/aimux-providers/tests/openai_compatible_test.rs @@ -15,6 +15,7 @@ //! Because the wrappers are structurally identical, a single macro generates //! the four tests per provider — DRY without hiding what is asserted. +use aimux_core::tool::RawToolCall; use futures::StreamExt; use serde_json::{Value, json}; use wiremock::matchers::{method, path}; @@ -592,10 +593,10 @@ macro_rules! openai_compatible_tool_tests { assert_eq!(result.content.len(), 1); match &result.content[0] { - GenerateContent::ToolCall { tool_call_id, tool_name, input, .. } => { + GenerateContent::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. }) => { assert_eq!(tool_call_id, "call_abc"); assert_eq!(tool_name, "get-weather"); - assert_eq!(input, &Value::String(r#"{"city":"SF"}"#.into())); + assert_eq!(input, r#"{"city":"SF"}"#); } other => panic!("expected ToolCall, got {:?}", other), } @@ -630,7 +631,7 @@ macro_rules! openai_compatible_tool_tests { let parts = collect_stream(result).await; let tool_call = parts.iter().find_map(|p| match p { - StreamPart::ToolCall { tool_call_id, tool_name, input, .. } => { + StreamPart::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. }) => { Some((tool_call_id.clone(), tool_name.clone(), input.clone())) } _ => None, @@ -638,7 +639,7 @@ macro_rules! openai_compatible_tool_tests { let (id, name, input) = tool_call.expect("should have ToolCall"); assert_eq!(id, "call_abc"); assert_eq!(name, "get-weather"); - assert_eq!(input, Value::String(r#"{"city":"SF"}"#.into())); + assert_eq!(input, r#"{"city":"SF"}"#); } } } @@ -688,10 +689,10 @@ macro_rules! openai_compatible_tool_tests { assert_eq!(result.content.len(), 1); match &result.content[0] { - GenerateContent::ToolCall { tool_call_id, tool_name, input, .. } => { + GenerateContent::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. }) => { assert_eq!(tool_call_id, "call_abc"); assert_eq!(tool_name, "get-weather"); - assert_eq!(input, &Value::String(r#"{"city":"SF"}"#.into())); + assert_eq!(input, r#"{"city":"SF"}"#); } other => panic!("expected ToolCall, got {:?}", other), } @@ -726,7 +727,7 @@ macro_rules! openai_compatible_tool_tests { let parts = collect_stream(result).await; let tool_call = parts.iter().find_map(|p| match p { - StreamPart::ToolCall { tool_call_id, tool_name, input, .. } => { + StreamPart::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. }) => { Some((tool_call_id.clone(), tool_name.clone(), input.clone())) } _ => None, @@ -734,7 +735,7 @@ macro_rules! openai_compatible_tool_tests { let (id, name, input) = tool_call.expect("should have ToolCall"); assert_eq!(id, "call_abc"); assert_eq!(name, "get-weather"); - assert_eq!(input, Value::String(r#"{"city":"SF"}"#.into())); + assert_eq!(input, r#"{"city":"SF"}"#); } } } diff --git a/aimux-providers/tests/openrouter_test.rs b/aimux-providers/tests/openrouter_test.rs index 320979af6..76d4f3bbd 100644 --- a/aimux-providers/tests/openrouter_test.rs +++ b/aimux-providers/tests/openrouter_test.rs @@ -1,4 +1,4 @@ -//! Provider-specific tests for the OpenRouter provider. +//! Provider-specific tests for the OpenRouter provider. //! //! OpenRouter is a thin OpenAI-compatible wrapper over [`OpenAIProvider`]. The //! behaviours verified here are the ones the wrapper is responsible for: @@ -17,6 +17,7 @@ //! `/api/v1/chat/completions` — matching the path the real OpenRouter API //! (base `https://openrouter.ai/api/v1`) records in the cassettes. +use aimux_core::tool::RawToolCall; use futures::StreamExt; use serde_json::{Value, json}; use serial_test::serial; @@ -28,7 +29,7 @@ use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; use aimux_core::message::Role; -use aimux_core::options::{CallOptions, Tool, ToolChoice}; +use aimux_core::options::{CallOptions, Tool}; use aimux_core::provider::Provider; use aimux_core::result::GenerateContent; use aimux_core::stream_part::StreamPart; @@ -399,47 +400,21 @@ async fn do_generate_extracts_tool_call() { assert_eq!(result.content.len(), 1); match &result.content[0] { - GenerateContent::ToolCall { + GenerateContent::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. - } => { + }) => { assert_eq!(tool_call_id, "call_abc"); assert_eq!(tool_name, "get-weather"); - assert_eq!(input, &Value::String(r#"{"city":"SF"}"#.into())); + assert_eq!(input, r#"{"city":"SF"}"#); } other => panic!("expected ToolCall, got {other:?}"), } assert_eq!(result.finish_reason.unified, FinishReasonUnified::ToolCalls); } -/// Tool choice `required` is forwarded as `"required"`. -#[tokio::test] -async fn tool_choice_required_forwarded() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/chat/completions")) - .respond_with(ResponseTemplate::new(200).set_body_json(text_completion_body())) - .mount(&server) - .await; - - let provider = make_provider(&server); - let model = provider.model("openai/gpt-4o-mini"); - - let tool = FunctionTool::new("get-weather", json!({})).with_description("Test"); - let options = CallOptions { - tools: Some(vec![Tool::from(tool)]), - tool_choice: ToolChoice::Required, - ..default_options(test_prompt()) - }; - model.do_generate(&options).await.unwrap(); - - let requests = server.received_requests().await.expect("requests recorded"); - let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - assert_eq!(body["tool_choice"], "required"); -} - /// do_stream returns text deltas. #[tokio::test] async fn do_stream_returns_text() { @@ -527,18 +502,18 @@ async fn do_stream_emits_tool_call() { let parts = collect_stream(result).await; let tool_call = parts.iter().find_map(|p| match p { - StreamPart::ToolCall { + StreamPart::ToolCall(RawToolCall { tool_call_id, tool_name, input, .. - } => Some((tool_call_id.clone(), tool_name.clone(), input.clone())), + }) => Some((tool_call_id.clone(), tool_name.clone(), input.clone())), _ => None, }); let (id, name, input) = tool_call.expect("should have ToolCall"); assert_eq!(id, "call_abc"); assert_eq!(name, "get-weather"); - assert_eq!(input, Value::String(r#"{"city":"SF"}"#.into())); + assert_eq!(input, r#"{"city":"SF"}"#); } /// A 401 response maps to `AiMuxError::ApiCall` (401 in `status_code`). @@ -645,7 +620,7 @@ mod conformance { fn has_tool_call(content: &[GenerateContent]) -> bool { content .iter() - .any(|c| matches!(c, GenerateContent::ToolCall { .. })) + .any(|c| matches!(c, GenerateContent::ToolCall(_))) } fn has_reasoning(content: &[GenerateContent]) -> bool { diff --git a/bindings/flutter/lib/types.dart b/bindings/flutter/lib/types.dart index 2eb4beb2f..52c564fc9 100644 --- a/bindings/flutter/lib/types.dart +++ b/bindings/flutter/lib/types.dart @@ -161,8 +161,6 @@ class ToolCall { // and mapped to the `dynamic` JSON key. @JsonKey(name: 'dynamic') final bool? isDynamic; - @JsonKey(name: 'thought_signature') - final String? thoughtSignature; /// Additional provider-specific metadata associated with this call. @JsonKey(name: 'provider_metadata', includeIfNull: false) final dynamic providerMetadata; @@ -177,7 +175,6 @@ class ToolCall { required this.input, this.providerExecuted, this.isDynamic, - this.thoughtSignature, this.providerMetadata, this.invalid, this.error, @@ -217,10 +214,6 @@ class RawToolCall { /// identifier, so the field is `isDynamic` and maps to the `dynamic` key. final bool? isDynamic; - /// Provider-assigned thought signature (e.g. Gemini `thoughtSignature`), - /// echoed back verbatim on the next turn. - final String? thoughtSignature; - /// Additional provider-specific metadata associated with this call. final dynamic providerMetadata; @@ -230,7 +223,6 @@ class RawToolCall { required this.input, this.providerExecuted, this.isDynamic, - this.thoughtSignature, this.providerMetadata, }); @@ -240,7 +232,6 @@ class RawToolCall { input: json['input'] as String, providerExecuted: json['provider_executed'] as bool?, isDynamic: json['dynamic'] as bool?, - thoughtSignature: json['thought_signature'] as String?, providerMetadata: json['provider_metadata'], ); @@ -253,7 +244,6 @@ class RawToolCall { input: input ?? this.input, providerExecuted: providerExecuted, isDynamic: isDynamic, - thoughtSignature: thoughtSignature, providerMetadata: providerMetadata, ); @@ -263,7 +253,6 @@ class RawToolCall { 'input': input, if (providerExecuted != null) 'provider_executed': providerExecuted, if (isDynamic != null) 'dynamic': isDynamic, - if (thoughtSignature != null) 'thought_signature': thoughtSignature, if (providerMetadata != null) 'provider_metadata': providerMetadata, }; } @@ -339,6 +328,14 @@ typedef RepairToolCall = FutureOr Function( // Tools // ───────────────────────────────────────────────────────────────────────────── +class FunctionToolInputExample { + final Map input; + FunctionToolInputExample({required this.input}); + factory FunctionToolInputExample.fromJson(Map json) => + FunctionToolInputExample(input: json['input'] as Map); + Map toJson() => {'input': input}; +} + /// A function tool definition. Mirrors `FunctionTool.ts`. @JsonSerializable() class FunctionTool { @@ -350,7 +347,7 @@ class FunctionTool { @JsonKey(name: 'provider_options') final Map? providerOptions; @JsonKey(name: 'input_examples') - final List>? inputExamples; + final List? inputExamples; FunctionTool({ required this.name, @@ -498,7 +495,6 @@ final class GenerateContentToolCall extends GenerateContent { final dynamic input; final bool? providerExecuted; final bool? isDynamic; - final String? thoughtSignature; final Map? providerMetadata; GenerateContentToolCall({ @@ -507,7 +503,6 @@ final class GenerateContentToolCall extends GenerateContent { required this.input, this.providerExecuted, this.isDynamic, - this.thoughtSignature, this.providerMetadata, }); @@ -521,7 +516,6 @@ final class GenerateContentToolCall extends GenerateContent { input: json['input'], providerExecuted: json['provider_executed'] as bool?, isDynamic: json['dynamic'] as bool?, - thoughtSignature: json['thought_signature'] as String?, providerMetadata: json['provider_metadata'] as Map?, ); @@ -534,7 +528,6 @@ final class GenerateContentToolCall extends GenerateContent { 'input': input, if (providerExecuted != null) 'provider_executed': providerExecuted, if (isDynamic != null) 'dynamic': isDynamic, - if (thoughtSignature != null) 'thought_signature': thoughtSignature, if (providerMetadata != null) 'provider_metadata': providerMetadata, }, }; @@ -1427,7 +1420,6 @@ final class StreamPartToolCall extends StreamPart { final dynamic input; final bool? providerExecuted; final bool? isDynamic; - final String? thoughtSignature; final Map? providerMetadata; /// Set by Core when the tool call stays invalid after optional repair. final bool? invalid; @@ -1440,7 +1432,6 @@ final class StreamPartToolCall extends StreamPart { required this.input, this.providerExecuted, this.isDynamic, - this.thoughtSignature, this.providerMetadata, this.invalid, this.error, @@ -1453,7 +1444,6 @@ final class StreamPartToolCall extends StreamPart { input: json['input'], providerExecuted: json['provider_executed'] as bool?, isDynamic: json['dynamic'] as bool?, - thoughtSignature: json['thought_signature'] as String?, providerMetadata: json['provider_metadata'] as Map?, invalid: json['invalid'] as bool?, @@ -1468,7 +1458,6 @@ final class StreamPartToolCall extends StreamPart { 'input': input, if (providerExecuted != null) 'provider_executed': providerExecuted, if (isDynamic != null) 'dynamic': isDynamic, - if (thoughtSignature != null) 'thought_signature': thoughtSignature, if (providerMetadata != null) 'provider_metadata': providerMetadata, if (invalid != null) 'invalid': invalid, if (error != null) 'error': error, @@ -1987,17 +1976,15 @@ final class ContentPartToolCall extends ContentPart { final String toolName; final dynamic input; final bool? providerExecuted; - final String? thoughtSignature; final Map? providerOptions; const ContentPartToolCall( - {required this.toolCallId, required this.toolName, required this.input, this.providerExecuted, this.thoughtSignature, this.providerOptions}); + {required this.toolCallId, required this.toolName, required this.input, this.providerExecuted, this.providerOptions}); static ContentPartToolCall fromJson(Map json) => ContentPartToolCall( toolCallId: json['tool_call_id'] as String, toolName: json['tool_name'] as String, input: json['input'], providerExecuted: json['provider_executed'] as bool?, - thoughtSignature: json['thought_signature'] as String?, providerOptions: json['provider_options'] as Map?, ); @override @@ -2007,7 +1994,6 @@ final class ContentPartToolCall extends ContentPart { 'tool_name': toolName, 'input': input, if (providerExecuted != null) 'provider_executed': providerExecuted, - if (thoughtSignature != null) 'thought_signature': thoughtSignature, if (providerOptions != null) 'provider_options': providerOptions, }; } diff --git a/bindings/flutter/lib/types.g.dart b/bindings/flutter/lib/types.g.dart index a0a8f56fa..1004fd158 100644 --- a/bindings/flutter/lib/types.g.dart +++ b/bindings/flutter/lib/types.g.dart @@ -57,7 +57,6 @@ ToolCall _$ToolCallFromJson(Map json) => ToolCall( input: json['input'], providerExecuted: json['provider_executed'] as bool?, isDynamic: json['dynamic'] as bool?, - thoughtSignature: json['thought_signature'] as String?, providerMetadata: json['provider_metadata'], invalid: json['invalid'] as bool?, error: json['error'], @@ -69,7 +68,6 @@ Map _$ToolCallToJson(ToolCall instance) => { 'input': instance.input, 'provider_executed': instance.providerExecuted, 'dynamic': instance.isDynamic, - 'thought_signature': instance.thoughtSignature, if (instance.providerMetadata != null) 'provider_metadata': instance.providerMetadata, 'invalid': instance.invalid, @@ -83,7 +81,7 @@ FunctionTool _$FunctionToolFromJson(Map json) => FunctionTool( strict: json['strict'] as bool?, providerOptions: json['provider_options'] as Map?, inputExamples: (json['input_examples'] as List?) - ?.map((e) => e as Map) + ?.map((e) => FunctionToolInputExample.fromJson(e as Map)) .toList(), ); @@ -94,7 +92,7 @@ Map _$FunctionToolToJson(FunctionTool instance) => 'input_schema': instance.inputSchema, 'strict': instance.strict, 'provider_options': instance.providerOptions, - 'input_examples': instance.inputExamples, + 'input_examples': instance.inputExamples?.map((e) => e.toJson()).toList(), }; ResponseMetadata _$ResponseMetadataFromJson(Map json) => diff --git a/bindings/flutter/test/tool_call_repair_typed_test.dart b/bindings/flutter/test/tool_call_repair_typed_test.dart index 0a0cfd35e..181cb5c5b 100644 --- a/bindings/flutter/test/tool_call_repair_typed_test.dart +++ b/bindings/flutter/test/tool_call_repair_typed_test.dart @@ -116,7 +116,6 @@ void main() { 'input': '{"town":"Tokyo"}', 'provider_executed': true, 'dynamic': true, - 'thought_signature': 'sig', 'provider_metadata': { 'google': {'x': 1}, }, @@ -127,7 +126,6 @@ void main() { 'input': '{"city":"Tokyo"}', 'provider_executed': true, 'dynamic': true, - 'thought_signature': 'sig', 'provider_metadata': { 'google': {'x': 1}, }, diff --git a/bindings/go/repair.go b/bindings/go/repair.go index 5c12383fb..017eb7925 100644 --- a/bindings/go/repair.go +++ b/bindings/go/repair.go @@ -29,7 +29,6 @@ type RawToolCall struct { Input string `json:"input"` ProviderExecuted *bool `json:"provider_executed,omitempty"` Dynamic *bool `json:"dynamic,omitempty"` - ThoughtSignature *string `json:"thought_signature,omitempty"` ProviderMetadata json.RawMessage `json:"provider_metadata,omitempty"` } diff --git a/bindings/go/types.go b/bindings/go/types.go index b63089e6c..d504c528e 100644 --- a/bindings/go/types.go +++ b/bindings/go/types.go @@ -108,7 +108,6 @@ type ToolCall struct { Input json.RawMessage `json:"input,omitempty"` ProviderExecuted *bool `json:"provider_executed,omitempty"` Dynamic *bool `json:"dynamic,omitempty"` - ThoughtSignature *string `json:"thought_signature,omitempty"` // ProviderMetadata carries provider-specific data associated with this call. ProviderMetadata json.RawMessage `json:"provider_metadata,omitempty"` // Invalid is set by Core when the tool call stays invalid after optional repair. @@ -309,6 +308,10 @@ type GenerateTextOptions struct { RepairToolCall RepairToolCallFunc `json:"-"` } +type FunctionToolInputExample struct { + Input map[string]json.RawMessage `json:"input"` +} + // Tool is a function tool definition (the "function" variant). // Mirrors Kotlin Tool.Function (Types.kt:209-230). type Tool struct { @@ -318,7 +321,7 @@ type Tool struct { InputSchema json.RawMessage `json:"input_schema,omitempty"` Strict *bool `json:"strict,omitempty"` ProviderOptions map[string]json.RawMessage `json:"provider_options,omitempty"` - InputExamples []json.RawMessage `json:"input_examples,omitempty"` + InputExamples []FunctionToolInputExample `json:"input_examples,omitempty"` } // MarshalOptions serializes GenerateTextOptions to JSON. Returns "" for nil opts. diff --git a/bindings/java/src/main/java/ai/arcships/aimux/Types.java b/bindings/java/src/main/java/ai/arcships/aimux/Types.java index 62d272e89..f39084310 100644 --- a/bindings/java/src/main/java/ai/arcships/aimux/Types.java +++ b/bindings/java/src/main/java/ai/arcships/aimux/Types.java @@ -410,7 +410,6 @@ public static class ToolCall { @JsonProperty("input") private JsonNode input = emptyObject(); @JsonProperty("provider_executed") private Boolean providerExecuted; @JsonProperty("dynamic") private Boolean dynamic; - @JsonProperty("thought_signature") private String thoughtSignature; @JsonProperty("provider_metadata") private JsonNode providerMetadata; @JsonProperty("invalid") private Boolean invalid; @JsonProperty("error") private JsonNode error; @@ -419,13 +418,12 @@ public static class ToolCall { ToolCall() {} private ToolCall(String toolCallId, String toolName, JsonNode input, Boolean providerExecuted, Boolean dynamic, - String thoughtSignature, JsonNode providerMetadata, Boolean invalid, JsonNode error) { + JsonNode providerMetadata, Boolean invalid, JsonNode error) { this.toolCallId = toolCallId; this.toolName = toolName; this.input = input; this.providerExecuted = providerExecuted; this.dynamic = dynamic; - this.thoughtSignature = thoughtSignature; this.providerMetadata = providerMetadata; this.invalid = invalid; this.error = error; @@ -436,7 +434,6 @@ private ToolCall(String toolCallId, String toolName, JsonNode input, Boolean pro public JsonNode getInput() { return input; } public Boolean getProviderExecuted() { return providerExecuted; } public Boolean getDynamic() { return dynamic; } - public String getThoughtSignature() { return thoughtSignature; } public JsonNode getProviderMetadata() { return providerMetadata; } /** Set by Core when the tool call stays invalid after optional repair. */ public Boolean getInvalid() { return invalid; } @@ -451,7 +448,6 @@ public static class Builder { private JsonNode input = emptyObject(); private Boolean providerExecuted; private Boolean dynamic; - private String thoughtSignature; private JsonNode providerMetadata; private Boolean invalid; private JsonNode error; @@ -461,13 +457,12 @@ public static class Builder { public Builder input(JsonNode v) { this.input = v; return this; } public Builder providerExecuted(Boolean v) { this.providerExecuted = v; return this; } public Builder dynamic(Boolean v) { this.dynamic = v; return this; } - public Builder thoughtSignature(String v) { this.thoughtSignature = v; return this; } public Builder providerMetadata(JsonNode v) { this.providerMetadata = v; return this; } public Builder invalid(Boolean v) { this.invalid = v; return this; } public Builder error(JsonNode v) { this.error = v; return this; } public ToolCall build() { - return new ToolCall(toolCallId, toolName, input, providerExecuted, dynamic, thoughtSignature, providerMetadata, + return new ToolCall(toolCallId, toolName, input, providerExecuted, dynamic, providerMetadata, invalid, error); } } @@ -482,7 +477,6 @@ public boolean equals(Object o) { && Objects.equals(input, that.input) && Objects.equals(providerExecuted, that.providerExecuted) && Objects.equals(dynamic, that.dynamic) - && Objects.equals(thoughtSignature, that.thoughtSignature) && Objects.equals(providerMetadata, that.providerMetadata) && Objects.equals(invalid, that.invalid) && Objects.equals(error, that.error); @@ -490,7 +484,7 @@ public boolean equals(Object o) { @Override public int hashCode() { - return Objects.hash(toolCallId, toolName, input, providerExecuted, dynamic, thoughtSignature, providerMetadata, + return Objects.hash(toolCallId, toolName, input, providerExecuted, dynamic, providerMetadata, invalid, error); } } @@ -501,6 +495,25 @@ public int hashCode() { // `Tool` is an internally-tagged union on `type` (`"function" | "provider"`). // ───────────────────────────────────────────────────────────────────────────── + public static class FunctionToolInputExample { + @JsonProperty("input") private ObjectNode input = emptyObject(); + + @JsonCreator + FunctionToolInputExample() {} + + public FunctionToolInputExample(ObjectNode input) { this.input = input; } + public ObjectNode getInput() { return input; } + + @Override + public boolean equals(Object o) { + return o instanceof FunctionToolInputExample + && Objects.equals(input, ((FunctionToolInputExample) o).input); + } + + @Override + public int hashCode() { return Objects.hash(input); } + } + /** * A user-defined function tool. * @@ -512,13 +525,13 @@ public static class FunctionTool { @JsonProperty("input_schema") private JsonNode inputSchema = emptyObject(); @JsonProperty("strict") private Boolean strict; @JsonProperty("provider_options") private Map providerOptions; - @JsonProperty("input_examples") private List inputExamples; + @JsonProperty("input_examples") private List inputExamples; @JsonCreator FunctionTool() {} private FunctionTool(String name, String description, JsonNode inputSchema, Boolean strict, - Map providerOptions, List inputExamples) { + Map providerOptions, List inputExamples) { this.name = name; this.description = description; this.inputSchema = inputSchema; @@ -532,7 +545,7 @@ private FunctionTool(String name, String description, JsonNode inputSchema, Bool public JsonNode getInputSchema() { return inputSchema; } public Boolean getStrict() { return strict; } public Map getProviderOptions() { return providerOptions; } - public List getInputExamples() { return inputExamples; } + public List getInputExamples() { return inputExamples; } public static Builder builder() { return new Builder(); } @@ -542,14 +555,14 @@ public static class Builder { private JsonNode inputSchema = emptyObject(); private Boolean strict; private Map providerOptions; - private List inputExamples; + private List inputExamples; public Builder name(String v) { this.name = v; return this; } public Builder description(String v) { this.description = v; return this; } public Builder inputSchema(JsonNode v) { this.inputSchema = v; return this; } public Builder strict(Boolean v) { this.strict = v; return this; } public Builder providerOptions(Map v) { this.providerOptions = v; return this; } - public Builder inputExamples(List v) { this.inputExamples = v; return this; } + public Builder inputExamples(List v) { this.inputExamples = v; return this; } public FunctionTool build() { return new FunctionTool(name, description, inputSchema, strict, providerOptions, inputExamples); @@ -583,12 +596,12 @@ public int hashCode() { public static class ProviderTool { @JsonProperty("id") private String id = ""; @JsonProperty("name") private String name = ""; - @JsonProperty("args") private JsonNode args = emptyObject(); + @JsonProperty("args") private ObjectNode args = emptyObject(); @JsonCreator ProviderTool() {} - private ProviderTool(String id, String name, JsonNode args) { + private ProviderTool(String id, String name, ObjectNode args) { this.id = id; this.name = name; this.args = args; @@ -596,18 +609,18 @@ private ProviderTool(String id, String name, JsonNode args) { public String getId() { return id; } public String getName() { return name; } - public JsonNode getArgs() { return args; } + public ObjectNode getArgs() { return args; } public static Builder builder() { return new Builder(); } public static class Builder { private String id = ""; private String name = ""; - private JsonNode args = emptyObject(); + private ObjectNode args = emptyObject(); public Builder id(String v) { this.id = v; return this; } public Builder name(String v) { this.name = v; return this; } - public Builder args(JsonNode v) { this.args = v; return this; } + public Builder args(ObjectNode v) { this.args = v; return this; } public ProviderTool build() { return new ProviderTool(id, name, args); } } @@ -649,13 +662,13 @@ public static class Function extends Tool { @JsonProperty("input_schema") private JsonNode inputSchema = emptyObject(); @JsonProperty("strict") private Boolean strict; @JsonProperty("provider_options") private Map providerOptions; - @JsonProperty("input_examples") private List inputExamples; + @JsonProperty("input_examples") private List inputExamples; @JsonCreator Function() {} private Function(String name, String description, JsonNode inputSchema, Boolean strict, - Map providerOptions, List inputExamples) { + Map providerOptions, List inputExamples) { this.name = name; this.description = description; this.inputSchema = inputSchema; @@ -675,7 +688,7 @@ public static Function from(FunctionTool tool) { public JsonNode getInputSchema() { return inputSchema; } public Boolean getStrict() { return strict; } public Map getProviderOptions() { return providerOptions; } - public List getInputExamples() { return inputExamples; } + public List getInputExamples() { return inputExamples; } public static Builder builder() { return new Builder(); } @@ -685,14 +698,14 @@ public static class Builder { private JsonNode inputSchema = emptyObject(); private Boolean strict; private Map providerOptions; - private List inputExamples; + private List inputExamples; public Builder name(String v) { this.name = v; return this; } public Builder description(String v) { this.description = v; return this; } public Builder inputSchema(JsonNode v) { this.inputSchema = v; return this; } public Builder strict(Boolean v) { this.strict = v; return this; } public Builder providerOptions(Map v) { this.providerOptions = v; return this; } - public Builder inputExamples(List v) { this.inputExamples = v; return this; } + public Builder inputExamples(List v) { this.inputExamples = v; return this; } public Function build() { return new Function(name, description, inputSchema, strict, providerOptions, inputExamples); @@ -722,12 +735,12 @@ public int hashCode() { public static class Provider extends Tool { @JsonProperty("id") private String id = ""; @JsonProperty("name") private String name = ""; - @JsonProperty("args") private JsonNode args = emptyObject(); + @JsonProperty("args") private ObjectNode args = emptyObject(); @JsonCreator Provider() {} - private Provider(String id, String name, JsonNode args) { + private Provider(String id, String name, ObjectNode args) { this.id = id; this.name = name; this.args = args; @@ -740,18 +753,18 @@ public static Provider from(ProviderTool tool) { public String getId() { return id; } public String getName() { return name; } - public JsonNode getArgs() { return args; } + public ObjectNode getArgs() { return args; } public static Builder builder() { return new Builder(); } public static class Builder { private String id = ""; private String name = ""; - private JsonNode args = emptyObject(); + private ObjectNode args = emptyObject(); public Builder id(String v) { this.id = v; return this; } public Builder name(String v) { this.name = v; return this; } - public Builder args(JsonNode v) { this.args = v; return this; } + public Builder args(ObjectNode v) { this.args = v; return this; } public Provider build() { return new Provider(id, name, args); } } @@ -1273,19 +1286,17 @@ public static class ToolCall extends ContentPart { @JsonProperty("tool_name") private String toolName = ""; @JsonProperty("input") private JsonNode input = emptyObject(); @JsonProperty("provider_executed") private Boolean providerExecuted; - @JsonProperty("thought_signature") private String thoughtSignature; @JsonProperty("provider_options") private JsonNode providerOptions; @JsonCreator ToolCall() {} private ToolCall(String toolCallId, String toolName, JsonNode input, - Boolean providerExecuted, String thoughtSignature, JsonNode providerOptions) { + Boolean providerExecuted, JsonNode providerOptions) { this.toolCallId = toolCallId; this.toolName = toolName; this.input = input; this.providerExecuted = providerExecuted; - this.thoughtSignature = thoughtSignature; this.providerOptions = providerOptions; } @@ -1293,7 +1304,6 @@ private ToolCall(String toolCallId, String toolName, JsonNode input, public String getToolName() { return toolName; } public JsonNode getInput() { return input; } public Boolean getProviderExecuted() { return providerExecuted; } - public String getThoughtSignature() { return thoughtSignature; } public JsonNode getProviderOptions() { return providerOptions; } public static Builder builder() { return new Builder(); } @@ -1303,18 +1313,16 @@ public static class Builder { private String toolName = ""; private JsonNode input = emptyObject(); private Boolean providerExecuted; - private String thoughtSignature; private JsonNode providerOptions; public Builder toolCallId(String v) { this.toolCallId = v; return this; } public Builder toolName(String v) { this.toolName = v; return this; } public Builder input(JsonNode v) { this.input = v; return this; } public Builder providerExecuted(Boolean v) { this.providerExecuted = v; return this; } - public Builder thoughtSignature(String v) { this.thoughtSignature = v; return this; } public Builder providerOptions(JsonNode v) { this.providerOptions = v; return this; } public ToolCall build() { - return new ToolCall(toolCallId, toolName, input, providerExecuted, thoughtSignature, providerOptions); + return new ToolCall(toolCallId, toolName, input, providerExecuted, providerOptions); } } @@ -1327,13 +1335,12 @@ public boolean equals(Object o) { && Objects.equals(toolName, that.toolName) && Objects.equals(input, that.input) && Objects.equals(providerExecuted, that.providerExecuted) - && Objects.equals(thoughtSignature, that.thoughtSignature) && Objects.equals(providerOptions, that.providerOptions); } @Override public int hashCode() { - return Objects.hash(toolCallId, toolName, input, providerExecuted, thoughtSignature, providerOptions); + return Objects.hash(toolCallId, toolName, input, providerExecuted, providerOptions); } } @@ -1729,20 +1736,18 @@ public static class RawToolCall { @JsonProperty("input") private String input = ""; @JsonProperty("provider_executed") private Boolean providerExecuted; @JsonProperty("dynamic") private Boolean dynamic; - @JsonProperty("thought_signature") private String thoughtSignature; @JsonProperty("provider_metadata") private JsonNode providerMetadata; @JsonCreator RawToolCall() {} private RawToolCall(String toolCallId, String toolName, String input, Boolean providerExecuted, - Boolean dynamic, String thoughtSignature, JsonNode providerMetadata) { + Boolean dynamic, JsonNode providerMetadata) { this.toolCallId = toolCallId; this.toolName = toolName; this.input = input; this.providerExecuted = providerExecuted; this.dynamic = dynamic; - this.thoughtSignature = thoughtSignature; this.providerMetadata = providerMetadata; } @@ -1752,7 +1757,6 @@ private RawToolCall(String toolCallId, String toolName, String input, Boolean pr public String getInput() { return input; } public Boolean getProviderExecuted() { return providerExecuted; } public Boolean getDynamic() { return dynamic; } - public String getThoughtSignature() { return thoughtSignature; } public JsonNode getProviderMetadata() { return providerMetadata; } public static Builder builder() { return new Builder(); } @@ -1763,7 +1767,6 @@ public static class Builder { private String input = ""; private Boolean providerExecuted; private Boolean dynamic; - private String thoughtSignature; private JsonNode providerMetadata; public Builder toolCallId(String v) { this.toolCallId = v; return this; } @@ -1771,12 +1774,11 @@ public static class Builder { public Builder input(String v) { this.input = v; return this; } public Builder providerExecuted(Boolean v) { this.providerExecuted = v; return this; } public Builder dynamic(Boolean v) { this.dynamic = v; return this; } - public Builder thoughtSignature(String v) { this.thoughtSignature = v; return this; } public Builder providerMetadata(JsonNode v) { this.providerMetadata = v; return this; } public RawToolCall build() { return new RawToolCall(toolCallId, toolName, input, providerExecuted, dynamic, - thoughtSignature, providerMetadata); + providerMetadata); } } } @@ -2320,20 +2322,18 @@ public static class ToolCall extends GenerateContent { @JsonProperty("input") private JsonNode input = emptyObject(); @JsonProperty("provider_executed") private Boolean providerExecuted; @JsonProperty("dynamic") private Boolean dynamic; - @JsonProperty("thought_signature") private String thoughtSignature; @JsonProperty("provider_metadata") private JsonNode providerMetadata; @JsonCreator ToolCall() {} private ToolCall(String toolCallId, String toolName, JsonNode input, Boolean providerExecuted, - Boolean dynamic, String thoughtSignature, JsonNode providerMetadata) { + Boolean dynamic, JsonNode providerMetadata) { this.toolCallId = toolCallId; this.toolName = toolName; this.input = input; this.providerExecuted = providerExecuted; this.dynamic = dynamic; - this.thoughtSignature = thoughtSignature; this.providerMetadata = providerMetadata; } @@ -2342,7 +2342,6 @@ private ToolCall(String toolCallId, String toolName, JsonNode input, Boolean pro public JsonNode getInput() { return input; } public Boolean getProviderExecuted() { return providerExecuted; } public Boolean getDynamic() { return dynamic; } - public String getThoughtSignature() { return thoughtSignature; } public JsonNode getProviderMetadata() { return providerMetadata; } public static Builder builder() { return new Builder(); } @@ -2353,7 +2352,6 @@ public static class Builder { private JsonNode input = emptyObject(); private Boolean providerExecuted; private Boolean dynamic; - private String thoughtSignature; private JsonNode providerMetadata; public Builder toolCallId(String v) { this.toolCallId = v; return this; } @@ -2361,11 +2359,10 @@ public static class Builder { public Builder input(JsonNode v) { this.input = v; return this; } public Builder providerExecuted(Boolean v) { this.providerExecuted = v; return this; } public Builder dynamic(Boolean v) { this.dynamic = v; return this; } - public Builder thoughtSignature(String v) { this.thoughtSignature = v; return this; } public Builder providerMetadata(JsonNode v) { this.providerMetadata = v; return this; } public ToolCall build() { - return new ToolCall(toolCallId, toolName, input, providerExecuted, dynamic, thoughtSignature, + return new ToolCall(toolCallId, toolName, input, providerExecuted, dynamic, providerMetadata); } } @@ -2380,13 +2377,12 @@ public boolean equals(Object o) { && Objects.equals(input, that.input) && Objects.equals(providerExecuted, that.providerExecuted) && Objects.equals(dynamic, that.dynamic) - && Objects.equals(thoughtSignature, that.thoughtSignature) && Objects.equals(providerMetadata, that.providerMetadata); } @Override public int hashCode() { - return Objects.hash(toolCallId, toolName, input, providerExecuted, dynamic, thoughtSignature, + return Objects.hash(toolCallId, toolName, input, providerExecuted, dynamic, providerMetadata); } } @@ -3543,7 +3539,6 @@ public static class ToolCall extends StreamPart { @JsonProperty("input") private JsonNode input = emptyObject(); @JsonProperty("provider_executed") private Boolean providerExecuted; @JsonProperty("dynamic") private Boolean dynamic; - @JsonProperty("thought_signature") private String thoughtSignature; @JsonProperty("provider_metadata") private JsonNode providerMetadata; @JsonProperty("invalid") private Boolean invalid; @JsonProperty("error") private JsonNode error; @@ -3552,14 +3547,13 @@ public static class ToolCall extends StreamPart { ToolCall() {} private ToolCall(String toolCallId, String toolName, JsonNode input, Boolean providerExecuted, - Boolean dynamic, String thoughtSignature, JsonNode providerMetadata, Boolean invalid, + Boolean dynamic, JsonNode providerMetadata, Boolean invalid, JsonNode error) { this.toolCallId = toolCallId; this.toolName = toolName; this.input = input; this.providerExecuted = providerExecuted; this.dynamic = dynamic; - this.thoughtSignature = thoughtSignature; this.providerMetadata = providerMetadata; this.invalid = invalid; this.error = error; @@ -3570,7 +3564,6 @@ private ToolCall(String toolCallId, String toolName, JsonNode input, Boolean pro public JsonNode getInput() { return input; } public Boolean getProviderExecuted() { return providerExecuted; } public Boolean getDynamic() { return dynamic; } - public String getThoughtSignature() { return thoughtSignature; } public JsonNode getProviderMetadata() { return providerMetadata; } /** Set by Core when the tool call stays invalid after optional repair. */ public Boolean getInvalid() { return invalid; } @@ -3585,7 +3578,6 @@ public static class Builder { private JsonNode input = emptyObject(); private Boolean providerExecuted; private Boolean dynamic; - private String thoughtSignature; private JsonNode providerMetadata; private Boolean invalid; private JsonNode error; @@ -3595,13 +3587,12 @@ public static class Builder { public Builder input(JsonNode v) { this.input = v; return this; } public Builder providerExecuted(Boolean v) { this.providerExecuted = v; return this; } public Builder dynamic(Boolean v) { this.dynamic = v; return this; } - public Builder thoughtSignature(String v) { this.thoughtSignature = v; return this; } public Builder providerMetadata(JsonNode v) { this.providerMetadata = v; return this; } public Builder invalid(Boolean v) { this.invalid = v; return this; } public Builder error(JsonNode v) { this.error = v; return this; } public ToolCall build() { - return new ToolCall(toolCallId, toolName, input, providerExecuted, dynamic, thoughtSignature, + return new ToolCall(toolCallId, toolName, input, providerExecuted, dynamic, providerMetadata, invalid, error); } } @@ -3616,7 +3607,6 @@ public boolean equals(Object o) { && Objects.equals(input, that.input) && Objects.equals(providerExecuted, that.providerExecuted) && Objects.equals(dynamic, that.dynamic) - && Objects.equals(thoughtSignature, that.thoughtSignature) && Objects.equals(providerMetadata, that.providerMetadata) && Objects.equals(invalid, that.invalid) && Objects.equals(error, that.error); @@ -3624,7 +3614,7 @@ public boolean equals(Object o) { @Override public int hashCode() { - return Objects.hash(toolCallId, toolName, input, providerExecuted, dynamic, thoughtSignature, + return Objects.hash(toolCallId, toolName, input, providerExecuted, dynamic, providerMetadata, invalid, error); } } diff --git a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/ToolCallRepair.kt b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/ToolCallRepair.kt index 5e0f4f685..dd913f937 100644 --- a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/ToolCallRepair.kt +++ b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/ToolCallRepair.kt @@ -45,7 +45,6 @@ data class RawToolCall( val input: String, @SerialName("provider_executed") val providerExecuted: Boolean? = null, val dynamic: Boolean? = null, - @SerialName("thought_signature") val thoughtSignature: String? = null, @SerialName("provider_metadata") val providerMetadata: JsonElement? = null, ) diff --git a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Types.kt b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Types.kt index 2820ef158..d14e4fcfe 100644 --- a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Types.kt +++ b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Types.kt @@ -166,7 +166,6 @@ data class ToolCall( val input: JsonElement = JsonObject(emptyMap()), @SerialName("provider_executed") val providerExecuted: Boolean? = null, @SerialName("dynamic") val dynamic: Boolean? = null, - @SerialName("thought_signature") val thoughtSignature: String? = null, @SerialName("provider_metadata") val providerMetadata: JsonElement? = null, val invalid: Boolean? = null, val error: JsonElement? = null, @@ -180,6 +179,9 @@ data class ToolCall( // the `"type"` discriminator — no custom serializer needed. // ───────────────────────────────────────────────────────────────────────────── +@Serializable +data class FunctionToolInputExample(val input: JsonObject) + /** * A user-defined function tool. * @@ -192,7 +194,7 @@ data class FunctionTool( @SerialName("input_schema") val inputSchema: JsonElement, val strict: Boolean? = null, @SerialName("provider_options") val providerOptions: Map? = null, - @SerialName("input_examples") val inputExamples: List? = null, + @SerialName("input_examples") val inputExamples: List? = null, ) /** @@ -204,7 +206,7 @@ data class FunctionTool( data class ProviderTool( val id: String, val name: String, - val args: JsonElement = JsonObject(emptyMap()), + val args: JsonObject = JsonObject(emptyMap()), ) /** @@ -223,7 +225,7 @@ sealed interface Tool { @SerialName("input_schema") val inputSchema: JsonElement, val strict: Boolean? = null, @SerialName("provider_options") val providerOptions: Map? = null, - @SerialName("input_examples") val inputExamples: List? = null, + @SerialName("input_examples") val inputExamples: List? = null, ) : Tool { companion object { /** Convenience constructor from a [FunctionTool]. */ @@ -243,7 +245,7 @@ sealed interface Tool { data class Provider( val id: String, val name: String, - val args: JsonElement = JsonObject(emptyMap()), + val args: JsonObject = JsonObject(emptyMap()), ) : Tool { companion object { fun from(tool: ProviderTool): Provider = Provider(tool.id, tool.name, tool.args) @@ -391,7 +393,6 @@ sealed interface ContentPart { @SerialName("tool_call_id") val toolCallId: String = "", @SerialName("tool_name") val toolName: String = "", val input: JsonElement = JsonObject(emptyMap()), - @SerialName("thought_signature") val thoughtSignature: String? = null, @SerialName("provider_options") val providerOptions: JsonElement? = null, @SerialName("provider_executed") val providerExecuted: Boolean? = null, ) : ContentPart @@ -807,7 +808,6 @@ sealed interface GenerateContent { val input: JsonElement = JsonObject(emptyMap()), @SerialName("provider_executed") val providerExecuted: Boolean? = null, @SerialName("dynamic") val dynamic: Boolean? = null, - @SerialName("thought_signature") val thoughtSignature: String? = null, @SerialName("provider_metadata") val providerMetadata: JsonElement? = null, ) : GenerateContent @@ -1099,7 +1099,6 @@ sealed interface StreamPart { val input: JsonElement = JsonObject(emptyMap()), @SerialName("provider_executed") val providerExecuted: Boolean? = null, @SerialName("dynamic") val dynamic: Boolean? = null, - @SerialName("thought_signature") val thoughtSignature: String? = null, @SerialName("provider_metadata") val providerMetadata: JsonElement? = null, val invalid: Boolean? = null, val error: JsonElement? = null, diff --git a/bindings/node/src/lib.rs b/bindings/node/src/lib.rs index e4e98be9a..76dd6c89a 100644 --- a/bindings/node/src/lib.rs +++ b/bindings/node/src/lib.rs @@ -20,7 +20,7 @@ use std::future::Future; use std::sync::Arc; use crate::error::{ - AimuxResult, BindingError, AiMuxBindingError, MResult, StreamItem, parse_wire_json, + AiMuxBindingError, AimuxResult, BindingError, MResult, StreamItem, parse_wire_json, serialize_result, }; use aimux_core::AiMuxError; @@ -327,8 +327,9 @@ impl Model { // move into the spawned task. let abort_signal = bridge.map(|b| b.core_signal()); - let (tx, rx) = - tokio::sync::mpsc::channel::>(64); + let (tx, rx) = tokio::sync::mpsc::channel::< + std::result::Result, + >(64); // Spawn the stream-driving task immediately on napi's tokio runtime. napi::tokio::spawn(async move { @@ -379,7 +380,9 @@ impl Model { // frame; deliver it as a StreamPart::Error data item // and keep pumping. match serde_json::to_string( - &aimux_core::stream_part::StreamPart::Error { error: e }, + &aimux_core::stream_part::TextStreamPart::Error { + error: e, + }, ) { Ok(json) => { if tx.send(Ok(json)).await.is_err() { @@ -471,8 +474,9 @@ impl Model { let model = self.inner.clone(); let abort_signal = bridge.map(|b| b.core_signal()); - let (tx, rx) = - tokio::sync::mpsc::channel::>(64); + let (tx, rx) = tokio::sync::mpsc::channel::< + std::result::Result, + >(64); napi::tokio::spawn(async move { let prompt = match parse_prompt(&prompt) { @@ -565,6 +569,8 @@ impl Model { } } +type StreamReceiver = tokio::sync::mpsc::Receiver>; + /// AsyncGenerator that yields StreamPart JSON strings. /// /// The stream is started eagerly in `stream_text()` — a tokio task drives @@ -572,11 +578,7 @@ impl Model { /// receives from the channel. #[napi(async_iterator)] pub struct StreamTextGenerator { - rx: std::sync::Arc< - tokio::sync::Mutex< - Option>>, - >, - >, + rx: std::sync::Arc>>, } #[napi] @@ -945,19 +947,19 @@ struct RouterFfiConfig { /// or no store is registered. #[napi] pub fn session_calls(session_id: String) -> AimuxResult { - AimuxResult((|| -> crate::error::MResult { + AimuxResult({ let calls = aimux_core::session::session_calls(&session_id); serialize_result(&calls) - })()) + }) } /// Query: all known sessions (RFC-0024), as a JSON-serialized `SessionView[]`. #[napi] pub fn list_sessions() -> AimuxResult { - AimuxResult((|| -> crate::error::MResult { + AimuxResult({ let views = aimux_core::session::list_sessions(); serialize_result(&views) - })()) + }) } /// Create an OpenAI model instance. @@ -1377,10 +1379,10 @@ pub async fn azure( } None => {} } - if let Some(version) = api_version { - if !version.is_empty() { - cfg = cfg.with_api_version(version); - } + if let Some(version) = api_version + && !version.is_empty() + { + cfg = cfg.with_api_version(version); } if !resource_name.is_empty() { cfg = cfg.with_resource_name(resource_name); diff --git a/bindings/node/src/multimodal.rs b/bindings/node/src/multimodal.rs index 1d6095704..00666cc09 100644 --- a/bindings/node/src/multimodal.rs +++ b/bindings/node/src/multimodal.rs @@ -3,7 +3,9 @@ //! Each modality is a napi class wrapping the Rust trait object. //! All cross-boundary data uses JSON strings (base64 for binary). -use crate::error::{AimuxResult, BindingError, AiMuxBindingError, parse_wire_json, serialize_result}; +use crate::error::{ + AiMuxBindingError, AimuxResult, BindingError, parse_wire_json, serialize_result, +}; use std::sync::Arc; use aimux_core::AiMuxError; @@ -55,8 +57,7 @@ fn parse_opts_json( /// Shape-only stand-in for `TranscriptionCallOptions::audio`, built from the /// real enum so a variant rename is a compile error rather than a runtime one. fn audio_placeholder() -> serde_json::Value { - serde_json::to_value(AudioInput::Base64(String::new())) - .expect("AudioInput always serializes") + serde_json::to_value(AudioInput::Base64(String::new())).expect("AudioInput always serializes") } /// Shape-only stand-in for `RerankingCallOptions::documents`. @@ -208,27 +209,27 @@ impl TranscriptionModel { let __r: crate::error::MResult = async { let mut opts = TranscriptionCallOptions::new(AudioInput::Base64(audio_base64), media_type); - if let Some(s) = opts_json.as_deref() { - if !s.trim().is_empty() && s.trim() != "null" { - let mut parsed: TranscriptionCallOptions = parse_opts_json( - "opts_json", - s, - &[ - ("audio", audio_placeholder()), - ("media_type", serde_json::Value::from("")), - ], - )?; - // Required data comes from the explicit arguments; all - // operation policy and optional fields come from JSON. - parsed.audio = opts.audio; - parsed.media_type = opts.media_type; - opts = parsed; - } + if let Some(s) = opts_json.as_deref() + && !s.trim().is_empty() + && s.trim() != "null" + { + let mut parsed: TranscriptionCallOptions = parse_opts_json( + "opts_json", + s, + &[ + ("audio", audio_placeholder()), + ("media_type", serde_json::Value::from("")), + ], + )?; + // Required data comes from the explicit arguments; all + // operation policy and optional fields come from JSON. + parsed.audio = opts.audio; + parsed.media_type = opts.media_type; + opts = parsed; } opts.abort_signal = bridge.map(|b| b.core_signal()); - let result = - aimux_core::transcription_model::transcribe(self.inner.as_ref(), opts) + let result = aimux_core::transcription_model::transcribe(self.inner.as_ref(), opts) .await .map_err(|e| AiMuxBindingError::from(&e))?; serialize_result(&result) @@ -268,22 +269,23 @@ impl RerankingModel { let docs: RerankingDocuments = parse_wire_json("docs_json", &docs_json)?; let mut opts = RerankingCallOptions::new(query, docs); - if let Some(s) = opts_json.as_deref() { - if !s.trim().is_empty() && s.trim() != "null" { - let mut parsed: RerankingCallOptions = parse_opts_json( - "opts_json", - s, - &[ - ("query", serde_json::Value::from("")), - ("documents", documents_placeholder()), - ], - )?; - // Required data comes from the explicit arguments; all - // operation policy and optional fields come from JSON. - parsed.query = opts.query; - parsed.documents = opts.documents; - opts = parsed; - } + if let Some(s) = opts_json.as_deref() + && !s.trim().is_empty() + && s.trim() != "null" + { + let mut parsed: RerankingCallOptions = parse_opts_json( + "opts_json", + s, + &[ + ("query", serde_json::Value::from("")), + ("documents", documents_placeholder()), + ], + )?; + // Required data comes from the explicit arguments; all + // operation policy and optional fields come from JSON. + parsed.query = opts.query; + parsed.documents = opts.documents; + opts = parsed; } opts.abort_signal = bridge.map(|b| b.core_signal()); @@ -359,18 +361,16 @@ impl SearchModel { let __r: crate::error::MResult = async { use aimux_core::search_model::SearchCallOptions; let mut opts = SearchCallOptions::new(query); - if let Some(s) = opts_json.as_deref() { - if !s.trim().is_empty() && s.trim() != "null" { - // Required data comes from the explicit arguments; all - // operation policy and optional fields come from JSON. - let mut parsed: SearchCallOptions = parse_opts_json( - "opts_json", - s, - &[("query", serde_json::Value::from(""))], - )?; - parsed.query = opts.query; - opts = parsed; - } + if let Some(s) = opts_json.as_deref() + && !s.trim().is_empty() + && s.trim() != "null" + { + // Required data comes from the explicit arguments; all + // operation policy and optional fields come from JSON. + let mut parsed: SearchCallOptions = + parse_opts_json("opts_json", s, &[("query", serde_json::Value::from(""))])?; + parsed.query = opts.query; + opts = parsed; } opts.abort_signal = bridge.map(|b| b.core_signal()); @@ -416,12 +416,13 @@ impl Files { }, media_type, ); - if let Some(s) = opts_json.as_deref() { - if !s.trim().is_empty() && s.trim() != "null" { - let parsed: UploadFileCallOptions = parse_wire_json("opts_json", s)?; - opts.filename = parsed.filename; - opts.provider_options = parsed.provider_options; - } + if let Some(s) = opts_json.as_deref() + && !s.trim().is_empty() + && s.trim() != "null" + { + let parsed: UploadFileCallOptions = parse_wire_json("opts_json", s)?; + opts.filename = parsed.filename; + opts.provider_options = parsed.provider_options; } let result = self @@ -719,8 +720,9 @@ pub async fn tavily_search(api_key: String, base_url: Option) -> AimuxRe #[napi] pub struct TranscriptionSession { audio_tx: std::sync::Mutex>>, - parts_rx: - tokio::sync::Mutex>>, + parts_rx: tokio::sync::Mutex< + tokio::sync::mpsc::Receiver>, + >, token: aimux_core::AbortSignal, } @@ -827,46 +829,42 @@ pub async fn start_transcription_session( return; } }; - loop { - match tx.try_send(Ok(json.clone())) { - Ok(()) => break, - Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { - // Full: block on send (waits for - // capacity) unless aborted. - tokio::select! { - _ = effective.cancelled() => return, - res = tx.send(Ok(json.clone())) => { - if res.is_err() { return; } - break; - } + match tx.try_send(Ok(json.clone())) { + Ok(()) => {}, + Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { + // Full: block on send (waits for + // capacity) unless aborted. + tokio::select! { + _ = effective.cancelled() => return, + res = tx.send(Ok(json.clone())) => { + if res.is_err() { return; } } } - Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => { - return; - } + } + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => { + return; } } + } } Err(e) => { // Connect failure: deliver as the first channel item. // try_send + abort-select (a full channel must not // stall; abort covers the session-drop path). - loop { - match tx.try_send(Err(AiMuxBindingError::from(&e))) { - Ok(()) => break, - Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { - tokio::select! { - _ = effective.cancelled() => return, - res = tx.send(Err(AiMuxBindingError::from(&e))) => { - if res.is_err() { return; } - break; - } + match tx.try_send(Err(AiMuxBindingError::from(&e))) { + Ok(()) => {}, + Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { + tokio::select! { + _ = effective.cancelled() => return, + res = tx.send(Err(AiMuxBindingError::from(&e))) => { + if res.is_err() { return; } } } - Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => return, } + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => return, } + } } }); diff --git a/bindings/node/src/types.ts b/bindings/node/src/types.ts index 6b86bf3bd..07041be8f 100644 --- a/bindings/node/src/types.ts +++ b/bindings/node/src/types.ts @@ -49,6 +49,7 @@ export type * from './types/Fingerprint' export type * from './types/FinishReason' export type * from './types/FinishReasonUnified' export type * from './types/FunctionTool' +export type * from './types/FunctionToolInputExample' export type * from './types/GenerateContent' export type * from './types/GenerateObjectResult' export type * from './types/GenerateResult' diff --git a/bindings/node/src/types/CallOptions.ts b/bindings/node/src/types/CallOptions.ts index 53c072370..40f9151ad 100644 --- a/bindings/node/src/types/CallOptions.ts +++ b/bindings/node/src/types/CallOptions.ts @@ -62,7 +62,7 @@ tools: Array | null, /** * How the model should choose tools. */ -tool_choice: ToolChoice, +tool_choice: ToolChoice | null, /** * Extra HTTP headers. */ diff --git a/bindings/node/src/types/ContentPart.ts b/bindings/node/src/types/ContentPart.ts index a7538a154..16ca615db 100644 --- a/bindings/node/src/types/ContentPart.ts +++ b/bindings/node/src/types/ContentPart.ts @@ -27,12 +27,6 @@ input: JsonValue, * need it to replay server tool calls on a later turn. */ provider_executed?: boolean | null, -/** - * Provider-assigned thought signature (e.g. Google Gemini - * `thoughtSignature`). Must be echoed back verbatim on the follow-up - * turn when the tool result is sent. - */ -thought_signature?: string | null, /** * Provider-specific options for this part (e.g. * `anthropic.cacheControl`). diff --git a/bindings/node/src/types/FunctionTool.ts b/bindings/node/src/types/FunctionTool.ts index a79887e0c..efdcfaccf 100644 --- a/bindings/node/src/types/FunctionTool.ts +++ b/bindings/node/src/types/FunctionTool.ts @@ -1,4 +1,5 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { FunctionToolInputExample } from "./FunctionToolInputExample"; import type { JsonValue } from "./serde_json/JsonValue"; /** @@ -31,4 +32,4 @@ provider_options?: { [key in string]: JsonValue } | null, * Example inputs for the tool (V4 `inputExamples`), used by some providers * to emit `input_examples` in the request body. */ -input_examples?: Array | null, }; +input_examples?: Array | null, }; diff --git a/bindings/node/src/types/FunctionToolInputExample.ts b/bindings/node/src/types/FunctionToolInputExample.ts new file mode 100644 index 000000000..588fc7ad4 --- /dev/null +++ b/bindings/node/src/types/FunctionToolInputExample.ts @@ -0,0 +1,7 @@ +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { JsonValue } from "./serde_json/JsonValue"; + +/** + * An example input for a function tool. + */ +export type FunctionToolInputExample = { input: { [key in string]: JsonValue }, }; diff --git a/bindings/node/src/types/GenerateContent.ts b/bindings/node/src/types/GenerateContent.ts index f6da76302..aa72530dc 100644 --- a/bindings/node/src/types/GenerateContent.ts +++ b/bindings/node/src/types/GenerateContent.ts @@ -1,55 +1,10 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. import type { FileData } from "./FileData"; +import type { RawToolCall } from "./RawToolCall"; +import type { ToolResult } from "./ToolResult"; import type { JsonValue } from "./serde_json/JsonValue"; /** * A content item in the generation result. */ -export type GenerateContent = { "Text": { text: string, provider_metadata?: JsonValue | null, } } | { "ToolCall": { tool_call_id: string, tool_name: string, -/** - * The model's raw argument text, exactly as the provider delivered - * it (possibly malformed). Providers never parse it — `generate_text` - * owns parsing, schema validation, and repair. - * - * Always serializes as a JSON string. Deserializes a JSON string - * (the current wire shape) unchanged, and also re-serializes the - * pre-refactor shape — an already-parsed JSON value — to its compact - * JSON text, so a `GenerateResult` persisted before this field - * became a `String` keeps loading. See docs/api/gaps.md §9. - */ -input: string, -/** - * Whether the tool call will be executed by the provider. - * If false/unset, the tool call is executed by the client. - */ -provider_executed?: boolean | null, -/** - * Whether the tool is dynamic (defined at runtime, e.g. MCP tools). - */ -dynamic?: boolean | null, -/** - * Provider-assigned thought signature (e.g. Google Gemini - * `thoughtSignature`). Must be echoed back verbatim on the follow-up - * turn when the tool result is sent. - */ -thought_signature?: string | null, -/** - * Additional provider-specific metadata for the tool call. - */ -provider_metadata?: JsonValue | null, } } | { "Source": { id: string, source_type: string, url: string | null, title: string | null, provider_metadata?: JsonValue | null, } } | { "Reasoning": { text: string, provider_metadata: JsonValue | null, } } | { "File": { data: FileData, media_type: string, provider_metadata?: JsonValue | null, } } | { "ToolResult": { tool_call_id: string, tool_name: string, result: JsonValue, -/** - * Whether the result is an error or error message. - */ -is_error?: boolean | null, -/** - * Whether the result is preliminary (replaces prior, e.g. image previews). - */ -preliminary?: boolean | null, -/** - * Whether the tool is dynamic (defined at runtime, e.g. MCP tools). - */ -dynamic?: boolean | null, -/** - * Additional provider-specific metadata for the tool result. - */ -provider_metadata?: JsonValue | null, } }; +export type GenerateContent = { "Text": { text: string, provider_metadata?: JsonValue | null, } } | { "ToolCall": RawToolCall } | { "Source": { id: string, source_type: string, url: string | null, title: string | null, provider_metadata?: JsonValue | null, } } | { "Reasoning": { text: string, provider_metadata: JsonValue | null, } } | { "File": { data: FileData, media_type: string, provider_metadata?: JsonValue | null, } } | { "ToolResult": ToolResult }; diff --git a/bindings/node/src/types/ProviderTool.ts b/bindings/node/src/types/ProviderTool.ts index c7770fef8..425e83930 100644 --- a/bindings/node/src/types/ProviderTool.ts +++ b/bindings/node/src/types/ProviderTool.ts @@ -22,4 +22,4 @@ name: string, * The arguments for configuring the tool. Must match the expected * arguments defined by the provider for this tool. */ -args: JsonValue, }; +args: { [key in string]: JsonValue }, }; diff --git a/bindings/node/src/types/RawToolCall.ts b/bindings/node/src/types/RawToolCall.ts index bab390b0e..64adeade3 100644 --- a/bindings/node/src/types/RawToolCall.ts +++ b/bindings/node/src/types/RawToolCall.ts @@ -7,4 +7,4 @@ import type { JsonValue } from "./serde_json/JsonValue"; * The wire shape matches [`ToolCall`] field for field, except that `input` * is the provider's raw argument *text* rather than a parsed value. */ -export type RawToolCall = { tool_call_id: string, tool_name: string, input: string, provider_executed?: boolean | null, dynamic?: boolean | null, thought_signature?: string | null, provider_metadata?: JsonValue | null, }; +export type RawToolCall = { tool_call_id: string, tool_name: string, input: string, provider_executed?: boolean | null, dynamic?: boolean | null, provider_metadata?: JsonValue | null, }; diff --git a/bindings/node/src/types/StreamPart.ts b/bindings/node/src/types/StreamPart.ts index a04fe8385..153fddc19 100644 --- a/bindings/node/src/types/StreamPart.ts +++ b/bindings/node/src/types/StreamPart.ts @@ -2,12 +2,18 @@ import type { AiMuxError } from "./AiMuxError"; import type { FileData } from "./FileData"; import type { FinishReason } from "./FinishReason"; +import type { ToolCall } from "./ToolCall"; +import type { ToolResult } from "./ToolResult"; import type { Usage } from "./Usage"; import type { Warning } from "./Warning"; import type { JsonValue } from "./serde_json/JsonValue"; /** * A single chunk in the stream returned by `do_stream`. + * + * Generic over the tool-call payload: providers emit the default + * [`StreamPart`] (raw argument text); `stream_text` emits [`TextStreamPart`] + * (parsed and validated [`ToolCall`]). */ export type StreamPart = { "TextStart": { id: string, provider_metadata?: JsonValue | null, } } | { "TextDelta": { id: string, delta: string, provider_metadata?: JsonValue | null, } } | { "TextEnd": { id: string, provider_metadata?: JsonValue | null, } } | { "StreamStart": { warnings: Array, } } | { "Finish": { finish_reason: FinishReason, usage: Usage, provider_metadata: JsonValue | null, } } | { "Error": { error: AiMuxError, } } | { "ToolInputStart": { id: string, tool_name: string, /** @@ -21,50 +27,7 @@ dynamic?: boolean | null, /** * Optional title for the tool call. */ -title?: string | null, provider_metadata?: JsonValue | null, } } | { "ToolInputDelta": { id: string, delta: string, provider_metadata?: JsonValue | null, } } | { "ToolInputEnd": { id: string, provider_metadata?: JsonValue | null, } } | { "ToolCall": { tool_call_id: string, tool_name: string, -/** - * Serialized argument text in a `Value::String` from `do_stream`; - * parsed input after `stream_text`. - */ -input: JsonValue, -/** - * Whether the tool call will be executed by the provider. - */ -provider_executed?: boolean | null, -/** - * Whether the tool is dynamic (defined at runtime, e.g. MCP tools). - */ -dynamic?: boolean | null, -/** - * Provider-assigned thought signature (e.g. Google Gemini - * `thoughtSignature`). Must be echoed back verbatim on the follow-up - * turn when the tool result is sent. - */ -thought_signature?: string | null, -/** - * Set by Core when the call remains invalid after optional repair. - */ -invalid?: boolean | null, -/** - * Typed lookup, parsing, schema, or repair failure. - */ -error?: AiMuxError | null, provider_metadata?: JsonValue | null, } } | { "ToolResult": { tool_call_id: string, tool_name: string, result: JsonValue, -/** - * Whether the result is an error or error message. - */ -is_error?: boolean | null, -/** - * Whether the result is preliminary (replaces prior, e.g. image previews). - */ -preliminary?: boolean | null, -/** - * Whether the tool is dynamic (defined at runtime, e.g. MCP tools). - */ -dynamic?: boolean | null, -/** - * Additional provider-specific metadata for the tool result. - */ -provider_metadata?: JsonValue | null, } } | { "File": { data: FileData, media_type: string, provider_metadata?: JsonValue | null, } } | { "ReasoningStart": { id: string, +title?: string | null, provider_metadata?: JsonValue | null, } } | { "ToolInputDelta": { id: string, delta: string, provider_metadata?: JsonValue | null, } } | { "ToolInputEnd": { id: string, provider_metadata?: JsonValue | null, } } | { "ToolCall": ToolCall } | { "ToolResult": ToolResult } | { "File": { data: FileData, media_type: string, provider_metadata?: JsonValue | null, } } | { "ReasoningStart": { id: string, /** * Provider-specific metadata (e.g. xAI `itemId`). */ diff --git a/bindings/node/src/types/ToolCall.ts b/bindings/node/src/types/ToolCall.ts index c6b72530d..605cb733d 100644 --- a/bindings/node/src/types/ToolCall.ts +++ b/bindings/node/src/types/ToolCall.ts @@ -27,13 +27,6 @@ provider_executed?: boolean | null, * Whether the tool is dynamic (defined at runtime, e.g. MCP tools). */ dynamic?: boolean | null, -/** - * Provider-assigned thought signature (e.g. Google Gemini - * `thoughtSignature`). Must be echoed back verbatim on the follow-up turn - * when the tool result is sent; thinking models reject the request - * otherwise. - */ -thought_signature?: string | null, /** * Additional provider-specific metadata associated with this call. */ diff --git a/bindings/node/src/types/ToolResult.ts b/bindings/node/src/types/ToolResult.ts index 1398d1a46..9c5039b52 100644 --- a/bindings/node/src/types/ToolResult.ts +++ b/bindings/node/src/types/ToolResult.ts @@ -2,13 +2,14 @@ import type { JsonValue } from "./serde_json/JsonValue"; /** - * The result of executing a tool call, to be sent back to the model. + * A tool result produced by a provider-executed tool (e.g. xAI `file_search`, + * web search), carried next to the [`RawToolCall`] it answers. */ export type ToolResult = { /** - * Must match the corresponding `ToolCall::tool_call_id`. + * Must match the corresponding call's `tool_call_id`. */ -tool_call_id: string, +tool_call_id: string, tool_name: string, /** * The tool's output (usually a JSON-serializable value or plain text). */ @@ -20,4 +21,12 @@ is_error?: boolean | null, /** * Whether the result is preliminary (replaces prior, e.g. image previews). */ -preliminary?: boolean | null, }; +preliminary?: boolean | null, +/** + * Whether the tool is dynamic (defined at runtime, e.g. MCP tools). + */ +dynamic?: boolean | null, +/** + * Additional provider-specific metadata for the tool result. + */ +provider_metadata?: JsonValue | null, }; diff --git a/bindings/python/python/aimux/wrapper.py b/bindings/python/python/aimux/wrapper.py index 86e683262..a6e4e0fc2 100644 --- a/bindings/python/python/aimux/wrapper.py +++ b/bindings/python/python/aimux/wrapper.py @@ -54,6 +54,7 @@ "RepairToolCall", "ModelMessage", "FunctionTool", + "FunctionToolInputExample", "ProviderTool", "TextContentPart", "GenerateTextOptions", @@ -154,7 +155,6 @@ class ToolCall(BaseModel): input: Any provider_executed: Optional[bool] = None dynamic: Optional[bool] = None - thought_signature: Optional[str] = None provider_metadata: Optional[Any] = None invalid: Optional[bool] = None error: Optional[AiMuxErrorValue] = None @@ -173,7 +173,6 @@ class RawToolCall(BaseModel): input: str provider_executed: Optional[bool] = None dynamic: Optional[bool] = None - thought_signature: Optional[str] = None provider_metadata: Optional[Any] = None @@ -378,7 +377,6 @@ class _ContentToolCall(BaseModel): input: Any provider_executed: Optional[bool] = None dynamic: Optional[bool] = None - thought_signature: Optional[str] = None provider_metadata: Optional[Any] = None @@ -511,7 +509,6 @@ class _SPToolCall(BaseModel): input: Any provider_executed: Optional[bool] = None dynamic: Optional[bool] = None - thought_signature: Optional[str] = None provider_metadata: Optional[Dict[str, Any]] = None invalid: Optional[bool] = None error: Optional[AiMuxErrorValue] = None @@ -665,6 +662,10 @@ def _ser(self, handler): # type: ignore[no-untyped-def] # union with a ``type`` discriminator works directly. # ───────────────────────────────────────────────────────────────────────────── +class FunctionToolInputExample(BaseModel): + input: Dict[str, Any] + + class FunctionTool(BaseModel): """A user-defined function tool (Rust ``FunctionTool``).""" @@ -674,7 +675,7 @@ class FunctionTool(BaseModel): input_schema: Any strict: Optional[bool] = None provider_options: Optional[Dict[str, Any]] = None - input_examples: Optional[List[Any]] = None + input_examples: Optional[List[FunctionToolInputExample]] = None class ProviderTool(BaseModel): @@ -683,7 +684,7 @@ class ProviderTool(BaseModel): type: Literal["provider"] = "provider" id: str name: str - args: Any + args: Dict[str, Any] Tool = Annotated[ @@ -753,7 +754,6 @@ class _ToolCallContentPart(BaseModel): tool_name: str input: Any provider_executed: Optional[bool] = None - thought_signature: Optional[str] = None provider_options: Optional[Any] = None diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 527e468a4..834f118de 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -6,6 +6,8 @@ // pyo3 0.22 macros generate unsafe-op-in-unsafe-fn calls that trigger // edition-2024 lint. Suppress until pyo3 0.23+ lands. #![allow(unsafe_op_in_unsafe_fn)] +// PyO3 0.22 wraps PyResult returns in an identity conversion. +#![allow(clippy::useless_conversion)] mod decision; mod error; @@ -243,7 +245,7 @@ impl Model { // deliver it as a StreamPart::Error data item and keep // pumping. match serde_json::to_string( - &aimux_core::stream_part::StreamPart::Error { error: e }, + &aimux_core::stream_part::TextStreamPart::Error { error: e }, ) { Ok(json) => { if tx.send(Ok(json)).await.is_err() { @@ -410,7 +412,7 @@ impl StreamIterator { let item = py.allow_threads(|| runtime().block_on(self.rx.recv())); match item { - Some(Ok(json)) => Ok(Some(json.to_object(py).into())), + Some(Ok(json)) => Ok(Some(json.to_object(py))), Some(Err(f)) => Err(f.to_py_err()), None => Ok(None), // stream finished } @@ -663,10 +665,10 @@ fn azure( if let Some(url) = base_url { config = config.with_base_url(url); } - if let Some(version) = api_version { - if !version.is_empty() { - config = config.with_api_version(version); - } + if let Some(version) = api_version + && !version.is_empty() + { + config = config.with_api_version(version); } if !resource_name.is_empty() { config = config.with_resource_name(resource_name); @@ -786,6 +788,7 @@ fn create_provider( /// string (serialized `Catalogue`). `source_url` defaults to the anya2a /// `dist/all.json`. #[pyfunction] +#[pyo3(signature = (source_url=None))] fn get_model_specs(source_url: Option<&str>) -> PyResult { let catalogue = runtime() .block_on(async { aimux_providers::get_model_specs(source_url).await }) diff --git a/bindings/python/src/multimodal.rs b/bindings/python/src/multimodal.rs index 3c7289617..0bef00f90 100644 --- a/bindings/python/src/multimodal.rs +++ b/bindings/python/src/multimodal.rs @@ -14,7 +14,7 @@ use std::sync::Arc; use crate::error::{ - BindingError, AiMuxBindingError, binding_py_err, serialize_result, to_py_err, wire_error, + AiMuxBindingError, BindingError, binding_py_err, serialize_result, to_py_err, wire_error, wire_json, }; use aimux_core::AiMuxError; @@ -66,8 +66,7 @@ fn parse_opts_json( /// Shape-only stand-in for `TranscriptionCallOptions::audio`, built from the /// real enum so a variant rename is a compile error rather than a runtime one. fn audio_placeholder() -> serde_json::Value { - serde_json::to_value(AudioInput::Base64(String::new())) - .expect("AudioInput always serializes") + serde_json::to_value(AudioInput::Base64(String::new())).expect("AudioInput always serializes") } /// Shape-only stand-in for `RerankingCallOptions::documents`. @@ -192,22 +191,23 @@ impl TranscriptionModel { AudioInput::Base64(audio_base64.to_string()), media_type.to_string(), ); - if let Some(s) = opts_json { - if !s.trim().is_empty() && s.trim() != "null" { - // Take every caller option; audio and media_type always come - // from the explicit args. - let mut parsed: TranscriptionCallOptions = parse_opts_json( - "opts_json", - s, - &[ - ("audio", audio_placeholder()), - ("media_type", serde_json::Value::from("")), - ], - )?; - parsed.audio = opts.audio; - parsed.media_type = opts.media_type; - opts = parsed; - } + if let Some(s) = opts_json + && !s.trim().is_empty() + && s.trim() != "null" + { + // Take every caller option; audio and media_type always come + // from the explicit args. + let mut parsed: TranscriptionCallOptions = parse_opts_json( + "opts_json", + s, + &[ + ("audio", audio_placeholder()), + ("media_type", serde_json::Value::from("")), + ], + )?; + parsed.audio = opts.audio; + parsed.media_type = opts.media_type; + opts = parsed; } let result = crate::runtime().block_on(async move { @@ -246,22 +246,23 @@ impl RerankingModel { let docs: RerankingDocuments = wire_json("docs_json", docs_json)?; let mut opts = RerankingCallOptions::new(query.to_string(), docs); - if let Some(s) = opts_json { - if !s.trim().is_empty() && s.trim() != "null" { - // Take every caller option; query and documents always come - // from the explicit args. - let mut parsed: RerankingCallOptions = parse_opts_json( - "opts_json", - s, - &[ - ("query", serde_json::Value::from("")), - ("documents", documents_placeholder()), - ], - )?; - parsed.query = opts.query; - parsed.documents = opts.documents; - opts = parsed; - } + if let Some(s) = opts_json + && !s.trim().is_empty() + && s.trim() != "null" + { + // Take every caller option; query and documents always come + // from the explicit args. + let mut parsed: RerankingCallOptions = parse_opts_json( + "opts_json", + s, + &[ + ("query", serde_json::Value::from("")), + ("documents", documents_placeholder()), + ], + )?; + parsed.query = opts.query; + parsed.documents = opts.documents; + opts = parsed; } let result = crate::runtime().block_on(async move { @@ -318,18 +319,16 @@ impl SearchModel { #[pyo3(signature = (query, opts_json=None))] pub fn search(&self, query: &str, opts_json: Option<&str>) -> PyResult { let mut opts = SearchCallOptions::new(query.to_string()); - if let Some(s) = opts_json { - if !s.trim().is_empty() && s.trim() != "null" { - // Take every caller option; the query always comes from the - // explicit arg. - let mut parsed: SearchCallOptions = parse_opts_json( - "opts_json", - s, - &[("query", serde_json::Value::from(""))], - )?; - parsed.query = opts.query; - opts = parsed; - } + if let Some(s) = opts_json + && !s.trim().is_empty() + && s.trim() != "null" + { + // Take every caller option; the query always comes from the + // explicit arg. + let mut parsed: SearchCallOptions = + parse_opts_json("opts_json", s, &[("query", serde_json::Value::from(""))])?; + parsed.query = opts.query; + opts = parsed; } let result = crate::runtime().block_on(async move { @@ -372,12 +371,13 @@ impl Files { }, media_type.to_string(), ); - if let Some(s) = opts_json { - if !s.trim().is_empty() && s.trim() != "null" { - let parsed: UploadFileCallOptions = wire_json("opts_json", s)?; - opts.filename = parsed.filename; - opts.provider_options = parsed.provider_options; - } + if let Some(s) = opts_json + && !s.trim().is_empty() + && s.trim() != "null" + { + let parsed: UploadFileCallOptions = wire_json("opts_json", s)?; + opts.filename = parsed.filename; + opts.provider_options = parsed.provider_options; } let result = crate::runtime().block_on(async move { self.inner.upload_file(&opts).await }); @@ -672,8 +672,8 @@ pub fn start_transcription_session( include_raw_chunks: opts.include_raw_chunks.unwrap_or(false), timeout: opts.timeout, }; - let result = aimux_core::transcription_model::stream_transcribe(model.as_ref(), options) - .await; + let result = + aimux_core::transcription_model::stream_transcribe(model.as_ref(), options).await; match result { Ok(stream_result) => { use futures::StreamExt; @@ -712,21 +712,18 @@ pub fn start_transcription_session( return; } }; - loop { - match tx.try_send(Ok(json.clone())) { - Ok(()) => break, - Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { - tokio::select! { - _ = effective.cancelled() => return, - res = tx.send(Ok(json.clone())) => { - if res.is_err() { return; } - break; - } + match tx.try_send(Ok(json.clone())) { + Ok(()) => {} + Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { + tokio::select! { + _ = effective.cancelled() => return, + res = tx.send(Ok(json.clone())) => { + if res.is_err() { return; } } } - Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => { - return; - } + } + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => { + return; } } } @@ -734,20 +731,17 @@ pub fn start_transcription_session( Err(e) => { // Connect failure: deliver as the first channel item // (try_send + abort-select; mirrors the FFI driver). - loop { - match tx.try_send(Err(e.clone().into())) { - Ok(()) => break, - Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { - tokio::select! { - _ = effective.cancelled() => return, - res = tx.send(Err(e.clone().into())) => { - if res.is_err() { return; } - break; - } + match tx.try_send(Err(e.clone().into())) { + Ok(()) => {} + Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { + tokio::select! { + _ = effective.cancelled() => return, + res = tx.send(Err(e.clone().into())) => { + if res.is_err() { return; } } } - Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => return, } + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => return, } } } diff --git a/bindings/swift/Sources/Aimux/ToolCallRepair.swift b/bindings/swift/Sources/Aimux/ToolCallRepair.swift index 7be8e8141..af5bd09a7 100644 --- a/bindings/swift/Sources/Aimux/ToolCallRepair.swift +++ b/bindings/swift/Sources/Aimux/ToolCallRepair.swift @@ -33,7 +33,6 @@ public struct RawToolCall: Codable, Equatable { public var input: String public var providerExecuted: Bool? public var dynamic: Bool? - public var thoughtSignature: String? public var providerMetadata: JSONValue? enum CodingKeys: String, CodingKey { @@ -42,16 +41,15 @@ public struct RawToolCall: Codable, Equatable { case input case providerExecuted = "provider_executed" case dynamic - case thoughtSignature = "thought_signature" case providerMetadata = "provider_metadata" } public init(toolCallId: String, toolName: String, input: String, providerExecuted: Bool? = nil, dynamic: Bool? = nil, - thoughtSignature: String? = nil, providerMetadata: JSONValue? = nil) { + providerMetadata: JSONValue? = nil) { self.toolCallId = toolCallId; self.toolName = toolName; self.input = input self.providerExecuted = providerExecuted; self.dynamic = dynamic - self.thoughtSignature = thoughtSignature; self.providerMetadata = providerMetadata + self.providerMetadata = providerMetadata } } diff --git a/bindings/swift/Sources/Aimux/Types.swift b/bindings/swift/Sources/Aimux/Types.swift index 13b0f41d4..529bd71ac 100644 --- a/bindings/swift/Sources/Aimux/Types.swift +++ b/bindings/swift/Sources/Aimux/Types.swift @@ -251,6 +251,12 @@ public enum Warning: Codable, Equatable { // MARK: - Tool / ToolChoice +public struct FunctionToolInputExample: Codable, Equatable { + public var input: [String: JSONValue] + + public init(input: [String: JSONValue]) { self.input = input } +} + /// A user-defined function tool definition. public struct FunctionTool: Codable, Equatable { public var name: String @@ -259,7 +265,7 @@ public struct FunctionTool: Codable, Equatable { public var inputSchema: JSONValue public var strict: Bool? public var providerOptions: JSONValue? - public var inputExamples: [JSONValue]? + public var inputExamples: [FunctionToolInputExample]? enum CodingKeys: String, CodingKey { case name, description @@ -271,7 +277,7 @@ public struct FunctionTool: Codable, Equatable { public init(name: String, inputSchema: JSONValue, description: String? = nil, strict: Bool? = nil, providerOptions: JSONValue? = nil, - inputExamples: [JSONValue]? = nil) { + inputExamples: [FunctionToolInputExample]? = nil) { self.name = name; self.inputSchema = inputSchema; self.description = description self.strict = strict; self.providerOptions = providerOptions; self.inputExamples = inputExamples } @@ -281,9 +287,9 @@ public struct FunctionTool: Codable, Equatable { public struct ProviderTool: Codable, Equatable { public var id: String public var name: String - public var args: JSONValue + public var args: [String: JSONValue] - public init(id: String, name: String, args: JSONValue) { + public init(id: String, name: String, args: [String: JSONValue]) { self.id = id; self.name = name; self.args = args } } @@ -307,12 +313,12 @@ public enum Tool: Codable, Equatable { description: try c.decodeIfPresent(String.self, forKey: AnyCodingKey("description")), strict: try c.decodeIfPresent(Bool.self, forKey: AnyCodingKey("strict")), providerOptions: try c.decodeIfPresent(JSONValue.self, forKey: AnyCodingKey("provider_options")), - inputExamples: try c.decodeIfPresent([JSONValue].self, forKey: AnyCodingKey("input_examples")))) + inputExamples: try c.decodeIfPresent([FunctionToolInputExample].self, forKey: AnyCodingKey("input_examples")))) case "provider": self = .provider(ProviderTool( id: try c.decode(String.self, forKey: AnyCodingKey("id")), name: try c.decode(String.self, forKey: AnyCodingKey("name")), - args: try c.decode(JSONValue.self, forKey: AnyCodingKey("args")))) + args: try c.decode([String: JSONValue].self, forKey: AnyCodingKey("args")))) default: throw aimuxDecodingError(c.codingPath, "unknown tool type \(type)") } @@ -436,7 +442,7 @@ public enum ContentPart: Codable, Equatable { case fileReference(mediaType: String, reference: JSONValue, filename: String?, providerOptions: JSONValue?) case reasoning(text: String, signature: String?, providerOptions: JSONValue?) case toolCall(toolCallId: String, toolName: String, input: JSONValue, - providerExecuted: Bool? = nil, thoughtSignature: String? = nil, + providerExecuted: Bool? = nil, providerOptions: JSONValue?) case toolResult(toolCallId: String, result: JSONValue, toolName: String?, isError: Bool?, preliminary: Bool?, dynamic: Bool?, providerOptions: JSONValue?) @@ -444,7 +450,7 @@ public enum ContentPart: Codable, Equatable { private enum Field: String, CodingKey { case text, image, data, mediaType = "media_type", filename, url, reference case signature, toolCallId = "tool_call_id", toolName = "tool_name", input, result - case providerExecuted = "provider_executed", thoughtSignature = "thought_signature" + case providerExecuted = "provider_executed" case isError = "is_error", preliminary, dynamic case providerOptions = "provider_options" } @@ -489,7 +495,6 @@ public enum ContentPart: Codable, Equatable { toolName: try c.decode(String.self, forKey: AnyCodingKey("tool_name")), input: try c.decode(JSONValue.self, forKey: AnyCodingKey("input")), providerExecuted: try c.decodeIfPresent(Bool.self, forKey: AnyCodingKey("provider_executed")), - thoughtSignature: try c.decodeIfPresent(String.self, forKey: AnyCodingKey("thought_signature")), providerOptions: try po()) case "tool_result": self = .toolResult(toolCallId: try c.decode(String.self, forKey: AnyCodingKey("tool_call_id")), @@ -544,13 +549,12 @@ public enum ContentPart: Codable, Equatable { try c.encode(text, forKey: AnyCodingKey("text")) try c.encodeIfPresent(signature, forKey: AnyCodingKey("signature")) try c.encodeIfPresent(po, forKey: AnyCodingKey("provider_options")) - case .toolCall(let toolCallId, let toolName, let input, let providerExecuted, let thoughtSignature, let po): + case .toolCall(let toolCallId, let toolName, let input, let providerExecuted, let po): try c.encode("tool_call", forKey: AnyCodingKey("type")) try c.encode(toolCallId, forKey: AnyCodingKey("tool_call_id")) try c.encode(toolName, forKey: AnyCodingKey("tool_name")) try c.encode(input, forKey: AnyCodingKey("input")) try c.encodeIfPresent(providerExecuted, forKey: AnyCodingKey("provider_executed")) - try c.encodeIfPresent(thoughtSignature, forKey: AnyCodingKey("thought_signature")) try c.encodeIfPresent(po, forKey: AnyCodingKey("provider_options")) case .toolResult(let toolCallId, let result, let toolName, let isError, let preliminary, let dynamic, let po): try c.encode("tool_result", forKey: AnyCodingKey("type")) @@ -633,9 +637,6 @@ public struct ToolCall: Codable, Equatable { public var input: JSONValue public var providerExecuted: Bool? public var dynamic: Bool? - /// Provider-assigned thought signature (e.g. Google Gemini - /// `thoughtSignature`); must be echoed back verbatim on follow-up turns. - public var thoughtSignature: String? /// Additional provider-specific metadata associated with this call. public var providerMetadata: JSONValue? /// Set by Core when the tool call stays invalid after optional repair. @@ -649,7 +650,6 @@ public struct ToolCall: Codable, Equatable { case input case providerExecuted = "provider_executed" case dynamic - case thoughtSignature = "thought_signature" case providerMetadata = "provider_metadata" case invalid case error @@ -657,12 +657,11 @@ public struct ToolCall: Codable, Equatable { public init(toolCallId: String, toolName: String, input: JSONValue, providerExecuted: Bool? = nil, dynamic: Bool? = nil, - thoughtSignature: String? = nil, providerMetadata: JSONValue? = nil, + providerMetadata: JSONValue? = nil, invalid: Bool? = nil, error: JSONValue? = nil) { self.toolCallId = toolCallId; self.toolName = toolName; self.input = input self.providerExecuted = providerExecuted; self.dynamic = dynamic - self.thoughtSignature = thoughtSignature self.providerMetadata = providerMetadata self.invalid = invalid; self.error = error } @@ -763,7 +762,7 @@ public enum FileData: Codable, Equatable { public enum GenerateContent: Codable, Equatable { case text(text: String, providerMetadata: JSONValue?) case toolCall(toolCallId: String, toolName: String, input: JSONValue, - providerExecuted: Bool?, dynamic: Bool?, thoughtSignature: String? = nil, + providerExecuted: Bool?, dynamic: Bool?, providerMetadata: JSONValue?) case source(id: String, sourceType: String, url: String?, title: String?, providerMetadata: JSONValue?) @@ -776,7 +775,7 @@ public enum GenerateContent: Codable, Equatable { case text case toolCallId = "tool_call_id", toolName = "tool_name", input, result case providerExecuted = "provider_executed", dynamic - case thoughtSignature = "thought_signature", providerMetadata = "provider_metadata" + case providerMetadata = "provider_metadata" case id, sourceType = "source_type", url, title case isError = "is_error", preliminary case data, mediaType = "media_type" @@ -798,7 +797,6 @@ public enum GenerateContent: Codable, Equatable { input: try n.decode(JSONValue.self, forKey: .input), providerExecuted: try n.decodeIfPresent(Bool.self, forKey: .providerExecuted), dynamic: try n.decodeIfPresent(Bool.self, forKey: .dynamic), - thoughtSignature: try n.decodeIfPresent(String.self, forKey: .thoughtSignature), providerMetadata: try n.decodeIfPresent(JSONValue.self, forKey: .providerMetadata)) case "Source": self = .source(id: try n.decode(String.self, forKey: .id), @@ -833,14 +831,13 @@ public enum GenerateContent: Codable, Equatable { var n = c.nestedContainer(keyedBy: Field.self, forKey: AnyCodingKey("Text")) try n.encode(text, forKey: .text) try n.encodeIfPresent(pm, forKey: .providerMetadata) - case .toolCall(let toolCallId, let toolName, let input, let pe, let dyn, let signature, let pm): + case .toolCall(let toolCallId, let toolName, let input, let pe, let dyn, let pm): var n = c.nestedContainer(keyedBy: Field.self, forKey: AnyCodingKey("ToolCall")) try n.encode(toolCallId, forKey: .toolCallId) try n.encode(toolName, forKey: .toolName) try n.encode(input, forKey: .input) try n.encodeIfPresent(pe, forKey: .providerExecuted) try n.encodeIfPresent(dyn, forKey: .dynamic) - try n.encodeIfPresent(signature, forKey: .thoughtSignature) try n.encodeIfPresent(pm, forKey: .providerMetadata) case .source(let id, let sourceType, let url, let title, let pm): var n = c.nestedContainer(keyedBy: Field.self, forKey: AnyCodingKey("Source")) @@ -1246,7 +1243,7 @@ public enum StreamPart: Codable, Equatable { case toolInputDelta(id: String, delta: String, providerMetadata: JSONValue?) case toolInputEnd(id: String, providerMetadata: JSONValue?) case toolCall(toolCallId: String, toolName: String, input: JSONValue, - providerExecuted: Bool?, dynamic: Bool?, thoughtSignature: String? = nil, + providerExecuted: Bool?, dynamic: Bool?, providerMetadata: JSONValue?, invalid: Bool?, error: JSONValue?) case toolResult(toolCallId: String, toolName: String, result: JSONValue, @@ -1268,7 +1265,7 @@ public enum StreamPart: Codable, Equatable { case finishReason = "finish_reason", providerMetadata = "provider_metadata" case error case toolName = "tool_name", toolCallId = "tool_call_id", input, result - case providerExecuted = "provider_executed", dynamic, thoughtSignature = "thought_signature", invalid + case providerExecuted = "provider_executed", dynamic, invalid case isError = "is_error", preliminary case timestamp, modelId = "model_id" case sourceType = "source_type", url, title @@ -1323,7 +1320,6 @@ public enum StreamPart: Codable, Equatable { input: try n.decode(JSONValue.self, forKey: .input), providerExecuted: try n.decodeIfPresent(Bool.self, forKey: .providerExecuted), dynamic: try n.decodeIfPresent(Bool.self, forKey: .dynamic), - thoughtSignature: try n.decodeIfPresent(String.self, forKey: .thoughtSignature), providerMetadata: try n.decodeIfPresent(JSONValue.self, forKey: .providerMetadata), invalid: try n.decodeIfPresent(Bool.self, forKey: .invalid), error: try n.decodeIfPresent(JSONValue.self, forKey: .error)) @@ -1401,12 +1397,11 @@ public enum StreamPart: Codable, Equatable { case .toolInputEnd(let id, let pm): var n = c.nestedContainer(keyedBy: Field.self, forKey: AnyCodingKey("ToolInputEnd")) try n.encode(id, forKey: .id); try n.encodeIfPresent(pm, forKey: .providerMetadata) - case .toolCall(let toolCallId, let toolName, let input, let pe, let dyn, let signature, let pm, let inv, let err): + case .toolCall(let toolCallId, let toolName, let input, let pe, let dyn, let pm, let inv, let err): var n = c.nestedContainer(keyedBy: Field.self, forKey: AnyCodingKey("ToolCall")) try n.encode(toolCallId, forKey: .toolCallId); try n.encode(toolName, forKey: .toolName) try n.encode(input, forKey: .input) try n.encodeIfPresent(pe, forKey: .providerExecuted); try n.encodeIfPresent(dyn, forKey: .dynamic) - try n.encodeIfPresent(signature, forKey: .thoughtSignature) try n.encodeIfPresent(pm, forKey: .providerMetadata) try n.encodeIfPresent(inv, forKey: .invalid); try n.encodeIfPresent(err, forKey: .error) case .toolResult(let toolCallId, let toolName, let result, let ie, let prel, let dyn, let pm): diff --git a/bindings/swift/Tests/AimuxTests/WrapperTests.swift b/bindings/swift/Tests/AimuxTests/WrapperTests.swift index 3241b0309..428579bbf 100644 --- a/bindings/swift/Tests/AimuxTests/WrapperTests.swift +++ b/bindings/swift/Tests/AimuxTests/WrapperTests.swift @@ -89,7 +89,7 @@ final class WrapperTests: XCTestCase { // Structured raw.content contains a ToolCall variant mirroring the call. let toolContents = result.raw.content.compactMap { part -> ToolCall? in - if case .toolCall(let id, let name, let input, _, _, _, _) = part { + if case .toolCall(let id, let name, let input, _, _, _) = part { return ToolCall(toolCallId: id, toolName: name, input: input) } return nil @@ -254,7 +254,7 @@ final class WrapperTests: XCTestCase { // The complete ToolCall part carries the tool name and structured input. let toolCall = parts.compactMap { part -> (String, JSONValue)? in - if case .toolCall(_, let name, let input, _, _, _, _, _, _) = part { return (name, input) } + if case .toolCall(_, let name, let input, _, _, _, _, _) = part { return (name, input) } return nil }.first XCTAssertEqual(toolCall?.0, "get_weather") diff --git a/docs/api/gaps.md b/docs/api/gaps.md index 5bdcec9fb..07aebde20 100644 --- a/docs/api/gaps.md +++ b/docs/api/gaps.md @@ -197,12 +197,15 @@ The wire shape changed with it: |---|---|---| | Wire JSON | `"input": {"city":"Paris"}` | `"input": "{\"city\":\"Paris\"}"` | -Serialization always writes the JSON string; nothing writes the old shape any -more. Deserialization accepts both — a JSON string loads unchanged, and any -other JSON value (object, array, number, bool, null) is re-serialized to its -compact JSON text (`deserialize_tool_call_input` in -`aimux-core/src/result.rs`). A `GenerateResult` persisted or recorded before -the refactor therefore keeps loading, with no migration step. +Serialization always writes the JSON string, and deserialization accepts only +the JSON string. The pre-refactor shape (an object, array, number, bool or +null in `input`) no longer loads: the compatibility deserializer was removed +when the tool-call types were unified, and a `GenerateResult` persisted before +PR #165 must be re-recorded. + +The same rule holds for a provider's stream: `StreamPart::ToolCall` carries a +`RawToolCall` whose `input` is the raw text. After `stream_text` the part is a +`TextStreamPart`, whose `ToolCall` carries the parsed `tool::ToolCall`. The parsed arguments still reach callers as a `Value`, on `GenerateTextResult.tool_calls[].input` — that field is unchanged. diff --git a/tools/aimux-web/src/wire.rs b/tools/aimux-web/src/wire.rs index 0c97f2ad6..4222b79d9 100644 --- a/tools/aimux-web/src/wire.rs +++ b/tools/aimux-web/src/wire.rs @@ -179,7 +179,6 @@ fn to_content_part(p: &WireContentPart) -> ContentPart { tool_name: tool_name.clone(), input: input.clone(), provider_executed: *provider_executed, - thought_signature: None, provider_options: None, }, WireContentPart::ToolResult {