diff --git a/aimux-providers/layout.toml b/aimux-providers/layout.toml index b4ce587cb..63d52ad09 100644 --- a/aimux-providers/layout.toml +++ b/aimux-providers/layout.toml @@ -2165,3 +2165,278 @@ reason = "the native Decisions API; upstream decides through the Responses model rust = "openai/responses/tool_args.rs" status = "aimux" reason = "the provider tool argument schemas of tool/*.ts, validated in one place" + +[bedrock] +package = "@ai-sdk/amazon-bedrock" +# amazon-bedrock-chat-language-model-event-stream.test.ts and amazon-bedrock-usage-raw.test.ts +# have no source file of their own name; every case in them is ported, so they need no +# exclusions. + +[[bedrock.files]] +upstream = "src/amazon-bedrock-anthropic-model-support.ts" +rust = "bedrock/convert.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/amazon-bedrock-api-types.ts" +rust = "bedrock/types.rs" +status = "mapped" + +[bedrock.files.tests] +sub-feature = [ + { case = "createAmazonBedrockCachePoint › should return default cache point when no TTL is provided", issue = 275, reason = "no createAmazonBedrockCachePoint builder; a cache point is passed verbatim through providerOptions.amazonBedrock.cachePoint" }, + { case = "createAmazonBedrockCachePoint › should create cache point with 5m TTL", issue = 275, reason = "no createAmazonBedrockCachePoint builder; a cache point is passed verbatim through providerOptions.amazonBedrock.cachePoint" }, + { case = "createAmazonBedrockCachePoint › should create cache point with 1h TTL", issue = 275, reason = "no createAmazonBedrockCachePoint builder; a cache point is passed verbatim through providerOptions.amazonBedrock.cachePoint" }, +] + +[[bedrock.files]] +upstream = "src/amazon-bedrock-chat-language-model-options.ts" +rust = "bedrock/convert.rs" +status = "mapped" + +[bedrock.files.tests] +ts-only = [ + { case = "amazonBedrockLanguageModelChatOptions › type inference › infers AmazonBedrockLanguageModelChatOptions type correctly", reason = "checks a TypeScript type annotation; the Rust options are untyped JSON validated when the request is built" }, +] + +[[bedrock.files]] +upstream = "src/amazon-bedrock-chat-language-model.ts" +rust = "bedrock/model.rs" +status = "mapped" + +[bedrock.files.tests] +ts-only = [ + { case = "doGenerate › should handle unsupported response format types", reason = "passes responseFormat { type: 'xml' } through a TypeScript `as any`; ResponseFormat is a closed Rust enum (Text, Json)" }, +] + +[[bedrock.files]] +upstream = "src/amazon-bedrock-embedding-model-options.ts" +rust = "bedrock/embedding.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/amazon-bedrock-embedding-model.ts" +rust = "bedrock/embedding.rs" +status = "mapped" + +[bedrock.files.tests] +ts-only = [ + { case = "doEmbed › should return NaN usage for Cohere models when the token count header is absent", reason = "NaN is a JS number; EmbeddingUsage.tokens is a u32" }, +] + +[[bedrock.files]] +upstream = "src/amazon-bedrock-error.ts" +rust = "bedrock/mod.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/amazon-bedrock-event-stream-decoder.ts" +rust = "bedrock/event_stream.rs" +status = "mapped" + +[bedrock.files.tests] +ts-only = [ + { case = "createAmazonBedrockEventStreamDecoder › surfaces processEvent errors", reason = "the TS decoder takes a processEvent callback that can throw; decode_stream yields messages and has no callback" }, +] + +[[bedrock.files]] +upstream = "src/amazon-bedrock-event-stream-response-handler.ts" +rust = "bedrock/event_stream.rs" +status = "mapped" + +[bedrock.files.tests] +ts-only = [ + { case = "createEventSourceResponseHandler › throws EmptyResponseBodyError when response body is null", reason = "a FetchResponse always has a body stream; there is no null body" }, + { case = "createEventSourceResponseHandler › handles schema validation failures", reason = "mocks EventStreamCodec and passes a zod chunk schema to the handler; the Rust handler takes no schema" }, +] + +[[bedrock.files]] +upstream = "src/amazon-bedrock-image-model-options.ts" +rust = "bedrock/image.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/amazon-bedrock-image-model.ts" +rust = "bedrock/image.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/amazon-bedrock-image-settings.ts" +rust = "bedrock/image.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/amazon-bedrock-prepare-tools.ts" +rust = "bedrock/tools.rs" +status = "mapped" + +[bedrock.files.tests] +ts-only = [ + { case = "prepareTools › tool choice › models that reject forced tool use › should pass the forced-tool capability to Anthropic provider tool preparation", reason = "asserts the arguments of a vi.mock spy on the Anthropic prepareTools; there is no seam to observe in Rust" }, +] + +[[bedrock.files]] +upstream = "src/amazon-bedrock-provider.ts" +rust = "bedrock/mod.rs" +status = "mapped" + +[bedrock.files.tests] +sub-feature = [ + { case = "AmazonBedrockProvider › provider methods › should expose anthropicTools", issue = 275, reason = "AmazonBedrockProvider has no tools member (#275 厂商内置工具工厂)" }, +] +ts-only = [ + { case = "AmazonBedrockProvider › createAmazonBedrock › should throw error when called with new keyword", reason = "Rust has no `new` on a callable; the provider is a struct" }, +] + +[[bedrock.files]] +upstream = "src/amazon-bedrock-reasoning-metadata.ts" +rust = "bedrock/model.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/amazon-bedrock-sigv4-fetch.ts" +rust = "bedrock/mod.rs" +status = "mapped" + +[bedrock.files.tests] +registered = [ + { case = "createSigV4FetchFunction › should bypass signing for non-POST requests", ref = "AWS SigV4 fetch 也给 GET 请求签名（§9.1 已知差距）", reason = "aimux signs every request that carries no API key, GET included, so that list_models is accepted" }, + { case = "createSigV4FetchFunction › should bypass signing if POST request has no body", ref = "AWS SigV4 fetch 也给 GET 请求签名（§9.1 已知差距）", reason = "aimux signs every request that carries no API key, a bodyless POST included" }, + { case = "createSigV4FetchFunction › should call original fetch if init is undefined", ref = "AWS SigV4 fetch 也给 GET 请求签名（§9.1 已知差距）", reason = "a request with no init is a GET, which upstream forwards unsigned; aimux signs it" }, +] +ts-only = [ + { case = "createSigV4FetchFunction › should handle non-string body by stringifying it", reason = "FetchRequest.body is bytes; there is no JS object body for JSON.stringify to coerce" }, + { case = "createApiKeyFetchFunction › should handle empty string API key", reason = "createApiKeyFetchFunction is a TS module export; the Rust provider sends a non-blank key as a header and has no fetch factory to hand an empty key" }, +] + +[[bedrock.files]] +upstream = "src/amazon-bedrock-stream-error.ts" +rust = "bedrock/model.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/anthropic/amazon-bedrock-anthropic-fetch.ts" +status = "not-applicable" +reason = "Bedrock-Anthropic (InvokeModel with the Anthropic Messages body) is not implemented; alignment doc §5 product choice row \"Bedrock 上的 Anthropic InvokeModel\" (anthropic_aws is Claude Platform on AWS, not this)" + +[[bedrock.files]] +upstream = "src/anthropic/amazon-bedrock-anthropic-options.ts" +status = "not-applicable" +reason = "Bedrock-Anthropic (InvokeModel with the Anthropic Messages body) is not implemented; alignment doc §5 product choice row \"Bedrock 上的 Anthropic InvokeModel\" (anthropic_aws is Claude Platform on AWS, not this)" + +[[bedrock.files]] +upstream = "src/anthropic/amazon-bedrock-anthropic-provider.ts" +status = "not-applicable" +reason = "Bedrock-Anthropic (InvokeModel with the Anthropic Messages body) is not implemented; alignment doc §5 product choice row \"Bedrock 上的 Anthropic InvokeModel\" (anthropic_aws is Claude Platform on AWS, not this)" + +[[bedrock.files]] +upstream = "src/anthropic/index.ts" +status = "not-applicable" +reason = "Bedrock-Anthropic (InvokeModel with the Anthropic Messages body) is not implemented; alignment doc §5 product choice row \"Bedrock 上的 Anthropic InvokeModel\" (anthropic_aws is Claude Platform on AWS, not this)" + +[[bedrock.files]] +upstream = "src/convert-amazon-bedrock-usage.ts" +rust = "bedrock/convert.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/convert-to-amazon-bedrock-chat-messages.ts" +rust = "bedrock/convert.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/index.ts" +rust = "bedrock/mod.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/inject-fetch-headers.ts" +status = "not-applicable" +reason = "a test helper of the package (not exported); the ported tests keep a Rust copy in bedrock/tests/support.rs" + +[[bedrock.files]] +upstream = "src/mantle/bedrock-mantle-options.ts" +status = "missing" +issue = 275 +reason = "Bedrock Mantle is a registry preset (OpenAI-compatible chat), not the full createBedrockMantle provider; #275 \"Bedrock Mantle 改为完整 provider\"" + +[[bedrock.files]] +upstream = "src/mantle/bedrock-mantle-provider.ts" +status = "missing" +issue = 275 +reason = "Bedrock Mantle is a registry preset (OpenAI-compatible chat), not the full createBedrockMantle provider; #275 \"Bedrock Mantle 改为完整 provider\"" + +[[bedrock.files]] +upstream = "src/mantle/index.ts" +status = "missing" +issue = 275 +reason = "Bedrock Mantle is a registry preset (OpenAI-compatible chat), not the full createBedrockMantle provider; #275 \"Bedrock Mantle 改为完整 provider\"" + +[[bedrock.files]] +upstream = "src/map-amazon-bedrock-finish-reason.ts" +rust = "bedrock/convert.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/normalize-tool-call-id.ts" +rust = "bedrock/convert.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/reranking/amazon-bedrock-reranking-api.ts" +rust = "bedrock/reranking.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/reranking/amazon-bedrock-reranking-model-options.ts" +rust = "bedrock/reranking.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/reranking/amazon-bedrock-reranking-model.ts" +rust = "bedrock/reranking.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/resolve-amazon-bedrock-base-url.ts" +rust = "bedrock/mod.rs" +status = "mapped" + +[bedrock.files.tests] + +[[bedrock.files]] +upstream = "src/version.ts" +status = "not-applicable" +reason = "VERSION is the package version string; aimux pins it in shared/aisdk_pins.rs, which the user-agent tests read" diff --git a/aimux-providers/src/bedrock/convert.rs b/aimux-providers/src/bedrock/convert.rs index 8f6ee7251..3c3ac020d 100644 --- a/aimux-providers/src/bedrock/convert.rs +++ b/aimux-providers/src/bedrock/convert.rs @@ -48,7 +48,7 @@ pub fn convert_prompt_to_bedrock( convert_prompt(prompt, false) } -fn convert_prompt( +pub(super) fn convert_prompt( prompt: &LanguageModelPrompt, is_mistral: bool, ) -> Result<(Vec, Vec), aimux_core::AiMuxError> { diff --git a/aimux-providers/src/bedrock/embedding.rs b/aimux-providers/src/bedrock/embedding.rs index b35677ea0..37018d3bd 100644 --- a/aimux-providers/src/bedrock/embedding.rs +++ b/aimux-providers/src/bedrock/embedding.rs @@ -30,7 +30,7 @@ use crate::shared::EndpointConfig; pub struct BedrockEmbeddingModel { model_id: String, config: EndpointConfig, - model_family: Option, + pub(super) model_family: Option, } impl BedrockEmbeddingModel { diff --git a/aimux-providers/src/bedrock/event_stream.rs b/aimux-providers/src/bedrock/event_stream.rs index db4fd3099..573bc01ed 100644 --- a/aimux-providers/src/bedrock/event_stream.rs +++ b/aimux-providers/src/bedrock/event_stream.rs @@ -362,6 +362,7 @@ mod tests { ); } + /// partial: `surfaces frame decode errors instead of stranding later frames`; full port in src/bedrock/tests/event_stream_decoder.rs. #[test] fn corrupt_message_crc_is_rejected() { let payload = r#"{"delta":{"text":"Hi"}}"#; diff --git a/aimux-providers/src/bedrock/mod.rs b/aimux-providers/src/bedrock/mod.rs index 899b86f12..defb2f864 100644 --- a/aimux-providers/src/bedrock/mod.rs +++ b/aimux-providers/src/bedrock/mod.rs @@ -718,3 +718,6 @@ impl ProviderDiscovery for AmazonBedrockProvider { }) } } + +#[cfg(test)] +mod tests; diff --git a/aimux-providers/src/bedrock/model.rs b/aimux-providers/src/bedrock/model.rs index 3c00292d5..006164e1c 100644 --- a/aimux-providers/src/bedrock/model.rs +++ b/aimux-providers/src/bedrock/model.rs @@ -95,7 +95,7 @@ fn bedrock_event_stream_response_handler() pub struct BedrockModel { model_id: String, config: EndpointConfig, - model_family: Option, + pub(super) model_family: Option, } impl BedrockModel { diff --git a/aimux-providers/src/bedrock/tests/chat_language_model.rs b/aimux-providers/src/bedrock/tests/chat_language_model.rs new file mode 100644 index 000000000..6f99ac1ed --- /dev/null +++ b/aimux-providers/src/bedrock/tests/chat_language_model.rs @@ -0,0 +1,7712 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use futures::StreamExt; +use reqwest::header::{HeaderName, HeaderValue}; +use serde_json::{Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::language_model::LanguageModel; +use aimux_core::options::CallOptions; +use aimux_core::result::StreamResult; +use aimux_core::stream_part::StreamPart; +use aimux_provider_utils::{ + Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, HeaderMapOpt, Resolvable, +}; + +use crate::bedrock::{ + AmazonBedrockChatModelFamily, AmazonBedrockChatModelSettings, AmazonBedrockProvider, + AmazonBedrockProviderSettings, BedrockModel, create_amazon_bedrock, +}; +use crate::test_support::mock_fetch::{Canned, MockFetch, RouteFetch, Seen}; + +/// `modelId` of the upstream test file. +const MODEL_ID: &str = "anthropic.claude-3-haiku-20240307-v1:0"; +/// `baseUrl` of the upstream test file. +const BASE_URL: &str = "https://bedrock-runtime.us-east-1.amazonaws.com"; + +/// `TEST_PROMPT`: a system message and one user text part. +fn test_prompt() -> Value { + json!([ + { "role": "system", "content": "System Prompt" }, + { "role": "user", "content": [{ "type": "text", "text": "Hello" }] }, + ]) +} + +/// Call options from their AI SDK JSON (`LanguageModelV4CallOptions`). +fn call_options(options: Value) -> CallOptions { + serde_json::from_value(options).expect("valid call options") +} + +/// `{ prompt: TEST_PROMPT, ...extra }`. +fn options_with(extra: Value) -> CallOptions { + let mut options = json!({ "prompt": test_prompt() }); + if let (Some(target), Value::Object(extra)) = (options.as_object_mut(), extra) { + target.extend(extra); + } + call_options(options) +} + +// ── AWS event-stream frames ────────────────────────────────────────────────── + +fn crc32(data: &[u8]) -> u32 { + let mut crc = 0xFFFF_FFFFu32; + for &byte in data { + crc ^= u32::from(byte); + for _ in 0..8 { + crc = if crc & 1 == 0 { + crc >> 1 + } else { + (crc >> 1) ^ 0xEDB8_8320 + }; + } + } + !crc +} + +/// One AWS event-stream frame with string headers (the `EventStreamCodec.encode` +/// of `@smithy/eventstream-codec`). +fn encode_frame(headers: &[(&str, &str)], body: &[u8]) -> Vec { + let mut header_bytes = Vec::new(); + for (name, value) in headers { + header_bytes.push(u8::try_from(name.len()).expect("short header name")); + header_bytes.extend_from_slice(name.as_bytes()); + header_bytes.push(7); + header_bytes.extend_from_slice( + &u16::try_from(value.len()) + .expect("short value") + .to_be_bytes(), + ); + header_bytes.extend_from_slice(value.as_bytes()); + } + let total = u32::try_from(16 + header_bytes.len() + body.len()).expect("small frame"); + let headers_length = u32::try_from(header_bytes.len()).expect("small headers"); + let mut frame = Vec::new(); + frame.extend_from_slice(&total.to_be_bytes()); + frame.extend_from_slice(&headers_length.to_be_bytes()); + let prelude_crc = crc32(&frame); + frame.extend_from_slice(&prelude_crc.to_be_bytes()); + frame.extend_from_slice(&header_bytes); + frame.extend_from_slice(body); + let message_crc = crc32(&frame); + frame.extend_from_slice(&message_crc.to_be_bytes()); + frame +} + +/// The frame the Converse stream carries for one upstream chunk +/// `{ : }`: an `event` frame, or an `exception` frame for an +/// `…Exception` type (the decoder wraps either payload under its type again). +fn chunk_frame(chunk: &Value) -> Vec { + let (kind, payload) = chunk + .as_object() + .and_then(|object| object.iter().next()) + .expect("a chunk is a one-key object"); + let body = serde_json::to_vec(payload).expect("serializable payload"); + if kind.ends_with("Exception") { + encode_frame( + &[ + (":message-type", "exception"), + (":exception-type", kind), + (":content-type", "application/json"), + ], + &body, + ) + } else { + encode_frame( + &[ + (":message-type", "event"), + (":event-type", kind), + (":content-type", "application/json"), + ], + &body, + ) + } +} + +/// The event-stream body of a list of upstream chunks. +fn event_stream_body(chunks: &[Value]) -> Vec { + chunks.iter().flat_map(chunk_frame).collect() +} + +/// The chunks of a `__fixtures__/*.chunks.txt` file (one JSON chunk per line). +fn fixture_chunks(text: &str) -> Vec { + text.lines() + .filter(|line| !line.is_empty()) + .map(|line| serde_json::from_str(line).expect("fixture chunk is JSON")) + .collect() +} + +// ── Mock responses and transports ──────────────────────────────────────────── + +/// A `createTestServer` `stream-chunks` response: its default headers, then +/// `headers`. +fn stream_response(body: Vec, headers: &[(&str, &str)]) -> Canned { + let mut all: Vec<(String, String)> = vec![ + ("content-type".into(), "text/event-stream".into()), + ("cache-control".into(), "no-cache".into()), + ("connection".into(), "keep-alive".into()), + ]; + for (name, value) in headers { + all.push(((*name).into(), (*value).into())); + } + Canned { + status: 200, + headers: all, + body, + } +} + +/// A `stream-chunks` response carrying `chunks` as event-stream frames. +fn stream_chunks(chunks: &[Value]) -> Canned { + stream_response(event_stream_body(chunks), &[]) +} + +/// `injectFetchHeaders`: a transport that sets `headers` on every request +/// (over the ones already there) before handing it to `inner`. +struct InjectFetchHeaders { + headers: Vec<(&'static str, &'static str)>, + inner: FetchFunction, +} + +#[async_trait] +impl Fetch for InjectFetchHeaders { + async fn fetch(&self, mut request: FetchRequest) -> Result { + for (name, value) in &self.headers { + request.headers.insert( + HeaderName::from_static(name), + HeaderValue::from_static(value), + ); + } + self.inner.fetch(request).await + } +} + +fn inject_fetch_headers( + headers: Vec<(&'static str, &'static str)>, + inner: FetchFunction, +) -> FetchFunction { + Arc::new(InjectFetchHeaders { headers, inner }) +} + +/// `fakeFetchWithAuth`: `injectFetchHeaders({ 'x-amz-auth': 'test-auth' })` +/// over the mock server. +fn fake_fetch_with_auth(mock: &Arc) -> FetchFunction { + inject_fetch_headers(vec![("x-amz-auth", "test-auth")], mock.transport()) +} + +/// A provider on `BASE_URL` whose requests go through `fetch`. An API key +/// stands in for the upstream model config's absent signing layer, so +/// nothing is signed. +fn provider(fetch: FetchFunction, headers: Option) -> AmazonBedrockProvider { + create_amazon_bedrock(AmazonBedrockProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".into())), + region: Some("us-east-1".into()), + base_url: Some(BASE_URL.into()), + headers, + fetch: Some(fetch), + ..Default::default() + }) + .expect("valid settings") +} + +/// The chat model `model_id` over `fetch`. +fn chat_model(model_id: &str, fetch: FetchFunction) -> BedrockModel { + provider(fetch, None).chat(model_id) +} + +/// The upstream `model`: `MODEL_ID` over `fakeFetchWithAuth`. +fn model(mock: &Arc) -> BedrockModel { + chat_model(MODEL_ID, fake_fetch_with_auth(mock)) +} + +/// The JSON body of the `index`th request the mock saw. +fn request_body(seen: &[Seen], index: usize) -> Value { + seen.get(index).expect("request was sent").json_body() +} + +// ── Stream output ──────────────────────────────────────────────────────────── + +/// `convertReadableStreamToArray`, keeping errors the stream threw. +async fn stream_items(result: StreamResult) -> Vec> { + result.stream.collect().await +} + +/// Every part of a stream that must not throw. +async fn stream_parts(result: StreamResult) -> Vec { + stream_items(result) + .await + .into_iter() + .map(|item| item.expect("stream does not throw")) + .collect() +} + +/// The upstream form of an error part's error: a provider stream error +/// (`message`, `type`, `statusCode`, `isRetryable`, `data`) for an API call +/// error, `{ message }` for anything else. +fn error_json(error: &AiMuxError) -> Value { + match error { + AiMuxError::ApiCall(error) => { + let mut object = serde_json::Map::new(); + object.insert("message".into(), json!(error.message)); + if let Some(kind) = &error.provider_code { + object.insert("type".into(), json!(kind)); + } + if let Some(status) = error.status_code { + object.insert("statusCode".into(), json!(status)); + } + object.insert("isRetryable".into(), json!(error.is_retryable)); + if let Some(data) = &error.data { + object.insert("data".into(), data.clone()); + } + Value::Object(object) + } + other => json!({ "message": other.to_string() }), + } +} + +/// A stream part as the AI SDK JSON the upstream snapshots print. +fn part_json(part: &StreamPart) -> Value { + match part { + StreamPart::Error { error } => json!({ "type": "error", "error": error_json(error) }), + other => serde_json::to_value(other).expect("serializable part"), + } +} + +fn parts_json(parts: &[StreamPart]) -> Value { + Value::Array(parts.iter().map(part_json).collect()) +} + +/// The first part of type `kind`, as JSON. +fn find_part(parts: &[StreamPart], kind: &str) -> Option { + parts + .iter() + .map(part_json) + .find(|part| part["type"] == json!(kind)) +} + +/// `expect(actual).toMatchObject(expected)`: every key of an expected object +/// is present and matches; arrays match element by element. +fn assert_match_object(actual: &Value, expected: &Value) { + fn check(actual: &Value, expected: &Value, path: &str) { + match (actual, expected) { + (Value::Object(actual), Value::Object(expected)) => { + for (key, value) in expected { + let at = format!("{path}.{key}"); + let Some(found) = actual.get(key) else { + panic!("missing {at}; expected {value}"); + }; + check(found, value, &at); + } + } + (Value::Array(actual), Value::Array(expected)) => { + assert_eq!(actual.len(), expected.len(), "length of {path}"); + for (index, (found, value)) in actual.iter().zip(expected).enumerate() { + check(found, value, &format!("{path}[{index}]")); + } + } + _ => assert_eq!(actual, expected, "at {path}"), + } + } + check(actual, expected, "$"); +} + +mod supported_urls { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › supportedUrls › should support S3 URLs for image and video parts + #[tokio::test] + async fn should_support_s3_urls_for_image_and_video_parts() { + let mock = MockFetch::new(vec![]); + let urls = model(&mock).supported_urls().await; + + let patterns: std::collections::BTreeMap<&str, Vec<&str>> = urls + .0 + .iter() + .map(|(media_type, patterns)| { + ( + media_type.as_str(), + patterns.iter().map(regex::Regex::as_str).collect(), + ) + }) + .collect(); + assert_eq!( + patterns, + [("image/*", vec!["^s3://"]), ("video/*", vec!["^s3://"])] + .into_iter() + .collect() + ); + } +} + +mod application_inference_profile_reasoning { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › application inference profile reasoning › returns reasoning for an Anthropic application inference profile ARN when budgetTokens is configured + #[tokio::test] + async fn returns_reasoning_for_an_anthropic_application_inference_profile_arn() { + let application_profile_arn = + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/custom-profile"; + let fetch = RouteFetch::new(|seen| { + let request_body = seen.json_body(); + let content = if request_body["additionalModelRequestFields"]["thinking"].is_null() { + json!([{ "text": "OK" }]) + } else { + json!([ + { + "reasoningContent": { + "reasoningText": { + "text": "The response should be OK.", + "signature": "test-signature", + }, + }, + }, + { "text": "OK" }, + ]) + }; + Canned::json(&json!({ + "output": { "message": { "role": "assistant", "content": content } }, + "stopReason": "end_turn", + "usage": { "inputTokens": 1, "outputTokens": 1, "totalTokens": 2 }, + })) + }); + let application_profile_model = chat_model(application_profile_arn, fetch.transport()); + + let result = application_profile_model + .do_generate(&options_with(json!({ + "maxOutputTokens": 1100, + "providerOptions": { + "bedrock": { + "reasoningConfig": { "type": "enabled", "budgetTokens": 1024 }, + }, + }, + }))) + .await + .expect("generate succeeds"); + + let content = serde_json::to_value(&result.content).expect("serializable content"); + assert!( + content + .as_array() + .expect("content array") + .iter() + .any(|part| { + part["type"] == json!("reasoning") + && part["text"] == json!("The response should be OK.") + && part["providerMetadata"]["amazonBedrock"] + == json!({ "signature": "test-signature" }) + }), + "no reasoning part in {content}" + ); + let request_body = request_body(&fetch.seen(), 0); + assert_eq!( + request_body["additionalModelRequestFields"]["thinking"], + json!({ "type": "enabled", "budget_tokens": 1024 }) + ); + let warnings = serde_json::to_value(&result.warnings).expect("serializable warnings"); + assert!( + !warnings + .as_array() + .expect("warnings array") + .iter() + .any(|warning| { + warning["type"] == json!("unsupported") + && warning["feature"] == json!("budgetTokens") + }), + "unexpected budgetTokens warning in {warnings}" + ); + } +} + +mod application_inference_profile_structured_output { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › application inference profile structured output › uses native structured output in %s mode when the Anthropic family is declared + #[tokio::test] + async fn uses_native_structured_output_in_mode_when_the_anthropic_family_is_declared() { + for structured_output_mode in ["outputFormat", "auto"] { + let application_profile_arn = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/custom-profile"; + let fetch = RouteFetch::new(|_| { + Canned::json(&json!({ + "output": { + "message": { + "role": "assistant", + "content": [{ "text": "{\"name\":\"Test\"}" }], + }, + }, + "stopReason": "end_turn", + "usage": { "inputTokens": 1, "outputTokens": 1, "totalTokens": 2 }, + })) + }); + let application_profile_model = provider(fetch.transport(), None).chat_with_settings( + application_profile_arn, + &AmazonBedrockChatModelSettings { + model_family: Some(AmazonBedrockChatModelFamily::Anthropic), + }, + ); + + application_profile_model + .do_generate(&options_with(json!({ + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { "name": { "type": "string" } }, + "required": ["name"], + }, + }, + "providerOptions": { + "amazonBedrock": { "structuredOutputMode": structured_output_mode }, + }, + }))) + .await + .expect("generate succeeds"); + + let request_body = request_body(&fetch.seen(), 0); + assert_eq!( + request_body.get("toolConfig"), + None, + "{structured_output_mode}" + ); + assert_eq!( + request_body["additionalModelRequestFields"]["output_config"]["format"], + json!({ + "type": "json_schema", + "schema": { + "type": "object", + "additionalProperties": false, + "properties": { "name": { "type": "string" } }, + "required": ["name"], + }, + }), + "{structured_output_mode}" + ); + } + } +} + +mod do_generate_request_metadata { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate request metadata › should return the request body + #[tokio::test] + async fn should_return_the_request_body() { + let fixture: Value = + serde_json::from_str(include_str!("fixtures/amazon-bedrock-text.json")).unwrap(); + let mock = MockFetch::new(vec![Canned::json(&fixture)]); + + let result = model(&mock) + .do_generate(&options_with(json!({}))) + .await + .expect("generate succeeds"); + + assert_eq!( + result.request.and_then(|request| request.body), + Some(json!({ + "additionalModelResponseFieldPaths": [ + "/delta/stop_sequence" + ], + "messages": [ + { + "content": [ + { + "text": "Hello" + } + ], + "role": "user" + } + ], + "system": [ + { + "text": "System Prompt" + } + ] + })) + ); + } +} + +mod request_url { + use super::*; + + mod arn_model_ids_containing_a_slash { + use super::*; + + const INFERENCE_PROFILE_ARN: &str = + "arn:aws:bedrock:eu-west-1:474668406012:inference-profile/eu.amazon.nova-lite-v1:0"; + + /// `encodeURIComponent(inferenceProfileArn)`. + const ENCODED_INFERENCE_PROFILE_ARN: &str = "arn%3Aaws%3Abedrock%3Aeu-west-1%3A474668406012%3Ainference-profile%2Feu.amazon.nova-lite-v1%3A0"; + + /// `createInferenceProfileModel`: answers the encoded Converse routes + /// and an `UnknownOperationException` body anywhere else. + fn create_inference_profile_model() -> BedrockModel { + let fetch = RouteFetch::new(|seen| { + let url = &seen.url; + let is_encoded_route = url.contains(ENCODED_INFERENCE_PROFILE_ARN); + let unknown_operation_response = json!({ + "Output": { "__type": "com.amazon.coral.service#UnknownOperationException" }, + "Version": "1.0", + }); + let converse_response = json!({ + "output": { "message": { "role": "assistant", "content": [{ "text": "Hello!" }] } }, + "stopReason": "end_turn", + "usage": { "inputTokens": 1, "outputTokens": 1, "totalTokens": 2 }, + }); + let converse_stream_response = [ + json!({ "messageStart": { "role": "assistant" } }), + json!({ "contentBlockDelta": { "contentBlockIndex": 0, "delta": { "text": "Hello!" } } }), + json!({ "contentBlockStop": { "contentBlockIndex": 0 } }), + json!({ "messageStop": { "stopReason": "end_turn" } }), + json!({ "metadata": { "usage": { "inputTokens": 1, "outputTokens": 1, "totalTokens": 2 } } }), + ]; + if is_encoded_route && url.ends_with("/converse-stream") { + Canned { + status: 200, + headers: vec![( + "content-type".into(), + "application/vnd.amazon.eventstream".into(), + )], + body: event_stream_body(&converse_stream_response), + } + } else if is_encoded_route { + Canned::json(&converse_response) + } else { + Canned::json(&unknown_operation_response) + } + }); + chat_model(INFERENCE_PROFILE_ARN, fetch.transport()) + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › request URL › ARN model IDs containing a slash › should generate text through the encoded Converse route + #[tokio::test] + async fn should_generate_text_through_the_encoded_converse_route() { + let result = create_inference_profile_model() + .do_generate(&options_with(json!({}))) + .await + .expect("generate succeeds"); + + assert_eq!( + serde_json::to_value(&result.content).unwrap(), + json!([ + { + "text": "Hello!", + "type": "text" + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › request URL › ARN model IDs containing a slash › should stream text through the encoded Converse route + #[tokio::test] + async fn should_stream_text_through_the_encoded_converse_route() { + let result = create_inference_profile_model() + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + + assert_eq!( + parts_json(&stream_parts(result).await), + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "arn:aws:bedrock:eu-west-1:474668406012:inference-profile/eu.amazon.nova-lite-v1:0", + "type": "response-metadata" + }, + { + "id": "0", + "type": "text-start" + }, + { + "delta": "Hello!", + "id": "0", + "type": "text-delta" + }, + { + "id": "0", + "type": "text-end" + }, + { + "finishReason": { + "raw": "end_turn", + "unified": "stop" + }, + "type": "finish", + "usage": { + "inputTokens": { + "cacheRead": 0, + "cacheWrite": 0, + "noCache": 1, + "total": 1 + }, + "outputTokens": { + "text": 1, + "total": 1 + }, + "raw": { + "inputTokens": 1, + "outputTokens": 1, + "totalTokens": 2 + } + } + } + ]) + ); + } + } +} + +mod do_stream { + use super::*; + + mod text { + use super::*; + + fn text_fixture_response(headers: &[(&str, &str)]) -> Canned { + stream_response( + event_stream_body(&fixture_chunks(include_str!( + "fixtures/amazon-bedrock-text.chunks.txt" + ))), + headers, + ) + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › text › should stream text deltas + #[tokio::test] + async fn should_stream_text_deltas() { + let mock = MockFetch::new(vec![text_fixture_response(&[])]); + + let result = model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + + let expected: Value = serde_json::from_str(include_str!( + "snapshots/do_stream_text_should_stream_text_deltas.json" + )) + .unwrap(); + assert_eq!(parts_json(&stream_parts(result).await), expected); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › text › should expose the raw response headers + #[tokio::test] + async fn should_expose_the_raw_response_headers() { + let mock = MockFetch::new(vec![text_fixture_response(&[( + "test-header", + "test-value", + )])]); + + let result = model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + + assert_eq!( + serde_json::to_value(result.response.and_then(|response| response.headers)) + .unwrap(), + json!({ + "cache-control": "no-cache", + "connection": "keep-alive", + "content-type": "text/event-stream", + "test-header": "test-value" + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › text › should return the request body + #[tokio::test] + async fn should_return_the_request_body() { + let mock = MockFetch::new(vec![text_fixture_response(&[])]); + + let result = model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + + assert_eq!( + result.request.and_then(|request| request.body), + Some(json!({ + "additionalModelResponseFieldPaths": [ + "/delta/stop_sequence" + ], + "messages": [ + { + "content": [ + { + "text": "Hello" + } + ], + "role": "user" + } + ], + "system": [ + { + "text": "System Prompt" + } + ] + })) + ); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should accept citation deltas + #[tokio::test] + async fn should_accept_citation_deltas() { + let mock = MockFetch::new(vec![stream_chunks(&[json!({ + "contentBlockDelta": { + "contentBlockIndex": 1, + "delta": { + "citation": { + "location": { + "documentPage": { "documentIndex": 0, "start": 1, "end": 2 }, + }, + "sourceContent": [{ "text": "Source content" }], + "title": "document", + }, + }, + }, + })])]); + + let result = model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + + let parts = stream_parts(result).await; + let errors: Vec = parts + .iter() + .map(part_json) + .filter(|part| part["type"] == json!("error")) + .collect(); + assert_eq!(errors, Vec::::new()); + } + + mod reasoning { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › reasoning › should stream reasoning and text parts + #[tokio::test] + async fn should_stream_reasoning_and_text_parts() { + let mock = MockFetch::new(vec![stream_chunks(&fixture_chunks(include_str!( + "fixtures/amazon-bedrock-reasoning.chunks.txt" + )))]); + + let result = model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + + let expected: Value = serde_json::from_str(include_str!( + "snapshots/do_stream_reasoning_should_stream_reasoning_and_text_parts.json" + )) + .unwrap(); + assert_eq!(parts_json(&stream_parts(result).await), expected); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should stream text deltas with metadata and usage + #[tokio::test] + async fn should_stream_text_deltas_with_metadata_and_usage() { + let mock = MockFetch::new(vec![stream_chunks(&[ + json!({ "contentBlockDelta": { "contentBlockIndex": 0, "delta": { "text": "Hello" } } }), + json!({ "contentBlockDelta": { "contentBlockIndex": 1, "delta": { "text": ", " } } }), + json!({ "contentBlockDelta": { "contentBlockIndex": 2, "delta": { "text": "World!" } } }), + json!({ + "metadata": { + "usage": { "inputTokens": 4, "outputTokens": 34, "totalTokens": 38 }, + "metrics": { "latencyMs": 10 }, + }, + }), + json!({ "messageStop": { "stopReason": "stop_sequence" } }), + ])]); + + let result = model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + + assert_eq!( + parts_json(&stream_parts(result).await), + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "id": "0", + "type": "text-start" + }, + { + "delta": "Hello", + "id": "0", + "type": "text-delta" + }, + { + "id": "1", + "type": "text-start" + }, + { + "delta": ", ", + "id": "1", + "type": "text-delta" + }, + { + "id": "2", + "type": "text-start" + }, + { + "delta": "World!", + "id": "2", + "type": "text-delta" + }, + { + "finishReason": { + "raw": "stop_sequence", + "unified": "stop" + }, + "type": "finish", + "usage": { + "inputTokens": { + "cacheRead": 0, + "cacheWrite": 0, + "noCache": 4, + "total": 4 + }, + "outputTokens": { + "text": 34, + "total": 34 + }, + "raw": { + "inputTokens": 4, + "outputTokens": 34, + "totalTokens": 38 + } + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should stream tool deltas + #[tokio::test] + async fn should_stream_tool_deltas() { + let mock = MockFetch::new(vec![stream_chunks(&[ + json!({ + "contentBlockStart": { + "contentBlockIndex": 0, + "start": { "toolUse": { "toolUseId": "tool-use-id", "name": "test-tool" } }, + }, + }), + json!({ + "contentBlockDelta": { + "contentBlockIndex": 0, + "delta": { "toolUse": { "input": "{\"value\":" } }, + }, + }), + json!({ + "contentBlockDelta": { + "contentBlockIndex": 0, + "delta": { "toolUse": { "input": "\"Sparkle Day\"}" } }, + }, + }), + json!({ "contentBlockStop": { "contentBlockIndex": 0 } }), + json!({ "messageStop": { "stopReason": "tool_use" } }), + ])]); + + let result = model(&mock) + .do_stream(&options_with(json!({ + "tools": [ + { + "type": "function", + "name": "test-tool", + "inputSchema": { + "type": "object", + "properties": { "value": { "type": "string" } }, + "required": ["value"], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#", + }, + }, + ], + "toolChoice": { "type": "tool", "toolName": "test-tool" }, + "includeRawChunks": false, + }))) + .await + .expect("stream starts"); + + assert_eq!( + parts_json(&stream_parts(result).await), + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "id": "tool-use-id", + "toolName": "test-tool", + "type": "tool-input-start" + }, + { + "delta": "{\"value\":", + "id": "tool-use-id", + "type": "tool-input-delta" + }, + { + "delta": "\"Sparkle Day\"}", + "id": "tool-use-id", + "type": "tool-input-delta" + }, + { + "id": "tool-use-id", + "type": "tool-input-end" + }, + { + "input": "{\"value\":\"Sparkle Day\"}", + "toolCallId": "tool-use-id", + "toolName": "test-tool", + "type": "tool-call" + }, + { + "finishReason": { + "raw": "tool_use", + "unified": "tool-calls" + }, + "type": "finish", + "usage": { + "inputTokens": {}, + "outputTokens": {} + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should stream parallel tool calls + #[tokio::test] + async fn should_stream_parallel_tool_calls() { + let mock = MockFetch::new(vec![stream_chunks(&[ + json!({ + "contentBlockStart": { + "contentBlockIndex": 0, + "start": { "toolUse": { "toolUseId": "tool-use-id-1", "name": "test-tool-1" } }, + }, + }), + json!({ + "contentBlockDelta": { + "contentBlockIndex": 0, + "delta": { "toolUse": { "input": "{\"value1\":" } }, + }, + }), + json!({ + "contentBlockStart": { + "contentBlockIndex": 1, + "start": { "toolUse": { "toolUseId": "tool-use-id-2", "name": "test-tool-2" } }, + }, + }), + json!({ + "contentBlockDelta": { + "contentBlockIndex": 1, + "delta": { "toolUse": { "input": "{\"value2\":" } }, + }, + }), + json!({ + "contentBlockDelta": { + "contentBlockIndex": 1, + "delta": { "toolUse": { "input": "\"Sparkle Day\"}" } }, + }, + }), + json!({ + "contentBlockDelta": { + "contentBlockIndex": 0, + "delta": { "toolUse": { "input": "\"Sparkle Day\"}" } }, + }, + }), + json!({ "contentBlockStop": { "contentBlockIndex": 0 } }), + json!({ "contentBlockStop": { "contentBlockIndex": 1 } }), + json!({ "messageStop": { "stopReason": "tool_use" } }), + ])]); + + let result = model(&mock) + .do_stream(&options_with(json!({ + "tools": [ + { + "type": "function", + "name": "test-tool-1", + "inputSchema": { + "type": "object", + "properties": { "value1": { "type": "string" } }, + "required": ["value"], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#", + }, + }, + { + "type": "function", + "name": "test-tool-2", + "inputSchema": { + "type": "object", + "properties": { "value2": { "type": "string" } }, + "required": ["value"], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#", + }, + }, + ], + "toolChoice": { "type": "tool", "toolName": "test-tool" }, + "includeRawChunks": false, + }))) + .await + .expect("stream starts"); + + assert_eq!( + parts_json(&stream_parts(result).await), + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "id": "tool-use-id-1", + "toolName": "test-tool-1", + "type": "tool-input-start" + }, + { + "delta": "{\"value1\":", + "id": "tool-use-id-1", + "type": "tool-input-delta" + }, + { + "id": "tool-use-id-2", + "toolName": "test-tool-2", + "type": "tool-input-start" + }, + { + "delta": "{\"value2\":", + "id": "tool-use-id-2", + "type": "tool-input-delta" + }, + { + "delta": "\"Sparkle Day\"}", + "id": "tool-use-id-2", + "type": "tool-input-delta" + }, + { + "delta": "\"Sparkle Day\"}", + "id": "tool-use-id-1", + "type": "tool-input-delta" + }, + { + "id": "tool-use-id-1", + "type": "tool-input-end" + }, + { + "input": "{\"value1\":\"Sparkle Day\"}", + "toolCallId": "tool-use-id-1", + "toolName": "test-tool-1", + "type": "tool-call" + }, + { + "id": "tool-use-id-2", + "type": "tool-input-end" + }, + { + "input": "{\"value2\":\"Sparkle Day\"}", + "toolCallId": "tool-use-id-2", + "toolName": "test-tool-2", + "type": "tool-call" + }, + { + "finishReason": { + "raw": "tool_use", + "unified": "tool-calls" + }, + "type": "finish", + "usage": { + "inputTokens": {}, + "outputTokens": {} + } + } + ]) + ); + } + + /// The parts of a stream that carries one `{ : }` chunk. + async fn exception_stream(kind: &str, message: &str, name: &str) -> Value { + let mock = MockFetch::new(vec![stream_chunks(&[json!({ + kind: { "message": message, "name": name, "$fault": "server", "$metadata": {} }, + })])]); + let result = model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + parts_json(&stream_parts(result).await) + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should handle error stream parts + #[tokio::test] + async fn should_handle_error_stream_parts() { + assert_eq!( + exception_stream( + "internalServerException", + "Internal Server Error", + "InternalServerException" + ) + .await, + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "error": { + "data": { + "internalServerException": { + "$fault": "server", + "$metadata": {}, + "message": "Internal Server Error", + "name": "InternalServerException" + } + }, + "isRetryable": true, + "message": "Internal Server Error", + "statusCode": 500, + "type": "internalServerException" + }, + "type": "error" + }, + { + "finishReason": { + "unified": "error" + }, + "type": "finish", + "usage": { + "inputTokens": {}, + "outputTokens": {} + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should handle modelStreamErrorException error + #[tokio::test] + async fn should_handle_model_stream_error_exception_error() { + assert_eq!( + exception_stream( + "modelStreamErrorException", + "Model Stream Error", + "ModelStreamErrorException" + ) + .await, + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "error": { + "data": { + "modelStreamErrorException": { + "$fault": "server", + "$metadata": {}, + "message": "Model Stream Error", + "name": "ModelStreamErrorException" + } + }, + "isRetryable": true, + "message": "Model Stream Error", + "statusCode": 424, + "type": "modelStreamErrorException" + }, + "type": "error" + }, + { + "finishReason": { + "unified": "error" + }, + "type": "finish", + "usage": { + "inputTokens": {}, + "outputTokens": {} + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should handle throttlingException error + #[tokio::test] + async fn should_handle_throttling_exception_error() { + assert_eq!( + exception_stream( + "throttlingException", + "Throttling Error", + "ThrottlingException" + ) + .await, + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "error": { + "data": { + "throttlingException": { + "$fault": "server", + "$metadata": {}, + "message": "Throttling Error", + "name": "ThrottlingException" + } + }, + "isRetryable": true, + "message": "Throttling Error", + "statusCode": 429, + "type": "throttlingException" + }, + "type": "error" + }, + { + "finishReason": { + "unified": "error" + }, + "type": "finish", + "usage": { + "inputTokens": {}, + "outputTokens": {} + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should handle validationException error + #[tokio::test] + async fn should_handle_validation_exception_error() { + assert_eq!( + exception_stream( + "validationException", + "Validation Error", + "ValidationException" + ) + .await, + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "error": { + "data": { + "validationException": { + "$fault": "server", + "$metadata": {}, + "message": "Validation Error", + "name": "ValidationException" + } + }, + "isRetryable": false, + "message": "Validation Error", + "statusCode": 400, + "type": "validationException" + }, + "type": "error" + }, + { + "finishReason": { + "unified": "error" + }, + "type": "finish", + "usage": { + "inputTokens": {}, + "outputTokens": {} + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should handle failed chunk parsing + /// + /// Upstream mocks the response handler to emit a failed parse result + /// carrying an injected error value (`{ message: 'Chunk Parsing Failed' }`). + /// Rust has no seam for that value, so the chunk that fails to parse is a + /// frame whose payload is not JSON; the parts are compared with the error + /// value left out of both sides. + #[tokio::test] + async fn should_handle_failed_chunk_parsing() { + let body = encode_frame( + &[ + (":message-type", "event"), + (":event-type", "contentBlockDelta"), + (":content-type", "application/json"), + ], + b"Chunk Parsing Failed", + ); + let mock = MockFetch::new(vec![stream_response(body, &[])]); + + let result = model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + + let mut actual = parts_json(&stream_parts(result).await); + let mut expected = json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "error": { + "message": "Chunk Parsing Failed" + }, + "type": "error" + }, + { + "finishReason": { + "unified": "error" + }, + "type": "finish", + "usage": { + "inputTokens": {}, + "outputTokens": {} + } + } + ]); + for parts in [&mut actual, &mut expected] { + for part in parts.as_array_mut().expect("parts array") { + if part["type"] == json!("error") { + part.as_object_mut().unwrap().remove("error"); + } + } + } + assert_eq!(actual, expected); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should pass the messages and the model + #[tokio::test] + async fn should_pass_the_messages_and_the_model() { + let mock = MockFetch::new(vec![stream_chunks(&[])]); + + model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + + assert_eq!( + request_body(&mock.seen(), 0), + json!({ + "additionalModelResponseFieldPaths": [ + "/delta/stop_sequence" + ], + "messages": [ + { + "content": [ + { + "text": "Hello" + } + ], + "role": "user" + } + ], + "system": [ + { + "text": "System Prompt" + } + ] + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should support guardrails + #[tokio::test] + async fn should_support_guardrails() { + let mock = MockFetch::new(vec![stream_chunks(&[])]); + + model(&mock) + .do_stream(&options_with(json!({ + "includeRawChunks": false, + "providerOptions": { + "bedrock": { + "guardrailConfig": { + "guardrailIdentifier": "-1", + "guardrailVersion": "1", + "trace": "enabled", + "streamProcessingMode": "async", + }, + }, + }, + }))) + .await + .expect("stream starts"); + + assert_match_object( + &request_body(&mock.seen(), 0), + &json!({ + "messages": [{ "role": "user", "content": [{ "text": "Hello" }] }], + "system": [{ "text": "System Prompt" }], + "guardrailConfig": { + "guardrailIdentifier": "-1", + "guardrailVersion": "1", + "trace": "enabled", + "streamProcessingMode": "async", + }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should accept the new `amazonBedrock` providerOptions key + #[tokio::test] + async fn should_accept_the_new_amazon_bedrock_provider_options_key() { + let mock = MockFetch::new(vec![stream_chunks(&[])]); + + model(&mock) + .do_stream(&options_with(json!({ + "includeRawChunks": false, + "providerOptions": { + "amazonBedrock": { + "guardrailConfig": { + "guardrailIdentifier": "-1", + "guardrailVersion": "1", + "trace": "enabled", + "streamProcessingMode": "async", + }, + }, + }, + }))) + .await + .expect("stream starts"); + + assert_match_object( + &request_body(&mock.seen(), 0), + &json!({ + "guardrailConfig": { + "guardrailIdentifier": "-1", + "guardrailVersion": "1", + "trace": "enabled", + "streamProcessingMode": "async", + }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should prefer `amazonBedrock` providerOptions over the legacy `bedrock` key + #[tokio::test] + async fn should_prefer_amazon_bedrock_provider_options_over_the_legacy_bedrock_key() { + let mock = MockFetch::new(vec![stream_chunks(&[])]); + + model(&mock) + .do_stream(&options_with(json!({ + "includeRawChunks": false, + "providerOptions": { + "bedrock": { + "guardrailConfig": { "guardrailIdentifier": "legacy", "guardrailVersion": "1" }, + }, + "amazonBedrock": { + "guardrailConfig": { "guardrailIdentifier": "preferred", "guardrailVersion": "2" }, + }, + }, + }))) + .await + .expect("stream starts"); + + assert_match_object( + &request_body(&mock.seen(), 0), + &json!({ + "guardrailConfig": { "guardrailIdentifier": "preferred", "guardrailVersion": "2" }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should send text parts as guardContent blocks in streaming requests + #[tokio::test] + async fn should_send_text_parts_as_guard_content_blocks_in_streaming_requests() { + let mock = MockFetch::new(vec![stream_chunks(&[])]); + + model(&mock) + .do_stream(&call_options(json!({ + "prompt": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "What is the capital of Japan?", + "providerOptions": { + "bedrock": { + "guardContent": true, + "guardContentQualifiers": ["query"], + }, + }, + }, + ], + }, + ], + "includeRawChunks": false, + "providerOptions": { + "bedrock": { + "guardrailConfig": { + "guardrailIdentifier": "-1", + "guardrailVersion": "1", + "trace": "enabled", + }, + }, + }, + }))) + .await + .expect("stream starts"); + + assert_match_object( + &request_body(&mock.seen(), 0), + &json!({ + "messages": [ + { + "role": "user", + "content": [ + { + "guardContent": { + "text": { + "text": "What is the capital of Japan?", + "qualifiers": ["query"], + }, + }, + }, + ], + }, + ], + "guardrailConfig": { + "guardrailIdentifier": "-1", + "guardrailVersion": "1", + "trace": "enabled", + }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should map guardrail_intervened messageStop to a content-filter finish reason in streaming + #[tokio::test] + async fn should_map_guardrail_intervened_message_stop_to_a_content_filter_finish_reason() { + let mock = MockFetch::new(vec![stream_chunks(&[ + json!({ + "contentBlockDelta": { + "contentBlockIndex": 0, + "delta": { "text": "Sorry, the model cannot answer this question." }, + }, + }), + json!({ + "metadata": { + "usage": { "inputTokens": 18, "outputTokens": 12, "totalTokens": 30 }, + "metrics": { "latencyMs": 412 }, + "trace": { + "guardrail": { + "inputAssessment": { + "1abcd2ef34gh": { + "topicPolicy": { + "topics": [ + { "action": "BLOCKED", "name": "Restricted topic", "type": "DENY" }, + ], + }, + }, + }, + }, + }, + }, + }), + json!({ "messageStop": { "stopReason": "guardrail_intervened" } }), + ])]); + + let result = model(&mock) + .do_stream(&call_options(json!({ + "prompt": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Tell me something restricted.", + "providerOptions": { + "bedrock": { + "guardContent": true, + "guardContentQualifiers": ["query"], + }, + }, + }, + ], + }, + ], + "includeRawChunks": false, + }))) + .await + .expect("stream starts"); + + let finish = find_part(&stream_parts(result).await, "finish").expect("finish part"); + assert_eq!( + finish["finishReason"], + json!({ + "raw": "guardrail_intervened", + "unified": "content-filter" + }) + ); + assert_eq!( + finish["providerMetadata"]["bedrock"]["trace"], + json!({ + "guardrail": { + "inputAssessment": { + "1abcd2ef34gh": { + "topicPolicy": { + "topics": [ + { + "action": "BLOCKED", + "name": "Restricted topic", + "type": "DENY" + } + ] + } + } + } + } + }) + ); + } + + /// `mockTrace` of the upstream test file. + fn mock_trace() -> Value { + json!({ + "guardrail": { + "inputAssessment": { + "1abcd2ef34gh": { + "contentPolicy": { + "filters": [{ "action": "BLOCKED", "confidence": "LOW", "type": "INSULTS" }], + }, + "wordPolicy": { + "managedWordLists": [ + { "action": "BLOCKED", "match": "", "type": "PROFANITY" }, + ], + }, + }, + }, + }, + }) + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should include trace information in providerMetadata + #[tokio::test] + async fn should_include_trace_information_in_provider_metadata() { + let mock = MockFetch::new(vec![stream_chunks(&[ + json!({ "contentBlockDelta": { "contentBlockIndex": 0, "delta": { "text": "Hello" } } }), + json!({ + "metadata": { + "usage": { "inputTokens": 4, "outputTokens": 34, "totalTokens": 38 }, + "metrics": { "latencyMs": 10 }, + "trace": mock_trace(), + }, + }), + json!({ "messageStop": { "stopReason": "stop_sequence" } }), + ])]); + + let result = model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + + assert_eq!( + parts_json(&stream_parts(result).await), + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "id": "0", + "type": "text-start" + }, + { + "delta": "Hello", + "id": "0", + "type": "text-delta" + }, + { + "finishReason": { + "raw": "stop_sequence", + "unified": "stop" + }, + "providerMetadata": { + "amazonBedrock": { + "trace": { + "guardrail": { + "inputAssessment": { + "1abcd2ef34gh": { + "contentPolicy": { + "filters": [ + { + "action": "BLOCKED", + "confidence": "LOW", + "type": "INSULTS" + } + ] + }, + "wordPolicy": { + "managedWordLists": [ + { + "action": "BLOCKED", + "match": "", + "type": "PROFANITY" + } + ] + } + } + } + } + } + }, + "bedrock": { + "trace": { + "guardrail": { + "inputAssessment": { + "1abcd2ef34gh": { + "contentPolicy": { + "filters": [ + { + "action": "BLOCKED", + "confidence": "LOW", + "type": "INSULTS" + } + ] + }, + "wordPolicy": { + "managedWordLists": [ + { + "action": "BLOCKED", + "match": "", + "type": "PROFANITY" + } + ] + } + } + } + } + } + } + }, + "type": "finish", + "usage": { + "inputTokens": { + "cacheRead": 0, + "cacheWrite": 0, + "noCache": 4, + "total": 4 + }, + "outputTokens": { + "text": 34, + "total": 34 + }, + "raw": { + "inputTokens": 4, + "outputTokens": 34, + "totalTokens": 38 + } + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should include stop_sequence in provider metadata + #[tokio::test] + async fn should_include_stop_sequence_in_provider_metadata() { + let mock = MockFetch::new(vec![stream_chunks(&[ + json!({ "contentBlockDelta": { "contentBlockIndex": 0, "delta": { "text": "Hello" } } }), + json!({ "metadata": { "usage": { "inputTokens": 4, "outputTokens": 34, "totalTokens": 38 } } }), + json!({ + "messageStop": { + "stopReason": "stop_sequence", + "additionalModelResponseFields": { "delta": { "stop_sequence": "STOP" } }, + }, + }), + ])]); + + let result = model(&mock) + .do_stream(&options_with(json!({ "stopSequences": ["STOP"] }))) + .await + .expect("stream starts"); + + let finishes: Vec = stream_parts(result) + .await + .iter() + .map(part_json) + .filter(|part| part["type"] == json!("finish")) + .collect(); + assert_eq!( + Value::Array(finishes), + json!([ + { + "finishReason": { + "raw": "stop_sequence", + "unified": "stop" + }, + "providerMetadata": { + "amazonBedrock": { + "stopSequence": "STOP" + }, + "bedrock": { + "stopSequence": "STOP" + } + }, + "type": "finish", + "usage": { + "inputTokens": { + "cacheRead": 0, + "cacheWrite": 0, + "noCache": 4, + "total": 4 + }, + "outputTokens": { + "text": 34, + "total": 34 + }, + "raw": { + "inputTokens": 4, + "outputTokens": 34, + "totalTokens": 38 + } + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should correctly parse delta.stop_sequence structure in streaming with additional fields + #[tokio::test] + async fn should_correctly_parse_delta_stop_sequence_structure_in_streaming() { + let mock = MockFetch::new(vec![stream_chunks(&[ + json!({ "contentBlockDelta": { "contentBlockIndex": 0, "delta": { "text": "Response" } } }), + json!({ "metadata": { "usage": { "inputTokens": 5, "outputTokens": 10, "totalTokens": 15 } } }), + json!({ + "messageStop": { + "stopReason": "stop_sequence", + "additionalModelResponseFields": { + "delta": { "stop_sequence": "CUSTOM_END" }, + "otherField": "value", + }, + }, + }), + ])]); + + let result = model(&mock) + .do_stream(&options_with(json!({ "stopSequences": ["CUSTOM_END"] }))) + .await + .expect("stream starts"); + + let finish = find_part(&stream_parts(result).await, "finish").expect("finish part"); + assert_eq!( + finish["providerMetadata"]["bedrock"]["stopSequence"], + json!("CUSTOM_END") + ); + assert_eq!( + finish["finishReason"], + json!({ "unified": "stop", "raw": "stop_sequence" }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should include response headers in rawResponse + #[tokio::test] + async fn should_include_response_headers_in_raw_response() { + let mock = MockFetch::new(vec![stream_response( + event_stream_body(&[ + json!({ "contentBlockDelta": { "contentBlockIndex": 0, "delta": { "text": "Hello" } } }), + json!({ "messageStop": { "stopReason": "stop_sequence" } }), + ]), + &[ + ("x-amzn-requestid", "test-request-id"), + ("x-amzn-trace-id", "test-trace-id"), + ], + )]); + + let result = model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + + assert_eq!( + serde_json::to_value(result.response.and_then(|response| response.headers)).unwrap(), + json!({ + "cache-control": "no-cache", + "connection": "keep-alive", + "content-type": "text/event-stream", + "x-amzn-requestid": "test-request-id", + "x-amzn-trace-id": "test-trace-id", + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should extract response metadata + #[tokio::test] + async fn should_extract_response_metadata() { + let mock = MockFetch::new(vec![stream_response( + event_stream_body(&[ + json!({ "contentBlockDelta": { "contentBlockIndex": 0, "delta": { "text": "Hello" } } }), + json!({ "messageStop": { "stopReason": "end_turn" } }), + ]), + &[ + ("x-amzn-requestid", "test-request-id"), + ("date", "Wed, 01 Jan 2025 00:00:00 GMT"), + ], + )]); + + let result = model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + + let metadata = find_part(&stream_parts(result).await, "response-metadata"); + assert_eq!( + metadata, + Some(json!({ + "id": "test-request-id", + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "timestamp": "2025-01-01T00:00:00.000Z", + "type": "response-metadata" + })) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should properly combine headers from all sources + #[tokio::test] + async fn should_properly_combine_headers_from_all_sources() { + let mock = MockFetch::new(vec![stream_response( + event_stream_body(&[ + json!({ "contentBlockDelta": { "contentBlockIndex": 0, "delta": { "text": "Hello" } } }), + json!({ "messageStop": { "stopReason": "stop_sequence" } }), + ]), + &[ + ("x-amzn-requestid", "test-request-id"), + ("x-amzn-trace-id", "test-trace-id"), + ], + )]); + let model = provider( + inject_fetch_headers( + vec![ + ("options-header", "options-value"), + ("model-header", "model-value"), + ("shared-header", "options-shared"), + ("signed-header", "signed-value"), + ("authorization", "AWS4-HMAC-SHA256..."), + ], + mock.transport(), + ), + Some( + [ + ("model-header".to_string(), Some("model-value".to_string())), + ( + "shared-header".to_string(), + Some("model-shared".to_string()), + ), + ] + .into_iter() + .collect(), + ), + ) + .chat(MODEL_ID); + + let mut options = options_with(json!({ "includeRawChunks": false })); + options.headers = Some( + [ + ("options-header".to_string(), "options-value".to_string()), + ("shared-header".to_string(), "options-shared".to_string()), + ] + .into_iter() + .collect(), + ); + model.do_stream(&options).await.expect("stream starts"); + + let request_headers = &mock.seen()[0].headers; + assert_eq!( + request_headers.get("options-header").map(String::as_str), + Some("options-value") + ); + assert_eq!( + request_headers.get("model-header").map(String::as_str), + Some("model-value") + ); + assert_eq!( + request_headers.get("signed-header").map(String::as_str), + Some("signed-value") + ); + assert_eq!( + request_headers.get("authorization").map(String::as_str), + Some("AWS4-HMAC-SHA256...") + ); + assert_eq!( + request_headers.get("shared-header").map(String::as_str), + Some("options-shared") + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should work with partial headers + #[tokio::test] + async fn should_work_with_partial_headers() { + let mock = MockFetch::new(vec![stream_chunks(&[])]); + let model = provider( + inject_fetch_headers( + vec![ + ("model-header", "model-value"), + ("signed-header", "signed-value"), + ("authorization", "AWS4-HMAC-SHA256..."), + ], + mock.transport(), + ), + Some( + [("model-header".to_string(), Some("model-value".to_string()))] + .into_iter() + .collect(), + ), + ) + .chat(MODEL_ID); + + model + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + + let request_headers = &mock.seen()[0].headers; + assert_eq!( + request_headers.get("model-header").map(String::as_str), + Some("model-value") + ); + assert_eq!( + request_headers.get("signed-header").map(String::as_str), + Some("signed-value") + ); + assert_eq!( + request_headers.get("authorization").map(String::as_str), + Some("AWS4-HMAC-SHA256...") + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should include providerOptions in the request for streaming calls + #[tokio::test] + async fn should_include_provider_options_in_the_request_for_streaming_calls() { + let mock = MockFetch::new(vec![stream_chunks(&[ + json!({ "contentBlockDelta": { "contentBlockIndex": 0, "delta": { "text": "Dummy" } } }), + json!({ "messageStop": { "stopReason": "stop_sequence" } }), + ])]); + + model(&mock) + .do_stream(&options_with(json!({ + "includeRawChunks": false, + "providerOptions": { "bedrock": { "foo": "bar" } }, + }))) + .await + .expect("stream starts"); + + assert_match_object(&request_body(&mock.seen(), 0), &json!({ "foo": "bar" })); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should include cache token usage in providerMetadata + #[tokio::test] + async fn should_include_cache_token_usage_in_provider_metadata() { + let mock = MockFetch::new(vec![stream_chunks(&[ + json!({ "contentBlockDelta": { "contentBlockIndex": 0, "delta": { "text": "Hello" } } }), + json!({ + "metadata": { + "usage": { + "inputTokens": 4, + "outputTokens": 34, + "totalTokens": 38, + "cacheReadInputTokens": 2, + "cacheWriteInputTokens": 3, + }, + }, + }), + json!({ "messageStop": { "stopReason": "stop_sequence" } }), + ])]); + + let result = model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + + assert_eq!( + parts_json(&stream_parts(result).await), + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "id": "0", + "type": "text-start" + }, + { + "delta": "Hello", + "id": "0", + "type": "text-delta" + }, + { + "finishReason": { + "raw": "stop_sequence", + "unified": "stop" + }, + "providerMetadata": { + "amazonBedrock": { + "usage": { + "cacheWriteInputTokens": 3 + } + }, + "bedrock": { + "usage": { + "cacheWriteInputTokens": 3 + } + } + }, + "type": "finish", + "usage": { + "inputTokens": { + "cacheRead": 2, + "cacheWrite": 3, + "noCache": 4, + "total": 9 + }, + "outputTokens": { + "text": 34, + "total": 34 + }, + "raw": { + "cacheReadInputTokens": 2, + "cacheWriteInputTokens": 3, + "inputTokens": 4, + "outputTokens": 34, + "totalTokens": 38 + } + } + } + ]) + ); + } + + /// The finish part's provider metadata of a `Hello` stream whose + /// `metadata` chunk is `metadata`. + async fn finish_provider_metadata(metadata: Value) -> Value { + let mock = MockFetch::new(vec![stream_chunks(&[ + json!({ "contentBlockDelta": { "contentBlockIndex": 0, "delta": { "text": "Hello" } } }), + json!({ "metadata": metadata }), + json!({ "messageStop": { "stopReason": "end_turn" } }), + ])]); + let result = model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + let finish = find_part(&stream_parts(result).await, "finish").expect("finish part"); + finish["providerMetadata"].clone() + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should include performanceConfig in providerMetadata + #[tokio::test] + async fn should_include_performance_config_in_provider_metadata() { + let provider_metadata = finish_provider_metadata(json!({ + "usage": { "inputTokens": 4, "outputTokens": 34 }, + "performanceConfig": { "latency": "optimized" }, + })) + .await; + + assert_eq!( + provider_metadata, + json!({ + "amazonBedrock": { + "performanceConfig": { + "latency": "optimized" + } + }, + "bedrock": { + "performanceConfig": { + "latency": "optimized" + } + } + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should include serviceTier in providerMetadata + #[tokio::test] + async fn should_include_service_tier_in_provider_metadata() { + let provider_metadata = finish_provider_metadata(json!({ + "usage": { "inputTokens": 4, "outputTokens": 34 }, + "serviceTier": { "type": "on-demand" }, + })) + .await; + + assert_eq!( + provider_metadata, + json!({ + "amazonBedrock": { + "serviceTier": { + "type": "on-demand" + } + }, + "bedrock": { + "serviceTier": { + "type": "on-demand" + } + } + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should include cacheDetails in providerMetadata + #[tokio::test] + async fn should_include_cache_details_in_provider_metadata() { + let provider_metadata = finish_provider_metadata(json!({ + "usage": { + "inputTokens": 4, + "outputTokens": 34, + "cacheDetails": [{ "inputTokens": 100, "ttl": "T5M" }], + }, + })) + .await; + + assert_eq!( + provider_metadata, + json!({ + "amazonBedrock": { + "usage": { + "cacheDetails": [ + { + "inputTokens": 100, + "ttl": "T5M" + } + ] + } + }, + "bedrock": { + "usage": { + "cacheDetails": [ + { + "inputTokens": 100, + "ttl": "T5M" + } + ] + } + } + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should handle system messages with cache points + #[tokio::test] + async fn should_handle_system_messages_with_cache_points() { + let mock = MockFetch::new(vec![stream_chunks(&[ + json!({ "contentBlockDelta": { "contentBlockIndex": 0, "delta": { "text": "Hello" } } }), + json!({ "messageStop": { "stopReason": "stop_sequence" } }), + ])]); + + model(&mock) + .do_stream(&call_options(json!({ + "prompt": [ + { + "role": "system", + "content": "System Prompt", + "providerOptions": { "bedrock": { "cachePoint": { "type": "default" } } }, + }, + { "role": "user", "content": [{ "type": "text", "text": "Hello" }] }, + ], + "includeRawChunks": false, + }))) + .await + .expect("stream starts"); + + assert_match_object( + &request_body(&mock.seen(), 0), + &json!({ + "system": [{ "text": "System Prompt" }, { "cachePoint": { "type": "default" } }], + "messages": [{ "role": "user", "content": [{ "text": "Hello" }] }], + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should stream reasoning text deltas + #[tokio::test] + async fn should_stream_reasoning_text_deltas() { + let mock = MockFetch::new(vec![stream_chunks(&[ + json!({ + "contentBlockDelta": { + "contentBlockIndex": 0, + "delta": { "reasoningContent": { "text": "I am thinking" } }, + }, + }), + json!({ + "contentBlockDelta": { + "contentBlockIndex": 0, + "delta": { "reasoningContent": { "text": " about this problem..." } }, + }, + }), + json!({ + "contentBlockDelta": { + "contentBlockIndex": 0, + "delta": { "reasoningContent": { "signature": "abc123signature" } }, + }, + }), + json!({ + "contentBlockDelta": { + "contentBlockIndex": 1, + "delta": { "text": "Based on my reasoning, the answer is 42." }, + }, + }), + json!({ "messageStop": { "stopReason": "stop_sequence" } }), + ])]); + + let result = model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + + assert_eq!( + parts_json(&stream_parts(result).await), + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "id": "0", + "type": "reasoning-start" + }, + { + "delta": "I am thinking", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " about this problem...", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": "", + "id": "0", + "providerMetadata": { + "amazonBedrock": { + "signature": "abc123signature" + }, + "bedrock": { + "signature": "abc123signature" + } + }, + "type": "reasoning-delta" + }, + { + "id": "1", + "type": "text-start" + }, + { + "delta": "Based on my reasoning, the answer is 42.", + "id": "1", + "type": "text-delta" + }, + { + "finishReason": { + "raw": "stop_sequence", + "unified": "stop" + }, + "type": "finish", + "usage": { + "inputTokens": {}, + "outputTokens": {} + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should stream redacted reasoning + #[tokio::test] + async fn should_stream_redacted_reasoning() { + let mock = MockFetch::new(vec![stream_chunks(&[ + json!({ + "contentBlockDelta": { + "contentBlockIndex": 0, + "delta": { "reasoningContent": { "data": "redacted-reasoning-data" } }, + }, + }), + json!({ + "contentBlockDelta": { + "contentBlockIndex": 1, + "delta": { "text": "Here is my answer." }, + }, + }), + json!({ "messageStop": { "stopReason": "stop_sequence" } }), + ])]); + + let result = model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": false }))) + .await + .expect("stream starts"); + + assert_eq!( + parts_json(&stream_parts(result).await), + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "id": "0", + "type": "reasoning-start" + }, + { + "delta": "", + "id": "0", + "providerMetadata": { + "amazonBedrock": { + "redactedData": "redacted-reasoning-data" + }, + "bedrock": { + "redactedData": "redacted-reasoning-data" + } + }, + "type": "reasoning-delta" + }, + { + "id": "1", + "type": "text-start" + }, + { + "delta": "Here is my answer.", + "id": "1", + "type": "text-delta" + }, + { + "finishReason": { + "raw": "stop_sequence", + "unified": "stop" + }, + "type": "finish", + "usage": { + "inputTokens": {}, + "outputTokens": {} + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should include raw chunks when includeRawChunks is true + #[tokio::test] + async fn should_include_raw_chunks_when_include_raw_chunks_is_true() { + let mock = MockFetch::new(vec![stream_chunks(&[ + json!({ "contentBlockDelta": { "contentBlockIndex": 0, "delta": { "text": "Hello" } } }), + json!({ "messageStop": { "stopReason": "stop_sequence" } }), + ])]); + + let result = model(&mock) + .do_stream(&options_with(json!({ "includeRawChunks": true }))) + .await + .expect("stream starts"); + + assert_eq!( + parts_json(&stream_parts(result).await), + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "rawValue": { + "contentBlockDelta": { + "contentBlockIndex": 0, + "delta": { + "text": "Hello" + } + } + }, + "type": "raw" + }, + { + "id": "0", + "type": "text-start" + }, + { + "delta": "Hello", + "id": "0", + "type": "text-delta" + }, + { + "rawValue": { + "messageStop": { + "stopReason": "stop_sequence" + } + }, + "type": "raw" + }, + { + "finishReason": { + "raw": "stop_sequence", + "unified": "stop" + }, + "type": "finish", + "usage": { + "inputTokens": {}, + "outputTokens": {} + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should transform reasoningConfig to thinking in stream requests + #[tokio::test] + async fn should_transform_reasoning_config_to_thinking_in_stream_requests() { + let mock = MockFetch::new(vec![stream_chunks(&[ + json!({ "contentBlockDelta": { "contentBlockIndex": 0, "delta": { "text": "Hello" } } }), + json!({ "messageStop": { "stopReason": "stop_sequence" } }), + ])]); + + model(&mock) + .do_stream(&options_with(json!({ + "maxOutputTokens": 100, + "includeRawChunks": false, + "providerOptions": { + "bedrock": { "reasoningConfig": { "type": "enabled", "budgetTokens": 2000 } }, + }, + }))) + .await + .expect("stream starts"); + + let request_body = request_body(&mock.seen(), 0); + assert_match_object( + &request_body, + &json!({ + "additionalModelRequestFields": { + "thinking": { "type": "enabled", "budget_tokens": 2000 }, + }, + "inferenceConfig": { "maxTokens": 2100 }, + }), + ); + assert_eq!(request_body.get("reasoningConfig"), None); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › merges user additionalModelRequestFields with derived thinking (stream) + #[tokio::test] + async fn merges_user_additional_model_request_fields_with_derived_thinking_stream() { + let mock = MockFetch::new(vec![stream_chunks(&[json!({ + "messageStop": { "stopReason": "stop_sequence" }, + })])]); + + model(&mock) + .do_stream(&options_with(json!({ + "includeRawChunks": false, + "providerOptions": { + "bedrock": { + "reasoningConfig": { "type": "enabled", "budgetTokens": 500 }, + "additionalModelRequestFields": { "foo": "bar", "custom": 42 }, + }, + }, + }))) + .await + .expect("stream starts"); + + let body = request_body(&mock.seen(), 0); + assert_eq!(body.get("reasoningConfig"), None); + assert_match_object( + &body["additionalModelRequestFields"], + &json!({ + "foo": "bar", + "custom": 42, + "thinking": { "type": "enabled", "budget_tokens": 500 }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should pass serviceTier provider option in stream requests + #[tokio::test] + async fn should_pass_service_tier_provider_option_in_stream_requests() { + let mock = MockFetch::new(vec![stream_chunks(&[json!({ + "messageStop": { "stopReason": "stop_sequence" }, + })])]); + + model(&mock) + .do_stream(&options_with(json!({ + "includeRawChunks": false, + "providerOptions": { "bedrock": { "serviceTier": "priority" } }, + }))) + .await + .expect("stream starts"); + + let request_body = request_body(&mock.seen(), 0); + assert_match_object( + &request_body, + &json!({ "serviceTier": { "type": "priority" } }), + ); + assert_eq!( + request_body + .get("additionalModelRequestFields") + .and_then(|fields| fields.get("serviceTier")), + None + ); + } + + /// The parts of a fixture stream answered for `options`. + async fn fixture_stream(chunks: &str, options: Value) -> Value { + let mock = MockFetch::new(vec![stream_chunks(&fixture_chunks(chunks))]); + let result = model(&mock) + .do_stream(&call_options(options)) + .await + .expect("stream starts"); + parts_json(&stream_parts(result).await) + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should handle JSON response format in streaming + #[tokio::test] + async fn should_handle_json_response_format_in_streaming() { + let parts = fixture_stream( + include_str!("fixtures/amazon-bedrock-json-tool.1.chunks.txt"), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Generate JSON" }] }], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { "value": { "type": "string" } }, + "required": ["value"], + }, + }, + "includeRawChunks": false, + }), + ) + .await; + + assert_eq!( + parts, + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "id": "0", + "type": "text-start" + }, + { + "delta": "{\"value\":\"test\"}", + "id": "0", + "type": "text-delta" + }, + { + "id": "0", + "type": "text-end" + }, + { + "finishReason": { + "raw": "tool_use", + "unified": "stop" + }, + "providerMetadata": { + "amazonBedrock": { + "isJsonResponseFromTool": true, + "stopSequence": null + }, + "bedrock": { + "isJsonResponseFromTool": true, + "stopSequence": null + } + }, + "type": "finish", + "usage": { + "inputTokens": {}, + "outputTokens": {} + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should include text content before JSON tool call in streaming + #[tokio::test] + async fn should_include_text_content_before_json_tool_call_in_streaming() { + let parts = fixture_stream( + include_str!("fixtures/amazon-bedrock-json-tool.2.chunks.txt"), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Generate JSON" }] }], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { "value": { "type": "string" } }, + "required": ["value"], + }, + }, + "includeRawChunks": false, + }), + ) + .await; + + assert_eq!( + parts, + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "id": "0", + "type": "text-start" + }, + { + "delta": "Let me generate", + "id": "0", + "type": "text-delta" + }, + { + "delta": " that JSON for you.", + "id": "0", + "type": "text-delta" + }, + { + "id": "1", + "type": "text-start" + }, + { + "delta": "{\"value\":\"test\"}", + "id": "1", + "type": "text-delta" + }, + { + "id": "1", + "type": "text-end" + }, + { + "finishReason": { + "raw": "tool_use", + "unified": "stop" + }, + "providerMetadata": { + "amazonBedrock": { + "isJsonResponseFromTool": true, + "stopSequence": null + }, + "bedrock": { + "isJsonResponseFromTool": true, + "stopSequence": null + } + }, + "type": "finish", + "usage": { + "inputTokens": {}, + "outputTokens": {} + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should stream text introduction before JSON-only response + #[tokio::test] + async fn should_stream_text_introduction_before_json_only_response() { + let parts = fixture_stream( + include_str!("fixtures/amazon-bedrock-json-only-text-first.1.chunks.txt"), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Return name data" }] }], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { "name": { "type": "string" } }, + "required": ["name"], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#", + }, + }, + "includeRawChunks": false, + }), + ) + .await; + + assert_eq!( + parts, + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "id": "0", + "type": "text-start" + }, + { + "delta": "Here is your data:", + "id": "0", + "type": "text-delta" + }, + { + "id": "0", + "type": "text-end" + }, + { + "id": "1", + "type": "text-start" + }, + { + "delta": "{\"name\":\"John\"}", + "id": "1", + "type": "text-delta" + }, + { + "id": "1", + "type": "text-end" + }, + { + "finishReason": { + "raw": "tool_use", + "unified": "stop" + }, + "providerMetadata": { + "amazonBedrock": { + "isJsonResponseFromTool": true, + "stopSequence": null + }, + "bedrock": { + "isJsonResponseFromTool": true, + "stopSequence": null + } + }, + "type": "finish", + "usage": { + "inputTokens": { + "cacheRead": 0, + "cacheWrite": 0, + "noCache": 100, + "total": 100 + }, + "outputTokens": { + "text": 20, + "total": 20 + }, + "raw": { + "inputTokens": 100, + "outputTokens": 20, + "totalTokens": 120 + } + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should include multiple text blocks before JSON tool call in streaming + #[tokio::test] + async fn should_include_multiple_text_blocks_before_json_tool_call_in_streaming() { + let parts = fixture_stream( + include_str!("fixtures/amazon-bedrock-json-tool.3.chunks.txt"), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Generate data" }] }], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { "result": { "type": "number" } }, + "required": ["result"], + }, + }, + "includeRawChunks": false, + }), + ) + .await; + + assert_eq!( + parts, + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "id": "0", + "type": "text-start" + }, + { + "delta": "First text block.", + "id": "0", + "type": "text-delta" + }, + { + "id": "1", + "type": "text-start" + }, + { + "delta": "Second text block.", + "id": "1", + "type": "text-delta" + }, + { + "id": "2", + "type": "text-start" + }, + { + "delta": "Third block.", + "id": "2", + "type": "text-delta" + }, + { + "id": "3", + "type": "text-start" + }, + { + "delta": "{\"result\":42}", + "id": "3", + "type": "text-delta" + }, + { + "id": "3", + "type": "text-end" + }, + { + "finishReason": { + "raw": "tool_use", + "unified": "stop" + }, + "providerMetadata": { + "amazonBedrock": { + "isJsonResponseFromTool": true, + "stopSequence": null + }, + "bedrock": { + "isJsonResponseFromTool": true, + "stopSequence": null + } + }, + "type": "finish", + "usage": { + "inputTokens": {}, + "outputTokens": {} + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should handle regular tool call before JSON tool call in streaming + #[tokio::test] + async fn should_handle_regular_tool_call_before_json_tool_call_in_streaming() { + let parts = fixture_stream( + include_str!("fixtures/amazon-bedrock-json-with-tool.1.chunks.txt"), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Test" }] }], + "tools": [ + { + "type": "function", + "name": "get-weather", + "inputSchema": { + "type": "object", + "properties": { "location": { "type": "string" } }, + "required": ["location"], + }, + }, + ], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { "result": { "type": "string" } }, + "required": ["result"], + }, + }, + "includeRawChunks": false, + }), + ) + .await; + + assert_eq!( + parts, + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "id": "0", + "type": "text-start" + }, + { + "delta": "Calling a tool first.", + "id": "0", + "type": "text-delta" + }, + { + "id": "tool-1", + "toolName": "get-weather", + "type": "tool-input-start" + }, + { + "delta": "{\"location\":\"SF\"}", + "id": "tool-1", + "type": "tool-input-delta" + }, + { + "id": "tool-1", + "type": "tool-input-end" + }, + { + "input": "{\"location\":\"SF\"}", + "toolCallId": "tool-1", + "toolName": "get-weather", + "type": "tool-call" + }, + { + "id": "2", + "type": "text-start" + }, + { + "delta": "Now JSON.", + "id": "2", + "type": "text-delta" + }, + { + "id": "3", + "type": "text-start" + }, + { + "delta": "{\"result\":\"data\"}", + "id": "3", + "type": "text-delta" + }, + { + "id": "3", + "type": "text-end" + }, + { + "finishReason": { + "raw": "tool_use", + "unified": "stop" + }, + "providerMetadata": { + "amazonBedrock": { + "isJsonResponseFromTool": true, + "stopSequence": null + }, + "bedrock": { + "isJsonResponseFromTool": true, + "stopSequence": null + } + }, + "type": "finish", + "usage": { + "inputTokens": {}, + "outputTokens": {} + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should handle multiple regular tool calls before JSON tool call in streaming + #[tokio::test] + async fn should_handle_multiple_regular_tool_calls_before_json_tool_call_in_streaming() { + let parts = fixture_stream( + include_str!("fixtures/amazon-bedrock-json-with-tools.1.chunks.txt"), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Test" }] }], + "tools": [ + { + "type": "function", + "name": "tool-a", + "inputSchema": { "type": "object", "properties": { "a": { "type": "string" } } }, + }, + { + "type": "function", + "name": "tool-b", + "inputSchema": { "type": "object", "properties": { "b": { "type": "string" } } }, + }, + ], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { "final": { "type": "string" } }, + "required": ["final"], + }, + }, + "includeRawChunks": false, + }), + ) + .await; + + assert_eq!( + parts, + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "id": "0", + "type": "text-start" + }, + { + "delta": "Multiple tools.", + "id": "0", + "type": "text-delta" + }, + { + "id": "tool-1", + "toolName": "tool-a", + "type": "tool-input-start" + }, + { + "delta": "{\"a\":\"1\"}", + "id": "tool-1", + "type": "tool-input-delta" + }, + { + "id": "tool-1", + "type": "tool-input-end" + }, + { + "input": "{\"a\":\"1\"}", + "toolCallId": "tool-1", + "toolName": "tool-a", + "type": "tool-call" + }, + { + "id": "tool-2", + "toolName": "tool-b", + "type": "tool-input-start" + }, + { + "delta": "{\"b\":\"2\"}", + "id": "tool-2", + "type": "tool-input-delta" + }, + { + "id": "tool-2", + "type": "tool-input-end" + }, + { + "input": "{\"b\":\"2\"}", + "toolCallId": "tool-2", + "toolName": "tool-b", + "type": "tool-call" + }, + { + "id": "3", + "type": "text-start" + }, + { + "delta": "Final JSON.", + "id": "3", + "type": "text-delta" + }, + { + "id": "4", + "type": "text-start" + }, + { + "delta": "{\"final\":\"result\"}", + "id": "4", + "type": "text-delta" + }, + { + "id": "4", + "type": "text-end" + }, + { + "finishReason": { + "raw": "tool_use", + "unified": "stop" + }, + "providerMetadata": { + "amazonBedrock": { + "isJsonResponseFromTool": true, + "stopSequence": null + }, + "bedrock": { + "isJsonResponseFromTool": true, + "stopSequence": null + } + }, + "type": "finish", + "usage": { + "inputTokens": {}, + "outputTokens": {} + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should stream text, then regular tool calls, with JSON response format available + #[tokio::test] + async fn should_stream_text_then_regular_tool_calls_with_json_response_format_available() { + let parts = fixture_stream( + include_str!("fixtures/amazon-bedrock-json-tool-text-then-weather-then-json.1.chunks.txt"), + json!({ + "prompt": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "What's 2+2? Also check the weather in San Francisco and London.", + }, + ], + }, + ], + "tools": [ + { + "type": "function", + "name": "weather", + "inputSchema": { + "type": "object", + "properties": { "location": { "type": "string" } }, + "required": ["location"], + }, + }, + ], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { + "locations": { + "type": "array", + "items": { + "type": "object", + "properties": { + "location": { "type": "string" }, + "temperature": { "type": "number" }, + }, + }, + }, + }, + }, + }, + "includeRawChunks": false, + }), + ) + .await; + + assert_eq!( + parts, + json!([ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "id": "0", + "type": "text-start" + }, + { + "delta": "2 + 2 equals 4. Now let me check", + "id": "0", + "type": "text-delta" + }, + { + "delta": " the weather for you.", + "id": "0", + "type": "text-delta" + }, + { + "id": "0", + "type": "text-end" + }, + { + "id": "weather-tool-1", + "toolName": "weather", + "type": "tool-input-start" + }, + { + "delta": "{\"location\":", + "id": "weather-tool-1", + "type": "tool-input-delta" + }, + { + "delta": "\"San Francisco\"}", + "id": "weather-tool-1", + "type": "tool-input-delta" + }, + { + "id": "weather-tool-1", + "type": "tool-input-end" + }, + { + "input": "{\"location\":\"San Francisco\"}", + "toolCallId": "weather-tool-1", + "toolName": "weather", + "type": "tool-call" + }, + { + "id": "weather-tool-2", + "toolName": "weather", + "type": "tool-input-start" + }, + { + "delta": "{\"location\":\"London\"}", + "id": "weather-tool-2", + "type": "tool-input-delta" + }, + { + "id": "weather-tool-2", + "type": "tool-input-end" + }, + { + "input": "{\"location\":\"London\"}", + "toolCallId": "weather-tool-2", + "toolName": "weather", + "type": "tool-call" + }, + { + "finishReason": { + "raw": "tool_use", + "unified": "tool-calls" + }, + "type": "finish", + "usage": { + "inputTokens": { + "cacheRead": 0, + "cacheWrite": 0, + "noCache": 500, + "total": 500 + }, + "outputTokens": { + "text": 100, + "total": 100 + }, + "raw": { + "inputTokens": 500, + "outputTokens": 100, + "totalTokens": 600 + } + } + } + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should warn when Anthropic model receives maxReasoningEffort in stream + #[tokio::test] + async fn should_warn_when_anthropic_model_receives_max_reasoning_effort_in_stream() { + let mock = MockFetch::new(vec![stream_chunks(&[json!({ + "messageStop": { "stopReason": "stop_sequence" }, + })])]); + + let result = model(&mock) + .do_stream(&options_with(json!({ + "includeRawChunks": false, + "providerOptions": { + "bedrock": { + "reasoningConfig": { "type": "enabled", "maxReasoningEffort": "medium" }, + }, + }, + }))) + .await + .expect("stream starts"); + stream_items(result).await; + + let request_body = request_body(&mock.seen(), 0); + assert_eq!( + request_body + .get("additionalModelRequestFields") + .and_then(|fields| fields.get("reasoningConfig")), + None + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doStream › should support tool calls with empty input (no arguments) + #[tokio::test] + async fn should_support_tool_calls_with_empty_input_no_arguments() { + let mock = MockFetch::new(vec![stream_chunks(&fixture_chunks(include_str!( + "fixtures/amazon-bedrock-tool-no-args.chunks.txt" + )))]); + + let result = model(&mock) + .do_stream(&options_with(json!({ + "tools": [ + { + "type": "function", + "name": "updateIssueList", + "inputSchema": { + "type": "object", + "properties": {}, + "required": [], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#", + }, + }, + ], + "includeRawChunks": false, + }))) + .await + .expect("stream starts"); + + let tool_call_part = find_part(&stream_parts(result).await, "tool-call"); + assert!(tool_call_part.is_some()); + assert_eq!(tool_call_part.unwrap()["input"], json!("{}")); + } +} + +mod do_generate { + #[allow(unused_imports)] + use super::*; + + use std::collections::BTreeMap; + use std::sync::Arc; + + use aimux_core::language_model::LanguageModel; + use aimux_core::options::CallOptions; + use aimux_core::result::GenerateResult; + use aimux_provider_utils::{ + Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, HeaderMapOpt, Resolvable, + }; + use async_trait::async_trait; + use reqwest::header::{HeaderName, HeaderValue}; + use serde_json::{Value, json}; + + use crate::bedrock::{AmazonBedrockProviderSettings, create_amazon_bedrock}; + use crate::test_support::mock_fetch::{Canned, MockFetch, Seen}; + + const BASE_URL: &str = "https://bedrock-runtime.us-east-1.amazonaws.com"; + const MODEL_ID: &str = "anthropic.claude-3-haiku-20240307-v1:0"; + const ANTHROPIC_MODEL_ID: &str = "anthropic.claude-3-5-sonnet-20240620-v1:0"; + const LEGACY_ANTHROPIC_37_MODEL_ID: &str = "us.anthropic.claude-3-7-sonnet-20250219-v1:0"; + const NOVA_MODEL_ID: &str = "us.amazon.nova-2-lite-v1:0"; + const NOVA_MICRO_MODEL_ID: &str = "us.amazon.nova-micro-v1:0"; + const OPENAI_MODEL_ID: &str = "openai.gpt-oss-120b-1:0"; + const US_OPENAI_MODEL_ID: &str = "us.openai.gpt-5.6-luna"; + const GLOBAL_OPENAI_MODEL_ID: &str = "global.openai.gpt-5.6-luna"; + const CUSTOM_OPENAI_SUBSTRING_MODEL_ID: &str = "custom-openai.gpt-5.6-luna"; + const NEWER_ANTHROPIC_MODEL_ID: &str = "anthropic.claude-sonnet-4-6-v1"; + const HAIKU45_ANTHROPIC_MODEL_ID: &str = "us.anthropic.claude-haiku-4-5-20251001-v1:0"; + const NATIVE_STRUCTURED_OUTPUT_ANTHROPIC_MODEL_ID: &str = + "anthropic.claude-sonnet-4-5-20250929-v1:0"; + const OPUS_ANTHROPIC_MODEL_ID: &str = "us.anthropic.claude-opus-4-8"; + const OPUS5_ANTHROPIC_MODEL_ID: &str = "us.anthropic.claude-opus-5"; + const OPUS55_ANTHROPIC_MODEL_ID: &str = "us.anthropic.claude-opus-5-5"; + const SONNET5_ANTHROPIC_MODEL_ID: &str = "us.anthropic.claude-sonnet-5"; + + /// Upstream `TEST_PROMPT`. + fn test_prompt() -> Value { + json!([ + { "role": "system", "content": "System Prompt" }, + { "role": "user", "content": [{ "type": "text", "text": "Hello" }] }, + ]) + } + + /// Upstream `mockTrace`. + fn mock_trace() -> Value { + json!({ + "guardrail": { + "inputAssessment": { + "1abcd2ef34gh": { + "contentPolicy": { + "filters": [ + { "action": "BLOCKED", "confidence": "LOW", "type": "INSULTS" }, + ], + }, + "wordPolicy": { + "managedWordLists": [ + { "action": "BLOCKED", "match": "", "type": "PROFANITY" }, + ], + }, + }, + }, + }, + }) + } + + /// An upstream `__fixtures__/.json` file. + fn fixture(name: &str) -> Value { + let text = match name { + "amazon-bedrock-text" => include_str!("fixtures/amazon-bedrock-text.json"), + "amazon-bedrock-reasoning" => include_str!("fixtures/amazon-bedrock-reasoning.json"), + "amazon-bedrock-guard-content-intervened" => { + include_str!("fixtures/amazon-bedrock-guard-content-intervened.json") + } + "amazon-bedrock-json-tool.1" => { + include_str!("fixtures/amazon-bedrock-json-tool.1.json") + } + "amazon-bedrock-json-tool.2" => { + include_str!("fixtures/amazon-bedrock-json-tool.2.json") + } + "amazon-bedrock-json-tool.3" => { + include_str!("fixtures/amazon-bedrock-json-tool.3.json") + } + "amazon-bedrock-json-tool-with-answer.1" => { + include_str!("fixtures/amazon-bedrock-json-tool-with-answer.1.json") + } + "amazon-bedrock-json-other-tool.1" => { + include_str!("fixtures/amazon-bedrock-json-other-tool.1.json") + } + "amazon-bedrock-json-with-tool.1" => { + include_str!("fixtures/amazon-bedrock-json-with-tool.1.json") + } + "amazon-bedrock-json-with-tools.1" => { + include_str!("fixtures/amazon-bedrock-json-with-tools.1.json") + } + "amazon-bedrock-tool-no-args" => { + include_str!("fixtures/amazon-bedrock-tool-no-args.json") + } + other => panic!("unknown fixture {other}"), + }; + serde_json::from_str(text).expect("fixture is JSON") + } + + /// The test server's `json-value` response: the compact JSON body with + /// `content-type` and `content-length`, plus any extra headers. + fn json_response_with_headers(body: &Value, headers: &[(&str, &str)]) -> Canned { + let mut canned = Canned::json(body); + canned + .headers + .push(("content-length".into(), canned.body.len().to_string())); + for (name, value) in headers { + canned.headers.push(((*name).into(), (*value).into())); + } + canned + } + + fn json_response(body: &Value) -> Canned { + json_response_with_headers(body, &[]) + } + + /// Upstream `prepareJsonFixtureResponse(name)`. + fn fixture_response(name: &str) -> Canned { + json_response(&fixture(name)) + } + + /// The `simpleResponse` of the reasoning describes. + fn simple_response() -> Canned { + hello_response("stop_sequence") + } + + /// `{ output: { content: [{ text: 'Hello' }] }, stopReason }` with usage 1/1/2. + fn hello_response(stop_reason: &str) -> Canned { + json_response(&json!({ + "output": { "message": { "content": [{ "text": "Hello" }], "role": "assistant" } }, + "stopReason": stop_reason, + "usage": { "inputTokens": 1, "outputTokens": 1, "totalTokens": 2 }, + })) + } + + /// A text-only response with usage 4/34/38 plus extra top-level fields. + fn hello_world_response(extra: Value) -> Canned { + let mut body = json!({ + "output": { "message": { "role": "assistant", "content": [{ "type": "text", "text": "Hello, World!" }] } }, + "usage": { "inputTokens": 4, "outputTokens": 34, "totalTokens": 38 }, + "stopReason": "stop_sequence", + }); + body.as_object_mut() + .unwrap() + .extend(extra.as_object().unwrap().clone()); + json_response(&body) + } + + /// A response whose message content is `content` (usage 4/34/38, + /// stopReason `stop_sequence`). + fn content_response(content: Value) -> Canned { + json_response(&json!({ + "output": { "message": { "role": "assistant", "content": content } }, + "usage": { "inputTokens": 4, "outputTokens": 34, "totalTokens": 38 }, + "stopReason": "stop_sequence", + })) + } + + /// The JSON response of a JSON-tool call: `{ name: 'Test' }`. + fn json_tool_name_response() -> Canned { + json_response(&json!({ + "output": { + "message": { + "role": "assistant", + "content": [ + { "toolUse": { "toolUseId": "json-tool-id", "name": "json", "input": { "name": "Test" } } }, + ], + }, + }, + "usage": { "inputTokens": 4, "outputTokens": 10, "totalTokens": 14 }, + "stopReason": "tool_use", + })) + } + + /// The text response of the Anthropic beta cases. + fn test_response_stop() -> Canned { + json_response(&json!({ + "output": { "message": { "role": "assistant", "content": [{ "type": "text", "text": "test response" }] } }, + "usage": { "inputTokens": 10, "outputTokens": 20, "totalTokens": 30 }, + "stopReason": "stop", + })) + } + + /// The `{ name: string }` response format schema many cases use. + fn name_response_format() -> Value { + json!({ + "type": "json", + "schema": { + "type": "object", + "properties": { "name": { "type": "string" } }, + "required": ["name"], + }, + }) + } + + /// The tool-call conversation of the "no active tools" cases. + fn conversation_with_tool_calls() -> Value { + json!([ + { "role": "user", "content": [{ "type": "text", "text": "What is the weather in Toronto?" }] }, + { + "role": "assistant", + "content": [ + { "type": "tool-call", "toolCallId": "tool-call-1", "toolName": "weather", "input": { "city": "Toronto" } }, + ], + }, + { + "role": "tool", + "content": [ + { + "type": "tool-result", + "toolCallId": "tool-call-1", + "toolName": "weather", + "output": { "type": "text", "value": "The weather in Toronto is 20°C." }, + }, + ], + }, + { "role": "user", "content": [{ "type": "text", "text": "Now give me a summary." }] }, + ]) + } + + /// The `bash_20241022` provider tool. + fn bash_tool() -> Value { + json!({ "type": "provider", "id": "anthropic.bash_20241022", "name": "bash", "args": {} }) + } + + /// Upstream `injectFetchHeaders`: sets the headers on every request, then + /// forwards it. + struct InjectHeaders { + headers: Vec<(&'static str, &'static str)>, + inner: FetchFunction, + } + + #[async_trait] + impl Fetch for InjectHeaders { + async fn fetch(&self, mut request: FetchRequest) -> Result { + for (name, value) in &self.headers { + request.headers.insert( + HeaderName::from_static(name), + HeaderValue::from_static(value), + ); + } + self.inner.fetch(request).await + } + } + + /// The settings of upstream's module-level models (`baseUrl`, a fetch + /// that authenticates the request, no extra headers). + fn settings(fetch: FetchFunction) -> AmazonBedrockProviderSettings { + AmazonBedrockProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".into())), + region: Some("us-east-1".into()), + base_url: Some(BASE_URL.into()), + fetch: Some(fetch), + ..Default::default() + } + } + + /// Upstream `LanguageModelV4CallOptions` (camelCase JSON). + fn call_options(options: Value) -> CallOptions { + serde_json::from_value(options).expect("valid call options") + } + + struct Generated { + result: GenerateResult, + seen: Vec, + } + + impl Generated { + /// `server.calls[0].requestBodyJson`. + fn request_body(&self) -> Value { + self.seen[0].json_body() + } + + /// `server.calls[0].requestHeaders`. + fn request_headers(&self) -> &BTreeMap { + &self.seen[0].headers + } + + fn content(&self) -> Value { + to_json(&self.result.content) + } + + fn warnings(&self) -> Value { + to_json(&self.result.warnings) + } + + fn finish_reason(&self) -> Value { + to_json(&self.result.finish_reason) + } + + /// `result.providerMetadata?.bedrock` at a JSON pointer below it. + fn bedrock_metadata(&self, pointer: &str) -> Option { + to_json(&self.result.provider_metadata) + .pointer(&format!("/bedrock{pointer}")) + .cloned() + } + } + + async fn generate_with( + settings: AmazonBedrockProviderSettings, + fetch: &Arc, + model_id: &str, + options: Value, + ) -> Generated { + let model = create_amazon_bedrock(settings) + .expect("provider") + .chat(model_id); + let result = model + .do_generate(&call_options(options)) + .await + .expect("doGenerate succeeds"); + Generated { + result, + seen: fetch.seen(), + } + } + + /// `doGenerate` on a model built like upstream's module-level models. + async fn generate(model_id: &str, response: Canned, options: Value) -> Generated { + let fetch = MockFetch::new(vec![response]); + generate_with(settings(fetch.transport()), &fetch, model_id, options).await + } + + fn to_json(value: &T) -> Value { + serde_json::to_value(value).expect("serializable") + } + + /// Deep equality with JavaScript number semantics (`1` equals `1.0`). + fn js_equal(actual: &Value, expected: &Value) -> bool { + match (actual, expected) { + (Value::Number(a), Value::Number(e)) => a.as_f64() == e.as_f64(), + (Value::Object(a), Value::Object(e)) => { + a.len() == e.len() + && e.iter() + .all(|(key, ev)| a.get(key).is_some_and(|av| js_equal(av, ev))) + } + (Value::Array(a), Value::Array(e)) => { + a.len() == e.len() && a.iter().zip(e).all(|(av, ev)| js_equal(av, ev)) + } + _ => actual == expected, + } + } + + /// `toMatchObject`: every expected key matches, recursively. + fn js_match(actual: &Value, expected: &Value) -> bool { + match (actual, expected) { + (Value::Object(a), Value::Object(e)) => e + .iter() + .all(|(key, ev)| a.get(key).is_some_and(|av| js_match(av, ev))), + (Value::Array(a), Value::Array(e)) => { + a.len() == e.len() && a.iter().zip(e).all(|(av, ev)| js_match(av, ev)) + } + _ => js_equal(actual, expected), + } + } + + /// `expect(actual).toEqual(expected)`. + fn assert_json_eq(actual: &Value, expected: &Value) { + assert!( + js_equal(actual, expected), + "toEqual failed\n actual: {actual:#}\nexpected: {expected:#}" + ); + } + + /// `expect(actual).toMatchObject(expected)`. + fn assert_match_object(actual: &Value, expected: &Value) { + assert!( + js_match(actual, expected), + "toMatchObject failed\n actual: {actual:#}\nexpected: {expected:#}" + ); + } + + /// `expect(list).toContainEqual(expected)`. + fn assert_contains_equal(list: &Value, expected: &Value) { + assert!( + list.as_array() + .is_some_and(|items| items.iter().any(|item| js_equal(item, expected))), + "toContainEqual failed\n list: {list:#}\nexpected: {expected:#}" + ); + } + + /// Whether some warning has this `feature` + /// (`expect.objectContaining({ feature })`). + fn has_warning_feature(g: &Generated, feature: &str) -> bool { + g.warnings() + .as_array() + .unwrap() + .iter() + .any(|warning| warning["feature"] == feature) + } + + /// The value at a JSON pointer (the `?.` chains); `None` is `undefined`. + fn at<'a>(value: &'a Value, pointer: &str) -> Option<&'a Value> { + value.pointer(pointer) + } + + /// The `toolSpec.name` of every tool in the request's `toolConfig`. + fn tool_names(body: &Value) -> Vec { + body["toolConfig"]["tools"] + .as_array() + .expect("toolConfig.tools") + .iter() + .map(|tool| { + tool["toolSpec"]["name"] + .as_str() + .expect("toolSpec.name") + .to_string() + }) + .collect() + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should generate an ID when toolUseId is empty + #[tokio::test] + async fn should_generate_an_id_when_tool_use_id_is_empty() { + let g = generate( + MODEL_ID, + json_response(&json!({ + "output": { + "message": { + "role": "assistant", + "content": [ + { "toolUse": { "toolUseId": "", "name": "test-tool", "input": { "value": "test" } } }, + ], + }, + }, + "stopReason": "tool_use", + "usage": { "inputTokens": 1, "outputTokens": 1, "totalTokens": 2 }, + })), + json!({ "prompt": test_prompt() }), + ) + .await; + + // Upstream injects `generateId: () => 'test-id'`; aimux has no id + // generator seam, so the generated id must be non-empty and the rest + // of the part is compared with the id replaced by `test-id`. + let content = g.content(); + let call = content + .as_array() + .unwrap() + .iter() + .find(|part| part["type"] == "tool-call") + .expect("a tool-call part"); + assert!( + call["toolCallId"].as_str().is_some_and(|id| !id.is_empty()), + "expected a generated toolCallId, got {call:#}" + ); + let mut call = call.clone(); + call["toolCallId"] = json!("test-id"); + assert_contains_equal( + &json!([call]), + &json!({ + "type": "tool-call", + "toolCallId": "test-id", + "toolName": "test-tool", + "input": "{\"value\":\"test\"}", + }), + ); + } + + mod text { + use super::*; + + /// The describe's `beforeEach`. + fn text_response() -> Canned { + fixture_response("amazon-bedrock-text") + } + + fn citations_response(citations_content: Value) -> Canned { + json_response(&json!({ + "output": { + "message": { + "role": "assistant", + "content": [{ "citationsContent": citations_content }], + }, + }, + "stopReason": "end_turn", + "usage": { "inputTokens": 1, "outputTokens": 2, "totalTokens": 3 }, + })) + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › text › should extract text response + #[tokio::test] + async fn should_extract_text_response() { + let g = generate( + MODEL_ID, + text_response(), + json!({ "prompt": test_prompt() }), + ) + .await; + + let snapshot: Value = serde_json::from_str(include_str!( + "snapshots/do_generate_text_should_extract_text_response.json" + )) + .unwrap(); + assert_json_eq(&to_json(&g.result), &snapshot); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › text › should extract text from citation content + #[tokio::test] + async fn should_extract_text_from_citation_content() { + let g = generate( + MODEL_ID, + citations_response(json!({ + "content": [{ "text": "Citation " }, { "text": "response" }], + "citations": [], + })), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &g.content(), + &json!([ + { "type": "text", "text": "Citation " }, + { "type": "text", "text": "response" }, + ]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › text › should ignore citation content when generated content is %s + #[tokio::test] + async fn should_ignore_citation_content_when_generated_content_is_s() { + // `undefined` leaves the key out; `null` sends it as null. + for generated_content in [None, Some(Value::Null)] { + let mut citations_content = json!({ "citations": [] }); + if let Some(value) = generated_content { + citations_content["content"] = value; + } + let g = generate( + MODEL_ID, + citations_response(citations_content), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq(&g.content(), &json!([])); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › text › should ignore citation generated content without text + #[tokio::test] + async fn should_ignore_citation_generated_content_without_text() { + let g = generate( + MODEL_ID, + citations_response(json!({ + "content": [{}, { "text": null }, { "text": "Response" }], + "citations": [], + })), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &g.content(), + &json!([{ "type": "text", "text": "Response" }]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › text › should prefer regular text when citation content is also present + #[tokio::test] + async fn should_prefer_regular_text_when_citation_content_is_also_present() { + let g = generate( + MODEL_ID, + json_response(&json!({ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "text": "Canonical response", + "citationsContent": { + "content": [{ "text": "Duplicate response" }], + "citations": [], + }, + }, + ], + }, + }, + "stopReason": "end_turn", + "usage": { "inputTokens": 1, "outputTokens": 2, "totalTokens": 3 }, + })), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &g.content(), + &json!([{ "type": "text", "text": "Canonical response" }]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › text › should extract usage + #[tokio::test] + async fn should_extract_usage() { + let g = generate( + MODEL_ID, + text_response(), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &to_json(&g.result.usage), + &json!({ + "inputTokens": { "cacheRead": 0, "cacheWrite": 0, "noCache": 22, "total": 22 }, + "outputTokens": { "text": 57, "total": 57 }, + "raw": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 22, + "outputTokens": 57, + "serverToolUsage": {}, + "totalTokens": 79, + }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › text › should send additional response information + #[tokio::test] + async fn should_send_additional_response_information() { + let g = generate( + MODEL_ID, + text_response(), + json!({ "prompt": test_prompt() }), + ) + .await; + + let response = g.result.response.as_ref().expect("response"); + assert_eq!(response.id, None); + assert_eq!(response.timestamp, None); + assert_eq!(response.model_id.as_deref(), Some(MODEL_ID)); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › text › should expose the raw response headers + #[tokio::test] + async fn should_expose_the_raw_response_headers() { + let g = generate( + MODEL_ID, + json_response_with_headers( + &fixture("amazon-bedrock-text"), + &[("test-header", "test-value")], + ), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &to_json(&g.result.response.as_ref().unwrap().headers), + &json!({ + "content-length": "435", + "content-type": "application/json", + "test-header": "test-value", + }), + ); + } + } + + mod reasoning { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › reasoning › should extract reasoning and text response + #[tokio::test] + async fn should_extract_reasoning_and_text_response() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-reasoning"), + json!({ "prompt": test_prompt() }), + ) + .await; + + let snapshot: Value = serde_json::from_str(include_str!( + "snapshots/do_generate_reasoning_should_extract_reasoning_and_text_response.json" + )) + .unwrap(); + assert_json_eq(&to_json(&g.result), &snapshot); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should extract finish reason + #[tokio::test] + async fn should_extract_finish_reason() { + let g = generate( + MODEL_ID, + hello_world_response(json!({})), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &g.finish_reason(), + &json!({ "raw": "stop_sequence", "unified": "stop" }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should support unknown finish reason + #[tokio::test] + async fn should_support_unknown_finish_reason() { + let g = generate( + MODEL_ID, + hello_world_response(json!({ "stopReason": "eos" })), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &g.finish_reason(), + &json!({ "raw": "eos", "unified": "other" }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should pass the model and the messages + #[tokio::test] + async fn should_pass_the_model_and_the_messages() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &g.request_body(), + &json!({ + "additionalModelResponseFieldPaths": ["/delta/stop_sequence"], + "messages": [{ "content": [{ "text": "Hello" }], "role": "user" }], + "system": [{ "text": "System Prompt" }], + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should pass settings + #[tokio::test] + async fn should_pass_settings() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "prompt": test_prompt(), + "maxOutputTokens": 100, + "temperature": 0.5, + "topP": 0.5, + "topK": 1, + }), + ) + .await; + + assert_match_object( + &g.request_body(), + &json!({ + "inferenceConfig": { "maxTokens": 100, "temperature": 0.5, "topP": 0.5, "topK": 1 }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should omit unsupported sampling parameters for %s + #[tokio::test] + async fn should_omit_unsupported_sampling_parameters_for_s() { + for sampling_model_id in [ + "global.anthropic.claude-opus-4-7", + "eu.anthropic.claude-opus-4-8", + "us.anthropic.claude-opus-5", + ] { + let g = generate( + sampling_model_id, + json_response(&json!({ + "output": { "message": { "role": "assistant", "content": [{ "text": "OK" }] } }, + "stopReason": "end_turn", + "usage": { "inputTokens": 1, "outputTokens": 1, "totalTokens": 2 }, + })), + json!({ "prompt": test_prompt(), "temperature": 0.5, "topP": 0.7, "topK": 10 }), + ) + .await; + + let body = g.request_body(); + assert_eq!(at(&body, "/inferenceConfig/temperature"), None); + assert_eq!(at(&body, "/inferenceConfig/topP"), None); + assert_eq!(at(&body, "/inferenceConfig/topK"), None); + assert_json_eq( + &g.warnings(), + &json!([ + { + "type": "unsupported", + "feature": "temperature", + "details": format!("temperature is not supported by {sampling_model_id} and will be ignored"), + }, + { + "type": "unsupported", + "feature": "topK", + "details": format!("topK is not supported by {sampling_model_id} and will be ignored"), + }, + { + "type": "unsupported", + "feature": "topP", + "details": format!("topP is not supported by {sampling_model_id} and will be ignored"), + }, + ]), + ); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should support guardrails + #[tokio::test] + async fn should_support_guardrails() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "prompt": test_prompt(), + "providerOptions": { + "bedrock": { + "guardrailConfig": { "guardrailIdentifier": "-1", "guardrailVersion": "1", "trace": "enabled" }, + }, + }, + }), + ) + .await; + + assert_match_object( + &g.request_body(), + &json!({ + "guardrailConfig": { "guardrailIdentifier": "-1", "guardrailVersion": "1", "trace": "enabled" }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should send text parts as guardContent blocks when providerOptions.bedrock.guardContent is set + #[tokio::test] + async fn should_send_text_parts_as_guard_content_blocks() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "prompt": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "London is the capital of UK. Tokyo is the capital of Japan.", + "providerOptions": { + "bedrock": { "guardContent": true, "guardContentQualifiers": ["grounding_source"] }, + }, + }, + { "type": "text", "text": "Some additional background information." }, + { + "type": "text", + "text": "What is the capital of Japan?", + "providerOptions": { + "bedrock": { "guardContent": true, "guardContentQualifiers": ["query"] }, + }, + }, + ], + }, + ], + "providerOptions": { + "bedrock": { + "guardrailConfig": { "guardrailIdentifier": "-1", "guardrailVersion": "1", "trace": "enabled" }, + }, + }, + }), + ) + .await; + + assert_match_object( + &g.request_body(), + &json!({ + "messages": [ + { + "role": "user", + "content": [ + { + "guardContent": { + "text": { + "text": "London is the capital of UK. Tokyo is the capital of Japan.", + "qualifiers": ["grounding_source"], + }, + }, + }, + { "text": "Some additional background information." }, + { + "guardContent": { + "text": { "text": "What is the capital of Japan?", "qualifiers": ["query"] }, + }, + }, + ], + }, + ], + "guardrailConfig": { "guardrailIdentifier": "-1", "guardrailVersion": "1", "trace": "enabled" }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should send image parts as guardContent blocks when providerOptions.bedrock.guardContent is set + #[tokio::test] + async fn should_send_image_parts_as_guard_content_blocks() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "prompt": [ + { + "role": "user", + "content": [ + { + "type": "file", + "data": { "type": "data", "data": "AAECAw==" }, + "mediaType": "image/png", + "providerOptions": { "bedrock": { "guardContent": true } }, + }, + { "type": "text", "text": "Describe this image." }, + ], + }, + ], + }), + ) + .await; + + assert_match_object( + &g.request_body(), + &json!({ + "messages": [ + { + "role": "user", + "content": [ + { "guardContent": { "image": { "format": "png", "source": { "bytes": "AAECAw==" } } } }, + { "text": "Describe this image." }, + ], + }, + ], + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should surface guardrail intervention as a content-filter finish reason with trace metadata + #[tokio::test] + async fn should_surface_guardrail_intervention_as_a_content_filter_finish_reason() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-guard-content-intervened"), + json!({ + "prompt": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Tell me something restricted.", + "providerOptions": { + "bedrock": { "guardContent": true, "guardContentQualifiers": ["query"] }, + }, + }, + ], + }, + ], + "providerOptions": { + "bedrock": { + "guardrailConfig": { + "guardrailIdentifier": "1abcd2ef34gh", + "guardrailVersion": "1", + "trace": "enabled", + }, + }, + }, + }), + ) + .await; + + assert_json_eq( + &g.finish_reason(), + &json!({ "raw": "guardrail_intervened", "unified": "content-filter" }), + ); + assert_json_eq( + &g.content(), + &json!([{ "text": "Sorry, the model cannot answer this question.", "type": "text" }]), + ); + assert_json_eq( + &g.bedrock_metadata("/trace").expect("trace"), + &json!({ + "guardrail": { + "inputAssessment": { + "1abcd2ef34gh": { + "topicPolicy": { + "topics": [{ "action": "BLOCKED", "name": "Restricted topic", "type": "DENY" }], + }, + }, + }, + }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should include trace information in providerMetadata + #[tokio::test] + async fn should_include_trace_information_in_provider_metadata() { + let g = generate( + MODEL_ID, + hello_world_response(json!({ "trace": mock_trace() })), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_match_object(&g.bedrock_metadata("/trace").expect("trace"), &mock_trace()); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should include stop_sequence in provider metadata + #[tokio::test] + async fn should_include_stop_sequence_in_provider_metadata() { + let g = generate( + MODEL_ID, + json_response(&json!({ + "output": { "message": { "role": "assistant", "content": [{ "text": "Hello, World!" }] } }, + "stopReason": "stop_sequence", + "additionalModelResponseFields": { "delta": { "stop_sequence": "STOP" } }, + "usage": { "inputTokens": 4, "outputTokens": 30, "totalTokens": 34 }, + })), + json!({ "prompt": test_prompt(), "stopSequences": ["STOP"] }), + ) + .await; + + assert_json_eq( + &to_json(&g.result.provider_metadata), + &json!({ + "amazonBedrock": { "stopSequence": "STOP" }, + "bedrock": { "stopSequence": "STOP" }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should handle stop_sequence: null when stopReason is tool_use + #[tokio::test] + async fn should_handle_stop_sequence_null_when_stop_reason_is_tool_use() { + let g = generate( + MODEL_ID, + json_response(&json!({ + "output": { + "message": { + "role": "assistant", + "content": [ + { "text": "I'll query your Salesforce..." }, + { + "toolUse": { + "toolUseId": "tool-use-id", + "name": "querySalesforce", + "input": { "query": "SELECT Id FROM Account" }, + }, + }, + ], + }, + }, + "stopReason": "tool_use", + "additionalModelResponseFields": { "delta": { "stop_sequence": null } }, + "usage": { "inputTokens": 10, "outputTokens": 20, "totalTokens": 30 }, + })), + json!({ + "prompt": test_prompt(), + "tools": [ + { + "type": "function", + "name": "querySalesforce", + "description": "Query Salesforce", + "inputSchema": { + "type": "object", + "properties": { "query": { "type": "string" } }, + "required": ["query"], + "additionalProperties": false, + }, + }, + ], + }), + ) + .await; + + assert_json_eq( + &g.finish_reason(), + &json!({ "unified": "tool-calls", "raw": "tool_use" }), + ); + assert_json_eq( + &g.content(), + &json!([ + { "text": "I'll query your Salesforce...", "type": "text" }, + { + "input": "{\"query\":\"SELECT Id FROM Account\"}", + "toolCallId": "tool-use-id", + "toolName": "querySalesforce", + "type": "tool-call", + }, + ]), + ); + assert_json_eq( + &to_json(&g.result.provider_metadata), + &json!({ + "amazonBedrock": { "stopSequence": null }, + "bedrock": { "stopSequence": null }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should correctly parse delta.stop_sequence structure from additionalModelResponseFields + #[tokio::test] + async fn should_correctly_parse_delta_stop_sequence_structure() { + let g = generate( + MODEL_ID, + json_response(&json!({ + "output": { "message": { "role": "assistant", "content": [{ "text": "Response text" }] } }, + "stopReason": "stop_sequence", + "additionalModelResponseFields": { + "delta": { "stop_sequence": "CUSTOM_STOP" }, + "otherField": "value", + }, + "usage": { "inputTokens": 10, "outputTokens": 20, "totalTokens": 30 }, + })), + json!({ "prompt": test_prompt(), "stopSequences": ["CUSTOM_STOP"] }), + ) + .await; + + assert_eq!( + g.bedrock_metadata("/stopSequence"), + Some(json!("CUSTOM_STOP")) + ); + assert_json_eq( + &g.finish_reason(), + &json!({ "unified": "stop", "raw": "stop_sequence" }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should include response headers in rawResponse + #[tokio::test] + async fn should_include_response_headers_in_raw_response() { + let g = generate( + MODEL_ID, + json_response_with_headers( + &json!({ + "output": { "message": { "role": "assistant", "content": [{ "text": "Testing" }] } }, + "usage": { "inputTokens": 4, "outputTokens": 34, "totalTokens": 38 }, + "stopReason": "stop_sequence", + }), + &[ + ("x-amzn-requestid", "test-request-id"), + ("x-amzn-trace-id", "test-trace-id"), + ], + ), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &to_json(&g.result.response.as_ref().unwrap().headers), + &json!({ + "x-amzn-requestid": "test-request-id", + "x-amzn-trace-id": "test-trace-id", + "content-type": "application/json", + "content-length": "164", + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should extract response metadata + #[tokio::test] + async fn should_extract_response_metadata() { + let g = generate( + MODEL_ID, + json_response_with_headers( + &json!({ + "output": { "message": { "role": "assistant", "content": [{ "text": "Testing" }] } }, + "usage": { "inputTokens": 4, "outputTokens": 34, "totalTokens": 38 }, + "stopReason": "end_turn", + }), + &[ + ("x-amzn-requestid", "test-request-id"), + ("date", "Wed, 01 Jan 2025 00:00:00 GMT"), + ], + ), + json!({ "prompt": test_prompt() }), + ) + .await; + + let response = g.result.response.as_ref().unwrap(); + assert_eq!(response.id.as_deref(), Some("test-request-id")); + assert_eq!( + response.timestamp.map(|timestamp| timestamp.to_string()), + Some("2025-01-01T00:00:00.000Z".to_string()) + ); + assert_eq!(response.model_id.as_deref(), Some(MODEL_ID)); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should pass tools and tool choice correctly + #[tokio::test] + async fn should_pass_tools_and_tool_choice_correctly() { + let schema = json!({ + "type": "object", + "properties": { "param1": { "type": "string" }, "param2": { "type": "number" } }, + "required": ["param1"], + "additionalProperties": false, + }); + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "tools": [ + { "type": "function", "name": "test-tool-1", "description": "A test tool", "inputSchema": schema }, + ], + "toolChoice": { "type": "auto" }, + "prompt": test_prompt(), + }), + ) + .await; + + assert_match_object( + &g.request_body(), + &json!({ + "toolConfig": { + "tools": [ + { + "toolSpec": { + "name": "test-tool-1", + "description": "A test tool", + "inputSchema": { "json": schema }, + }, + }, + ], + }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should disable parallel tool use without sending conflicting tool choices + #[tokio::test] + async fn should_disable_parallel_tool_use_without_sending_conflicting_tool_choices() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "tools": [ + { + "type": "function", + "name": "test-tool", + "description": "A test tool", + "inputSchema": { "type": "object", "properties": {}, "additionalProperties": false }, + }, + ], + "toolChoice": { "type": "auto" }, + "prompt": test_prompt(), + "providerOptions": { "anthropic": { "disableParallelToolUse": true } }, + }), + ) + .await; + + let body = g.request_body(); + assert_match_object( + &body["additionalModelRequestFields"], + &json!({ "tool_choice": { "type": "auto", "disable_parallel_tool_use": true } }), + ); + assert_eq!(at(&body, "/toolConfig/toolChoice"), None); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should only send the forced tool when toolChoice specifies a specific tool + #[tokio::test] + async fn should_only_send_the_forced_tool_when_tool_choice_specifies_a_specific_tool() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "tools": [ + { + "type": "function", + "name": "getWeatherByCity", + "description": "Get weather by city", + "inputSchema": { + "type": "object", + "properties": { "city": { "type": "string" } }, + "required": ["city"], + "additionalProperties": false, + }, + }, + { + "type": "function", + "name": "getCurrentWeather", + "description": "Get current weather", + "inputSchema": { + "type": "object", + "properties": { "location": { "type": "string" } }, + "required": ["location"], + "additionalProperties": false, + }, + }, + { + "type": "function", + "name": "getWeatherForecast", + "description": "Get forecast", + "inputSchema": { + "type": "object", + "properties": { "location": { "type": "string" }, "days": { "type": "number" } }, + "required": ["location", "days"], + "additionalProperties": false, + }, + }, + ], + "toolChoice": { "type": "tool", "toolName": "getWeatherByCity" }, + "prompt": test_prompt(), + }), + ) + .await; + + let body = g.request_body(); + assert_eq!(tool_names(&body), vec!["getWeatherByCity"]); + assert_json_eq( + &body["toolConfig"]["toolChoice"], + &json!({ "tool": { "name": "getWeatherByCity" } }), + ); + } + + mod models_that_reject_forced_tool_use { + use super::*; + + fn weather_tool() -> Value { + json!({ "type": "function", "name": "getWeather", "description": "Get weather", "inputSchema": { "type": "object" } }) + } + + fn time_tool() -> Value { + json!({ "type": "function", "name": "getTime", "description": "Get time", "inputSchema": { "type": "object" } }) + } + + /// The response of upstream `createModelThatCapturesRequest`'s fetch. + fn done_response() -> Canned { + json_response(&json!({ + "output": { "message": { "role": "assistant", "content": [{ "text": "Done" }] } }, + "stopReason": "end_turn", + "usage": { "inputTokens": 1, "outputTokens": 1, "totalTokens": 2 }, + })) + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › models that reject forced tool use › should build an auto tool choice for required choice on %s + #[tokio::test] + async fn should_build_an_auto_tool_choice_for_required_choice_on_s() { + for model_id in [ + "anthropic.claude-opus-5-5", + "us.anthropic.claude-opus-5-5", + "global.anthropic.claude-opus-5-5", + "anthropic.claude-sonnet-5-5", + "us.anthropic.claude-sonnet-5-5", + ] { + let g = generate( + model_id, + done_response(), + json!({ + "tools": [weather_tool()], + "toolChoice": { "type": "required" }, + "prompt": test_prompt(), + }), + ) + .await; + + assert_json_eq( + &g.request_body()["toolConfig"]["toolChoice"], + &json!({ "auto": {} }), + ); + assert_json_eq( + &g.warnings(), + &json!([ + { + "type": "unsupported", + "feature": "toolChoice", + "details": "toolChoice 'required' is not supported by this model because it rejects forced tool use. Using 'auto' instead. Instruct the model to use a tool in the prompt and verify that a tool call was made.", + }, + ]), + ); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › models that reject forced tool use › should build an auto choice containing only the named tool + #[tokio::test] + async fn should_build_an_auto_choice_containing_only_the_named_tool() { + let g = generate( + OPUS55_ANTHROPIC_MODEL_ID, + done_response(), + json!({ + "tools": [weather_tool(), time_tool()], + "toolChoice": { "type": "tool", "toolName": "getWeather" }, + "prompt": test_prompt(), + }), + ) + .await; + + let body = g.request_body(); + assert_json_eq(&body["toolConfig"]["toolChoice"], &json!({ "auto": {} })); + assert_eq!(tool_names(&body), vec!["getWeather"]); + assert_json_eq( + &g.warnings(), + &json!([ + { + "type": "unsupported", + "feature": "toolChoice", + "details": "toolChoice 'tool' is not supported by this model because it rejects forced tool use. Only the 'getWeather' tool is sent with 'auto' tool choice. Instruct the model to use the tool in the prompt and verify that a tool call was made.", + }, + ]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › models that reject forced tool use › should build an Anthropic auto choice when parallel tool use is disabled + #[tokio::test] + async fn should_build_an_anthropic_auto_choice_when_parallel_tool_use_is_disabled() { + let g = generate( + OPUS55_ANTHROPIC_MODEL_ID, + done_response(), + json!({ + "tools": [weather_tool()], + "toolChoice": { "type": "required" }, + "prompt": test_prompt(), + "providerOptions": { "anthropic": { "disableParallelToolUse": true } }, + }), + ) + .await; + + let body = g.request_body(); + assert_json_eq( + at(&body, "/additionalModelRequestFields/tool_choice").expect("tool_choice"), + &json!({ "type": "auto", "disable_parallel_tool_use": true }), + ); + assert_eq!(at(&body, "/toolConfig/toolChoice"), None); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › models that reject forced tool use › should keep building forced tool choices for models that support them + #[tokio::test] + async fn should_keep_building_forced_tool_choices_for_models_that_support_them() { + let g = generate( + OPUS5_ANTHROPIC_MODEL_ID, + done_response(), + json!({ + "tools": [weather_tool()], + "toolChoice": { "type": "required" }, + "prompt": test_prompt(), + }), + ) + .await; + + assert_json_eq( + &g.request_body()["toolConfig"]["toolChoice"], + &json!({ "any": {} }), + ); + assert_json_eq(&g.warnings(), &json!([])); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should send all tools when toolChoice is auto + #[tokio::test] + async fn should_send_all_tools_when_tool_choice_is_auto() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "tools": [ + { + "type": "function", + "name": "getWeatherByCity", + "description": "Get weather by city", + "inputSchema": { + "type": "object", + "properties": { "city": { "type": "string" } }, + "required": ["city"], + "additionalProperties": false, + }, + }, + { + "type": "function", + "name": "getCurrentWeather", + "description": "Get current weather", + "inputSchema": { + "type": "object", + "properties": { "location": { "type": "string" } }, + "required": ["location"], + "additionalProperties": false, + }, + }, + ], + "toolChoice": { "type": "auto" }, + "prompt": test_prompt(), + }), + ) + .await; + + assert_eq!( + tool_names(&g.request_body()), + vec!["getWeatherByCity", "getCurrentWeather"] + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should omit empty tool descriptions to avoid Bedrock validation errors + #[tokio::test] + async fn should_omit_empty_tool_descriptions_to_avoid_bedrock_validation_errors() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "tools": [ + { + "type": "function", + "name": "tool-with-empty-desc", + "description": "", + "inputSchema": { + "type": "object", + "properties": { "param1": { "type": "string" } }, + "required": ["param1"], + "additionalProperties": false, + }, + }, + { + "type": "function", + "name": "tool-with-whitespace-desc", + "description": " ", + "inputSchema": { + "type": "object", + "properties": { "param2": { "type": "number" } }, + "required": ["param2"], + "additionalProperties": false, + }, + }, + { + "type": "function", + "name": "tool-with-valid-desc", + "description": "Valid description", + "inputSchema": { + "type": "object", + "properties": { "param3": { "type": "boolean" } }, + "required": ["param3"], + "additionalProperties": false, + }, + }, + ], + "toolChoice": { "type": "auto" }, + "prompt": test_prompt(), + }), + ) + .await; + + let body = g.request_body(); + let tools = &body["toolConfig"]["tools"]; + assert_eq!(tools[0]["toolSpec"].get("description"), None); + assert_eq!(tools[0]["toolSpec"]["name"], "tool-with-empty-desc"); + assert_eq!(tools[1]["toolSpec"].get("description"), None); + assert_eq!(tools[1]["toolSpec"]["name"], "tool-with-whitespace-desc"); + assert_eq!(tools[2]["toolSpec"]["description"], "Valid description"); + assert_eq!(tools[2]["toolSpec"]["name"], "tool-with-valid-desc"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should handle Anthropic provider-defined tools + #[tokio::test] + async fn should_handle_anthropic_provider_defined_tools() { + // Upstream mocks Anthropic `prepareTools` to return what the real + // function returns for `bash_20241022` (tool choice `auto`, beta + // `computer-use-2024-10-22`); aimux runs the real one. + let g = generate( + ANTHROPIC_MODEL_ID, + json_response(&json!({ + "output": { + "message": { + "role": "assistant", + "content": [ + { "toolUse": { "toolUseId": "tool-use-id", "name": "bash", "input": { "command": "ls -l" } } }, + ], + }, + }, + "usage": { "inputTokens": 10, "outputTokens": 20, "totalTokens": 30 }, + "stopReason": "tool_use", + })), + json!({ + "prompt": test_prompt(), + "tools": [bash_tool()], + "toolChoice": { "type": "auto" }, + }), + ) + .await; + + let body = g.request_body(); + assert_json_eq( + &body["additionalModelRequestFields"], + &json!({ + "tool_choice": { "type": "auto" }, + "anthropic_beta": ["computer-use-2024-10-22"], + }), + ); + assert!(body.get("toolConfig").is_some()); + assert_eq!(tool_names(&body), vec!["bash"]); + assert_json_eq( + &body["toolConfig"]["tools"][0]["toolSpec"]["inputSchema"]["json"], + &json!({ + "$schema": "http://json-schema.org/draft-07/schema#", + "type": "object", + "properties": { "command": { "type": "string" }, "restart": { "type": "boolean" } }, + "required": ["command"], + "additionalProperties": false, + }), + ); + assert_json_eq(&g.warnings(), &json!([])); + assert_json_eq( + &g.content(), + &json!([ + { + "input": "{\"command\":\"ls -l\"}", + "toolCallId": "tool-use-id", + "toolName": "bash", + "type": "tool-call", + }, + ]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should include anthropic_beta in additionalModelRequestFields when using extended context + #[tokio::test] + async fn should_include_anthropic_beta_in_additional_model_request_fields() { + let g = generate( + ANTHROPIC_MODEL_ID, + test_response_stop(), + json!({ + "prompt": test_prompt(), + "providerOptions": { "bedrock": { "anthropicBeta": ["context-1m-2025-08-07"] } }, + }), + ) + .await; + + assert_json_eq( + &g.request_body()["additionalModelRequestFields"], + &json!({ "anthropic_beta": ["context-1m-2025-08-07"] }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should not include anthropic-beta in HTTP headers + #[tokio::test] + async fn should_not_include_anthropic_beta_in_http_headers() { + let g = generate( + ANTHROPIC_MODEL_ID, + test_response_stop(), + json!({ "prompt": test_prompt(), "tools": [bash_tool()] }), + ) + .await; + + assert_eq!(g.request_headers().get("anthropic-beta"), None); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should combine user-provided and tool-generated betas in body + #[tokio::test] + async fn should_combine_user_provided_and_tool_generated_betas_in_body() { + let g = generate( + ANTHROPIC_MODEL_ID, + json_response(&json!({ + "output": { + "message": { + "role": "assistant", + "content": [ + { "type": "tool_use", "id": "tool-use-id", "name": "bash", "input": { "command": "ls -l" } }, + ], + }, + }, + "usage": { "inputTokens": 10, "outputTokens": 20, "totalTokens": 30 }, + "stopReason": "tool_use", + })), + json!({ + "prompt": test_prompt(), + "tools": [bash_tool()], + "providerOptions": { "bedrock": { "anthropicBeta": ["context-1m-2025-08-07"] } }, + }), + ) + .await; + + assert_json_eq( + &g.request_body()["additionalModelRequestFields"]["anthropic_beta"], + &json!(["context-1m-2025-08-07", "computer-use-2024-10-22"]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should properly combine headers from all sources + #[tokio::test] + async fn should_properly_combine_headers_from_all_sources() { + let fetch = MockFetch::new(vec![fixture_response("amazon-bedrock-text")]); + let inject: FetchFunction = Arc::new(InjectHeaders { + headers: vec![ + ("options-header", "options-value"), + ("model-header", "model-value"), + ("shared-header", "options-shared"), + ("signed-header", "signed-value"), + ("authorization", "AWS4-HMAC-SHA256..."), + ], + inner: fetch.transport(), + }); + let model_headers: HeaderMapOpt = [ + ("model-header".to_string(), Some("model-value".to_string())), + ( + "shared-header".to_string(), + Some("model-shared".to_string()), + ), + ] + .into_iter() + .collect(); + let g = generate_with( + AmazonBedrockProviderSettings { + headers: Some(model_headers), + ..settings(inject) + }, + &fetch, + MODEL_ID, + json!({ + "prompt": test_prompt(), + "headers": { "options-header": "options-value", "shared-header": "options-shared" }, + }), + ) + .await; + + let headers = g.request_headers(); + assert_eq!(headers["options-header"], "options-value"); + assert_eq!(headers["model-header"], "model-value"); + assert_eq!(headers["signed-header"], "signed-value"); + assert_eq!(headers["authorization"], "AWS4-HMAC-SHA256..."); + assert_eq!(headers["shared-header"], "options-shared"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should work with partial headers + #[tokio::test] + async fn should_work_with_partial_headers() { + let fetch = MockFetch::new(vec![fixture_response("amazon-bedrock-text")]); + let inject: FetchFunction = Arc::new(InjectHeaders { + headers: vec![ + ("model-header", "model-value"), + ("signed-header", "signed-value"), + ("authorization", "AWS4-HMAC-SHA256..."), + ], + inner: fetch.transport(), + }); + let model_headers: HeaderMapOpt = + [("model-header".to_string(), Some("model-value".to_string()))] + .into_iter() + .collect(); + let g = generate_with( + AmazonBedrockProviderSettings { + headers: Some(model_headers), + ..settings(inject) + }, + &fetch, + MODEL_ID, + json!({ "prompt": test_prompt() }), + ) + .await; + + let headers = g.request_headers(); + assert_eq!(headers["model-header"], "model-value"); + assert_eq!(headers["signed-header"], "signed-value"); + assert_eq!(headers["authorization"], "AWS4-HMAC-SHA256..."); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should include providerOptions in the request for generate calls + #[tokio::test] + async fn should_include_provider_options_in_the_request_for_generate_calls() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ "prompt": test_prompt(), "providerOptions": { "bedrock": { "foo": "bar" } } }), + ) + .await; + + assert_match_object(&g.request_body(), &json!({ "foo": "bar" })); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should include cache token usage in providerMetadata + #[tokio::test] + async fn should_include_cache_token_usage_in_provider_metadata() { + let g = generate( + MODEL_ID, + json_response(&json!({ + "output": { "message": { "role": "assistant", "content": [{ "type": "text", "text": "Testing" }] } }, + "usage": { + "inputTokens": 4, + "outputTokens": 34, + "totalTokens": 38, + "cacheReadInputTokens": 2, + "cacheWriteInputTokens": 3, + }, + "stopReason": "stop_sequence", + })), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &to_json(&g.result.provider_metadata), + &json!({ + "amazonBedrock": { "stopSequence": null, "usage": { "cacheWriteInputTokens": 3 } }, + "bedrock": { "stopSequence": null, "usage": { "cacheWriteInputTokens": 3 } }, + }), + ); + assert_json_eq( + &to_json(&g.result.usage), + &json!({ + "inputTokens": { "cacheRead": 2, "cacheWrite": 3, "noCache": 4, "total": 9 }, + "outputTokens": { "text": 34, "total": 34 }, + "raw": { + "cacheReadInputTokens": 2, + "cacheWriteInputTokens": 3, + "inputTokens": 4, + "outputTokens": 34, + "totalTokens": 38, + }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should include performanceConfig in providerMetadata + #[tokio::test] + async fn should_include_performance_config_in_provider_metadata() { + let g = generate( + MODEL_ID, + json_response(&json!({ + "output": { "message": { "role": "assistant", "content": [{ "type": "text", "text": "Hello" }] } }, + "usage": { "inputTokens": 4, "outputTokens": 34, "totalTokens": 38 }, + "stopReason": "end_turn", + "performanceConfig": { "latency": "optimized" }, + })), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &to_json(&g.result.provider_metadata), + &json!({ + "amazonBedrock": { "performanceConfig": { "latency": "optimized" }, "stopSequence": null }, + "bedrock": { "performanceConfig": { "latency": "optimized" }, "stopSequence": null }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should include serviceTier in providerMetadata + #[tokio::test] + async fn should_include_service_tier_in_provider_metadata() { + let g = generate( + MODEL_ID, + json_response(&json!({ + "output": { "message": { "role": "assistant", "content": [{ "type": "text", "text": "Hello" }] } }, + "usage": { "inputTokens": 4, "outputTokens": 34, "totalTokens": 38 }, + "stopReason": "end_turn", + "serviceTier": { "type": "on-demand" }, + })), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &to_json(&g.result.provider_metadata), + &json!({ + "amazonBedrock": { "serviceTier": { "type": "on-demand" }, "stopSequence": null }, + "bedrock": { "serviceTier": { "type": "on-demand" }, "stopSequence": null }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should include cacheDetails in providerMetadata + #[tokio::test] + async fn should_include_cache_details_in_provider_metadata() { + let g = generate( + MODEL_ID, + json_response(&json!({ + "output": { "message": { "role": "assistant", "content": [{ "type": "text", "text": "Hello" }] } }, + "usage": { + "inputTokens": 4, + "outputTokens": 34, + "totalTokens": 38, + "cacheDetails": [{ "inputTokens": 100, "ttl": "T5M" }], + }, + "stopReason": "end_turn", + })), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &to_json(&g.result.provider_metadata), + &json!({ + "amazonBedrock": { + "stopSequence": null, + "usage": { "cacheDetails": [{ "inputTokens": 100, "ttl": "T5M" }] }, + }, + "bedrock": { + "stopSequence": null, + "usage": { "cacheDetails": [{ "inputTokens": 100, "ttl": "T5M" }] }, + }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should handle system messages with cache points + #[tokio::test] + async fn should_handle_system_messages_with_cache_points() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "prompt": [ + { + "role": "system", + "content": "System Prompt", + "providerOptions": { "bedrock": { "cachePoint": { "type": "default" } } }, + }, + { "role": "user", "content": [{ "type": "text", "text": "Hello" }] }, + ], + }), + ) + .await; + + assert_match_object( + &g.request_body(), + &json!({ + "system": [{ "text": "System Prompt" }, { "cachePoint": { "type": "default" } }], + "messages": [{ "role": "user", "content": [{ "text": "Hello" }] }], + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should transform reasoningConfig to thinking in additionalModelRequestFields + #[tokio::test] + async fn should_transform_reasoning_config_to_thinking_in_additional_model_request_fields() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "prompt": test_prompt(), + "maxOutputTokens": 100, + "providerOptions": { + "bedrock": { "reasoningConfig": { "type": "enabled", "budgetTokens": 2000 } }, + }, + }), + ) + .await; + + let body = g.request_body(); + assert_match_object( + &body, + &json!({ + "additionalModelRequestFields": { "thinking": { "type": "enabled", "budget_tokens": 2000 } }, + "inferenceConfig": { "maxTokens": 2100 }, + }), + ); + assert_eq!(body.get("reasoningConfig"), None); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › merges user additionalModelRequestFields with derived thinking (generate) + #[tokio::test] + async fn merges_user_additional_model_request_fields_with_derived_thinking_generate() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "prompt": test_prompt(), + "providerOptions": { + "bedrock": { + "reasoningConfig": { "type": "enabled", "budgetTokens": 1234 }, + "additionalModelRequestFields": { "foo": "bar", "custom": 42 }, + }, + }, + }), + ) + .await; + + let body = g.request_body(); + assert_eq!(body.get("reasoningConfig"), None); + assert_match_object( + &body["additionalModelRequestFields"], + &json!({ + "foo": "bar", + "custom": 42, + "thinking": { "type": "enabled", "budget_tokens": 1234 }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should forward display in adaptive reasoningConfig to thinking + #[tokio::test] + async fn should_forward_display_in_adaptive_reasoning_config_to_thinking() { + let g = generate( + NEWER_ANTHROPIC_MODEL_ID, + hello_response("stop_sequence"), + json!({ + "prompt": test_prompt(), + "providerOptions": { + "bedrock": { "reasoningConfig": { "type": "adaptive", "display": "summarized" } }, + }, + }), + ) + .await; + + assert_json_eq( + at(&g.request_body(), "/additionalModelRequestFields/thinking").expect("thinking"), + &json!({ "type": "adaptive", "display": "summarized" }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should pass serviceTier provider option in generate requests + #[tokio::test] + async fn should_pass_service_tier_provider_option_in_generate_requests() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ "prompt": test_prompt(), "providerOptions": { "bedrock": { "serviceTier": "priority" } } }), + ) + .await; + + let body = g.request_body(); + assert_match_object(&body, &json!({ "serviceTier": { "type": "priority" } })); + assert_eq!(at(&body, "/additionalModelRequestFields/serviceTier"), None); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › maps maxReasoningEffort for Nova without thinking (generate) + #[tokio::test] + async fn maps_max_reasoning_effort_for_nova_without_thinking_generate() { + let g = generate( + NOVA_MODEL_ID, + hello_response("stop_sequence"), + json!({ + "prompt": test_prompt(), + "providerOptions": { + "bedrock": { + "reasoningConfig": { "type": "enabled", "maxReasoningEffort": "medium", "budgetTokens": 2048 }, + }, + }, + }), + ) + .await; + + let body = g.request_body(); + assert_match_object( + &body, + &json!({ + "additionalModelRequestFields": { + "reasoningConfig": { "type": "enabled", "maxReasoningEffort": "medium" }, + }, + }), + ); + assert_eq!(at(&body, "/additionalModelRequestFields/thinking"), None); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › maps maxReasoningEffort to reasoning_effort for OpenAI gpt-oss models (generate) + #[tokio::test] + async fn maps_max_reasoning_effort_to_reasoning_effort_for_openai_gpt_oss_models() { + let g = generate( + OPENAI_MODEL_ID, + hello_response("stop_sequence"), + json!({ + "prompt": test_prompt(), + "providerOptions": { "bedrock": { "reasoningConfig": { "maxReasoningEffort": "medium" } } }, + }), + ) + .await; + + let body = g.request_body(); + assert_match_object( + &body, + &json!({ "additionalModelRequestFields": { "reasoning_effort": "medium" } }), + ); + assert_eq!( + at(&body, "/additionalModelRequestFields/reasoningConfig"), + None + ); + assert_eq!(at(&body, "/additionalModelRequestFields/thinking"), None); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › maps maxReasoningEffort to nested reasoning.effort for CRIS model %s (generate) + #[tokio::test] + async fn maps_max_reasoning_effort_to_nested_reasoning_effort_for_cris_model_s() { + for cris_model_id in [US_OPENAI_MODEL_ID, GLOBAL_OPENAI_MODEL_ID] { + let g = generate( + cris_model_id, + hello_response("stop_sequence"), + json!({ + "prompt": test_prompt(), + "providerOptions": { "bedrock": { "reasoningConfig": { "maxReasoningEffort": "medium" } } }, + }), + ) + .await; + + let body = g.request_body(); + assert_match_object( + &body, + &json!({ "additionalModelRequestFields": { "reasoning": { "effort": "medium" } } }), + ); + assert_eq!( + at(&body, "/additionalModelRequestFields/reasoning_effort"), + None + ); + assert_eq!( + at(&body, "/additionalModelRequestFields/reasoningConfig"), + None + ); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › does not classify custom model IDs containing openai. as OpenAI models + #[tokio::test] + async fn does_not_classify_custom_model_ids_containing_openai_as_openai_models() { + let g = generate( + CUSTOM_OPENAI_SUBSTRING_MODEL_ID, + hello_response("stop_sequence"), + json!({ + "prompt": test_prompt(), + "providerOptions": { "bedrock": { "reasoningConfig": { "maxReasoningEffort": "medium" } } }, + }), + ) + .await; + + let body = g.request_body(); + assert_match_object( + &body, + &json!({ + "additionalModelRequestFields": { "reasoningConfig": { "maxReasoningEffort": "medium" } }, + }), + ); + assert_eq!(at(&body, "/additionalModelRequestFields/reasoning"), None); + assert_eq!( + at(&body, "/additionalModelRequestFields/reasoning_effort"), + None + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › strips unsupported sampling settings for OpenAI model %s + #[tokio::test] + async fn strips_unsupported_sampling_settings_for_openai_model_s() { + for cris_model_id in [US_OPENAI_MODEL_ID, GLOBAL_OPENAI_MODEL_ID] { + let g = generate( + cris_model_id, + hello_response("end_turn"), + json!({ + "prompt": test_prompt(), + "temperature": 0, + "topP": 0.5, + "topK": 5, + "stopSequences": [], + "maxOutputTokens": 100, + }), + ) + .await; + + assert_json_eq( + &g.request_body()["inferenceConfig"], + &json!({ "maxTokens": 100, "topK": 5 }), + ); + assert_json_eq( + &g.warnings(), + &json!([ + { + "type": "unsupported", + "feature": "temperature", + "details": "temperature is not supported by this OpenAI model on the Converse API", + }, + { + "type": "unsupported", + "feature": "topP", + "details": "topP is not supported by this OpenAI model on the Converse API", + }, + { + "type": "unsupported", + "feature": "stopSequences", + "details": "stopSequences is not supported by this OpenAI model on the Converse API", + }, + ]), + ); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › does not warn about clamping an unsupported OpenAI temperature + #[tokio::test] + async fn does_not_warn_about_clamping_an_unsupported_openai_temperature() { + let g = generate( + US_OPENAI_MODEL_ID, + hello_response("end_turn"), + json!({ "prompt": test_prompt(), "temperature": 2 }), + ) + .await; + + assert_eq!(g.request_body().get("inferenceConfig"), None); + assert_json_eq( + &g.warnings(), + &json!([ + { + "type": "unsupported", + "feature": "temperature", + "details": "temperature is not supported by this OpenAI model on the Converse API", + }, + ]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › keeps supported sampling settings for OpenAI gpt-oss models + #[tokio::test] + async fn keeps_supported_sampling_settings_for_openai_gpt_oss_models() { + let g = generate( + OPENAI_MODEL_ID, + hello_response("end_turn"), + json!({ "prompt": test_prompt(), "temperature": 0.2, "topP": 0.5, "stopSequences": ["END"] }), + ) + .await; + + assert_json_eq( + &g.request_body()["inferenceConfig"], + &json!({ "temperature": 0.2, "topP": 0.5 }), + ); + assert_json_eq( + &g.warnings(), + &json!([ + { + "type": "unsupported", + "feature": "stopSequences", + "details": "stopSequences is not supported by this OpenAI model on the Converse API", + }, + ]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should pass maxReasoningEffort as output_config.effort for Anthropic models (generate) + #[tokio::test] + async fn should_pass_max_reasoning_effort_as_output_config_effort_for_anthropic_models() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "prompt": test_prompt(), + "providerOptions": { + "bedrock": { "reasoningConfig": { "type": "enabled", "maxReasoningEffort": "medium" } }, + }, + }), + ) + .await; + + let body = g.request_body(); + assert_eq!( + at(&body, "/additionalModelRequestFields/reasoningConfig"), + None + ); + assert_json_eq( + at(&body, "/additionalModelRequestFields/output_config").expect("output_config"), + &json!({ "effort": "medium" }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should use native output_config.format instead of json tool when thinking is enabled with structured output + #[tokio::test] + async fn should_use_native_output_config_format_when_thinking_is_enabled() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Generate a recipe" }] }], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { + "recipe": { + "type": "object", + "properties": { + "name": { "type": "string" }, + "ingredients": { "type": "array", "items": { "type": "string" } }, + }, + "required": ["name", "ingredients"], + }, + }, + "required": ["recipe"], + }, + }, + "providerOptions": { + "bedrock": { "reasoningConfig": { "type": "enabled", "budgetTokens": 2000 } }, + }, + }), + ) + .await; + + let body = g.request_body(); + assert_eq!(body.get("toolConfig"), None); + assert_json_eq( + at(&body, "/additionalModelRequestFields/output_config").expect("output_config"), + &json!({ + "format": { + "type": "json_schema", + "schema": { + "type": "object", + "additionalProperties": false, + "properties": { + "recipe": { + "type": "object", + "additionalProperties": false, + "properties": { + "name": { "type": "string" }, + "ingredients": { "type": "array", "items": { "type": "string" } }, + }, + "required": ["name", "ingredients"], + }, + }, + "required": ["recipe"], + }, + }, + }), + ); + assert!(at(&body, "/additionalModelRequestFields/thinking").is_some()); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should merge output_config.effort and output_config.format when thinking with maxReasoningEffort and structured output + #[tokio::test] + async fn should_merge_output_config_effort_and_output_config_format() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Generate a recipe" }] }], + "responseFormat": name_response_format(), + "providerOptions": { + "bedrock": { + "reasoningConfig": { "type": "enabled", "budgetTokens": 2000, "maxReasoningEffort": "medium" }, + }, + }, + }), + ) + .await; + + let body = g.request_body(); + assert_eq!(body.get("toolConfig"), None); + assert_json_eq( + at(&body, "/additionalModelRequestFields/output_config").expect("output_config"), + &json!({ + "effort": "medium", + "format": { + "type": "json_schema", + "schema": { + "type": "object", + "additionalProperties": false, + "properties": { "name": { "type": "string" } }, + "required": ["name"], + }, + }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should still use json tool fallback for structured output without thinking enabled + #[tokio::test] + async fn should_still_use_json_tool_fallback_for_structured_output_without_thinking() { + let g = generate( + MODEL_ID, + json_tool_name_response(), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Generate a name" }] }], + "responseFormat": name_response_format(), + }), + ) + .await; + + let body = g.request_body(); + assert!(body.get("toolConfig").is_some()); + assert_eq!(tool_names(&body), vec!["json"]); + assert_json_eq(&body["toolConfig"]["toolChoice"], &json!({ "any": {} })); + assert_eq!( + at(&body, "/additionalModelRequestFields/output_config/format"), + None + ); + assert_json_eq( + &g.content(), + &json!([{ "text": "{\"name\":\"Test\"}", "type": "text" }]), + ); + assert_eq!( + g.bedrock_metadata("/isJsonResponseFromTool"), + Some(json!(true)) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should use the json tool fallback for claude-opus-5 (Bedrock rejects output_config.format) + #[tokio::test] + async fn should_use_the_json_tool_fallback_for_claude_opus_5() { + let g = generate( + OPUS5_ANTHROPIC_MODEL_ID, + json_tool_name_response(), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Generate a name" }] }], + "responseFormat": name_response_format(), + }), + ) + .await; + + let body = g.request_body(); + assert_eq!(body["toolConfig"]["tools"][0]["toolSpec"]["name"], "json"); + assert_eq!( + at(&body, "/additionalModelRequestFields/output_config"), + None + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should use JSON instructions instead of forced tool use for claude-opus-5-5 with structuredOutputMode %s + #[tokio::test] + async fn should_use_json_instructions_instead_of_forced_tool_use_for_claude_opus_5_5() { + for structured_output_mode in [None, Some("jsonTool")] { + let mut options = json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Return an answer of OK." }] }], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { "answer": { "type": "string" } }, + "required": ["answer"], + }, + }, + }); + if let Some(mode) = structured_output_mode { + options["providerOptions"] = + json!({ "amazonBedrock": { "structuredOutputMode": mode } }); + } + let g = generate( + OPUS55_ANTHROPIC_MODEL_ID, + json_response(&json!({ + "output": { "message": { "role": "assistant", "content": [{ "text": "{\"answer\":\"OK\"}" }] } }, + "usage": { "inputTokens": 4, "outputTokens": 5, "totalTokens": 9 }, + "stopReason": "end_turn", + })), + options, + ) + .await; + + let body = g.request_body(); + assert_eq!(body.get("toolConfig"), None); + assert_eq!( + at(&body, "/additionalModelRequestFields/output_config"), + None + ); + assert_json_eq( + &body["system"], + &json!([ + { + "text": "JSON schema:\n{\"type\":\"object\",\"properties\":{\"answer\":{\"type\":\"string\"}},\"required\":[\"answer\"]}\nYou MUST answer with only a JSON object that matches the JSON schema above. Do not wrap it in markdown fences or include any other text.", + }, + ]), + ); + assert_json_eq( + &g.content(), + &json!([{ "type": "text", "text": "{\"answer\":\"OK\"}" }]), + ); + assert_eq!(g.bedrock_metadata("/isJsonResponseFromTool"), None); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should default to the json tool for $modelId + #[tokio::test] + async fn should_default_to_the_json_tool_for_model_id() { + for model_id in [NEWER_ANTHROPIC_MODEL_ID, HAIKU45_ANTHROPIC_MODEL_ID] { + let g = generate( + model_id, + json_tool_name_response(), + json!({ "prompt": test_prompt(), "responseFormat": name_response_format() }), + ) + .await; + + let body = g.request_body(); + assert_eq!(tool_names(&body), vec!["json"]); + assert_json_eq(&body["toolConfig"]["toolChoice"], &json!({ "any": {} })); + assert_eq!( + at(&body, "/additionalModelRequestFields/output_config/format"), + None + ); + assert_json_eq( + &g.content(), + &json!([{ "type": "text", "text": "{\"name\":\"Test\"}" }]), + ); + assert_eq!(g.finish_reason()["unified"], "stop"); + assert_eq!( + g.bedrock_metadata("/isJsonResponseFromTool"), + Some(json!(true)) + ); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should force the json tool wire format with $providerOptionsName provider options + #[tokio::test] + async fn should_force_the_json_tool_wire_format_with_provider_options_name_provider_options() { + for provider_options in [ + json!({ "amazonBedrock": { "structuredOutputMode": "jsonTool" } }), + json!({ "bedrock": { "structuredOutputMode": "jsonTool" } }), + json!({ "anthropic": { "structuredOutputMode": "jsonTool" } }), + ] { + let g = generate( + NEWER_ANTHROPIC_MODEL_ID, + json_tool_name_response(), + json!({ + "prompt": test_prompt(), + "providerOptions": provider_options, + "responseFormat": name_response_format(), + }), + ) + .await; + + let body = g.request_body(); + assert_eq!(tool_names(&body), vec!["json"]); + assert_json_eq(&body["toolConfig"]["toolChoice"], &json!({ "any": {} })); + assert_eq!(body.get("structuredOutputMode"), None); + assert_eq!( + at(&body, "/additionalModelRequestFields/output_config/format"), + None + ); + assert_json_eq( + &g.content(), + &json!([{ "type": "text", "text": "{\"name\":\"Test\"}" }]), + ); + assert_eq!(g.finish_reason()["unified"], "stop"); + assert_eq!( + g.bedrock_metadata("/isJsonResponseFromTool"), + Some(json!(true)) + ); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should remove a manually supplied output_config.format in jsonTool mode while preserving sibling fields + #[tokio::test] + async fn should_remove_a_manually_supplied_output_config_format_in_json_tool_mode() { + let g = generate( + NEWER_ANTHROPIC_MODEL_ID, + json_tool_name_response(), + json!({ + "prompt": test_prompt(), + "providerOptions": { + "amazonBedrock": { + "structuredOutputMode": "jsonTool", + "additionalModelRequestFields": { + "output_config": { + "effort": "medium", + "format": { "type": "manually-supplied-format" }, + }, + }, + }, + }, + "responseFormat": name_response_format(), + }), + ) + .await; + + assert_json_eq( + at( + &g.request_body(), + "/additionalModelRequestFields/output_config", + ) + .expect("output_config"), + &json!({ "effort": "medium" }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should force output_config.format for models that auto mode routes to the json tool + #[tokio::test] + async fn should_force_output_config_format_for_models_that_auto_mode_routes_to_json_tool() { + let g = generate( + OPUS5_ANTHROPIC_MODEL_ID, + json_response(&json!({ + "output": { "message": { "role": "assistant", "content": [{ "text": "{\"name\":\"Test\"}" }] } }, + "usage": { "inputTokens": 4, "outputTokens": 10, "totalTokens": 14 }, + "stopReason": "end_turn", + })), + json!({ + "prompt": test_prompt(), + "providerOptions": { "amazonBedrock": { "structuredOutputMode": "outputFormat" } }, + "responseFormat": name_response_format(), + }), + ) + .await; + + let body = g.request_body(); + assert_eq!(body.get("toolConfig"), None); + assert_json_eq( + at(&body, "/additionalModelRequestFields/output_config/format").expect("format"), + &json!({ + "type": "json_schema", + "schema": { + "type": "object", + "additionalProperties": false, + "properties": { "name": { "type": "string" } }, + "required": ["name"], + }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should prefer amazonBedrock structuredOutputMode over anthropic structuredOutputMode + #[tokio::test] + async fn should_prefer_amazon_bedrock_structured_output_mode_over_anthropic() { + let g = generate( + NEWER_ANTHROPIC_MODEL_ID, + json_tool_name_response(), + json!({ + "prompt": test_prompt(), + "providerOptions": { + "amazonBedrock": { "structuredOutputMode": "jsonTool" }, + "anthropic": { "structuredOutputMode": "outputFormat" }, + }, + "responseFormat": name_response_format(), + }), + ) + .await; + + let body = g.request_body(); + assert_json_eq(&body["toolConfig"]["toolChoice"], &json!({ "any": {} })); + assert_eq!( + at(&body, "/additionalModelRequestFields/output_config/format"), + None + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should use native output_config.format for models with structured output support even without thinking enabled + #[tokio::test] + async fn should_use_native_output_config_format_for_models_with_structured_output_support() { + let g = generate( + NATIVE_STRUCTURED_OUTPUT_ANTHROPIC_MODEL_ID, + json_response(&json!({ + "output": { "message": { "content": [{ "text": "{\"name\":\"Test\"}" }], "role": "assistant" } }, + "usage": { "inputTokens": 4, "outputTokens": 10, "totalTokens": 14 }, + "stopReason": "end_turn", + })), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Generate a name" }] }], + "responseFormat": name_response_format(), + }), + ) + .await; + + let body = g.request_body(); + assert_eq!(body.get("toolConfig"), None); + assert_json_eq( + at(&body, "/additionalModelRequestFields/output_config").expect("output_config"), + &json!({ + "format": { + "schema": { + "additionalProperties": false, + "properties": { "name": { "type": "string" } }, + "required": ["name"], + "type": "object", + }, + "type": "json_schema", + }, + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should sanitize unsupported JSON schema keywords for native structured output + #[tokio::test] + async fn should_sanitize_unsupported_json_schema_keywords_for_native_structured_output() { + let g = generate( + NATIVE_STRUCTURED_OUTPUT_ANTHROPIC_MODEL_ID, + json_response(&json!({ + "output": { + "message": { "content": [{ "text": "{\"labels\":[\"Spring\",\"Summer\",\"Autumn\"]}" }], "role": "assistant" }, + }, + "usage": { "inputTokens": 4, "outputTokens": 10, "totalTokens": 14 }, + "stopReason": "end_turn", + })), + json!({ + "prompt": test_prompt(), + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { + "labels": { "type": "array", "maxItems": 3, "items": { "type": "string" } }, + }, + "required": ["labels"], + "additionalProperties": false, + }, + }, + }), + ) + .await; + + assert_json_eq( + &g.request_body()["additionalModelRequestFields"]["output_config"]["format"]["schema"], + &json!({ + "additionalProperties": false, + "properties": { + "labels": { "description": "max items: 3.", "items": { "type": "string" }, "type": "array" }, + }, + "required": ["labels"], + "type": "object", + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should use JSON instructions instead of a response tool when structured output is combined with tools on %s + #[tokio::test] + async fn should_use_json_instructions_when_structured_output_is_combined_with_tools_on_s() { + for model_id in [OPUS_ANTHROPIC_MODEL_ID, SONNET5_ANTHROPIC_MODEL_ID] { + let g = generate( + model_id, + json_response(&json!({ + "output": { + "message": { "content": [{ "text": "```json\n{\"name\":\"Test\"}\n```." }], "role": "assistant" }, + }, + "usage": { "inputTokens": 4, "outputTokens": 10, "totalTokens": 14 }, + "stopReason": "end_turn", + })), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Look up and generate a name" }] }], + "responseFormat": name_response_format(), + "tools": [ + { "type": "function", "name": "lookupName", "inputSchema": { "type": "object", "properties": {} } }, + ], + }), + ) + .await; + + let body = g.request_body(); + assert_json_eq( + &body["toolConfig"], + &json!({ + "tools": [ + { + "toolSpec": { + "inputSchema": { "json": { "properties": {}, "type": "object" } }, + "name": "lookupName", + }, + }, + ], + }), + ); + assert_eq!( + at(&body, "/additionalModelRequestFields/output_config"), + None + ); + assert_json_eq( + &body["system"], + &json!([ + { + "text": "JSON schema:\n{\"type\":\"object\",\"properties\":{\"name\":{\"type\":\"string\"}},\"required\":[\"name\"]}\nYou MUST answer with only a JSON object that matches the JSON schema above. Do not wrap it in markdown fences or include any other text.", + }, + ]), + ); + assert_json_eq( + &g.content(), + &json!([{ "text": "{\"name\":\"Test\"}", "type": "text" }]), + ); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should extract reasoning text with signature + #[tokio::test] + async fn should_extract_reasoning_text_with_signature() { + let g = generate( + MODEL_ID, + content_response(json!([ + { + "reasoningContent": { + "reasoningText": { + "text": "I need to think about this problem carefully...", + "signature": "abc123signature", + }, + }, + }, + { "type": "text", "text": "The answer is 42." }, + ])), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &g.content(), + &json!([ + { + "providerMetadata": { + "amazonBedrock": { "signature": "abc123signature" }, + "bedrock": { "signature": "abc123signature" }, + }, + "text": "I need to think about this problem carefully...", + "type": "reasoning", + }, + { "text": "The answer is 42.", "type": "text" }, + ]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should preserve empty text blocks between reasoning blocks + #[tokio::test] + async fn should_preserve_empty_text_blocks_between_reasoning_blocks() { + let g = generate( + MODEL_ID, + content_response(json!([ + { "reasoningContent": { "reasoningText": { "text": "thinking...", "signature": "sig-1" } } }, + { "text": "" }, + { "reasoningContent": { "reasoningText": { "text": "more thinking...", "signature": "sig-2" } } }, + { "text": "The answer is 42." }, + ])), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &g.content(), + &json!([ + { + "providerMetadata": { + "amazonBedrock": { "signature": "sig-1" }, + "bedrock": { "signature": "sig-1" }, + }, + "text": "thinking...", + "type": "reasoning", + }, + { "text": "", "type": "text" }, + { + "providerMetadata": { + "amazonBedrock": { "signature": "sig-2" }, + "bedrock": { "signature": "sig-2" }, + }, + "text": "more thinking...", + "type": "reasoning", + }, + { "text": "The answer is 42.", "type": "text" }, + ]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should extract reasoning text without signature + #[tokio::test] + async fn should_extract_reasoning_text_without_signature() { + let g = generate( + MODEL_ID, + content_response(json!([ + { "reasoningContent": { "reasoningText": { "text": "I need to think about this problem carefully..." } } }, + { "type": "text", "text": "The answer is 42." }, + ])), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &g.content(), + &json!([ + { "text": "I need to think about this problem carefully...", "type": "reasoning" }, + { "text": "The answer is 42.", "type": "text" }, + ]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should extract redacted reasoning + #[tokio::test] + async fn should_extract_redacted_reasoning() { + let g = generate( + MODEL_ID, + content_response(json!([ + { "reasoningContent": { "redactedReasoning": { "data": "redacted-reasoning-data" } } }, + { "type": "text", "text": "The answer is 42." }, + ])), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &g.content(), + &json!([ + { + "providerMetadata": { + "amazonBedrock": { "redactedData": "redacted-reasoning-data" }, + "bedrock": { "redactedData": "redacted-reasoning-data" }, + }, + "text": "", + "type": "reasoning", + }, + { "text": "The answer is 42.", "type": "text" }, + ]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should expose reasoning redacted as `redactedContent` for replay + #[tokio::test] + async fn should_expose_reasoning_redacted_as_redacted_content_for_replay() { + let g = generate( + MODEL_ID, + content_response(json!([ + { "reasoningContent": { "redactedContent": "encrypted-reasoning-payload" } }, + { "type": "text", "text": "The answer is 42." }, + ])), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &g.content(), + &json!([ + { + "providerMetadata": { + "amazonBedrock": { "redactedContent": "encrypted-reasoning-payload" }, + "bedrock": { "redactedContent": "encrypted-reasoning-payload" }, + }, + "text": "", + "type": "reasoning", + }, + { "text": "The answer is 42.", "type": "text" }, + ]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should handle multiple reasoning blocks + #[tokio::test] + async fn should_handle_multiple_reasoning_blocks() { + let g = generate( + MODEL_ID, + content_response(json!([ + { "reasoningContent": { "reasoningText": { "text": "First reasoning block", "signature": "sig1" } } }, + { "reasoningContent": { "redactedReasoning": { "data": "redacted-data" } } }, + { "type": "text", "text": "The answer is 42." }, + ])), + json!({ "prompt": test_prompt() }), + ) + .await; + + assert_json_eq( + &g.content(), + &json!([ + { + "providerMetadata": { + "amazonBedrock": { "signature": "sig1" }, + "bedrock": { "signature": "sig1" }, + }, + "text": "First reasoning block", + "type": "reasoning", + }, + { + "providerMetadata": { + "amazonBedrock": { "redactedData": "redacted-data" }, + "bedrock": { "redactedData": "redacted-data" }, + }, + "text": "", + "type": "reasoning", + }, + { "text": "The answer is 42.", "type": "text" }, + ]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should omit toolConfig and filter tool content when conversation has tool calls but no active tools + #[tokio::test] + async fn should_omit_tool_config_and_filter_tool_content_without_active_tools() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ "prompt": conversation_with_tool_calls(), "tools": [] }), + ) + .await; + + let body = g.request_body(); + assert_eq!(body.get("toolConfig"), None); + assert_json_eq( + &body["messages"], + &json!([ + { + "content": [ + { "text": "What is the weather in Toronto?" }, + { "text": "Now give me a summary." }, + ], + "role": "user", + }, + ]), + ); + assert_json_eq( + &g.warnings(), + &json!([ + { + "details": "Tool calls and results removed from conversation because Bedrock does not support tool content without active tools.", + "feature": "toolContent", + "type": "unsupported", + }, + ]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should handle JSON response format with schema + #[tokio::test] + async fn should_handle_json_response_format_with_schema() { + let g = generate( + MODEL_ID, + json_response(&json!({ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "json-tool-id", + "name": "json", + "input": { "recipe": { "name": "Lasagna", "ingredients": ["pasta", "cheese"] } }, + }, + }, + ], + }, + }, + "usage": { "inputTokens": 4, "outputTokens": 34, "totalTokens": 38 }, + "stopReason": "tool_use", + })), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Generate a recipe" }] }], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { + "recipe": { + "type": "object", + "properties": { + "name": { "type": "string" }, + "ingredients": { "type": "array", "items": { "type": "string" } }, + }, + "required": ["name", "ingredients"], + }, + }, + "required": ["recipe"], + }, + }, + }), + ) + .await; + + assert_json_eq( + &g.content(), + &json!([ + { + "text": "{\"recipe\":{\"name\":\"Lasagna\",\"ingredients\":[\"pasta\",\"cheese\"]}}", + "type": "text", + }, + ]), + ); + assert_eq!( + g.bedrock_metadata("/isJsonResponseFromTool"), + Some(json!(true)) + ); + assert_json_eq( + &g.finish_reason(), + &json!({ "raw": "tool_use", "unified": "stop" }), + ); + let body = g.request_body(); + assert_eq!(tool_names(&body), vec!["json"]); + assert_eq!( + body["toolConfig"]["tools"][0]["toolSpec"]["description"], + "Respond with a JSON object." + ); + assert_json_eq(&body["toolConfig"]["toolChoice"], &json!({ "any": {} })); + } + + mod json_schema_response_format_with_json_tool_response { + use super::*; + + /// The describe's `beforeEach`. + async fn generate_json_tool_response() -> Generated { + generate( + MODEL_ID, + fixture_response("amazon-bedrock-json-tool.1"), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Test" }] }], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { + "elements": { + "type": "array", + "items": { + "type": "object", + "properties": { + "location": { "type": "string" }, + "temperature": { "type": "number" }, + "condition": { "type": "string" }, + }, + "required": ["location", "temperature", "condition"], + }, + }, + }, + "required": ["elements"], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#", + }, + }, + }), + ) + .await + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › json schema response format with json tool response › should pass json schema response format as a tool + #[tokio::test] + async fn should_pass_json_schema_response_format_as_a_tool() { + let g = generate_json_tool_response().await; + + let body = g.request_body(); + assert!(body.get("toolConfig").is_some()); + assert_eq!(tool_names(&body), vec!["json"]); + assert_eq!( + body["toolConfig"]["tools"][0]["toolSpec"]["description"], + "Respond with a JSON object." + ); + assert_json_eq(&body["toolConfig"]["toolChoice"], &json!({ "any": {} })); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › json schema response format with json tool response › should return the json response as text + #[tokio::test] + async fn should_return_the_json_response_as_text() { + let g = generate_json_tool_response().await; + + assert_json_eq( + &g.content(), + &json!([ + { + "text": "{\"elements\":[{\"location\":\"San Francisco\",\"temperature\":-5,\"condition\":\"snowy\"},{\"location\":\"London\",\"temperature\":0,\"condition\":\"snowy\"}]}", + "type": "text", + }, + ]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › json schema response format with json tool response › should send stop finish reason when json tool is used + #[tokio::test] + async fn should_send_stop_finish_reason_when_json_tool_is_used() { + let g = generate_json_tool_response().await; + + assert_json_eq( + &g.finish_reason(), + &json!({ "raw": "tool_use", "unified": "stop" }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › json schema response format with json tool response › should set isJsonResponseFromTool in provider metadata + #[tokio::test] + async fn should_set_is_json_response_from_tool_in_provider_metadata() { + let g = generate_json_tool_response().await; + + assert_eq!( + g.bedrock_metadata("/isJsonResponseFromTool"), + Some(json!(true)) + ); + } + } + + mod json_schema_response_format_with_other_tool_response { + use super::*; + + /// The describe's `beforeEach`. + async fn generate_other_tool_response() -> Generated { + generate( + MODEL_ID, + fixture_response("amazon-bedrock-json-other-tool.1"), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Test" }] }], + "tools": [ + { + "type": "function", + "name": "get-weather", + "description": "Get the weather in a location", + "inputSchema": { + "type": "object", + "properties": { "location": { "type": "string" } }, + "required": ["location"], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#", + }, + }, + ], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { + "weather": { "type": "string" }, + "temperature": { "type": "number" }, + }, + "required": ["weather", "temperature"], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#", + }, + }, + }), + ) + .await + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › json schema response format with other tool response › should pass the tool and the json schema response format as tools + #[tokio::test] + async fn should_pass_the_tool_and_the_json_schema_response_format_as_tools() { + let g = generate_other_tool_response().await; + + let body = g.request_body(); + assert!(body.get("toolConfig").is_some()); + assert_eq!(tool_names(&body), vec!["get-weather", "json"]); + assert_json_eq(&body["toolConfig"]["toolChoice"], &json!({ "any": {} })); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › json schema response format with other tool response › should return the regular tool call + #[tokio::test] + async fn should_return_the_regular_tool_call() { + let g = generate_other_tool_response().await; + + assert_json_eq( + &g.content(), + &json!([ + { + "input": "{\"location\":\"San Francisco\"}", + "toolCallId": "toolu_01PQjhxo3eirCdKNvCJrKc8f", + "toolName": "get-weather", + "type": "tool-call", + }, + ]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › json schema response format with other tool response › should send tool-calls finish reason + #[tokio::test] + async fn should_send_tool_calls_finish_reason() { + let g = generate_other_tool_response().await; + + assert_json_eq( + &g.finish_reason(), + &json!({ "raw": "tool_use", "unified": "tool-calls" }), + ); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should omit toolConfig when conversation has tool calls but toolChoice is none + #[tokio::test] + async fn should_omit_tool_config_when_conversation_has_tool_calls_but_tool_choice_is_none() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "prompt": conversation_with_tool_calls(), + "tools": [ + { + "type": "function", + "name": "weather", + "inputSchema": { + "type": "object", + "properties": { "city": { "type": "string" } }, + "required": ["city"], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#", + }, + }, + ], + "toolChoice": { "type": "none" }, + }), + ) + .await; + + assert_eq!(g.request_body().get("toolConfig"), None); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should clamp temperature above 1 to 1 and add warning + #[tokio::test] + async fn should_clamp_temperature_above_1_to_1_and_add_warning() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ "prompt": test_prompt(), "temperature": 1.5 }), + ) + .await; + + assert_json_eq( + &g.request_body()["inferenceConfig"]["temperature"], + &json!(1), + ); + assert_json_eq( + &g.warnings(), + &json!([ + { + "details": "1.5 exceeds bedrock maximum of 1.0. clamped to 1.0", + "feature": "temperature", + "type": "unsupported", + }, + ]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should clamp temperature below 0 to 0 and add warning + #[tokio::test] + async fn should_clamp_temperature_below_0_to_0_and_add_warning() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ "prompt": test_prompt(), "temperature": -0.5 }), + ) + .await; + + assert_json_eq( + &g.request_body()["inferenceConfig"]["temperature"], + &json!(0), + ); + assert_json_eq( + &g.warnings(), + &json!([ + { + "details": "-0.5 is below bedrock minimum of 0. clamped to 0", + "feature": "temperature", + "type": "unsupported", + }, + ]), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should not clamp valid temperature between 0 and 1 + #[tokio::test] + async fn should_not_clamp_valid_temperature_between_0_and_1() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ "prompt": test_prompt(), "temperature": 0.7 }), + ) + .await; + + assert_json_eq( + &g.request_body()["inferenceConfig"]["temperature"], + &json!(0.7), + ); + assert_json_eq(&g.warnings(), &json!([])); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should include text content before JSON tool call in doGenerate + #[tokio::test] + async fn should_include_text_content_before_json_tool_call_in_do_generate() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-json-tool.2"), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Generate JSON" }] }], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { "value": { "type": "string" } }, + "required": ["value"], + }, + }, + }), + ) + .await; + + assert_json_eq( + &g.content(), + &json!([ + { "text": "Let me generate that JSON for you.", "type": "text" }, + { "text": "{\"value\":\"test data\"}", "type": "text" }, + ]), + ); + assert_eq!( + g.bedrock_metadata("/isJsonResponseFromTool"), + Some(json!(true)) + ); + assert_json_eq( + &g.finish_reason(), + &json!({ "raw": "tool_use", "unified": "stop" }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should preserve text response before JSON output (answering question then returning structured data) + #[tokio::test] + async fn should_preserve_text_response_before_json_output() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-json-tool-with-answer.1"), + json!({ + "prompt": [ + { + "role": "user", + "content": [{ "type": "text", "text": "What's 2+2? Also the name to use is John Doe" }], + }, + ], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { "firstName": { "type": "string" }, "lastName": { "type": "string" } }, + "required": ["firstName", "lastName"], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#", + }, + }, + }), + ) + .await; + + assert_json_eq( + &g.content(), + &json!([ + { "text": "2 + 2 equals 4.", "type": "text" }, + { "text": "{\"firstName\":\"John\",\"lastName\":\"Doe\"}", "type": "text" }, + ]), + ); + assert_json_eq( + &g.finish_reason(), + &json!({ "raw": "tool_use", "unified": "stop" }), + ); + assert_eq!( + g.bedrock_metadata("/isJsonResponseFromTool"), + Some(json!(true)) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should include multiple text blocks before JSON tool call in doGenerate + #[tokio::test] + async fn should_include_multiple_text_blocks_before_json_tool_call_in_do_generate() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-json-tool.3"), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Generate data" }] }], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { "result": { "type": "number" } }, + "required": ["result"], + }, + }, + }), + ) + .await; + + assert_json_eq( + &g.content(), + &json!([ + { "text": "First text block.", "type": "text" }, + { "text": "Second text block.", "type": "text" }, + { "text": "Third text block before JSON.", "type": "text" }, + { "text": "{\"result\":42}", "type": "text" }, + ]), + ); + assert_eq!( + g.bedrock_metadata("/isJsonResponseFromTool"), + Some(json!(true)) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should handle regular tool call before JSON tool call in doGenerate + #[tokio::test] + async fn should_handle_regular_tool_call_before_json_tool_call_in_do_generate() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-json-with-tool.1"), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Test" }] }], + "tools": [ + { + "type": "function", + "name": "get-weather", + "inputSchema": { + "type": "object", + "properties": { "location": { "type": "string" } }, + "required": ["location"], + }, + }, + ], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { "result": { "type": "string" } }, + "required": ["result"], + }, + }, + }), + ) + .await; + + assert_json_eq( + &g.content(), + &json!([ + { "text": "Let me call a tool first.", "type": "text" }, + { + "input": "{\"location\":\"SF\"}", + "toolCallId": "tool-1", + "toolName": "get-weather", + "type": "tool-call", + }, + { "text": "Now the JSON.", "type": "text" }, + { "text": "{\"result\":\"data\"}", "type": "text" }, + ]), + ); + assert_eq!( + g.bedrock_metadata("/isJsonResponseFromTool"), + Some(json!(true)) + ); + assert_json_eq( + &g.finish_reason(), + &json!({ "raw": "tool_use", "unified": "stop" }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should handle multiple regular tool calls before JSON tool call in doGenerate + #[tokio::test] + async fn should_handle_multiple_regular_tool_calls_before_json_tool_call_in_do_generate() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-json-with-tools.1"), + json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Test" }] }], + "tools": [ + { + "type": "function", + "name": "tool-a", + "inputSchema": { "type": "object", "properties": { "param": { "type": "string" } } }, + }, + { + "type": "function", + "name": "tool-b", + "inputSchema": { "type": "object", "properties": { "param": { "type": "string" } } }, + }, + ], + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { "final": { "type": "string" } }, + "required": ["final"], + }, + }, + }), + ) + .await; + + assert_json_eq( + &g.content(), + &json!([ + { "text": "Calling multiple tools.", "type": "text" }, + { "input": "{\"param\":\"a\"}", "toolCallId": "tool-1", "toolName": "tool-a", "type": "tool-call" }, + { "input": "{\"param\":\"b\"}", "toolCallId": "tool-2", "toolName": "tool-b", "type": "tool-call" }, + { "text": "Finally the JSON response.", "type": "text" }, + { "text": "{\"final\":\"result\"}", "type": "text" }, + ]), + ); + assert_eq!( + g.bedrock_metadata("/isJsonResponseFromTool"), + Some(json!(true)) + ); + assert_json_eq( + &g.finish_reason(), + &json!({ "raw": "tool_use", "unified": "stop" }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › should support tool calls with empty input (no arguments) + #[tokio::test] + async fn should_support_tool_calls_with_empty_input_no_arguments() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-tool-no-args"), + json!({ + "tools": [ + { + "type": "function", + "name": "updateIssueList", + "inputSchema": { + "type": "object", + "properties": {}, + "required": [], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#", + }, + }, + ], + "prompt": test_prompt(), + }), + ) + .await; + + assert_json_eq( + &g.content(), + &json!([ + { "text": "I'll update the issue list for you.", "type": "text" }, + { "input": "{}", "toolCallId": "tool-use-id", "toolName": "updateIssueList", "type": "tool-call" }, + ]), + ); + } + + mod legacy_anthropic_model_capabilities { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › legacy Anthropic model capabilities › should use budget-based reasoning for a platform-prefixed Claude 3.5 model + #[tokio::test] + async fn should_use_budget_based_reasoning_for_a_platform_prefixed_claude_3_5_model() { + let g = generate( + ANTHROPIC_MODEL_ID, + simple_response(), + json!({ "prompt": test_prompt(), "reasoning": "high" }), + ) + .await; + + let body = g.request_body(); + assert_json_eq( + at(&body, "/additionalModelRequestFields/thinking").expect("thinking"), + &json!({ "type": "enabled", "budget_tokens": (4096.0_f64 * 0.6).round() }), + ); + assert_eq!( + at(&body, "/additionalModelRequestFields/output_config"), + None + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › legacy Anthropic model capabilities › should use the JSON tool fallback for a platform-prefixed Claude 3.7 model + #[tokio::test] + async fn should_use_the_json_tool_fallback_for_a_platform_prefixed_claude_3_7_model() { + let g = generate( + LEGACY_ANTHROPIC_37_MODEL_ID, + simple_response(), + json!({ + "prompt": test_prompt(), + "responseFormat": { + "type": "json", + "schema": { + "type": "object", + "properties": { "name": { "type": "string" } }, + "required": ["name"], + "additionalProperties": false, + }, + }, + }), + ) + .await; + + let body = g.request_body(); + assert_eq!(tool_names(&body), vec!["json"]); + assert_eq!( + at(&body, "/additionalModelRequestFields/output_config"), + None + ); + } + } + + mod top_level_reasoning_parameter { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should not set reasoning config when reasoning is "provider-default" for newer Anthropic models + #[tokio::test] + async fn should_not_set_reasoning_config_when_reasoning_is_provider_default() { + let g = generate( + NEWER_ANTHROPIC_MODEL_ID, + simple_response(), + json!({ "prompt": test_prompt(), "reasoning": "provider-default" }), + ) + .await; + + let body = g.request_body(); + assert_eq!(at(&body, "/additionalModelRequestFields/thinking"), None); + assert_eq!( + at(&body, "/additionalModelRequestFields/output_config"), + None + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should map reasoning to adaptive thinking with effort for newer Anthropic models + #[tokio::test] + async fn should_map_reasoning_to_adaptive_thinking_with_effort_for_newer_anthropic_models() + { + let g = generate( + NEWER_ANTHROPIC_MODEL_ID, + simple_response(), + json!({ "prompt": test_prompt(), "reasoning": "high" }), + ) + .await; + + let body = g.request_body(); + assert_json_eq( + at(&body, "/additionalModelRequestFields/thinking").expect("thinking"), + &json!({ "type": "adaptive" }), + ); + assert_eq!( + at(&body, "/additionalModelRequestFields/output_config/effort"), + Some(&json!("high")) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should map reasoning "xhigh" to effort "max" for newer Anthropic models + #[tokio::test] + async fn should_map_reasoning_xhigh_to_effort_max_for_newer_anthropic_models() { + let g = generate( + NEWER_ANTHROPIC_MODEL_ID, + simple_response(), + json!({ "prompt": test_prompt(), "reasoning": "xhigh" }), + ) + .await; + + let body = g.request_body(); + assert_json_eq( + at(&body, "/additionalModelRequestFields/thinking").expect("thinking"), + &json!({ "type": "adaptive" }), + ); + assert_eq!( + at(&body, "/additionalModelRequestFields/output_config/effort"), + Some(&json!("max")) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should warn when reasoning "minimal" is mapped for newer Anthropic models + #[tokio::test] + async fn should_warn_when_reasoning_minimal_is_mapped_for_newer_anthropic_models() { + let g = generate( + NEWER_ANTHROPIC_MODEL_ID, + simple_response(), + json!({ "prompt": test_prompt(), "reasoning": "minimal" }), + ) + .await; + + assert_contains_equal( + &g.warnings(), + &json!({ + "type": "compatibility", + "feature": "reasoning", + "details": "reasoning \"minimal\" is not directly supported by this model. mapped to effort \"low\".", + }), + ); + assert_eq!( + at( + &g.request_body(), + "/additionalModelRequestFields/output_config/effort" + ), + Some(&json!("low")) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should map reasoning to budget-based thinking for older Anthropic models + #[tokio::test] + async fn should_map_reasoning_to_budget_based_thinking_for_older_anthropic_models() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ "prompt": test_prompt(), "reasoning": "high" }), + ) + .await; + + assert_json_eq( + at(&g.request_body(), "/additionalModelRequestFields/thinking").expect("thinking"), + &json!({ "type": "enabled", "budget_tokens": (4096.0_f64 * 0.6).round().max(1024.0) }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should map reasoning to budget with minimum of 1024 tokens for older Anthropic models + #[tokio::test] + async fn should_map_reasoning_to_budget_with_minimum_of_1024_tokens_for_older_anthropic() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ "prompt": test_prompt(), "reasoning": "minimal" }), + ) + .await; + + assert_json_eq( + at( + &g.request_body(), + "/additionalModelRequestFields/thinking/budget_tokens", + ) + .expect("budget_tokens"), + &json!(1024), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should map reasoning directly to reasoning_effort for OpenAI gpt-oss models + #[tokio::test] + async fn should_map_reasoning_directly_to_reasoning_effort_for_openai_gpt_oss_models() { + let g = generate( + OPENAI_MODEL_ID, + simple_response(), + json!({ "prompt": test_prompt(), "reasoning": "medium" }), + ) + .await; + + let body = g.request_body(); + assert_eq!( + at(&body, "/additionalModelRequestFields/reasoning_effort"), + Some(&json!("medium")) + ); + assert_eq!(at(&body, "/additionalModelRequestFields/thinking"), None); + assert_eq!( + at(&body, "/additionalModelRequestFields/reasoningConfig"), + None + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should map portable reasoning for Nova 2 + #[tokio::test] + async fn should_map_portable_reasoning_for_nova_2() { + let g = generate( + NOVA_MODEL_ID, + simple_response(), + json!({ "prompt": test_prompt(), "reasoning": "high" }), + ) + .await; + + assert_json_eq( + at( + &g.request_body(), + "/additionalModelRequestFields/reasoningConfig", + ) + .expect("reasoningConfig"), + &json!({ "type": "enabled", "maxReasoningEffort": "high" }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should omit maxOutputTokens for Nova 2 high reasoning + #[tokio::test] + async fn should_omit_max_output_tokens_for_nova_2_high_reasoning() { + let g = generate( + NOVA_MODEL_ID, + simple_response(), + json!({ "prompt": test_prompt(), "reasoning": "high", "maxOutputTokens": 1024 }), + ) + .await; + + let body = g.request_body(); + assert_eq!(at(&body, "/inferenceConfig/maxTokens"), None); + assert_json_eq( + at(&body, "/additionalModelRequestFields/reasoningConfig") + .expect("reasoningConfig"), + &json!({ "type": "enabled", "maxReasoningEffort": "high" }), + ); + assert_contains_equal( + &g.warnings(), + &json!({ + "type": "unsupported", + "feature": "maxOutputTokens", + "details": format!("maxOutputTokens is not supported by {NOVA_MODEL_ID} when high reasoning is enabled and will be ignored"), + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should preserve maxOutputTokens for Nova 2 medium reasoning + #[tokio::test] + async fn should_preserve_max_output_tokens_for_nova_2_medium_reasoning() { + let g = generate( + NOVA_MODEL_ID, + simple_response(), + json!({ "prompt": test_prompt(), "reasoning": "medium", "maxOutputTokens": 1024 }), + ) + .await; + + let body = g.request_body(); + assert_json_eq( + at(&body, "/inferenceConfig/maxTokens").expect("maxTokens"), + &json!(1024), + ); + assert_json_eq( + at(&body, "/additionalModelRequestFields/reasoningConfig") + .expect("reasoningConfig"), + &json!({ "type": "enabled", "maxReasoningEffort": "medium" }), + ); + assert!( + !has_warning_feature(&g, "maxOutputTokens"), + "unexpected maxOutputTokens warning: {:#}", + g.warnings() + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should omit maxOutputTokens for Nova 2 xhigh portable reasoning (max effort) + #[tokio::test] + async fn should_omit_max_output_tokens_for_nova_2_xhigh_portable_reasoning_max_effort() { + let g = generate( + NOVA_MODEL_ID, + simple_response(), + json!({ "prompt": test_prompt(), "reasoning": "xhigh", "maxOutputTokens": 1024 }), + ) + .await; + + let body = g.request_body(); + assert_eq!(at(&body, "/inferenceConfig/maxTokens"), None); + assert_json_eq( + at(&body, "/additionalModelRequestFields/reasoningConfig") + .expect("reasoningConfig"), + &json!({ "type": "enabled", "maxReasoningEffort": "max" }), + ); + assert!( + has_warning_feature(&g, "maxOutputTokens"), + "expected a maxOutputTokens warning: {:#}", + g.warnings() + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should omit maxOutputTokens for Nova 2 explicit max reasoning effort + #[tokio::test] + async fn should_omit_max_output_tokens_for_nova_2_explicit_max_reasoning_effort() { + let g = generate( + NOVA_MODEL_ID, + simple_response(), + json!({ + "prompt": test_prompt(), + "maxOutputTokens": 1024, + "providerOptions": { + "amazonBedrock": { "reasoningConfig": { "type": "enabled", "maxReasoningEffort": "max" } }, + }, + }), + ) + .await; + + let body = g.request_body(); + assert_eq!(at(&body, "/inferenceConfig/maxTokens"), None); + assert_json_eq( + at(&body, "/additionalModelRequestFields/reasoningConfig") + .expect("reasoningConfig"), + &json!({ "type": "enabled", "maxReasoningEffort": "max" }), + ); + assert!( + has_warning_feature(&g, "maxOutputTokens"), + "expected a maxOutputTokens warning: {:#}", + g.warnings() + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should map reasoning to reasoningConfig when explicitly enabled for Nova 2 + #[tokio::test] + async fn should_map_reasoning_to_reasoning_config_when_explicitly_enabled_for_nova_2() { + let g = generate( + NOVA_MODEL_ID, + simple_response(), + json!({ + "prompt": test_prompt(), + "reasoning": "high", + "providerOptions": { "amazonBedrock": { "reasoningConfig": { "type": "enabled" } } }, + }), + ) + .await; + + assert_json_eq( + at( + &g.request_body(), + "/additionalModelRequestFields/reasoningConfig", + ) + .expect("reasoningConfig"), + &json!({ "type": "enabled", "maxReasoningEffort": "high" }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should ignore portable reasoning for models without known reasoning support + #[tokio::test] + async fn should_ignore_portable_reasoning_for_models_without_known_reasoning_support() { + let g = generate( + NOVA_MICRO_MODEL_ID, + simple_response(), + json!({ "prompt": test_prompt(), "reasoning": "high" }), + ) + .await; + + assert_eq!( + at( + &g.request_body(), + "/additionalModelRequestFields/reasoningConfig" + ), + None + ); + assert_contains_equal( + &g.warnings(), + &json!({ + "type": "unsupported", + "feature": "reasoning", + "details": "Portable reasoning is not supported for this model and will be ignored. If the model supports a provider-specific reasoning configuration, use providerOptions.amazonBedrock.reasoningConfig.", + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should forward explicit reasoningConfig for models without known reasoning support + #[tokio::test] + async fn should_forward_explicit_reasoning_config_for_models_without_reasoning_support() { + let g = generate( + NOVA_MICRO_MODEL_ID, + simple_response(), + json!({ + "prompt": test_prompt(), + "providerOptions": { + "amazonBedrock": { "reasoningConfig": { "type": "enabled", "maxReasoningEffort": "high" } }, + }, + }), + ) + .await; + + assert_json_eq( + at( + &g.request_body(), + "/additionalModelRequestFields/reasoningConfig", + ) + .expect("reasoningConfig"), + &json!({ "type": "enabled", "maxReasoningEffort": "high" }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should let explicit reasoningConfig fields win over derived reasoning values + #[tokio::test] + async fn should_let_explicit_reasoning_config_fields_win_over_derived_reasoning_values() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "prompt": test_prompt(), + "reasoning": "high", + "providerOptions": { + "bedrock": { "reasoningConfig": { "type": "enabled", "budgetTokens": 5000 } }, + }, + }), + ) + .await; + + assert_json_eq( + at(&g.request_body(), "/additionalModelRequestFields/thinking").expect("thinking"), + &json!({ "type": "enabled", "budget_tokens": 5000 }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should merge top-level reasoning with partial reasoningConfig for newer Anthropic models + #[tokio::test] + async fn should_merge_top_level_reasoning_with_partial_reasoning_config_for_newer_models() { + let g = generate( + NEWER_ANTHROPIC_MODEL_ID, + simple_response(), + json!({ + "prompt": test_prompt(), + "reasoning": "high", + "providerOptions": { "bedrock": { "reasoningConfig": { "display": "summarized" } } }, + }), + ) + .await; + + let body = g.request_body(); + assert_json_eq( + at(&body, "/additionalModelRequestFields/thinking").expect("thinking"), + &json!({ "type": "adaptive", "display": "summarized" }), + ); + assert_eq!( + at(&body, "/additionalModelRequestFields/output_config/effort"), + Some(&json!("high")) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should honor reasoning "none" even when partial reasoningConfig is provided + #[tokio::test] + async fn should_honor_reasoning_none_even_when_partial_reasoning_config_is_provided() { + let g = generate( + NEWER_ANTHROPIC_MODEL_ID, + simple_response(), + json!({ + "prompt": test_prompt(), + "reasoning": "none", + "providerOptions": { "bedrock": { "reasoningConfig": { "display": "summarized" } } }, + }), + ) + .await; + + let body = g.request_body(); + assert_eq!(at(&body, "/additionalModelRequestFields/thinking"), None); + assert_eq!( + at(&body, "/additionalModelRequestFields/output_config"), + None + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should let user-specified type win while still deriving maxReasoningEffort from reasoning + #[tokio::test] + async fn should_let_user_specified_type_win_while_still_deriving_max_reasoning_effort() { + let g = generate( + NEWER_ANTHROPIC_MODEL_ID, + simple_response(), + json!({ + "prompt": test_prompt(), + "reasoning": "high", + "providerOptions": { + "bedrock": { "reasoningConfig": { "type": "enabled", "budgetTokens": 3000 } }, + }, + }), + ) + .await; + + let body = g.request_body(); + assert_json_eq( + at(&body, "/additionalModelRequestFields/thinking").expect("thinking"), + &json!({ "type": "enabled", "budget_tokens": 3000 }), + ); + assert_eq!( + at(&body, "/additionalModelRequestFields/output_config/effort"), + Some(&json!("high")) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should let user-specified maxReasoningEffort win over derived for non-Anthropic models + #[tokio::test] + async fn should_let_user_specified_max_reasoning_effort_win_for_non_anthropic_models() { + let g = generate( + NOVA_MODEL_ID, + simple_response(), + json!({ + "prompt": test_prompt(), + "reasoning": "high", + "providerOptions": { "bedrock": { "reasoningConfig": { "maxReasoningEffort": "low" } } }, + }), + ) + .await; + + assert_eq!( + at( + &g.request_body(), + "/additionalModelRequestFields/reasoningConfig/maxReasoningEffort" + ), + Some(&json!("low")) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should strip temperature, topP, topK for Anthropic models when reasoning enables thinking + #[tokio::test] + async fn should_strip_temperature_top_p_top_k_for_anthropic_models_when_thinking() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ + "prompt": test_prompt(), + "reasoning": "high", + "temperature": 0.7, + "topP": 0.9, + "topK": 5, + }), + ) + .await; + + let body = g.request_body(); + assert_eq!(at(&body, "/inferenceConfig/temperature"), None); + assert_eq!(at(&body, "/inferenceConfig/topP"), None); + assert_eq!(at(&body, "/inferenceConfig/topK"), None); + let warnings = g.warnings(); + assert_contains_equal( + &warnings, + &json!({ + "type": "unsupported", + "feature": "temperature", + "details": "temperature is not supported when thinking is enabled", + }), + ); + assert_contains_equal( + &warnings, + &json!({ + "type": "unsupported", + "feature": "topP", + "details": "topP is not supported when thinking is enabled", + }), + ); + assert_contains_equal( + &warnings, + &json!({ + "type": "unsupported", + "feature": "topK", + "details": "topK is not supported when thinking is enabled", + }), + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should handle reasoning "none" for Anthropic models + #[tokio::test] + async fn should_handle_reasoning_none_for_anthropic_models() { + let g = generate( + MODEL_ID, + fixture_response("amazon-bedrock-text"), + json!({ "prompt": test_prompt(), "reasoning": "none" }), + ) + .await; + + assert_eq!( + at(&g.request_body(), "/additionalModelRequestFields/thinking"), + None + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model.test.ts › doGenerate › top-level reasoning parameter › should handle reasoning "none" for OpenAI models + #[tokio::test] + async fn should_handle_reasoning_none_for_openai_models() { + let g = generate( + OPENAI_MODEL_ID, + simple_response(), + json!({ "prompt": test_prompt(), "reasoning": "none" }), + ) + .await; + + assert_eq!( + at( + &g.request_body(), + "/additionalModelRequestFields/reasoning_effort" + ), + None + ); + } + } +} diff --git a/aimux-providers/src/bedrock/tests/chat_language_model_event_stream.rs b/aimux-providers/src/bedrock/tests/chat_language_model_event_stream.rs new file mode 100644 index 000000000..5058c104f --- /dev/null +++ b/aimux-providers/src/bedrock/tests/chat_language_model_event_stream.rs @@ -0,0 +1,439 @@ +use futures::StreamExt; +use serde_json::{Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::language_model::LanguageModel; +use aimux_core::options::CallOptions; +use aimux_core::stream_part::StreamPart; +use aimux_provider_utils::Resolvable; + +use crate::bedrock::{AmazonBedrockProviderSettings, BedrockModel, create_amazon_bedrock}; +use crate::test_support::mock_fetch::{Canned, MockFetch}; + +/// `TEST_PROMPT`: one user text part. +fn test_options() -> CallOptions { + serde_json::from_value(json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Hello" }] }], + "includeRawChunks": false, + })) + .expect("valid call options") +} + +fn crc32(data: &[u8]) -> u32 { + let mut crc = 0xFFFF_FFFFu32; + for &byte in data { + crc ^= u32::from(byte); + for _ in 0..8 { + crc = if crc & 1 == 0 { + crc >> 1 + } else { + (crc >> 1) ^ 0xEDB8_8320 + }; + } + } + !crc +} + +/// `codec.encode({ headers, body })` of `@smithy/eventstream-codec`, with +/// string headers. +fn encode(headers: &[(&str, &str)], body: &[u8]) -> Vec { + let mut header_bytes = Vec::new(); + for (name, value) in headers { + header_bytes.push(u8::try_from(name.len()).expect("short header name")); + header_bytes.extend_from_slice(name.as_bytes()); + header_bytes.push(7); + header_bytes.extend_from_slice( + &u16::try_from(value.len()) + .expect("short value") + .to_be_bytes(), + ); + header_bytes.extend_from_slice(value.as_bytes()); + } + let total = u32::try_from(16 + header_bytes.len() + body.len()).expect("small frame"); + let headers_length = u32::try_from(header_bytes.len()).expect("small headers"); + let mut frame = Vec::new(); + frame.extend_from_slice(&total.to_be_bytes()); + frame.extend_from_slice(&headers_length.to_be_bytes()); + let prelude_crc = crc32(&frame); + frame.extend_from_slice(&prelude_crc.to_be_bytes()); + frame.extend_from_slice(&header_bytes); + frame.extend_from_slice(body); + let message_crc = crc32(&frame); + frame.extend_from_slice(&message_crc.to_be_bytes()); + frame +} + +/// `createEvent(eventType, data)`. +fn create_event(event_type: &str, data: &str) -> Vec { + encode( + &[(":message-type", "event"), (":event-type", event_type)], + data.as_bytes(), + ) +} + +/// A model whose every request is answered with `frames` as an +/// `application/vnd.amazon.eventstream` body. +fn model_streaming(model_id: &str, frames: &[Vec]) -> BedrockModel { + let mock = MockFetch::new(vec![Canned { + status: 200, + headers: vec![( + "content-type".into(), + "application/vnd.amazon.eventstream".into(), + )], + body: frames.concat(), + }]); + create_amazon_bedrock(AmazonBedrockProviderSettings { + api_key: Some(Resolvable::Value("test-api-key".into())), + region: Some("us-east-1".into()), + base_url: Some("https://bedrock-runtime.us-east-1.amazonaws.com".into()), + fetch: Some(mock.transport()), + ..Default::default() + }) + .expect("valid settings") + .chat(model_id) +} + +/// `convertReadableStreamToArray(stream)`, keeping what the stream threw. +async fn stream_items(model: &BedrockModel) -> Vec> { + model + .do_stream(&test_options()) + .await + .expect("stream starts") + .stream + .collect() + .await +} + +/// The parts of a stream that must not throw, as AI SDK JSON. +async fn stream_parts(model: &BedrockModel) -> Vec { + stream_items(model) + .await + .into_iter() + .map(|item| serde_json::to_value(item.expect("stream does not throw")).unwrap()) + .collect() +} + +/// `rejects.toThrow(message)`: the stream throws an error whose message +/// contains `message`. +fn assert_throws(items: &[Result], message: &str) { + let error = items + .iter() + .find_map(|item| item.as_ref().err()) + .unwrap_or_else(|| panic!("stream did not throw; expected {message:?}")); + assert!( + error.to_string().contains(message), + "expected the error to contain {message:?}, got {error}" + ); +} + +/// `expect.objectContaining(expected)`: every expected key is present and +/// equal. +fn object_containing(actual: &Value, expected: &Value) -> bool { + expected + .as_object() + .expect("expected object") + .iter() + .all(|(key, value)| actual.get(key) == Some(value)) +} + +/// `expect(parts).toEqual(expect.arrayContaining(expected))` of +/// `objectContaining` elements. +fn assert_contains_all(parts: &[Value], expected: &[Value]) { + for wanted in expected { + assert!( + parts.iter().any(|part| object_containing(part, wanted)), + "no part matches {wanted} in {parts:?}" + ); + } +} + +/// `expect(actual).toMatchObject(expected)`. +fn assert_match_object(actual: &Value, expected: &Value) { + match (actual, expected) { + (Value::Object(actual), Value::Object(expected)) => { + for (key, value) in expected { + let found = actual + .get(key) + .unwrap_or_else(|| panic!("missing {key}; expected {value}")); + assert_match_object(found, value); + } + } + (Value::Array(actual), Value::Array(expected)) => { + assert_eq!(actual.len(), expected.len()); + for (found, value) in actual.iter().zip(expected) { + assert_match_object(found, value); + } + } + _ => assert_eq!(actual, expected), + } +} + +mod amazon_bedrock_chat_language_model_do_stream { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model-event-stream.test.ts › AmazonBedrockChatLanguageModel doStream › surfaces event stream decoding failures + #[tokio::test] + async fn surfaces_event_stream_decoding_failures() { + let mut corrupted_frame = create_event("contentBlockDelta", "corrupted"); + let last = corrupted_frame.len() - 1; + corrupted_frame[last] ^= 0xff; + let model = model_streaming( + "anthropic.claude-3-haiku-20240307-v1:0", + &[ + corrupted_frame, + create_event( + "contentBlockDelta", + &json!({ "contentBlockIndex": 0, "delta": { "text": "later" } }).to_string(), + ), + ], + ); + + assert_throws(&stream_items(&model).await, "The message checksum"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model-event-stream.test.ts › AmazonBedrockChatLanguageModel doStream › rejects a truncated event stream frame at EOF + #[tokio::test] + async fn rejects_a_truncated_event_stream_frame_at_eof() { + let text_frame = create_event( + "contentBlockDelta", + &json!({ "contentBlockIndex": 0, "delta": { "text": "partial result" } }).to_string(), + ); + let message_stop_frame = create_event( + "messageStop", + &json!({ "stopReason": "end_turn" }).to_string(), + ); + let truncated = message_stop_frame[..message_stop_frame.len() - 1].to_vec(); + let model = model_streaming( + "anthropic.claude-3-haiku-20240307-v1:0", + &[text_frame, truncated], + ); + + assert_throws( + &stream_items(&model).await, + "Incomplete Amazon Bedrock event-stream frame", + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model-event-stream.test.ts › AmazonBedrockChatLanguageModel doStream › allows EOF at a complete frame boundary without messageStop + #[tokio::test] + async fn allows_eof_at_a_complete_frame_boundary_without_message_stop() { + let model = model_streaming( + "anthropic.claude-3-haiku-20240307-v1:0", + &[create_event( + "contentBlockDelta", + &json!({ "contentBlockIndex": 0, "delta": { "text": "complete frame" } }) + .to_string(), + )], + ); + + assert_contains_all( + &stream_parts(&model).await, + &[ + json!({ "type": "text-delta", "delta": "complete frame" }), + json!({ "type": "finish", "finishReason": { "unified": "other" } }), + ], + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model-event-stream.test.ts › AmazonBedrockChatLanguageModel doStream › surfaces $exceptionType frames as stream errors + #[tokio::test] + async fn surfaces_exception_type_frames_as_stream_errors() { + for (exception_type, status_code, is_retryable) in [ + ("internalServerException", 500, true), + ("modelStreamErrorException", 424, true), + ("serviceUnavailableException", 503, true), + ("throttlingException", 429, true), + ("validationException", 400, false), + ] { + let exception = json!({ "message": format!("Modeled exception: {exception_type}") }); + let exception_frame = encode( + &[ + (":message-type", "exception"), + (":exception-type", exception_type), + (":content-type", "application/json"), + ], + exception.to_string().as_bytes(), + ); + let model = + model_streaming("anthropic.claude-3-haiku-20240307-v1:0", &[exception_frame]); + + let parts: Vec = stream_items(&model) + .await + .into_iter() + .map(|item| item.expect("stream does not throw")) + .collect(); + let error_part = &parts[parts.len() - 2]; + + let StreamPart::Error { error } = error_part else { + panic!("Expected an error part, got {error_part:?}"); + }; + // `isProviderStreamError`: the provider's typed stream error. + let AiMuxError::ApiCall(error) = error else { + panic!("expected a provider stream error, got {error:?}"); + }; + assert_match_object( + &json!({ + "message": error.message, + "type": error.provider_code, + "statusCode": error.status_code, + "isRetryable": error.is_retryable, + "data": error.data, + }), + &json!({ + "message": exception["message"], + "type": exception_type, + "statusCode": status_code, + "isRetryable": is_retryable, + "data": { exception_type: exception }, + }), + ); + assert_match_object( + &serde_json::to_value(&parts[parts.len() - 1]).unwrap(), + &json!({ "type": "finish", "finishReason": { "unified": "error" } }), + ); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model-event-stream.test.ts › AmazonBedrockChatLanguageModel doStream › streams reasoning redacted as `redactedContent` for replay + #[tokio::test] + async fn streams_reasoning_redacted_as_redacted_content_for_replay() { + let model = model_streaming( + "us.openai.gpt-5.6-luna", + &[ + create_event( + "contentBlockDelta", + &json!({ + "contentBlockIndex": 0, + "delta": { "reasoningContent": { "redactedContent": "encrypted-reasoning-" } }, + }) + .to_string(), + ), + create_event( + "contentBlockDelta", + &json!({ + "contentBlockIndex": 0, + "delta": { "reasoningContent": { "redactedContent": "payload" } }, + }) + .to_string(), + ), + create_event( + "contentBlockStop", + &json!({ "contentBlockIndex": 0 }).to_string(), + ), + create_event( + "contentBlockDelta", + &json!({ "contentBlockIndex": 1, "delta": { "text": "The answer is 42." } }) + .to_string(), + ), + create_event( + "contentBlockStop", + &json!({ "contentBlockIndex": 1 }).to_string(), + ), + create_event("messageStop", &json!({ "stopReason": "end_turn" }).to_string()), + ], + ); + + let parts = stream_parts(&model).await; + + let errors: Vec<&Value> = parts + .iter() + .filter(|part| part["type"] == json!("error")) + .collect(); + assert_eq!(errors, Vec::<&Value>::new()); + assert_contains_all( + &parts, + &[ + json!({ "type": "reasoning-start", "id": "0" }), + json!({ + "type": "reasoning-end", + "id": "0", + "providerMetadata": { + "amazonBedrock": { "redactedContent": "encrypted-reasoning-payload" }, + "bedrock": { "redactedContent": "encrypted-reasoning-payload" }, + }, + }), + json!({ "type": "text-delta", "delta": "The answer is 42." }), + ], + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model-event-stream.test.ts › AmazonBedrockChatLanguageModel doStream › keeps multiple redacted reasoning blocks separate + #[tokio::test] + async fn keeps_multiple_redacted_reasoning_blocks_separate() { + let model = model_streaming( + "us.openai.gpt-5.6-luna", + &[ + create_event( + "contentBlockDelta", + &json!({ + "contentBlockIndex": 0, + "delta": { "reasoningContent": { "redactedContent": "first-payload" } }, + }) + .to_string(), + ), + create_event( + "contentBlockStop", + &json!({ "contentBlockIndex": 0 }).to_string(), + ), + create_event( + "contentBlockDelta", + &json!({ + "contentBlockIndex": 1, + "delta": { "reasoningContent": { "redactedContent": "second-payload" } }, + }) + .to_string(), + ), + create_event( + "contentBlockStop", + &json!({ "contentBlockIndex": 1 }).to_string(), + ), + create_event( + "contentBlockDelta", + &json!({ "contentBlockIndex": 2, "delta": { "text": "The answer is 42." } }) + .to_string(), + ), + create_event( + "contentBlockStop", + &json!({ "contentBlockIndex": 2 }).to_string(), + ), + create_event( + "messageStop", + &json!({ "stopReason": "end_turn" }).to_string(), + ), + ], + ); + + let parts = stream_parts(&model).await; + + let errors: Vec<&Value> = parts + .iter() + .filter(|part| part["type"] == json!("error")) + .collect(); + assert_eq!(errors, Vec::<&Value>::new()); + let reasoning_ends: Vec<&Value> = parts + .iter() + .filter(|part| part["type"] == json!("reasoning-end")) + .collect(); + assert_eq!( + reasoning_ends, + vec![ + &json!({ + "type": "reasoning-end", + "id": "0", + "providerMetadata": { + "amazonBedrock": { "redactedContent": "first-payload" }, + "bedrock": { "redactedContent": "first-payload" }, + }, + }), + &json!({ + "type": "reasoning-end", + "id": "1", + "providerMetadata": { + "amazonBedrock": { "redactedContent": "second-payload" }, + "bedrock": { "redactedContent": "second-payload" }, + }, + }), + ] + ); + } +} diff --git a/aimux-providers/src/bedrock/tests/chat_language_model_options.rs b/aimux-providers/src/bedrock/tests/chat_language_model_options.rs new file mode 100644 index 000000000..6319d2114 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/chat_language_model_options.rs @@ -0,0 +1,83 @@ +//! `amazonBedrockLanguageModelChatOptions` is a zod schema upstream. aimux +//! validates the same options where the request body is built +//! (`build_request_body_checked`), so `safeParse(options)` is that call with +//! `{ amazonBedrock: options }` as the provider options: `Ok` is +//! `success: true`, an `Err` is `success: false`. + +use serde_json::{Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::options::CallOptions; + +use crate::bedrock::convert::build_request_body_checked; + +/// `amazonBedrockLanguageModelChatOptions.safeParse(options)`: the request +/// body built with `options`, or the validation error. +fn safe_parse(options: Value) -> Result { + let call: CallOptions = serde_json::from_value(json!({ + "prompt": [{ "role": "user", "content": [{ "type": "text", "text": "Hello" }] }], + "providerOptions": { "amazonBedrock": options }, + })) + .expect("valid call options"); + build_request_body_checked("anthropic.claude-3-haiku-20240307-v1:0", &call, None) + .map(|(body, ..)| body) +} + +mod amazon_bedrock_language_model_chat_options { + use super::*; + + mod structured_output_mode { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model-options.test.ts › amazonBedrockLanguageModelChatOptions › structuredOutputMode › accepts %s + /// + /// The parsed value (`result.data`) has no Rust form: the mode only + /// steers how the body is built, so only acceptance is asserted. + #[test] + fn accepts() { + for structured_output_mode in ["outputFormat", "jsonTool", "auto"] { + let result = safe_parse(json!({ "structuredOutputMode": structured_output_mode })); + + assert!(result.is_ok(), "{structured_output_mode}: {result:?}"); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model-options.test.ts › amazonBedrockLanguageModelChatOptions › structuredOutputMode › rejects invalid values + #[test] + fn rejects_invalid_values() { + let result = safe_parse(json!({ "structuredOutputMode": "native" })); + + assert!(result.is_err(), "{result:?}"); + } + } + + mod service_tier { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model-options.test.ts › amazonBedrockLanguageModelChatOptions › serviceTier › accepts valid service tier values + #[test] + fn accepts_valid_service_tier_values() { + let valid_values = ["reserved", "priority", "default", "flex"]; + + for value in valid_values { + let result = safe_parse(json!({ "serviceTier": value })); + + let body = result.unwrap_or_else(|error| panic!("{value}: {error}")); + // `result.data.serviceTier` is what the body carries. + assert_eq!(body["serviceTier"]["type"], json!(value)); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-chat-language-model-options.test.ts › amazonBedrockLanguageModelChatOptions › serviceTier › rejects invalid service tier values + #[test] + fn rejects_invalid_service_tier_values() { + let invalid_values = ["on-demand", "auto", "standard", "", "PRIORITY"]; + + for value in invalid_values { + let result = safe_parse(json!({ "serviceTier": value })); + + assert!(result.is_err(), "{value:?}: {result:?}"); + } + } + } +} diff --git a/aimux-providers/src/bedrock/tests/convert_to_chat_messages.rs b/aimux-providers/src/bedrock/tests/convert_to_chat_messages.rs new file mode 100644 index 000000000..efea26a82 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/convert_to_chat_messages.rs @@ -0,0 +1,2409 @@ +use aimux_core::AiMuxError; +use aimux_core::language_model_message::LanguageModelPrompt; +use serde_json::{Value, json}; + +use crate::bedrock::convert::{convert_prompt, convert_prompt_to_bedrock}; + +/// Upstream prompts are written as `LanguageModelV4Prompt` JSON literals. +fn prompt(value: Value) -> LanguageModelPrompt { + serde_json::from_value(value).expect("valid prompt") +} + +fn to_result((system, messages): (Vec, Vec)) -> Value { + json!({ "system": system, "messages": messages }) +} + +/// `convertToAmazonBedrockChatMessages(prompt)` (isMistral not provided). +fn convert(value: Value) -> Result { + convert_prompt_to_bedrock(&prompt(value)).map(to_result) +} + +/// `convertToAmazonBedrockChatMessages(prompt, isMistral)`. +fn convert_with(value: Value, is_mistral: bool) -> Result { + convert_prompt(&prompt(value), is_mistral).map(to_result) +} + +/// The message of an `AI_UnsupportedFunctionalityError`. +/// The `functionality` of an `UnsupportedFunctionalityError`. +/// +/// aimux's `UnsupportedFunctionality` carries one string, upstream's +/// `functionality`; it has no separate `message`. Upstream's default message +/// is `'' functionality not supported.`, and the mime type +/// errors set their own message around the same functionality +/// (`convert-to-amazon-bedrock-chat-messages.ts:701-743`), so these cases +/// compare the functionality the upstream source throws with. +fn unsupported_functionality(error: AiMuxError) -> String { + match error { + AiMuxError::UnsupportedFunctionality(functionality) => functionality, + other => panic!("expected AI_UnsupportedFunctionalityError, got {other:?}"), + } +} + +mod system_messages { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › system messages › should throw an error if a system message is provided after a non-system message + #[test] + fn should_throw_an_error_if_a_system_message_is_provided_after_a_non_system_message() { + assert!( + convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "Hello" }] }, + { "role": "system", "content": "World" }, + ])) + .is_err() + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › system messages › should set isSystemCachePoint when system message has cache point + #[test] + fn should_set_is_system_cache_point_when_system_message_has_cache_point() { + let result = convert(json!([{ + "role": "system", + "content": "Hello", + "providerOptions": { "bedrock": { "cachePoint": { "type": "default" } } }, + }])) + .unwrap(); + + assert_eq!( + result, + json!({ + "system": [{ "text": "Hello" }, { "cachePoint": { "type": "default" } }], + "messages": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › system messages › should add cache point with 5m TTL to system message + #[test] + fn should_add_cache_point_with_5m_ttl_to_system_message() { + let result = convert(json!([{ + "role": "system", + "content": "Hello", + "providerOptions": { + "bedrock": { "cachePoint": { "type": "default", "ttl": "5m" } }, + }, + }])) + .unwrap(); + + assert_eq!( + result, + json!({ + "system": [ + { "text": "Hello" }, + { "cachePoint": { "type": "default", "ttl": "5m" } }, + ], + "messages": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › system messages › should add cache point with 1h TTL to system message + #[test] + fn should_add_cache_point_with_1h_ttl_to_system_message() { + let result = convert(json!([{ + "role": "system", + "content": "Hello", + "providerOptions": { + "bedrock": { "cachePoint": { "type": "default", "ttl": "1h" } }, + }, + }])) + .unwrap(); + + assert_eq!( + result, + json!({ + "system": [ + { "text": "Hello" }, + { "cachePoint": { "type": "default", "ttl": "1h" } }, + ], + "messages": [], + }) + ); + } +} + +mod user_messages { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should convert image parts with S3 URLs + #[test] + fn should_convert_image_parts_with_s3_urls() { + let result = convert(json!([{ + "role": "user", + "content": [ + { "type": "text", "text": "Describe the image" }, + { + "type": "file", + "data": { "type": "url", "url": "s3://my-test-bucket/path/to/image.png" }, + "mediaType": "image/png", + }, + ], + }])) + .unwrap(); + + assert_eq!( + result["messages"], + json!([{ + "role": "user", + "content": [ + { "text": "Describe the image" }, + { + "image": { + "format": "png", + "source": { + "s3Location": { "uri": "s3://my-test-bucket/path/to/image.png" }, + }, + }, + }, + ], + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should convert messages with video parts + #[test] + fn should_convert_messages_with_video_parts() { + let result = convert(json!([{ + "role": "user", + "content": [ + { "type": "text", "text": "Describe the video" }, + { + "type": "file", + "data": { "type": "data", "data": "AAECAw==" }, + "mediaType": "video/mp4", + }, + ], + }])) + .unwrap(); + + assert_eq!( + result["messages"], + json!([{ + "role": "user", + "content": [ + { "text": "Describe the video" }, + { "video": { "format": "mp4", "source": { "bytes": "AAECAw==" } } }, + ], + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should convert video parts with S3 URLs + #[test] + fn should_convert_video_parts_with_s3_urls() { + let result = convert(json!([{ + "role": "user", + "content": [ + { "type": "text", "text": "Describe the video" }, + { + "type": "file", + "data": { "type": "url", "url": "s3://my-test-bucket/path/to/video.mp4" }, + "mediaType": "video/mp4", + }, + ], + }])) + .unwrap(); + + assert_eq!( + result["messages"], + json!([{ + "role": "user", + "content": [ + { "text": "Describe the video" }, + { + "video": { + "format": "mp4", + "source": { + "s3Location": { "uri": "s3://my-test-bucket/path/to/video.mp4" }, + }, + }, + }, + ], + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should throw for unsupported video mime types + #[test] + fn should_throw_for_unsupported_video_mime_types() { + let error = convert(json!([{ + "role": "user", + "content": [{ + "type": "file", + "data": { "type": "data", "data": "base64data" }, + "mediaType": "video/unsupported", + }], + }])) + .unwrap_err(); + + assert_eq!( + unsupported_functionality(error), + "video mime type: video/unsupported" + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should convert messages with document parts + #[test] + fn should_convert_messages_with_document_parts() { + let result = convert(json!([{ + "role": "user", + "content": [ + { "type": "text", "text": "Hello" }, + { + "type": "file", + "data": { "type": "data", "data": "AAECAw==" }, + "mediaType": "application/pdf", + }, + ], + }])) + .unwrap(); + + assert_eq!( + result["messages"], + json!([{ + "content": [ + { "text": "Hello" }, + { + "document": { + "format": "pdf", + "name": "document-1", + "source": { "bytes": "AAECAw==" }, + }, + }, + ], + "role": "user", + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should strip file extension when filename is provided + #[test] + fn should_strip_file_extension_when_filename_is_provided() { + let result = convert(json!([{ + "role": "user", + "content": [ + { "type": "text", "text": "Hello" }, + { + "type": "file", + "data": { "type": "data", "data": "AAECAw==" }, + "mediaType": "application/pdf", + "filename": "custom-filename.pdf", + }, + ], + }])) + .unwrap(); + + assert_eq!( + result["messages"], + json!([{ + "content": [ + { "text": "Hello" }, + { + "document": { + "format": "pdf", + "name": "custom-filename", + "source": { "bytes": "AAECAw==" }, + }, + }, + ], + "role": "user", + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should preserve filename without extension when provided + #[test] + fn should_preserve_filename_without_extension_when_provided() { + let result = convert(json!([{ + "role": "user", + "content": [{ + "type": "file", + "data": { "type": "data", "data": "AAECAw==" }, + "mediaType": "application/pdf", + "filename": "custom-filename", + }], + }])) + .unwrap(); + + assert_eq!( + result["messages"], + json!([{ + "content": [{ + "document": { + "format": "pdf", + "name": "custom-filename", + "source": { "bytes": "AAECAw==" }, + }, + }], + "role": "user", + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should sanitize document filenames to Bedrock-compatible names + #[test] + fn should_sanitize_document_filenames_to_bedrock_compatible_names() { + let long_name = format!("{}.txt", "a".repeat(201)); + let filenames = [ + "John's report.txt", + "invoice #123.txt", + "a&b.txt", + "report,2026.txt", + "résumé.txt", + "분기보고서.txt", + "Report - Final.txt", + "a\tb.txt", + long_name.as_str(), + ".txt", + "report (final) [v2]_draft.txt", + ]; + + let result = convert(json!([{ + "role": "user", + "content": filenames + .iter() + .map(|filename| json!({ + "type": "file", + "data": { "type": "data", "data": "base64data" }, + "mediaType": "application/pdf", + "filename": filename, + })) + .collect::>(), + }])) + .unwrap(); + + let names: Vec = result["messages"][0]["content"] + .as_array() + .unwrap() + .iter() + .map(|content| content["document"]["name"].clone()) + .collect(); + assert_eq!( + Value::Array(names), + json!([ + "Johns report", + "invoice 123", + "ab", + "report2026", + "rsum", + "document-1", + "Report - Final", + "a b", + "a".repeat(200), + "document-2", + "report (final) [v2]draft", + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should sanitize filenames for text document data + #[test] + fn should_sanitize_filenames_for_text_document_data() { + let result = convert(json!([{ + "role": "user", + "content": [{ + "type": "file", + "data": { "type": "text", "text": "Hello" }, + "mediaType": "text/plain", + "filename": "John's report.txt", + }], + }])) + .unwrap(); + + assert_eq!( + result["messages"][0]["content"][0], + json!({ + "document": { + "format": "txt", + "name": "Johns report", + "source": { "bytes": "SGVsbG8=" }, + }, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should add cache point to user message content when specified + #[test] + fn should_add_cache_point_to_user_message_content_when_specified() { + let result = convert(json!([{ + "role": "user", + "content": [{ "type": "text", "text": "Hello" }], + "providerOptions": { "bedrock": { "cachePoint": { "type": "default" } } }, + }])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [{ + "role": "user", + "content": [{ "text": "Hello" }, { "cachePoint": { "type": "default" } }], + }], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should add cache point with 5m TTL to user message + #[test] + fn should_add_cache_point_with_5m_ttl_to_user_message() { + let result = convert(json!([{ + "role": "user", + "content": [{ "type": "text", "text": "Hello" }], + "providerOptions": { + "bedrock": { "cachePoint": { "type": "default", "ttl": "5m" } }, + }, + }])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [{ + "role": "user", + "content": [ + { "text": "Hello" }, + { "cachePoint": { "type": "default", "ttl": "5m" } }, + ], + }], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should add cache point with 1h TTL to user message + #[test] + fn should_add_cache_point_with_1h_ttl_to_user_message() { + let result = convert(json!([{ + "role": "user", + "content": [{ "type": "text", "text": "Hello" }], + "providerOptions": { + "bedrock": { "cachePoint": { "type": "default", "ttl": "1h" } }, + }, + }])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [{ + "role": "user", + "content": [ + { "text": "Hello" }, + { "cachePoint": { "type": "default", "ttl": "1h" } }, + ], + }], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should convert text part to guardContent when guardContent provider option is true + #[test] + fn should_convert_text_part_to_guard_content_when_guard_content_option_is_true() { + let result = convert(json!([{ + "role": "user", + "content": [{ + "type": "text", + "text": "Grounding text", + "providerOptions": { "amazonBedrock": { "guardContent": true } }, + }], + }])) + .unwrap(); + + assert_eq!( + result["messages"], + json!([{ + "role": "user", + "content": [{ "guardContent": { "text": { "text": "Grounding text" } } }], + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should convert text part to guardContent with qualifiers + #[test] + fn should_convert_text_part_to_guard_content_with_qualifiers() { + let result = convert(json!([{ + "role": "user", + "content": [{ + "type": "text", + "text": "Grounding text", + "providerOptions": { + "amazonBedrock": { + "guardContent": true, + "guardContentQualifiers": ["grounding_source"], + }, + }, + }], + }])) + .unwrap(); + + assert_eq!( + result["messages"], + json!([{ + "role": "user", + "content": [{ + "guardContent": { + "text": { "text": "Grounding text", "qualifiers": ["grounding_source"] }, + }, + }], + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should convert text part to guardContent with multiple qualifiers + #[test] + fn should_convert_text_part_to_guard_content_with_multiple_qualifiers() { + let result = convert(json!([{ + "role": "user", + "content": [{ + "type": "text", + "text": "Query text", + "providerOptions": { + "amazonBedrock": { + "guardContent": true, + "guardContentQualifiers": ["grounding_source", "query"], + }, + }, + }], + }])) + .unwrap(); + + assert_eq!( + result["messages"], + json!([{ + "role": "user", + "content": [{ + "guardContent": { + "text": { + "text": "Query text", + "qualifiers": ["grounding_source", "query"], + }, + }, + }], + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should convert text part as normal text when guardContent is false + #[test] + fn should_convert_text_part_as_normal_text_when_guard_content_is_false() { + let result = convert(json!([{ + "role": "user", + "content": [{ + "type": "text", + "text": "Normal text", + "providerOptions": { "amazonBedrock": { "guardContent": false } }, + }], + }])) + .unwrap(); + + assert_eq!( + result["messages"], + json!([{ "role": "user", "content": [{ "text": "Normal text" }] }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should convert text part as normal text when no provider options + #[test] + fn should_convert_text_part_as_normal_text_when_no_provider_options() { + let result = convert(json!([{ + "role": "user", + "content": [{ "type": "text", "text": "Normal text" }], + }])) + .unwrap(); + + assert_eq!( + result["messages"], + json!([{ "role": "user", "content": [{ "text": "Normal text" }] }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should convert image part to guardContent when guardContent provider option is true + #[test] + fn should_convert_image_part_to_guard_content_when_guard_content_option_is_true() { + let result = convert(json!([{ + "role": "user", + "content": [{ + "type": "file", + "data": { "type": "data", "data": "AAECAw==" }, + "mediaType": "image/png", + "providerOptions": { "amazonBedrock": { "guardContent": true } }, + }], + }])) + .unwrap(); + + assert_eq!( + result["messages"], + json!([{ + "role": "user", + "content": [{ + "guardContent": { + "image": { "format": "png", "source": { "bytes": "AAECAw==" } }, + }, + }], + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should convert image part as normal image when guardContent is false + #[test] + fn should_convert_image_part_as_normal_image_when_guard_content_is_false() { + let result = convert(json!([{ + "role": "user", + "content": [{ + "type": "file", + "data": { "type": "data", "data": "AAECAw==" }, + "mediaType": "image/png", + "providerOptions": { "amazonBedrock": { "guardContent": false } }, + }], + }])) + .unwrap(); + + assert_eq!( + result["messages"], + json!([{ + "role": "user", + "content": [{ "image": { "format": "png", "source": { "bytes": "AAECAw==" } } }], + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should convert image part as normal image when no provider options + #[test] + fn should_convert_image_part_as_normal_image_when_no_provider_options() { + let result = convert(json!([{ + "role": "user", + "content": [{ + "type": "file", + "data": { "type": "data", "data": "AAECAw==" }, + "mediaType": "image/png", + }], + }])) + .unwrap(); + + assert_eq!( + result["messages"], + json!([{ + "role": "user", + "content": [{ "image": { "format": "png", "source": { "bytes": "AAECAw==" } } }], + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should throw for file parts with provider references + #[test] + fn should_throw_for_file_parts_with_provider_references() { + let error = convert(json!([{ + "role": "user", + "content": [{ + "type": "file", + "data": { "type": "reference", "reference": { "bedrock": "file-ref-123" } }, + "mediaType": "image/png", + }], + }])) + .unwrap_err(); + + assert_eq!( + unsupported_functionality(error), + "file parts with provider references" + ); + } +} + +mod assistant_messages { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should preserve the order of provider-executed tool calls and results + #[test] + fn should_preserve_the_order_of_provider_executed_tool_calls_and_results() { + let result = convert(json!([ + { + "role": "user", + "content": [{ "type": "text", "text": "Add 2 and 2, then add 3 and 3." }], + }, + { + "role": "assistant", + "content": [ + { "type": "text", "text": "Running the additions." }, + { + "type": "tool-call", + "toolCallId": "call-1", + "toolName": "add", + "input": { "a": 2, "b": 2 }, + "providerExecuted": true, + }, + { + "type": "tool-result", + "toolCallId": "call-1", + "toolName": "add", + "output": { "type": "json", "value": { "sum": 4 } }, + }, + { + "type": "tool-call", + "toolCallId": "call-2", + "toolName": "add", + "input": { "a": 3, "b": 3 }, + "providerExecuted": true, + }, + { + "type": "tool-result", + "toolCallId": "call-2", + "toolName": "add", + "output": { "type": "json", "value": { "sum": 6 } }, + }, + { "type": "text", "text": "The sums are 4 and 6." }, + ], + }, + { + "role": "user", + "content": [{ "type": "text", "text": "Now add those sums together." }], + }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "system": [], + "messages": [ + { "role": "user", "content": [{ "text": "Add 2 and 2, then add 3 and 3." }] }, + { + "role": "assistant", + "content": [ + { "text": "Running the additions." }, + { + "toolUse": { + "toolUseId": "call-1", + "name": "add", + "input": { "a": 2, "b": 2 }, + }, + }, + ], + }, + { + "role": "user", + "content": [{ + "toolResult": { + "toolUseId": "call-1", + "content": [{ "text": "{\"sum\":4}" }], + }, + }], + }, + { + "role": "assistant", + "content": [{ + "toolUse": { + "toolUseId": "call-2", + "name": "add", + "input": { "a": 3, "b": 3 }, + }, + }], + }, + { + "role": "user", + "content": [{ + "toolResult": { + "toolUseId": "call-2", + "content": [{ "text": "{\"sum\":6}" }], + }, + }], + }, + { "role": "assistant", "content": [{ "text": "The sums are 4 and 6." }] }, + { "role": "user", "content": [{ "text": "Now add those sums together." }] }, + ], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should combine a trailing provider-executed tool result with the next user message + #[test] + fn should_combine_a_trailing_provider_executed_tool_result_with_the_next_user_message() { + let result = convert(json!([ + { + "role": "assistant", + "content": [ + { + "type": "tool-call", + "toolCallId": "call-1", + "toolName": "add", + "input": { "a": 2, "b": 2 }, + "providerExecuted": true, + }, + { + "type": "tool-result", + "toolCallId": "call-1", + "toolName": "add", + "output": { "type": "json", "value": { "sum": 4 } }, + }, + ], + }, + { "role": "user", "content": [{ "type": "text", "text": "Now double that." }] }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "system": [], + "messages": [ + { + "role": "assistant", + "content": [{ + "toolUse": { + "toolUseId": "call-1", + "name": "add", + "input": { "a": 2, "b": 2 }, + }, + }], + }, + { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "call-1", + "content": [{ "text": "{\"sum\":4}" }], + }, + }, + { "text": "Now double that." }, + ], + }, + ], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should remove trailing whitespace from last assistant message when there is no further user message + #[test] + fn should_remove_trailing_whitespace_from_last_assistant_message() { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "user content" }] }, + { + "role": "assistant", + "content": [{ "type": "text", "text": "assistant content " }], + }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "role": "user", "content": [{ "text": "user content" }] }, + { "role": "assistant", "content": [{ "text": "assistant content" }] }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should remove trailing whitespace from last assistant message with multi-part content when there is no further user message + #[test] + fn should_remove_trailing_whitespace_from_last_assistant_message_with_multi_part_content() { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "user content" }] }, + { + "role": "assistant", + "content": [ + { "type": "text", "text": "assistant " }, + { "type": "text", "text": "content " }, + ], + }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "role": "user", "content": [{ "text": "user content" }] }, + { + "role": "assistant", + "content": [{ "text": "assistant " }, { "text": "content" }], + }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should keep trailing whitespace from assistant message when there is a further user message + #[test] + fn should_keep_trailing_whitespace_from_assistant_message_when_there_is_a_further_user_message() + { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "user content" }] }, + { + "role": "assistant", + "content": [{ "type": "text", "text": "assistant content " }], + }, + { "role": "user", "content": [{ "type": "text", "text": "user content 2" }] }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "role": "user", "content": [{ "text": "user content" }] }, + { "role": "assistant", "content": [{ "text": "assistant content " }] }, + { "role": "user", "content": [{ "text": "user content 2" }] }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should combine multiple sequential assistant messages into a single message + #[test] + fn should_combine_multiple_sequential_assistant_messages_into_a_single_message() { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "Hi!" }] }, + { "role": "assistant", "content": [{ "type": "text", "text": "Hello" }] }, + { "role": "assistant", "content": [{ "type": "text", "text": "World" }] }, + { "role": "assistant", "content": [{ "type": "text", "text": "!" }] }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "role": "user", "content": [{ "text": "Hi!" }] }, + { + "role": "assistant", + "content": [{ "text": "Hello" }, { "text": "World" }, { "text": "!" }], + }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should add cache point to assistant message content when specified + #[test] + fn should_add_cache_point_to_assistant_message_content_when_specified() { + let result = convert(json!([{ + "role": "assistant", + "content": [{ "type": "text", "text": "Hello" }], + "providerOptions": { "bedrock": { "cachePoint": { "type": "default" } } }, + }])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [{ + "role": "assistant", + "content": [{ "text": "Hello" }, { "cachePoint": { "type": "default" } }], + }], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should add cache point with 5m TTL to assistant message + #[test] + fn should_add_cache_point_with_5m_ttl_to_assistant_message() { + let result = convert(json!([{ + "role": "assistant", + "content": [{ "type": "text", "text": "Hello" }], + "providerOptions": { + "bedrock": { "cachePoint": { "type": "default", "ttl": "5m" } }, + }, + }])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [{ + "role": "assistant", + "content": [ + { "text": "Hello" }, + { "cachePoint": { "type": "default", "ttl": "5m" } }, + ], + }], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should add cache point with 1h TTL to assistant message + #[test] + fn should_add_cache_point_with_1h_ttl_to_assistant_message() { + let result = convert(json!([{ + "role": "assistant", + "content": [{ "type": "text", "text": "Hello" }], + "providerOptions": { + "bedrock": { "cachePoint": { "type": "default", "ttl": "1h" } }, + }, + }])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [{ + "role": "assistant", + "content": [ + { "text": "Hello" }, + { "cachePoint": { "type": "default", "ttl": "1h" } }, + ], + }], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should add cache point to assistant content part when specified + #[test] + fn should_add_cache_point_to_assistant_content_part_when_specified() { + let result = convert(json!([{ + "role": "assistant", + "content": [ + { "type": "text", "text": "Hello" }, + { + "type": "text", + "text": "cached", + "providerOptions": { + "bedrock": { "cachePoint": { "type": "default", "ttl": "1h" } }, + }, + }, + { "type": "text", "text": "World" }, + ], + }])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [{ + "role": "assistant", + "content": [ + { "text": "Hello" }, + { "text": "cached" }, + { "cachePoint": { "type": "default", "ttl": "1h" } }, + { "text": "World" }, + ], + }], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should properly convert reasoning content type + #[test] + fn should_properly_convert_reasoning_content_type() { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "Explain your reasoning" }] }, + { + "role": "assistant", + "content": [{ + "type": "reasoning", + "text": "This is my step-by-step reasoning process", + "providerOptions": { "bedrock": { "signature": "test-signature" } }, + }], + }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "role": "user", "content": [{ "text": "Explain your reasoning" }] }, + { + "role": "assistant", + "content": [{ + "reasoningContent": { + "reasoningText": { + "text": "This is my step-by-step reasoning process", + "signature": "test-signature", + }, + }, + }], + }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should properly convert redacted-reasoning content type + #[test] + fn should_properly_convert_redacted_reasoning_content_type() { + let reasoning_data = "Redacted reasoning information"; + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "Explain your reasoning" }] }, + { + "role": "assistant", + "content": [{ + "type": "reasoning", + "text": "", + "providerOptions": { "bedrock": { "redactedData": reasoning_data } }, + }], + }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "role": "user", "content": [{ "text": "Explain your reasoning" }] }, + { + "role": "assistant", + "content": [{ + "reasoningContent": { "redactedReasoning": { "data": reasoning_data } }, + }], + }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should replay reasoning redacted as `redactedContent` + #[test] + fn should_replay_reasoning_redacted_as_redacted_content() { + let redacted_content = "encrypted-reasoning-payload"; + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "Explain your reasoning" }] }, + { + "role": "assistant", + "content": [{ + "type": "reasoning", + "text": "", + "providerOptions": { "bedrock": { "redactedContent": redacted_content } }, + }], + }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "role": "user", "content": [{ "text": "Explain your reasoning" }] }, + { + "role": "assistant", + "content": [{ "reasoningContent": { "redactedContent": redacted_content } }], + }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should omit assistant message reasoning parts signed by a foreign provider + #[test] + fn should_omit_assistant_message_reasoning_parts_signed_by_a_foreign_provider() { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "Explain your reasoning" }] }, + { + "role": "assistant", + "content": [ + { + "type": "reasoning", + "text": "Anthropic-signed reasoning replayed to Bedrock", + "providerOptions": { "anthropic": { "signature": "anthropic-signature" } }, + }, + { "type": "text", "text": "final answer" }, + ], + }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "role": "user", "content": [{ "text": "Explain your reasoning" }] }, + { "role": "assistant", "content": [{ "text": "final answer" }] }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should preserve assistant message reasoning parts with amazonBedrock providerOptions + #[test] + fn should_preserve_assistant_message_reasoning_parts_with_amazon_bedrock_provider_options() { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "Explain your reasoning" }] }, + { + "role": "assistant", + "content": [ + { + "type": "reasoning", + "text": "Bedrock-signed reasoning round-tripped to Bedrock", + "providerOptions": { "amazonBedrock": { "signature": "bedrock-signature" } }, + }, + { "type": "text", "text": "final answer" }, + ], + }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "role": "user", "content": [{ "text": "Explain your reasoning" }] }, + { + "role": "assistant", + "content": [ + { + "reasoningContent": { + "reasoningText": { + "text": "Bedrock-signed reasoning round-tripped to Bedrock", + "signature": "bedrock-signature", + }, + }, + }, + { "text": "final answer" }, + ], + }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should not trim reasoning text when a signature is present + #[test] + fn should_not_trim_reasoning_text_when_a_signature_is_present() { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "Explain your reasoning" }] }, + { + "role": "assistant", + "content": [{ + "type": "reasoning", + "text": "This is my reasoning with trailing space ", + "providerOptions": { "bedrock": { "signature": "test-signature" } }, + }], + }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "content": [{ "text": "Explain your reasoning" }], "role": "user" }, + { + "content": [{ + "reasoningContent": { + "reasoningText": { + "signature": "test-signature", + "text": "This is my reasoning with trailing space ", + }, + }, + }], + "role": "assistant", + }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should omit reasoning content without signature + #[test] + fn should_omit_reasoning_content_without_signature() { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "Explain your reasoning" }] }, + { + "role": "assistant", + "content": [ + { "type": "reasoning", "text": "This is my reasoning with trailing space " }, + { "type": "text", "text": "final answer" }, + ], + }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "content": [{ "text": "Explain your reasoning" }], "role": "user" }, + { "content": [{ "text": "final answer" }], "role": "assistant" }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should omit an assistant message when unsigned reasoning is its only content + #[test] + fn should_omit_an_assistant_message_when_unsigned_reasoning_is_its_only_content() { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "Think hard then answer" }] }, + { + "role": "assistant", + "content": [{ "type": "reasoning", "text": "Let me consider the options" }], + }, + { "role": "user", "content": [{ "type": "text", "text": "Hello?" }] }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "role": "user", "content": [{ "text": "Think hard then answer" }] }, + { "role": "user", "content": [{ "text": "Hello?" }] }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should omit an assistant message when only a cache point remains after filtering unsigned reasoning + #[test] + fn should_omit_an_assistant_message_when_only_a_cache_point_remains() { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "Think hard then answer" }] }, + { + "role": "assistant", + "content": [{ "type": "reasoning", "text": "Let me consider the options" }], + "providerOptions": { "bedrock": { "cachePoint": { "type": "default" } } }, + }, + { "role": "user", "content": [{ "type": "text", "text": "Hello?" }] }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "role": "user", "content": [{ "text": "Think hard then answer" }] }, + { "role": "user", "content": [{ "text": "Hello?" }] }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should omit multiple reasoning parts without signatures + #[test] + fn should_omit_multiple_reasoning_parts_without_signatures() { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "Explain your reasoning" }] }, + { + "role": "assistant", + "content": [ + { "type": "reasoning", "text": "First reasoning with trailing space " }, + { "type": "reasoning", "text": "Second reasoning with trailing space " }, + { "type": "text", "text": "final answer" }, + ], + }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "content": [{ "text": "Explain your reasoning" }], "role": "user" }, + { "content": [{ "text": "final answer" }], "role": "assistant" }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should omit unsigned reasoning while preserving tool calls in multi-turn tool use + #[test] + fn should_omit_unsigned_reasoning_while_preserving_tool_calls_in_multi_turn_tool_use() { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "What is the weather?" }] }, + { + "role": "assistant", + "content": [ + { "type": "reasoning", "text": "I should call the weather tool." }, + { + "type": "tool-call", + "toolCallId": "call-1", + "toolName": "getWeather", + "input": { "city": "SF" }, + }, + ], + }, + { + "role": "tool", + "content": [{ + "type": "tool-result", + "toolCallId": "call-1", + "toolName": "getWeather", + "output": { "type": "text", "value": "Sunny, 72F" }, + }], + }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "content": [{ "text": "What is the weather?" }], "role": "user" }, + { + "content": [{ + "toolUse": { + "input": { "city": "SF" }, + "name": "getWeather", + "toolUseId": "call-1", + }, + }], + "role": "assistant", + }, + { + "content": [{ + "toolResult": { + "content": [{ "text": "Sunny, 72F" }], + "toolUseId": "call-1", + }, + }], + "role": "user", + }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should preserve reasoning text with signature in multi-turn tool use + #[test] + fn should_preserve_reasoning_text_with_signature_in_multi_turn_tool_use() { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "What is the weather?" }] }, + { + "role": "assistant", + "content": [ + { + "type": "reasoning", + "text": "Let me check the weather API.\n", + "providerOptions": { "bedrock": { "signature": "sig-abc123" } }, + }, + { + "type": "tool-call", + "toolCallId": "call-1", + "toolName": "getWeather", + "input": { "city": "SF" }, + }, + ], + }, + { + "role": "tool", + "content": [{ + "type": "tool-result", + "toolCallId": "call-1", + "toolName": "getWeather", + "output": { "type": "text", "value": "Sunny, 72F" }, + }], + }, + { + "role": "assistant", + "content": [ + { + "type": "reasoning", + "text": "The weather is sunny and warm.\n", + "providerOptions": { "bedrock": { "signature": "sig-def456" } }, + }, + { "type": "text", "text": "It is sunny and 72F in SF." }, + ], + }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "content": [{ "text": "What is the weather?" }], "role": "user" }, + { + "content": [ + { + "reasoningContent": { + "reasoningText": { + "signature": "sig-abc123", + "text": "Let me check the weather API.\n", + }, + }, + }, + { + "toolUse": { + "input": { "city": "SF" }, + "name": "getWeather", + "toolUseId": "call-1", + }, + }, + ], + "role": "assistant", + }, + { + "content": [{ + "toolResult": { + "content": [{ "text": "Sunny, 72F" }], + "toolUseId": "call-1", + }, + }], + "role": "user", + }, + { + "content": [ + { + "reasoningContent": { + "reasoningText": { + "signature": "sig-def456", + "text": "The weather is sunny and warm.\n", + }, + }, + }, + { "text": "It is sunny and 72F in SF." }, + ], + "role": "assistant", + }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should handle a mix of text and reasoning content types + #[test] + fn should_handle_a_mix_of_text_and_reasoning_content_types() { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "Explain your reasoning" }] }, + { + "role": "assistant", + "content": [ + { "type": "text", "text": "My answer is 42." }, + { + "type": "reasoning", + "text": "I calculated this by analyzing the meaning of life", + "providerOptions": { "bedrock": { "signature": "reasoning-process" } }, + }, + ], + }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "role": "user", "content": [{ "text": "Explain your reasoning" }] }, + { + "role": "assistant", + "content": [ + { "text": "My answer is 42." }, + { + "reasoningContent": { + "reasoningText": { + "text": "I calculated this by analyzing the meaning of life", + "signature": "reasoning-process", + }, + }, + }, + ], + }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should filter out empty text blocks in assistant messages + #[test] + fn should_filter_out_empty_text_blocks_in_assistant_messages() { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "Hello" }] }, + { + "role": "assistant", + "content": [ + { "type": "text", "text": "\n\n" }, + { "type": "tool-call", "toolCallId": "call-123", "toolName": "test", "input": {} }, + { "type": "text", "text": " " }, + { "type": "text", "text": "actual content" }, + ], + }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "role": "user", "content": [{ "text": "Hello" }] }, + { + "role": "assistant", + "content": [ + { "toolUse": { "toolUseId": "call-123", "name": "test", "input": {} } }, + { "text": "actual content" }, + ], + }, + ], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should wrap non-object (invalid) tool call input in an object + #[test] + fn should_wrap_non_object_invalid_tool_call_input_in_an_object() { + let result = convert(json!([{ + "role": "assistant", + "content": [{ + "type": "tool-call", + "toolCallId": "call-1", + "toolName": "cityAttractions", + // malformed JSON the model produced, kept as a raw string + "input": "{ \"city\": \"San Francisco\", }", + }], + }])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [{ + "role": "assistant", + "content": [{ + "toolUse": { + "toolUseId": "call-1", + "name": "cityAttractions", + "input": { "rawInvalidInput": "{ \"city\": \"San Francisco\", }" }, + }, + }], + }], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should strip invalid characters from tool call names + #[test] + fn should_strip_invalid_characters_from_tool_call_names() { + let result = convert(json!([{ + "role": "assistant", + "content": [ + { "type": "tool-call", "toolCallId": "call-1", "toolName": "$READFILE", "input": {} }, + { + "type": "tool-call", + "toolCallId": "call-2", + "toolName": "exchange_delivered_order_items<|channel|>", + "input": {}, + }, + { "type": "tool-call", "toolCallId": "call-3", "toolName": "$", "input": {} }, + ], + }])) + .unwrap(); + + assert_eq!( + result["messages"], + json!([{ + "role": "assistant", + "content": [ + { "toolUse": { "toolUseId": "call-1", "name": "READFILE", "input": {} } }, + { + "toolUse": { + "toolUseId": "call-2", + "name": "exchange_delivered_order_itemschannel", + "input": {}, + }, + }, + { "toolUse": { "toolUseId": "call-3", "name": "_", "input": {} } }, + ], + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › assistant messages › should preserve empty text blocks when reasoning blocks are present + #[test] + fn should_preserve_empty_text_blocks_when_reasoning_blocks_are_present() { + let result = convert(json!([ + { "role": "user", "content": [{ "type": "text", "text": "Hello" }] }, + { + "role": "assistant", + "content": [ + { + "type": "reasoning", + "text": "thinking...", + "providerOptions": { "bedrock": { "signature": "sig-1" } }, + }, + { "type": "text", "text": "" }, + { + "type": "reasoning", + "text": "more thinking...", + "providerOptions": { "bedrock": { "signature": "sig-2" } }, + }, + { "type": "text", "text": "response text" }, + { "type": "tool-call", "toolCallId": "call-123", "toolName": "test", "input": {} }, + ], + }, + ])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [ + { "role": "user", "content": [{ "text": "Hello" }] }, + { + "role": "assistant", + "content": [ + { + "reasoningContent": { + "reasoningText": { "text": "thinking...", "signature": "sig-1" }, + }, + }, + { "text": "" }, + { + "reasoningContent": { + "reasoningText": { + "text": "more thinking...", + "signature": "sig-2", + }, + }, + }, + { "text": "response text" }, + { "toolUse": { "toolUseId": "call-123", "name": "test", "input": {} } }, + ], + }, + ], + "system": [], + }) + ); + } +} + +mod tool_messages { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › tool messages › should convert tool result with content array containing image + #[test] + fn should_convert_tool_result_with_content_array_containing_image() { + let result = convert(json!([{ + "role": "tool", + "content": [{ + "type": "tool-result", + "toolCallId": "call-123", + "toolName": "image-generator", + "output": { + "type": "content", + "value": [{ + "type": "file", + "data": { "type": "data", "data": "base64data" }, + "mediaType": "image/jpeg", + }], + }, + }], + }])) + .unwrap(); + + assert_eq!( + result["messages"][0], + json!({ + "role": "user", + "content": [{ + "toolResult": { + "toolUseId": "call-123", + "content": [ + { "image": { "format": "jpeg", "source": { "bytes": "base64data" } } }, + ], + }, + }], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › tool messages › should convert tool result images with S3 URLs + #[test] + fn should_convert_tool_result_images_with_s3_urls() { + let result = convert(json!([{ + "role": "tool", + "content": [{ + "type": "tool-result", + "toolCallId": "call-123", + "toolName": "image-generator", + "output": { + "type": "content", + "value": [{ + "type": "file", + "data": { "type": "url", "url": "s3://my-test-bucket/path/to/image.png" }, + "mediaType": "image/png", + }], + }, + }], + }])) + .unwrap(); + + assert_eq!( + result["messages"][0], + json!({ + "role": "user", + "content": [{ + "toolResult": { + "toolUseId": "call-123", + "content": [{ + "image": { + "format": "png", + "source": { + "s3Location": { "uri": "s3://my-test-bucket/path/to/image.png" }, + }, + }, + }], + }, + }], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › tool messages › should convert tool result with content array containing video + #[test] + fn should_convert_tool_result_with_content_array_containing_video() { + let result = convert(json!([{ + "role": "tool", + "content": [{ + "type": "tool-result", + "toolCallId": "call-123", + "toolName": "video-analyzer", + "output": { + "type": "content", + "value": [{ + "type": "file", + "data": { "type": "data", "data": "base64data" }, + "mediaType": "video/mp4", + }], + }, + }], + }])) + .unwrap(); + + assert_eq!( + result["messages"][0], + json!({ + "role": "user", + "content": [{ + "toolResult": { + "toolUseId": "call-123", + "content": [ + { "video": { "format": "mp4", "source": { "bytes": "base64data" } } }, + ], + }, + }], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › tool messages › should convert tool result videos with S3 URLs + #[test] + fn should_convert_tool_result_videos_with_s3_urls() { + let result = convert(json!([{ + "role": "tool", + "content": [{ + "type": "tool-result", + "toolCallId": "call-123", + "toolName": "video-analyzer", + "output": { + "type": "content", + "value": [{ + "type": "file", + "data": { "type": "url", "url": "s3://my-test-bucket/path/to/video.mp4" }, + "mediaType": "video/mp4", + }], + }, + }], + }])) + .unwrap(); + + assert_eq!( + result["messages"][0], + json!({ + "role": "user", + "content": [{ + "toolResult": { + "toolUseId": "call-123", + "content": [{ + "video": { + "format": "mp4", + "source": { + "s3Location": { "uri": "s3://my-test-bucket/path/to/video.mp4" }, + }, + }, + }], + }, + }], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › tool messages › should convert tool result with content array containing PDF + #[test] + fn should_convert_tool_result_with_content_array_containing_pdf() { + let result = convert(json!([{ + "role": "tool", + "content": [{ + "type": "tool-result", + "toolCallId": "call-123", + "toolName": "document-reader", + "output": { + "type": "content", + "value": [{ + "type": "file", + "data": { "type": "data", "data": "base64data" }, + "mediaType": "application/pdf", + "filename": "tool's result.pdf", + }], + }, + }], + }])) + .unwrap(); + + assert_eq!( + result["messages"][0], + json!({ + "role": "user", + "content": [{ + "toolResult": { + "toolUseId": "call-123", + "content": [{ + "document": { + "format": "pdf", + "name": "tools result", + "source": { "bytes": "base64data" }, + }, + }], + }, + }], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › tool messages › should throw error for unsupported image format in tool result content + #[test] + fn should_throw_error_for_unsupported_image_format_in_tool_result_content() { + let error = convert(json!([{ + "role": "tool", + "content": [{ + "type": "tool-result", + "toolCallId": "call-123", + "toolName": "image-generator", + "output": { + "type": "content", + "value": [{ + "type": "file", + "data": { "type": "data", "data": "base64data" }, + "mediaType": "image/avif", + }], + }, + }], + }])) + .unwrap_err(); + + assert_eq!( + unsupported_functionality(error), + "image mime type: image/avif" + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › tool messages › should throw error for unsupported mime type in tool result file content + #[test] + fn should_throw_error_for_unsupported_mime_type_in_tool_result_file_content() { + let error = convert(json!([{ + "role": "tool", + "content": [{ + "type": "tool-result", + "toolCallId": "call-123", + "toolName": "image-generator", + "output": { + "type": "content", + "value": [{ + "type": "file", + "data": { "type": "data", "data": "base64data" }, + "mediaType": "unsupported/mime-type", + }], + }, + }], + }])) + .unwrap_err(); + + assert_eq!( + unsupported_functionality(error), + "file mime type: unsupported/mime-type" + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › tool messages › should fallback to stringified result when content is undefined + #[test] + fn should_fallback_to_stringified_result_when_content_is_undefined() { + let result = convert(json!([{ + "role": "tool", + "content": [{ + "type": "tool-result", + "toolCallId": "call-123", + "toolName": "calculator", + "output": { "type": "json", "value": { "value": 42 } }, + }], + }])) + .unwrap(); + + assert_eq!( + result["messages"][0], + json!({ + "role": "user", + "content": [{ + "toolResult": { + "toolUseId": "call-123", + "content": [{ "text": "{\"value\":42}" }], + }, + }], + }) + ); + } +} + +mod additional_file_format_tests { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › additional file format tests › should throw an error for unsupported file mime type in user message content + #[test] + fn should_throw_an_error_for_unsupported_file_mime_type_in_user_message_content() { + let error = convert(json!([{ + "role": "user", + "content": [{ + "type": "file", + "data": { "type": "data", "data": "base64data" }, + "mediaType": "application/rtf", + }], + }])) + .unwrap_err(); + + assert_eq!( + unsupported_functionality(error), + "file mime type: application/rtf" + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › additional file format tests › should handle xlsx files correctly + #[test] + fn should_handle_xlsx_files_correctly() { + let result = convert(json!([{ + "role": "user", + "content": [{ + "type": "file", + "data": { "type": "data", "data": "base64data" }, + "mediaType": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + }], + }])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [{ + "content": [{ + "document": { + "format": "xlsx", + "name": "document-1", + "source": { "bytes": "base64data" }, + }, + }], + "role": "user", + }], + "system": [], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › additional file format tests › should handle docx files correctly + #[test] + fn should_handle_docx_files_correctly() { + let result = convert(json!([{ + "role": "user", + "content": [{ + "type": "file", + "data": { "type": "data", "data": "base64data" }, + "mediaType": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + }], + }])) + .unwrap(); + + assert_eq!( + result, + json!({ + "messages": [{ + "content": [{ + "document": { + "format": "docx", + "name": "document-1", + "source": { "bytes": "base64data" }, + }, + }], + "role": "user", + }], + "system": [], + }) + ); + } +} + +mod mistral_tool_call_id_normalization { + use super::*; + + fn calculator_result(tool_call_id: &str) -> Value { + json!([{ + "role": "tool", + "content": [{ + "type": "tool-result", + "toolCallId": tool_call_id, + "toolName": "calculator", + "output": { "type": "text", "value": "The result is 42" }, + }], + }]) + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › Mistral tool call ID normalization › should normalize tool call IDs in tool results when isMistral is true + #[test] + fn should_normalize_tool_call_ids_in_tool_results_when_is_mistral_is_true() { + let result = + convert_with(calculator_result("tooluse_bpe71yCfRu2b5i-nKGDr5g"), true).unwrap(); + + assert_eq!( + result["messages"][0], + json!({ + "role": "user", + "content": [{ + "toolResult": { + "toolUseId": "8eHypBDcw", + "content": [{ "text": "The result is 42" }], + }, + }], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › Mistral tool call ID normalization › should normalize tool call IDs in tool calls when isMistral is true + #[test] + fn should_normalize_tool_call_ids_in_tool_calls_when_is_mistral_is_true() { + let result = convert_with( + json!([{ + "role": "assistant", + "content": [{ + "type": "tool-call", + "toolCallId": "tooluse_xyz123ABC456-def", + "toolName": "test-tool", + "input": { "query": "test" }, + }], + }]), + true, + ) + .unwrap(); + + assert_eq!( + result["messages"][0], + json!({ + "role": "assistant", + "content": [{ + "toolUse": { + "toolUseId": "De61BW1Dz", + "name": "test-tool", + "input": { "query": "test" }, + }, + }], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › Mistral tool call ID normalization › should not normalize tool call IDs when isMistral is false + #[test] + fn should_not_normalize_tool_call_ids_when_is_mistral_is_false() { + let original_id = "tooluse_bpe71yCfRu2b5i-nKGDr5g"; + let result = convert_with(calculator_result(original_id), false).unwrap(); + + assert_eq!( + result["messages"][0], + json!({ + "role": "user", + "content": [{ + "toolResult": { + "toolUseId": original_id, + "content": [{ "text": "The result is 42" }], + }, + }], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › Mistral tool call ID normalization › should default to not normalizing when isMistral is not provided + #[test] + fn should_default_to_not_normalizing_when_is_mistral_is_not_provided() { + let original_id = "tooluse_bpe71yCfRu2b5i-nKGDr5g"; + let result = convert(calculator_result(original_id)).unwrap(); + + assert_eq!( + result["messages"][0], + json!({ + "role": "user", + "content": [{ + "toolResult": { + "toolUseId": original_id, + "content": [{ "text": "The result is 42" }], + }, + }], + }) + ); + } +} + +mod top_level_only_media_type_resolution { + use super::*; + + /// base64 of `[0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]`. + const PNG_BYTES: &str = "iVBORw0KGgo="; + /// base64 of `[0x25, 0x50, 0x44, 0x46, 0x2d, 0x31, 0x2e, 0x34]`. + const PDF_BYTES: &str = "JVBERi0xLjQ="; + /// base64 of `[0x00, 0x00, 0x00, 0x18, 0x66, 0x74, 0x79, 0x70]`. + const MP4_BYTES: &str = "AAAAGGZ0eXA="; + + fn single_file(data: Value, media_type: &str) -> Value { + json!([{ + "role": "user", + "content": [{ "type": "file", "data": data, "mediaType": media_type }], + }]) + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › top-level-only mediaType resolution › should detect subtype from inline bytes when mediaType is top-level-only (image) + #[test] + fn should_detect_subtype_from_inline_bytes_when_media_type_is_top_level_only_image() { + let result = convert(single_file( + json!({ "type": "data", "data": PNG_BYTES }), + "image", + )) + .unwrap(); + + assert_eq!( + result["messages"][0], + json!({ + "role": "user", + "content": [{ "image": { "format": "png", "source": { "bytes": "iVBORw0KGgo=" } } }], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › top-level-only mediaType resolution › should detect subtype from inline bytes when mediaType is top-level-only (video) + #[test] + fn should_detect_subtype_from_inline_bytes_when_media_type_is_top_level_only_video() { + let result = convert(single_file( + json!({ "type": "data", "data": MP4_BYTES }), + "video", + )) + .unwrap(); + + assert_eq!( + result["messages"][0], + json!({ + "role": "user", + "content": [{ "video": { "format": "mp4", "source": { "bytes": "AAAAGGZ0eXA=" } } }], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › top-level-only mediaType resolution › should detect subtype from inline bytes when mediaType is top-level-only (application/pdf) + #[test] + fn should_detect_subtype_from_inline_bytes_when_media_type_is_top_level_only_pdf() { + let result = convert(single_file( + json!({ "type": "data", "data": PDF_BYTES }), + "application", + )) + .unwrap(); + + assert_eq!( + result["messages"][0], + json!({ + "role": "user", + "content": [{ + "document": { + "format": "pdf", + "name": "document-1", + "source": { "bytes": "JVBERi0xLjQ=" }, + }, + }], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › top-level-only mediaType resolution › should route to document slot for non-image top-level type via getTopLevelMediaType + #[test] + fn should_route_to_document_slot_for_non_image_top_level_type() { + let result = convert(single_file( + json!({ "type": "data", "data": "base64data" }), + "text/plain", + )) + .unwrap(); + + assert_eq!( + result["messages"][0], + json!({ + "role": "user", + "content": [{ + "document": { + "format": "txt", + "name": "document-1", + "source": { "bytes": "base64data" }, + }, + }], + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › top-level-only mediaType resolution › should throw UnsupportedFunctionalityError for URL data (File URL) + #[test] + fn should_throw_unsupported_functionality_error_for_url_data_file_url() { + let error = convert(single_file( + json!({ "type": "url", "url": "https://example.com/image.png" }), + "image", + )) + .unwrap_err(); + + assert_eq!(unsupported_functionality(error), "File URL data"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › top-level-only mediaType resolution › should throw UnsupportedFunctionalityError for unsupported full image mediaType + #[test] + fn should_throw_unsupported_functionality_error_for_unsupported_full_image_media_type() { + let error = convert(single_file( + json!({ "type": "data", "data": "base64data" }), + "image/avif", + )) + .unwrap_err(); + + assert_eq!( + unsupported_functionality(error), + "image mime type: image/avif" + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › top-level-only mediaType resolution › should throw UnsupportedFunctionalityError when top-level-only bytes cannot be detected + #[test] + fn should_throw_unsupported_functionality_error_when_top_level_only_bytes_cannot_be_detected() { + // base64 of `[0x00, 0x01, 0x02]`. + let error = convert(single_file( + json!({ "type": "data", "data": "AAEC" }), + "image", + )) + .unwrap_err(); + + let message = error.to_string(); + assert!( + regex::Regex::new(r#"media type "image".*could not be auto-detected"#) + .unwrap() + .is_match(&message), + "{message}" + ); + } +} diff --git a/aimux-providers/src/bedrock/tests/convert_usage.rs b/aimux-providers/src/bedrock/tests/convert_usage.rs new file mode 100644 index 000000000..071f18b22 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/convert_usage.rs @@ -0,0 +1,142 @@ +//! `convertAmazonBedrockUsage(usage)` is `convert_usage`: the upstream usage +//! literal is parsed as the Converse `usage` object (`BedrockUsage`) and the +//! result is compared as JSON (keys whose value is `undefined` omitted). The +//! two `raw`-only cases are ported in +//! `aimux-providers/tests/bedrock_remaining_test.rs`. + +use serde_json::{Value, json}; + +use crate::bedrock::convert::{BedrockUsage, convert_usage}; + +fn convert(usage: Option) -> Value { + let usage: Option = + usage.map(|usage| serde_json::from_value(usage).expect("valid Bedrock usage")); + serde_json::to_value(convert_usage(usage.as_ref())).unwrap() +} + +mod convert_amazon_bedrock_usage { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/convert-amazon-bedrock-usage.test.ts › convertAmazonBedrockUsage › should convert basic usage without cache tokens + #[test] + fn should_convert_basic_usage_without_cache_tokens() { + let result = convert(Some(json!({ "inputTokens": 100, "outputTokens": 50 }))); + + assert_eq!( + result, + json!({ + "inputTokens": { "total": 100, "noCache": 100, "cacheRead": 0, "cacheWrite": 0 }, + "outputTokens": { "total": 50, "text": 50 }, + "raw": { "inputTokens": 100, "outputTokens": 50 }, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-amazon-bedrock-usage.test.ts › convertAmazonBedrockUsage › should convert usage with cache read tokens + #[test] + fn should_convert_usage_with_cache_read_tokens() { + let result = convert(Some(json!({ + "inputTokens": 100, + "outputTokens": 50, + "cacheReadInputTokens": 80, + }))); + + assert_eq!( + result, + json!({ + // 100 + 80 + 0 + "inputTokens": { "total": 180, "noCache": 100, "cacheRead": 80, "cacheWrite": 0 }, + "outputTokens": { "total": 50, "text": 50 }, + "raw": { "inputTokens": 100, "outputTokens": 50, "cacheReadInputTokens": 80 }, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-amazon-bedrock-usage.test.ts › convertAmazonBedrockUsage › should convert usage with cache write tokens + #[test] + fn should_convert_usage_with_cache_write_tokens() { + let result = convert(Some(json!({ + "inputTokens": 100, + "outputTokens": 50, + "cacheWriteInputTokens": 60, + }))); + + assert_eq!( + result, + json!({ + // 100 + 0 + 60 + "inputTokens": { "total": 160, "noCache": 100, "cacheRead": 0, "cacheWrite": 60 }, + "outputTokens": { "total": 50, "text": 50 }, + "raw": { "inputTokens": 100, "outputTokens": 50, "cacheWriteInputTokens": 60 }, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-amazon-bedrock-usage.test.ts › convertAmazonBedrockUsage › should convert usage with both cache read and write tokens + #[test] + fn should_convert_usage_with_both_cache_read_and_write_tokens() { + let result = convert(Some(json!({ + "inputTokens": 100, + "outputTokens": 50, + "cacheReadInputTokens": 80, + "cacheWriteInputTokens": 60, + }))); + + assert_eq!( + result, + json!({ + // 100 + 80 + 60 + "inputTokens": { "total": 240, "noCache": 100, "cacheRead": 80, "cacheWrite": 60 }, + "outputTokens": { "total": 50, "text": 50 }, + "raw": { + "inputTokens": 100, + "outputTokens": 50, + "cacheReadInputTokens": 80, + "cacheWriteInputTokens": 60, + }, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-amazon-bedrock-usage.test.ts › convertAmazonBedrockUsage › should handle null cache tokens + #[test] + fn should_handle_null_cache_tokens() { + let result = convert(Some(json!({ + "inputTokens": 100, + "outputTokens": 50, + "cacheReadInputTokens": null, + "cacheWriteInputTokens": null, + }))); + + assert_eq!( + result, + json!({ + "inputTokens": { "total": 100, "noCache": 100, "cacheRead": 0, "cacheWrite": 0 }, + "outputTokens": { "total": 50, "text": 50 }, + "raw": { + "inputTokens": 100, + "outputTokens": 50, + "cacheReadInputTokens": null, + "cacheWriteInputTokens": null, + }, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-amazon-bedrock-usage.test.ts › convertAmazonBedrockUsage › should handle null usage + #[test] + fn should_handle_null_usage() { + // `null` and `undefined` are both `None` here. + let result = convert(None); + + assert_eq!(result, json!({ "inputTokens": {}, "outputTokens": {} })); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-amazon-bedrock-usage.test.ts › convertAmazonBedrockUsage › should handle undefined usage + #[test] + fn should_handle_undefined_usage() { + let result = convert(None); + + assert_eq!(result, json!({ "inputTokens": {}, "outputTokens": {} })); + } +} diff --git a/aimux-providers/src/bedrock/tests/embedding_model.rs b/aimux-providers/src/bedrock/tests/embedding_model.rs new file mode 100644 index 000000000..7e917fc7c --- /dev/null +++ b/aimux-providers/src/bedrock/tests/embedding_model.rs @@ -0,0 +1,514 @@ +//! Ports of `@ai-sdk/amazon-bedrock` 5.0.100 +//! `src/amazon-bedrock-embedding-model.test.ts`. +//! +//! Upstream builds `AmazonBedrockEmbeddingModel` directly with a base URL, +//! plain config headers and `fetch: injectFetchHeaders(..)`, so no request is +//! signed. The port does the same: `BedrockEmbeddingModel::from_config` with a +//! fixed endpoint, the config headers and the +//! [`inject_fetch_headers`](super::support::inject_fetch_headers) transport over the +//! mock transport (which stands in for the test server). +//! +//! Two cases stay in `aimux-providers/tests/bedrock_embedding_test.rs`, whose +//! tests already cover them: "should expose model-specific max embeddings per +//! call" and "should handle single input value and extract usage". + +use std::collections::HashMap; + +use async_trait::async_trait; +use serde_json::{Value, json}; + +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::error::AiMuxError; +use aimux_core::shared::SharedProviderOptions; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt}; + +use crate::bedrock::{AmazonBedrockEmbeddingModelFamily, BedrockEmbeddingModel}; +use crate::shared::{Endpoint, EndpointConfig, ResolveEndpoint}; +use crate::test_support::mock_fetch::{Canned, MockFetch}; + +use super::support::inject_fetch_headers; + +const BASE_URL: &str = "https://bedrock-runtime.us-east-1.amazonaws.com"; + +const TEST_VALUES: [&str; 2] = ["sunny day at the beach", "rainy day in the city"]; + +const COHERE_V4_APPLICATION_PROFILE_ARN: &str = + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/qibm5eutlkcy"; + +fn mock_embeddings() -> Vec> { + vec![ + vec![-0.09, 0.05, -0.02, 0.01, 0.04], + vec![-0.08, 0.06, -0.03, 0.02, 0.03], + ] +} + +/// The `/model/{id}/invoke` URL of `model_id` (`encodeURIComponent`). +fn invoke_url(model_id: &str) -> String { + format!( + "{BASE_URL}/model/{}/invoke", + model_id.replace(':', "%3A").replace('/', "%2F") + ) +} + +/// A fixed base URL and fixed headers: the `{ baseUrl, headers }` config of +/// the upstream model. +struct FixedEndpoint { + base_url: String, + headers: HeaderMapOpt, +} + +#[async_trait] +impl ResolveEndpoint for FixedEndpoint { + async fn resolve(&self) -> Result { + Ok(Endpoint { + base_url: self.base_url.clone(), + headers: self.headers.clone(), + }) + } +} + +/// `new AmazonBedrockEmbeddingModel(model_id, { baseUrl, headers, fetch })`. +fn embedding_model( + model_id: &str, + headers: &[(&str, &str)], + fetch: FetchFunction, +) -> BedrockEmbeddingModel { + BedrockEmbeddingModel::from_config( + model_id.to_string(), + EndpointConfig::new( + "amazon-bedrock".to_string(), + FixedEndpoint { + base_url: BASE_URL.to_string(), + headers: headers + .iter() + .map(|(name, value)| ((*name).to_string(), Some((*value).to_string()))) + .collect(), + }, + Some(fetch), + ), + ) +} + +/// The upstream `mockConfigHeaders`. +const MOCK_CONFIG_HEADERS: &[(&str, &str)] = &[ + ("config-header", "config-value"), + ("shared-header", "config-shared"), +]; + +/// `fakeFetchWithAuth` over `mock`. +fn fake_fetch_with_auth(mock: &std::sync::Arc) -> FetchFunction { + inject_fetch_headers(&[("x-amz-auth", "test-auth")], mock.transport()) +} + +/// A binary JSON response with extra headers. +fn json_response(body: &Value, headers: &[(&str, &str)]) -> Canned { + let mut canned = Canned::json(body); + canned.headers.extend( + headers + .iter() + .map(|(name, value)| ((*name).to_string(), (*value).to_string())), + ); + canned +} + +/// The Titan response of the upstream `embedUrl`. +fn titan_response() -> Canned { + json_response( + &json!({ "embedding": mock_embeddings()[0], "inputTextTokenCount": 8 }), + &[], + ) +} + +fn options(values: &[&str]) -> EmbeddingCallOptions { + EmbeddingCallOptions { + values: values.iter().map(|value| (*value).to_string()).collect(), + ..EmbeddingCallOptions::new("") + } +} + +fn provider_options(namespace: &str, options: Value) -> SharedProviderOptions { + aimux_core::shared::provider_namespace(namespace, options).expect("options are an object") +} + +mod do_embed { + use super::*; + + // "should expose model-specific max embeddings per call": the marked test + // in aimux-providers/tests/bedrock_embedding_test.rs. + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-embedding-model.test.ts › doEmbed › should handle single input value and return embeddings + #[tokio::test] + async fn should_handle_single_input_value_and_return_embeddings() { + let mock = MockFetch::new(vec![titan_response()]); + let model = embedding_model( + "amazon.titan-embed-text-v2:0", + MOCK_CONFIG_HEADERS, + fake_fetch_with_auth(&mock), + ); + + let result = model + .do_embed(&options(&[TEST_VALUES[0]])) + .await + .expect("embedding succeeds"); + + assert_eq!(result.embeddings.len(), 1); + assert_eq!(result.embeddings[0], vec![-0.09, 0.05, -0.02, 0.01, 0.04]); + + let seen = mock.seen(); + assert_eq!(seen[0].url, invoke_url("amazon.titan-embed-text-v2:0")); + assert_eq!(seen[0].json_body(), json!({ "inputText": TEST_VALUES[0] })); + } + + // "should handle single input value and extract usage": the marked test + // in aimux-providers/tests/bedrock_embedding_test.rs. + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-embedding-model.test.ts › doEmbed › should support Cohere embedding models + #[tokio::test] + async fn should_support_cohere_embedding_models() { + let mock = MockFetch::new(vec![json_response( + &json!({ "embeddings": [mock_embeddings()[0]] }), + &[("x-amzn-bedrock-input-token-count", "6")], + )]); + let cohere_model = embedding_model( + "cohere.embed-english-v3", + MOCK_CONFIG_HEADERS, + fake_fetch_with_auth(&mock), + ); + + let result = cohere_model + .do_embed(&options(&[TEST_VALUES[0]])) + .await + .expect("embedding succeeds"); + + assert_eq!(result.embeddings.len(), 1); + assert_eq!(result.embeddings[0], vec![-0.09, 0.05, -0.02, 0.01, 0.04]); + assert_eq!(result.usage.expect("usage").tokens, 6); + + let seen = mock.seen(); + assert_eq!(seen[0].url, invoke_url("cohere.embed-english-v3")); + assert_eq!( + seen[0].json_body(), + json!({ "input_type": "search_query", "texts": [TEST_VALUES[0]] }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-embedding-model.test.ts › doEmbed › should support Cohere v4 embedding models + #[tokio::test] + async fn should_support_cohere_v4_embedding_models() { + let mock = MockFetch::new(vec![json_response( + &json!({ "embeddings": { "float": [mock_embeddings()[0]] } }), + &[("x-amzn-bedrock-input-token-count", "6")], + )]); + let cohere_v4_model = embedding_model( + "cohere.embed-v4:0", + MOCK_CONFIG_HEADERS, + fake_fetch_with_auth(&mock), + ); + + let result = cohere_v4_model + .do_embed(&options(&[TEST_VALUES[0]])) + .await + .expect("embedding succeeds"); + + assert_eq!(result.embeddings.len(), 1); + assert_eq!(result.embeddings[0], vec![-0.09, 0.05, -0.02, 0.01, 0.04]); + assert_eq!(result.usage.expect("usage").tokens, 6); + + let seen = mock.seen(); + assert_eq!(seen[0].url, invoke_url("cohere.embed-v4:0")); + assert_eq!( + seen[0].json_body(), + json!({ "input_type": "search_query", "texts": [TEST_VALUES[0]] }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-embedding-model.test.ts › doEmbed › should send multiple values for Cohere embedding models + #[tokio::test] + async fn should_send_multiple_values_for_cohere_embedding_models() { + let mock = MockFetch::new(vec![json_response( + &json!({ "embeddings": { "float": mock_embeddings() } }), + &[("x-amzn-bedrock-input-token-count", "12")], + )]); + let cohere_v4_model = embedding_model( + "cohere.embed-v4:0", + MOCK_CONFIG_HEADERS, + fake_fetch_with_auth(&mock), + ); + + let result = cohere_v4_model + .do_embed(&options(&TEST_VALUES)) + .await + .expect("embedding succeeds"); + + assert_eq!(result.embeddings, mock_embeddings()); + assert_eq!(result.usage.expect("usage").tokens, 12); + + let seen = mock.seen(); + assert_eq!(seen[0].url, invoke_url("cohere.embed-v4:0")); + assert_eq!( + seen[0].json_body(), + json!({ "input_type": "search_query", "texts": TEST_VALUES }) + ); + } + + // Not ported (sub-feature): "should return NaN usage for Cohere models + // when the token count header is absent". `EmbeddingUsage::tokens` is a + // `u32`, which has no NaN (or other "unknown") value; aimux reports 0. + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-embedding-model.test.ts › doEmbed › should support Cohere models behind cross-region inference profile ids + #[tokio::test] + async fn should_support_cohere_models_behind_cross_region_inference_profile_ids() { + let mock = MockFetch::new(vec![json_response( + &json!({ "embeddings": { "float": [mock_embeddings()[0]] } }), + &[("x-amzn-bedrock-input-token-count", "6")], + )]); + let cohere_v4_us_profile_model = embedding_model( + "us.cohere.embed-v4:0", + MOCK_CONFIG_HEADERS, + fake_fetch_with_auth(&mock), + ); + + let result = cohere_v4_us_profile_model + .do_embed(&options(&[TEST_VALUES[0]])) + .await + .expect("embedding succeeds"); + + assert_eq!(result.embeddings.len(), 1); + assert_eq!(result.usage.expect("usage").tokens, 6); + + let seen = mock.seen(); + assert_eq!(seen[0].url, invoke_url("us.cohere.embed-v4:0")); + assert_eq!( + seen[0].json_body(), + json!({ "input_type": "search_query", "texts": [TEST_VALUES[0]] }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-embedding-model.test.ts › doEmbed › should support Cohere models behind application inference profile ARNs + #[tokio::test] + async fn should_support_cohere_models_behind_application_inference_profile_arns() { + let mock = MockFetch::new(vec![json_response( + &json!({ "embeddings": { "float": mock_embeddings() } }), + &[("x-amzn-bedrock-input-token-count", "12")], + )]); + let cohere_v4_application_profile_model = embedding_model( + COHERE_V4_APPLICATION_PROFILE_ARN, + MOCK_CONFIG_HEADERS, + fake_fetch_with_auth(&mock), + ) + .with_model_family(Some(AmazonBedrockEmbeddingModelFamily::Cohere)); + + let result = cohere_v4_application_profile_model + .do_embed(&options(&[TEST_VALUES[0]])) + .await + .expect("embedding succeeds"); + + assert_eq!(result.embeddings, mock_embeddings()); + assert_eq!(result.usage.expect("usage").tokens, 12); + + let seen = mock.seen(); + assert_eq!(seen[0].url, invoke_url(COHERE_V4_APPLICATION_PROFILE_ARN)); + assert_eq!( + seen[0].json_body(), + json!({ "input_type": "search_query", "texts": [TEST_VALUES[0]] }) + ); + assert_eq!( + cohere_v4_application_profile_model.max_embeddings_per_call(), + Some(96) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-embedding-model.test.ts › doEmbed › should pass outputDimension for Cohere v4 embedding models + #[tokio::test] + async fn should_pass_output_dimension_for_cohere_v4_embedding_models() { + let mock = MockFetch::new(vec![json_response( + &json!({ "embeddings": { "float": [mock_embeddings()[0]] } }), + &[("x-amzn-bedrock-input-token-count", "6")], + )]); + let cohere_v4_model = embedding_model( + "cohere.embed-v4:0", + MOCK_CONFIG_HEADERS, + fake_fetch_with_auth(&mock), + ); + + let result = cohere_v4_model + .do_embed(&EmbeddingCallOptions { + provider_options: Some(provider_options( + "bedrock", + json!({ "outputDimension": 256 }), + )), + ..options(&[TEST_VALUES[0]]) + }) + .await + .expect("embedding succeeds"); + + assert_eq!(result.embeddings.len(), 1); + assert_eq!(result.embeddings[0], vec![-0.09, 0.05, -0.02, 0.01, 0.04]); + + let seen = mock.seen(); + assert_eq!(seen[0].url, invoke_url("cohere.embed-v4:0")); + assert_eq!( + seen[0].json_body(), + json!({ + "input_type": "search_query", + "texts": [TEST_VALUES[0]], + "output_dimension": 256, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-embedding-model.test.ts › doEmbed › should properly combine headers from all sources + #[tokio::test] + async fn should_properly_combine_headers_from_all_sources() { + let mock = MockFetch::new(vec![titan_response()]); + let model_with_headers = embedding_model( + "amazon.titan-embed-text-v2:0", + &[ + ("model-header", "model-value"), + ("shared-header", "model-shared"), + ], + inject_fetch_headers( + &[ + ("signed-header", "signed-value"), + ("authorization", "AWS4-HMAC-SHA256..."), + ], + mock.transport(), + ), + ); + + model_with_headers + .do_embed(&EmbeddingCallOptions { + headers: Some(HashMap::from([ + ("options-header".to_string(), "options-value".to_string()), + ("shared-header".to_string(), "options-shared".to_string()), + ])), + ..options(&[TEST_VALUES[0]]) + }) + .await + .expect("embedding succeeds"); + + let request_headers = &mock.seen()[0].headers; + assert_eq!(request_headers["options-header"], "options-value"); + assert_eq!(request_headers["model-header"], "model-value"); + assert_eq!(request_headers["signed-header"], "signed-value"); + assert_eq!(request_headers["authorization"], "AWS4-HMAC-SHA256..."); + assert_eq!(request_headers["shared-header"], "options-shared"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-embedding-model.test.ts › doEmbed › should work with partial headers + #[tokio::test] + async fn should_work_with_partial_headers() { + let mock = MockFetch::new(vec![titan_response()]); + let model_with_partial_headers = embedding_model( + "amazon.titan-embed-text-v2:0", + &[("model-header", "model-value")], + inject_fetch_headers( + &[ + ("signed-header", "signed-value"), + ("authorization", "AWS4-HMAC-SHA256..."), + ], + mock.transport(), + ), + ); + + model_with_partial_headers + .do_embed(&options(&[TEST_VALUES[0]])) + .await + .expect("embedding succeeds"); + + let request_headers = &mock.seen()[0].headers; + assert_eq!(request_headers["model-header"], "model-value"); + assert_eq!(request_headers["signed-header"], "signed-value"); + assert_eq!(request_headers["authorization"], "AWS4-HMAC-SHA256..."); + } +} + +mod should_support_nova_embeddings { + use super::*; + + const NOVA_MODEL_ID: &str = "amazon.nova-2-multimodal-embeddings-v1:0"; + + fn nova_response() -> Canned { + json_response( + &json!({ + "embeddings": [{ "embeddingType": "TEXT", "embedding": mock_embeddings()[0] }], + "inputTokenCount": 8, + }), + &[], + ) + } + + fn nova_model(mock: &std::sync::Arc) -> BedrockEmbeddingModel { + embedding_model( + NOVA_MODEL_ID, + &[("config-header", "config-value")], + fake_fetch_with_auth(mock), + ) + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-embedding-model.test.ts › should support Nova embeddings › should send SINGLE_EMBEDDING payload for Nova embeddings + #[tokio::test] + async fn should_send_single_embedding_payload_for_nova_embeddings() { + let mock = MockFetch::new(vec![nova_response()]); + + let result = nova_model(&mock) + .do_embed(&options(&[TEST_VALUES[0]])) + .await + .expect("embedding succeeds"); + + assert_eq!(result.embeddings[0], vec![-0.09, 0.05, -0.02, 0.01, 0.04]); + + let seen = mock.seen(); + assert_eq!(seen[0].url, invoke_url(NOVA_MODEL_ID)); + assert_eq!( + seen[0].json_body(), + json!({ + "taskType": "SINGLE_EMBEDDING", + "singleEmbeddingParams": { + "embeddingPurpose": "GENERIC_INDEX", + "embeddingDimension": 1024, + "text": { + "truncationMode": "END", + "value": TEST_VALUES[0], + }, + }, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-embedding-model.test.ts › should support Nova embeddings › should pass embeddingDimension for Nova embeddings + #[tokio::test] + async fn should_pass_embedding_dimension_for_nova_embeddings() { + let mock = MockFetch::new(vec![nova_response()]); + + let result = nova_model(&mock) + .do_embed(&EmbeddingCallOptions { + provider_options: Some(provider_options( + "bedrock", + json!({ "embeddingDimension": 256 }), + )), + ..options(&[TEST_VALUES[0]]) + }) + .await + .expect("embedding succeeds"); + + assert_eq!(result.embeddings[0], vec![-0.09, 0.05, -0.02, 0.01, 0.04]); + + let seen = mock.seen(); + assert_eq!(seen[0].url, invoke_url(NOVA_MODEL_ID)); + assert_eq!( + seen[0].json_body(), + json!({ + "taskType": "SINGLE_EMBEDDING", + "singleEmbeddingParams": { + "embeddingPurpose": "GENERIC_INDEX", + "embeddingDimension": 256, + "text": { + "truncationMode": "END", + "value": TEST_VALUES[0], + }, + }, + }) + ); + } +} diff --git a/aimux-providers/src/bedrock/tests/error.rs b/aimux-providers/src/bedrock/tests/error.rs new file mode 100644 index 000000000..ba01a2811 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/error.rs @@ -0,0 +1,69 @@ +//! Ports of `@ai-sdk/amazon-bedrock` 5.0.100 `src/amazon-bedrock-error.test.ts`. +//! +//! `amazonBedrockFailedResponseHandler` is `bedrock_failed_response_handler` +//! in `bedrock/mod.rs`. + +use bytes::Bytes; +use reqwest::StatusCode; +use reqwest::header::{CONTENT_TYPE, HeaderMap, HeaderValue}; +use serde_json::{Value, json}; +use url::Url; + +use aimux_core::error::AiMuxError; +use aimux_provider_utils::{FetchResponse, ResponseHandlerInput}; + +use crate::bedrock::bedrock_failed_response_handler; + +const URL: &str = "https://bedrock-runtime.us-east-1.amazonaws.com/model/test/invoke"; + +/// A `400 Bad Request` JSON response with `body`. +fn make_response(body: &Value) -> FetchResponse { + let mut headers = HeaderMap::new(); + headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); + FetchResponse::from_bytes( + StatusCode::BAD_REQUEST, + headers, + Url::parse(URL).expect("valid URL"), + Bytes::from(serde_json::to_vec(body).expect("JSON body")), + ) +} + +/// The message of the `APICallError` the handler builds for `body`. +async fn get_error_message(body: &Value) -> String { + let output = bedrock_failed_response_handler() + .handle(ResponseHandlerInput { + url: URL.to_string(), + request_body_values: json!({}), + response: make_response(body), + abort_signal: None, + max_json_response_bytes: None, + }) + .await + .expect("the handler produces an error value"); + match output.value { + AiMuxError::ApiCall(error) => error.message, + other => panic!("expected an APICallError, got {other:?}"), + } +} + +mod amazon_bedrock_failed_response_handler { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-error.test.ts › amazonBedrockFailedResponseHandler › preserves the provider message when the error type is omitted + #[tokio::test] + async fn preserves_the_provider_message_when_the_error_type_is_omitted() { + assert_eq!( + get_error_message(&json!({ "message": "boom" })).await, + "boom" + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-error.test.ts › amazonBedrockFailedResponseHandler › prefixes the provider message when the error type is present + #[tokio::test] + async fn prefixes_the_provider_message_when_the_error_type_is_present() { + assert_eq!( + get_error_message(&json!({ "type": "ValidationException", "message": "boom" })).await, + "ValidationException: boom" + ); + } +} diff --git a/aimux-providers/src/bedrock/tests/event_stream_decoder.rs b/aimux-providers/src/bedrock/tests/event_stream_decoder.rs new file mode 100644 index 000000000..46d93a679 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/event_stream_decoder.rs @@ -0,0 +1,109 @@ +//! Ports of `@ai-sdk/amazon-bedrock` 5.0.100 +//! `src/amazon-bedrock-event-stream-decoder.test.ts`. +//! +//! `createAmazonBedrockEventStreamDecoder(body, processEvent)` is +//! `event_stream::decode_stream(body)`: it yields each decoded frame (or the +//! decoding error) instead of handing it to a `processEvent` callback, so the +//! cases collect what it yields. The upstream "processEvent" that enqueues +//! `event.data` is a `map` over the yielded messages. + +use bytes::Bytes; +use futures::StreamExt; +use futures::stream::BoxStream; + +use aimux_core::error::AiMuxError; + +use crate::bedrock::event_stream::{EventStreamMessage, decode_stream, encode_message}; + +/// `createEvent(data)`: one `event` frame of event type `chunk`. +fn create_event(data: &str) -> Vec { + encode_message("event", "chunk", data) +} + +/// `createStream(chunks)`: a body that delivers `chunks` and ends. +fn create_stream(chunks: Vec>) -> BoxStream<'static, Result> { + futures::stream::iter(chunks.into_iter().map(|chunk| Ok(Bytes::from(chunk)))).boxed() +} + +/// Everything the decoder yields for `chunks`. +async fn decode_all(chunks: Vec>) -> Vec> { + decode_stream(create_stream(chunks)).collect().await +} + +/// The `data` of every decoded frame, failing on the first error. +fn data_of(items: Vec>) -> Vec { + items + .into_iter() + .map(|item| item.expect("the frame decodes").data) + .collect() +} + +mod create_amazon_bedrock_event_stream_decoder { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-event-stream-decoder.test.ts › createAmazonBedrockEventStreamDecoder › surfaces frame decode errors instead of stranding later frames + #[tokio::test] + async fn surfaces_frame_decode_errors_instead_of_stranding_later_frames() { + let mut corrupted_frame = create_event("corrupted"); + let last = corrupted_frame.len() - 1; + corrupted_frame[last] ^= 0xff; + + let items = decode_all(vec![corrupted_frame, create_event("later")]).await; + + // The stream rejects, and no frame reaches the event processing. + assert!( + items.iter().any(Result::is_err), + "the decode error is surfaced: {items:?}" + ); + assert!( + items.iter().all(Result::is_err), + "no frame is processed: {items:?}" + ); + } + + // Not ported (sub-feature): "surfaces processEvent errors". `decode_stream` + // takes no `processEvent` callback; its consumer is the chat stream's + // `Transformer`, which reports problems as stream parts and has no + // throwing callback whose error the decoder could surface. + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-event-stream-decoder.test.ts › createAmazonBedrockEventStreamDecoder › processes frames split across chunks + #[tokio::test] + async fn processes_frames_split_across_chunks() { + let frame = create_event("data"); + let midpoint = frame.len() / 2; + + let items = decode_all(vec![frame[..midpoint].to_vec(), frame[midpoint..].to_vec()]).await; + + assert_eq!(data_of(items), vec!["data".to_string()]); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-event-stream-decoder.test.ts › createAmazonBedrockEventStreamDecoder › rejects when EOF leaves an incomplete frame + #[tokio::test] + async fn rejects_when_eof_leaves_an_incomplete_frame() { + let frame = create_event("data"); + let incomplete_frame = frame[..frame.len() - 1].to_vec(); + let incomplete_length = incomplete_frame.len(); + + let items = decode_all(vec![incomplete_frame]).await; + + let error = items + .into_iter() + .find_map(Result::err) + .expect("the stream rejects"); + let expected = format!( + "Incomplete Amazon Bedrock event-stream frame: {incomplete_length} buffered bytes remain at end of stream." + ); + assert!( + error.to_string().contains(&expected), + "unexpected error: {error}" + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-event-stream-decoder.test.ts › createAmazonBedrockEventStreamDecoder › allows EOF at a complete frame boundary + #[tokio::test] + async fn allows_eof_at_a_complete_frame_boundary() { + let items = decode_all(vec![create_event("data")]).await; + + assert_eq!(data_of(items), vec!["data".to_string()]); + } +} diff --git a/aimux-providers/src/bedrock/tests/event_stream_response_handler.rs b/aimux-providers/src/bedrock/tests/event_stream_response_handler.rs new file mode 100644 index 000000000..8c266386d --- /dev/null +++ b/aimux-providers/src/bedrock/tests/event_stream_response_handler.rs @@ -0,0 +1,174 @@ +//! Ports of `@ai-sdk/amazon-bedrock` 5.0.100 +//! `src/amazon-bedrock-event-stream-response-handler.test.ts`. +//! +//! `createAmazonBedrockEventStreamResponseHandler(chunkSchema)` decodes the +//! event-stream body and parses each frame's JSON as +//! `{ [eventType]: payload }`. In aimux these steps are split: the response +//! handler of `BedrockModel::do_stream` hands the body on as bytes, +//! `event_stream::decode_stream` decodes the frames, and the chat stream's +//! `Transformer` parses the JSON (emitting the wrapped value as the `raw` +//! part, and a parse failure as an `error` part). So the cases run through +//! `do_stream` of a directly built `BedrockModel` (a fixed endpoint over the +//! mock transport) with raw chunks on. +//! +//! Upstream mocks `EventStreamCodec.decode`, so its frames are dummy bytes +//! with only a valid length prefix. The port sends real frames (encoded with +//! `event_stream::encode_message`) carrying the decoded messages upstream's +//! mock returns. + +use async_trait::async_trait; +use futures::StreamExt; +use serde_json::json; + +use aimux_core::error::AiMuxError; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::{LanguageModelMessage, TextPart, UserPart}; +use aimux_core::options::CallOptions; +use aimux_core::stream_part::StreamPart; +use aimux_provider_utils::HeaderMapOpt; + +use crate::bedrock::BedrockModel; +use crate::bedrock::event_stream::encode_message; +use crate::shared::{Endpoint, EndpointConfig, ResolveEndpoint}; +use crate::test_support::mock_fetch::{Canned, MockFetch}; + +/// A fixed base URL and no headers. +struct FixedEndpoint; + +#[async_trait] +impl ResolveEndpoint for FixedEndpoint { + async fn resolve(&self) -> Result { + Ok(Endpoint { + base_url: "https://bedrock-runtime.us-east-1.amazonaws.com".to_string(), + headers: HeaderMapOpt::new(), + }) + } +} + +/// A `200` event-stream response with `body`. +fn event_stream_response(body: Vec) -> Canned { + Canned { + status: 200, + headers: vec![( + "content-type".to_string(), + "application/vnd.amazon.eventstream".to_string(), + )], + body, + } +} + +/// Everything the stream yields when the response body is `body`. +async fn stream_parts(body: Vec) -> Vec> { + let mock = MockFetch::new(vec![event_stream_response(body)]); + let model = BedrockModel::from_config( + "anthropic.claude-3-haiku-20240307-v1:0".to_string(), + EndpointConfig::new( + "amazon-bedrock".to_string(), + FixedEndpoint, + Some(mock.transport()), + ), + ); + let mut options = CallOptions::new(vec![LanguageModelMessage::User { + content: vec![UserPart::Text(TextPart { + text: "Hello".into(), + provider_options: None, + })], + provider_options: None, + }]); + options.include_raw_chunks = Some(true); + let result = model.do_stream(&options).await.expect("the stream starts"); + result.stream.collect().await +} + +/// A frame `{ ':message-type': 'event', ':event-type': 'chunk' }` with `body`. +fn chunk_frame(body: &str) -> Vec { + encode_message("event", "chunk", body) +} + +mod create_event_source_response_handler { + use super::*; + + // Not ported (ts-only): "throws EmptyResponseBodyError when response body + // is null". A `FetchResponse` always has a body stream (an empty one for + // no content); there is no null body to reject. + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-event-stream-response-handler.test.ts › createEventSourceResponseHandler › successfully processes valid event stream data + /// + /// The `rawValue` of the parse result is the `raw` part. aimux has no + /// schema-validated `value` to compare: it does not validate chunks + /// against a schema. + #[tokio::test] + async fn successfully_processes_valid_event_stream_data() { + let parts = stream_parts(chunk_frame(r#"{"content":"test message"}"#)).await; + + let raw_values: Vec<_> = parts + .iter() + .filter_map(|part| match part { + Ok(StreamPart::Raw { raw_value }) => Some(raw_value.clone()), + _ => None, + }) + .collect(); + assert_eq!( + raw_values, + vec![json!({ "chunk": { "content": "test message" } })], + "{parts:?}" + ); + assert!(parts.iter().all(Result::is_ok), "{parts:?}"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-event-stream-response-handler.test.ts › createEventSourceResponseHandler › handles invalid JSON data + /// + /// The failed parse result (`success: false` with an `error`) is an + /// `error` part. + #[tokio::test] + async fn handles_invalid_json_data() { + let parts = stream_parts(chunk_frame("invalid json")).await; + + assert!( + parts + .iter() + .any(|part| matches!(part, Ok(StreamPart::Error { .. }))), + "{parts:?}" + ); + assert!( + !parts + .iter() + .any(|part| matches!(part, Ok(StreamPart::Raw { .. }))), + "{parts:?}" + ); + } + + // Not ported (sub-feature): "handles schema validation failures". The + // response handler takes no chunk schema (upstream's `chunkSchema` + // parameter): the chat stream reads the event payloads leniently, so a + // payload that does not match a schema is not reported. + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-event-stream-response-handler.test.ts › createEventSourceResponseHandler › surfaces event stream decoding failures + /// + /// Upstream's mocked codec throws on `createFrame([13, 14])`. The real + /// decoder rejects the same bytes (a six-byte frame is too short to be + /// one), and the stream ends with that error. + #[tokio::test] + async fn surfaces_event_stream_decoding_failures() { + let frame = vec![0, 0, 0, 6, 13, 14]; + + let parts = stream_parts(frame).await; + + let last = parts.last().expect("the stream yields"); + let error = last + .as_ref() + .expect_err("the stream ends with the decode error"); + assert!( + error + .to_string() + .contains("Invalid Amazon Bedrock event-stream frame"), + "unexpected error: {error}" + ); + assert!( + !parts + .iter() + .any(|part| matches!(part, Ok(StreamPart::Raw { .. }))), + "{parts:?}" + ); + } +} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-guard-content-intervened.json b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-guard-content-intervened.json new file mode 100644 index 000000000..79cb53dd0 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-guard-content-intervened.json @@ -0,0 +1,43 @@ +{ + "metrics": { + "latencyMs": 412 + }, + "output": { + "message": { + "content": [ + { + "text": "Sorry, the model cannot answer this question." + } + ], + "role": "assistant" + } + }, + "stopReason": "guardrail_intervened", + "trace": { + "guardrail": { + "inputAssessment": { + "1abcd2ef34gh": { + "topicPolicy": { + "topics": [ + { + "action": "BLOCKED", + "name": "Restricted topic", + "type": "DENY" + } + ] + } + } + } + } + }, + "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 18, + "outputTokens": 12, + "serverToolUsage": {}, + "totalTokens": 30 + } +} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-only-text-first.1.chunks.txt b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-only-text-first.1.chunks.txt new file mode 100644 index 000000000..a7c903ef8 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-only-text-first.1.chunks.txt @@ -0,0 +1,7 @@ +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"Here is your data:"}}} +{"contentBlockStop":{"contentBlockIndex":0}} +{"contentBlockStart":{"contentBlockIndex":1,"start":{"toolUse":{"toolUseId":"json-tool-1","name":"json"}}}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"toolUse":{"input":"{\"name\":\"John\"}"}}}} +{"contentBlockStop":{"contentBlockIndex":1}} +{"metadata":{"usage":{"inputTokens":100,"outputTokens":20,"totalTokens":120}}} +{"messageStop":{"stopReason":"tool_use"}} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-other-tool.1.json b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-other-tool.1.json new file mode 100644 index 000000000..d948f3713 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-other-tool.1.json @@ -0,0 +1,24 @@ +{ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "toolu_01PQjhxo3eirCdKNvCJrKc8f", + "name": "get-weather", + "input": { + "location": "San Francisco" + } + } + } + ] + } + }, + "usage": { + "inputTokens": 843, + "outputTokens": 28, + "totalTokens": 871 + }, + "stopReason": "tool_use" +} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool-text-then-weather-then-json.1.chunks.txt b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool-text-then-weather-then-json.1.chunks.txt new file mode 100644 index 000000000..9d4f3eb8f --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool-text-then-weather-then-json.1.chunks.txt @@ -0,0 +1,12 @@ +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"2 + 2 equals 4. Now let me check"}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":" the weather for you."}}} +{"contentBlockStop":{"contentBlockIndex":0}} +{"contentBlockStart":{"contentBlockIndex":1,"start":{"toolUse":{"toolUseId":"weather-tool-1","name":"weather"}}}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"toolUse":{"input":"{\"location\":"}}}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"toolUse":{"input":"\"San Francisco\"}"}}}} +{"contentBlockStop":{"contentBlockIndex":1}} +{"contentBlockStart":{"contentBlockIndex":2,"start":{"toolUse":{"toolUseId":"weather-tool-2","name":"weather"}}}} +{"contentBlockDelta":{"contentBlockIndex":2,"delta":{"toolUse":{"input":"{\"location\":\"London\"}"}}}} +{"contentBlockStop":{"contentBlockIndex":2}} +{"metadata":{"usage":{"inputTokens":500,"outputTokens":100,"totalTokens":600}}} +{"messageStop":{"stopReason":"tool_use"}} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool-with-answer.1.json b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool-with-answer.1.json new file mode 100644 index 000000000..37b0f7623 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool-with-answer.1.json @@ -0,0 +1,29 @@ +{ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "2 + 2 equals 4." + }, + { + "toolUse": { + "toolUseId": "json-tool-id-123", + "name": "json", + "input": { + "firstName": "John", + "lastName": "Doe" + } + } + } + ] + } + }, + "usage": { + "inputTokens": 156, + "outputTokens": 35, + "totalTokens": 191 + }, + "stopReason": "tool_use" +} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.1.chunks.txt b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.1.chunks.txt new file mode 100644 index 000000000..d72e039dc --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.1.chunks.txt @@ -0,0 +1,4 @@ +{"contentBlockStart":{"contentBlockIndex":0,"start":{"toolUse":{"toolUseId":"json-tool-id","name":"json"}}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"toolUse":{"input":"{\"value\":\"test\"}"}}}} +{"contentBlockStop":{"contentBlockIndex":0}} +{"messageStop":{"stopReason":"tool_use"}} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.1.json b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.1.json new file mode 100644 index 000000000..4790fb283 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.1.json @@ -0,0 +1,35 @@ +{ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "json-tool-id", + "name": "json", + "input": { + "elements": [ + { + "location": "San Francisco", + "temperature": -5, + "condition": "snowy" + }, + { + "location": "London", + "temperature": 0, + "condition": "snowy" + } + ] + } + } + } + ] + } + }, + "usage": { + "inputTokens": 15, + "outputTokens": 42, + "totalTokens": 57 + }, + "stopReason": "tool_use" +} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.2.chunks.txt b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.2.chunks.txt new file mode 100644 index 000000000..62e7c735a --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.2.chunks.txt @@ -0,0 +1,6 @@ +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"Let me generate"}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":" that JSON for you."}}} +{"contentBlockStart":{"contentBlockIndex":1,"start":{"toolUse":{"toolUseId":"json-tool-id","name":"json"}}}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"toolUse":{"input":"{\"value\":\"test\"}"}}}} +{"contentBlockStop":{"contentBlockIndex":1}} +{"messageStop":{"stopReason":"tool_use"}} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.2.json b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.2.json new file mode 100644 index 000000000..37beb8029 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.2.json @@ -0,0 +1,28 @@ +{ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Let me generate that JSON for you." + }, + { + "toolUse": { + "toolUseId": "json-tool-id", + "name": "json", + "input": { + "value": "test data" + } + } + } + ] + } + }, + "usage": { + "inputTokens": 12, + "outputTokens": 28, + "totalTokens": 40 + }, + "stopReason": "tool_use" +} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.3.chunks.txt b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.3.chunks.txt new file mode 100644 index 000000000..d3c9f5f3e --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.3.chunks.txt @@ -0,0 +1,7 @@ +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"First text block."}}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"text":"Second text block."}}} +{"contentBlockDelta":{"contentBlockIndex":2,"delta":{"text":"Third block."}}} +{"contentBlockStart":{"contentBlockIndex":3,"start":{"toolUse":{"toolUseId":"json-tool-id","name":"json"}}}} +{"contentBlockDelta":{"contentBlockIndex":3,"delta":{"toolUse":{"input":"{\"result\":42}"}}}} +{"contentBlockStop":{"contentBlockIndex":3}} +{"messageStop":{"stopReason":"tool_use"}} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.3.json b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.3.json new file mode 100644 index 000000000..fc89ff08d --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-tool.3.json @@ -0,0 +1,36 @@ +{ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "First text block." + }, + { + "type": "text", + "text": "Second text block." + }, + { + "type": "text", + "text": "Third text block before JSON." + }, + { + "toolUse": { + "toolUseId": "json-tool-id", + "name": "json", + "input": { + "result": 42 + } + } + } + ] + } + }, + "usage": { + "inputTokens": 18, + "outputTokens": 35, + "totalTokens": 53 + }, + "stopReason": "tool_use" +} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-with-tool.1.chunks.txt b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-with-tool.1.chunks.txt new file mode 100644 index 000000000..a829f3267 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-with-tool.1.chunks.txt @@ -0,0 +1,9 @@ +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"Calling a tool first."}}} +{"contentBlockStart":{"contentBlockIndex":1,"start":{"toolUse":{"toolUseId":"tool-1","name":"get-weather"}}}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"toolUse":{"input":"{\"location\":\"SF\"}"}}}} +{"contentBlockStop":{"contentBlockIndex":1}} +{"contentBlockDelta":{"contentBlockIndex":2,"delta":{"text":"Now JSON."}}} +{"contentBlockStart":{"contentBlockIndex":3,"start":{"toolUse":{"toolUseId":"json-tool-id","name":"json"}}}} +{"contentBlockDelta":{"contentBlockIndex":3,"delta":{"toolUse":{"input":"{\"result\":\"data\"}"}}}} +{"contentBlockStop":{"contentBlockIndex":3}} +{"messageStop":{"stopReason":"tool_use"}} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-with-tool.1.json b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-with-tool.1.json new file mode 100644 index 000000000..f58981057 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-with-tool.1.json @@ -0,0 +1,41 @@ +{ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Let me call a tool first." + }, + { + "toolUse": { + "toolUseId": "tool-1", + "name": "get-weather", + "input": { + "location": "SF" + } + } + }, + { + "type": "text", + "text": "Now the JSON." + }, + { + "toolUse": { + "toolUseId": "json-tool-id", + "name": "json", + "input": { + "result": "data" + } + } + } + ] + } + }, + "usage": { + "inputTokens": 22, + "outputTokens": 48, + "totalTokens": 70 + }, + "stopReason": "tool_use" +} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-with-tools.1.chunks.txt b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-with-tools.1.chunks.txt new file mode 100644 index 000000000..57c7bb1ac --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-with-tools.1.chunks.txt @@ -0,0 +1,12 @@ +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"Multiple tools."}}} +{"contentBlockStart":{"contentBlockIndex":1,"start":{"toolUse":{"toolUseId":"tool-1","name":"tool-a"}}}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"toolUse":{"input":"{\"a\":\"1\"}"}}}} +{"contentBlockStop":{"contentBlockIndex":1}} +{"contentBlockStart":{"contentBlockIndex":2,"start":{"toolUse":{"toolUseId":"tool-2","name":"tool-b"}}}} +{"contentBlockDelta":{"contentBlockIndex":2,"delta":{"toolUse":{"input":"{\"b\":\"2\"}"}}}} +{"contentBlockStop":{"contentBlockIndex":2}} +{"contentBlockDelta":{"contentBlockIndex":3,"delta":{"text":"Final JSON."}}} +{"contentBlockStart":{"contentBlockIndex":4,"start":{"toolUse":{"toolUseId":"json-tool-id","name":"json"}}}} +{"contentBlockDelta":{"contentBlockIndex":4,"delta":{"toolUse":{"input":"{\"final\":\"result\"}"}}}} +{"contentBlockStop":{"contentBlockIndex":4}} +{"messageStop":{"stopReason":"tool_use"}} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-with-tools.1.json b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-with-tools.1.json new file mode 100644 index 000000000..4b2b9c585 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-json-with-tools.1.json @@ -0,0 +1,50 @@ +{ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Calling multiple tools." + }, + { + "toolUse": { + "toolUseId": "tool-1", + "name": "tool-a", + "input": { + "param": "a" + } + } + }, + { + "toolUse": { + "toolUseId": "tool-2", + "name": "tool-b", + "input": { + "param": "b" + } + } + }, + { + "type": "text", + "text": "Finally the JSON response." + }, + { + "toolUse": { + "toolUseId": "json-tool-id", + "name": "json", + "input": { + "final": "result" + } + } + } + ] + } + }, + "usage": { + "inputTokens": 28, + "outputTokens": 62, + "totalTokens": 90 + }, + "stopReason": "tool_use" +} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-reasoning.chunks.txt b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-reasoning.chunks.txt new file mode 100644 index 000000000..1db22185b --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-reasoning.chunks.txt @@ -0,0 +1,26 @@ +{"messageStart":{"role":"assistant"}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"reasoningContent":{"text":"Let me count the r"}}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"reasoningContent":{"text":"'s in \""}}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"reasoningContent":{"text":"strawberry\":"}}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"reasoningContent":{"text":"\n\ns-t-r-a"}}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"reasoningContent":{"text":"-w-b-e-r"}}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"reasoningContent":{"text":"-r-y\n\nr"}}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"reasoningContent":{"text":" appears at positions "}}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"reasoningContent":{"text":"3, 8"}}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"reasoningContent":{"text":", and 9.\n\nSo there"}}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"reasoningContent":{"text":" are 3 r's."}}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"reasoningContent":{"text":""}}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"reasoningContent":{"signature":"Ep0CCkgICxABGAIqQOWPB6/PmA5SW9jC6FvaNq3E+U9ev4FMWcFWuAho+VGLCtazKc5WDjQ5i0MuxsY0o5pKDSVWVKii8KJDusXH4eASDK7jyzuk8iij7fJNihoMxHO9haYzt48R36HVIjCb/EmIFrJLIXShqN6DN//T6vZBtO9qj1QhNWJa3CGm8VZoq80S2/Ok4U0aVaIDiZcqggHC2b8BHuv8BHZrmsR0wjU1ynansBGMdfjnG+iIv8R5lPpRmYGhSVwNybwP3aQZ6o8Dr48Rau8TJfdsArW+r7bvL7bPs4f5nnlp2vG7WkMzWwABHK3fdM44zZ1GZQaWyECNWR2GfY6dXiklo94vgpFTPuZ97mfiN3LY6uYyBwL8RkDaGAE="}}}} +{"contentBlockStop":{"contentBlockIndex":0}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"text":"There"}}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"text":" are **3** r's in \""}}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"text":"strawberry\":"}}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"text":"\n\n1"}}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"text":". st"}}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"text":"**"}}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"text":"r**awbe"}}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"text":"**r****"}}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"text":"r**y"}}} +{"contentBlockStop":{"contentBlockIndex":1}} +{"messageStop":{"additionalModelResponseFields":{"delta":{"stop_sequence":null}},"stopReason":"end_turn"}} +{"metadata":{"metrics":{"latencyMs":2281},"usage":{"inputTokens":51,"outputTokens":94,"serverToolUsage":{},"totalTokens":145}}} \ No newline at end of file diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-reasoning.json b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-reasoning.json new file mode 100644 index 000000000..9aca48889 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-reasoning.json @@ -0,0 +1,34 @@ +{ + "metrics": { + "latencyMs": 2202 + }, + "output": { + "message": { + "content": [ + { + "reasoningContent": { + "reasoningText": { + "signature": "EvYBCkgICxABGAIqQA9E9VC377UnbjdfXCw4RwQaaIXsqocZKzI3WwWtXT/wBjzAkBOfIiVkP/mhvsPGb5oQW/2j+E1tZcNFoXIb26YSDBTSbDrbHwXJoAl3/hoMA/Qo4rIHvObP/yBrIjAT+l69mh9k7evFO+w+ransSD5YGq4hYrieOyJH3k/zHv8nBJwiz3Nlrb3jNFi2Ib0qXJH7BIK0hR85yYIoccReoJUpeLEMxeWbsO6f0EpOSMSi3YDLd1U9NAw7Itjeay1fEBPHwvawt2M4e/rn52SsuEz3p2DfOZm+N3bnL63rg0Cb9lFcp3k5kuU1h6hcGAE=", + "text": "Let me count the r's in \"strawberry\":\n\ns-t-r-a-w-b-e-r-r-y\n\nThere are 3 r's." + } + } + }, + { + "text": "There are **3** r's in \"strawberry\":\n\n1. st**r**awbe**r****r**y" + } + ], + "role": "assistant" + } + }, + "stopReason": "end_turn", + "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 51, + "outputTokens": 78, + "serverToolUsage": {}, + "totalTokens": 129 + } +} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-text.chunks.txt b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-text.chunks.txt new file mode 100644 index 000000000..cba6ad06f --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-text.chunks.txt @@ -0,0 +1,16 @@ +{"messageStart":{"role":"assistant"}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"Let"}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":" me count the \""}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"r\"s in \""}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"strawberry\":\n\ns-t-"}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"**"}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"r**-a-w-b"}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"-e-**"}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"r**-**"}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"r**-y\n\nThere"}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":" are **3"}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"** r"}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"'s in \"strawberry.\""}}} +{"contentBlockStop":{"contentBlockIndex":0}} +{"messageStop":{"additionalModelResponseFields":{"delta":{"stop_sequence":null}},"stopReason":"end_turn"}} +{"metadata":{"metrics":{"latencyMs":2040},"usage":{"inputTokens":22,"outputTokens":55,"serverToolUsage":{},"totalTokens":77}}} \ No newline at end of file diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-text.json b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-text.json new file mode 100644 index 000000000..884cbfa9f --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-text.json @@ -0,0 +1,26 @@ +{ + "metrics": { + "latencyMs": 1864 + }, + "output": { + "message": { + "content": [ + { + "text": "Let me count the \"r\"s in \"strawberry\":\n\ns-t-**r**-a-w-b-e-**r**-**r**-y\n\nThere are **3** \"r\"s in \"strawberry.\"" + } + ], + "role": "assistant" + } + }, + "stopReason": "end_turn", + "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 22, + "outputTokens": 57, + "serverToolUsage": {}, + "totalTokens": 79 + } +} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-tool-no-args.chunks.txt b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-tool-no-args.chunks.txt new file mode 100644 index 000000000..2ec6e7bf8 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-tool-no-args.chunks.txt @@ -0,0 +1,8 @@ +{"contentBlockStart":{"contentBlockIndex":0,"start":{}}} +{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"I'll update the issue list for you."}}} +{"contentBlockStop":{"contentBlockIndex":0}} +{"contentBlockStart":{"contentBlockIndex":1,"start":{"toolUse":{"toolUseId":"tool-use-id","name":"updateIssueList"}}}} +{"contentBlockDelta":{"contentBlockIndex":1,"delta":{"toolUse":{"input":""}}}} +{"contentBlockStop":{"contentBlockIndex":1}} +{"metadata":{"usage":{"inputTokens":100,"outputTokens":25,"totalTokens":125},"metrics":{"latencyMs":150}}} +{"messageStop":{"stopReason":"tool_use"}} diff --git a/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-tool-no-args.json b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-tool-no-args.json new file mode 100644 index 000000000..e2e295a8d --- /dev/null +++ b/aimux-providers/src/bedrock/tests/fixtures/amazon-bedrock-tool-no-args.json @@ -0,0 +1,25 @@ +{ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "text": "I'll update the issue list for you." + }, + { + "toolUse": { + "toolUseId": "tool-use-id", + "name": "updateIssueList", + "input": {} + } + } + ] + } + }, + "usage": { + "inputTokens": 100, + "outputTokens": 25, + "totalTokens": 125 + }, + "stopReason": "tool_use" +} diff --git a/aimux-providers/src/bedrock/tests/image_model.rs b/aimux-providers/src/bedrock/tests/image_model.rs new file mode 100644 index 000000000..a9a051019 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/image_model.rs @@ -0,0 +1,1019 @@ +//! Ports of `@ai-sdk/amazon-bedrock` 5.0.100 +//! `src/amazon-bedrock-image-model.test.ts`. +//! +//! Upstream builds `AmazonBedrockImageModel` directly with a base URL, plain +//! config headers and `fetch: injectFetchHeaders(..)`, so no request is +//! signed. The port does the same: `BedrockImageModel::from_config` with a +//! fixed endpoint, the config headers and the +//! [`inject_fetch_headers`](super::support::inject_fetch_headers) transport over the +//! mock transport (which stands in for the test server). +//! +//! One case stays in `aimux-providers/tests/amazon_bedrock_image_test.rs`, +//! whose test already covers it: "should throw error when request is +//! moderated". + +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; +use serde_json::{Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::image_model::{ + ImageCallOptions, ImageFile, ImageFileData, ImageModel, ImageResult, +}; +use aimux_core::shared::{AspectRatio, SharedProviderOptions, Size}; +use aimux_provider_utils::{FetchFunction, HeaderMapOpt}; + +use crate::bedrock::{AmazonBedrockProviderSettings, BedrockImageModel, create_amazon_bedrock}; +use crate::shared::{Endpoint, EndpointConfig, ResolveEndpoint}; +use crate::test_support::mock_fetch::{Canned, MockFetch}; + +use super::support::inject_fetch_headers; + +const PROMPT: &str = "A cute baby sea otter"; + +const BASE_URL: &str = "https://bedrock-runtime.us-east-1.amazonaws.com"; + +const INVOKE_URL: &str = + "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.nova-canvas-v1%3A0/invoke"; + +/// A fixed base URL and fixed headers: the `{ baseUrl, headers }` config of +/// the upstream model. +struct FixedEndpoint { + base_url: String, + headers: HeaderMapOpt, +} + +#[async_trait] +impl ResolveEndpoint for FixedEndpoint { + async fn resolve(&self) -> Result { + Ok(Endpoint { + base_url: self.base_url.clone(), + headers: self.headers.clone(), + }) + } +} + +/// `new AmazonBedrockImageModel(model_id, { baseUrl, headers, fetch })`. +fn image_model( + model_id: &str, + base_url: &str, + headers: &[(&str, &str)], + fetch: Option, +) -> BedrockImageModel { + BedrockImageModel::from_config( + model_id.to_string(), + EndpointConfig::new( + "amazon-bedrock".to_string(), + FixedEndpoint { + base_url: base_url.to_string(), + headers: headers + .iter() + .map(|(name, value)| ((*name).to_string(), Some((*value).to_string()))) + .collect(), + }, + fetch, + ), + ) +} + +/// `fakeFetchWithAuth` over `mock`. +fn fake_fetch_with_auth(mock: &Arc) -> FetchFunction { + inject_fetch_headers(&[("x-amz-auth", "test-auth")], mock.transport()) +} + +/// A binary JSON response, with the `content-length` the test server sends. +fn json_response(body: &Value) -> Canned { + let mut canned = Canned::json(body); + let length = canned.body.len().to_string(); + canned.headers.push(("content-length".to_string(), length)); + canned +} + +/// Options with every field the upstream calls leave `undefined` unset. +fn options(prompt: Option<&str>) -> ImageCallOptions { + ImageCallOptions { + prompt: prompt.map(str::to_string), + ..ImageCallOptions::new("") + } +} + +fn provider_options(namespace: &str, options: Value) -> SharedProviderOptions { + aimux_core::shared::provider_namespace(namespace, options).expect("options are an object") +} + +fn png(bytes: &[u8], media_type: &str) -> ImageFile { + ImageFile::File { + media_type: media_type.to_string(), + data: ImageFileData::Binary(bytes.to_vec()), + } +} + +/// PNG magic bytes. +const IMAGE_DATA: [u8; 4] = [137, 80, 78, 71]; +const MASK_DATA: [u8; 4] = [255, 255, 255, 0]; + +fn images_of(result: &ImageResult) -> Value { + serde_json::to_value(&result.images).expect("images serialize") +} + +mod capabilities { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › capabilities › advertises file=$supportsFileInputs and mask=$supportsMaskInputs for $modelId + #[test] + fn advertises_file_and_mask_inputs_for_model_id() { + for (model_id, supports_file_inputs, supports_mask_inputs) in [ + ("amazon.nova-canvas-v1:0", Some(true), Some(true)), + ("custom-image-model", None, None), + ] { + let model = image_model(model_id, "https://api.example.com", &[], None); + + assert_eq!( + model.supports_file_inputs(), + supports_file_inputs, + "{model_id}" + ); + assert_eq!( + model.supports_mask_inputs(), + supports_mask_inputs, + "{model_id}" + ); + } + } +} + +mod do_generate { + use super::*; + + /// The upstream `mockConfigHeaders`. + const MOCK_CONFIG_HEADERS: &[(&str, &str)] = &[ + ("config-header", "config-value"), + ("shared-header", "config-shared"), + ]; + + fn default_response() -> Canned { + json_response(&json!({ "images": ["base64-image-1", "base64-image-2"] })) + } + + /// The upstream `model` over a mock that answers `responses`. + fn model_and_mock(responses: Vec) -> (BedrockImageModel, Arc) { + let mock = MockFetch::new(responses); + let model = image_model( + "amazon.nova-canvas-v1:0", + BASE_URL, + MOCK_CONFIG_HEADERS, + Some(fake_fetch_with_auth(&mock)), + ); + (model, mock) + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › doGenerate › should pass the model and the settings + #[tokio::test] + async fn should_pass_the_model_and_the_settings() { + let (model, mock) = model_and_mock(vec![default_response()]); + + model + .do_generate(&ImageCallOptions { + n: 1, + size: Some(Size::new(1024, 1024)), + seed: Some(1234), + provider_options: provider_options( + "bedrock", + json!({ "negativeText": "bad", "quality": "premium", "cfgScale": 1.2 }), + ), + ..options(Some(PROMPT)) + }) + .await + .expect("generation succeeds"); + + let seen = mock.seen(); + assert_eq!(seen[0].url, INVOKE_URL); + assert_eq!( + seen[0].json_body(), + json!({ + "imageGenerationConfig": { + "cfgScale": 1.2, + "height": 1024, + "numberOfImages": 1, + "quality": "premium", + "seed": 1234, + "width": 1024, + }, + "taskType": "TEXT_IMAGE", + "textToImageParams": { + "negativeText": "bad", + "text": "A cute baby sea otter", + }, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › doGenerate › should accept the new `amazonBedrock` providerOptions key + #[tokio::test] + async fn should_accept_the_new_amazon_bedrock_provider_options_key() { + let (model, mock) = model_and_mock(vec![default_response()]); + + model + .do_generate(&ImageCallOptions { + n: 1, + size: Some(Size::new(1024, 1024)), + seed: Some(1234), + provider_options: provider_options( + "amazonBedrock", + json!({ + "negativeText": "bad", + "quality": "premium", + "cfgScale": 1.2, + "style": "PHOTOREALISM", + }), + ), + ..options(Some(PROMPT)) + }) + .await + .expect("generation succeeds"); + + assert_eq!( + mock.seen()[0].json_body(), + json!({ + "imageGenerationConfig": { + "cfgScale": 1.2, + "height": 1024, + "numberOfImages": 1, + "quality": "premium", + "seed": 1234, + "width": 1024, + }, + "taskType": "TEXT_IMAGE", + "textToImageParams": { + "negativeText": "bad", + "style": "PHOTOREALISM", + "text": "A cute baby sea otter", + }, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › doGenerate › should prefer `amazonBedrock` providerOptions over the legacy `bedrock` key + #[tokio::test] + async fn should_prefer_amazon_bedrock_provider_options_over_the_legacy_bedrock_key() { + let (model, mock) = model_and_mock(vec![default_response()]); + let mut provider_options = provider_options( + "bedrock", + json!({ + "negativeText": "legacy", + "quality": "standard", + "cfgScale": 1, + "style": "DESIGN_SKETCH", + }), + ); + provider_options.extend(super::provider_options( + "amazonBedrock", + json!({ + "negativeText": "preferred", + "quality": "premium", + "cfgScale": 2, + "style": "PHOTOREALISM", + }), + )); + + model + .do_generate(&ImageCallOptions { + n: 1, + provider_options, + ..options(Some(PROMPT)) + }) + .await + .expect("generation succeeds"); + + assert_eq!( + mock.seen()[0].json_body(), + json!({ + "imageGenerationConfig": { + "cfgScale": 2, + "numberOfImages": 1, + "quality": "premium", + }, + "taskType": "TEXT_IMAGE", + "textToImageParams": { + "negativeText": "preferred", + "style": "PHOTOREALISM", + "text": "A cute baby sea otter", + }, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › doGenerate › should reject invalid provider options + #[tokio::test] + async fn should_reject_invalid_provider_options() { + let (model, _mock) = model_and_mock(vec![default_response()]); + + let error = model + .do_generate(&ImageCallOptions { + n: 1, + provider_options: provider_options("amazonBedrock", json!({ "quality": "ultra" })), + ..options(Some(PROMPT)) + }) + .await + .expect_err("the options are rejected"); + + assert!( + error + .to_string() + .contains("invalid amazonBedrock provider options"), + "unexpected error: {error}" + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › doGenerate › should properly combine headers from all sources + #[tokio::test] + async fn should_properly_combine_headers_from_all_sources() { + let mock = MockFetch::new(vec![default_response()]); + let model_with_headers = image_model( + "amazon.nova-canvas-v1:0", + BASE_URL, + &[ + ("model-header", "model-value"), + ("shared-header", "model-shared"), + ], + Some(inject_fetch_headers( + &[ + ("signed-header", "signed-value"), + ("authorization", "AWS4-HMAC-SHA256..."), + ], + mock.transport(), + )), + ); + + model_with_headers + .do_generate(&ImageCallOptions { + n: 1, + headers: Some(HashMap::from([ + ("options-header".to_string(), "options-value".to_string()), + ("shared-header".to_string(), "options-shared".to_string()), + ])), + ..options(Some(PROMPT)) + }) + .await + .expect("generation succeeds"); + + let request_headers = &mock.seen()[0].headers; + assert_eq!(request_headers["options-header"], "options-value"); + assert_eq!(request_headers["model-header"], "model-value"); + assert_eq!(request_headers["signed-header"], "signed-value"); + assert_eq!(request_headers["authorization"], "AWS4-HMAC-SHA256..."); + assert_eq!(request_headers["shared-header"], "options-shared"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › doGenerate › should respect maxImagesPerCall setting + #[test] + fn should_respect_max_images_per_call_setting() { + let provider = create_amazon_bedrock(AmazonBedrockProviderSettings::default()) + .expect("default settings are valid"); + + let default_model = provider.image("amazon.nova-canvas-v1:0"); + assert_eq!(default_model.max_images_per_call(), Some(5)); + + let unknown_model = provider.image("unknown-model"); + assert_eq!(unknown_model.max_images_per_call(), Some(1)); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › doGenerate › should return warnings for unsupported settings + #[tokio::test] + async fn should_return_warnings_for_unsupported_settings() { + let (model, _mock) = model_and_mock(vec![default_response()]); + + let result = model + .do_generate(&ImageCallOptions { + n: 1, + size: Some(Size::new(1024, 1024)), + aspect_ratio: Some(AspectRatio::new(1, 1)), + ..options(Some(PROMPT)) + }) + .await + .expect("generation succeeds"); + + assert_eq!( + serde_json::to_value(&result.warnings).expect("warnings serialize"), + json!([ + { + "details": "This model does not support aspect ratio. Use `size` instead.", + "feature": "aspectRatio", + "type": "unsupported", + }, + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › doGenerate › should extract the generated images + #[tokio::test] + async fn should_extract_the_generated_images() { + let (model, _mock) = model_and_mock(vec![default_response()]); + + let result = model + .do_generate(&ImageCallOptions { + n: 1, + ..options(Some(PROMPT)) + }) + .await + .expect("generation succeeds"); + + assert_eq!( + images_of(&result), + json!(["base64-image-1", "base64-image-2"]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › doGenerate › should include response data with timestamp, modelId and headers + /// + /// Upstream fixes the timestamp through `_internal.currentDate`. + /// `BedrockImageModel` stamps `Timestamp::now()` and takes no date + /// provider, so the timestamp is asserted to be the call time instead of + /// `2024-03-15T12:00:00.000Z`. + #[tokio::test] + async fn should_include_response_data_with_timestamp_model_id_and_headers() { + let mock = MockFetch::new(vec![default_response()]); + let custom_model = image_model( + "amazon.nova-canvas-v1:0", + BASE_URL, + &[], + Some(mock.transport()), + ); + + let before = aimux_core::Timestamp::now(); + let result = custom_model + .do_generate(&ImageCallOptions { + n: 1, + size: Some(Size::new(1024, 1024)), + ..options(Some(PROMPT)) + }) + .await + .expect("generation succeeds"); + let after = aimux_core::Timestamp::now(); + + assert_eq!( + result.response.headers, + Some(HashMap::from([ + ("content-length".to_string(), "46".to_string()), + ("content-type".to_string(), "application/json".to_string()), + ])) + ); + assert_eq!( + result.response.model_id.as_deref(), + Some("amazon.nova-canvas-v1:0") + ); + let timestamp = result.response.timestamp.expect("timestamp").unix_millis(); + assert!(before.unix_millis() <= timestamp && timestamp <= after.unix_millis()); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › doGenerate › should use real date when no custom date provider is specified + #[tokio::test] + async fn should_use_real_date_when_no_custom_date_provider_is_specified() { + let (model, _mock) = model_and_mock(vec![default_response()]); + + let before_date = aimux_core::Timestamp::now(); + let result = model + .do_generate(&ImageCallOptions { + n: 1, + seed: Some(1234), + ..options(Some(PROMPT)) + }) + .await + .expect("generation succeeds"); + let after_date = aimux_core::Timestamp::now(); + + let timestamp = result.response.timestamp.expect("timestamp").unix_millis(); + assert!(timestamp >= before_date.unix_millis()); + assert!(timestamp <= after_date.unix_millis()); + assert_eq!( + result.response.model_id.as_deref(), + Some("amazon.nova-canvas-v1:0") + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › doGenerate › should pass the style parameter when provided + #[tokio::test] + async fn should_pass_the_style_parameter_when_provided() { + let (model, mock) = model_and_mock(vec![default_response()]); + + model + .do_generate(&ImageCallOptions { + n: 1, + size: Some(Size::new(1024, 1024)), + seed: Some(1234), + provider_options: provider_options( + "bedrock", + json!({ + "negativeText": "bad", + "quality": "premium", + "cfgScale": 1.2, + "style": "PHOTOREALISM", + }), + ), + ..options(Some(PROMPT)) + }) + .await + .expect("generation succeeds"); + + assert_eq!( + mock.seen()[0].json_body(), + json!({ + "imageGenerationConfig": { + "cfgScale": 1.2, + "height": 1024, + "numberOfImages": 1, + "quality": "premium", + "seed": 1234, + "width": 1024, + }, + "taskType": "TEXT_IMAGE", + "textToImageParams": { + "negativeText": "bad", + "style": "PHOTOREALISM", + "text": "A cute baby sea otter", + }, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › doGenerate › should not include style parameter when not provided + #[tokio::test] + async fn should_not_include_style_parameter_when_not_provided() { + let (model, mock) = model_and_mock(vec![default_response()]); + + model + .do_generate(&ImageCallOptions { + n: 1, + size: Some(Size::new(1024, 1024)), + seed: Some(1234), + provider_options: provider_options("bedrock", json!({ "quality": "standard" })), + ..options(Some(PROMPT)) + }) + .await + .expect("generation succeeds"); + + let request_body = mock.seen()[0].json_body(); + assert!( + request_body["textToImageParams"].get("style").is_none(), + "{request_body}" + ); + } + + // "should throw error when request is moderated": the marked test in + // aimux-providers/tests/amazon_bedrock_image_test.rs. + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › doGenerate › should throw error when no images are returned + #[tokio::test] + async fn should_throw_error_when_no_images_are_returned() { + let (model, _mock) = model_and_mock(vec![json_response(&json!({ "images": [] }))]); + + let error = model + .do_generate(&ImageCallOptions { + n: 1, + ..options(Some("Generate an image")) + }) + .await + .expect_err("generation fails"); + + assert!( + error + .to_string() + .contains("Amazon Bedrock returned no images"), + "unexpected error: {error}" + ); + } +} + +mod image_editing { + use super::*; + + fn edited_response() -> Canned { + json_response(&json!({ "images": ["edited-image-base64"] })) + } + + /// The upstream `model` of this describe over a mock that answers once. + fn model_and_mock() -> (BedrockImageModel, Arc) { + let mock = MockFetch::new(vec![edited_response()]); + let model = image_model( + "amazon.nova-canvas-v1:0", + BASE_URL, + &[], + Some(fake_fetch_with_auth(&mock)), + ); + (model, mock) + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › Image Editing › should send inpainting request with files and maskPrompt + #[tokio::test] + async fn should_send_inpainting_request_with_files_and_mask_prompt() { + let (model, mock) = model_and_mock(); + + model + .do_generate(&ImageCallOptions { + n: 1, + seed: Some(42), + files: Some(vec![png(&IMAGE_DATA, "image/png")]), + provider_options: provider_options( + "bedrock", + json!({ "maskPrompt": "cat", "quality": "standard", "cfgScale": 7.0 }), + ), + ..options(Some("a cute corgi dog")) + }) + .await + .expect("generation succeeds"); + + assert_eq!( + mock.seen()[0].json_body(), + json!({ + "imageGenerationConfig": { + "cfgScale": 7.0, + "numberOfImages": 1, + "quality": "standard", + "seed": 42, + }, + "inPaintingParams": { + "image": "iVBORw==", + "maskPrompt": "cat", + "text": "a cute corgi dog", + }, + "taskType": "INPAINTING", + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › Image Editing › should send inpainting request with files and mask image + #[tokio::test] + async fn should_send_inpainting_request_with_files_and_mask_image() { + let (model, mock) = model_and_mock(); + + model + .do_generate(&ImageCallOptions { + n: 1, + files: Some(vec![png(&IMAGE_DATA, "image/png")]), + mask: Some(png(&MASK_DATA, "image/png")), + provider_options: provider_options("bedrock", json!({ "quality": "standard" })), + ..options(Some( + "A sunlit indoor lounge area with a pool containing a flamingo", + )) + }) + .await + .expect("generation succeeds"); + + assert_eq!( + mock.seen()[0].json_body(), + json!({ + "imageGenerationConfig": { + "numberOfImages": 1, + "quality": "standard", + }, + "inPaintingParams": { + "image": "iVBORw==", + "maskImage": "////AA==", + "text": "A sunlit indoor lounge area with a pool containing a flamingo", + }, + "taskType": "INPAINTING", + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › Image Editing › should send inpainting request with base64 string data + #[tokio::test] + async fn should_send_inpainting_request_with_base64_string_data() { + let (model, mock) = model_and_mock(); + let base64_image = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk"; + + model + .do_generate(&ImageCallOptions { + n: 1, + files: Some(vec![ImageFile::File { + media_type: "image/png".to_string(), + data: ImageFileData::Base64(base64_image.to_string()), + }]), + provider_options: provider_options( + "bedrock", + json!({ "maskPrompt": "background" }), + ), + ..options(Some("Edit this image")) + }) + .await + .expect("generation succeeds"); + + let request_body = mock.seen()[0].json_body(); + assert_eq!(request_body["taskType"], "INPAINTING"); + assert_eq!(request_body["inPaintingParams"]["image"], base64_image); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › Image Editing › should include negativeText in inpainting params + #[tokio::test] + async fn should_include_negative_text_in_inpainting_params() { + let (model, mock) = model_and_mock(); + + model + .do_generate(&ImageCallOptions { + n: 1, + files: Some(vec![png(&IMAGE_DATA, "image/png")]), + provider_options: provider_options( + "bedrock", + json!({ "maskPrompt": "sky", "negativeText": "clouds, rain" }), + ), + ..options(Some("a beautiful garden")) + }) + .await + .expect("generation succeeds"); + + assert_eq!( + mock.seen()[0].json_body(), + json!({ + "imageGenerationConfig": { + "numberOfImages": 1, + }, + "inPaintingParams": { + "image": "iVBORw==", + "maskPrompt": "sky", + "negativeText": "clouds, rain", + "text": "a beautiful garden", + }, + "taskType": "INPAINTING", + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › Image Editing › should extract edited images from response + #[tokio::test] + async fn should_extract_edited_images_from_response() { + let (model, _mock) = model_and_mock(); + + let result = model + .do_generate(&ImageCallOptions { + n: 1, + files: Some(vec![png(&IMAGE_DATA, "image/png")]), + provider_options: provider_options("bedrock", json!({ "maskPrompt": "object" })), + ..options(Some("Edit this image")) + }) + .await + .expect("generation succeeds"); + + assert_eq!(images_of(&result), json!(["edited-image-base64"])); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › Image Editing › should throw error for URL-based images + #[tokio::test] + async fn should_throw_error_for_url_based_images() { + let (model, _mock) = model_and_mock(); + + let error = model + .do_generate(&ImageCallOptions { + n: 1, + files: Some(vec![ImageFile::Url { + url: "https://example.com/image.png".to_string(), + }]), + ..options(Some("Edit this image")) + }) + .await + .expect_err("generation fails"); + + assert!( + error + .to_string() + .contains("URL-based images are not supported for Amazon Bedrock image editing."), + "unexpected error: {error}" + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › Image Editing › should send outpainting request with taskType OUTPAINTING + #[tokio::test] + async fn should_send_outpainting_request_with_task_type_outpainting() { + let (model, mock) = model_and_mock(); + + model + .do_generate(&ImageCallOptions { + n: 1, + files: Some(vec![png(&IMAGE_DATA, "image/png")]), + mask: Some(png(&MASK_DATA, "image/png")), + provider_options: provider_options( + "bedrock", + json!({ + "taskType": "OUTPAINTING", + "outPaintingMode": "DEFAULT", + "negativeText": "bad quality", + }), + ), + ..options(Some("Extend the background with a beautiful sunset")) + }) + .await + .expect("generation succeeds"); + + assert_eq!( + mock.seen()[0].json_body(), + json!({ + "imageGenerationConfig": { + "numberOfImages": 1, + }, + "outPaintingParams": { + "image": "iVBORw==", + "maskImage": "////AA==", + "negativeText": "bad quality", + "outPaintingMode": "DEFAULT", + "text": "Extend the background with a beautiful sunset", + }, + "taskType": "OUTPAINTING", + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › Image Editing › should send outpainting request with maskPrompt + #[tokio::test] + async fn should_send_outpainting_request_with_mask_prompt() { + let (model, mock) = model_and_mock(); + + model + .do_generate(&ImageCallOptions { + n: 1, + files: Some(vec![png(&IMAGE_DATA, "image/png")]), + provider_options: provider_options( + "bedrock", + json!({ "taskType": "OUTPAINTING", "maskPrompt": "background" }), + ), + ..options(Some("Replace the background with mountains")) + }) + .await + .expect("generation succeeds"); + + assert_eq!( + mock.seen()[0].json_body(), + json!({ + "imageGenerationConfig": { + "numberOfImages": 1, + }, + "outPaintingParams": { + "image": "iVBORw==", + "maskPrompt": "background", + "text": "Replace the background with mountains", + }, + "taskType": "OUTPAINTING", + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › Image Editing › should send background removal request + #[tokio::test] + async fn should_send_background_removal_request() { + let (model, mock) = model_and_mock(); + + model + .do_generate(&ImageCallOptions { + n: 1, + files: Some(vec![png(&IMAGE_DATA, "image/png")]), + provider_options: provider_options( + "bedrock", + json!({ "taskType": "BACKGROUND_REMOVAL" }), + ), + ..options(None) + }) + .await + .expect("generation succeeds"); + + assert_eq!( + mock.seen()[0].json_body(), + json!({ + "backgroundRemovalParams": { + "image": "iVBORw==", + }, + "taskType": "BACKGROUND_REMOVAL", + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › Image Editing › should send image variation request with single image + #[tokio::test] + async fn should_send_image_variation_request_with_single_image() { + let (model, mock) = model_and_mock(); + + model + .do_generate(&ImageCallOptions { + n: 3, + size: Some(Size::new(512, 512)), + files: Some(vec![png(&IMAGE_DATA, "image/png")]), + provider_options: provider_options( + "bedrock", + json!({ + "taskType": "IMAGE_VARIATION", + "similarityStrength": 0.7, + "negativeText": "bad quality, low resolution", + }), + ), + ..options(Some("Create a variation in anime style")) + }) + .await + .expect("generation succeeds"); + + assert_eq!( + mock.seen()[0].json_body(), + json!({ + "imageGenerationConfig": { + "height": 512, + "numberOfImages": 3, + "width": 512, + }, + "imageVariationParams": { + "images": ["iVBORw=="], + "negativeText": "bad quality, low resolution", + "similarityStrength": 0.7, + "text": "Create a variation in anime style", + }, + "taskType": "IMAGE_VARIATION", + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › Image Editing › should send image variation request with multiple images + #[tokio::test] + async fn should_send_image_variation_request_with_multiple_images() { + let (model, mock) = model_and_mock(); + + model + .do_generate(&ImageCallOptions { + n: 1, + files: Some(vec![ + png(&IMAGE_DATA, "image/png"), + png(&[255, 216, 255, 224], "image/jpeg"), + ]), + provider_options: provider_options( + "bedrock", + json!({ "taskType": "IMAGE_VARIATION" }), + ), + ..options(Some("Combine these images into one cohesive scene")) + }) + .await + .expect("generation succeeds"); + + assert_eq!( + mock.seen()[0].json_body(), + json!({ + "imageGenerationConfig": { + "numberOfImages": 1, + }, + "imageVariationParams": { + "images": ["iVBORw==", "/9j/4A=="], + "text": "Combine these images into one cohesive scene", + }, + "taskType": "IMAGE_VARIATION", + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › Image Editing › should default to IMAGE_VARIATION when files provided without mask or maskPrompt + #[tokio::test] + async fn should_default_to_image_variation_when_files_provided_without_mask_or_mask_prompt() { + let (model, mock) = model_and_mock(); + + model + .do_generate(&ImageCallOptions { + n: 1, + files: Some(vec![png(&IMAGE_DATA, "image/png")]), + ..options(Some("Create variations")) + }) + .await + .expect("generation succeeds"); + + assert_eq!(mock.seen()[0].json_body()["taskType"], "IMAGE_VARIATION"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › Image Editing › should default to INPAINTING when files provided with mask + #[tokio::test] + async fn should_default_to_inpainting_when_files_provided_with_mask() { + let (model, mock) = model_and_mock(); + + model + .do_generate(&ImageCallOptions { + n: 1, + files: Some(vec![png(&IMAGE_DATA, "image/png")]), + mask: Some(png(&MASK_DATA, "image/png")), + ..options(Some("Edit masked area")) + }) + .await + .expect("generation succeeds"); + + assert_eq!(mock.seen()[0].json_body()["taskType"], "INPAINTING"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › Image Editing › should default to INPAINTING when files provided with maskPrompt + #[tokio::test] + async fn should_default_to_inpainting_when_files_provided_with_mask_prompt() { + let (model, mock) = model_and_mock(); + + model + .do_generate(&ImageCallOptions { + n: 1, + files: Some(vec![png(&IMAGE_DATA, "image/png")]), + provider_options: provider_options("bedrock", json!({ "maskPrompt": "cat" })), + ..options(Some("Edit the cat")) + }) + .await + .expect("generation succeeds"); + + assert_eq!(mock.seen()[0].json_body()["taskType"], "INPAINTING"); + } +} diff --git a/aimux-providers/src/bedrock/tests/mod.rs b/aimux-providers/src/bedrock/tests/mod.rs new file mode 100644 index 000000000..14f633766 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/mod.rs @@ -0,0 +1,19 @@ +//! Tests ported from `@ai-sdk/amazon-bedrock` 5.0.100, one file per upstream test file. + +mod chat_language_model; +mod chat_language_model_event_stream; +mod chat_language_model_options; +mod convert_to_chat_messages; +mod convert_usage; +mod embedding_model; +mod error; +mod event_stream_decoder; +mod event_stream_response_handler; +mod image_model; +mod normalize_tool_call_id; +mod prepare_tools; +mod provider; +mod reranking; +mod sigv4_fetch; +mod support; +mod usage_raw; diff --git a/aimux-providers/src/bedrock/tests/normalize_tool_call_id.rs b/aimux-providers/src/bedrock/tests/normalize_tool_call_id.rs new file mode 100644 index 000000000..28ca38def --- /dev/null +++ b/aimux-providers/src/bedrock/tests/normalize_tool_call_id.rs @@ -0,0 +1,21 @@ +//! The other `isMistralModel` / `normalizeToolCallId` cases are ported in +//! `aimux-providers/tests/bedrock_remaining_test.rs`. + +use crate::bedrock::convert::normalize_tool_call_id; + +#[allow(clippy::module_inception)] // the describe has the file's name +mod normalize_tool_call_id { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/normalize-tool-call-id.test.ts › normalizeToolCallId › should keep distinct Bedrock IDs distinct when their prefixes match + #[test] + fn should_keep_distinct_bedrock_ids_distinct_when_their_prefixes_match() { + assert_eq!( + [ + normalize_tool_call_id("tooluse_Ac1Xq9ZklmNoPq", true), + normalize_tool_call_id("tooluse_Ac2Yt7WrstUvWx", true), + ], + ["7rDVWRig0", "bNZvZKNBZ"] + ); + } +} diff --git a/aimux-providers/src/bedrock/tests/prepare_tools.rs b/aimux-providers/src/bedrock/tests/prepare_tools.rs new file mode 100644 index 000000000..01f5a2265 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/prepare_tools.rs @@ -0,0 +1,831 @@ +//! Upstream unit-tests `prepareTools` directly. aimux folds tool preparation +//! into `build_request_body_checked`, so the harness calls it with an empty +//! prompt and reads upstream's result fields back from the request body: +//! `toolConfig` is `body.toolConfig` (absent → `{}`), `additionalTools` is the +//! `tool_choice` that lands in `body.additionalModelRequestFields`, `betas` is +//! `additionalModelRequestFields.anthropic_beta`, and `toolWarnings` are the +//! returned warnings. `disableParallelToolUse` is passed the way the chat +//! model receives it (`providerOptions.anthropic.disableParallelToolUse`); +//! `rejectsForcedToolUse` is derived by aimux from the model id (every model +//! upstream passes it for rejects forced tool use). + +use aimux_core::options::CallOptions; +use aimux_core::shared::provider_namespace; +use serde_json::{Value, json}; + +use crate::bedrock::convert::build_request_body_checked; + +const NON_ANTHROPIC_MODEL: &str = "meta.llama3-70b-instruct-v1:0"; +const ANTHROPIC_MODEL: &str = "anthropic.claude-sonnet-4-5-20250929-v1:0"; + +#[derive(Default)] +struct PrepareToolsArgs { + tools: Option, + tool_choice: Option, + model_id: &'static str, + model_family: Option<&'static str>, + disable_parallel_tool_use: bool, +} + +struct PreparedTools { + tool_config: Value, + additional_tools: Option, + betas: Vec, + tool_warnings: Value, +} + +fn prepare_tools(args: PrepareToolsArgs) -> PreparedTools { + let mut call = CallOptions::new(vec![]); + call.tools = args + .tools + .map(|tools| serde_json::from_value(tools).expect("valid tools")); + call.tool_choice = args + .tool_choice + .map(|choice| serde_json::from_value(choice).expect("valid tool choice")); + if args.disable_parallel_tool_use { + call.provider_options = Some( + provider_namespace("anthropic", json!({ "disableParallelToolUse": true })).unwrap(), + ); + } + let (body, warnings, _, _) = + build_request_body_checked(args.model_id, &call, args.model_family).unwrap(); + let fields = &body["additionalModelRequestFields"]; + PreparedTools { + tool_config: body.get("toolConfig").cloned().unwrap_or_else(|| json!({})), + additional_tools: fields + .get("tool_choice") + .map(|choice| json!({ "tool_choice": choice })), + betas: fields["anthropic_beta"] + .as_array() + .cloned() + .unwrap_or_default(), + tool_warnings: serde_json::to_value(warnings).unwrap(), + } +} + +fn tool_spec(result: &PreparedTools, index: usize) -> &serde_json::Map { + result.tool_config["tools"][index]["toolSpec"] + .as_object() + .expect("toolSpec") +} + +#[allow(clippy::module_inception)] // the describe has the file's name +mod prepare_tools { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › should return empty toolConfig when tools are undefined + #[test] + fn should_return_empty_tool_config_when_tools_are_undefined() { + let result = prepare_tools(PrepareToolsArgs { + tools: None, + model_id: ANTHROPIC_MODEL, + ..Default::default() + }); + + assert_eq!(result.tool_config, json!({})); + assert_eq!(result.additional_tools, None); + assert!(result.betas.is_empty()); + assert_eq!(result.tool_warnings, json!([])); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › should return empty toolConfig when tools are empty + #[test] + fn should_return_empty_tool_config_when_tools_are_empty() { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([])), + model_id: ANTHROPIC_MODEL, + ..Default::default() + }); + + assert_eq!(result.tool_config, json!({})); + assert_eq!(result.additional_tools, None); + assert!(result.betas.is_empty()); + assert_eq!(result.tool_warnings, json!([])); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › should correctly prepare function tools + #[test] + fn should_correctly_prepare_function_tools() { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([{ + "type": "function", + "name": "testFunction", + "description": "A test function", + "inputSchema": { "type": "object", "properties": {} }, + }])), + model_id: ANTHROPIC_MODEL, + ..Default::default() + }); + + assert_eq!( + result.tool_config["tools"], + json!([{ + "toolSpec": { + "name": "testFunction", + "description": "A test function", + "inputSchema": { "json": { "type": "object", "properties": {} } }, + }, + }]) + ); + assert_eq!(result.tool_warnings, json!([])); + } + + mod tool_description_handling { + use super::*; + + fn with_description(description: &str) -> PreparedTools { + prepare_tools(PrepareToolsArgs { + tools: Some(json!([{ + "type": "function", + "name": "testFunction", + "description": description, + "inputSchema": { "type": "object", "properties": {} }, + }])), + model_id: ANTHROPIC_MODEL, + ..Default::default() + }) + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › tool description handling › should exclude description when it is empty string + #[test] + fn should_exclude_description_when_it_is_empty_string() { + let result = with_description(""); + + assert!(!tool_spec(&result, 0).contains_key("description")); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › tool description handling › should exclude description when it is whitespace-only + #[test] + fn should_exclude_description_when_it_is_whitespace_only() { + let result = with_description(" "); + + assert!(!tool_spec(&result, 0).contains_key("description")); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › tool description handling › should include description when it has content + #[test] + fn should_include_description_when_it_has_content() { + let result = with_description("Valid description"); + + assert_eq!(tool_spec(&result, 0)["description"], "Valid description"); + } + } + + mod unsupported_provider_defined_tools { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › unsupported provider-defined tools › should warn for provider-defined tools on non-anthropic models + #[test] + fn should_warn_for_provider_defined_tools_on_non_anthropic_models() { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([{ + "type": "provider", + "id": "some.custom_tool", + "name": "custom_tool", + "args": {}, + }])), + model_id: NON_ANTHROPIC_MODEL, + ..Default::default() + }); + + assert_eq!(result.tool_config, json!({})); + assert_eq!( + result.tool_warnings, + json!([{ "type": "unsupported", "feature": "tool some.custom_tool" }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › unsupported provider-defined tools › should warn and filter out web_search_20250305 tool + #[test] + fn should_warn_and_filter_out_web_search_20250305_tool() { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([{ + "type": "provider", + "id": "anthropic.web_search_20250305", + "name": "web_search", + "args": {}, + }])), + model_id: ANTHROPIC_MODEL, + ..Default::default() + }); + + assert_eq!(result.tool_config, json!({})); + assert_eq!( + result.tool_warnings, + json!([{ + "type": "unsupported", + "feature": "web_search_20250305 tool", + "details": "The web_search_20250305 tool is not supported on Amazon Bedrock.", + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › unsupported provider-defined tools › should warn and filter out unsupported %s tool + #[test] + fn should_warn_and_filter_out_unsupported_tool() { + for tool_id in [ + "anthropic.web_search_20260318", + "anthropic.web_fetch_20260318", + ] { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([{ + "type": "provider", + "id": tool_id, + "name": if tool_id.contains("search") { "web_search" } else { "web_fetch" }, + "args": {}, + }])), + model_id: ANTHROPIC_MODEL, + ..Default::default() + }); + + let tool_type = &tool_id["anthropic.".len()..]; + assert_eq!(result.tool_config, json!({}), "{tool_id}"); + assert_eq!( + result.tool_warnings, + json!([{ + "type": "unsupported", + "feature": format!("{tool_type} tool"), + "details": format!("The {tool_type} tool is not supported on Amazon Bedrock."), + }]), + "{tool_id}" + ); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › unsupported provider-defined tools › should return empty toolConfig when all tools are filtered out + #[test] + fn should_return_empty_tool_config_when_all_tools_are_filtered_out() { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([{ + "type": "provider", + "id": "anthropic.web_search_20250305", + "name": "web_search", + "args": {}, + }])), + model_id: ANTHROPIC_MODEL, + ..Default::default() + }); + + assert_eq!(result.tool_config, json!({})); + } + } + + mod tool_choice { + use super::*; + + fn test_function() -> Value { + json!([{ + "type": "function", + "name": "testFunction", + "description": "Test", + "inputSchema": {}, + }]) + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › tool choice › should use Anthropic tool choice fields for a declared Anthropic application inference profile + #[test] + fn should_use_anthropic_tool_choice_fields_for_a_declared_anthropic_inference_profile() { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(test_function()), + model_id: "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/custom-profile", + model_family: Some("anthropic"), + disable_parallel_tool_use: true, + ..Default::default() + }); + + assert_eq!( + result.additional_tools, + Some(json!({ + "tool_choice": { "type": "auto", "disable_parallel_tool_use": true }, + })) + ); + assert_eq!(result.tool_config.get("toolChoice"), None); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › tool choice › should use Anthropic tool choice fields when parallel tool use is disabled + #[test] + fn should_use_anthropic_tool_choice_fields_when_parallel_tool_use_is_disabled() { + for (tool_choice, expected) in [ + ( + None, + json!({ "type": "auto", "disable_parallel_tool_use": true }), + ), + ( + Some(json!({ "type": "auto" })), + json!({ "type": "auto", "disable_parallel_tool_use": true }), + ), + ( + Some(json!({ "type": "required" })), + json!({ "type": "any", "disable_parallel_tool_use": true }), + ), + ( + Some(json!({ "type": "tool", "toolName": "testFunction" })), + json!({ + "type": "tool", + "name": "testFunction", + "disable_parallel_tool_use": true, + }), + ), + ] { + let label = format!("{tool_choice:?}"); + let result = prepare_tools(PrepareToolsArgs { + tools: Some(test_function()), + tool_choice, + model_id: ANTHROPIC_MODEL, + disable_parallel_tool_use: true, + ..Default::default() + }); + + assert_eq!( + result.additional_tools, + Some(json!({ "tool_choice": expected })), + "{label}" + ); + assert_eq!(result.tool_config.get("toolChoice"), None, "{label}"); + } + } + + mod models_that_reject_forced_tool_use { + use super::*; + + fn weather_tool() -> Value { + json!({ + "type": "function", + "name": "getWeather", + "description": "Get weather", + "inputSchema": { "type": "object" }, + }) + } + + fn time_tool() -> Value { + json!({ + "type": "function", + "name": "getTime", + "description": "Get time", + "inputSchema": { "type": "object" }, + }) + } + + fn tool_names(result: &PreparedTools) -> Value { + result.tool_config["tools"] + .as_array() + .unwrap() + .iter() + .map(|tool| tool["toolSpec"]["name"].clone()) + .collect() + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › tool choice › models that reject forced tool use › should fall back to auto for tool choice "required" on %s + #[test] + fn should_fall_back_to_auto_for_tool_choice_required_on() { + for model_id in [ + "anthropic.claude-opus-5-5", + "us.anthropic.claude-opus-5-5", + "global.anthropic.claude-opus-5-5", + "anthropic.claude-sonnet-5-5", + "us.anthropic.claude-sonnet-5-5", + ] { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([weather_tool()])), + tool_choice: Some(json!({ "type": "required" })), + model_id, + ..Default::default() + }); + + assert_eq!( + result.tool_config["toolChoice"], + json!({ "auto": {} }), + "{model_id}" + ); + assert_eq!( + result.tool_warnings, + json!([{ + "type": "unsupported", + "feature": "toolChoice", + "details": "toolChoice 'required' is not supported by this model because it rejects forced tool use. Using 'auto' instead. Instruct the model to use a tool in the prompt and verify that a tool call was made.", + }]), + "{model_id}" + ); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › tool choice › models that reject forced tool use › should fall back to auto and send only the named function tool for tool choice "tool" + #[test] + fn should_fall_back_to_auto_and_send_only_the_named_function_tool() { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([weather_tool(), time_tool()])), + tool_choice: Some(json!({ "type": "tool", "toolName": "getWeather" })), + model_id: "us.anthropic.claude-opus-5-5", + ..Default::default() + }); + + assert_eq!(result.tool_config["toolChoice"], json!({ "auto": {} })); + assert_eq!(tool_names(&result), json!(["getWeather"])); + assert_eq!( + result.tool_warnings, + json!([{ + "type": "unsupported", + "feature": "toolChoice", + "details": "toolChoice 'tool' is not supported by this model because it rejects forced tool use. Only the 'getWeather' tool is sent with 'auto' tool choice. Instruct the model to use the tool in the prompt and verify that a tool call was made.", + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › tool choice › models that reject forced tool use › should drop provider tools that are not the named tool for tool choice "tool" + #[test] + fn should_drop_provider_tools_that_are_not_the_named_tool_for_tool_choice_tool() { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([ + weather_tool(), + { + "type": "provider", + "id": "anthropic.bash_20250124", + "name": "bash", + "args": {}, + }, + ])), + tool_choice: Some(json!({ "type": "tool", "toolName": "getWeather" })), + model_id: "us.anthropic.claude-opus-5-5", + ..Default::default() + }); + + assert_eq!(result.tool_config["toolChoice"], json!({ "auto": {} })); + assert_eq!(tool_names(&result), json!(["getWeather"])); + assert_eq!(result.additional_tools, None); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › tool choice › models that reject forced tool use › should preserve disabled parallel tool use in the auto fallback + #[test] + fn should_preserve_disabled_parallel_tool_use_in_the_auto_fallback() { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([weather_tool()])), + tool_choice: Some(json!({ "type": "required" })), + model_id: "us.anthropic.claude-opus-5-5", + disable_parallel_tool_use: true, + ..Default::default() + }); + + assert_eq!( + result.additional_tools, + Some(json!({ + "tool_choice": { "type": "auto", "disable_parallel_tool_use": true }, + })) + ); + assert_eq!(result.tool_config.get("toolChoice"), None); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › tool choice › models that reject forced tool use › should keep forced tool choice for models that support it + #[test] + fn should_keep_forced_tool_choice_for_models_that_support_it() { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([weather_tool()])), + tool_choice: Some(json!({ "type": "required" })), + model_id: "us.anthropic.claude-opus-5", + ..Default::default() + }); + + assert_eq!(result.tool_config["toolChoice"], json!({ "any": {} })); + assert_eq!(result.tool_warnings, json!([])); + } + } + } + + mod strict_mode_for_function_tools { + use super::*; + + const STRICT_WARNING: &str = "but Amazon Bedrock requires every object in a strict tool schema to set additionalProperties: false. The strict property will be ignored."; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › strict mode for function tools › should pass through strict mode when strict is true + #[test] + fn should_pass_through_strict_mode_when_strict_is_true() { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([{ + "type": "function", + "name": "testFunction", + "description": "A test function", + "inputSchema": { + "type": "object", + "properties": {}, + "additionalProperties": false, + }, + "strict": true, + }])), + model_id: ANTHROPIC_MODEL, + ..Default::default() + }); + + assert_eq!( + result.tool_config["tools"], + json!([{ + "toolSpec": { + "name": "testFunction", + "description": "A test function", + "strict": true, + "inputSchema": { + "json": { + "type": "object", + "properties": {}, + "additionalProperties": false, + }, + }, + }, + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › strict mode for function tools › should keep strict mode enabled for %s + #[test] + fn should_keep_strict_mode_enabled_for() { + for model_id in [ + "anthropic.claude-sonnet-4-6-v1", + "us.anthropic.claude-haiku-4-5-20251001-v1:0", + ] { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([{ + "type": "function", + "name": "testFunction", + "description": "A test function", + "inputSchema": { + "type": "object", + "properties": {}, + "additionalProperties": false, + }, + "strict": true, + }])), + model_id, + ..Default::default() + }); + + assert_eq!(tool_spec(&result, 0)["strict"], json!(true), "{model_id}"); + assert_eq!(result.tool_warnings, json!([]), "{model_id}"); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › strict mode for function tools › should omit strict mode when the top-level object schema is open + #[test] + fn should_omit_strict_mode_when_the_top_level_object_schema_is_open() { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([{ + "type": "function", + "name": "testFunction", + "description": "A test function", + "inputSchema": { "type": "object", "properties": {} }, + "strict": true, + }])), + model_id: ANTHROPIC_MODEL, + ..Default::default() + }); + + assert!(!tool_spec(&result, 0).contains_key("strict")); + assert_eq!( + result.tool_warnings, + json!([{ + "type": "unsupported", + "feature": "strict", + "details": format!("Tool 'testFunction' has strict: true, {STRICT_WARNING}"), + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › strict mode for function tools › should omit strict mode when a nested object schema is open + #[test] + fn should_omit_strict_mode_when_a_nested_object_schema_is_open() { + let input_schema = json!({ + "type": "object", + "properties": { + "location": { + "type": "object", + "properties": { "city": { "type": "string" } }, + "required": ["city"], + }, + }, + "required": ["location"], + "additionalProperties": false, + }); + + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([{ + "type": "function", + "name": "getWeather", + "inputSchema": input_schema, + "strict": true, + }])), + model_id: ANTHROPIC_MODEL, + ..Default::default() + }); + + assert_eq!( + Value::Object(tool_spec(&result, 0).clone()), + json!({ "name": "getWeather", "inputSchema": { "json": input_schema } }) + ); + assert_eq!( + result.tool_warnings, + json!([{ + "type": "unsupported", + "feature": "strict", + "details": format!("Tool 'getWeather' has strict: true, {STRICT_WARNING}"), + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › strict mode for function tools › should omit strict mode when a referenced object schema in $defs is open + #[test] + fn should_omit_strict_mode_when_a_referenced_object_schema_in_defs_is_open() { + let input_schema = json!({ + "$ref": "#/$defs/location", + "$defs": { + "location": { + "type": "object", + "properties": { "city": { "type": "string" } }, + "required": ["city"], + }, + }, + }); + + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([{ + "type": "function", + "name": "getWeather", + "inputSchema": input_schema, + "strict": true, + }])), + model_id: ANTHROPIC_MODEL, + ..Default::default() + }); + + assert_eq!( + Value::Object(tool_spec(&result, 0).clone()), + json!({ "name": "getWeather", "inputSchema": { "json": input_schema } }) + ); + assert_eq!( + result.tool_warnings, + json!([{ + "type": "unsupported", + "feature": "strict", + "details": format!("Tool 'getWeather' has strict: true, {STRICT_WARNING}"), + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › strict mode for function tools › should pass through strict mode when all nested object schemas are closed + #[test] + fn should_pass_through_strict_mode_when_all_nested_object_schemas_are_closed() { + let input_schema = json!({ + "type": "object", + "properties": { + "location": { + "type": "object", + "properties": { "city": { "type": "string" } }, + "required": ["city"], + "additionalProperties": false, + }, + }, + "required": ["location"], + "additionalProperties": false, + }); + + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([{ + "type": "function", + "name": "getWeather", + "inputSchema": input_schema, + "strict": true, + }])), + model_id: ANTHROPIC_MODEL, + ..Default::default() + }); + + assert_eq!( + Value::Object(tool_spec(&result, 0).clone()), + json!({ + "name": "getWeather", + "strict": true, + "inputSchema": { "json": input_schema }, + }) + ); + assert_eq!(result.tool_warnings, json!([])); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › strict mode for function tools › should pass through strict mode when strict is false + #[test] + fn should_pass_through_strict_mode_when_strict_is_false() { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([{ + "type": "function", + "name": "testFunction", + "description": "A test function", + "inputSchema": { "type": "object", "properties": {} }, + "strict": false, + }])), + model_id: ANTHROPIC_MODEL, + ..Default::default() + }); + + assert_eq!( + result.tool_config["tools"], + json!([{ + "toolSpec": { + "name": "testFunction", + "description": "A test function", + "strict": false, + "inputSchema": { "json": { "type": "object", "properties": {} } }, + }, + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › strict mode for function tools › should not include strict when strict is undefined + #[test] + fn should_not_include_strict_when_strict_is_undefined() { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([{ + "type": "function", + "name": "testFunction", + "description": "A test function", + "inputSchema": { "type": "object", "properties": {} }, + }])), + model_id: ANTHROPIC_MODEL, + ..Default::default() + }); + + assert!(!tool_spec(&result, 0).contains_key("strict")); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › strict mode for function tools › should pass through strict mode for multiple tools with different strict settings + #[test] + fn should_pass_through_strict_mode_for_multiple_tools_with_different_strict_settings() { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([ + { + "type": "function", + "name": "strictTool", + "description": "A strict tool", + "inputSchema": { + "type": "object", + "properties": {}, + "additionalProperties": false, + }, + "strict": true, + }, + { + "type": "function", + "name": "nonStrictTool", + "description": "A non-strict tool", + "inputSchema": { "type": "object", "properties": {} }, + "strict": false, + }, + { + "type": "function", + "name": "defaultTool", + "description": "A tool without strict setting", + "inputSchema": { "type": "object", "properties": {} }, + }, + ])), + model_id: ANTHROPIC_MODEL, + ..Default::default() + }); + + assert_eq!(tool_spec(&result, 0)["strict"], json!(true)); + assert_eq!(tool_spec(&result, 1)["strict"], json!(false)); + assert!(!tool_spec(&result, 2).contains_key("strict")); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › strict mode for function tools › should warn when strict is omitted for %s + #[test] + fn should_warn_when_strict_is_omitted_for() { + for model_id in [ + "us.anthropic.claude-opus-4-7", + "anthropic.claude-opus-4-8", + "us.anthropic.claude-opus-5", + "anthropic.claude-sonnet-5", + "eu.anthropic.claude-fable-5", + "anthropic.claude-fable-5-1", + "us.anthropic.claude-fable-5-1", + "global.anthropic.claude-fable-5-1", + ] { + let result = prepare_tools(PrepareToolsArgs { + tools: Some(json!([{ + "type": "function", + "name": "testFunction", + "description": "A test function", + "inputSchema": { "type": "object", "properties": {} }, + "strict": true, + }])), + model_id, + ..Default::default() + }); + + assert!(!tool_spec(&result, 0).contains_key("strict"), "{model_id}"); + assert_eq!( + result.tool_warnings, + json!([{ + "type": "unsupported", + "feature": "strict", + "details": "Tool 'testFunction' has strict: true, but strict mode is not supported by this model on Amazon Bedrock. The strict property will be ignored.", + }]), + "{model_id}" + ); + } + } + } +} diff --git a/aimux-providers/src/bedrock/tests/provider.rs b/aimux-providers/src/bedrock/tests/provider.rs new file mode 100644 index 000000000..63ad932ab --- /dev/null +++ b/aimux-providers/src/bedrock/tests/provider.rs @@ -0,0 +1,892 @@ +//! Ports of `@ai-sdk/amazon-bedrock` 5.0.100 +//! `src/amazon-bedrock-provider.test.ts`. +//! +//! Upstream mocks the model classes, the fetch factories and the settings +//! loaders, then inspects the arguments the provider passes them. aimux has +//! none of these seams (the models take a private `EndpointConfig`, and the +//! transport is chosen per request), so each case makes one request through +//! the provider with the mock transport and asserts what upstream's mocks +//! would have received, as it shows on the wire: the `baseUrl()` is the +//! request URL's prefix, `headers()` are the request headers, an API-key +//! fetch is an unsigned request with `Authorization: Bearer `, and a +//! SigV4 fetch is a signed one whose credential names the access key. +//! +//! Upstream's `loadSetting` mock returns the setting or `us-east-1`, and its +//! `loadOptionalSetting` mock returns the setting and never the environment. +//! The cases that depend on the environment emulate that with environment +//! variables (`AWS_REGION`, `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY` set +//! to `us-east-1`, the optional variables unset; a case's own +//! `mockImplementation` sets more) and are `#[serial]`. +//! +//! Upstream mocks the version as `0.0.0-test`; the user-agent part checked +//! here is the pinned `ai-sdk-amazon-bedrock/5.0.100`. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use serde_json::json; + +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::image_model::{ImageCallOptions, ImageModel}; +use aimux_core::language_model_message::{LanguageModelMessage, TextPart, UserPart}; +use aimux_core::options::CallOptions; +use aimux_core::provider::Provider; +use aimux_core::reranking_model::{RerankingCallOptions, RerankingDocuments, RerankingModel}; +use aimux_provider_utils::{AwsCredentials, FetchFunction, HeaderMapOpt, Resolvable}; + +use crate::bedrock::{ + AmazonBedrockChatModelFamily, AmazonBedrockChatModelSettings, + AmazonBedrockEmbeddingModelFamily, AmazonBedrockEmbeddingModelSettings, AmazonBedrockProvider, + AmazonBedrockProviderSettings, create_amazon_bedrock, +}; +use crate::test_support::mock_fetch::{Canned, EnvVar, MockFetch, Seen}; + +/// The pinned `ai-sdk-amazon-bedrock/` user-agent part. +const USER_AGENT_PART: &str = "ai-sdk-amazon-bedrock/5.0.100"; + +/// The environment variables the provider reads. +const ENV_VARS: [&str; 8] = [ + "AWS_REGION", + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "AWS_BEARER_TOKEN_BEDROCK", + "AWS_ENDPOINT_URL", + "AWS_ENDPOINT_URL_BEDROCK_RUNTIME", + "AWS_ENDPOINT_URL_BEDROCK_AGENT_RUNTIME", +]; + +/// Sets the environment the upstream settings mocks describe, with +/// `overrides` (`None` unsets a variable), and restores it when dropped. +#[must_use] +fn mocked_env(overrides: &[(&'static str, Option<&str>)]) -> Vec { + ENV_VARS + .iter() + .map(|&name| { + let base = match name { + // `loadSetting` returns `settingValue || 'us-east-1'`. + "AWS_REGION" | "AWS_ACCESS_KEY_ID" | "AWS_SECRET_ACCESS_KEY" => Some("us-east-1"), + // `loadOptionalSetting` returns only `settingValue`. + _ => None, + }; + let value = overrides + .iter() + .find(|(overridden, _)| *overridden == name) + .map_or(base, |(_, value)| *value); + EnvVar::set(name, value) + }) + .collect() +} + +/// The `loadOptionalSetting(() => undefined)` / no-environment cases: only +/// the explicit settings count. +#[must_use] +fn empty_env() -> Vec { + ENV_VARS + .iter() + .map(|&name| EnvVar::set(name, None)) + .collect() +} + +/// A Converse response. +fn converse_response() -> Canned { + Canned::json(&json!({ + "output": { "message": { "role": "assistant", "content": [{ "text": "ok" }] } }, + "stopReason": "end_turn", + "usage": { "inputTokens": 1, "outputTokens": 1, "totalTokens": 2 }, + })) +} + +fn chat_options() -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::User { + content: vec![UserPart::Text(TextPart { + text: "Hello".into(), + provider_options: None, + })], + provider_options: None, + }]) +} + +/// `create_amazon_bedrock(settings)` with the mock transport. +fn provider( + settings: AmazonBedrockProviderSettings, + mock: &Arc, +) -> AmazonBedrockProvider { + let fetch: FetchFunction = mock.transport(); + create_amazon_bedrock(AmazonBedrockProviderSettings { + fetch: Some(fetch), + ..settings + }) + .expect("valid Bedrock settings") +} + +/// One `doGenerate` of `provider(model_id)`; the request it sends. +async fn generate(provider: &AmazonBedrockProvider, mock: &Arc, model_id: &str) -> Seen { + let model = provider.call(model_id); + assert_eq!(model.model_id(), model_id); + model + .do_generate(&chat_options()) + .await + .expect("generation succeeds"); + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + seen.into_iter().next().expect("one request") +} + +/// One `doEmbed` of `provider.embedding(model_id)`; the request it sends. +async fn embed(provider: &AmazonBedrockProvider, mock: &Arc, model_id: &str) -> Seen { + let model = provider.embedding(model_id); + assert_eq!(model.model_id(), model_id); + model + .do_embed(&EmbeddingCallOptions::new("x")) + .await + .expect("embedding succeeds"); + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + seen.into_iter().next().expect("one request") +} + +fn titan_response() -> Canned { + Canned::json(&json!({ "embedding": [0.1], "inputTextTokenCount": 1 })) +} + +fn rerank_response() -> Canned { + Canned::json(&json!({ "results": [{ "index": 0, "relevanceScore": 0.5 }] })) +} + +/// One `doRerank` of `provider.reranking(model_id)`; the request it sends. +async fn rerank(provider: &AmazonBedrockProvider, mock: &Arc, model_id: &str) -> Seen { + let model = provider.reranking(model_id); + assert_eq!(model.model_id(), model_id); + model + .do_rerank(&RerankingCallOptions::new( + "query", + RerankingDocuments::Text { + values: vec!["document".to_string()], + }, + )) + .await + .expect("reranking succeeds"); + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + seen.into_iter().next().expect("one request") +} + +/// The base URL a Converse request of `model_id` was sent to. +fn base_url_of(seen: &Seen, path: &str) -> String { + seen.url + .strip_suffix(path) + .unwrap_or_else(|| panic!("{} does not end with {path}", seen.url)) + .to_string() +} + +fn converse_base_url(seen: &Seen) -> String { + base_url_of(seen, "/model/anthropic.claude-v2/converse") +} + +/// `headers()['user-agent']` contains `part`. +fn assert_user_agent_contains(headers: &BTreeMap, part: &str) { + let user_agent = headers.get("user-agent").expect("user-agent is sent"); + assert!( + user_agent.contains(part), + "user-agent {user_agent} lacks {part}" + ); +} + +/// The request went through the API-key fetch: `Bearer `, not signed. +fn assert_api_key(seen: &Seen, api_key: &str) { + assert_eq!( + seen.headers.get("authorization").map(String::as_str), + Some(format!("Bearer {api_key}").as_str()) + ); + assert!( + !seen.headers.contains_key("x-amz-date"), + "{:?}", + seen.headers + ); +} + +/// The request went through the SigV4 fetch, signed with `access_key_id`. +fn assert_sigv4(seen: &Seen, access_key_id: &str) { + let authorization = seen.headers.get("authorization").expect("authorization"); + assert!( + authorization.starts_with(&format!("AWS4-HMAC-SHA256 Credential={access_key_id}/")), + "unexpected authorization: {authorization}" + ); + assert!( + seen.headers.contains_key("x-amz-date"), + "{:?}", + seen.headers + ); +} + +fn custom_headers() -> Option { + Some(HeaderMapOpt::from([( + "custom-header".to_string(), + Some("value".to_string()), + )])) +} + +/// A credential provider that resolves dynamic credentials (no region, as +/// upstream's). +fn dynamic_credential_provider() -> Resolvable { + Resolvable::from_async_fn(|| async { + Ok(AwsCredentials { + access_key_id: "dynamic-access-key".to_string(), + secret_access_key: "dynamic-secret-key".to_string(), + session_token: Some("dynamic-session-token".to_string()), + region: String::new(), + }) + }) +} + +mod amazon_bedrock_provider { + use super::*; + + mod create_amazon_bedrock { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › should create a provider instance with default options + #[tokio::test] + #[serial_test::serial] + async fn should_create_a_provider_instance_with_default_options() { + let _env = mocked_env(&[]); + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider(AmazonBedrockProviderSettings::default(), &mock); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_user_agent_contains(&seen.headers, USER_AGENT_PART); + assert_eq!( + converse_base_url(&seen), + "https://bedrock-runtime.us-east-1.amazonaws.com" + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › should create a provider instance with custom options + #[tokio::test] + #[serial_test::serial] + async fn should_create_a_provider_instance_with_custom_options() { + let _env = mocked_env(&[]); + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + region: Some("eu-west-1".to_string()), + base_url: Some("https://custom.url".to_string()), + headers: custom_headers(), + ..Default::default() + }, + &mock, + ); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_eq!(seen.headers["custom-header"], "value"); + assert_user_agent_contains(&seen.headers, USER_AGENT_PART); + assert_eq!(converse_base_url(&seen), "https://custom.url"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › resolves the Bedrock Runtime endpoint for %s + #[tokio::test] + #[serial_test::serial] + async fn resolves_the_bedrock_runtime_endpoint_for_region() { + let _env = mocked_env(&[]); + for (region, dns_suffix) in [ + ("cn-north-1", "amazonaws.com.cn"), + ("us-gov-west-1", "amazonaws.com"), + ("us-iso-east-1", "c2s.ic.gov"), + ("us-isob-east-1", "sc2s.sgov.gov"), + ("eu-isoe-west-1", "cloud.adc-e.uk"), + ("us-isof-south-1", "csp.hci.ic.gov"), + ("eusc-de-east-1", "amazonaws.eu"), + ] { + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + region: Some(region.to_string()), + ..Default::default() + }, + &mock, + ); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_eq!( + converse_base_url(&seen), + format!("https://bedrock-runtime.{region}.{dns_suffix}") + ); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › uses an explicit base URL without loading a region + /// + /// `loadSetting` not being called is the request succeeding with no + /// region anywhere. + #[tokio::test] + #[serial_test::serial] + async fn uses_an_explicit_base_url_without_loading_a_region() { + let _env = mocked_env(&[ + ( + "AWS_ENDPOINT_URL_BEDROCK_RUNTIME", + Some("https://runtime.example.com"), + ), + ("AWS_ENDPOINT_URL", Some("https://global.example.com")), + ("AWS_REGION", None), + ]); + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + api_key: Some("test-api-key".to_string().into()), + base_url: Some("https://explicit.example.com/".to_string()), + ..Default::default() + }, + &mock, + ); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_eq!(converse_base_url(&seen), "https://explicit.example.com"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › prefers the service-specific Bedrock Runtime endpoint over the global endpoint + #[tokio::test] + #[serial_test::serial] + async fn prefers_the_service_specific_bedrock_runtime_endpoint_over_the_global_endpoint() { + let _env = mocked_env(&[ + ( + "AWS_ENDPOINT_URL_BEDROCK_RUNTIME", + Some("https://runtime.example.com/"), + ), + ("AWS_ENDPOINT_URL", Some("https://global.example.com")), + ("AWS_REGION", None), + ]); + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + api_key: Some("test-api-key".to_string().into()), + ..Default::default() + }, + &mock, + ); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_eq!(converse_base_url(&seen), "https://runtime.example.com"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › uses the global endpoint override when no service-specific endpoint is set + #[tokio::test] + #[serial_test::serial] + async fn uses_the_global_endpoint_override_when_no_service_specific_endpoint_is_set() { + let _env = mocked_env(&[ + ("AWS_ENDPOINT_URL", Some("https://global.example.com/")), + ("AWS_REGION", None), + ]); + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + api_key: Some("test-api-key".to_string().into()), + ..Default::default() + }, + &mock, + ); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_eq!(converse_base_url(&seen), "https://global.example.com"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › uses a separate service-specific endpoint for Bedrock Agent Runtime + #[tokio::test] + #[serial_test::serial] + async fn uses_a_separate_service_specific_endpoint_for_bedrock_agent_runtime() { + let _env = mocked_env(&[ + ( + "AWS_ENDPOINT_URL_BEDROCK_RUNTIME", + Some("https://runtime.example.com"), + ), + ( + "AWS_ENDPOINT_URL_BEDROCK_AGENT_RUNTIME", + Some("https://agent-runtime.example.com"), + ), + ]); + let mock = MockFetch::new(vec![rerank_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + region: Some("us-east-1".to_string()), + ..Default::default() + }, + &mock, + ); + + let seen = rerank(&provider, &mock, "amazon.rerank-v1:0").await; + + assert_eq!( + base_url_of(&seen, "/rerank"), + "https://agent-runtime.example.com" + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › resolves the Bedrock Agent Runtime endpoint for a non-standard AWS partition + #[tokio::test] + #[serial_test::serial] + async fn resolves_the_bedrock_agent_runtime_endpoint_for_a_non_standard_aws_partition() { + let _env = mocked_env(&[]); + let mock = MockFetch::new(vec![rerank_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + region: Some("eu-isoe-west-1".to_string()), + ..Default::default() + }, + &mock, + ); + + let seen = rerank(&provider, &mock, "amazon.rerank-v1:0").await; + + assert_eq!( + base_url_of(&seen, "/rerank"), + "https://bedrock-agent-runtime.eu-isoe-west-1.cloud.adc-e.uk" + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › should accept a credentialProvider in options + #[tokio::test] + #[serial_test::serial] + async fn should_accept_a_credential_provider_in_options() { + let _env = mocked_env(&[]); + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + credential_provider: Some(dynamic_credential_provider()), + ..Default::default() + }, + &mock, + ); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_user_agent_contains(&seen.headers, USER_AGENT_PART); + assert_eq!( + converse_base_url(&seen), + "https://bedrock-runtime.us-east-1.amazonaws.com" + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › should prioritize credentialProvider over static credentials + #[test] + fn should_prioritize_credential_provider_over_static_credentials() { + let provider = create_amazon_bedrock(AmazonBedrockProviderSettings { + access_key_id: Some("static-access-key".to_string()), + secret_access_key: Some("static-secret-key".to_string()), + session_token: Some("static-session-token".to_string()), + credential_provider: Some(dynamic_credential_provider()), + ..Default::default() + }) + .expect("valid Bedrock settings"); + + let model = provider.call("anthropic.claude-v2"); + + assert_eq!(model.model_id(), "anthropic.claude-v2"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › should pass headers to embedding model + #[tokio::test] + #[serial_test::serial] + async fn should_pass_headers_to_embedding_model() { + let _env = mocked_env(&[]); + let mock = MockFetch::new(vec![titan_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + headers: custom_headers(), + ..Default::default() + }, + &mock, + ); + + let seen = embed(&provider, &mock, "amazon.titan-embed-text-v1").await; + + assert_eq!(seen.headers["custom-header"], "value"); + assert_user_agent_contains(&seen.headers, USER_AGENT_PART); + } + + // Not ported (ts-only): "should throw error when called with new + // keyword". A Rust provider is not a JS function and has no `new`. + + mod api_key_authentication { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › API Key Authentication › should use API key when provided in options + #[tokio::test] + #[serial_test::serial] + async fn should_use_api_key_when_provided_in_options() { + let _env = mocked_env(&[]); + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + api_key: Some("test-api-key".to_string().into()), + region: Some("us-east-1".to_string()), + ..Default::default() + }, + &mock, + ); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_api_key(&seen, "test-api-key"); + assert_user_agent_contains(&seen.headers, "ai-sdk-amazon-bedrock/"); + assert_eq!( + converse_base_url(&seen), + "https://bedrock-runtime.us-east-1.amazonaws.com" + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › API Key Authentication › should use API key from environment variable + #[tokio::test] + #[serial_test::serial] + async fn should_use_api_key_from_the_environment_variable() { + let _env = mocked_env(&[("AWS_BEARER_TOKEN_BEDROCK", Some("env-api-key"))]); + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + region: Some("us-east-1".to_string()), + ..Default::default() + }, + &mock, + ); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_api_key(&seen, "env-api-key"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › API Key Authentication › should prioritize options.apiKey over environment variable + #[tokio::test] + #[serial_test::serial] + async fn should_prioritize_options_api_key_over_environment_variable() { + let _env = mocked_env(&[("AWS_BEARER_TOKEN_BEDROCK", Some("env-api-key"))]); + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + api_key: Some("options-api-key".to_string().into()), + region: Some("us-east-1".to_string()), + ..Default::default() + }, + &mock, + ); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_api_key(&seen, "options-api-key"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › API Key Authentication › should fall back to SigV4 when no API key provided + #[tokio::test] + #[serial_test::serial] + async fn should_fall_back_to_sig_v4_when_no_api_key_provided() { + let _env = empty_env(); + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + region: Some("us-east-1".to_string()), + access_key_id: Some("test-access-key".to_string()), + secret_access_key: Some("test-secret-key".to_string()), + ..Default::default() + }, + &mock, + ); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_sigv4(&seen, "test-access-key"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › API Key Authentication › should pass custom fetch function to API key authentication + /// + /// The custom fetch is the mock transport: the API-key request + /// reaches it. + #[tokio::test] + #[serial_test::serial] + async fn should_pass_custom_fetch_function_to_api_key_authentication() { + let _env = mocked_env(&[]); + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + api_key: Some("test-api-key".to_string().into()), + region: Some("us-east-1".to_string()), + ..Default::default() + }, + &mock, + ); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_api_key(&seen, "test-api-key"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › API Key Authentication › should pass custom fetch function to SigV4 authentication + /// + /// The custom fetch is the mock transport: the signed request + /// reaches it. + #[tokio::test] + #[serial_test::serial] + async fn should_pass_custom_fetch_function_to_sig_v4_authentication() { + let _env = empty_env(); + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + region: Some("us-east-1".to_string()), + access_key_id: Some("test-access-key".to_string()), + secret_access_key: Some("test-secret-key".to_string()), + ..Default::default() + }, + &mock, + ); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_sigv4(&seen, "test-access-key"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › API Key Authentication › should work with embedding models when using API key + #[tokio::test] + #[serial_test::serial] + async fn should_work_with_embedding_models_when_using_api_key() { + let _env = mocked_env(&[]); + let mock = MockFetch::new(vec![titan_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + api_key: Some("test-api-key".to_string().into()), + region: Some("us-east-1".to_string()), + headers: custom_headers(), + ..Default::default() + }, + &mock, + ); + + let seen = embed(&provider, &mock, "amazon.titan-embed-text-v1").await; + + assert_eq!(seen.headers["custom-header"], "value"); + assert_user_agent_contains(&seen.headers, USER_AGENT_PART); + assert_api_key(&seen, "test-api-key"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › API Key Authentication › should work with image models when using API key + #[tokio::test] + #[serial_test::serial] + async fn should_work_with_image_models_when_using_api_key() { + let _env = mocked_env(&[]); + let mock = MockFetch::new(vec![Canned::json(&json!({ "images": ["image"] }))]); + let provider = provider( + AmazonBedrockProviderSettings { + api_key: Some("test-api-key".to_string().into()), + region: Some("us-east-1".to_string()), + headers: custom_headers(), + ..Default::default() + }, + &mock, + ); + + let model = provider.image("amazon.titan-image-generator"); + assert_eq!(model.model_id(), "amazon.titan-image-generator"); + model + .do_generate(&ImageCallOptions::new("prompt")) + .await + .expect("generation succeeds"); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].headers["custom-header"], "value"); + assert_user_agent_contains(&seen[0].headers, USER_AGENT_PART); + assert_api_key(&seen[0], "test-api-key"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › API Key Authentication › should maintain backward compatibility with existing SigV4 authentication + #[tokio::test] + #[serial_test::serial] + async fn should_maintain_backward_compatibility_with_existing_sig_v4_authentication() { + let _env = empty_env(); + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + region: Some("eu-west-1".to_string()), + access_key_id: Some("test-access-key".to_string()), + secret_access_key: Some("test-secret-key".to_string()), + session_token: Some("test-session-token".to_string()), + ..Default::default() + }, + &mock, + ); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_sigv4(&seen, "test-access-key"); + assert_eq!( + converse_base_url(&seen), + "https://bedrock-runtime.eu-west-1.amazonaws.com" + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › API Key Authentication › does not use AWS_SESSION_TOKEN from env when both access keys are passed as options + /// + /// The resolved credentials show as the signature's access key and + /// the (absent) `x-amz-security-token`; the secret key is not + /// observable on the wire. + #[tokio::test] + #[serial_test::serial] + async fn does_not_use_aws_session_token_from_the_env_when_both_access_keys_are_passed() + { + let _env = mocked_env(&[( + "AWS_SESSION_TOKEN", + Some("env-session-token-should-not-be-used"), + )]); + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + region: Some("us-east-1".to_string()), + access_key_id: Some("from-options-ak".to_string()), + secret_access_key: Some("from-options-sk".to_string()), + ..Default::default() + }, + &mock, + ); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_sigv4(&seen, "from-options-ak"); + assert!( + !seen.headers.contains_key("x-amz-security-token"), + "{:?}", + seen.headers + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › API Key Authentication › uses options.sessionToken when both access keys are passed as options + #[tokio::test] + #[serial_test::serial] + async fn uses_options_session_token_when_both_access_keys_are_passed_as_options() { + let _env = empty_env(); + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + region: Some("us-east-1".to_string()), + access_key_id: Some("from-options-ak".to_string()), + secret_access_key: Some("from-options-sk".to_string()), + session_token: Some("explicit-session".to_string()), + ..Default::default() + }, + &mock, + ); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_eq!(seen.headers["x-amz-security-token"], "explicit-session"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › createAmazonBedrock › API Key Authentication › should work with credential provider when no API key is provided + #[tokio::test] + #[serial_test::serial] + async fn should_work_with_credential_provider_when_no_api_key_is_provided() { + let _env = empty_env(); + let mock = MockFetch::new(vec![converse_response()]); + let provider = provider( + AmazonBedrockProviderSettings { + region: Some("us-east-1".to_string()), + credential_provider: Some(dynamic_credential_provider()), + ..Default::default() + }, + &mock, + ); + + let seen = generate(&provider, &mock, "anthropic.claude-v2").await; + + assert_sigv4(&seen, "dynamic-access-key"); + } + } + } + + mod provider_methods { + use super::*; + + const APPLICATION_PROFILE_ARN: &str = + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/qibm5eutlkcy"; + + fn default_provider() -> AmazonBedrockProvider { + create_amazon_bedrock(AmazonBedrockProviderSettings::default()) + .expect("default settings are valid") + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › provider methods › should create an embedding model + /// + /// `toBeInstanceOf` is the return type, `BedrockEmbeddingModel`. + #[test] + fn should_create_an_embedding_model() { + let model_id = "amazon.titan-embed-text-v1"; + + let model: crate::bedrock::BedrockEmbeddingModel = + default_provider().embedding(model_id); + + assert_eq!(model.model_id(), model_id); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › provider methods › should pass chat model settings to the model + #[test] + fn should_pass_chat_model_settings_to_the_model() { + let model = default_provider().chat_with_settings( + APPLICATION_PROFILE_ARN, + &AmazonBedrockChatModelSettings { + model_family: Some(AmazonBedrockChatModelFamily::Anthropic), + }, + ); + + assert_eq!(model.model_family.as_deref(), Some("anthropic")); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › provider methods › should pass embedding model settings to the model + #[test] + fn should_pass_embedding_model_settings_to_the_model() { + let model = default_provider().embedding_with_settings( + APPLICATION_PROFILE_ARN, + &AmazonBedrockEmbeddingModelSettings { + model_family: Some(AmazonBedrockEmbeddingModelFamily::Cohere), + }, + ); + + assert_eq!( + model.model_family, + Some(AmazonBedrockEmbeddingModelFamily::Cohere) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › provider methods › should create an image model + /// + /// `toBeInstanceOf` is the return type, `BedrockImageModel`. + #[test] + fn should_create_an_image_model() { + let model_id = "amazon.titan-image-generator"; + + let model: crate::bedrock::BedrockImageModel = default_provider().image(model_id); + + assert_eq!(model.model_id(), model_id); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-provider.test.ts › AmazonBedrockProvider › provider methods › should create an image model via imageModel method + /// + /// `imageModel` is `Provider::image_model`, which returns the model as + /// a trait object; its instance type is its provider and model id. + #[test] + fn should_create_an_image_model_via_image_model_method() { + let model_id = "amazon.titan-image-generator"; + + let model = default_provider() + .image_model(model_id) + .expect("an image model"); + + assert_eq!(model.model_id(), model_id); + assert_eq!(model.provider(), "amazon-bedrock"); + } + + // Not ported (sub-feature): "should expose anthropicTools". + // `AmazonBedrockProvider` has no `tools` (the Anthropic provider tools) + // member. + } +} diff --git a/aimux-providers/src/bedrock/tests/reranking/fixtures/amazon-bedrock-reranking.1.json b/aimux-providers/src/bedrock/tests/reranking/fixtures/amazon-bedrock-reranking.1.json new file mode 100644 index 000000000..311752bb9 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/reranking/fixtures/amazon-bedrock-reranking.1.json @@ -0,0 +1,12 @@ +{ + "results": [ + { + "index": 0, + "relevanceScore": 0.5110583305358887 + }, + { + "index": 5, + "relevanceScore": 0.30241215229034424 + } + ] +} diff --git a/aimux-providers/src/bedrock/tests/reranking/mod.rs b/aimux-providers/src/bedrock/tests/reranking/mod.rs new file mode 100644 index 000000000..b3b703cab --- /dev/null +++ b/aimux-providers/src/bedrock/tests/reranking/mod.rs @@ -0,0 +1 @@ +mod reranking_model; diff --git a/aimux-providers/src/bedrock/tests/reranking/reranking_model.rs b/aimux-providers/src/bedrock/tests/reranking/reranking_model.rs new file mode 100644 index 000000000..600322614 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/reranking/reranking_model.rs @@ -0,0 +1,391 @@ +//! Ports of `@ai-sdk/amazon-bedrock` 5.0.100 +//! `src/reranking/amazon-bedrock-reranking-model.test.ts`. +//! +//! Upstream builds `AmazonBedrockRerankingModel` directly with a base URL, a +//! region, plain config headers and `fetch: injectFetchHeaders(..)`, so no +//! request is signed. The port does the same: `BedrockRerankingModel::from_config` +//! with a fixed endpoint, the config headers, an `Auth` that names only the +//! region, and the `inject_fetch_headers` transport over the mock transport +//! (which stands in for the test server). Each upstream `beforeEach` request is +//! made by the test itself. +//! +//! The test server's recorded request headers leave out `user-agent`; it is +//! compared separately. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use async_trait::async_trait; +use serde_json::{Value, json}; + +use aimux_core::error::AiMuxError; +use aimux_core::reranking_model::{ + RerankingCallOptions, RerankingDocuments, RerankingModel, RerankingResult, +}; +use aimux_provider_utils::HeaderMapOpt; + +use crate::bedrock::{Auth, BedrockRerankingModel}; +use crate::shared::{Endpoint, EndpointConfig, ResolveEndpoint}; +use crate::test_support::mock_fetch::{Canned, MockFetch, Seen}; + +use super::super::support::{BEDROCK_USER_AGENT_PART, inject_fetch_headers}; + +const RERANK_URL: &str = "https://bedrock-agent-runtime.us-east-1.amazonaws.com/rerank"; + +/// A fixed base URL and fixed headers: the `{ baseUrl, headers }` config of +/// the upstream model. +struct FixedEndpoint { + base_url: String, + headers: HeaderMapOpt, +} + +#[async_trait] +impl ResolveEndpoint for FixedEndpoint { + async fn resolve(&self) -> Result { + Ok(Endpoint { + base_url: self.base_url.clone(), + headers: self.headers.clone(), + }) + } +} + +/// The upstream `model`: `cohere.rerank-v3-5:0` on the us-east-1 Agent +/// Runtime, region `us-west-2`, with `fakeFetchWithAuth`. +fn model(mock: &Arc) -> BedrockRerankingModel { + BedrockRerankingModel::from_config( + "cohere.rerank-v3-5:0".to_string(), + EndpointConfig::new( + "amazon-bedrock".to_string(), + FixedEndpoint { + base_url: "https://bedrock-agent-runtime.us-east-1.amazonaws.com".to_string(), + headers: HeaderMapOpt::from([ + ( + "config-header".to_string(), + Some("config-value".to_string()), + ), + ( + "shared-header".to_string(), + Some("config-shared".to_string()), + ), + ]), + }, + Some(inject_fetch_headers( + &[("x-amz-auth", "test-auth")], + mock.transport(), + )), + ), + Arc::new(Auth { + region: Some("us-west-2".to_string()), + api_key: None, + access_key_id: None, + secret_access_key: None, + session_token: None, + credential_provider: None, + }), + ) +} + +/// `prepareJsonFixtureResponse(filename)`. +fn prepare_json_fixture_response() -> Canned { + let body = include_str!("fixtures/amazon-bedrock-reranking.1.json"); + Canned { + status: 200, + headers: vec![ + ("content-type".to_string(), "application/json".to_string()), + ("content-length".to_string(), body.len().to_string()), + ], + body: body.as_bytes().to_vec(), + } +} + +/// The `beforeEach` call of a describe: rerank `documents` for `rainy day` +/// with the upstream provider options. +async fn do_rerank(documents: RerankingDocuments) -> (RerankingResult, Seen) { + let mock = MockFetch::new(vec![prepare_json_fixture_response()]); + let result = model(&mock) + .do_rerank(&RerankingCallOptions { + top_n: Some(2), + provider_options: Some( + aimux_core::shared::provider_namespace( + "bedrock", + json!({ + "nextToken": "test-token", + "additionalModelRequestFields": { "test": "test-value" }, + }), + ) + .expect("options are an object"), + ), + ..RerankingCallOptions::new("rainy day", documents) + }) + .await + .expect("reranking succeeds"); + let mut seen = mock.seen(); + assert_eq!(seen.len(), 1); + let seen = seen.remove(0); + assert_eq!(seen.url, RERANK_URL); + (result, seen) +} + +/// The request headers the test server records: all but `user-agent`, which +/// is checked to carry the package part. +fn recorded_headers(seen: &Seen) -> BTreeMap { + let mut headers = seen.headers.clone(); + let user_agent = headers.remove("user-agent").expect("user-agent is sent"); + assert!( + user_agent + .split(' ') + .any(|part| part == BEDROCK_USER_AGENT_PART), + "user-agent {user_agent} lacks {BEDROCK_USER_AGENT_PART}" + ); + headers +} + +fn expected_ranking() -> Value { + json!([ + { "index": 0, "relevanceScore": 0.5110583305358887 }, + { "index": 5, "relevanceScore": 0.30241215229034424 }, + ]) +} + +fn expected_response() -> Value { + json!({ + "body": { + "results": [ + { "index": 0, "relevanceScore": 0.5110583305358887 }, + { "index": 5, "relevanceScore": 0.30241215229034424 }, + ], + }, + "headers": { + "content-length": "171", + "content-type": "application/json", + }, + }) +} + +fn expected_headers() -> BTreeMap { + BTreeMap::from([ + ("config-header".to_string(), "config-value".to_string()), + ("content-type".to_string(), "application/json".to_string()), + ("shared-header".to_string(), "config-shared".to_string()), + ("x-amz-auth".to_string(), "test-auth".to_string()), + ]) +} + +/// The request body of both describes, with `sources`. +fn expected_body(sources: Value) -> Value { + json!({ + "nextToken": "test-token", + "queries": [ + { + "textQuery": { "text": "rainy day" }, + "type": "TEXT", + }, + ], + "rerankingConfiguration": { + "bedrockRerankingConfiguration": { + "modelConfiguration": { + "additionalModelRequestFields": { "test": "test-value" }, + "modelArn": "arn:aws:bedrock:us-west-2::foundation-model/cohere.rerank-v3-5:0", + }, + "numberOfResults": 2, + }, + "type": "BEDROCK_RERANKING_MODEL", + }, + "sources": sources, + }) +} + +mod do_rerank { + use super::*; + + mod json_documents { + use super::*; + + fn documents() -> RerankingDocuments { + RerankingDocuments::Object { + values: vec![ + json!({ "example": "sunny day at the beach" }), + json!({ "example": "rainy day in the city" }), + ], + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/reranking/amazon-bedrock-reranking-model.test.ts › doRerank › json documents › should send request with stringified json documents + #[tokio::test] + async fn should_send_request_with_stringified_json_documents() { + let (_result, seen) = do_rerank(documents()).await; + + assert_eq!( + seen.json_body(), + expected_body(json!([ + { + "inlineDocumentSource": { + "jsonDocument": { "example": "sunny day at the beach" }, + "type": "JSON", + }, + "type": "INLINE", + }, + { + "inlineDocumentSource": { + "jsonDocument": { "example": "rainy day in the city" }, + "type": "JSON", + }, + "type": "INLINE", + }, + ])) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/reranking/amazon-bedrock-reranking-model.test.ts › doRerank › json documents › should send the AWS-required `bedrockRerankingConfiguration` wire key + #[tokio::test] + async fn should_send_the_aws_required_bedrock_reranking_configuration_wire_key() { + let (_result, seen) = do_rerank(documents()).await; + + let body = seen.json_body(); + assert!( + body.pointer("/rerankingConfiguration/bedrockRerankingConfiguration") + .is_some(), + "{body}" + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/reranking/amazon-bedrock-reranking-model.test.ts › doRerank › json documents › should send request with the correct headers + #[tokio::test] + async fn should_send_request_with_the_correct_headers() { + let (_result, seen) = do_rerank(documents()).await; + + assert_eq!(recorded_headers(&seen), expected_headers()); + } + + /// upstream: @ai-sdk/amazon-bedrock src/reranking/amazon-bedrock-reranking-model.test.ts › doRerank › json documents › should return result with warnings + #[tokio::test] + async fn should_return_result_with_warnings() { + let (result, _seen) = do_rerank(documents()).await; + + assert!(result.warnings.is_none(), "{:?}", result.warnings); + } + + /// upstream: @ai-sdk/amazon-bedrock src/reranking/amazon-bedrock-reranking-model.test.ts › doRerank › json documents › should return result with the correct ranking + #[tokio::test] + async fn should_return_result_with_the_correct_ranking() { + let (result, _seen) = do_rerank(documents()).await; + + assert_eq!( + serde_json::to_value(&result.ranking).expect("ranking serializes"), + expected_ranking() + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/reranking/amazon-bedrock-reranking-model.test.ts › doRerank › json documents › should not return provider metadata (use response body instead) + #[tokio::test] + async fn should_not_return_provider_metadata_use_response_body_instead() { + let (result, _seen) = do_rerank(documents()).await; + + assert!( + result.provider_metadata.is_none(), + "{:?}", + result.provider_metadata + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/reranking/amazon-bedrock-reranking-model.test.ts › doRerank › json documents › should return result with the correct response + #[tokio::test] + async fn should_return_result_with_the_correct_response() { + let (result, _seen) = do_rerank(documents()).await; + + assert_eq!( + serde_json::to_value(&result.response).expect("response serializes"), + expected_response() + ); + } + } + + mod text_documents { + use super::*; + + fn documents() -> RerankingDocuments { + RerankingDocuments::Text { + values: vec![ + "sunny day at the beach".to_string(), + "rainy day in the city".to_string(), + ], + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/reranking/amazon-bedrock-reranking-model.test.ts › doRerank › text documents › should send request with text documents + #[tokio::test] + async fn should_send_request_with_text_documents() { + let (_result, seen) = do_rerank(documents()).await; + + assert_eq!( + seen.json_body(), + expected_body(json!([ + { + "inlineDocumentSource": { + "textDocument": { "text": "sunny day at the beach" }, + "type": "TEXT", + }, + "type": "INLINE", + }, + { + "inlineDocumentSource": { + "textDocument": { "text": "rainy day in the city" }, + "type": "TEXT", + }, + "type": "INLINE", + }, + ])) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/reranking/amazon-bedrock-reranking-model.test.ts › doRerank › text documents › should send request with the correct headers + #[tokio::test] + async fn should_send_request_with_the_correct_headers() { + let (_result, seen) = do_rerank(documents()).await; + + assert_eq!(recorded_headers(&seen), expected_headers()); + } + + /// upstream: @ai-sdk/amazon-bedrock src/reranking/amazon-bedrock-reranking-model.test.ts › doRerank › text documents › should return result without warnings + #[tokio::test] + async fn should_return_result_without_warnings() { + let (result, _seen) = do_rerank(documents()).await; + + assert!(result.warnings.is_none(), "{:?}", result.warnings); + } + + /// upstream: @ai-sdk/amazon-bedrock src/reranking/amazon-bedrock-reranking-model.test.ts › doRerank › text documents › should return result with the correct ranking + #[tokio::test] + async fn should_return_result_with_the_correct_ranking() { + let (result, _seen) = do_rerank(documents()).await; + + assert_eq!( + serde_json::to_value(&result.ranking).expect("ranking serializes"), + expected_ranking() + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/reranking/amazon-bedrock-reranking-model.test.ts › doRerank › text documents › should not return provider metadata (use response body instead) + #[tokio::test] + async fn should_not_return_provider_metadata_use_response_body_instead() { + let (result, _seen) = do_rerank(documents()).await; + + assert!( + result.provider_metadata.is_none(), + "{:?}", + result.provider_metadata + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/reranking/amazon-bedrock-reranking-model.test.ts › doRerank › text documents › should return result with the correct response + #[tokio::test] + async fn should_return_result_with_the_correct_response() { + let (result, _seen) = do_rerank(documents()).await; + + assert_eq!( + serde_json::to_value(&result.response).expect("response serializes"), + expected_response() + ); + } + } +} diff --git a/aimux-providers/src/bedrock/tests/sigv4_fetch.rs b/aimux-providers/src/bedrock/tests/sigv4_fetch.rs new file mode 100644 index 000000000..3a6feea41 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/sigv4_fetch.rs @@ -0,0 +1,1056 @@ +//! Ports of `@ai-sdk/amazon-bedrock` 5.0.100 +//! `src/amazon-bedrock-sigv4-fetch.test.ts`. +//! +//! aimux has no `createSigV4FetchFunction` / `createApiKeyFetchFunction`. The +//! provider's transport is `BedrockAuthFetch` (`bedrock/mod.rs`): it adds the +//! `ai-sdk-amazon-bedrock/` user-agent part, then signs with +//! `aimux_provider_utils::SigV4Fetch` (service `bedrock`, aws4fetch's header +//! rule) unless the request already carries `Authorization`. The API key is +//! put into the provider headers, not by the transport, so the cases about +//! adding, merging or replacing it run through `create_amazon_bedrock` and the +//! mock transport; the cases about what the transport does with a request that +//! carries the key start from such a request. +//! +//! Upstream mocks the `aws4fetch` signer: its `x-amz-date` is fixed and its +//! `authorization` is `AWS4-HMAC-SHA256 Credential=test`. The port signs for +//! real on a fixed clock (`SigV4Fetch::with_clock`, the upstream date), so +//! `x-amz-date` is compared exactly and `authorization` by its credential +//! prefix. Upstream also mocks the runtime part of the user-agent +//! (`testenv`), which aimux does not send. +//! +//! Not ported (registered, `AWS SigV4 fetch 也给 GET 请求签名（§9.1 已知差距）`): +//! "should bypass signing for non-POST requests", "should bypass signing if +//! POST request has no body" and "should call original fetch if init is +//! undefined" (a `GET`). aimux signs every request, GETs and bodyless POSTs +//! included (`list_models` is a signed GET). + +use std::collections::BTreeMap; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use bytes::Bytes; +use chrono::{DateTime, TimeZone, Utc}; +use futures::StreamExt; +use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; +use reqwest::{Method, StatusCode}; +use serde_json::{Value, json}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use url::Url; + +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::error::AiMuxError; +use aimux_provider_utils::{ + AwsCredentials, Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse, HeaderMapOpt, + ProxyConfig, RequestCredentials, ReqwestFetch, Resolvable, SigV4Fetch, SigV4Signing, + set_default_transport, +}; + +use crate::bedrock::{ + AmazonBedrockProvider, AmazonBedrockProviderSettings, BedrockAuthFetch, SIGV4_SERVICE, + create_amazon_bedrock, +}; +use crate::test_support::mock_fetch::{Canned, MockFetch, Seen}; + +/// The user-agent the transport sends: the pinned package part. +const USER_AGENT: &str = "ai-sdk-amazon-bedrock/5.0.100"; + +/// The signing date of upstream's mocked signer. +const AMZ_DATE: &str = "20240315T000000Z"; + +fn fixed_clock() -> DateTime { + Utc.with_ymd_and_hms(2024, 3, 15, 0, 0, 0) + .single() + .expect("valid date") +} + +/// The credentials of the upstream `createFetchFunction`. +fn credentials(session_token: Option<&str>) -> AwsCredentials { + AwsCredentials { + access_key_id: "test-access-key".into(), + secret_access_key: "test-secret".into(), + session_token: session_token.map(str::to_string), + region: "us-west-2".into(), + } +} + +/// `createSigV4FetchFunction(getCredentials, inner, service)`: the provider's +/// transport, signing for `service`. +fn sigv4_fetch_for( + credentials: Resolvable, + inner: FetchFunction, + service: &str, +) -> FetchFunction { + Arc::new(BedrockAuthFetch { + signing: SigV4Fetch::new(inner.clone(), credentials, service, SigV4Signing::Aws4Fetch) + .with_clock(fixed_clock), + inner, + }) +} + +/// `createSigV4FetchFunction(getCredentials, inner)`: the provider's +/// transport with its own service name. +fn sigv4_fetch(credentials: Resolvable, inner: FetchFunction) -> FetchFunction { + sigv4_fetch_for(credentials, inner, SIGV4_SERVICE) +} + +/// The provider's transport for a request that already carries the API key +/// as `Authorization` (which the provider headers put there): forwarded +/// unsigned. Its signer would fail, so a signed request would show. +fn api_key_fetch(inner: FetchFunction) -> FetchFunction { + sigv4_fetch( + Resolvable::from_fn(|| Err(AiMuxError::other("no credentials"))), + inner, + ) +} + +/// A request with the given headers and body. The JS `Request`, `Headers`, +/// header array and `ArrayBuffer` / `Uint8Array` inputs all become one of these. +fn request(method: Method, url: &str, headers: &[(&str, &str)], body: &[u8]) -> FetchRequest { + let mut request = FetchRequest::new(method, Url::parse(url).expect("valid URL")); + for (name, value) in headers { + request.headers.insert( + HeaderName::from_bytes(name.as_bytes()).expect("valid header name"), + HeaderValue::from_str(value).expect("valid header value"), + ); + } + request.body = Bytes::copy_from_slice(body); + request +} + +/// A canned `200 OK` with a text body. +fn ok(body: &str) -> Canned { + Canned { + status: 200, + headers: vec![], + body: body.as_bytes().to_vec(), + } +} + +async fn body_text(response: FetchResponse) -> String { + let chunks: Vec> = response.bytes_stream().collect().await; + let bytes: Vec = chunks + .into_iter() + .flat_map(|chunk| chunk.expect("body chunk").to_vec()) + .collect(); + String::from_utf8(bytes).expect("UTF-8 body") +} + +/// Asserts the SigV4 headers of a forwarded request: the fixed signing date +/// and a SigV4 credential of `access_key_id` for `region` and `service`. +fn assert_signed(seen: &Seen, access_key_id: &str, region: &str, service: &str) { + assert_eq!(seen.headers["x-amz-date"], AMZ_DATE); + let authorization = &seen.headers["authorization"]; + let credential = format!( + "AWS4-HMAC-SHA256 Credential={access_key_id}/20240315/{region}/{service}/aws4_request, " + ); + assert!( + authorization.starts_with(&credential), + "unexpected authorization: {authorization}" + ); +} + +/// The headers a provider request sends without the `user-agent`, which is +/// checked on its own to carry the package part. +fn caller_visible(headers: &BTreeMap) -> Value { + let mut headers = headers.clone(); + let user_agent = headers.remove("user-agent").expect("user-agent is sent"); + assert!( + user_agent.split(' ').any(|part| part == USER_AGENT), + "user-agent {user_agent} lacks {USER_AGENT}" + ); + json!(headers) +} + +/// The Titan embedding response, for the cases that go through the provider. +fn titan_response() -> Canned { + Canned::json(&json!({ "embedding": [0.1], "inputTextTokenCount": 1 })) +} + +const TITAN: &str = "amazon.titan-embed-text-v2:0"; + +/// `create_amazon_bedrock` with an API key, the given headers, transport and +/// base URL. +fn api_key_provider( + api_key: &str, + headers: Option, + fetch: Option, + base_url: &str, +) -> AmazonBedrockProvider { + create_amazon_bedrock(AmazonBedrockProviderSettings { + api_key: Some(api_key.to_string().into()), + region: Some("us-east-1".to_string()), + base_url: Some(base_url.to_string()), + headers, + fetch, + ..Default::default() + }) + .expect("valid Bedrock settings") +} + +/// One `POST /model/{titan}/invoke` with a JSON body through `provider`. +async fn embed(provider: &AmazonBedrockProvider) { + provider + .embedding(TITAN) + .do_embed(&EmbeddingCallOptions::new("x")) + .await + .expect("embedding succeeds"); +} + +fn header_map(headers: &HeaderMap) -> BTreeMap { + headers + .iter() + .map(|(name, value)| { + ( + name.as_str().to_string(), + value.to_str().expect("ASCII header").to_string(), + ) + }) + .collect() +} + +/// Records the whole request the transport forwards, including the fields +/// [`Seen`] does not carry. +#[derive(Default)] +struct Recorder { + requests: Mutex>, +} + +#[async_trait] +impl Fetch for Recorder { + async fn fetch(&self, request: FetchRequest) -> Result { + let url = request.url.clone(); + self.requests.lock().unwrap().push(request); + Ok(FetchResponse::from_bytes( + StatusCode::OK, + HeaderMap::new(), + url, + Bytes::from_static(b"OK"), + )) + } +} + +/// One request as the fake proxy received it. +#[derive(Clone, Debug)] +struct SeenByProxy { + request_line: String, + headers: BTreeMap, +} + +/// A plain-HTTP stand-in for a proxy. The default transport sends plain-HTTP +/// requests to it in absolute form. It records each request head and answers +/// `200` with a Titan embedding. +#[derive(Clone)] +struct FakeProxy { + port: u16, + seen: Arc>>, +} + +impl FakeProxy { + async fn start() -> Self { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind fake proxy"); + let port = listener.local_addr().expect("local address").port(); + let seen = Arc::new(Mutex::new(Vec::new())); + let recorded = Arc::clone(&seen); + tokio::spawn(async move { + loop { + let Ok((stream, _)) = listener.accept().await else { + continue; + }; + let recorded = Arc::clone(&recorded); + tokio::spawn(async move { serve_one(stream, recorded).await }); + } + }); + Self { port, seen } + } + + /// Proxies plain-HTTP requests, except to loopback hosts (other tests' + /// local servers). + fn config(&self) -> ProxyConfig { + ProxyConfig { + http_url: Some(format!("http://127.0.0.1:{}", self.port)), + no_proxy: Some("localhost,127.0.0.1".to_string()), + ..ProxyConfig::default() + } + } + + fn seen(&self) -> Vec { + self.seen.lock().unwrap().clone() + } +} + +async fn serve_one(mut stream: TcpStream, recorded: Arc>>) { + let mut head = Vec::new(); + let mut chunk = [0u8; 1024]; + while !head.windows(4).any(|window| window == b"\r\n\r\n") { + match stream.read(&mut chunk).await { + Ok(0) | Err(_) => return, + Ok(n) => head.extend_from_slice(&chunk[..n]), + } + } + let text = String::from_utf8_lossy(&head).into_owned(); + let mut lines = text.split("\r\n"); + let request_line = lines.next().unwrap_or_default().to_string(); + let headers = lines + .take_while(|line| !line.is_empty()) + .filter_map(|line| line.split_once(':')) + .map(|(name, value)| (name.trim().to_ascii_lowercase(), value.trim().to_string())) + .collect(); + recorded.lock().unwrap().push(SeenByProxy { + request_line, + headers, + }); + let body = r#"{"embedding":[0.1],"inputTextTokenCount":1}"#; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + let _ = stream.write_all(response.as_bytes()).await; + let _ = stream.shutdown().await; +} + +/// Routes the process default transport through `proxy`. Tests that call it +/// are `#[serial]` and hold a [`RestoreDefaultTransport`]. +fn route_default_transport_to(proxy: &FakeProxy) { + set_default_transport(ReqwestFetch::new(proxy.config())); +} + +/// Restores the built-in process default transport when dropped. +struct RestoreDefaultTransport; + +impl Drop for RestoreDefaultTransport { + fn drop(&mut self) { + set_default_transport(ReqwestFetch::default()); + } +} + +/// The request line of the provider's Titan request through the proxy. +const PROXIED_REQUEST_LINE: &str = + "POST http://example.com/model/amazon.titan-embed-text-v2%3A0/invoke "; + +mod create_sig_v4_fetch_function { + use super::*; + + // "should bypass signing for non-POST requests" and "should bypass signing + // if POST request has no body": registered (§9.1), see the file comment. + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createSigV4FetchFunction › should handle a POST request with a string body and merge signed headers including user-agent + #[tokio::test] + async fn should_handle_a_post_request_with_a_string_body_and_merge_signed_headers() { + let mock = MockFetch::new(vec![ok("Signed")]); + // The session token makes the signer add `x-amz-security-token`. + let fetch = sigv4_fetch( + Resolvable::from_fn(|| Ok(credentials(Some("test-session-token")))), + mock.transport(), + ); + + fetch + .fetch(request( + Method::POST, + "http://example.com", + &[ + ("Content-Type", "application/json"), + ("Custom-Header", "value"), + ], + br#"{"test": "data"}"#, + )) + .await + .expect("fetch succeeds"); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + let headers = &seen[0].headers; + assert_eq!(headers["content-type"], "application/json"); + assert_eq!(headers["custom-header"], "value"); + assert!(!headers.contains_key("empty-header")); + assert_signed(&seen[0], "test-access-key", "us-west-2", "bedrock"); + assert_eq!(headers["x-amz-security-token"], "test-session-token"); + assert_eq!(headers["user-agent"], USER_AGENT); + // Body is left unmodified for a string body. + assert_eq!(seen[0].body, br#"{"test": "data"}"#); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createSigV4FetchFunction › should handle a POST request with a Request object + /// + /// The JS `Request` object and the `init` headers are two layers the + /// transport merges. A Rust `FetchRequest` is already one request, so both + /// layers are given as its headers. + #[tokio::test] + async fn should_handle_a_post_request_with_a_request_object() { + let mock = MockFetch::new(vec![ok("Signed")]); + let fetch = sigv4_fetch( + Resolvable::from_fn(|| Ok(credentials(None))), + mock.transport(), + ); + + fetch + .fetch(request( + Method::POST, + "http://example.com", + &[ + ("X-From-Request", "from-request"), + ("Content-Type", "application/json"), + ("Custom-Header", "value"), + ], + br#"{"test": "data"}"#, + )) + .await + .expect("fetch succeeds"); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].url, "http://example.com/"); + assert_eq!(seen[0].body, br#"{"test": "data"}"#); + let headers = &seen[0].headers; + assert_eq!(headers["content-type"], "application/json"); + assert_eq!(headers["custom-header"], "value"); + assert_eq!(headers["x-from-request"], "from-request"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createSigV4FetchFunction › should sign when input is a POST Request with body and no init + #[tokio::test] + async fn should_sign_when_input_is_a_post_request_with_body_and_no_init() { + let mock = MockFetch::new(vec![ok("Signed")]); + let fetch = sigv4_fetch( + Resolvable::from_fn(|| Ok(credentials(None))), + mock.transport(), + ); + + fetch + .fetch(request( + Method::POST, + "http://example.com", + &[ + ("Content-Type", "application/json"), + ("X-From-Request", "from-request"), + ], + br#"{"test": "data"}"#, + )) + .await + .expect("fetch succeeds"); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].body, br#"{"test": "data"}"#); + let headers = &seen[0].headers; + assert_eq!(headers["content-type"], "application/json"); + assert_eq!(headers["x-from-request"], "from-request"); + assert_signed(&seen[0], "test-access-key", "us-west-2", "bedrock"); + assert_eq!(headers["user-agent"], USER_AGENT); + } + + // Not ported (ts-only): "should handle non-string body by stringifying + // it". Its only subject is a JS object passed as `body` and coerced with + // `JSON.stringify`. `FetchRequest.body` is bytes, so there is no + // non-string body to coerce. + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createSigV4FetchFunction › should handle Uint8Array body + /// + /// A `Uint8Array` body and an `ArrayBuffer` body are both bytes in Rust, so + /// this case and the next send the same kind of request; the body is + /// forwarded as the same bytes (upstream: decoded to the same string). + #[tokio::test] + async fn should_handle_uint8array_body() { + let mock = MockFetch::new(vec![ok("Signed")]); + let fetch = sigv4_fetch( + Resolvable::from_fn(|| Ok(credentials(None))), + mock.transport(), + ); + + fetch + .fetch(request( + Method::POST, + "http://example.com", + &[], + b"binaryTest", + )) + .await + .expect("fetch succeeds"); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].body, b"binaryTest"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createSigV4FetchFunction › should handle ArrayBuffer body + #[tokio::test] + async fn should_handle_arraybuffer_body() { + let mock = MockFetch::new(vec![ok("Signed")]); + let fetch = sigv4_fetch( + Resolvable::from_fn(|| Ok(credentials(None))), + mock.transport(), + ); + + fetch + .fetch(request( + Method::POST, + "http://example.com", + &[], + b"bufferTest", + )) + .await + .expect("fetch succeeds"); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].body, b"bufferTest"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createSigV4FetchFunction › should extract headers from a Headers instance + #[tokio::test] + async fn should_extract_headers_from_a_headers_instance() { + let mock = MockFetch::new(vec![ok("Signed")]); + let fetch = sigv4_fetch( + Resolvable::from_fn(|| Ok(credentials(None))), + mock.transport(), + ); + + fetch + .fetch(request( + Method::POST, + "http://example.com", + &[("A", "value-a"), ("B", "value-b")], + br#"{"test": "data"}"#, + )) + .await + .expect("fetch succeeds"); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].headers["a"], "value-a"); + assert_eq!(seen[0].headers["b"], "value-b"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createSigV4FetchFunction › should handle headers provided as an array + #[tokio::test] + async fn should_handle_headers_provided_as_an_array() { + let mock = MockFetch::new(vec![ok("Signed")]); + let fetch = sigv4_fetch( + Resolvable::from_fn(|| Ok(credentials(None))), + mock.transport(), + ); + + fetch + .fetch(request( + Method::POST, + "http://example.com", + &[ + ("Array-Header", "array-value"), + ("Another-Header", "another-value"), + ], + br#"{"test": "data"}"#, + )) + .await + .expect("fetch succeeds"); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].headers["array-header"], "array-value"); + assert_eq!(seen[0].headers["another-header"], "another-value"); + // Also check that the signing headers are included. + assert_signed(&seen[0], "test-access-key", "us-west-2", "bedrock"); + } + + // "should call original fetch if init is undefined": registered (§9.1), + // see the file comment. + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createSigV4FetchFunction › should correctly handle async credential providers + #[tokio::test] + async fn should_correctly_handle_async_credential_providers() { + let mock = MockFetch::new(vec![ok("Signed")]); + let fetch = sigv4_fetch( + Resolvable::from_async_fn(|| async { + Ok(AwsCredentials { + access_key_id: "async-access-key".into(), + secret_access_key: "async-secret-key".into(), + session_token: Some("async-session-token".into()), + region: "us-east-1".into(), + }) + }), + mock.transport(), + ); + + fetch + .fetch(request( + Method::POST, + "http://example.com", + &[("Content-Type", "application/json")], + br#"{"test": "async"}"#, + )) + .await + .expect("fetch succeeds"); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_signed(&seen[0], "async-access-key", "us-east-1", "bedrock"); + assert_eq!( + seen[0].headers["x-amz-security-token"], + "async-session-token" + ); + assert_eq!(seen[0].headers["content-type"], "application/json"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createSigV4FetchFunction › should handle async credential providers that reject + #[tokio::test] + async fn should_handle_async_credential_providers_that_reject() { + let mock = MockFetch::new(vec![]); + let fetch = sigv4_fetch( + Resolvable::from_async_fn(|| async { + Err(AiMuxError::other("Failed to get credentials")) + }), + mock.transport(), + ); + + let error = fetch + .fetch(request( + Method::POST, + "http://example.com", + &[], + br#"{"test": "data"}"#, + )) + .await + .expect_err("the fetch rejects"); + assert!( + error.to_string().contains("Failed to get credentials"), + "unexpected error: {error}" + ); + // The underlying fetch should not be called. + assert!(mock.seen().is_empty()); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createSigV4FetchFunction › should use default service name "bedrock" when no service parameter is provided + /// + /// The signer options are not observable; the service is the one the + /// signature's credential scope names. The provider's service is its + /// default. + #[tokio::test] + async fn should_use_default_service_name_bedrock_when_no_service_parameter_is_provided() { + let mock = MockFetch::new(vec![ok("Signed")]); + let fetch = sigv4_fetch( + Resolvable::from_fn(|| Ok(credentials(None))), + mock.transport(), + ); + + fetch + .fetch(request( + Method::POST, + "http://example.com", + &[], + br#"{"test": "data"}"#, + )) + .await + .expect("fetch succeeds"); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_signed(&seen[0], "test-access-key", "us-west-2", "bedrock"); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createSigV4FetchFunction › should use custom service name when service parameter is provided + #[tokio::test] + async fn should_use_custom_service_name_when_service_parameter_is_provided() { + let mock = MockFetch::new(vec![ok("Signed")]); + let fetch = sigv4_fetch_for( + Resolvable::from_fn(|| Ok(credentials(None))), + mock.transport(), + "bedrock-mantle", + ); + + fetch + .fetch(request( + Method::POST, + "http://example.com", + &[], + br#"{"test": "data"}"#, + )) + .await + .expect("fetch succeeds"); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_signed(&seen[0], "test-access-key", "us-west-2", "bedrock-mantle"); + } +} + +mod create_api_key_fetch_function { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createApiKeyFetchFunction › should add Authorization header with Bearer token and user-agent + /// + /// The key is added by the provider headers, so this runs through the + /// provider: a `POST` with a JSON body, whose `content-type` the provider + /// sets. + #[tokio::test] + async fn should_add_authorization_header_with_bearer_token_and_user_agent() { + let mock = MockFetch::new(vec![titan_response()]); + let provider = api_key_provider( + "test-api-key-123", + None, + Some(mock.transport()), + "http://example.com", + ); + + embed(&provider).await; + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].method, "POST"); + assert_eq!(seen[0].json_body(), json!({ "inputText": "x" })); + assert_eq!( + caller_visible(&seen[0].headers), + json!({ + "content-type": "application/json", + "authorization": "Bearer test-api-key-123", + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createApiKeyFetchFunction › should merge Authorization header with existing headers + #[tokio::test] + async fn should_merge_authorization_header_with_existing_headers() { + let mock = MockFetch::new(vec![titan_response()]); + let provider = api_key_provider( + "test-api-key-456", + Some(HeaderMapOpt::from([ + ( + "Content-Type".to_string(), + Some("application/json".to_string()), + ), + ( + "Custom-Header".to_string(), + Some("custom-value".to_string()), + ), + ("X-Request-ID".to_string(), Some("req-123".to_string())), + ])), + Some(mock.transport()), + "http://example.com", + ); + + embed(&provider).await; + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_eq!( + caller_visible(&seen[0].headers), + json!({ + "content-type": "application/json", + "custom-header": "custom-value", + "x-request-id": "req-123", + "authorization": "Bearer test-api-key-456", + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createApiKeyFetchFunction › should work with Headers instance + /// + /// The request already carries the key, as the provider puts it there. + #[tokio::test] + async fn should_work_with_headers_instance() { + let mock = MockFetch::new(vec![ok("OK")]); + let fetch = api_key_fetch(mock.transport()); + + fetch + .fetch(request( + Method::POST, + "http://example.com", + &[ + ("Content-Type", "application/json"), + ("X-Custom", "value"), + ("Authorization", "Bearer test-api-key-789"), + ], + br#"{"test": "data"}"#, + )) + .await + .expect("fetch succeeds"); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].method, "POST"); + assert_eq!(seen[0].url, "http://example.com/"); + assert_eq!(seen[0].body, br#"{"test": "data"}"#); + assert_eq!( + json!(seen[0].headers), + json!({ + "content-type": "application/json", + "x-custom": "value", + "authorization": "Bearer test-api-key-789", + "user-agent": USER_AGENT, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createApiKeyFetchFunction › should work with headers as array + #[tokio::test] + async fn should_work_with_headers_as_array() { + let mock = MockFetch::new(vec![ok("OK")]); + let fetch = api_key_fetch(mock.transport()); + + fetch + .fetch(request( + Method::POST, + "http://example.com", + &[ + ("Content-Type", "application/json"), + ("X-Array-Header", "array-value"), + ("Authorization", "Bearer test-api-key-array"), + ], + br#"{"test": "data"}"#, + )) + .await + .expect("fetch succeeds"); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].method, "POST"); + assert_eq!(seen[0].body, br#"{"test": "data"}"#); + assert_eq!( + json!(seen[0].headers), + json!({ + "content-type": "application/json", + "x-array-header": "array-value", + "authorization": "Bearer test-api-key-array", + "user-agent": USER_AGENT, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createApiKeyFetchFunction › should work with GET requests + #[tokio::test] + async fn should_work_with_get_requests() { + let mock = MockFetch::new(vec![ok("OK")]); + let fetch = api_key_fetch(mock.transport()); + + fetch + .fetch(request( + Method::GET, + "http://example.com", + &[ + ("Accept", "application/json"), + ("Authorization", "Bearer test-api-key-get"), + ], + b"", + )) + .await + .expect("fetch succeeds"); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].method, "GET"); + assert_eq!( + json!(seen[0].headers), + json!({ + "accept": "application/json", + "authorization": "Bearer test-api-key-get", + "user-agent": USER_AGENT, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createApiKeyFetchFunction › should work when no headers are provided + #[tokio::test] + async fn should_work_when_no_headers_are_provided() { + let mock = MockFetch::new(vec![ok("OK")]); + let fetch = api_key_fetch(mock.transport()); + + fetch + .fetch(request( + Method::POST, + "http://example.com", + &[("Authorization", "Bearer test-api-key-no-headers")], + br#"{"test": "data"}"#, + )) + .await + .expect("fetch succeeds"); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].method, "POST"); + assert_eq!(seen[0].body, br#"{"test": "data"}"#); + assert_eq!( + json!(seen[0].headers), + json!({ + "authorization": "Bearer test-api-key-no-headers", + "user-agent": USER_AGENT, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createApiKeyFetchFunction › should work when init is undefined + /// + /// No `init` is the `fetch` default, a `GET` without headers (but the key). + #[tokio::test] + async fn should_work_when_init_is_undefined() { + let mock = MockFetch::new(vec![ok("OK")]); + let fetch = api_key_fetch(mock.transport()); + + let response = fetch + .fetch(request( + Method::GET, + "http://example.com", + &[("Authorization", "Bearer test-api-key-undefined")], + b"", + )) + .await + .expect("fetch succeeds"); + assert_eq!(body_text(response).await, "OK"); + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].url, "http://example.com/"); + assert_eq!( + json!(seen[0].headers), + json!({ + "authorization": "Bearer test-api-key-undefined", + "user-agent": USER_AGENT, + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createApiKeyFetchFunction › should override existing Authorization header + /// + /// The caller's `Authorization` is a provider header, which the key + /// replaces. Upstream's expected header object also keeps a second, + /// lower-case `authorization: Bearer old-token` key (a JS object is + /// case-sensitive); a Rust header map is not, so that key is not asserted. + #[tokio::test] + async fn should_override_existing_authorization_header() { + let mock = MockFetch::new(vec![titan_response()]); + let provider = api_key_provider( + "test-api-key-override", + Some(HeaderMapOpt::from([ + ( + "Content-Type".to_string(), + Some("application/json".to_string()), + ), + ( + "Authorization".to_string(), + Some("Bearer old-token".to_string()), + ), + ])), + Some(mock.transport()), + "http://example.com", + ); + + embed(&provider).await; + + let seen = mock.seen(); + assert_eq!(seen.len(), 1); + assert_eq!( + caller_visible(&seen[0].headers), + json!({ + "content-type": "application/json", + "authorization": "Bearer test-api-key-override", + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createApiKeyFetchFunction › should use default fetch when no custom fetch provided + /// + /// The process default transport is routed to a fake proxy, so the request + /// reaches it as the default transport would send it. + #[tokio::test] + #[serial_test::serial] + async fn should_use_default_fetch_when_no_custom_fetch_provided() { + let proxy = FakeProxy::start().await; + let _restore = RestoreDefaultTransport; + route_default_transport_to(&proxy); + let provider = api_key_provider("test-api-key-default", None, None, "http://example.com"); + + embed(&provider).await; + + let seen = proxy.seen(); + assert_eq!(seen.len(), 1); + assert!( + seen[0].request_line.starts_with(PROXIED_REQUEST_LINE), + "unexpected request line: {}", + seen[0].request_line + ); + assert_eq!( + seen[0].headers["authorization"], + "Bearer test-api-key-default" + ); + assert!( + seen[0].headers["user-agent"] + .split(' ') + .any(|part| part == USER_AGENT) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createApiKeyFetchFunction › should resolve default fetch lazily when no custom fetch provided + /// + /// The provider is created while `initial` is the process default, and the + /// default is replaced afterwards. The request must use the default that + /// is current when it is made. + #[tokio::test] + #[serial_test::serial] + async fn should_resolve_default_fetch_lazily_when_no_custom_fetch_provided() { + let initial = FakeProxy::start().await; + let patched = FakeProxy::start().await; + let _restore = RestoreDefaultTransport; + route_default_transport_to(&initial); + let provider = api_key_provider("test-api-key-lazy", None, None, "http://example.com"); + route_default_transport_to(&patched); + + embed(&provider).await; + + assert!( + initial.seen().is_empty(), + "the default transport captured at creation was used: {:?}", + initial.seen() + ); + let seen = patched.seen(); + assert_eq!(seen.len(), 1, "the replaced default transport was not used"); + assert_eq!(seen[0].headers["authorization"], "Bearer test-api-key-lazy"); + assert!( + seen[0].request_line.starts_with(PROXIED_REQUEST_LINE), + "unexpected request line: {}", + seen[0].request_line + ); + } + + // Not ported (sub-feature): "should handle empty string API key". aimux + // has no `createApiKeyFetchFunction` that could be handed an empty key: + // the key is a provider header, set only for a non-blank key (as + // upstream's `createAmazonBedrock` only calls the function for one). + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-sigv4-fetch.test.ts › createApiKeyFetchFunction › should preserve request body and other properties + /// + /// `cache: 'no-cache'` has no Rust counterpart on `FetchRequest` and is not + /// asserted. `credentials: 'include'` is `RequestCredentials::Include`. + #[tokio::test] + async fn should_preserve_request_body_and_other_properties() { + let recorder = Arc::new(Recorder::default()); + let fetch = api_key_fetch(recorder.clone()); + let request_body = json!({ "data": "test" }).to_string(); + let mut sent = request( + Method::PUT, + "http://example.com", + &[ + ("Content-Type", "application/json"), + ("Authorization", "Bearer test-api-key-preserve"), + ], + request_body.as_bytes(), + ); + sent.credentials = RequestCredentials::Include; + + fetch.fetch(sent).await.expect("fetch succeeds"); + + let requests = recorder.requests.lock().unwrap(); + assert_eq!(requests.len(), 1); + let forwarded = &requests[0]; + assert_eq!(forwarded.method, Method::PUT); + assert_eq!(forwarded.url.as_str(), "http://example.com/"); + assert_eq!(forwarded.body, Bytes::from(request_body)); + assert_eq!(forwarded.credentials, RequestCredentials::Include); + assert_eq!( + json!(header_map(&forwarded.headers)), + json!({ + "content-type": "application/json", + "authorization": "Bearer test-api-key-preserve", + "user-agent": USER_AGENT, + }) + ); + } +} diff --git a/aimux-providers/src/bedrock/tests/snapshots/do_generate_reasoning_should_extract_reasoning_and_text_response.json b/aimux-providers/src/bedrock/tests/snapshots/do_generate_reasoning_should_extract_reasoning_and_text_response.json new file mode 100644 index 000000000..77e7c9322 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/snapshots/do_generate_reasoning_should_extract_reasoning_and_text_response.json @@ -0,0 +1,90 @@ +{ + "content": [ + { + "providerMetadata": { + "amazonBedrock": { + "signature": "EvYBCkgICxABGAIqQA9E9VC377UnbjdfXCw4RwQaaIXsqocZKzI3WwWtXT/wBjzAkBOfIiVkP/mhvsPGb5oQW/2j+E1tZcNFoXIb26YSDBTSbDrbHwXJoAl3/hoMA/Qo4rIHvObP/yBrIjAT+l69mh9k7evFO+w+ransSD5YGq4hYrieOyJH3k/zHv8nBJwiz3Nlrb3jNFi2Ib0qXJH7BIK0hR85yYIoccReoJUpeLEMxeWbsO6f0EpOSMSi3YDLd1U9NAw7Itjeay1fEBPHwvawt2M4e/rn52SsuEz3p2DfOZm+N3bnL63rg0Cb9lFcp3k5kuU1h6hcGAE=" + }, + "bedrock": { + "signature": "EvYBCkgICxABGAIqQA9E9VC377UnbjdfXCw4RwQaaIXsqocZKzI3WwWtXT/wBjzAkBOfIiVkP/mhvsPGb5oQW/2j+E1tZcNFoXIb26YSDBTSbDrbHwXJoAl3/hoMA/Qo4rIHvObP/yBrIjAT+l69mh9k7evFO+w+ransSD5YGq4hYrieOyJH3k/zHv8nBJwiz3Nlrb3jNFi2Ib0qXJH7BIK0hR85yYIoccReoJUpeLEMxeWbsO6f0EpOSMSi3YDLd1U9NAw7Itjeay1fEBPHwvawt2M4e/rn52SsuEz3p2DfOZm+N3bnL63rg0Cb9lFcp3k5kuU1h6hcGAE=" + } + }, + "text": "Let me count the r's in \"strawberry\":\n\ns-t-r-a-w-b-e-r-r-y\n\nThere are 3 r's.", + "type": "reasoning" + }, + { + "text": "There are **3** r's in \"strawberry\":\n\n1. st**r**awbe**r****r**y", + "type": "text" + } + ], + "finishReason": { + "raw": "end_turn", + "unified": "stop" + }, + "providerMetadata": { + "amazonBedrock": { + "stopSequence": null, + "usage": { + "cacheWriteInputTokens": 0 + } + }, + "bedrock": { + "stopSequence": null, + "usage": { + "cacheWriteInputTokens": 0 + } + } + }, + "request": { + "body": { + "additionalModelResponseFieldPaths": [ + "/delta/stop_sequence" + ], + "messages": [ + { + "content": [ + { + "text": "Hello" + } + ], + "role": "user" + } + ], + "system": [ + { + "text": "System Prompt" + } + ] + } + }, + "response": { + "headers": { + "content-length": "865", + "content-type": "application/json" + }, + "modelId": "anthropic.claude-3-haiku-20240307-v1:0" + }, + "usage": { + "inputTokens": { + "cacheRead": 0, + "cacheWrite": 0, + "noCache": 51, + "total": 51 + }, + "outputTokens": { + "text": 78, + "total": 78 + }, + "raw": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 51, + "outputTokens": 78, + "serverToolUsage": {}, + "totalTokens": 129 + } + }, + "warnings": [] +} diff --git a/aimux-providers/src/bedrock/tests/snapshots/do_generate_text_should_extract_text_response.json b/aimux-providers/src/bedrock/tests/snapshots/do_generate_text_should_extract_text_response.json new file mode 100644 index 000000000..9d65ebe92 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/snapshots/do_generate_text_should_extract_text_response.json @@ -0,0 +1,78 @@ +{ + "content": [ + { + "text": "Let me count the \"r\"s in \"strawberry\":\n\ns-t-**r**-a-w-b-e-**r**-**r**-y\n\nThere are **3** \"r\"s in \"strawberry.\"", + "type": "text" + } + ], + "finishReason": { + "raw": "end_turn", + "unified": "stop" + }, + "providerMetadata": { + "amazonBedrock": { + "stopSequence": null, + "usage": { + "cacheWriteInputTokens": 0 + } + }, + "bedrock": { + "stopSequence": null, + "usage": { + "cacheWriteInputTokens": 0 + } + } + }, + "request": { + "body": { + "additionalModelResponseFieldPaths": [ + "/delta/stop_sequence" + ], + "messages": [ + { + "content": [ + { + "text": "Hello" + } + ], + "role": "user" + } + ], + "system": [ + { + "text": "System Prompt" + } + ] + } + }, + "response": { + "headers": { + "content-length": "435", + "content-type": "application/json" + }, + "modelId": "anthropic.claude-3-haiku-20240307-v1:0" + }, + "usage": { + "inputTokens": { + "cacheRead": 0, + "cacheWrite": 0, + "noCache": 22, + "total": 22 + }, + "outputTokens": { + "text": 57, + "total": 57 + }, + "raw": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 22, + "outputTokens": 57, + "serverToolUsage": {}, + "totalTokens": 79 + } + }, + "warnings": [] +} diff --git a/aimux-providers/src/bedrock/tests/snapshots/do_stream_reasoning_should_stream_reasoning_and_text_parts.json b/aimux-providers/src/bedrock/tests/snapshots/do_stream_reasoning_should_stream_reasoning_and_text_parts.json new file mode 100644 index 000000000..1598b68c6 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/snapshots/do_stream_reasoning_should_stream_reasoning_and_text_parts.json @@ -0,0 +1,159 @@ +[ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "id": "0", + "type": "reasoning-start" + }, + { + "delta": "Let me count the r", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": "'s in \"", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": "strawberry\":", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": "\n\ns-t-r-a", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": "-w-b-e-r", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": "-r-y\n\nr", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " appears at positions ", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": "3, 8", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": ", and 9.\n\nSo there", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " are 3 r's.", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": "", + "id": "0", + "providerMetadata": { + "amazonBedrock": { + "signature": "Ep0CCkgICxABGAIqQOWPB6/PmA5SW9jC6FvaNq3E+U9ev4FMWcFWuAho+VGLCtazKc5WDjQ5i0MuxsY0o5pKDSVWVKii8KJDusXH4eASDK7jyzuk8iij7fJNihoMxHO9haYzt48R36HVIjCb/EmIFrJLIXShqN6DN//T6vZBtO9qj1QhNWJa3CGm8VZoq80S2/Ok4U0aVaIDiZcqggHC2b8BHuv8BHZrmsR0wjU1ynansBGMdfjnG+iIv8R5lPpRmYGhSVwNybwP3aQZ6o8Dr48Rau8TJfdsArW+r7bvL7bPs4f5nnlp2vG7WkMzWwABHK3fdM44zZ1GZQaWyECNWR2GfY6dXiklo94vgpFTPuZ97mfiN3LY6uYyBwL8RkDaGAE=" + }, + "bedrock": { + "signature": "Ep0CCkgICxABGAIqQOWPB6/PmA5SW9jC6FvaNq3E+U9ev4FMWcFWuAho+VGLCtazKc5WDjQ5i0MuxsY0o5pKDSVWVKii8KJDusXH4eASDK7jyzuk8iij7fJNihoMxHO9haYzt48R36HVIjCb/EmIFrJLIXShqN6DN//T6vZBtO9qj1QhNWJa3CGm8VZoq80S2/Ok4U0aVaIDiZcqggHC2b8BHuv8BHZrmsR0wjU1ynansBGMdfjnG+iIv8R5lPpRmYGhSVwNybwP3aQZ6o8Dr48Rau8TJfdsArW+r7bvL7bPs4f5nnlp2vG7WkMzWwABHK3fdM44zZ1GZQaWyECNWR2GfY6dXiklo94vgpFTPuZ97mfiN3LY6uYyBwL8RkDaGAE=" + } + }, + "type": "reasoning-delta" + }, + { + "id": "0", + "type": "reasoning-end" + }, + { + "id": "1", + "type": "text-start" + }, + { + "delta": "There", + "id": "1", + "type": "text-delta" + }, + { + "delta": " are **3** r's in \"", + "id": "1", + "type": "text-delta" + }, + { + "delta": "strawberry\":", + "id": "1", + "type": "text-delta" + }, + { + "delta": "\n\n1", + "id": "1", + "type": "text-delta" + }, + { + "delta": ". st", + "id": "1", + "type": "text-delta" + }, + { + "delta": "**", + "id": "1", + "type": "text-delta" + }, + { + "delta": "r**awbe", + "id": "1", + "type": "text-delta" + }, + { + "delta": "**r****", + "id": "1", + "type": "text-delta" + }, + { + "delta": "r**y", + "id": "1", + "type": "text-delta" + }, + { + "id": "1", + "type": "text-end" + }, + { + "finishReason": { + "raw": "end_turn", + "unified": "stop" + }, + "type": "finish", + "usage": { + "inputTokens": { + "cacheRead": 0, + "cacheWrite": 0, + "noCache": 51, + "total": 51 + }, + "outputTokens": { + "text": 94, + "total": 94 + }, + "raw": { + "inputTokens": 51, + "outputTokens": 94, + "serverToolUsage": {}, + "totalTokens": 145 + } + } + } +] diff --git a/aimux-providers/src/bedrock/tests/snapshots/do_stream_text_should_stream_text_deltas.json b/aimux-providers/src/bedrock/tests/snapshots/do_stream_text_should_stream_text_deltas.json new file mode 100644 index 000000000..176b0b3a3 --- /dev/null +++ b/aimux-providers/src/bedrock/tests/snapshots/do_stream_text_should_stream_text_deltas.json @@ -0,0 +1,103 @@ +[ + { + "type": "stream-start", + "warnings": [] + }, + { + "modelId": "anthropic.claude-3-haiku-20240307-v1:0", + "type": "response-metadata" + }, + { + "id": "0", + "type": "text-start" + }, + { + "delta": "Let", + "id": "0", + "type": "text-delta" + }, + { + "delta": " me count the \"", + "id": "0", + "type": "text-delta" + }, + { + "delta": "r\"s in \"", + "id": "0", + "type": "text-delta" + }, + { + "delta": "strawberry\":\n\ns-t-", + "id": "0", + "type": "text-delta" + }, + { + "delta": "**", + "id": "0", + "type": "text-delta" + }, + { + "delta": "r**-a-w-b", + "id": "0", + "type": "text-delta" + }, + { + "delta": "-e-**", + "id": "0", + "type": "text-delta" + }, + { + "delta": "r**-**", + "id": "0", + "type": "text-delta" + }, + { + "delta": "r**-y\n\nThere", + "id": "0", + "type": "text-delta" + }, + { + "delta": " are **3", + "id": "0", + "type": "text-delta" + }, + { + "delta": "** r", + "id": "0", + "type": "text-delta" + }, + { + "delta": "'s in \"strawberry.\"", + "id": "0", + "type": "text-delta" + }, + { + "id": "0", + "type": "text-end" + }, + { + "finishReason": { + "raw": "end_turn", + "unified": "stop" + }, + "type": "finish", + "usage": { + "inputTokens": { + "cacheRead": 0, + "cacheWrite": 0, + "noCache": 22, + "total": 22 + }, + "outputTokens": { + "text": 55, + "total": 55 + }, + "raw": { + "inputTokens": 22, + "outputTokens": 55, + "serverToolUsage": {}, + "totalTokens": 77 + } + } + } +] diff --git a/aimux-providers/src/bedrock/tests/support.rs b/aimux-providers/src/bedrock/tests/support.rs new file mode 100644 index 000000000..bd5fa026a --- /dev/null +++ b/aimux-providers/src/bedrock/tests/support.rs @@ -0,0 +1,59 @@ +//! Helpers shared by the ported tests. +//! +//! [`inject_fetch_headers`] is the Rust form of upstream's +//! `src/inject-fetch-headers.ts`, a test helper of `@ai-sdk/amazon-bedrock` +//! (not exported): a fetch that merges fixed headers over the request's, +//! appends the `ai-sdk-amazon-bedrock/` user-agent part and calls the +//! wrapped transport, which stands in for `globalThis.fetch`. The embedding, +//! image and reranking ports use it the way their upstream tests do. + +use std::sync::Arc; + +use async_trait::async_trait; +use reqwest::header::{HeaderName, HeaderValue}; + +use aimux_provider_utils::{Fetch, FetchError, FetchFunction, FetchRequest, FetchResponse}; + +/// The `ai-sdk-amazon-bedrock/` user-agent part, pinned. +pub(super) const BEDROCK_USER_AGENT_PART: &str = "ai-sdk-amazon-bedrock/5.0.100"; + +/// `injectFetchHeaders(customHeaders)` over `inner`. +struct InjectFetchHeaders { + custom_headers: Vec<(HeaderName, HeaderValue)>, + inner: FetchFunction, +} + +#[async_trait] +impl Fetch for InjectFetchHeaders { + async fn fetch(&self, mut request: FetchRequest) -> Result { + for (name, value) in &self.custom_headers { + request.headers.insert(name.clone(), value.clone()); + } + crate::shared::transport_user_agent( + &mut request.headers, + None, + Some(BEDROCK_USER_AGENT_PART), + ); + self.inner.fetch(request).await + } +} + +/// A transport that adds `custom_headers` (and the package user-agent part) +/// to every request and forwards it to `inner`. +pub(super) fn inject_fetch_headers( + custom_headers: &[(&str, &str)], + inner: FetchFunction, +) -> FetchFunction { + Arc::new(InjectFetchHeaders { + custom_headers: custom_headers + .iter() + .map(|(name, value)| { + ( + HeaderName::from_bytes(name.as_bytes()).expect("valid header name"), + HeaderValue::from_str(value).expect("valid header value"), + ) + }) + .collect(), + inner, + }) +} diff --git a/aimux-providers/src/bedrock/tests/usage_raw.rs b/aimux-providers/src/bedrock/tests/usage_raw.rs new file mode 100644 index 000000000..da9d454af --- /dev/null +++ b/aimux-providers/src/bedrock/tests/usage_raw.rs @@ -0,0 +1,187 @@ +use std::sync::Arc; + +use aimux_core::AiMuxError; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::LanguageModelPrompt; +use aimux_core::options::CallOptions; +use aimux_core::stream_part::StreamPart; +use aimux_core::types::Usage; +use aimux_provider_utils::Resolvable; +use futures::StreamExt; +use serde_json::{Value, json}; + +use crate::bedrock::event_stream::encode_message; +use crate::bedrock::{AmazonBedrockProviderSettings, BedrockModel, create_amazon_bedrock}; +use crate::test_support::mock_fetch::{Canned, RouteFetch}; + +fn prompt() -> LanguageModelPrompt { + serde_json::from_value(json!([ + { "role": "user", "content": [{ "type": "text", "text": "Hello" }] }, + ])) + .unwrap() +} + +fn complete_usage() -> Value { + json!({ + "inputTokens": 13, + "outputTokens": 5, + "totalTokens": 18, + "cacheReadInputTokens": 3, + "cacheWriteInputTokens": 2, + "cacheDetails": [ + { + "inputTokens": 2, + "ttl": "T5M", + "providerCacheMetadata": { "source": "provider" }, + }, + ], + "providerUsageMetadata": { "serverToolUsage": {} }, + }) +} + +fn create_event_frame(event_type: &str, body: Value) -> Vec { + encode_message("event", event_type, &body.to_string()) +} + +/// `createModel({ generateUsage, streamUsage })`: `/converse` answers a JSON +/// response, every other route an event stream. +fn create_model(generate_usage: Value, stream_usage: Value) -> BedrockModel { + let fetch = RouteFetch::new(move |seen| { + if seen.url.ends_with("/converse") { + Canned::json(&json!({ + "output": { "message": { "role": "assistant", "content": [] } }, + "stopReason": "end_turn", + "usage": generate_usage, + })) + } else { + let mut body = create_event_frame("messageStop", json!({ "stopReason": "end_turn" })); + body.extend(create_event_frame( + "metadata", + json!({ "metrics": { "latencyMs": 1 }, "usage": stream_usage }), + )); + Canned { + status: 200, + headers: vec![( + "content-type".into(), + "application/vnd.amazon.eventstream".into(), + )], + body, + } + } + }); + create_amazon_bedrock(AmazonBedrockProviderSettings { + base_url: Some("https://bedrock-runtime.us-east-1.amazonaws.com".into()), + api_key: Some(Resolvable::Value("test-api-key".into())), + fetch: Some(Arc::clone(&fetch).transport()), + ..Default::default() + }) + .unwrap() + .chat("anthropic.claude-3-haiku-20240307-v1:0") +} + +fn expect_normalized_usage(usage: &Usage) { + assert_eq!( + serde_json::to_value(&usage.input_tokens).unwrap(), + json!({ "total": 18, "noCache": 13, "cacheRead": 3, "cacheWrite": 2 }) + ); + assert_eq!( + serde_json::to_value(&usage.output_tokens).unwrap(), + json!({ "total": 5, "text": 5 }) + ); +} + +/// `expectInvalidResponseField`: an `APICallError` with message +/// `Invalid JSON response` whose validation cause names the field. aimux +/// carries the cause inside the message (`Invalid JSON response: `), +/// and its cause is serde's text (`inputTokens: invalid type: …`), not zod's +/// issue list, where the field appears as a quoted JSON string. +fn expect_invalid_response_field(result: Result, field: &str) { + let error = result.expect_err("expected an error"); + let AiMuxError::ApiCall(error) = error else { + panic!("expected APICallError, got {error:?}"); + }; + let (message, cause) = error + .message + .split_once(": ") + .unwrap_or((error.message.as_str(), "")); + assert_eq!(message, "Invalid JSON response"); + assert!( + cause.contains(field), + "cause does not name \"{field}\": {cause}" + ); +} + +async fn stream_parts(model: &BedrockModel) -> Vec { + let mut options = CallOptions::new(prompt()); + options.include_raw_chunks = Some(false); + let result = model.do_stream(&options).await.unwrap(); + result + .stream + .map(|part| part.expect("stream item")) + .collect() + .await +} + +mod amazon_bedrock_chat_language_model_raw_usage { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-usage-raw.test.ts › AmazonBedrockChatLanguageModel raw usage › preserves complete raw usage through doGenerate parsing + #[tokio::test] + async fn preserves_complete_raw_usage_through_do_generate_parsing() { + let result = create_model(complete_usage(), complete_usage()) + .do_generate(&CallOptions::new(prompt())) + .await + .unwrap(); + + assert_eq!( + result.usage.raw.clone().map(Value::Object), + Some(complete_usage()) + ); + expect_normalized_usage(&result.usage); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-usage-raw.test.ts › AmazonBedrockChatLanguageModel raw usage › preserves complete raw usage through the event-stream parser + #[tokio::test] + async fn preserves_complete_raw_usage_through_the_event_stream_parser() { + let parts = stream_parts(&create_model(complete_usage(), complete_usage())).await; + let usage = parts + .iter() + .find_map(|part| match part { + StreamPart::Finish { usage, .. } => Some(usage), + _ => None, + }) + .expect("Expected a finish part"); + + assert_eq!(usage.raw.clone().map(Value::Object), Some(complete_usage())); + expect_normalized_usage(usage); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-usage-raw.test.ts › AmazonBedrockChatLanguageModel raw usage › validates known cache detail fields during doGenerate parsing + #[tokio::test] + async fn validates_known_cache_detail_fields_during_do_generate_parsing() { + let mut generate_usage = complete_usage(); + generate_usage["cacheDetails"] = json!([{ "inputTokens": "2", "ttl": "T5M" }]); + let model = create_model(generate_usage, complete_usage()); + + expect_invalid_response_field( + model.do_generate(&CallOptions::new(prompt())).await, + "inputTokens", + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-usage-raw.test.ts › AmazonBedrockChatLanguageModel raw usage › validates totalTokens in streaming usage + #[tokio::test] + async fn validates_total_tokens_in_streaming_usage() { + let mut stream_usage = complete_usage(); + stream_usage["totalTokens"] = json!("18"); + let parts = stream_parts(&create_model(complete_usage(), stream_usage)).await; + + assert!( + parts.iter().any(|part| matches!( + part, + StreamPart::Error { error } if error.to_string().contains("totalTokens") + )), + "{parts:?}" + ); + } +} diff --git a/aimux-providers/tests/amazon_bedrock_image_test.rs b/aimux-providers/tests/amazon_bedrock_image_test.rs index a7f355ad6..903e239b1 100644 --- a/aimux-providers/tests/amazon_bedrock_image_test.rs +++ b/aimux-providers/tests/amazon_bedrock_image_test.rs @@ -36,6 +36,7 @@ fn options(prompt: &str) -> ImageCallOptions { ImageCallOptions::new(prompt.to_string()) } +/// partial: `should extract the generated images`; full port in src/bedrock/tests/image_model.rs. #[tokio::test] async fn should_extract_generated_images() { let server = MockServer::start().await; @@ -52,6 +53,7 @@ async fn should_extract_generated_images() { } } +/// partial: `should pass the model and the settings`; full port in src/bedrock/tests/image_model.rs. #[tokio::test] async fn should_pass_prompt() { let server = MockServer::start().await; @@ -65,6 +67,7 @@ async fn should_pass_prompt() { assert_eq!(body["textToImageParams"]["text"], PROMPT); } +/// partial: `should pass the model and the settings`; full port in src/bedrock/tests/image_model.rs. #[tokio::test] async fn should_pass_size_and_n() { let server = MockServer::start().await; @@ -82,6 +85,7 @@ async fn should_pass_size_and_n() { assert_eq!(body["imageGenerationConfig"]["numberOfImages"], 2); } +/// partial: `should pass the model and the settings`; full port in src/bedrock/tests/image_model.rs. #[tokio::test] async fn should_pass_seed() { let server = MockServer::start().await; @@ -97,6 +101,7 @@ async fn should_pass_seed() { assert_eq!(body["imageGenerationConfig"]["seed"], 42); } +/// partial: `should return warnings for unsupported settings`; full port in src/bedrock/tests/image_model.rs. #[tokio::test] async fn should_warn_for_aspect_ratio() { let server = MockServer::start().await; @@ -110,6 +115,7 @@ async fn should_warn_for_aspect_ratio() { assert!(result.warnings.iter().any(|w| matches!(w, aimux_core::types::Warning::Unsupported { feature, .. } if feature == "aspectRatio"))); } +/// partial: `should default to IMAGE_VARIATION when files provided without mask or maskPrompt`; full port in src/bedrock/tests/image_model.rs. #[tokio::test] async fn should_support_image_variation() { let server = MockServer::start().await; @@ -129,6 +135,7 @@ async fn should_support_image_variation() { assert!(body["imageVariationParams"]["images"].is_array()); } +/// partial: `should default to INPAINTING when files provided with mask`; full port in src/bedrock/tests/image_model.rs. #[tokio::test] async fn should_support_inpainting() { let server = MockServer::start().await; @@ -152,25 +159,30 @@ async fn should_support_inpainting() { assert!(body["inPaintingParams"]["maskImage"].is_string()); } -#[tokio::test] -async fn should_handle_moderated_response() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/model/amazon.nova-canvas-v1%3A0/invoke")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "id": "fe7256d1-50d9-4663-8592-85eaf002e80c", - "status": "Request Moderated", - "result": null, - "progress": null, - "details": { "Moderation Reasons": ["Derivative Works Filter"] }, - "preview": null - }))) - .mount(&server) - .await; - let provider = create_amazon_bedrock(config(&server)).unwrap(); - let model = provider.image_model("amazon.nova-canvas-v1:0").unwrap(); - let result = model.do_generate(&options(PROMPT)).await; - assert!(result.is_err()); - let err = result.unwrap_err().to_string(); - assert!(err.contains("Amazon Bedrock request was moderated: Derivative Works Filter")); +mod do_generate { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-image-model.test.ts › doGenerate › should throw error when request is moderated + #[tokio::test] + async fn should_handle_moderated_response() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/model/amazon.nova-canvas-v1%3A0/invoke")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "fe7256d1-50d9-4663-8592-85eaf002e80c", + "status": "Request Moderated", + "result": null, + "progress": null, + "details": { "Moderation Reasons": ["Derivative Works Filter"] }, + "preview": null + }))) + .mount(&server) + .await; + let provider = create_amazon_bedrock(config(&server)).unwrap(); + let model = provider.image_model("amazon.nova-canvas-v1:0").unwrap(); + let result = model.do_generate(&options(PROMPT)).await; + assert!(result.is_err()); + let err = result.unwrap_err().to_string(); + assert!(err.contains("Amazon Bedrock request was moderated: Derivative Works Filter")); + } } diff --git a/aimux-providers/tests/bedrock_convert_test.rs b/aimux-providers/tests/bedrock_convert_test.rs index 892a47a0f..dc5a2a001 100644 --- a/aimux-providers/tests/bedrock_convert_test.rs +++ b/aimux-providers/tests/bedrock_convert_test.rs @@ -129,47 +129,271 @@ fn tool_result(id: &str, output: ToolResultOutput) -> ToolPart { // ── system messages ───────────────────────────────────────────────────────── -/// TS: "should combine multiple leading system messages into a single system message" -#[test] -fn system_combine_multiple_leading() { - let (system, _) = - convert_prompt_to_bedrock(&vec![system_msg("Hello"), system_msg("World")]).unwrap(); - assert_eq!( - Value::Array(system), - json!([{ "text": "Hello" }, { "text": "World" }]) - ); +mod system_messages { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › system messages › should combine multiple leading system messages into a single system message + #[test] + fn system_combine_multiple_leading() { + let (system, _) = + convert_prompt_to_bedrock(&vec![system_msg("Hello"), system_msg("World")]).unwrap(); + assert_eq!( + Value::Array(system), + json!([{ "text": "Hello" }, { "text": "World" }]) + ); + } } -/// TS: "should extract the system message" -#[test] -fn system_extract_single() { - let (system, _) = convert_prompt_to_bedrock(&vec![system_msg("Hello")]).unwrap(); - assert_eq!(Value::Array(system), json!([{ "text": "Hello" }])); +mod user_messages { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should extract the system message + #[test] + fn system_extract_single() { + let (system, _) = convert_prompt_to_bedrock(&vec![system_msg("Hello")]).unwrap(); + assert_eq!(Value::Array(system), json!([{ "text": "Hello" }])); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should convert messages with image parts + #[test] + fn user_convert_image_parts() { + let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![ + user_text("Hello"), + file_base64("AAECAw==", "image/png", None, None), + ])]) + .unwrap(); + assert_eq!( + Value::Array(messages), + json!([{ + "role": "user", + "content": [ + { "text": "Hello" }, + { "image": { "format": "png", "source": { "bytes": "AAECAw==" } } }, + ] + }]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should use consistent document names for prompt cache effectiveness + #[test] + fn user_consistent_document_names() { + let (_, messages) = convert_prompt_to_bedrock(&vec![ + user(vec![ + file_base64("AAECAw==", "application/pdf", None, None), + file_base64("BAUGBw==", "application/pdf", None, None), + ]), + assistant(vec![assistant_text("OK")]), + user(vec![file_base64("AAECAw==", "application/pdf", None, None)]), + ]) + .unwrap(); + assert_eq!( + Value::Array(messages), + json!([ + { "role": "user", "content": [ + { "document": { "format": "pdf", "name": "document-1", "source": { "bytes": "AAECAw==" } } }, + { "document": { "format": "pdf", "name": "document-2", "source": { "bytes": "BAUGBw==" } } }, + ]}, + { "role": "assistant", "content": [{ "text": "OK" }] }, + { "role": "user", "content": [ + { "document": { "format": "pdf", "name": "document-3", "source": { "bytes": "AAECAw==" } } }, + ]}, + ]) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › user messages › should add cache point to user content part when specified + #[test] + fn user_content_part_cache_point() { + let (system, messages) = convert_prompt_to_bedrock(&vec![user(vec![ + user_text("Hello"), + UserPart::Text(text_with_cache("cached", "default", Some("5m"))), + user_text("World"), + ])]) + .unwrap(); + assert!(system.is_empty()); + assert_eq!( + Value::Array(messages), + json!([{ + "role": "user", + "content": [ + { "text": "Hello" }, + { "text": "cached" }, + { "cachePoint": { "type": "default", "ttl": "5m" } }, + { "text": "World" }, + ] + }]) + ); + } } -// ── user messages ─────────────────────────────────────────────────────────── +mod tool_messages { + use super::*; -/// TS: "should convert messages with image parts" -#[test] -fn user_convert_image_parts() { - let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![ - user_text("Hello"), - file_base64("AAECAw==", "image/png", None, None), - ])]) - .unwrap(); - assert_eq!( - Value::Array(messages), - json!([{ - "role": "user", - "content": [ - { "text": "Hello" }, - { "image": { "format": "png", "source": { "bytes": "AAECAw==" } } }, - ] - }]) - ); + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › tool messages › should convert tool result with content array containing text + #[test] + fn tool_result_content_text() { + let (_, messages) = convert_prompt_to_bedrock(&vec![tool_msg(vec![tool_result( + "call-123", + ToolResultOutput::Content { + value: vec![ToolResultContent::Text(TextPart { + text: "The result is 42".to_string(), + provider_options: None, + })], + }, + )])]) + .unwrap(); + assert_eq!( + messages[0], + json!({ + "role": "user", + "content": [{ + "toolResult": { + "toolUseId": "call-123", + "content": [{ "text": "The result is 42" }] + } + }] + }) + ); + } } -/// TS: "should convert messages with document parts" +mod citations { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › citations › should handle citations enabled for PDF + #[test] + fn citations_enabled_for_pdf() { + let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![file_base64( + "AAECAw==", + "application/pdf", + None, + Some( + provider_namespace("bedrock", json!({ "citations": { "enabled": true } })).unwrap(), + ), + )])]) + .unwrap(); + assert_eq!( + messages[0]["content"][0], + json!({ + "document": { + "format": "pdf", + "name": "document-1", + "source": { "bytes": "AAECAw==" }, + "citations": { "enabled": true } + } + }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › citations › should handle citations disabled for PDF + #[test] + fn citations_disabled_for_pdf() { + let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![file_base64( + "AAECAw==", + "application/pdf", + None, + Some( + provider_namespace("bedrock", json!({ "citations": { "enabled": false } })) + .unwrap(), + ), + )])]) + .unwrap(); + assert_eq!( + messages[0]["content"][0], + json!({ + "document": { "format": "pdf", "name": "document-1", "source": { "bytes": "AAECAw==" } } + }) + ); + assert!( + messages[0]["content"][0]["document"] + .get("citations") + .is_none() + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › citations › should handle no citations specified for PDF (default) + #[test] + fn citations_default_for_pdf() { + let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![file_base64( + "AAECAw==", + "application/pdf", + None, + None, + )])]) + .unwrap(); + assert_eq!( + messages[0]["content"][0], + json!({ + "document": { "format": "pdf", "name": "document-1", "source": { "bytes": "AAECAw==" } } + }) + ); + assert!( + messages[0]["content"][0]["document"] + .get("citations") + .is_none() + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › citations › should handle multiple PDFs with different citation settings + #[test] + fn citations_multiple_pdfs() { + let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![ + file_base64( + "AAECAw==", + "application/pdf", + None, + Some( + provider_namespace("bedrock", json!({ "citations": { "enabled": true } })) + .unwrap(), + ), + ), + file_base64( + "BAUGBw==", + "application/pdf", + None, + Some( + provider_namespace("bedrock", json!({ "citations": { "enabled": false } })) + .unwrap(), + ), + ), + ])]) + .unwrap(); + assert_eq!( + messages[0]["content"], + json!([ + { "document": { "format": "pdf", "name": "document-1", "source": { "bytes": "AAECAw==" }, "citations": { "enabled": true } } }, + { "document": { "format": "pdf", "name": "document-2", "source": { "bytes": "BAUGBw==" } } }, + ]) + ); + } +} + +mod top_level_only_media_type_resolution { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/convert-to-amazon-bedrock-chat-messages.test.ts › top-level-only mediaType resolution › should pass through a full image mediaType unchanged + #[test] + fn media_type_pass_through_full_image() { + let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![file_base64( + "iVBORw0KGgo=", + "image/png", + None, + None, + )])]) + .unwrap(); + assert_eq!( + Value::Array(messages), + json!([{ + "role": "user", + "content": [{ "image": { "format": "png", "source": { "bytes": "iVBORw0KGgo=" } } }] + }]) + ); + } +} + +// ── user messages ─────────────────────────────────────────────────────────── + +/// partial: `should convert messages with document parts`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn user_convert_document_parts() { let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![ @@ -186,7 +410,7 @@ fn user_convert_document_parts() { ); } -/// TS: "should strip file extension when filename is provided" +/// partial: `should strip file extension when filename is provided`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn user_strip_file_extension() { let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![ @@ -205,7 +429,7 @@ fn user_strip_file_extension() { ); } -/// TS: "should preserve filename without extension when provided" +/// partial: `should preserve filename without extension when provided`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn user_preserve_filename_without_extension() { let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![file_base64( @@ -221,60 +445,9 @@ fn user_preserve_filename_without_extension() { ); } -/// TS: "should use consistent document names for prompt cache effectiveness" -#[test] -fn user_consistent_document_names() { - let (_, messages) = convert_prompt_to_bedrock(&vec![ - user(vec![ - file_base64("AAECAw==", "application/pdf", None, None), - file_base64("BAUGBw==", "application/pdf", None, None), - ]), - assistant(vec![assistant_text("OK")]), - user(vec![file_base64("AAECAw==", "application/pdf", None, None)]), - ]) - .unwrap(); - assert_eq!( - Value::Array(messages), - json!([ - { "role": "user", "content": [ - { "document": { "format": "pdf", "name": "document-1", "source": { "bytes": "AAECAw==" } } }, - { "document": { "format": "pdf", "name": "document-2", "source": { "bytes": "BAUGBw==" } } }, - ]}, - { "role": "assistant", "content": [{ "text": "OK" }] }, - { "role": "user", "content": [ - { "document": { "format": "pdf", "name": "document-3", "source": { "bytes": "AAECAw==" } } }, - ]}, - ]) - ); -} - -/// TS: "should add cache point to user content part when specified" -#[test] -fn user_content_part_cache_point() { - let (system, messages) = convert_prompt_to_bedrock(&vec![user(vec![ - user_text("Hello"), - UserPart::Text(text_with_cache("cached", "default", Some("5m"))), - user_text("World"), - ])]) - .unwrap(); - assert!(system.is_empty()); - assert_eq!( - Value::Array(messages), - json!([{ - "role": "user", - "content": [ - { "text": "Hello" }, - { "text": "cached" }, - { "cachePoint": { "type": "default", "ttl": "5m" } }, - { "text": "World" }, - ] - }]) - ); -} - // ── assistant messages ────────────────────────────────────────────────────── -/// TS: "should remove trailing whitespace from last assistant message when there is no further user message" +/// partial: `should remove trailing whitespace from last assistant message when there is no further user message`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_trim_trailing_whitespace_last() { let (_, messages) = convert_prompt_to_bedrock(&vec![ @@ -291,7 +464,7 @@ fn assistant_trim_trailing_whitespace_last() { ); } -/// TS: "should remove trailing whitespace from last assistant message with multi-part content when there is no further user message" +/// partial: `should remove trailing whitespace from last assistant message with multi-part content when there is no further user message`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_trim_trailing_whitespace_multi_part() { let (_, messages) = convert_prompt_to_bedrock(&vec![ @@ -308,7 +481,7 @@ fn assistant_trim_trailing_whitespace_multi_part() { ); } -/// TS: "should keep trailing whitespace from assistant message when there is a further user message" +/// partial: `should keep trailing whitespace from assistant message when there is a further user message`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_keep_trailing_whitespace_with_further_user() { let (_, messages) = convert_prompt_to_bedrock(&vec![ @@ -323,7 +496,7 @@ fn assistant_keep_trailing_whitespace_with_further_user() { ); } -/// TS: "should combine multiple sequential assistant messages into a single message" +/// partial: `should combine multiple sequential assistant messages into a single message`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_combine_sequential() { let (_, messages) = convert_prompt_to_bedrock(&vec![ @@ -342,7 +515,7 @@ fn assistant_combine_sequential() { ); } -/// TS: "should add cache point to assistant content part when specified" +/// partial: `should add cache point to assistant content part when specified`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_content_part_cache_point() { let (_, messages) = convert_prompt_to_bedrock(&vec![assistant(vec![ @@ -365,7 +538,7 @@ fn assistant_content_part_cache_point() { ); } -/// TS: "should properly convert reasoning content type" +/// partial: `should properly convert reasoning content type`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_reasoning_with_signature() { let (_, messages) = convert_prompt_to_bedrock(&vec![ @@ -389,7 +562,7 @@ fn assistant_reasoning_with_signature() { ); } -/// TS: "should preserve assistant message reasoning parts with amazonBedrock providerOptions" +/// partial: `should preserve assistant message reasoning parts with amazonBedrock providerOptions`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_preserve_amazon_bedrock_reasoning() { let (_, messages) = convert_prompt_to_bedrock(&vec![ @@ -415,7 +588,7 @@ fn assistant_preserve_amazon_bedrock_reasoning() { ); } -/// TS: "should not trim reasoning text when a signature is present" +/// partial: `should not trim reasoning text when a signature is present`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_no_trim_reasoning_with_signature() { let (_, messages) = convert_prompt_to_bedrock(&vec![ @@ -432,7 +605,7 @@ fn assistant_no_trim_reasoning_with_signature() { ); } -/// TS: "should omit reasoning content without signature" +/// partial: `should omit reasoning content without signature`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_omit_reasoning_without_signature() { let (_, messages) = convert_prompt_to_bedrock(&vec![ @@ -446,7 +619,7 @@ fn assistant_omit_reasoning_without_signature() { assert_eq!(messages[1]["content"], json!([{ "text": "final answer" }])); } -/// TS: "should omit multiple reasoning parts without signatures" +/// partial: `should omit multiple reasoning parts without signatures`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_omit_multiple_reasoning_without_signatures() { let (_, messages) = convert_prompt_to_bedrock(&vec![ @@ -461,7 +634,7 @@ fn assistant_omit_multiple_reasoning_without_signatures() { assert_eq!(messages[1]["content"], json!([{ "text": "final answer" }])); } -/// TS: "should omit unsigned reasoning while preserving tool calls in multi-turn tool use" +/// partial: `should omit unsigned reasoning while preserving tool calls in multi-turn tool use`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_omit_unsigned_reasoning_preserving_tool_calls() { let (_, messages) = convert_prompt_to_bedrock(&vec![ @@ -493,7 +666,7 @@ fn assistant_omit_unsigned_reasoning_preserving_tool_calls() { ); } -/// TS: "should preserve reasoning text with signature in multi-turn tool use" +/// partial: `should preserve reasoning text with signature in multi-turn tool use`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_preserve_reasoning_multi_turn() { let (_, messages) = convert_prompt_to_bedrock(&vec![ @@ -538,7 +711,7 @@ fn assistant_preserve_reasoning_multi_turn() { ); } -/// TS: "should handle a mix of text and reasoning content types" +/// partial: `should handle a mix of text and reasoning content types`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_mix_text_and_reasoning() { let (_, messages) = convert_prompt_to_bedrock(&vec![ @@ -564,7 +737,7 @@ fn assistant_mix_text_and_reasoning() { ); } -/// TS: "should filter out empty text blocks in assistant messages" +/// partial: `should filter out empty text blocks in assistant messages`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_filter_empty_text_blocks() { let (_, messages) = convert_prompt_to_bedrock(&vec![ @@ -586,7 +759,7 @@ fn assistant_filter_empty_text_blocks() { ); } -/// TS: "should wrap non-object (invalid) tool call input in an object" +/// partial: `should wrap non-object (invalid) tool call input in an object`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_wrap_non_object_tool_input() { let (_, messages) = convert_prompt_to_bedrock(&vec![assistant(vec![tool_call( @@ -607,7 +780,7 @@ fn assistant_wrap_non_object_tool_input() { ); } -/// TS: "should strip invalid characters from tool call names" +/// partial: `should strip invalid characters from tool call names`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_strip_invalid_tool_name_chars() { let (_, messages) = convert_prompt_to_bedrock(&vec![assistant(vec![ @@ -630,7 +803,7 @@ fn assistant_strip_invalid_tool_name_chars() { ); } -/// TS: "should preserve empty text blocks when reasoning blocks are present" +/// partial: `should preserve empty text blocks when reasoning blocks are present`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn assistant_preserve_empty_text_with_reasoning() { let (_, messages) = convert_prompt_to_bedrock(&vec![ @@ -658,34 +831,7 @@ fn assistant_preserve_empty_text_with_reasoning() { // ── tool messages ─────────────────────────────────────────────────────────── -/// TS: "should convert tool result with content array containing text" -#[test] -fn tool_result_content_text() { - let (_, messages) = convert_prompt_to_bedrock(&vec![tool_msg(vec![tool_result( - "call-123", - ToolResultOutput::Content { - value: vec![ToolResultContent::Text(TextPart { - text: "The result is 42".to_string(), - provider_options: None, - })], - }, - )])]) - .unwrap(); - assert_eq!( - messages[0], - json!({ - "role": "user", - "content": [{ - "toolResult": { - "toolUseId": "call-123", - "content": [{ "text": "The result is 42" }] - } - }] - }) - ); -} - -/// TS: "should convert tool result with content array containing image" +/// partial: `should convert tool result with content array containing image`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn tool_result_content_image() { let (_, messages) = convert_prompt_to_bedrock(&vec![tool_msg(vec![tool_result( @@ -713,7 +859,7 @@ fn tool_result_content_image() { ); } -/// TS: "should convert tool result with content array containing PDF" +/// partial: `should convert tool result with content array containing PDF`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn tool_result_content_pdf() { let (_, messages) = convert_prompt_to_bedrock(&vec![tool_msg(vec![tool_result( @@ -747,7 +893,8 @@ fn tool_result_content_pdf() { // content" and "should throw error for unsupported mime type in tool result // file content"). -/// TS: "should fallback to stringified result when content is undefined" (json output) +/// partial: `should fallback to stringified result when content is undefined`; full port in src/bedrock/tests/convert_to_chat_messages.rs. +/// (json output) #[test] fn tool_result_json_output() { let (_, messages) = convert_prompt_to_bedrock(&vec![tool_msg(vec![tool_result( @@ -771,110 +918,9 @@ fn tool_result_json_output() { // ── citations ─────────────────────────────────────────────────────────────── -/// TS: "should handle citations enabled for PDF" -#[test] -fn citations_enabled_for_pdf() { - let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![file_base64( - "AAECAw==", - "application/pdf", - None, - Some(provider_namespace("bedrock", json!({ "citations": { "enabled": true } })).unwrap()), - )])]) - .unwrap(); - assert_eq!( - messages[0]["content"][0], - json!({ - "document": { - "format": "pdf", - "name": "document-1", - "source": { "bytes": "AAECAw==" }, - "citations": { "enabled": true } - } - }) - ); -} - -/// TS: "should handle citations disabled for PDF" -#[test] -fn citations_disabled_for_pdf() { - let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![file_base64( - "AAECAw==", - "application/pdf", - None, - Some(provider_namespace("bedrock", json!({ "citations": { "enabled": false } })).unwrap()), - )])]) - .unwrap(); - assert_eq!( - messages[0]["content"][0], - json!({ - "document": { "format": "pdf", "name": "document-1", "source": { "bytes": "AAECAw==" } } - }) - ); - assert!( - messages[0]["content"][0]["document"] - .get("citations") - .is_none() - ); -} - -/// TS: "should handle no citations specified for PDF (default)" -#[test] -fn citations_default_for_pdf() { - let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![file_base64( - "AAECAw==", - "application/pdf", - None, - None, - )])]) - .unwrap(); - assert_eq!( - messages[0]["content"][0], - json!({ - "document": { "format": "pdf", "name": "document-1", "source": { "bytes": "AAECAw==" } } - }) - ); - assert!( - messages[0]["content"][0]["document"] - .get("citations") - .is_none() - ); -} - -/// TS: "should handle multiple PDFs with different citation settings" -#[test] -fn citations_multiple_pdfs() { - let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![ - file_base64( - "AAECAw==", - "application/pdf", - None, - Some( - provider_namespace("bedrock", json!({ "citations": { "enabled": true } })).unwrap(), - ), - ), - file_base64( - "BAUGBw==", - "application/pdf", - None, - Some( - provider_namespace("bedrock", json!({ "citations": { "enabled": false } })) - .unwrap(), - ), - ), - ])]) - .unwrap(); - assert_eq!( - messages[0]["content"], - json!([ - { "document": { "format": "pdf", "name": "document-1", "source": { "bytes": "AAECAw==" }, "citations": { "enabled": true } } }, - { "document": { "format": "pdf", "name": "document-2", "source": { "bytes": "BAUGBw==" } } }, - ]) - ); -} - // ── additional file format tests ──────────────────────────────────────────── -/// TS: "should handle xlsx files correctly" +/// partial: `should handle xlsx files correctly`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn file_format_xlsx() { let (system, messages) = convert_prompt_to_bedrock(&vec![user(vec![file_base64( @@ -891,7 +937,7 @@ fn file_format_xlsx() { ); } -/// TS: "should handle docx files correctly" +/// partial: `should handle docx files correctly`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn file_format_docx() { let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![file_base64( @@ -909,7 +955,7 @@ fn file_format_docx() { // ── Mistral tool call ID normalization ────────────────────────────────────── -/// TS: "should not normalize tool call IDs when isMistral is false" +/// partial: `should not normalize tool call IDs when isMistral is false`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn mistral_no_normalize_when_false() { let original_id = "tooluse_bpe71yCfRu2b5i-nKGDr5g"; @@ -927,7 +973,7 @@ fn mistral_no_normalize_when_false() { ); } -/// TS: "should default to not normalizing when isMistral is not provided" +/// partial: `should default to not normalizing when isMistral is not provided`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn mistral_default_no_normalize() { let original_id = "tooluse_bpe71yCfRu2b5i-nKGDr5g"; @@ -947,26 +993,7 @@ fn mistral_default_no_normalize() { // ── top-level-only mediaType resolution ───────────────────────────────────── -/// TS: "should pass through a full image mediaType unchanged" -#[test] -fn media_type_pass_through_full_image() { - let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![file_base64( - "iVBORw0KGgo=", - "image/png", - None, - None, - )])]) - .unwrap(); - assert_eq!( - Value::Array(messages), - json!([{ - "role": "user", - "content": [{ "image": { "format": "png", "source": { "bytes": "iVBORw0KGgo=" } } }] - }]) - ); -} - -/// TS: "should route to document slot for non-image top-level type via getTopLevelMediaType" +/// partial: `should route to document slot for non-image top-level type via getTopLevelMediaType`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn media_type_route_to_document_text_plain() { let (_, messages) = convert_prompt_to_bedrock(&vec![user(vec![file_base64( @@ -999,79 +1026,87 @@ fn func_tool(name: &str, description: Option<&str>, input_schema: Value) -> Func } } -/// 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, - Some(&LanguageModelV4ToolChoice::Auto), - NON_ANTHROPIC_MODEL, - ); - assert_eq!(config["toolChoice"], json!({ "auto": {} })); -} - -/// TS: "should handle tool choice 'required'" -#[test] -fn prepare_tools_tool_choice_required() { - let tools = Some(vec![func_tool("testFunction", Some("Test"), json!({}))]); - let config = prepare_tools( - &tools, - Some(&LanguageModelV4ToolChoice::Required), - NON_ANTHROPIC_MODEL, - ); - assert_eq!(config["toolChoice"], json!({ "any": {} })); -} - -/// TS: "should handle tool choice 'none' by clearing tools" -#[test] -fn prepare_tools_tool_choice_none_clears() { - let tools = Some(vec![func_tool("testFunction", Some("Test"), json!({}))]); - let config = prepare_tools( - &tools, - Some(&LanguageModelV4ToolChoice::None), - NON_ANTHROPIC_MODEL, - ); - assert_eq!(config, json!({})); -} - -/// TS: "should handle tool choice 'tool'" -#[test] -fn prepare_tools_tool_choice_tool() { - let tools = Some(vec![func_tool("testFunction", Some("Test"), json!({}))]); - let config = prepare_tools( - &tools, - Some(&LanguageModelV4ToolChoice::Tool { - tool_name: "testFunction".to_string(), - }), - NON_ANTHROPIC_MODEL, - ); - assert_eq!( - config["toolChoice"], - json!({ "tool": { "name": "testFunction" } }) - ); -} - -/// TS: "should filter function tools to only the named tool when tool choice is 'tool'" -#[test] -fn prepare_tools_tool_choice_filters_to_named() { - let tools = Some(vec![ - func_tool( - "getWeather", - Some("Get weather"), - json!({ "type": "object" }), - ), - func_tool("getTime", Some("Get time"), json!({ "type": "object" })), - ]); - let config = prepare_tools( - &tools, - Some(&LanguageModelV4ToolChoice::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")); +mod prepare_tools { + use super::*; + + mod tool_choice { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › tool choice › 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, + Some(&LanguageModelV4ToolChoice::Auto), + NON_ANTHROPIC_MODEL, + ); + assert_eq!(config["toolChoice"], json!({ "auto": {} })); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › tool choice › should handle tool choice "required" + #[test] + fn prepare_tools_tool_choice_required() { + let tools = Some(vec![func_tool("testFunction", Some("Test"), json!({}))]); + let config = prepare_tools( + &tools, + Some(&LanguageModelV4ToolChoice::Required), + NON_ANTHROPIC_MODEL, + ); + assert_eq!(config["toolChoice"], json!({ "any": {} })); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › tool choice › should handle tool choice "none" by clearing tools + #[test] + fn prepare_tools_tool_choice_none_clears() { + let tools = Some(vec![func_tool("testFunction", Some("Test"), json!({}))]); + let config = prepare_tools( + &tools, + Some(&LanguageModelV4ToolChoice::None), + NON_ANTHROPIC_MODEL, + ); + assert_eq!(config, json!({})); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › tool choice › should handle tool choice "tool" + #[test] + fn prepare_tools_tool_choice_tool() { + let tools = Some(vec![func_tool("testFunction", Some("Test"), json!({}))]); + let config = prepare_tools( + &tools, + Some(&LanguageModelV4ToolChoice::Tool { + tool_name: "testFunction".to_string(), + }), + NON_ANTHROPIC_MODEL, + ); + assert_eq!( + config["toolChoice"], + json!({ "tool": { "name": "testFunction" } }) + ); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-prepare-tools.test.ts › prepareTools › tool choice › should filter function tools to only the named tool when tool choice is "tool" + #[test] + fn prepare_tools_tool_choice_filters_to_named() { + let tools = Some(vec![ + func_tool( + "getWeather", + Some("Get weather"), + json!({ "type": "object" }), + ), + func_tool("getTime", Some("Get time"), json!({ "type": "object" })), + ]); + let config = prepare_tools( + &tools, + Some(&LanguageModelV4ToolChoice::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")); + } + } } // ════════════════════════════════════════════════════════════════════════════ @@ -1097,7 +1132,7 @@ fn bedrock_usage( } } -/// TS: "should convert basic usage without cache tokens" +/// partial: `should convert basic usage without cache tokens`; full port in src/bedrock/tests/convert_usage.rs. #[test] fn usage_basic_without_cache() { let u = convert_usage(Some(&bedrock_usage(100, 50, None, None))); @@ -1109,7 +1144,7 @@ fn usage_basic_without_cache() { assert_eq!(u.output_tokens.text, Some(50)); } -/// TS: "should convert usage with cache read tokens" +/// partial: `should convert usage with cache read tokens`; full port in src/bedrock/tests/convert_usage.rs. #[test] fn usage_with_cache_read() { let u = convert_usage(Some(&bedrock_usage(100, 50, Some(80), None))); @@ -1121,7 +1156,7 @@ fn usage_with_cache_read() { assert_eq!(u.output_tokens.text, Some(50)); } -/// TS: "should convert usage with cache write tokens" +/// partial: `should convert usage with cache write tokens`; full port in src/bedrock/tests/convert_usage.rs. #[test] fn usage_with_cache_write() { let u = convert_usage(Some(&bedrock_usage(100, 50, None, Some(60)))); @@ -1133,7 +1168,7 @@ fn usage_with_cache_write() { assert_eq!(u.output_tokens.text, Some(50)); } -/// TS: "should convert usage with both cache read and write tokens" +/// partial: `should convert usage with both cache read and write tokens`; full port in src/bedrock/tests/convert_usage.rs. #[test] fn usage_with_both_cache() { let u = convert_usage(Some(&bedrock_usage(100, 50, Some(80), Some(60)))); @@ -1145,7 +1180,7 @@ fn usage_with_both_cache() { assert_eq!(u.output_tokens.text, Some(50)); } -/// TS: "should handle null cache tokens" +/// partial: `should handle null cache tokens`; full port in src/bedrock/tests/convert_usage.rs. #[test] fn usage_null_cache_tokens() { let u = convert_usage(Some(&bedrock_usage(100, 50, None, None))); @@ -1157,7 +1192,7 @@ fn usage_null_cache_tokens() { assert_eq!(u.output_tokens.text, Some(50)); } -/// TS: "should handle null usage" +/// partial: `should handle null usage`; full port in src/bedrock/tests/convert_usage.rs. #[test] fn usage_null() { let u = convert_usage(None); @@ -1169,7 +1204,7 @@ fn usage_null() { assert_eq!(u.output_tokens.text, None); } -/// TS: "should handle undefined usage" +/// partial: `should handle undefined usage`; full port in src/bedrock/tests/convert_usage.rs. #[test] fn usage_undefined() { let u = convert_usage(None); diff --git a/aimux-providers/tests/bedrock_embedding_test.rs b/aimux-providers/tests/bedrock_embedding_test.rs index b663074e5..642dc728b 100644 --- a/aimux-providers/tests/bedrock_embedding_test.rs +++ b/aimux-providers/tests/bedrock_embedding_test.rs @@ -46,29 +46,7 @@ fn default_options(values: Vec) -> EmbeddingCallOptions { } } -/// TS: "should expose model-specific max embeddings per call" -#[test] -fn should_expose_max_embeddings_per_call() { - let provider = test_provider("https://example.com".to_string()); - - let titan_model = provider - .embedding_model("amazon.titan-embed-text-v2:0") - .unwrap(); - assert_eq!(titan_model.max_embeddings_per_call(), Some(1)); - - let cohere_model = provider.embedding_model("cohere.embed-english-v3").unwrap(); - assert_eq!(cohere_model.max_embeddings_per_call(), Some(96)); - - let cohere_v4_us_model = provider.embedding_model("us.cohere.embed-v4:0").unwrap(); - assert_eq!(cohere_v4_us_model.max_embeddings_per_call(), Some(96)); - - let nova_model = provider - .embedding_model("amazon.nova-2-multimodal-embeddings-v1:0") - .unwrap(); - assert_eq!(nova_model.max_embeddings_per_call(), Some(1)); -} - -/// TS: "should handle single input value and return embeddings" (Titan) +/// partial: `should handle single input value and return embeddings`; full port in src/bedrock/tests/embedding_model.rs. #[tokio::test] async fn should_handle_titan_single_input() { let server = MockServer::start().await; @@ -103,34 +81,7 @@ async fn should_handle_titan_single_input() { assert_eq!(body["inputText"], TEST_VALUES[0]); } -/// TS: "should handle single input value and extract usage" (Titan) -#[tokio::test] -async fn should_extract_titan_usage() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/model/amazon.titan-embed-text-v2%3A0/invoke")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "embedding": mock_embeddings()[0], - "inputTextTokenCount": 8 - }))) - .mount(&server) - .await; - - let provider = test_provider(server.uri()); - let model = provider - .embedding_model("amazon.titan-embed-text-v2:0") - .unwrap(); - - let result = model - .do_embed(&default_options(vec![TEST_VALUES[0].to_string()])) - .await - .expect("should succeed"); - - let usage = result.usage.expect("usage present"); - assert_eq!(usage.tokens, 8); -} - -/// TS: "should support Cohere embedding models" (Cohere v3 response) +/// partial: `should support Cohere embedding models`; full port in src/bedrock/tests/embedding_model.rs. #[tokio::test] async fn should_support_cohere_v3() { let server = MockServer::start().await; @@ -164,7 +115,7 @@ async fn should_support_cohere_v3() { assert_eq!(body["texts"], json!([TEST_VALUES[0]])); } -/// TS: "should support Cohere v4 embedding models" +/// partial: `should support Cohere v4 embedding models`; full port in src/bedrock/tests/embedding_model.rs. #[tokio::test] async fn should_support_cohere_v4() { let server = MockServer::start().await; @@ -193,7 +144,7 @@ async fn should_support_cohere_v4() { assert_eq!(result.usage.unwrap().tokens, 6); } -/// TS: "should send multiple values for Cohere embedding models" +/// partial: `should send multiple values for Cohere embedding models`; full port in src/bedrock/tests/embedding_model.rs. #[tokio::test] async fn should_send_multiple_values_cohere_v4() { let server = MockServer::start().await; @@ -234,7 +185,7 @@ async fn should_send_multiple_values_cohere_v4() { ); } -/// TS: "should support Cohere models behind cross-region inference profile ids" +/// partial: `should support Cohere models behind cross-region inference profile ids`; full port in src/bedrock/tests/embedding_model.rs. #[tokio::test] async fn should_support_cross_region_cohere() { let server = MockServer::start().await; @@ -261,7 +212,7 @@ async fn should_support_cross_region_cohere() { assert_eq!(result.usage.unwrap().tokens, 6); } -/// TS: "should send SINGLE_EMBEDDING payload for Nova embeddings" +/// partial: `should send SINGLE_EMBEDDING payload for Nova embeddings`; full port in src/bedrock/tests/embedding_model.rs. #[tokio::test] async fn should_support_nova_embeddings() { let server = MockServer::start().await; @@ -307,3 +258,56 @@ async fn should_support_nova_embeddings() { TEST_VALUES[0] ); } + +mod do_embed { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-embedding-model.test.ts › doEmbed › should expose model-specific max embeddings per call + #[test] + fn should_expose_max_embeddings_per_call() { + let provider = test_provider("https://example.com".to_string()); + + let titan_model = provider + .embedding_model("amazon.titan-embed-text-v2:0") + .unwrap(); + assert_eq!(titan_model.max_embeddings_per_call(), Some(1)); + + let cohere_model = provider.embedding_model("cohere.embed-english-v3").unwrap(); + assert_eq!(cohere_model.max_embeddings_per_call(), Some(96)); + + let cohere_v4_us_model = provider.embedding_model("us.cohere.embed-v4:0").unwrap(); + assert_eq!(cohere_v4_us_model.max_embeddings_per_call(), Some(96)); + + let nova_model = provider + .embedding_model("amazon.nova-2-multimodal-embeddings-v1:0") + .unwrap(); + assert_eq!(nova_model.max_embeddings_per_call(), Some(1)); + } + + /// upstream: @ai-sdk/amazon-bedrock src/amazon-bedrock-embedding-model.test.ts › doEmbed › should handle single input value and extract usage + #[tokio::test] + async fn should_extract_titan_usage() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/model/amazon.titan-embed-text-v2%3A0/invoke")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "embedding": mock_embeddings()[0], + "inputTextTokenCount": 8 + }))) + .mount(&server) + .await; + + let provider = test_provider(server.uri()); + let model = provider + .embedding_model("amazon.titan-embed-text-v2:0") + .unwrap(); + + let result = model + .do_embed(&default_options(vec![TEST_VALUES[0].to_string()])) + .await + .expect("should succeed"); + + let usage = result.usage.expect("usage present"); + assert_eq!(usage.tokens, 8); + } +} diff --git a/aimux-providers/tests/bedrock_remaining_test.rs b/aimux-providers/tests/bedrock_remaining_test.rs index d070ee43c..50e512b62 100644 --- a/aimux-providers/tests/bedrock_remaining_test.rs +++ b/aimux-providers/tests/bedrock_remaining_test.rs @@ -163,83 +163,91 @@ fn clear_bedrock_env() { } } -/// TS: "should return true for mistral models" -#[test] -fn is_mistral_model_true_for_mistral() { - for id in [ - "mistral.mistral-7b-instruct-v0:2", - "mistral.mixtral-8x7b-instruct-v0:1", - "mistral.mistral-large-2402-v1:0", - "mistral.mistral-small-2402-v1:0", - "mistral.mistral-large-2407-v1:0", - "mistral.ministral-3-14b-instruct", - "mistral.ministral-3-8b-instruct", - ] { - assert!(model_normalizes_tool_id(id)); +mod is_mistral_model { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/normalize-tool-call-id.test.ts › isMistralModel › should return true for mistral models + #[test] + fn is_mistral_model_true_for_mistral() { + for id in [ + "mistral.mistral-7b-instruct-v0:2", + "mistral.mixtral-8x7b-instruct-v0:1", + "mistral.mistral-large-2402-v1:0", + "mistral.mistral-small-2402-v1:0", + "mistral.mistral-large-2407-v1:0", + "mistral.ministral-3-14b-instruct", + "mistral.ministral-3-8b-instruct", + ] { + assert!(model_normalizes_tool_id(id)); + } } -} -/// TS: "should return true for region-prefixed mistral models" -#[test] -fn is_mistral_model_region_prefixed() { - assert!(model_normalizes_tool_id( - "us.mistral.pixtral-large-2502-v1:0" - )); - assert!(model_normalizes_tool_id( - "eu.mistral.mistral-large-2407-v1:0" - )); -} + /// upstream: @ai-sdk/amazon-bedrock src/normalize-tool-call-id.test.ts › isMistralModel › should return true for region-prefixed mistral models + #[test] + fn is_mistral_model_region_prefixed() { + assert!(model_normalizes_tool_id( + "us.mistral.pixtral-large-2502-v1:0" + )); + assert!(model_normalizes_tool_id( + "eu.mistral.mistral-large-2407-v1:0" + )); + } -/// TS: "should return false for non-mistral models" -#[test] -fn is_mistral_model_false_for_non_mistral() { - for id in [ - "anthropic.claude-3-5-sonnet-20241022-v2:0", - "amazon.nova-pro-v1:0", - "openai.gpt-4o", - "meta.llama3-70b-instruct-v1:0", - ] { - assert!(!model_normalizes_tool_id(id)); + /// upstream: @ai-sdk/amazon-bedrock src/normalize-tool-call-id.test.ts › isMistralModel › should return false for non-mistral models + #[test] + fn is_mistral_model_false_for_non_mistral() { + for id in [ + "anthropic.claude-3-5-sonnet-20241022-v2:0", + "amazon.nova-pro-v1:0", + "openai.gpt-4o", + "meta.llama3-70b-instruct-v1:0", + ] { + assert!(!model_normalizes_tool_id(id)); + } } } -/// TS: "should return the original ID when not a Mistral model" -#[test] -fn normalize_tool_call_id_passthrough_when_not_mistral() { - let id = "tooluse_bpe71yCfRu2b5i-nKGDr5g"; - assert_eq!(normalize_tool_call_id(id, false), id); -} +mod normalize_tool_call_id { + use super::*; -/// TS: "should hash incompatible IDs deterministically for Mistral models" -#[test] -fn normalize_tool_call_id_hashes_deterministically() { - let id = "tooluse_bpe71yCfRu2b5i-nKGDr5g"; - assert_eq!(normalize_tool_call_id(id, true), "8eHypBDcw"); - assert_eq!(normalize_tool_call_id(id, true), "8eHypBDcw"); -} + /// upstream: @ai-sdk/amazon-bedrock src/normalize-tool-call-id.test.ts › normalizeToolCallId › should return the original ID when not a Mistral model + #[test] + fn normalize_tool_call_id_passthrough_when_not_mistral() { + let id = "tooluse_bpe71yCfRu2b5i-nKGDr5g"; + assert_eq!(normalize_tool_call_id(id, false), id); + } -/// TS: "should produce 9 alphanumeric characters for incompatible IDs" -#[test] -fn normalize_tool_call_id_special_chars() { - for (id, expected) in [ - ("tool-use_123ABC456", "hvVDqPNyj"), - ("___abc123DEF___", "TnzPqldGU"), - ("abc", "GRuIyUwcV"), - ("12345", "PceAgWDYe"), - ("___---___", "5C589HVqG"), - ] { - let normalized = normalize_tool_call_id(id, true); - assert_eq!(normalized, expected); - assert_eq!(normalized.len(), 9); - assert!(normalized.bytes().all(|c| c.is_ascii_alphanumeric())); + /// upstream: @ai-sdk/amazon-bedrock src/normalize-tool-call-id.test.ts › normalizeToolCallId › should hash incompatible IDs deterministically for Mistral models + #[test] + fn normalize_tool_call_id_hashes_deterministically() { + let id = "tooluse_bpe71yCfRu2b5i-nKGDr5g"; + assert_eq!(normalize_tool_call_id(id, true), "8eHypBDcw"); + assert_eq!(normalize_tool_call_id(id, true), "8eHypBDcw"); } -} -/// TS: "should preserve IDs that are already valid Mistral tool call IDs" -#[test] -fn normalize_tool_call_id_already_alphanumeric() { - assert_eq!(normalize_tool_call_id("abcdefghi", true), "abcdefghi"); - assert_eq!(normalize_tool_call_id("abc123XYZ", true), "abc123XYZ"); + /// upstream: @ai-sdk/amazon-bedrock src/normalize-tool-call-id.test.ts › normalizeToolCallId › should produce 9 alphanumeric characters for incompatible IDs + #[test] + fn normalize_tool_call_id_special_chars() { + for (id, expected) in [ + ("tool-use_123ABC456", "hvVDqPNyj"), + ("___abc123DEF___", "TnzPqldGU"), + ("abc", "GRuIyUwcV"), + ("12345", "PceAgWDYe"), + ("___---___", "5C589HVqG"), + ] { + let normalized = normalize_tool_call_id(id, true); + assert_eq!(normalized, expected); + assert_eq!(normalized.len(), 9); + assert!(normalized.bytes().all(|c| c.is_ascii_alphanumeric())); + } + } + + /// upstream: @ai-sdk/amazon-bedrock src/normalize-tool-call-id.test.ts › normalizeToolCallId › should preserve IDs that are already valid Mistral tool call IDs + #[test] + fn normalize_tool_call_id_already_alphanumeric() { + assert_eq!(normalize_tool_call_id("abcdefghi", true), "abcdefghi"); + assert_eq!(normalize_tool_call_id("abc123XYZ", true), "abc123XYZ"); + } } use aimux_providers::bedrock::convert::{BedrockUsage, convert_usage}; @@ -261,28 +269,32 @@ fn bedrock_usage_with_total( } } -/// TS: "should include totalTokens in raw when provided" -#[test] -fn usage_raw_includes_total_tokens() { - let raw = bedrock_usage_with_total(100, 50, Some(150), None, None); - let u = convert_usage(Some(&raw)); - assert_eq!( - u.raw, - Some( - json!({"inputTokens":100,"outputTokens":50,"totalTokens":150}) - .as_object() - .unwrap() - .clone() - ) - ); -} +mod convert_amazon_bedrock_usage { + use super::*; + + /// upstream: @ai-sdk/amazon-bedrock src/convert-amazon-bedrock-usage.test.ts › convertAmazonBedrockUsage › should include totalTokens in raw when provided + #[test] + fn usage_raw_includes_total_tokens() { + let raw = bedrock_usage_with_total(100, 50, Some(150), None, None); + let u = convert_usage(Some(&raw)); + assert_eq!( + u.raw, + Some( + json!({"inputTokens":100,"outputTokens":50,"totalTokens":150}) + .as_object() + .unwrap() + .clone() + ) + ); + } -/// TS: "should preserve raw usage data" -#[test] -fn usage_raw_preserved() { - let raw = bedrock_usage_with_total(100, 50, Some(150), Some(80), Some(60)); - let u = convert_usage(Some(&raw)); - assert_eq!(u.raw, Some(json!({"inputTokens":100,"outputTokens":50,"totalTokens":150,"cacheReadInputTokens":80,"cacheWriteInputTokens":60}).as_object().unwrap().clone())); + /// upstream: @ai-sdk/amazon-bedrock src/convert-amazon-bedrock-usage.test.ts › convertAmazonBedrockUsage › should preserve raw usage data + #[test] + fn usage_raw_preserved() { + let raw = bedrock_usage_with_total(100, 50, Some(150), Some(80), Some(60)); + let u = convert_usage(Some(&raw)); + assert_eq!(u.raw, Some(json!({"inputTokens":100,"outputTokens":50,"totalTokens":150,"cacheReadInputTokens":80,"cacheWriteInputTokens":60}).as_object().unwrap().clone())); + } } /// TS: "should handle null cache tokens" — `raw` is the usage object as @@ -306,8 +318,8 @@ fn usage_raw_keeps_null_and_unknown_keys() { // amazon-bedrock-provider.test.ts — provider configuration / auth // ════════════════════════════════════════════════════════════════════════════ -/// TS: "should create a provider instance with default options" — the mocked -/// upstream region loader returns `us-east-1`; pass it explicitly here. +/// partial: `should create a provider instance with default options`; full port in src/bedrock/tests/provider.rs. +/// The mocked upstream region loader returns `us-east-1`; pass it explicitly here. #[tokio::test] #[serial] async fn provider_default_region_and_base_url() { @@ -334,8 +346,8 @@ async fn provider_default_region_and_base_url() { assert!(seen.headers["authorization"].starts_with("AWS4-HMAC-SHA256")); } -/// TS: "should create a provider instance with custom options" — custom region -/// flows into the base URL. +/// partial: `should create a provider instance with custom options`; full port in src/bedrock/tests/provider.rs. +/// Custom region flows into the base URL. #[tokio::test] async fn provider_custom_region_base_url() { let fetch = MockFetch::new(vec![Canned::json(&ok_converse_body())]); @@ -360,7 +372,8 @@ async fn provider_custom_region_base_url() { ); } -/// TS: "should create a provider instance with custom options" — baseURL override. +/// partial: `should create a provider instance with custom options`; full port in src/bedrock/tests/provider.rs. +/// baseURL override. #[tokio::test] async fn provider_with_base_url_override() { let fetch = MockFetch::new(vec![Canned::json(&ok_converse_body())]); @@ -381,7 +394,8 @@ async fn provider_with_base_url_override() { assert!(seen.url.starts_with("https://custom.url/model/")); } -/// TS: "should use API key when provided in options" — bearer-token auth path. +/// partial: `should use API key when provided in options`; full port in src/bedrock/tests/provider.rs. +/// Bearer-token auth path. #[tokio::test] async fn provider_bearer_token_auth() { let fetch = MockFetch::new(vec![Canned::json(&ok_converse_body())]); @@ -401,8 +415,8 @@ async fn provider_bearer_token_auth() { assert_eq!(seen.headers["authorization"], "Bearer test-api-key"); } -/// TS: "should use API key from environment variable" — `AWS_BEARER_TOKEN_BEDROCK` -/// takes precedence over SigV4 environment credentials. +/// partial: `should use API key from environment variable`; full port in src/bedrock/tests/provider.rs. +/// `AWS_BEARER_TOKEN_BEDROCK` takes precedence over SigV4 environment credentials. #[tokio::test] #[serial] async fn provider_from_env_bearer_token_precedence() { @@ -431,8 +445,8 @@ async fn provider_from_env_bearer_token_precedence() { clear_bedrock_env(); } -/// TS: "should fall back to SigV4 when no API key provided" — the factory uses -/// SigV4 when `AWS_BEARER_TOKEN_BEDROCK` is absent. +/// partial: `should fall back to SigV4 when no API key provided`; full port in src/bedrock/tests/provider.rs. +/// The factory uses SigV4 when `AWS_BEARER_TOKEN_BEDROCK` is absent. #[tokio::test] #[serial] async fn provider_from_env_sigv4_fallback() { @@ -461,8 +475,8 @@ async fn provider_from_env_sigv4_fallback() { clear_bedrock_env(); } -/// TS: "should maintain backward compatibility with existing SigV4 -/// authentication" — `AWS_SESSION_TOKEN` is loaded from the environment. +/// partial: `should maintain backward compatibility with existing SigV4 authentication`; full port in src/bedrock/tests/provider.rs. +/// `AWS_SESSION_TOKEN` is loaded from the environment. #[tokio::test] #[serial] async fn provider_from_env_session_token() { @@ -502,8 +516,8 @@ fn provider_name_and_language_model() { assert_eq!(model.provider(), "amazon-bedrock"); } -/// TS: "should prioritize options.apiKey over environment variable" — explicit -/// bearer token construction does not consult env vars. +/// partial: `should prioritize options.apiKey over environment variable`; full port in src/bedrock/tests/provider.rs. +/// Explicit bearer token construction does not consult env vars. #[tokio::test] #[serial] async fn provider_explicit_bearer_over_env() { @@ -544,7 +558,7 @@ fn assistant_msg(content: Vec) -> LanguageModelMessage { } } -/// TS: "should normalize tool call IDs in tool results when isMistral is true" +/// partial: `should normalize tool call IDs in tool results when isMistral is true`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn convert_mistral_normalize_tool_result_id() { let original_id = "tooluse_bpe71yCfRu2b5i-nKGDr5g"; @@ -563,7 +577,7 @@ fn convert_mistral_normalize_tool_result_id() { ); } -/// TS: "should normalize tool call IDs in tool calls when isMistral is true" +/// partial: `should normalize tool call IDs in tool calls when isMistral is true`; full port in src/bedrock/tests/convert_to_chat_messages.rs. #[test] fn convert_mistral_normalize_tool_call_id() { let original_id = "tooluse_xyz123ABC456-def"; @@ -589,7 +603,7 @@ fn convert_mistral_normalize_tool_call_id() { // ── supportedUrls ──────────────────────────────────────────────────────────── -/// TS: "should support S3 URLs for image and video parts". +/// partial: `should support S3 URLs for image and video parts`; full port in src/bedrock/tests/chat_language_model.rs. #[tokio::test] async fn model_supported_urls_s3() { let model = create_amazon_bedrock(AmazonBedrockProviderSettings::default()) @@ -605,7 +619,7 @@ async fn model_supported_urls_s3() { // ── ARN model IDs containing a slash ───────────────────────────────────────── -/// TS: "should generate text through the encoded Converse route". +/// partial: `should generate text through the encoded Converse route`; full port in src/bedrock/tests/chat_language_model.rs. /// /// ARN inference-profile IDs contain a `/` which must be percent-encoded in the /// URL path. The upstream uses `encodeURIComponent(modelId)`. @@ -643,7 +657,7 @@ async fn arn_model_id_encoded_generate_route() { ); } -/// TS: "should stream text through the encoded Converse route". +/// partial: `should stream text through the encoded Converse route`; full port in src/bedrock/tests/chat_language_model.rs. #[tokio::test] async fn arn_model_id_encoded_stream_route() { let arn = "arn:aws:bedrock:eu-west-1:474668406012:inference-profile/eu.amazon.nova-lite-v1:0"; @@ -706,7 +720,7 @@ async fn arn_model_id_encoded_stream_route() { // ── temperature clamping ───────────────────────────────────────────────────── -/// TS: "should clamp temperature above 1 to 1 and add warning". +/// partial: `should clamp temperature above 1 to 1 and add warning`; full port in src/bedrock/tests/chat_language_model.rs. #[tokio::test] async fn temperature_clamped_above_1() { let server = mock_converse_json(200, ok_converse_body()); @@ -732,7 +746,7 @@ async fn temperature_clamped_above_1() { ); } -/// TS: "should clamp temperature below 0 to 0 and add warning". +/// partial: `should clamp temperature below 0 to 0 and add warning`; full port in src/bedrock/tests/chat_language_model.rs. #[tokio::test] async fn temperature_clamped_below_0() { let server = mock_converse_json(200, ok_converse_body()); @@ -757,7 +771,7 @@ async fn temperature_clamped_below_0() { ); } -/// TS: "should not clamp valid temperature between 0 and 1". +/// partial: `should not clamp valid temperature between 0 and 1`; full port in src/bedrock/tests/chat_language_model.rs. #[tokio::test] async fn temperature_not_clamped_in_range() { let server = mock_converse_json(200, ok_converse_body()); @@ -830,8 +844,8 @@ fn nova_high_reasoning_max_output_tokens_warning_names_the_model() { // ── guardrails ─────────────────────────────────────────────────────────────── -/// TS: "should support guardrails" — `providerOptions.bedrock.guardrailConfig` -/// is forwarded as a top-level `guardrailConfig` in the request body. +/// partial: `should support guardrails`; full port in src/bedrock/tests/chat_language_model.rs. +/// `providerOptions.bedrock.guardrailConfig` is forwarded as a top-level `guardrailConfig` in the request body. #[tokio::test] async fn guardrail_config_in_request_body() { let server = mock_converse_json(200, ok_converse_body()); @@ -865,8 +879,8 @@ async fn guardrail_config_in_request_body() { // ── trace in providerMetadata ──────────────────────────────────────────────── -/// TS: "should include trace information in providerMetadata" — the response -/// `trace` field is surfaced as `providerMetadata.bedrock.trace`. +/// partial: `should include trace information in providerMetadata`; full port in src/bedrock/tests/chat_language_model.rs. +/// The response `trace` field is surfaced as `providerMetadata.bedrock.trace`. #[tokio::test] async fn trace_in_provider_metadata() { let trace = json!({ @@ -909,9 +923,8 @@ async fn trace_in_provider_metadata() { // ── stop_sequence in providerMetadata ──────────────────────────────────────── -/// TS: "should include stop_sequence in provider metadata" — when the response -/// carries `additionalModelResponseFields.delta.stop_sequence`, it is surfaced -/// as `providerMetadata.bedrock.stopSequence` (and `amazonBedrock.stopSequence`). +/// partial: `should include stop_sequence in provider metadata`; full port in src/bedrock/tests/chat_language_model.rs. +/// When the response carries `additionalModelResponseFields.delta.stop_sequence`, it is surfaced as `providerMetadata.bedrock.stopSequence` (and `amazonBedrock.stopSequence`). #[tokio::test] async fn stop_sequence_in_provider_metadata() { let server = mock_converse_json( @@ -940,7 +953,8 @@ async fn stop_sequence_in_provider_metadata() { // ── tool calls with empty input ────────────────────────────────────────────── -/// TS: "should support tool calls with empty input (no arguments)" (stream). +/// partial: `should support tool calls with empty input (no arguments)`; full port in src/bedrock/tests/chat_language_model.rs. +/// (stream). /// /// When a `toolUse` block carries no input deltas, the stream event loop /// accumulates an empty string and yields `input: "{}"`. @@ -995,8 +1009,7 @@ async fn stream_tool_call_empty_input() { // ── omit toolConfig ────────────────────────────────────────────────────────── -/// TS: "should omit toolConfig and filter tool content when conversation has -/// tool calls but no active tools". +/// partial: `should omit toolConfig and filter tool content when conversation has tool calls but no active tools`; full port in src/bedrock/tests/chat_language_model.rs. /// #[tokio::test] async fn omit_tool_config_when_no_active_tools() { @@ -1042,8 +1055,8 @@ async fn omit_tool_config_when_no_active_tools() { ); } -/// TS: (same test) the assistant's `toolUse` block should be filtered out of -/// the messages when no tools are active. +/// partial: `should omit toolConfig and filter tool content when conversation has tool calls but no active tools`; full port in src/bedrock/tests/chat_language_model.rs. +/// (same test) the assistant's `toolUse` block should be filtered out of the messages when no tools are active. #[tokio::test] async fn filter_tool_content_when_no_active_tools() { let server = mock_converse_json(200, ok_converse_body()); @@ -1090,7 +1103,8 @@ async fn filter_tool_content_when_no_active_tools() { // ── doGenerate: tool call with empty input (non-streaming) ─────────────────── -/// TS: "should support tool calls with empty input (no arguments)" (generate). +/// partial: `should support tool calls with empty input (no arguments)`; full port in src/bedrock/tests/chat_language_model.rs. +/// (generate). #[tokio::test] async fn generate_tool_call_empty_input() { let server = mock_converse_json( @@ -1163,8 +1177,8 @@ async fn generate_tool_call_empty_tool_use_id() { // ── doGenerate: basic text + finish reason (sanity, already covered but // exercised here against the remaining-test helper set) ──────────────────── -/// TS: "should pass the model and the messages" — the request body carries the -/// model id implicitly via the URL path and the messages array. +/// partial: `should pass the model and the messages`; full port in src/bedrock/tests/chat_language_model.rs. +/// The request body carries the model id implicitly via the URL path and the messages array. #[tokio::test] async fn request_body_messages_shape() { let server = mock_converse_json(200, ok_converse_body()); @@ -1191,7 +1205,8 @@ async fn request_body_messages_shape() { assert_eq!(body["messages"][0]["content"][0]["text"], "Hello"); } -/// TS: "should extract finish reason" — `guardrail_intervened` → ContentFilter. +/// partial: `should extract finish reason`; full port in src/bedrock/tests/chat_language_model.rs. +/// `guardrail_intervened` → ContentFilter. #[tokio::test] async fn finish_reason_guardrail_intervened() { let server = mock_converse_json( @@ -1219,7 +1234,8 @@ async fn finish_reason_guardrail_intervened() { ); } -/// TS: "should handle throttlingException error" (stream). +/// partial: `should handle throttlingException error`; full port in src/bedrock/tests/chat_language_model.rs. +/// (stream). #[tokio::test] async fn stream_throttling_error() { let body_bytes = aimux_providers::bedrock::event_stream::encode_messages(&[( @@ -1247,7 +1263,8 @@ async fn stream_throttling_error() { assert!(parts.iter().any(|part| matches!(part, StreamPart::Finish { finish_reason, .. } if finish_reason.unified == FinishReasonUnified::Error))); } -/// TS: "should handle validationException error" (stream). +/// partial: `should handle validationException error`; full port in src/bedrock/tests/chat_language_model.rs. +/// (stream). #[tokio::test] async fn stream_validation_error() { let body_bytes = aimux_providers::bedrock::event_stream::encode_messages(&[( @@ -1279,8 +1296,8 @@ async fn stream_validation_error() { // ── doStream: request body parity ──────────────────────────────────────────── -/// TS: "should return the request body" (stream) — the stream result carries -/// the request body for debugging. +/// partial: `should return the request body`; full port in src/bedrock/tests/chat_language_model.rs. +/// (stream) — the stream result carries the request body for debugging. #[tokio::test] async fn stream_request_body_available() { let events: Vec<(&str, &str, &str)> = vec![ diff --git a/aimux-providers/tests/bedrock_upstream_e2e_test.rs b/aimux-providers/tests/bedrock_upstream_e2e_test.rs index 6e7cdf5af..c92e75f2a 100644 --- a/aimux-providers/tests/bedrock_upstream_e2e_test.rs +++ b/aimux-providers/tests/bedrock_upstream_e2e_test.rs @@ -56,7 +56,7 @@ fn stream_response(events: Value) -> Canned { } } -/// TS: "should extract reasoning text with signature" (amazon-bedrock/src/amazon-bedrock-chat-language-model.test.ts) +/// partial: `should extract reasoning text with signature`; full port in src/bedrock/tests/chat_language_model.rs. #[tokio::test] async fn generate_reasoning_signature_and_usage() { let response = json!({"output":{"message":{"role":"assistant","content":[ @@ -102,7 +102,7 @@ async fn generate_reasoning_signature_and_usage() { assert_eq!(result.usage.raw, response["usage"].as_object().cloned()); } -/// TS: "should stream text deltas with metadata and usage" (amazon-bedrock/src/amazon-bedrock-chat-language-model.test.ts) +/// partial: `should stream text deltas with metadata and usage`; full port in src/bedrock/tests/chat_language_model.rs. #[tokio::test] async fn stream_text_and_usage() { let fetch = MockFetch::new(vec![stream_response(json!([ @@ -165,7 +165,7 @@ async fn stream_text_and_usage() { ); } -/// TS: "should stream tool deltas" (amazon-bedrock/src/amazon-bedrock-chat-language-model.test.ts) +/// partial: `should stream tool deltas`; full port in src/bedrock/tests/chat_language_model.rs. #[tokio::test] async fn stream_partial_tool_arguments() { let fetch = MockFetch::new(vec![stream_response(json!([ @@ -217,7 +217,7 @@ async fn stream_partial_tool_arguments() { ); } -/// TS: "prefixes the provider message when the error type is present" (amazon-bedrock/src/amazon-bedrock-error.test.ts) +/// partial: `prefixes the provider message when the error type is present`; full port in src/bedrock/tests/error.rs. #[tokio::test] async fn generate_error_response() { let response = json!({"type":"ValidationException","message":"boom"}); @@ -242,7 +242,7 @@ async fn generate_error_response() { assert_eq!(error.data, Some(response)); } -/// TS: "should handle single input value and return embeddings" (amazon-bedrock/src/amazon-bedrock-embedding-model.test.ts) +/// partial: `should handle single input value and return embeddings`; full port in src/bedrock/tests/embedding_model.rs. #[tokio::test] async fn embed_single_value_and_usage() { let response = json!({"embedding":[-0.09,0.05,-0.02,0.01,0.04],"inputTextTokenCount":8}); @@ -268,7 +268,7 @@ async fn embed_single_value_and_usage() { assert!(result.response.is_none()); } -/// TS: "should pass the model and the settings" (amazon-bedrock/src/amazon-bedrock-image-model.test.ts) +/// partial: `should pass the model and the settings`; full port in src/bedrock/tests/image_model.rs. #[tokio::test] async fn generate_images_with_settings() { let fetch = MockFetch::new(vec![Canned::json( @@ -300,7 +300,7 @@ async fn generate_images_with_settings() { assert_eq!(images, vec!["base64-image-1", "base64-image-2"]); } -/// TS: "should throw error when request is moderated" (amazon-bedrock/src/amazon-bedrock-image-model.test.ts) +/// partial: `should throw error when request is moderated`; full port in src/bedrock/tests/image_model.rs. #[tokio::test] async fn image_moderation_error_response() { let fetch = MockFetch::new(vec![Canned::json( @@ -324,7 +324,7 @@ async fn image_moderation_error_response() { ); } -/// TS: "should send request with stringified json documents" (amazon-bedrock/src/reranking/amazon-bedrock-reranking-model.test.ts) +/// partial: `should send request with stringified json documents`; full port in src/bedrock/tests/reranking/reranking_model.rs. #[tokio::test] async fn rerank_json_documents_and_parse_ranking() { let response = json!({"results":[{"index":0,"relevanceScore":0.5110583305358887},{"index":5,"relevanceScore":0.30241215229034424}]}); @@ -359,7 +359,7 @@ async fn rerank_json_documents_and_parse_ranking() { assert_eq!(result.response.unwrap().body.unwrap(), response); } -/// TS: "should handle JSON response format in streaming" (amazon-bedrock-chat-language-model.test.ts). +/// partial: `should handle JSON response format in streaming`; full port in src/bedrock/tests/chat_language_model.rs. #[tokio::test] async fn stream_json_response_tool_as_text() { let fetch = MockFetch::new(vec![stream_response(json!([ @@ -472,7 +472,7 @@ async fn provider_tool_args_are_checked_before_the_request() { assert_eq!(result.warnings.len(), 1, "{:?}", result.warnings); } -/// TS: "should handle Anthropic provider-defined tools" (amazon-bedrock-chat-language-model.test.ts). +/// partial: `should handle Anthropic provider-defined tools`; full port in src/bedrock/tests/chat_language_model.rs. #[tokio::test] async fn generate_provider_defined_tool_with_input_schema() { let fetch = MockFetch::new(vec![Canned::json(&json!({