diff --git a/CHANGELOG.md b/CHANGELOG.md index 8157ad3a5..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 @@ -41,6 +56,43 @@ 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`, + `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)** + +- 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 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-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/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..48e0df249 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,10 @@ 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, 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 new file mode 100644 index 000000000..3de33b7f1 --- /dev/null +++ b/aimux-provider-utils/src/streaming_tool_call_argument_state.rs @@ -0,0 +1,115 @@ +//! Incremental structure tracking for streamed tool-call arguments. +//! +//! 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. + +/// Whether `value` (ignoring leading whitespace) starts with `{` or `[`. +#[must_use] +pub fn starts_with_structured_value(value: Option<&str>) -> bool { + value.is_some_and(|v| matches!(v.trim_start().as_bytes().first(), Some(b'{' | b'['))) +} + +/// Incrementally tracks whether streamed tool-call arguments contain a +/// complete structured JSON value. +#[derive(Debug, Clone)] +pub struct StreamingToolCallArgumentState { + 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 { + /// Create a state seeded with `initial_value`. + #[must_use] + pub fn new(initial_value: &str) -> Self { + let mut state = Self { + structure: Structure::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, Structure::Complete) + } + + /// Feed the next argument fragment. + pub fn append(&mut self, delta: &str) { + 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, + }; + } + Some(_) => { + self.structure = Structure::Other; + return; + } + } + } + + let Structure::Open { + stack, + in_string, + escaped, + } = &mut self.structure + else { + return; + }; + + // 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 new file mode 100644 index 000000000..fb862a71a --- /dev/null +++ b/aimux-provider-utils/src/streaming_tool_call_tracker.rs @@ -0,0 +1,393 @@ +//! Streaming tool call tracker for the OpenAI chat-completions wire format. +//! +//! 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). +//! +//! 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: +//! +//! | 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 | +//! +//! An ambiguous delta is dropped. +//! +//! Protocols whose streams carry explicit tool-call boundaries (Anthropic, +//! Google, Bedrock, Cohere, the Responses API) do not need this. + +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::{ + StreamingToolCallArgumentState, starts_with_structured_value, +}; + +/// 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"; + +/// One `tool_calls[i]` entry of a streaming chunk, borrowed from the +/// provider's own wire type. +/// +/// The upstream `function: { name, arguments }` object is flattened; a +/// missing `function` is the same as both fields being `None`. +#[derive(Debug, Clone, Default)] +pub struct StreamingToolCallDelta<'a> { + pub index: Option, + pub id: Option<&'a str>, + pub r#type: Option<&'a str>, + pub name: Option<&'a str>, + /// `None` is the wire `null` / absent; `Some("")` is an empty fragment. + pub arguments: Option<&'a str>, + /// Provider metadata for the `ToolCall` part (the upstream + /// `extractMetadata` hook). Read only from the delta that starts a call. + pub provider_metadata: Option, +} + +/// How to validate the `type` field on the delta that starts a call. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum TypeValidation { + /// No validation. + #[default] + None, + /// Reject a `type` that is present and not `"function"`. + IfPresent, + /// Reject a `type` that is not exactly `"function"`. + Required, +} + +/// A delta the tracker cannot accept; the upstream +/// `InvalidResponseDataError`. +#[derive(Debug, Error, PartialEq, Eq)] +pub enum TrackerError { + #[error("Expected 'function.name' to be a string.")] + MissingFunctionName, + #[error("Expected 'function' type.")] + InvalidType, +} + +impl From for AiMuxError { + fn from(error: TrackerError) -> Self { + AiMuxError::InvalidResponseData(error.to_string()) + } +} + +struct ToolCall { + id: String, + name: String, + arguments: String, + argument_state: StreamingToolCallArgumentState, + /// 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 Resolution { + Existing(usize), + New, + Ambiguous, +} + +/// Tracks the tool calls of one streamed response. See the module docs. +pub struct StreamingToolCallTracker { + calls: Vec, + generate_id: Box String + Send + Sync>, + type_validation: TypeValidation, +} + +impl Default for StreamingToolCallTracker { + fn default() -> Self { + Self::new() + } +} + +impl StreamingToolCallTracker { + /// A tracker without `type` validation whose generated ids are + /// `tool-call`, `tool-call-1`, `tool-call-2`, … + #[must_use] + pub fn new() -> Self { + Self { + calls: Vec::new(), + generate_id: Box::new(|| FALLBACK_TOOL_CALL_ID.to_string()), + type_validation: TypeValidation::None, + } + } + + /// Set the generator for calls without a usable wire id. Blank or + /// repeated outputs are turned into unique ids. + #[must_use] + pub fn with_generate_id(mut self, f: impl Fn() -> String + Send + Sync + 'static) -> Self { + self.generate_id = Box::new(f); + self + } + + #[must_use] + pub fn with_type_validation(mut self, validation: TypeValidation) -> Self { + self.type_validation = validation; + self + } + + /// Process one delta, appending the parts it produces to `out`. + /// + /// # Errors + /// + /// A delta that starts a call without a function name, or whose `type` + /// fails validation. + pub fn process( + &mut self, + delta: StreamingToolCallDelta<'_>, + out: &mut Vec, + ) -> Result<(), TrackerError> { + let StreamingToolCallDelta { + index, + 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, + ) + } + 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, + }); + } + call + } + }; + + let indices = &mut self.calls[call].indices; + if let Some(index) = index + && !indices.contains(&index) + { + indices.push(index); + } + Ok(()) + } + + /// 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 + } + for call in self.calls { + out.push(StreamPart::ToolInputEnd { + id: call.id.clone(), + provider_metadata: None, + }); + out.push(StreamPart::ToolCall(RawToolCall { + tool_call_id: call.id, + tool_name: call.name, + // Raw argument text; Core parses it. + input: call.arguments, + provider_executed: None, + dynamic: None, + provider_metadata: call.provider_metadata, + })); + } + } + + /// Apply the correlation table in the module docs. + fn resolve( + &self, + wire_id: Option<&str>, + index: Option, + name: Option<&str>, + 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 name.is_some() { + return Resolution::New; + } + return self.matching(|_| true, false); + }; + + 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 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() { + // 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; + } + if index_known { + // Conflicting labels on a continuation cannot be resolved safely. + return Resolution::Ambiguous; + } + self.matching(has_id, false) + } + + /// 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, + } + } + + 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 start( + &mut self, + wire_id: Option<&str>, + index: Option, + 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, + dynamic: None, + title: None, + provider_metadata: None, + }); + 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 + } + + /// 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 + && !used(wire_id) + { + return wire_id.to_string(); + } + let generated = (self.generate_id)(); + let generated = non_blank(Some(&generated)).unwrap_or(FALLBACK_TOOL_CALL_ID); + if !used(generated) { + return generated.to_string(); + } + // 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()) +} 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 new file mode 100644 index 000000000..ea6d87076 --- /dev/null +++ b/aimux-provider-utils/tests/streaming_tool_call_argument_state_test.rs @@ -0,0 +1,74 @@ +//! Port of `streaming-tool-call-argument-state.test.ts` from +//! `@ai-sdk/provider-utils`. + +use aimux_provider_utils::{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() { + // 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:?}"); + } +} + +#[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-provider-utils/tests/streaming_tool_call_tracker_test.rs b/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs new file mode 100644 index 000000000..a1ce8b16e --- /dev/null +++ b/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs @@ -0,0 +1,1097 @@ +//! 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_core::stream_part::StreamPart; +use aimux_core::tool::RawToolCall; +use aimux_provider_utils::{ + StreamingToolCallDelta, StreamingToolCallTracker, TrackerError, TypeValidation, +}; +use serde_json::{Value, json}; + +/// 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. + 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(RawToolCall { + tool_call_id, + tool_name, + input, + provider_executed: None, + dynamic: 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 { + /// `None` once flushed: `finish` consumes the tracker. + tracker: Option, + parts: Vec, +} + +impl Harness { + fn new() -> Self { + Self::with(StreamingToolCallTracker::new()) + } + + fn with(tracker: StreamingToolCallTracker) -> Self { + Self { + tracker: Some(tracker), + parts: Vec::new(), + } + } + + fn delta(&mut self, delta: StreamingToolCallDelta<'_>) -> Result<(), TrackerError> { + let mut parts = Vec::new(); + let tracker = self.tracker.as_mut().expect("delta after flush"); + let result = tracker.process(delta, &mut parts); + self.parts.extend(parts.into_iter().map(project)); + result + } + + fn flush(&mut self) { + let mut parts = Vec::new(); + self.tracker + .take() + .expect("flushed twice") + .finish(&mut parts); + self.parts.extend(parts.into_iter().map(project)); + } + + fn clear(&mut self) { + self.parts.clear(); + } + + /// `(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<'a>( + index: Option, + id: Option<&'a str>, + ty: Option<&'a str>, + name: Option<&'a str>, + arguments: Option<&'a str>, +) -> StreamingToolCallDelta<'a> { + StreamingToolCallDelta { + index, + id, + r#type: ty, + name, + arguments, + provider_metadata: None, + } +} + +/// Full call start: index, id, `type: 'function'`, name and arguments. +fn start<'a>( + index: usize, + id: &'a str, + name: &'a str, + arguments: &'a str, +) -> StreamingToolCallDelta<'a> { + 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_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<'static> { + delta( + Some(0), + Some("call_1"), + Some("custom"), + Some("fn"), + Some(""), + ) + } + + fn no_type() -> StreamingToolCallDelta<'static> { + 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"), + ] + ); + } +} + +mod metadata { + use super::*; + + // The upstream `extractMetadata` / `buildToolCallProviderMetadata` hooks + // become `provider_metadata` on the delta that starts a call. + + #[test] + fn includes_the_starting_delta_metadata_in_the_tool_call() { + let mut h = Harness::new(); + + h.delta(StreamingToolCallDelta { + provider_metadata: Some(json!({ "google": { "thoughtSignature": "sig123" } })), + ..start(0, "call_1", "fn", "{\"incomplete") + }) + .unwrap(); + 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!({ "google": { "thoughtSignature": "sig123" } })), + }) + ); + } + + #[test] + fn ignores_metadata_on_continuations() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "fn", "{")).unwrap(); + h.delta(StreamingToolCallDelta { + provider_metadata: Some(json!({ "custom": { "key": "value" } })), + ..cont(Some(0), "}") + }) + .unwrap(); + h.flush(); + + assert_eq!(h.parts.last(), Some(&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"]); + } +} 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 6445d52c2..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; @@ -20,7 +21,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; +use aimux_provider_utils::{HttpRequest, StreamingToolCallDelta, StreamingToolCallTracker}; use super::OpenAIConfig; use super::convert::{RequestBodyResult, build_request_body_with_warnings, parse_finish_reason}; @@ -176,15 +177,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 @@ -326,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. @@ -491,9 +482,10 @@ 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(); + let mut tool_parts = Vec::new(); // Process the first event (already peeked) then the rest. let mut event_iter = @@ -642,50 +634,25 @@ 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, - }); + for dtc in &tool_call_deltas { + let function = dtc.function.as_ref(); + let delta = StreamingToolCallDelta { + index: dtc.index, + id: dtc.id.as_deref(), + name: function.and_then(|f| f.name.as_deref()), + arguments: function.and_then(|f| f.arguments.as_deref()), + ..Default::default() + }; + // 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); } - - // 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, - }); - } } } @@ -763,27 +730,11 @@ 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. + tool_calls.finish(&mut tool_parts); + for part in tool_parts.drain(..) { + yield Ok(part); } // Build provider metadata for the Finish part. 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/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-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/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 e023a3204..3edc0c1ab 100644 --- a/aimux-stream/src/lib.rs +++ b/aimux-stream/src/lib.rs @@ -1,17 +1,13 @@ //! # 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) 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; -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); -} 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/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。 */ 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/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 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 {