diff --git a/aimux-providers/layout.toml b/aimux-providers/layout.toml index 6329564a7..6db6d45f2 100644 --- a/aimux-providers/layout.toml +++ b/aimux-providers/layout.toml @@ -812,3 +812,116 @@ issue = 275 upstream = "src/version.ts" status = "not-applicable" reason = "the crate version is used instead (RFC-0046 §4.2)" + +[cohere] +package = "@ai-sdk/cohere" + +[[cohere.files]] +upstream = "src/cohere-chat-language-model-options.ts" +rust = "cohere/convert.rs" +status = "mapped" + +[cohere.files.tests] + +[[cohere.files]] +upstream = "src/cohere-chat-language-model.ts" +rust = "cohere/model.rs" +status = "mapped" + +[cohere.files.tests] + +[[cohere.files]] +upstream = "src/cohere-chat-prompt.ts" +rust = "cohere/convert.rs" +status = "mapped" + +[cohere.files.tests] + +[[cohere.files]] +upstream = "src/cohere-embedding-model-options.ts" +rust = "cohere/embedding.rs" +status = "mapped" + +[cohere.files.tests] + +[[cohere.files]] +upstream = "src/cohere-embedding-model.ts" +rust = "cohere/embedding.rs" +status = "mapped" + +[cohere.files.tests] + +[[cohere.files]] +upstream = "src/cohere-error.ts" +rust = "cohere/mod.rs" +status = "mapped" + +[cohere.files.tests] + +[[cohere.files]] +upstream = "src/cohere-prepare-tools.ts" +rust = "cohere/convert.rs" +status = "mapped" + +[cohere.files.tests] + +[[cohere.files]] +upstream = "src/cohere-provider.ts" +rust = "cohere/mod.rs" +status = "mapped" + +[cohere.files.tests] + +[[cohere.files]] +upstream = "src/convert-cohere-usage.ts" +rust = "cohere/model.rs" +status = "mapped" + +[cohere.files.tests] + +[[cohere.files]] +upstream = "src/convert-to-cohere-chat-prompt.ts" +rust = "cohere/convert.rs" +status = "mapped" + +[cohere.files.tests] + +[[cohere.files]] +upstream = "src/index.ts" +rust = "cohere/mod.rs" +status = "mapped" + +[cohere.files.tests] + +[[cohere.files]] +upstream = "src/map-cohere-finish-reason.ts" +rust = "cohere/convert.rs" +status = "mapped" + +[cohere.files.tests] + +[[cohere.files]] +upstream = "src/reranking/cohere-reranking-api.ts" +rust = "cohere/reranking.rs" +status = "mapped" + +[cohere.files.tests] + +[[cohere.files]] +upstream = "src/reranking/cohere-reranking-model-options.ts" +rust = "cohere/reranking.rs" +status = "mapped" + +[cohere.files.tests] + +[[cohere.files]] +upstream = "src/reranking/cohere-reranking-model.ts" +rust = "cohere/reranking.rs" +status = "mapped" + +[cohere.files.tests] + +[[cohere.files]] +upstream = "src/version.ts" +status = "not-applicable" +reason = "the crate version is used instead (RFC-0046 §4.2)" diff --git a/aimux-providers/src/cohere/mod.rs b/aimux-providers/src/cohere/mod.rs index b9429e8aa..715dcfe8c 100644 --- a/aimux-providers/src/cohere/mod.rs +++ b/aimux-providers/src/cohere/mod.rs @@ -268,3 +268,6 @@ impl ProviderDiscovery for CohereProvider { }) } } + +#[cfg(test)] +mod tests; diff --git a/aimux-providers/src/cohere/tests/chat_language_model.rs b/aimux-providers/src/cohere/tests/chat_language_model.rs new file mode 100644 index 000000000..1c04b8abd --- /dev/null +++ b/aimux-providers/src/cohere/tests/chat_language_model.rs @@ -0,0 +1,1339 @@ +//! Ports of `src/cohere-chat-language-model.test.ts` from `@ai-sdk/cohere@4.0.52`. +//! +//! Harness differences from the TypeScript (expectations are unchanged): +//! - `toMatchSnapshot` values are the `.snap` entries in `snapshots/`, with +//! `undefined` properties absent (as `JSON.stringify` drops them). +//! - Inline snapshots and `toStrictEqual` that list `undefined` properties +//! are written without those keys, since the serialized results omit them. +//! - The TypeScript test server adds `content-length` to JSON responses and +//! `cache-control` / `connection` to event streams; the mock does the same. +//! The `content-length` is the byte length of the compact JSON body. +//! - The mock records `user-agent`, which the TypeScript `requestHeaders` do +//! not list; the header comparisons drop it. +//! +//! Cases that are not here (`should stream text deltas` and +//! `should handle unparsable stream parts`) are ported in +//! `tests/cohere_upstream_e2e_test.rs`. `should extract text response` and +//! `should extract tool calls` compare the whole snapshot here; the older +//! cases in that file check part of it. + +use std::collections::{BTreeMap, HashMap}; +use std::sync::Arc; + +use crate::cohere::{CohereProvider, CohereProviderSettings, create_cohere}; +use aimux_core::language_model::LanguageModel; +use aimux_core::language_model_message::{FilePart, LanguageModelMessage, TextPart, UserPart}; +use aimux_core::options::{CallOptions, ResponseFormat}; +use aimux_core::result::StreamResult; +use aimux_core::shared::{FileBytes, FileData, SharedProviderOptions}; +use aimux_core::tool::{FunctionTool, LanguageModelV4ToolChoice, Tool}; +use aimux_core::types::ReasoningEffort; +use aimux_provider_utils::HeaderMapOpt; +use futures::StreamExt; +use serde_json::{Value, json}; + +use crate::test_support::mock_fetch::{Canned, MockFetch}; + +const MODEL: &str = "command-r-plus"; +const TEXT_FIXTURE: &str = include_str!("fixtures/cohere-text.json"); +const MAX_TOKENS_FIXTURE: &str = include_str!("fixtures/cohere-max-tokens.json"); +const NULL_ARGS_FIXTURE: &str = include_str!("fixtures/cohere-null-args.json"); +const REASONING_FIXTURE: &str = include_str!("fixtures/cohere-reasoning.json"); +const CITATIONS_FIXTURE: &str = include_str!("fixtures/cohere-citations.json"); +const TOOL_CALL_FIXTURE: &str = include_str!("fixtures/cohere-tool-call.json"); +const TEXT_CHUNKS: &str = include_str!("fixtures/cohere-text.chunks.txt"); +const REASONING_CHUNKS: &str = include_str!("fixtures/cohere-reasoning.chunks.txt"); +const TOOL_CALL_CHUNKS: &str = include_str!("fixtures/cohere-tool-call.chunks.txt"); +const EMPTY_TOOL_CALL_CHUNKS: &str = include_str!("fixtures/cohere-empty-tool-call.chunks.txt"); + +const SNAPSHOT_DO_GENERATE_CITATIONS: &str = include_str!("snapshots/do_generate_citations.json"); +const SNAPSHOT_DO_GENERATE_NULL_TOOL_CALL_ARGS: &str = + include_str!("snapshots/do_generate_null_tool_call_args.json"); +const SNAPSHOT_DO_GENERATE_REASONING: &str = include_str!("snapshots/do_generate_reasoning.json"); +const SNAPSHOT_DO_GENERATE_TEXT: &str = include_str!("snapshots/do_generate_text.json"); +const SNAPSHOT_DO_GENERATE_TOOL_CALL: &str = include_str!("snapshots/do_generate_tool_call.json"); +const SNAPSHOT_DO_STREAM_EMPTY_TOOL_CALL: &str = + include_str!("snapshots/do_stream_empty_tool_call.json"); +const SNAPSHOT_DO_STREAM_REASONING: &str = include_str!("snapshots/do_stream_reasoning.json"); +const SNAPSHOT_DO_STREAM_RAW_CHUNKS: &str = include_str!("snapshots/do_stream_raw_chunks.json"); +const SNAPSHOT_DO_STREAM_TOOL_CALL: &str = include_str!("snapshots/do_stream_tool_call.json"); + +// ── helpers ───────────────────────────────────────────────────────────────── + +fn json_value(source: &str) -> Value { + serde_json::from_str(source).expect("fixture or snapshot is JSON") +} + +/// `createCohere({ apiKey, headers, generateId })` with a mock transport. +fn provider(fetch: &Arc) -> CohereProvider { + provider_with(fetch, None, None) +} + +fn provider_with( + fetch: &Arc, + headers: Option, + generate_id: Option String + Send + Sync>>, +) -> CohereProvider { + create_cohere(CohereProviderSettings { + api_key: Some("test-api-key".into()), + headers, + generate_id, + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() +} + +/// `TEST_PROMPT`: a system message and the user message "Hello". +fn options() -> CallOptions { + CallOptions::new(vec![ + LanguageModelMessage::System { + content: "you are a friendly bot!".into(), + provider_options: None, + }, + LanguageModelMessage::user_text("Hello"), + ]) +} + +/// A prompt with one user message made of the given parts. +fn user_options(parts: Vec) -> CallOptions { + CallOptions::new(vec![LanguageModelMessage::User { + content: parts, + provider_options: None, + }]) +} + +fn text_part(text: &str) -> UserPart { + UserPart::Text(TextPart { + text: text.into(), + provider_options: None, + }) +} + +/// A file part with inline bytes, as `Buffer.from(...)` or a string gives. +fn file_part(bytes: &[u8], media_type: &str, filename: &str) -> UserPart { + UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Binary(bytes.to_vec()), + }, + media_type: media_type.into(), + filename: Some(filename.into()), + provider_options: None, + }) +} + +fn test_tool() -> Tool { + Tool::Function(FunctionTool::new( + "test-tool", + json!({ + "type": "object", + "properties": {"value": {"type": "string"}}, + "required": ["value"], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#" + }), + )) +} + +fn provider_options(value: Value) -> SharedProviderOptions { + serde_json::from_value(value).expect("provider options are an object") +} + +fn map(pairs: &[(&str, &str)]) -> BTreeMap { + pairs + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() +} + +/// The request headers the TypeScript test checks with `toStrictEqual`. +/// The user agent is recorded by the mock and not compared here. +fn assert_sent_headers(fetch: &MockFetch, expected: &[(&str, &str)]) { + let mut headers = fetch.seen()[0].headers.clone(); + headers.remove("user-agent").expect("user-agent is sent"); + assert_eq!(headers, map(expected)); +} + +/// A non-streaming response with the headers the TypeScript test server sends. +fn json_response(body: &Value, extra_headers: &[(&str, &str)]) -> Canned { + let bytes = serde_json::to_vec(body).expect("serializable"); + let mut headers = vec![ + ("content-type".to_string(), "application/json".to_string()), + ("content-length".to_string(), bytes.len().to_string()), + ]; + headers.extend( + extra_headers + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())), + ); + Canned { + status: 200, + headers, + body: bytes, + } +} + +/// A single-response mock for a JSON fixture. +fn json_fetch(source: &str, extra_headers: &[(&str, &str)]) -> Arc { + MockFetch::new(vec![json_response(&json_value(source), extra_headers)]) +} + +/// One `event:` / `data:` frame, as the TypeScript helpers write it. +fn frame(event: &str, data: &str) -> String { + format!("event: {event}\ndata: {data}\n\n") +} + +/// `prepareChunksFixtureResponse`: every non-blank line of a `.chunks.txt` +/// fixture is one frame, named by its `type`. +fn chunks_frames(source: &str) -> Vec { + source + .lines() + .filter(|line| !line.trim().is_empty()) + .map(|line| { + let parsed: Value = serde_json::from_str(line).expect("fixture line is JSON"); + frame(parsed["type"].as_str().unwrap(), line) + }) + .collect() +} + +/// `prepareChunkLinesResponse`: each chunk is serialized into one frame. +fn chunk_frames(chunks: &[Value]) -> Vec { + chunks + .iter() + .map(|chunk| frame(chunk["type"].as_str().unwrap(), &chunk.to_string())) + .collect() +} + +/// An event-stream response with the headers the TypeScript test server sends. +fn sse_response(frames: &[String], extra_headers: &[(&str, &str)]) -> Canned { + let mut headers = vec![ + ("content-type".to_string(), "text/event-stream".to_string()), + ("cache-control".to_string(), "no-cache".to_string()), + ("connection".to_string(), "keep-alive".to_string()), + ]; + headers.extend( + extra_headers + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())), + ); + Canned { + status: 200, + headers, + body: frames.concat().into_bytes(), + } +} + +/// A single event-stream mock. +fn stream_fetch(frames: &[String], extra_headers: &[(&str, &str)]) -> Arc { + MockFetch::new(vec![sse_response(frames, extra_headers)]) +} + +/// `message-end` chunk with the given usage. +fn message_end(usage: Value) -> Value { + json!({ + "type": "message-end", + "delta": {"finish_reason": "COMPLETE", "usage": usage} + }) +} + +/// All parts of a stream, each serialized as the AI SDK writes it. +async fn stream_parts(result: StreamResult) -> Vec { + result + .stream + .map(|part| serde_json::to_value(part.expect("stream part")).unwrap()) + .collect() + .await +} + +/// `toMatchObject`: every property of `expected` is in `actual`, recursively; +/// arrays have the same length and match element by element. +fn assert_match_object(actual: &Value, expected: &Value) { + match (actual, expected) { + (Value::Object(actual_map), Value::Object(expected_map)) => { + for (key, value) in expected_map { + let got = actual_map + .get(key) + .unwrap_or_else(|| panic!("missing key {key} in {actual}")); + assert_match_object(got, value); + } + } + (Value::Array(actual_items), Value::Array(expected_items)) => { + assert_eq!(actual_items.len(), expected_items.len(), "in {actual}"); + for (got, want) in actual_items.iter().zip(expected_items) { + assert_match_object(got, want); + } + } + _ => assert_eq!(actual, expected), + } +} + +// ── doGenerate ────────────────────────────────────────────────────────────── + +mod do_generate { + mod text { + use super::super::*; + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › text › should extract text response + #[tokio::test] + async fn should_extract_text_response() { + let fetch = json_fetch(TEXT_FIXTURE, &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_generate(&options()) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(&result).unwrap(), + json_value(SNAPSHOT_DO_GENERATE_TEXT) + ); + } + } + + mod tool_call { + use super::super::*; + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › tool call › should extract tool calls + #[tokio::test] + async fn should_extract_tool_calls() { + let fetch = json_fetch(TOOL_CALL_FIXTURE, &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_generate(&CallOptions { + tools: Some(vec![test_tool()]), + ..options() + }) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(&result).unwrap(), + json_value(SNAPSHOT_DO_GENERATE_TOOL_CALL) + ); + } + } + + mod max_tokens { + use super::super::*; + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › max tokens › should map MAX_TOKENS finish reason to length + #[tokio::test] + async fn should_map_max_tokens_finish_reason_to_length() { + let fetch = json_fetch(MAX_TOKENS_FIXTURE, &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_generate(&options()) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(&result.finish_reason).unwrap(), + json!({"raw": "MAX_TOKENS", "unified": "length"}) + ); + } + } + + mod null_tool_call_arguments { + use super::super::*; + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › null tool call arguments › should handle string "null" tool call arguments + #[tokio::test] + async fn should_handle_string_null_tool_call_arguments() { + let fetch = json_fetch(NULL_ARGS_FIXTURE, &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_generate(&CallOptions { + tools: Some(vec![Tool::Function(FunctionTool::new( + "currentTime", + json!({ + "type": "object", + "properties": {}, + "required": [], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#" + }), + ))]), + ..CallOptions::new(vec![LanguageModelMessage::user_text( + "What is the current time?", + )]) + }) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(&result.content).unwrap(), + json_value(SNAPSHOT_DO_GENERATE_NULL_TOOL_CALL_ARGS) + ); + } + } + + mod reasoning { + use super::super::*; + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › reasoning › should extract reasoning from response + #[tokio::test] + async fn should_extract_reasoning_from_response() { + let fetch = json_fetch(REASONING_FIXTURE, &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_generate(&options()) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(&result).unwrap(), + json_value(SNAPSHOT_DO_GENERATE_REASONING) + ); + } + } + + mod top_level_reasoning { + use super::super::*; + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › top-level reasoning › should map top-level reasoning to thinking enabled with budget + #[tokio::test] + async fn should_map_top_level_reasoning_to_thinking_enabled_with_budget() { + let fetch = json_fetch(TEXT_FIXTURE, &[]); + provider(&fetch) + .chat(MODEL) + .do_generate(&CallOptions { + reasoning: Some(ReasoningEffort::High), + ..options() + }) + .await + .unwrap(); + let body = fetch.seen()[0].json_body(); + assert_eq!(body["thinking"]["type"], "enabled"); + assert!( + body["thinking"]["token_budget"] + .as_f64() + .is_some_and(|budget| budget > 0.0) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › top-level reasoning › should map top-level reasoning none to thinking disabled + #[tokio::test] + async fn should_map_top_level_reasoning_none_to_thinking_disabled() { + let fetch = json_fetch(TEXT_FIXTURE, &[]); + provider(&fetch) + .chat(MODEL) + .do_generate(&CallOptions { + reasoning: Some(ReasoningEffort::None), + ..options() + }) + .await + .unwrap(); + assert_eq!( + fetch.seen()[0].json_body()["thinking"], + json!({"type": "disabled"}) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › top-level reasoning › should prefer providerOptions over top-level reasoning + #[tokio::test] + async fn should_prefer_provider_options_over_top_level_reasoning() { + let fetch = json_fetch(TEXT_FIXTURE, &[]); + provider(&fetch) + .chat(MODEL) + .do_generate(&CallOptions { + reasoning: Some(ReasoningEffort::None), + provider_options: Some(provider_options(json!({ + "cohere": {"thinking": {"type": "enabled"}} + }))), + ..options() + }) + .await + .unwrap(); + assert_eq!(fetch.seen()[0].json_body()["thinking"]["type"], "enabled"); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › top-level reasoning › should not set thinking when reasoning is not specified + #[tokio::test] + async fn should_not_set_thinking_when_reasoning_is_not_specified() { + let fetch = json_fetch(TEXT_FIXTURE, &[]); + provider(&fetch) + .chat(MODEL) + .do_generate(&options()) + .await + .unwrap(); + assert!(fetch.seen()[0].json_body().get("thinking").is_none()); + } + } + + mod citations { + use super::super::*; + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › citations › should extract citations from response + #[tokio::test] + async fn should_extract_citations_from_response() { + let fetch = json_fetch(CITATIONS_FIXTURE, &[]); + let provider = provider_with( + &fetch, + None, + Some(Arc::new(|| "test-citation-id".to_string())), + ); + let result = provider + .chat(MODEL) + .do_generate(&user_options(vec![ + text_part("What are AI benefits?"), + file_part( + b"AI provides automation and efficiency.", + "text/plain", + "ai-benefits.txt", + ), + ])) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(&result).unwrap(), + json_value(SNAPSHOT_DO_GENERATE_CITATIONS) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › citations › should extract text documents and send to API + #[tokio::test] + async fn should_extract_text_documents_and_send_to_api() { + let fetch = json_fetch(CITATIONS_FIXTURE, &[]); + provider(&fetch) + .chat(MODEL) + .do_generate(&user_options(vec![ + text_part("What does this say?"), + file_part(b"This is a test document.", "text/plain", "test.txt"), + ])) + .await + .unwrap(); + assert_eq!( + fetch.seen()[0].json_body(), + json!({ + "documents": [ + {"data": {"text": "This is a test document.", "title": "test.txt"}} + ], + "messages": [ + {"content": "What does this say?", "role": "user"} + ], + "model": "command-r-plus" + }) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › citations › should extract multiple text documents + #[tokio::test] + async fn should_extract_multiple_text_documents() { + let fetch = json_fetch(CITATIONS_FIXTURE, &[]); + provider(&fetch) + .chat(MODEL) + .do_generate(&user_options(vec![ + text_part("What do these documents say?"), + file_part(b"First document content", "text/plain", "doc1.txt"), + file_part(b"Second document content", "text/plain", "doc2.txt"), + ])) + .await + .unwrap(); + assert_eq!( + fetch.seen()[0].json_body(), + json!({ + "documents": [ + {"data": {"text": "First document content", "title": "doc1.txt"}}, + {"data": {"text": "Second document content", "title": "doc2.txt"}} + ], + "messages": [ + {"content": "What do these documents say?", "role": "user"} + ], + "model": "command-r-plus" + }) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › citations › should support JSON files + #[tokio::test] + async fn should_support_json_files() { + let fetch = json_fetch(CITATIONS_FIXTURE, &[]); + provider(&fetch) + .chat(MODEL) + .do_generate(&user_options(vec![ + text_part("What is in this JSON?"), + file_part(b"{\"key\": \"value\"}", "application/json", "data.json"), + ])) + .await + .unwrap(); + assert_eq!( + fetch.seen()[0].json_body(), + json!({ + "documents": [ + {"data": {"text": "{\"key\": \"value\"}", "title": "data.json"}} + ], + "messages": [ + {"content": "What is in this JSON?", "role": "user"} + ], + "model": "command-r-plus" + }) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › citations › should not include mediaType in the outgoing payload (category D) + #[tokio::test] + async fn should_not_include_media_type_in_the_outgoing_payload() { + let fetch = json_fetch(CITATIONS_FIXTURE, &[]); + provider(&fetch) + .chat(MODEL) + .do_generate(&user_options(vec![ + text_part("What is this?"), + file_part(b"Some file content", "application/pdf", "document.pdf"), + ])) + .await + .unwrap(); + let body = fetch.seen()[0].json_body(); + assert_eq!( + body["documents"], + json!([{"data": {"text": "Some file content", "title": "document.pdf"}}]) + ); + let text = body.to_string(); + assert!(!text.contains("application/pdf")); + assert!(!text.contains("mediaType")); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › citations › should successfully process supported text media types + #[tokio::test] + async fn should_successfully_process_supported_text_media_types() { + let fetch = json_fetch(CITATIONS_FIXTURE, &[]); + provider(&fetch) + .chat(MODEL) + .do_generate(&user_options(vec![ + text_part("What is this?"), + file_part(b"This is plain text content", "text/plain", "text.txt"), + file_part(b"# Markdown Header\nContent", "text/markdown", "doc.md"), + ])) + .await + .unwrap(); + assert_match_object( + &fetch.seen()[0].json_body(), + &json!({ + "documents": [ + {"data": {"text": "This is plain text content", "title": "text.txt"}}, + {"data": {"text": "# Markdown Header\nContent", "title": "doc.md"}} + ] + }), + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › citations › should not include documents parameter when no files present + #[tokio::test] + async fn should_not_include_documents_parameter_when_no_files_present() { + let fetch = json_fetch(CITATIONS_FIXTURE, &[]); + provider(&fetch) + .chat(MODEL) + .do_generate(&options()) + .await + .unwrap(); + assert!(fetch.seen()[0].json_body().get("documents").is_none()); + } + } + + mod request { + use super::super::*; + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › request › should pass model and messages + #[tokio::test] + async fn should_pass_model_and_messages() { + let fetch = json_fetch(TEXT_FIXTURE, &[]); + provider(&fetch) + .chat(MODEL) + .do_generate(&options()) + .await + .unwrap(); + assert_eq!( + fetch.seen()[0].json_body(), + json!({ + "messages": [ + {"content": "you are a friendly bot!", "role": "system"}, + {"content": "Hello", "role": "user"} + ], + "model": "command-r-plus" + }) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › request › should pass tools + #[tokio::test] + async fn should_pass_tools() { + let fetch = json_fetch(TEXT_FIXTURE, &[]); + provider(&fetch) + .chat(MODEL) + .do_generate(&CallOptions { + tool_choice: Some(LanguageModelV4ToolChoice::None), + tools: Some(vec![test_tool()]), + ..options() + }) + .await + .unwrap(); + assert_eq!( + fetch.seen()[0].json_body(), + json!({ + "messages": [ + {"content": "you are a friendly bot!", "role": "system"}, + {"content": "Hello", "role": "user"} + ], + "model": "command-r-plus", + "tool_choice": "NONE", + "tools": [ + { + "function": { + "name": "test-tool", + "parameters": { + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "properties": {"value": {"type": "string"}}, + "required": ["value"], + "type": "object" + } + }, + "type": "function" + } + ] + }) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › request › should pass headers + #[tokio::test] + async fn should_pass_headers() { + let fetch = json_fetch(TEXT_FIXTURE, &[]); + let provider = provider_with( + &fetch, + Some(HashMap::from([( + "Custom-Provider-Header".to_string(), + Some("provider-header-value".to_string()), + )])), + None, + ); + provider + .chat(MODEL) + .do_generate(&CallOptions { + headers: Some(HashMap::from([( + "Custom-Request-Header".to_string(), + "request-header-value".to_string(), + )])), + ..options() + }) + .await + .unwrap(); + assert_sent_headers( + &fetch, + &[ + ("authorization", "Bearer test-api-key"), + ("content-type", "application/json"), + ("custom-provider-header", "provider-header-value"), + ("custom-request-header", "request-header-value"), + ], + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › request › should pass response format + #[tokio::test] + async fn should_pass_response_format() { + let fetch = json_fetch(TEXT_FIXTURE, &[]); + provider(&fetch) + .chat(MODEL) + .do_generate(&CallOptions { + response_format: Some(ResponseFormat::Json { + schema: Some(json!({ + "type": "object", + "properties": {"text": {"type": "string"}}, + "required": ["text"] + })), + name: None, + description: None, + }), + ..options() + }) + .await + .unwrap(); + assert_eq!( + fetch.seen()[0].json_body(), + json!({ + "messages": [ + {"content": "you are a friendly bot!", "role": "system"}, + {"content": "Hello", "role": "user"} + ], + "model": "command-r-plus", + "response_format": { + "json_schema": { + "properties": {"text": {"type": "string"}}, + "required": ["text"], + "type": "object" + }, + "type": "json_object" + } + }) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › request › should send request body + #[tokio::test] + async fn should_send_request_body() { + let fetch = json_fetch(TEXT_FIXTURE, &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_generate(&options()) + .await + .unwrap(); + assert_eq!( + result.request.unwrap().body, + Some(json!({ + "messages": [ + {"content": "you are a friendly bot!", "role": "system"}, + {"content": "Hello", "role": "user"} + ], + "model": "command-r-plus" + })) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › request › should expose the raw response headers + #[tokio::test] + async fn should_expose_the_raw_response_headers() { + let fetch = json_fetch(TEXT_FIXTURE, &[("test-header", "test-value")]); + let result = provider(&fetch) + .chat(MODEL) + .do_generate(&options()) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(result.response.unwrap().headers).unwrap(), + json!({ + "content-length": "304", + "content-type": "application/json", + "test-header": "test-value" + }) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › request › should extract usage + #[tokio::test] + async fn should_extract_usage() { + let fetch = json_fetch(TEXT_FIXTURE, &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_generate(&options()) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(&result.usage).unwrap(), + json!({ + "inputTokens": {"noCache": 507, "total": 507}, + "outputTokens": {"text": 10, "total": 10}, + "raw": { + "billed_units": {"input_tokens": 12, "output_tokens": 7}, + "cached_tokens": 448, + "tokens": {"input_tokens": 507, "output_tokens": 10} + } + }) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › request › should preserve extra top-level and nested fields in raw usage + #[tokio::test] + async fn should_preserve_extra_top_level_and_nested_fields_in_raw_usage() { + let fetch = json_fetch( + &json!({ + "message": { + "role": "assistant", + "content": [{"type": "text", "text": "Hello"}] + }, + "finish_reason": "COMPLETE", + "usage": { + "billed_units": { + "input_tokens": 12, + "output_tokens": 7, + "billing_tier": "standard" + }, + "tokens": { + "input_tokens": 507, + "output_tokens": 10, + "tokenizer": "command" + }, + "cached_tokens": 448, + "provider_usage_id": "usage-123" + } + }) + .to_string(), + &[], + ); + let result = provider(&fetch) + .chat(MODEL) + .do_generate(&options()) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(&result.usage).unwrap(), + json!({ + "inputTokens": {"total": 507, "noCache": 507}, + "outputTokens": {"total": 10, "text": 10}, + "raw": { + "billed_units": { + "input_tokens": 12, + "output_tokens": 7, + "billing_tier": "standard" + }, + "tokens": { + "input_tokens": 507, + "output_tokens": 10, + "tokenizer": "command" + }, + "cached_tokens": 448, + "provider_usage_id": "usage-123" + } + }) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › request › should validate cached token usage + #[tokio::test] + async fn should_validate_cached_token_usage() { + let fetch = json_fetch( + &json!({ + "message": { + "role": "assistant", + "content": [{"type": "text", "text": "Hello"}] + }, + "finish_reason": "COMPLETE", + "usage": { + "billed_units": {"input_tokens": 12, "output_tokens": 7}, + "tokens": {"input_tokens": 507, "output_tokens": 10}, + "cached_tokens": "448" + } + }) + .to_string(), + &[], + ); + assert!( + provider(&fetch) + .chat(MODEL) + .do_generate(&options()) + .await + .is_err() + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doGenerate › request › should send additional response information + #[tokio::test] + async fn should_send_additional_response_information() { + let fetch = json_fetch(TEXT_FIXTURE, &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_generate(&options()) + .await + .unwrap(); + let response = result.response.unwrap(); + assert!(response.id.is_none()); + assert!(response.timestamp.is_none()); + assert!(response.model_id.is_none()); + } + } +} + +// ── doStream ──────────────────────────────────────────────────────────────── + +mod do_stream { + mod text { + use super::super::*; + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doStream › text › should include raw chunks when includeRawChunks is enabled + #[tokio::test] + async fn should_include_raw_chunks_when_include_raw_chunks_is_enabled() { + let fetch = stream_fetch(&chunks_frames(TEXT_CHUNKS), &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_stream(&CallOptions { + include_raw_chunks: Some(true), + ..options() + }) + .await + .unwrap(); + let parts = stream_parts(result).await; + let raw: Vec = parts + .into_iter() + .filter(|part| part["type"] == "raw") + .collect(); + assert_eq!(Value::Array(raw), json_value(SNAPSHOT_DO_STREAM_RAW_CHUNKS)); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doStream › text › should not include raw chunks when includeRawChunks is false + #[tokio::test] + async fn should_not_include_raw_chunks_when_include_raw_chunks_is_false() { + let fetch = stream_fetch(&chunks_frames(TEXT_CHUNKS), &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_stream(&CallOptions { + include_raw_chunks: Some(false), + ..options() + }) + .await + .unwrap(); + let parts = stream_parts(result).await; + assert_eq!(parts.iter().filter(|part| part["type"] == "raw").count(), 0); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doStream › text › should preserve extra top-level and nested fields in raw usage + #[tokio::test] + async fn should_preserve_extra_top_level_and_nested_fields_in_raw_usage() { + let raw_usage = json!({ + "billed_units": { + "input_tokens": 12, + "output_tokens": 7, + "billing_tier": "standard" + }, + "tokens": { + "input_tokens": 507, + "output_tokens": 10, + "tokenizer": "command" + }, + "cached_tokens": 448, + "provider_usage_id": "usage-123" + }); + let fetch = stream_fetch(&chunk_frames(&[message_end(raw_usage.clone())]), &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_stream(&CallOptions { + include_raw_chunks: Some(false), + ..options() + }) + .await + .unwrap(); + let parts = stream_parts(result).await; + assert_eq!( + parts.last().unwrap(), + &json!({ + "type": "finish", + "finishReason": {"unified": "stop", "raw": "COMPLETE"}, + "usage": { + "inputTokens": {"total": 507, "noCache": 507}, + "outputTokens": {"total": 10, "text": 10}, + "raw": raw_usage + } + }) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doStream › text › should finish successfully when billed units are absent + #[tokio::test] + async fn should_finish_successfully_when_billed_units_are_absent() { + let raw_usage = json!({"tokens": {"input_tokens": 507, "output_tokens": 10}}); + let fetch = stream_fetch(&chunk_frames(&[message_end(raw_usage.clone())]), &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_stream(&CallOptions { + include_raw_chunks: Some(false), + ..options() + }) + .await + .unwrap(); + let parts = stream_parts(result).await; + assert_eq!( + parts.last().unwrap(), + &json!({ + "type": "finish", + "finishReason": {"unified": "stop", "raw": "COMPLETE"}, + "usage": { + "inputTokens": {"total": 507, "noCache": 507}, + "outputTokens": {"total": 10, "text": 10}, + "raw": raw_usage + } + }) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doStream › text › should validate streamed $field + #[tokio::test] + async fn should_validate_streamed_field() { + // The `it.each` rows of the TypeScript test, one loop over them. + let rows = [ + ( + "billed input tokens", + json!({ + "billed_units": {"input_tokens": "12", "output_tokens": 7}, + "tokens": {"input_tokens": 507, "output_tokens": 10}, + "cached_tokens": 448 + }), + ), + ( + "cached tokens", + json!({ + "billed_units": {"input_tokens": 12, "output_tokens": 7}, + "tokens": {"input_tokens": 507, "output_tokens": 10}, + "cached_tokens": "448" + }), + ), + ]; + for (field, usage) in rows { + let fetch = stream_fetch(&chunk_frames(&[message_end(usage)]), &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_stream(&CallOptions { + include_raw_chunks: Some(false), + ..options() + }) + .await + .unwrap(); + let parts = stream_parts(result).await; + assert_eq!(parts.len(), 3, "{field}"); + assert_eq!( + parts[0], + json!({"type": "stream-start", "warnings": []}), + "{field}" + ); + assert_eq!(parts[1]["type"], "error", "{field}"); + assert_eq!( + parts[2], + json!({ + "type": "finish", + "finishReason": {"unified": "error"}, + "usage": {"inputTokens": {}, "outputTokens": {}} + }), + "{field}" + ); + } + } + } + + mod reasoning { + use super::super::*; + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doStream › reasoning › should stream reasoning deltas + #[tokio::test] + async fn should_stream_reasoning_deltas() { + let fetch = stream_fetch(&chunks_frames(REASONING_CHUNKS), &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_stream(&CallOptions { + include_raw_chunks: Some(false), + ..options() + }) + .await + .unwrap(); + assert_eq!( + Value::Array(stream_parts(result).await), + json_value(SNAPSHOT_DO_STREAM_REASONING) + ); + } + } + + mod tool_call { + use super::super::*; + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doStream › tool call › should stream tool deltas + #[tokio::test] + async fn should_stream_tool_deltas() { + let fetch = stream_fetch(&chunks_frames(TOOL_CALL_CHUNKS), &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_stream(&CallOptions { + tools: Some(vec![test_tool()]), + include_raw_chunks: Some(false), + ..options() + }) + .await + .unwrap(); + assert_eq!( + Value::Array(stream_parts(result).await), + json_value(SNAPSHOT_DO_STREAM_TOOL_CALL) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doStream › tool call › rejects prototype keys in streamed tool call arguments + #[tokio::test] + async fn rejects_prototype_keys_in_streamed_tool_call_arguments() { + let fetch = stream_fetch( + &chunk_frames(&[ + json!({ + "type": "tool-call-start", + "index": 0, + "delta": {"message": {"tool_calls": { + "id": "test-tool-call", + "type": "function", + "function": {"name": "test-tool", "arguments": ""} + }}} + }), + json!({ + "type": "tool-call-delta", + "index": 0, + "delta": {"message": {"tool_calls": { + "function": {"arguments": "{\"__proto__\":{\"polluted\":true}}"} + }}} + }), + json!({"type": "tool-call-end", "index": 0}), + ]), + &[], + ); + let result = provider(&fetch) + .chat(MODEL) + .do_stream(&CallOptions { + tools: Some(vec![Tool::Function(FunctionTool::new( + "test-tool", + json!({ + "type": "object", + "properties": {}, + "additionalProperties": true, + "$schema": "http://json-schema.org/draft-07/schema#" + }), + ))]), + include_raw_chunks: Some(false), + ..options() + }) + .await + .unwrap(); + let items: Vec<_> = result.stream.collect().await; + let error = items + .iter() + .find_map(|item| item.as_ref().err()) + .expect("the stream fails"); + assert!( + error + .to_string() + .contains("Object contains forbidden prototype property") + ); + } + } + + mod empty_tool_call { + use super::super::*; + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doStream › empty tool call › should handle empty tool call arguments + #[tokio::test] + async fn should_handle_empty_tool_call_arguments() { + let fetch = stream_fetch(&chunks_frames(EMPTY_TOOL_CALL_CHUNKS), &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_stream(&CallOptions { + tools: Some(vec![Tool::Function(FunctionTool::new( + "test-tool", + json!({ + "type": "object", + "properties": {}, + "required": [], + "additionalProperties": false, + "$schema": "http://json-schema.org/draft-07/schema#" + }), + ))]), + include_raw_chunks: Some(false), + ..options() + }) + .await + .unwrap(); + assert_eq!( + Value::Array(stream_parts(result).await), + json_value(SNAPSHOT_DO_STREAM_EMPTY_TOOL_CALL) + ); + } + } + + mod request { + use super::super::*; + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doStream › request › should pass the messages and the model + #[tokio::test] + async fn should_pass_the_messages_and_the_model() { + let fetch = stream_fetch(&chunks_frames(TEXT_CHUNKS), &[]); + provider(&fetch) + .chat(MODEL) + .do_stream(&CallOptions { + include_raw_chunks: Some(false), + ..options() + }) + .await + .unwrap(); + assert_eq!( + fetch.seen()[0].json_body(), + json!({ + "messages": [ + {"content": "you are a friendly bot!", "role": "system"}, + {"content": "Hello", "role": "user"} + ], + "model": "command-r-plus", + "stream": true + }) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doStream › request › should pass headers + #[tokio::test] + async fn should_pass_headers() { + let fetch = stream_fetch(&chunks_frames(TEXT_CHUNKS), &[]); + let provider = provider_with( + &fetch, + Some(HashMap::from([( + "Custom-Provider-Header".to_string(), + Some("provider-header-value".to_string()), + )])), + None, + ); + provider + .chat(MODEL) + .do_stream(&CallOptions { + headers: Some(HashMap::from([( + "Custom-Request-Header".to_string(), + "request-header-value".to_string(), + )])), + include_raw_chunks: Some(false), + ..options() + }) + .await + .unwrap(); + assert_sent_headers( + &fetch, + &[ + ("authorization", "Bearer test-api-key"), + ("content-type", "application/json"), + ("custom-provider-header", "provider-header-value"), + ("custom-request-header", "request-header-value"), + ], + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doStream › request › should send request body + #[tokio::test] + async fn should_send_request_body() { + let fetch = stream_fetch(&chunks_frames(TEXT_CHUNKS), &[]); + let result = provider(&fetch) + .chat(MODEL) + .do_stream(&CallOptions { + include_raw_chunks: Some(false), + ..options() + }) + .await + .unwrap(); + assert_eq!( + result.request.unwrap().body, + Some(json!({ + "messages": [ + {"content": "you are a friendly bot!", "role": "system"}, + {"content": "Hello", "role": "user"} + ], + "model": "command-r-plus", + "stream": true + })) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doStream › request › should expose the raw response headers + #[tokio::test] + async fn should_expose_the_raw_response_headers() { + let fetch = stream_fetch( + &chunks_frames(TEXT_CHUNKS), + &[("test-header", "test-value")], + ); + let result = provider(&fetch) + .chat(MODEL) + .do_stream(&CallOptions { + include_raw_chunks: Some(false), + ..options() + }) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(result.response.unwrap().headers).unwrap(), + json!({ + "cache-control": "no-cache", + "connection": "keep-alive", + "content-type": "text/event-stream", + "test-header": "test-value" + }) + ); + } + } +} diff --git a/aimux-providers/src/cohere/tests/convert_to_chat_prompt.rs b/aimux-providers/src/cohere/tests/convert_to_chat_prompt.rs new file mode 100644 index 000000000..77d12d5de --- /dev/null +++ b/aimux-providers/src/cohere/tests/convert_to_chat_prompt.rs @@ -0,0 +1,541 @@ +//! Ports of `src/convert-to-cohere-chat-prompt.test.ts` from `@ai-sdk/cohere@4.0.52`. +//! +//! One nested module per `describe`, one `#[test]` per `it`. Each test carries +//! its upstream chain in an `upstream:` marker line. + +use crate::cohere::convert::{build_request_body, convert_prompt_to_cohere}; +use crate::cohere::options::NAMESPACE; +use aimux_core::AiMuxError; +use aimux_core::language_model_message::{ + AssistantPart, LanguageModelMessage, LanguageModelPrompt, TextPart, ToolCallPart, ToolPart, + ToolResultPart, UserPart, +}; +use aimux_core::language_model_message::{FilePart, ToolResultOutput}; +use aimux_core::options::CallOptions; +use aimux_core::shared::provider_namespace; +use aimux_core::shared::{FileBytes, FileData}; +use indexmap::IndexMap; +use serde_json::json; + +/// Upstream's `warnings: []`. `convert_prompt_to_cohere` returns no warnings; +/// the warnings a prompt can produce are collected by `build_request_body`, so +/// that path stands in for the TS `warnings` field. +fn assert_no_warnings(prompt: LanguageModelPrompt) { + let options = CallOptions { + prompt, + ..Default::default() + }; + let body = build_request_body("command-a", &options, false).expect("request body"); + assert!(body.warnings.is_empty(), "warnings: {:?}", body.warnings); +} + +mod convert_to_cohere_chat_prompt { + use super::*; + + mod file_processing { + use super::*; + + /// upstream: @ai-sdk/cohere src/convert-to-cohere-chat-prompt.test.ts › convert to cohere chat prompt › file processing › should extract documents from file parts + #[test] + fn should_extract_documents_from_file_parts() { + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::User { + content: vec![ + UserPart::Text(TextPart { + text: "Analyze this file: ".into(), + provider_options: None, + }), + UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Binary(b"This is file content".to_vec()), + }, + media_type: "text/plain".into(), + filename: Some("test.txt".into()), + provider_options: None, + }), + ], + provider_options: None, + }]; + + let result = convert_prompt_to_cohere(&prompt).unwrap(); + + assert_eq!( + result.messages, + vec![json!({ + "role": "user", + "content": "Analyze this file: ", + })] + ); + assert_eq!( + result.documents, + vec![json!({ + "data": { + "text": "This is file content", + "title": "test.txt", + }, + })] + ); + assert_no_warnings(prompt); + } + + /// upstream: @ai-sdk/cohere src/convert-to-cohere-chat-prompt.test.ts › convert to cohere chat prompt › file processing › should accept top-level-only mediaType without error (category D: mediaType not consumed) + #[test] + fn should_accept_top_level_only_media_type_without_error_category_d_media_type_not_consumed() + { + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Binary(b"This is file content".to_vec()), + }, + media_type: "text".into(), + filename: Some("test.txt".into()), + provider_options: None, + })], + provider_options: None, + }]; + + let result = convert_prompt_to_cohere(&prompt).unwrap(); + + assert_eq!( + result.messages, + vec![json!({ + "role": "user", + "content": "", + })] + ); + assert_eq!( + result.documents, + vec![json!({ + "data": { + "text": "This is file content", + "title": "test.txt", + }, + })] + ); + assert_no_warnings(prompt); + } + + /// upstream: @ai-sdk/cohere src/convert-to-cohere-chat-prompt.test.ts › convert to cohere chat prompt › file processing › should not read mediaType (document payload carries only text + title) + #[test] + fn should_not_read_media_type_document_payload_carries_only_text_title() { + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Binary(b"PDF-like content".to_vec()), + }, + media_type: "application/pdf".into(), + filename: Some("test.pdf".into()), + provider_options: None, + })], + provider_options: None, + }]; + + let result = convert_prompt_to_cohere(&prompt).unwrap(); + + assert_eq!( + result.documents, + vec![json!({ + "data": { + "text": "PDF-like content", + "title": "test.pdf", + }, + })] + ); + let payload = json!({ + "messages": result.messages, + "documents": result.documents, + }) + .to_string(); + assert!(!payload.contains("application/pdf"), "{payload}"); + assert!(!payload.contains("mediaType"), "{payload}"); + assert_no_warnings(prompt); + } + } + + mod image_processing { + use super::*; + + /// upstream: @ai-sdk/cohere src/convert-to-cohere-chat-prompt.test.ts › convert to cohere chat prompt › image processing › should convert image file with data bytes into image_url data URI + #[test] + fn should_convert_image_file_with_data_bytes_into_image_url_data_uri() { + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::User { + content: vec![ + UserPart::Text(TextPart { + text: "What is in this image?".into(), + provider_options: None, + }), + UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Binary(vec![0, 1, 2, 3]), + }, + media_type: "image/png".into(), + filename: None, + provider_options: None, + }), + ], + provider_options: None, + }]; + + let result = convert_prompt_to_cohere(&prompt).unwrap(); + + assert_eq!( + result.messages, + vec![json!({ + "role": "user", + "content": [ + { "type": "text", "text": "What is in this image?" }, + { + "type": "image_url", + "image_url": { "url": "data:image/png;base64,AAECAw==" }, + }, + ], + })] + ); + assert!(result.documents.is_empty()); + assert_no_warnings(prompt); + } + + /// upstream: @ai-sdk/cohere src/convert-to-cohere-chat-prompt.test.ts › convert to cohere chat prompt › image processing › should convert image file with URL data into image_url URL + #[test] + fn should_convert_image_file_with_url_data_into_image_url_url() { + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Url { + url: "https://example.com/cat.png".into(), + original_url: None, + }, + media_type: "image/png".into(), + filename: None, + provider_options: None, + })], + provider_options: None, + }]; + + let result = convert_prompt_to_cohere(&prompt).unwrap(); + + assert_eq!( + result.messages, + vec![json!({ + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { "url": "https://example.com/cat.png" }, + }, + ], + })] + ); + assert!(result.documents.is_empty()); + assert_no_warnings(prompt); + } + + /// upstream: @ai-sdk/cohere src/convert-to-cohere-chat-prompt.test.ts › convert to cohere chat prompt › image processing › should pass through detail provider option as image_url.detail + #[test] + fn should_pass_through_detail_provider_option_as_image_url_detail() { + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Binary(vec![0, 1, 2, 3]), + }, + media_type: "image/png".into(), + filename: None, + provider_options: Some( + provider_namespace(NAMESPACE, json!({ "detail": "high" })).unwrap(), + ), + })], + provider_options: None, + }]; + + let result = convert_prompt_to_cohere(&prompt).unwrap(); + + assert_eq!( + result.messages[0], + json!({ + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,AAECAw==", + "detail": "high", + }, + }, + ], + }) + ); + assert_no_warnings(prompt); + } + + /// upstream: @ai-sdk/cohere src/convert-to-cohere-chat-prompt.test.ts › convert to cohere chat prompt › image processing › should omit detail when no provider option is set + #[test] + fn should_omit_detail_when_no_provider_option_is_set() { + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Binary(vec![0, 1, 2, 3]), + }, + media_type: "image/png".into(), + filename: None, + provider_options: None, + })], + provider_options: None, + }]; + + let result = convert_prompt_to_cohere(&prompt).unwrap(); + + let user_message = &result.messages[0]; + assert_eq!(user_message["role"], "user"); + assert!(user_message["content"].is_array()); + let part = &user_message["content"][0]; + assert_eq!(part["type"], "image_url"); + assert!(part["image_url"].get("detail").is_none()); + assert_no_warnings(prompt); + } + + /// upstream: @ai-sdk/cohere src/convert-to-cohere-chat-prompt.test.ts › convert to cohere chat prompt › image processing › should send image inline and route non-image file to documents + #[test] + fn should_send_image_inline_and_route_non_image_file_to_documents() { + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::User { + content: vec![ + UserPart::Text(TextPart { + text: "See attached:".into(), + provider_options: None, + }), + UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Binary(vec![0, 1, 2, 3]), + }, + media_type: "image/png".into(), + filename: None, + provider_options: None, + }), + UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Binary(b"Doc text".to_vec()), + }, + media_type: "text/plain".into(), + filename: Some("note.txt".into()), + provider_options: None, + }), + ], + provider_options: None, + }]; + + let result = convert_prompt_to_cohere(&prompt).unwrap(); + + assert_eq!( + result.messages, + vec![json!({ + "role": "user", + "content": [ + { "type": "text", "text": "See attached:" }, + { + "type": "image_url", + "image_url": { "url": "data:image/png;base64,AAECAw==" }, + }, + ], + })] + ); + assert_eq!( + result.documents, + vec![json!({ + "data": { + "text": "Doc text", + "title": "note.txt", + }, + })] + ); + assert_no_warnings(prompt); + } + + /// upstream: @ai-sdk/cohere src/convert-to-cohere-chat-prompt.test.ts › convert to cohere chat prompt › image processing › should accept top-level "image" media type and detect full type from bytes + #[test] + fn should_accept_top_level_image_media_type_and_detect_full_type_from_bytes() { + let png_signature = vec![0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]; + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Data { + data: FileBytes::Binary(png_signature), + }, + media_type: "image".into(), + filename: None, + provider_options: None, + })], + provider_options: None, + }]; + + let result = convert_prompt_to_cohere(&prompt).unwrap(); + + let part = &result.messages[0]["content"][0]; + assert_eq!(part["type"], "image_url"); + let url = part["image_url"]["url"].as_str().unwrap(); + assert!( + url.starts_with("data:image/png;base64,"), + "unexpected url: {url}" + ); + assert_no_warnings(prompt); + } + } + + mod tool_messages { + use super::*; + + /// upstream: @ai-sdk/cohere src/convert-to-cohere-chat-prompt.test.ts › convert to cohere chat prompt › tool messages › should convert a tool call into a cohere chatbot message + #[test] + fn should_convert_a_tool_call_into_a_cohere_chatbot_message() { + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::Assistant { + content: vec![ + AssistantPart::Text(TextPart { + text: "Calling a tool".into(), + provider_options: None, + }), + AssistantPart::ToolCall(ToolCallPart { + tool_call_id: "tool-call-1".into(), + tool_name: "tool-1".into(), + input: json!({ "test": "This is a tool message" }), + provider_executed: None, + provider_options: None, + }), + ], + provider_options: None, + }]; + + let result = convert_prompt_to_cohere(&prompt).unwrap(); + + // Upstream `content: undefined` is absent from the message here. + assert_eq!( + result.messages, + vec![json!({ + "role": "assistant", + "tool_calls": [ + { + "id": "tool-call-1", + "type": "function", + "function": { + "name": "tool-1", + "arguments": json!({ "test": "This is a tool message" }).to_string(), + }, + }, + ], + })] + ); + assert!(result.documents.is_empty()); + assert_no_warnings(prompt); + } + + /// upstream: @ai-sdk/cohere src/convert-to-cohere-chat-prompt.test.ts › convert to cohere chat prompt › tool messages › should convert a single tool result into a cohere tool message + #[test] + fn should_convert_a_single_tool_result_into_a_cohere_tool_message() { + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::Tool { + content: vec![ToolPart::ToolResult(ToolResultPart { + tool_call_id: "tool-call-1".into(), + tool_name: "tool-1".into(), + output: ToolResultOutput::Json { + value: json!({ "test": "This is a tool message" }), + provider_options: None, + }, + provider_options: None, + })], + provider_options: None, + }]; + + let result = convert_prompt_to_cohere(&prompt).unwrap(); + + assert_eq!( + result.messages, + vec![json!({ + "role": "tool", + "content": json!({ "test": "This is a tool message" }).to_string(), + "tool_call_id": "tool-call-1", + })] + ); + assert!(result.documents.is_empty()); + assert_no_warnings(prompt); + } + + /// upstream: @ai-sdk/cohere src/convert-to-cohere-chat-prompt.test.ts › convert to cohere chat prompt › tool messages › should convert multiple tool results into a cohere tool message + #[test] + fn should_convert_multiple_tool_results_into_a_cohere_tool_message() { + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::Tool { + content: vec![ + ToolPart::ToolResult(ToolResultPart { + tool_call_id: "tool-call-1".into(), + tool_name: "tool-1".into(), + output: ToolResultOutput::Json { + value: json!({ "test": "This is a tool message" }), + provider_options: None, + }, + provider_options: None, + }), + ToolPart::ToolResult(ToolResultPart { + tool_call_id: "tool-call-2".into(), + tool_name: "tool-2".into(), + output: ToolResultOutput::Json { + value: json!({ "something": "else" }), + provider_options: None, + }, + provider_options: None, + }), + ], + provider_options: None, + }]; + + let result = convert_prompt_to_cohere(&prompt).unwrap(); + + assert_eq!( + result.messages, + vec![ + json!({ + "role": "tool", + "content": json!({ "test": "This is a tool message" }).to_string(), + "tool_call_id": "tool-call-1", + }), + json!({ + "role": "tool", + "content": json!({ "something": "else" }).to_string(), + "tool_call_id": "tool-call-2", + }), + ] + ); + assert!(result.documents.is_empty()); + assert_no_warnings(prompt); + } + } + + mod provider_reference { + use super::*; + + /// upstream: @ai-sdk/cohere src/convert-to-cohere-chat-prompt.test.ts › convert to cohere chat prompt › provider reference › should throw for file parts with provider references + #[test] + fn should_throw_for_file_parts_with_provider_references() { + let prompt: LanguageModelPrompt = vec![LanguageModelMessage::User { + content: vec![UserPart::File(FilePart { + data: FileData::Reference { + reference: IndexMap::from([( + NAMESPACE.to_string(), + "doc-ref-123".to_string(), + )]), + }, + media_type: "text/plain".into(), + filename: None, + provider_options: None, + })], + provider_options: None, + }]; + + // Upstream's message is "'file parts with provider references' + // functionality not supported"; the Rust error carries the + // functionality name, which is what is asserted here. + let Err(err) = convert_prompt_to_cohere(&prompt) else { + panic!("expected an error for a file part with a provider reference"); + }; + assert!( + matches!( + &err, + AiMuxError::UnsupportedFunctionality(name) + if name == "file parts with provider references" + ), + "unexpected error: {err}" + ); + } + } +} diff --git a/aimux-providers/src/cohere/tests/embedding_model.rs b/aimux-providers/src/cohere/tests/embedding_model.rs new file mode 100644 index 000000000..f45cdf864 --- /dev/null +++ b/aimux-providers/src/cohere/tests/embedding_model.rs @@ -0,0 +1,138 @@ +//! Ports of `src/cohere-embedding-model.test.ts` from `@ai-sdk/cohere@4.0.52`. +//! +//! Only the cases that the other tests do not already cover are here: +//! `should extract embedding`, `should extract usage`, `should pass the model +//! and the values` and `should pass headers` are in +//! `tests/cohere_embedding_test.rs` and `tests/cohere_upstream_e2e_test.rs`. +//! `should expose the raw response` compares the whole snapshot here; the +//! older case in `tests/cohere_embedding_test.rs` checks part of it. +//! +//! Harness differences from the TypeScript (expectations are unchanged): +//! - The TypeScript test server answers from the fixture file. The mock +//! answers with the fixture's content; `content-length` is the byte length +//! of the compact JSON body, as the TypeScript server sends it. +//! - The mock records `user-agent`; the TypeScript test checks it with +//! `vi.mock('./version')`. These cases do not check it. + +use std::sync::Arc; + +use crate::cohere::{CohereProvider, CohereProviderSettings, create_cohere}; +use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; +use aimux_core::shared::SharedProviderOptions; +use serde_json::{Value, json}; + +use crate::test_support::mock_fetch::{Canned, MockFetch}; + +const EMBEDDING_FIXTURE: &str = include_str!("fixtures/cohere-embedding.json"); +const SNAPSHOT_DO_EMBED_RAW_RESPONSE: &str = include_str!("snapshots/do_embed_raw_response.json"); + +const TEST_VALUES: [&str; 2] = ["sunny day at the beach", "rainy day in the city"]; + +fn provider(fetch: &Arc) -> CohereProvider { + create_cohere(CohereProviderSettings { + api_key: Some("test-api-key".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() +} + +/// A single-response mock for the embedding fixture. +fn embedding_fetch() -> Arc { + embedding_fetch_with(&[]) +} + +/// `prepareJsonFixtureResponse('cohere-embedding', headers)`. +fn embedding_fetch_with(extra_headers: &[(&str, &str)]) -> Arc { + let fixture: Value = serde_json::from_str(EMBEDDING_FIXTURE).expect("fixture is JSON"); + let bytes = serde_json::to_vec(&fixture).expect("serializable"); + let mut headers = vec![ + ("content-type".to_string(), "application/json".to_string()), + ("content-length".to_string(), bytes.len().to_string()), + ]; + headers.extend( + extra_headers + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())), + ); + MockFetch::new(vec![Canned { + status: 200, + headers, + body: bytes, + }]) +} + +/// `{ values: testValues, providerOptions }`. +fn options(provider_options: Value) -> EmbeddingCallOptions { + let mut options = EmbeddingCallOptions::new(TEST_VALUES[0]); + options.values = TEST_VALUES.iter().copied().map(String::from).collect(); + options.provider_options = + Some(serde_json::from_value::(provider_options).unwrap()); + options +} + +mod do_embed { + use super::*; + + /// upstream: @ai-sdk/cohere src/cohere-embedding-model.test.ts › doEmbed › should expose the raw response + #[tokio::test] + async fn should_expose_the_raw_response() { + let fetch = embedding_fetch_with(&[("test-header", "test-value")]); + let mut call = options(json!({})); + call.provider_options = None; + let result = provider(&fetch) + .embedding("embed-english-v3.0") + .do_embed(&call) + .await + .unwrap(); + let snapshot: Value = + serde_json::from_str(SNAPSHOT_DO_EMBED_RAW_RESPONSE).expect("snapshot is JSON"); + assert_eq!( + serde_json::to_value(result.response.expect("response")).unwrap(), + snapshot + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-embedding-model.test.ts › doEmbed › should pass the input_type setting + #[tokio::test] + async fn should_pass_the_input_type_setting() { + let fetch = embedding_fetch(); + provider(&fetch) + .embedding("embed-english-v3.0") + .do_embed(&options( + json!({"cohere": {"inputType": "search_document"}}), + )) + .await + .unwrap(); + assert_eq!( + fetch.seen()[0].json_body(), + json!({ + "embedding_types": ["float"], + "input_type": "search_document", + "model": "embed-english-v3.0", + "texts": ["sunny day at the beach", "rainy day in the city"] + }) + ); + } + + /// upstream: @ai-sdk/cohere src/cohere-embedding-model.test.ts › doEmbed › should pass the output_dimension setting + #[tokio::test] + async fn should_pass_the_output_dimension_setting() { + let fetch = embedding_fetch(); + provider(&fetch) + .embedding("embed-v4.0") + .do_embed(&options(json!({"cohere": {"outputDimension": 256}}))) + .await + .unwrap(); + assert_eq!( + fetch.seen()[0].json_body(), + json!({ + "embedding_types": ["float"], + "input_type": "search_query", + "model": "embed-v4.0", + "output_dimension": 256, + "texts": ["sunny day at the beach", "rainy day in the city"] + }) + ); + } +} diff --git a/aimux-providers/src/cohere/tests/fixtures/cohere-citations.json b/aimux-providers/src/cohere/tests/fixtures/cohere-citations.json new file mode 100644 index 000000000..c9a10134d --- /dev/null +++ b/aimux-providers/src/cohere/tests/fixtures/cohere-citations.json @@ -0,0 +1,77 @@ +{ + "id": "68475c80-574b-4c65-98a4-e81cebab5dce", + "message": { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "The key benefits mentioned in this document are:\n1. Automation of tasks\n2. Better decision-making\n3. Cost reduction" + } + ], + "citations": [ + { + "start": 52, + "end": 71, + "text": "Automation of tasks", + "sources": [ + { + "type": "document", + "id": "doc:0", + "document": { + "id": "doc:0", + "text": "AI provides: 1. Automation of tasks 2. Better decision-making 3. Cost reduction", + "title": "benefits.txt" + } + } + ], + "type": "TEXT_CONTENT" + }, + { + "start": 75, + "end": 97, + "text": "Better decision-making", + "sources": [ + { + "type": "document", + "id": "doc:0", + "document": { + "id": "doc:0", + "text": "AI provides: 1. Automation of tasks 2. Better decision-making 3. Cost reduction", + "title": "benefits.txt" + } + } + ], + "type": "TEXT_CONTENT" + }, + { + "start": 101, + "end": 115, + "text": "Cost reduction", + "sources": [ + { + "type": "document", + "id": "doc:0", + "document": { + "id": "doc:0", + "text": "AI provides: 1. Automation of tasks 2. Better decision-making 3. Cost reduction", + "title": "benefits.txt" + } + } + ], + "type": "TEXT_CONTENT" + } + ] + }, + "finish_reason": "COMPLETE", + "usage": { + "billed_units": { + "input_tokens": 39, + "output_tokens": 27 + }, + "tokens": { + "input_tokens": 1683, + "output_tokens": 62 + }, + "cached_tokens": 992 + } +} diff --git a/aimux-providers/src/cohere/tests/fixtures/cohere-embedding.json b/aimux-providers/src/cohere/tests/fixtures/cohere-embedding.json new file mode 100644 index 000000000..cea6fee98 --- /dev/null +++ b/aimux-providers/src/cohere/tests/fixtures/cohere-embedding.json @@ -0,0 +1,19 @@ +{ + "id": "f5aa3e7b-f011-4c5c-a825-f94669f760e5", + "texts": ["sunny day at the beach", "rainy day in the city"], + "embeddings": { + "float": [ + [0.03302002, 0.020904541, -0.019744873, -0.0625, 0.04437256], + [-0.04660034, 0.00037765503, -0.061157227, -0.08239746, -0.010360718] + ] + }, + "meta": { + "api_version": { + "version": "2" + }, + "billed_units": { + "input_tokens": 10 + } + }, + "response_type": "embeddings_by_type" +} diff --git a/aimux-providers/src/cohere/tests/fixtures/cohere-empty-tool-call.chunks.txt b/aimux-providers/src/cohere/tests/fixtures/cohere-empty-tool-call.chunks.txt new file mode 100644 index 000000000..d33339349 --- /dev/null +++ b/aimux-providers/src/cohere/tests/fixtures/cohere-empty-tool-call.chunks.txt @@ -0,0 +1,16 @@ +{"id":"66dec7d7-45e6-427c-8fd9-7d6375d12046","type":"message-start","delta":{"message":{"role":"assistant","content":[],"tool_plan":"","tool_calls":[],"citations":[]}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":"I"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" will"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" use"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" the"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" currentTime"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" tool"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" to"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" find"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" the"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" current"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" time"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":"."}}} +{"type":"tool-call-start","index":0,"delta":{"message":{"tool_calls":{"id":"currentTime_y46ar19t5gvw","type":"function","function":{"name":"currentTime","arguments":""}}}}} +{"type":"tool-call-end","index":0} +{"type":"message-end","delta":{"finish_reason":"TOOL_CALL","usage":{"billed_units":{"input_tokens":46,"output_tokens":14},"tokens":{"input_tokens":1445,"output_tokens":43},"cached_tokens":704}}} \ No newline at end of file diff --git a/aimux-providers/src/cohere/tests/fixtures/cohere-max-tokens.json b/aimux-providers/src/cohere/tests/fixtures/cohere-max-tokens.json new file mode 100644 index 000000000..4201b0334 --- /dev/null +++ b/aimux-providers/src/cohere/tests/fixtures/cohere-max-tokens.json @@ -0,0 +1,24 @@ +{ + "id": "039584d9-7236-4ecb-9dd7-f1bed57888bc", + "message": { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "**The History of" + } + ] + }, + "finish_reason": "MAX_TOKENS", + "usage": { + "billed_units": { + "input_tokens": 11, + "output_tokens": 4 + }, + "tokens": { + "input_tokens": 506, + "output_tokens": 5 + }, + "cached_tokens": 448 + } +} diff --git a/aimux-providers/src/cohere/tests/fixtures/cohere-null-args.json b/aimux-providers/src/cohere/tests/fixtures/cohere-null-args.json new file mode 100644 index 000000000..b30b75a7f --- /dev/null +++ b/aimux-providers/src/cohere/tests/fixtures/cohere-null-args.json @@ -0,0 +1,29 @@ +{ + "id": "316f0604-ff50-49f6-ba38-c64616e972b4", + "message": { + "role": "assistant", + "tool_plan": "I will use the currentTime tool to find the current time.", + "tool_calls": [ + { + "id": "currentTime_tf4dywn8wgnk", + "type": "function", + "function": { + "name": "currentTime", + "arguments": "null" + } + } + ] + }, + "finish_reason": "TOOL_CALL", + "usage": { + "billed_units": { + "input_tokens": 46, + "output_tokens": 14 + }, + "tokens": { + "input_tokens": 1445, + "output_tokens": 43 + }, + "cached_tokens": 992 + } +} diff --git a/aimux-providers/src/cohere/tests/fixtures/cohere-reasoning.chunks.txt b/aimux-providers/src/cohere/tests/fixtures/cohere-reasoning.chunks.txt new file mode 100644 index 000000000..91c788a92 --- /dev/null +++ b/aimux-providers/src/cohere/tests/fixtures/cohere-reasoning.chunks.txt @@ -0,0 +1,51 @@ +{"id":"c9117d7f-a7e4-499f-b643-a2a1e139687b","type":"message-start","delta":{"message":{"role":"assistant","content":[],"tool_plan":"","tool_calls":[],"citations":[]}}} +{"type":"content-start","index":0,"delta":{"message":{"content":{"type":"thinking","thinking":""}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":"The"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" user"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" is"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" asking"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" for"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" the"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" sum"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" of"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" 2"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" and"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" 2"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":"."}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" Since"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" this"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" is"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" a"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" straightforward"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" arithmetic"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" problem"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":","}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" I"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" don"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":"'t"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" need"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" to"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" use"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" any"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" tools"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":"."}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" I"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" can"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" calculate"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" the"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" answer"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":" directly"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"thinking":"."}}}} +{"type":"content-end","index":0} +{"type":"content-start","index":1,"delta":{"message":{"content":{"type":"text","text":""}}}} +{"type":"content-delta","index":1,"delta":{"message":{"content":{"text":"The"}}}} +{"type":"content-delta","index":1,"delta":{"message":{"content":{"text":" answer"}}}} +{"type":"content-delta","index":1,"delta":{"message":{"content":{"text":" to"}}}} +{"type":"content-delta","index":1,"delta":{"message":{"content":{"text":" 2"}}}} +{"type":"content-delta","index":1,"delta":{"message":{"content":{"text":" +"}}}} +{"type":"content-delta","index":1,"delta":{"message":{"content":{"text":" 2"}}}} +{"type":"content-delta","index":1,"delta":{"message":{"content":{"text":" is"}}}} +{"type":"content-delta","index":1,"delta":{"message":{"content":{"text":" 4"}}}} +{"type":"content-delta","index":1,"delta":{"message":{"content":{"text":"."}}}} +{"type":"content-end","index":1} +{"type":"message-end","delta":{"finish_reason":"COMPLETE","usage":{"billed_units":{"input_tokens":8,"output_tokens":50},"tokens":{"input_tokens":1394,"output_tokens":54},"cached_tokens":1360}}} \ No newline at end of file diff --git a/aimux-providers/src/cohere/tests/fixtures/cohere-reasoning.json b/aimux-providers/src/cohere/tests/fixtures/cohere-reasoning.json new file mode 100644 index 000000000..9b7986515 --- /dev/null +++ b/aimux-providers/src/cohere/tests/fixtures/cohere-reasoning.json @@ -0,0 +1,28 @@ +{ + "id": "53bcb235-5179-4a91-a578-cb372b5430bc", + "message": { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "Okay, so I need to figure out what 2 + 2 is. Let me start by recalling what addition means." + }, + { + "type": "text", + "text": "2 + 2 = 4" + } + ] + }, + "finish_reason": "COMPLETE", + "usage": { + "billed_units": { + "input_tokens": 8, + "output_tokens": 578 + }, + "tokens": { + "input_tokens": 1394, + "output_tokens": 582 + }, + "cached_tokens": 1344 + } +} diff --git a/aimux-providers/src/cohere/tests/fixtures/cohere-text.chunks.txt b/aimux-providers/src/cohere/tests/fixtures/cohere-text.chunks.txt new file mode 100644 index 000000000..0272b3995 --- /dev/null +++ b/aimux-providers/src/cohere/tests/fixtures/cohere-text.chunks.txt @@ -0,0 +1,11 @@ +{"id":"321d178c-2c12-44d3-ae42-2f5510f6b1cc","type":"message-start","delta":{"message":{"role":"assistant","content":[],"tool_plan":"","tool_calls":[],"citations":[]}}} +{"type":"content-start","index":0,"delta":{"message":{"content":{"type":"text","text":""}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"text":"The"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"text":" capital"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"text":" of"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"text":" France"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"text":" is"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"text":" Paris"}}}} +{"type":"content-delta","index":0,"delta":{"message":{"content":{"text":"."}}}} +{"type":"content-end","index":0} +{"type":"message-end","delta":{"finish_reason":"COMPLETE","usage":{"billed_units":{"input_tokens":12,"output_tokens":7},"tokens":{"input_tokens":507,"output_tokens":10},"cached_tokens":448}}} \ No newline at end of file diff --git a/aimux-providers/src/cohere/tests/fixtures/cohere-text.json b/aimux-providers/src/cohere/tests/fixtures/cohere-text.json new file mode 100644 index 000000000..fd3285bcd --- /dev/null +++ b/aimux-providers/src/cohere/tests/fixtures/cohere-text.json @@ -0,0 +1,24 @@ +{ + "id": "e7592632-1e3d-424f-b129-bd5f9f980f7b", + "message": { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "The capital of France is Paris." + } + ] + }, + "finish_reason": "COMPLETE", + "usage": { + "billed_units": { + "input_tokens": 12, + "output_tokens": 7 + }, + "tokens": { + "input_tokens": 507, + "output_tokens": 10 + }, + "cached_tokens": 448 + } +} diff --git a/aimux-providers/src/cohere/tests/fixtures/cohere-tool-call.chunks.txt b/aimux-providers/src/cohere/tests/fixtures/cohere-tool-call.chunks.txt new file mode 100644 index 000000000..8366d55e3 --- /dev/null +++ b/aimux-providers/src/cohere/tests/fixtures/cohere-tool-call.chunks.txt @@ -0,0 +1,47 @@ +{"id":"2941521a-b87a-45f6-9b0d-235fd66c3025","type":"message-start","delta":{"message":{"role":"assistant","content":[],"tool_plan":"","tool_calls":[],"citations":[]}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":"I"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" will"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" use"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" the"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" weather"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" tool"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" to"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" find"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" the"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" weather"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" in"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" San"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" Francisco"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" and"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" the"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" city"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":"At"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":"tract"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":"ions"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" tool"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" to"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" find"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" attractions"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" in"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" San"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":" Francisco"}}} +{"type":"tool-plan-delta","delta":{"message":{"tool_plan":"."}}} +{"type":"tool-call-start","index":0,"delta":{"message":{"tool_calls":{"id":"weather_e8p4pn45zt0t","type":"function","function":{"name":"weather","arguments":""}}}}} +{"type":"tool-call-delta","index":0,"delta":{"message":{"tool_calls":{"function":{"arguments":"{\""}}}}} +{"type":"tool-call-delta","index":0,"delta":{"message":{"tool_calls":{"function":{"arguments":"location"}}}}} +{"type":"tool-call-delta","index":0,"delta":{"message":{"tool_calls":{"function":{"arguments":"\":"}}}}} +{"type":"tool-call-delta","index":0,"delta":{"message":{"tool_calls":{"function":{"arguments":" \""}}}}} +{"type":"tool-call-delta","index":0,"delta":{"message":{"tool_calls":{"function":{"arguments":"San"}}}}} +{"type":"tool-call-delta","index":0,"delta":{"message":{"tool_calls":{"function":{"arguments":" Francisco"}}}}} +{"type":"tool-call-delta","index":0,"delta":{"message":{"tool_calls":{"function":{"arguments":"\"}"}}}}} +{"type":"tool-call-end","index":0} +{"type":"tool-call-start","index":1,"delta":{"message":{"tool_calls":{"id":"cityAttractions_pyxssbwnq9fq","type":"function","function":{"name":"cityAttractions","arguments":""}}}}} +{"type":"tool-call-delta","index":1,"delta":{"message":{"tool_calls":{"function":{"arguments":"{\""}}}}} +{"type":"tool-call-delta","index":1,"delta":{"message":{"tool_calls":{"function":{"arguments":"city"}}}}} +{"type":"tool-call-delta","index":1,"delta":{"message":{"tool_calls":{"function":{"arguments":"\":"}}}}} +{"type":"tool-call-delta","index":1,"delta":{"message":{"tool_calls":{"function":{"arguments":" \""}}}}} +{"type":"tool-call-delta","index":1,"delta":{"message":{"tool_calls":{"function":{"arguments":"San"}}}}} +{"type":"tool-call-delta","index":1,"delta":{"message":{"tool_calls":{"function":{"arguments":" Francisco"}}}}} +{"type":"tool-call-delta","index":1,"delta":{"message":{"tool_calls":{"function":{"arguments":"\"}"}}}}} +{"type":"tool-call-end","index":1} +{"type":"message-end","delta":{"finish_reason":"TOOL_CALL","usage":{"billed_units":{"input_tokens":119,"output_tokens":44},"tokens":{"input_tokens":1549,"output_tokens":95},"cached_tokens":1504}}} \ No newline at end of file diff --git a/aimux-providers/src/cohere/tests/fixtures/cohere-tool-call.json b/aimux-providers/src/cohere/tests/fixtures/cohere-tool-call.json new file mode 100644 index 000000000..21e302934 --- /dev/null +++ b/aimux-providers/src/cohere/tests/fixtures/cohere-tool-call.json @@ -0,0 +1,37 @@ +{ + "id": "f201af17-e24a-4396-8f6a-98e8bf9c3432", + "message": { + "role": "assistant", + "tool_plan": "I will use the weather tool to find out the weather in San Francisco. I will also use the cityAttractions tool to find out what attractions are in San Francisco.", + "tool_calls": [ + { + "id": "weather_dqgshstja6p9", + "type": "function", + "function": { + "name": "weather", + "arguments": "{\"location\":\"San Francisco\"}" + } + }, + { + "id": "cityAttractions_dcxfx4myvx68", + "type": "function", + "function": { + "name": "cityAttractions", + "arguments": "{\"city\":\"San Francisco\"}" + } + } + ] + }, + "finish_reason": "TOOL_CALL", + "usage": { + "billed_units": { + "input_tokens": 119, + "output_tokens": 52 + }, + "tokens": { + "input_tokens": 1549, + "output_tokens": 103 + }, + "cached_tokens": 992 + } +} diff --git a/aimux-providers/src/cohere/tests/mod.rs b/aimux-providers/src/cohere/tests/mod.rs new file mode 100644 index 000000000..29e4218fc --- /dev/null +++ b/aimux-providers/src/cohere/tests/mod.rs @@ -0,0 +1,7 @@ +//! Ports of the `@ai-sdk/cohere@4.0.52` unit tests, one module per upstream test file. + +mod chat_language_model; +mod convert_to_chat_prompt; +mod embedding_model; +mod prepare_tools; +mod reranking; diff --git a/aimux-providers/src/cohere/tests/prepare_tools.rs b/aimux-providers/src/cohere/tests/prepare_tools.rs new file mode 100644 index 000000000..184e5baa0 --- /dev/null +++ b/aimux-providers/src/cohere/tests/prepare_tools.rs @@ -0,0 +1,153 @@ +//! Ports of `src/cohere-prepare-tools.test.ts` from `@ai-sdk/cohere@4.0.52`. +//! +//! One nested module per `describe`, one `#[test]` per `it`. Each test carries +//! its upstream chain in an `upstream:` marker line. + +use crate::cohere::convert::prepare_tools; +use aimux_core::options::LanguageModelV4ToolChoice; +use aimux_core::tool::Tool; +use aimux_core::tool::{FunctionTool, ProviderTool}; +use aimux_core::types::NamespacedKind; +use serde_json::{Value, json}; + +/// upstream: @ai-sdk/cohere src/cohere-prepare-tools.test.ts › should return undefined tools when no tools are provided +#[test] +fn should_return_undefined_tools_when_no_tools_are_provided() { + let result = prepare_tools(&Some(Vec::new()), None); + + assert_eq!(result.tools, None); + assert_eq!(result.tool_choice, None); + assert!(result.tool_warnings.is_empty()); +} + +/// upstream: @ai-sdk/cohere src/cohere-prepare-tools.test.ts › should process function tools correctly +#[test] +fn should_process_function_tools_correctly() { + let function_tool = Tool::Function( + FunctionTool::new( + "testFunction", + json!({ "type": "object", "properties": {} }), + ) + .with_description("test description"), + ); + + let result = prepare_tools(&Some(vec![function_tool]), None); + + assert_eq!( + result.tools, + Some(vec![json!({ + "type": "function", + "function": { + "name": "testFunction", + "description": "test description", + "parameters": { "type": "object", "properties": {} }, + }, + })]) + ); + assert_eq!(result.tool_choice, None); + assert!(result.tool_warnings.is_empty()); +} + +/// upstream: @ai-sdk/cohere src/cohere-prepare-tools.test.ts › should add warnings for provider-defined tools +#[test] +fn should_add_warnings_for_provider_defined_tools() { + let provider_tool = Tool::Provider(ProviderTool { + id: NamespacedKind::new("provider.tool").unwrap(), + name: "tool".into(), + args: serde_json::Map::new(), + supports_deferred_results: None, + }); + + let result = prepare_tools(&Some(vec![provider_tool]), None); + + // Inline snapshot: `toolChoice: undefined`, `tools: []`, and one + // unsupported warning with no `details`. + assert_eq!(result.tools, Some(Vec::new())); + assert_eq!(result.tool_choice, None); + assert_eq!( + serde_json::to_value(&result.tool_warnings).unwrap(), + json!([ + { + "type": "unsupported", + "feature": "provider-defined tool provider.tool", + }, + ]) + ); +} + +mod tool_choice_handling { + use super::*; + + fn basic_tool() -> Tool { + Tool::Function( + FunctionTool::new( + "testFunction", + json!({ "type": "object", "properties": {} }), + ) + .with_description("test description"), + ) + } + + fn basic_tool_json() -> Value { + json!({ + "type": "function", + "function": { + "name": "testFunction", + "description": "test description", + "parameters": { "type": "object", "properties": {} }, + }, + }) + } + + /// upstream: @ai-sdk/cohere src/cohere-prepare-tools.test.ts › tool choice handling › should handle auto tool choice + #[test] + fn should_handle_auto_tool_choice() { + let result = prepare_tools( + &Some(vec![basic_tool()]), + Some(&LanguageModelV4ToolChoice::Auto), + ); + + assert_eq!(result.tool_choice, None); + } + + /// upstream: @ai-sdk/cohere src/cohere-prepare-tools.test.ts › tool choice handling › should handle none tool choice + #[test] + fn should_handle_none_tool_choice() { + let result = prepare_tools( + &Some(vec![basic_tool()]), + Some(&LanguageModelV4ToolChoice::None), + ); + + assert_eq!(result.tools, Some(vec![basic_tool_json()])); + assert_eq!(result.tool_choice, Some(json!("NONE"))); + assert!(result.tool_warnings.is_empty()); + } + + /// upstream: @ai-sdk/cohere src/cohere-prepare-tools.test.ts › tool choice handling › should handle required tool choice + #[test] + fn should_handle_required_tool_choice() { + let result = prepare_tools( + &Some(vec![basic_tool()]), + Some(&LanguageModelV4ToolChoice::Required), + ); + + assert_eq!(result.tools, Some(vec![basic_tool_json()])); + assert_eq!(result.tool_choice, Some(json!("REQUIRED"))); + assert!(result.tool_warnings.is_empty()); + } + + /// upstream: @ai-sdk/cohere src/cohere-prepare-tools.test.ts › tool choice handling › should handle tool type tool choice by filtering tools + #[test] + fn should_handle_tool_type_tool_choice_by_filtering_tools() { + let result = prepare_tools( + &Some(vec![basic_tool()]), + Some(&LanguageModelV4ToolChoice::Tool { + tool_name: "testFunction".into(), + }), + ); + + assert_eq!(result.tools, Some(vec![basic_tool_json()])); + assert_eq!(result.tool_choice, Some(json!("REQUIRED"))); + assert!(result.tool_warnings.is_empty()); + } +} diff --git a/aimux-providers/src/cohere/tests/reranking/fixtures/cohere-reranking.1.json b/aimux-providers/src/cohere/tests/reranking/fixtures/cohere-reranking.1.json new file mode 100644 index 000000000..cd520c1de --- /dev/null +++ b/aimux-providers/src/cohere/tests/reranking/fixtures/cohere-reranking.1.json @@ -0,0 +1,21 @@ +{ + "id": "b44fe75b-e3d3-489a-b61e-1a1aede3ef72", + "results": [ + { + "index": 1, + "relevance_score": 0.10183054 + }, + { + "index": 0, + "relevance_score": 0.03762639 + } + ], + "meta": { + "api_version": { + "version": "2" + }, + "billed_units": { + "search_units": 1 + } + } +} diff --git a/aimux-providers/src/cohere/tests/reranking/mod.rs b/aimux-providers/src/cohere/tests/reranking/mod.rs new file mode 100644 index 000000000..c556c2ff7 --- /dev/null +++ b/aimux-providers/src/cohere/tests/reranking/mod.rs @@ -0,0 +1,3 @@ +//! Upstream `src/reranking/` tests of `@ai-sdk/cohere@4.0.52`. + +mod reranking_model; diff --git a/aimux-providers/src/cohere/tests/reranking/reranking_model.rs b/aimux-providers/src/cohere/tests/reranking/reranking_model.rs new file mode 100644 index 000000000..531a8ec60 --- /dev/null +++ b/aimux-providers/src/cohere/tests/reranking/reranking_model.rs @@ -0,0 +1,239 @@ +//! Ports of `src/reranking/cohere-reranking-model.test.ts` from +//! `@ai-sdk/cohere@4.0.52`. +//! +//! Only the cases that `tests/cohere_upstream_e2e_test.rs` does not already +//! cover are here: the first request case of each describe is ported there. +//! +//! Harness differences from the TypeScript (expectations are unchanged): +//! - `content-length` is the byte length of the compact JSON body, as the +//! TypeScript test server sends it. +//! - The mock records `user-agent`, which the TypeScript `requestHeaders` do +//! not list; the header comparison drops it. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use crate::cohere::{CohereProvider, CohereProviderSettings, create_cohere}; +use aimux_core::reranking_model::{ + RerankingCallOptions, RerankingDocuments, RerankingModel, RerankingResult, +}; +use aimux_core::shared::{SharedProviderOptions, Warning}; +use serde_json::{Value, json}; + +use crate::test_support::mock_fetch::{Canned, MockFetch}; + +const RERANKING_FIXTURE: &str = include_str!("fixtures/cohere-reranking.1.json"); + +fn provider(fetch: &Arc) -> CohereProvider { + create_cohere(CohereProviderSettings { + api_key: Some("test-api-key".into()), + fetch: Some(fetch.transport()), + ..Default::default() + }) + .unwrap() +} + +/// A single-response mock for the reranking fixture. +fn reranking_fetch() -> Arc { + let fixture: Value = serde_json::from_str(RERANKING_FIXTURE).expect("fixture is JSON"); + let bytes = serde_json::to_vec(&fixture).expect("serializable"); + MockFetch::new(vec![Canned { + status: 200, + headers: vec![ + ("content-type".into(), "application/json".into()), + ("content-length".into(), bytes.len().to_string()), + ], + body: bytes, + }]) +} + +/// `doRerank` with the TypeScript call options of both describes. +async fn rerank(documents: RerankingDocuments) -> (Arc, RerankingResult) { + let fetch = reranking_fetch(); + let mut options = RerankingCallOptions::new("rainy day", documents); + options.top_n = Some(2); + options.provider_options = Some( + serde_json::from_value::( + json!({"cohere": {"maxTokensPerDoc": 1000, "priority": 1}}), + ) + .unwrap(), + ); + let result = provider(&fetch) + .reranking("rerank-english-v3.0") + .do_rerank(&options) + .await + .unwrap(); + (fetch, result) +} + +fn json_documents() -> RerankingDocuments { + RerankingDocuments::Object { + values: vec![ + json!({"example": "sunny day at the beach"}), + json!({"example": "rainy day in the city"}), + ], + } +} + +fn text_documents() -> RerankingDocuments { + RerankingDocuments::Text { + values: vec![ + "sunny day at the beach".to_string(), + "rainy day in the city".to_string(), + ], + } +} + +fn map(pairs: &[(&str, &str)]) -> BTreeMap { + pairs + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() +} + +/// The request headers the TypeScript test checks with `toMatchInlineSnapshot`. +/// The user agent is recorded by the mock and not compared here. +fn assert_sent_headers(fetch: &MockFetch, expected: &[(&str, &str)]) { + let mut headers = fetch.seen()[0].headers.clone(); + headers.remove("user-agent").expect("user-agent is sent"); + assert_eq!(headers, map(expected)); +} + +/// The `response` of the result, as the TypeScript snapshot lists it. +fn expected_response() -> Value { + json!({ + "body": { + "id": "b44fe75b-e3d3-489a-b61e-1a1aede3ef72", + "meta": { + "api_version": {"version": "2"}, + "billed_units": {"search_units": 1} + }, + "results": [ + {"index": 1, "relevance_score": 0.10183054}, + {"index": 0, "relevance_score": 0.03762639} + ] + }, + "headers": { + "content-length": "212", + "content-type": "application/json" + }, + "id": "b44fe75b-e3d3-489a-b61e-1a1aede3ef72" + }) +} + +fn expected_ranking() -> Value { + json!([ + {"index": 1, "relevanceScore": 0.10183054}, + {"index": 0, "relevanceScore": 0.03762639} + ]) +} + +mod do_rerank { + mod json_documents { + use super::super::*; + + /// upstream: @ai-sdk/cohere src/reranking/cohere-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 (fetch, _) = rerank(json_documents()).await; + assert_sent_headers( + &fetch, + &[ + ("authorization", "Bearer test-api-key"), + ("content-type", "application/json"), + ], + ); + } + + /// upstream: @ai-sdk/cohere src/reranking/cohere-reranking-model.test.ts › doRerank › json documents › should return result with warnings + #[tokio::test] + async fn should_return_result_with_warnings() { + let (_, result) = rerank(json_documents()).await; + let warnings = result.warnings.expect("warnings are present"); + assert!(matches!( + warnings.as_slice(), + [Warning::Compatibility { feature, details }] + if feature == "object documents" + && details.as_deref() == Some("Object documents are converted to strings.") + )); + } + + /// upstream: @ai-sdk/cohere src/reranking/cohere-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) = rerank(json_documents()).await; + assert_eq!( + serde_json::to_value(&result.ranking).unwrap(), + expected_ranking() + ); + } + + /// upstream: @ai-sdk/cohere src/reranking/cohere-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) = rerank(json_documents()).await; + assert!(result.provider_metadata.is_none()); + } + + /// upstream: @ai-sdk/cohere src/reranking/cohere-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) = rerank(json_documents()).await; + assert_eq!( + serde_json::to_value(&result.response).unwrap(), + expected_response() + ); + } + } + + mod text_documents { + use super::super::*; + + /// upstream: @ai-sdk/cohere src/reranking/cohere-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 (fetch, _) = rerank(text_documents()).await; + assert_sent_headers( + &fetch, + &[ + ("authorization", "Bearer test-api-key"), + ("content-type", "application/json"), + ], + ); + } + + /// upstream: @ai-sdk/cohere src/reranking/cohere-reranking-model.test.ts › doRerank › text documents › should return result without warnings + #[tokio::test] + async fn should_return_result_without_warnings() { + let (_, result) = rerank(text_documents()).await; + assert_eq!(serde_json::to_value(&result.warnings).unwrap(), json!([])); + } + + /// upstream: @ai-sdk/cohere src/reranking/cohere-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) = rerank(text_documents()).await; + assert_eq!( + serde_json::to_value(&result.ranking).unwrap(), + expected_ranking() + ); + } + + /// upstream: @ai-sdk/cohere src/reranking/cohere-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) = rerank(text_documents()).await; + assert!(result.provider_metadata.is_none()); + } + + /// upstream: @ai-sdk/cohere src/reranking/cohere-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) = rerank(text_documents()).await; + assert_eq!( + serde_json::to_value(&result.response).unwrap(), + expected_response() + ); + } + } +} diff --git a/aimux-providers/src/cohere/tests/snapshots/do_embed_raw_response.json b/aimux-providers/src/cohere/tests/snapshots/do_embed_raw_response.json new file mode 100644 index 000000000..bea118d37 --- /dev/null +++ b/aimux-providers/src/cohere/tests/snapshots/do_embed_raw_response.json @@ -0,0 +1,41 @@ +{ + "body": { + "embeddings": { + "float": [ + [ + 0.03302002, + 0.020904541, + -0.019744873, + -0.0625, + 0.04437256 + ], + [ + -0.04660034, + 0.00037765503, + -0.061157227, + -0.08239746, + -0.010360718 + ] + ] + }, + "id": "f5aa3e7b-f011-4c5c-a825-f94669f760e5", + "meta": { + "api_version": { + "version": "2" + }, + "billed_units": { + "input_tokens": 10 + } + }, + "response_type": "embeddings_by_type", + "texts": [ + "sunny day at the beach", + "rainy day in the city" + ] + }, + "headers": { + "content-length": "363", + "content-type": "application/json", + "test-header": "test-value" + } +} diff --git a/aimux-providers/src/cohere/tests/snapshots/do_generate_citations.json b/aimux-providers/src/cohere/tests/snapshots/do_generate_citations.json new file mode 100644 index 000000000..d61592389 --- /dev/null +++ b/aimux-providers/src/cohere/tests/snapshots/do_generate_citations.json @@ -0,0 +1,214 @@ +{ + "content": [ + { + "text": "The key benefits mentioned in this document are:\n1. Automation of tasks\n2. Better decision-making\n3. Cost reduction", + "type": "text" + }, + { + "id": "test-citation-id", + "mediaType": "text/plain", + "providerMetadata": { + "cohere": { + "citationType": "TEXT_CONTENT", + "end": 71, + "sources": [ + { + "document": { + "id": "doc:0", + "text": "AI provides: 1. Automation of tasks 2. Better decision-making 3. Cost reduction", + "title": "benefits.txt" + }, + "id": "doc:0", + "type": "document" + } + ], + "start": 52, + "text": "Automation of tasks" + } + }, + "sourceType": "document", + "title": "benefits.txt", + "type": "source" + }, + { + "id": "test-citation-id", + "mediaType": "text/plain", + "providerMetadata": { + "cohere": { + "citationType": "TEXT_CONTENT", + "end": 97, + "sources": [ + { + "document": { + "id": "doc:0", + "text": "AI provides: 1. Automation of tasks 2. Better decision-making 3. Cost reduction", + "title": "benefits.txt" + }, + "id": "doc:0", + "type": "document" + } + ], + "start": 75, + "text": "Better decision-making" + } + }, + "sourceType": "document", + "title": "benefits.txt", + "type": "source" + }, + { + "id": "test-citation-id", + "mediaType": "text/plain", + "providerMetadata": { + "cohere": { + "citationType": "TEXT_CONTENT", + "end": 115, + "sources": [ + { + "document": { + "id": "doc:0", + "text": "AI provides: 1. Automation of tasks 2. Better decision-making 3. Cost reduction", + "title": "benefits.txt" + }, + "id": "doc:0", + "type": "document" + } + ], + "start": 101, + "text": "Cost reduction" + } + }, + "sourceType": "document", + "title": "benefits.txt", + "type": "source" + } + ], + "finishReason": { + "raw": "COMPLETE", + "unified": "stop" + }, + "request": { + "body": { + "documents": [ + { + "data": { + "text": "AI provides automation and efficiency.", + "title": "ai-benefits.txt" + } + } + ], + "messages": [ + { + "content": "What are AI benefits?", + "role": "user" + } + ], + "model": "command-r-plus" + } + }, + "response": { + "body": { + "finish_reason": "COMPLETE", + "id": "68475c80-574b-4c65-98a4-e81cebab5dce", + "message": { + "citations": [ + { + "end": 71, + "sources": [ + { + "document": { + "id": "doc:0", + "text": "AI provides: 1. Automation of tasks 2. Better decision-making 3. Cost reduction", + "title": "benefits.txt" + }, + "id": "doc:0", + "type": "document" + } + ], + "start": 52, + "text": "Automation of tasks", + "type": "TEXT_CONTENT" + }, + { + "end": 97, + "sources": [ + { + "document": { + "id": "doc:0", + "text": "AI provides: 1. Automation of tasks 2. Better decision-making 3. Cost reduction", + "title": "benefits.txt" + }, + "id": "doc:0", + "type": "document" + } + ], + "start": 75, + "text": "Better decision-making", + "type": "TEXT_CONTENT" + }, + { + "end": 115, + "sources": [ + { + "document": { + "id": "doc:0", + "text": "AI provides: 1. Automation of tasks 2. Better decision-making 3. Cost reduction", + "title": "benefits.txt" + }, + "id": "doc:0", + "type": "document" + } + ], + "start": 101, + "text": "Cost reduction", + "type": "TEXT_CONTENT" + } + ], + "content": [ + { + "text": "The key benefits mentioned in this document are:\n1. Automation of tasks\n2. Better decision-making\n3. Cost reduction", + "type": "text" + } + ], + "role": "assistant" + }, + "usage": { + "billed_units": { + "input_tokens": 39, + "output_tokens": 27 + }, + "cached_tokens": 992, + "tokens": { + "input_tokens": 1683, + "output_tokens": 62 + } + } + }, + "headers": { + "content-length": "1175", + "content-type": "application/json" + } + }, + "usage": { + "inputTokens": { + "noCache": 1683, + "total": 1683 + }, + "outputTokens": { + "text": 62, + "total": 62 + }, + "raw": { + "billed_units": { + "input_tokens": 39, + "output_tokens": 27 + }, + "cached_tokens": 992, + "tokens": { + "input_tokens": 1683, + "output_tokens": 62 + } + } + }, + "warnings": [] +} diff --git a/aimux-providers/src/cohere/tests/snapshots/do_generate_null_tool_call_args.json b/aimux-providers/src/cohere/tests/snapshots/do_generate_null_tool_call_args.json new file mode 100644 index 000000000..a8d3524b2 --- /dev/null +++ b/aimux-providers/src/cohere/tests/snapshots/do_generate_null_tool_call_args.json @@ -0,0 +1,8 @@ +[ + { + "input": "{}", + "toolCallId": "currentTime_tf4dywn8wgnk", + "toolName": "currentTime", + "type": "tool-call" + } +] diff --git a/aimux-providers/src/cohere/tests/snapshots/do_generate_reasoning.json b/aimux-providers/src/cohere/tests/snapshots/do_generate_reasoning.json new file mode 100644 index 000000000..4f60eb85c --- /dev/null +++ b/aimux-providers/src/cohere/tests/snapshots/do_generate_reasoning.json @@ -0,0 +1,87 @@ +{ + "content": [ + { + "text": "Okay, so I need to figure out what 2 + 2 is. Let me start by recalling what addition means.", + "type": "reasoning" + }, + { + "text": "2 + 2 = 4", + "type": "text" + } + ], + "finishReason": { + "raw": "COMPLETE", + "unified": "stop" + }, + "request": { + "body": { + "messages": [ + { + "content": "you are a friendly bot!", + "role": "system" + }, + { + "content": "Hello", + "role": "user" + } + ], + "model": "command-r-plus" + } + }, + "response": { + "body": { + "finish_reason": "COMPLETE", + "id": "53bcb235-5179-4a91-a578-cb372b5430bc", + "message": { + "content": [ + { + "thinking": "Okay, so I need to figure out what 2 + 2 is. Let me start by recalling what addition means.", + "type": "thinking" + }, + { + "text": "2 + 2 = 4", + "type": "text" + } + ], + "role": "assistant" + }, + "usage": { + "billed_units": { + "input_tokens": 8, + "output_tokens": 578 + }, + "cached_tokens": 1344, + "tokens": { + "input_tokens": 1394, + "output_tokens": 582 + } + } + }, + "headers": { + "content-length": "411", + "content-type": "application/json" + } + }, + "usage": { + "inputTokens": { + "noCache": 1394, + "total": 1394 + }, + "outputTokens": { + "text": 582, + "total": 582 + }, + "raw": { + "billed_units": { + "input_tokens": 8, + "output_tokens": 578 + }, + "cached_tokens": 1344, + "tokens": { + "input_tokens": 1394, + "output_tokens": 582 + } + } + }, + "warnings": [] +} diff --git a/aimux-providers/src/cohere/tests/snapshots/do_generate_text.json b/aimux-providers/src/cohere/tests/snapshots/do_generate_text.json new file mode 100644 index 000000000..75d78b25d --- /dev/null +++ b/aimux-providers/src/cohere/tests/snapshots/do_generate_text.json @@ -0,0 +1,79 @@ +{ + "content": [ + { + "text": "The capital of France is Paris.", + "type": "text" + } + ], + "finishReason": { + "raw": "COMPLETE", + "unified": "stop" + }, + "request": { + "body": { + "messages": [ + { + "content": "you are a friendly bot!", + "role": "system" + }, + { + "content": "Hello", + "role": "user" + } + ], + "model": "command-r-plus" + } + }, + "response": { + "body": { + "finish_reason": "COMPLETE", + "id": "e7592632-1e3d-424f-b129-bd5f9f980f7b", + "message": { + "content": [ + { + "text": "The capital of France is Paris.", + "type": "text" + } + ], + "role": "assistant" + }, + "usage": { + "billed_units": { + "input_tokens": 12, + "output_tokens": 7 + }, + "cached_tokens": 448, + "tokens": { + "input_tokens": 507, + "output_tokens": 10 + } + } + }, + "headers": { + "content-length": "304", + "content-type": "application/json" + } + }, + "usage": { + "inputTokens": { + "noCache": 507, + "total": 507 + }, + "outputTokens": { + "text": 10, + "total": 10 + }, + "raw": { + "billed_units": { + "input_tokens": 12, + "output_tokens": 7 + }, + "cached_tokens": 448, + "tokens": { + "input_tokens": 507, + "output_tokens": 10 + } + } + }, + "warnings": [] +} diff --git a/aimux-providers/src/cohere/tests/snapshots/do_generate_tool_call.json b/aimux-providers/src/cohere/tests/snapshots/do_generate_tool_call.json new file mode 100644 index 000000000..d89f7d970 --- /dev/null +++ b/aimux-providers/src/cohere/tests/snapshots/do_generate_tool_call.json @@ -0,0 +1,121 @@ +{ + "content": [ + { + "input": "{\"location\":\"San Francisco\"}", + "toolCallId": "weather_dqgshstja6p9", + "toolName": "weather", + "type": "tool-call" + }, + { + "input": "{\"city\":\"San Francisco\"}", + "toolCallId": "cityAttractions_dcxfx4myvx68", + "toolName": "cityAttractions", + "type": "tool-call" + } + ], + "finishReason": { + "raw": "TOOL_CALL", + "unified": "tool-calls" + }, + "request": { + "body": { + "messages": [ + { + "content": "you are a friendly bot!", + "role": "system" + }, + { + "content": "Hello", + "role": "user" + } + ], + "model": "command-r-plus", + "tools": [ + { + "function": { + "name": "test-tool", + "parameters": { + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "properties": { + "value": { + "type": "string" + } + }, + "required": [ + "value" + ], + "type": "object" + } + }, + "type": "function" + } + ] + } + }, + "response": { + "body": { + "finish_reason": "TOOL_CALL", + "id": "f201af17-e24a-4396-8f6a-98e8bf9c3432", + "message": { + "role": "assistant", + "tool_calls": [ + { + "function": { + "arguments": "{\"location\":\"San Francisco\"}", + "name": "weather" + }, + "id": "weather_dqgshstja6p9", + "type": "function" + }, + { + "function": { + "arguments": "{\"city\":\"San Francisco\"}", + "name": "cityAttractions" + }, + "id": "cityAttractions_dcxfx4myvx68", + "type": "function" + } + ], + "tool_plan": "I will use the weather tool to find out the weather in San Francisco. I will also use the cityAttractions tool to find out what attractions are in San Francisco." + }, + "usage": { + "billed_units": { + "input_tokens": 119, + "output_tokens": 52 + }, + "cached_tokens": 992, + "tokens": { + "input_tokens": 1549, + "output_tokens": 103 + } + } + }, + "headers": { + "content-length": "693", + "content-type": "application/json" + } + }, + "usage": { + "inputTokens": { + "noCache": 1549, + "total": 1549 + }, + "outputTokens": { + "text": 103, + "total": 103 + }, + "raw": { + "billed_units": { + "input_tokens": 119, + "output_tokens": 52 + }, + "cached_tokens": 992, + "tokens": { + "input_tokens": 1549, + "output_tokens": 103 + } + } + }, + "warnings": [] +} diff --git a/aimux-providers/src/cohere/tests/snapshots/do_stream_empty_tool_call.json b/aimux-providers/src/cohere/tests/snapshots/do_stream_empty_tool_call.json new file mode 100644 index 000000000..211ab3cd9 --- /dev/null +++ b/aimux-providers/src/cohere/tests/snapshots/do_stream_empty_tool_call.json @@ -0,0 +1,53 @@ +[ + { + "type": "stream-start", + "warnings": [] + }, + { + "id": "66dec7d7-45e6-427c-8fd9-7d6375d12046", + "type": "response-metadata" + }, + { + "id": "currentTime_y46ar19t5gvw", + "toolName": "currentTime", + "type": "tool-input-start" + }, + { + "id": "currentTime_y46ar19t5gvw", + "type": "tool-input-end" + }, + { + "input": "{}", + "toolCallId": "currentTime_y46ar19t5gvw", + "toolName": "currentTime", + "type": "tool-call" + }, + { + "finishReason": { + "raw": "TOOL_CALL", + "unified": "tool-calls" + }, + "type": "finish", + "usage": { + "inputTokens": { + "noCache": 1445, + "total": 1445 + }, + "outputTokens": { + "text": 43, + "total": 43 + }, + "raw": { + "billed_units": { + "input_tokens": 46, + "output_tokens": 14 + }, + "cached_tokens": 704, + "tokens": { + "input_tokens": 1445, + "output_tokens": 43 + } + } + } + } +] diff --git a/aimux-providers/src/cohere/tests/snapshots/do_stream_raw_chunks.json b/aimux-providers/src/cohere/tests/snapshots/do_stream_raw_chunks.json new file mode 100644 index 000000000..901a5c113 --- /dev/null +++ b/aimux-providers/src/cohere/tests/snapshots/do_stream_raw_chunks.json @@ -0,0 +1,158 @@ +[ + { + "rawValue": { + "delta": { + "message": { + "citations": [], + "content": [], + "role": "assistant", + "tool_calls": [], + "tool_plan": "" + } + }, + "id": "321d178c-2c12-44d3-ae42-2f5510f6b1cc", + "type": "message-start" + }, + "type": "raw" + }, + { + "rawValue": { + "delta": { + "message": { + "content": { + "text": "", + "type": "text" + } + } + }, + "index": 0, + "type": "content-start" + }, + "type": "raw" + }, + { + "rawValue": { + "delta": { + "message": { + "content": { + "text": "The" + } + } + }, + "index": 0, + "type": "content-delta" + }, + "type": "raw" + }, + { + "rawValue": { + "delta": { + "message": { + "content": { + "text": " capital" + } + } + }, + "index": 0, + "type": "content-delta" + }, + "type": "raw" + }, + { + "rawValue": { + "delta": { + "message": { + "content": { + "text": " of" + } + } + }, + "index": 0, + "type": "content-delta" + }, + "type": "raw" + }, + { + "rawValue": { + "delta": { + "message": { + "content": { + "text": " France" + } + } + }, + "index": 0, + "type": "content-delta" + }, + "type": "raw" + }, + { + "rawValue": { + "delta": { + "message": { + "content": { + "text": " is" + } + } + }, + "index": 0, + "type": "content-delta" + }, + "type": "raw" + }, + { + "rawValue": { + "delta": { + "message": { + "content": { + "text": " Paris" + } + } + }, + "index": 0, + "type": "content-delta" + }, + "type": "raw" + }, + { + "rawValue": { + "delta": { + "message": { + "content": { + "text": "." + } + } + }, + "index": 0, + "type": "content-delta" + }, + "type": "raw" + }, + { + "rawValue": { + "index": 0, + "type": "content-end" + }, + "type": "raw" + }, + { + "rawValue": { + "delta": { + "finish_reason": "COMPLETE", + "usage": { + "billed_units": { + "input_tokens": 12, + "output_tokens": 7 + }, + "cached_tokens": 448, + "tokens": { + "input_tokens": 507, + "output_tokens": 10 + } + } + }, + "type": "message-end" + }, + "type": "raw" + } +] diff --git a/aimux-providers/src/cohere/tests/snapshots/do_stream_reasoning.json b/aimux-providers/src/cohere/tests/snapshots/do_stream_reasoning.json new file mode 100644 index 000000000..06785ad44 --- /dev/null +++ b/aimux-providers/src/cohere/tests/snapshots/do_stream_reasoning.json @@ -0,0 +1,279 @@ +[ + { + "type": "stream-start", + "warnings": [] + }, + { + "id": "c9117d7f-a7e4-499f-b643-a2a1e139687b", + "type": "response-metadata" + }, + { + "id": "0", + "type": "reasoning-start" + }, + { + "delta": "The", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " user", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " is", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " asking", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " for", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " the", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " sum", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " of", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " 2", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " and", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " 2", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": ".", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " Since", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " this", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " is", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " a", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " straightforward", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " arithmetic", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " problem", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": ",", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " I", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " don", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": "'t", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " need", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " to", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " use", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " any", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " tools", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": ".", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " I", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " can", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " calculate", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " the", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " answer", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": " directly", + "id": "0", + "type": "reasoning-delta" + }, + { + "delta": ".", + "id": "0", + "type": "reasoning-delta" + }, + { + "id": "0", + "type": "reasoning-end" + }, + { + "id": "1", + "type": "text-start" + }, + { + "delta": "The", + "id": "1", + "type": "text-delta" + }, + { + "delta": " answer", + "id": "1", + "type": "text-delta" + }, + { + "delta": " to", + "id": "1", + "type": "text-delta" + }, + { + "delta": " 2", + "id": "1", + "type": "text-delta" + }, + { + "delta": " +", + "id": "1", + "type": "text-delta" + }, + { + "delta": " 2", + "id": "1", + "type": "text-delta" + }, + { + "delta": " is", + "id": "1", + "type": "text-delta" + }, + { + "delta": " 4", + "id": "1", + "type": "text-delta" + }, + { + "delta": ".", + "id": "1", + "type": "text-delta" + }, + { + "id": "1", + "type": "text-end" + }, + { + "finishReason": { + "raw": "COMPLETE", + "unified": "stop" + }, + "type": "finish", + "usage": { + "inputTokens": { + "noCache": 1394, + "total": 1394 + }, + "outputTokens": { + "text": 54, + "total": 54 + }, + "raw": { + "billed_units": { + "input_tokens": 8, + "output_tokens": 50 + }, + "cached_tokens": 1360, + "tokens": { + "input_tokens": 1394, + "output_tokens": 54 + } + } + } + } +] diff --git a/aimux-providers/src/cohere/tests/snapshots/do_stream_tool_call.json b/aimux-providers/src/cohere/tests/snapshots/do_stream_tool_call.json new file mode 100644 index 000000000..b1c703516 --- /dev/null +++ b/aimux-providers/src/cohere/tests/snapshots/do_stream_tool_call.json @@ -0,0 +1,138 @@ +[ + { + "type": "stream-start", + "warnings": [] + }, + { + "id": "2941521a-b87a-45f6-9b0d-235fd66c3025", + "type": "response-metadata" + }, + { + "id": "weather_e8p4pn45zt0t", + "toolName": "weather", + "type": "tool-input-start" + }, + { + "delta": "{\"", + "id": "weather_e8p4pn45zt0t", + "type": "tool-input-delta" + }, + { + "delta": "location", + "id": "weather_e8p4pn45zt0t", + "type": "tool-input-delta" + }, + { + "delta": "\":", + "id": "weather_e8p4pn45zt0t", + "type": "tool-input-delta" + }, + { + "delta": " \"", + "id": "weather_e8p4pn45zt0t", + "type": "tool-input-delta" + }, + { + "delta": "San", + "id": "weather_e8p4pn45zt0t", + "type": "tool-input-delta" + }, + { + "delta": " Francisco", + "id": "weather_e8p4pn45zt0t", + "type": "tool-input-delta" + }, + { + "delta": "\"}", + "id": "weather_e8p4pn45zt0t", + "type": "tool-input-delta" + }, + { + "id": "weather_e8p4pn45zt0t", + "type": "tool-input-end" + }, + { + "input": "{\"location\":\"San Francisco\"}", + "toolCallId": "weather_e8p4pn45zt0t", + "toolName": "weather", + "type": "tool-call" + }, + { + "id": "cityAttractions_pyxssbwnq9fq", + "toolName": "cityAttractions", + "type": "tool-input-start" + }, + { + "delta": "{\"", + "id": "cityAttractions_pyxssbwnq9fq", + "type": "tool-input-delta" + }, + { + "delta": "city", + "id": "cityAttractions_pyxssbwnq9fq", + "type": "tool-input-delta" + }, + { + "delta": "\":", + "id": "cityAttractions_pyxssbwnq9fq", + "type": "tool-input-delta" + }, + { + "delta": " \"", + "id": "cityAttractions_pyxssbwnq9fq", + "type": "tool-input-delta" + }, + { + "delta": "San", + "id": "cityAttractions_pyxssbwnq9fq", + "type": "tool-input-delta" + }, + { + "delta": " Francisco", + "id": "cityAttractions_pyxssbwnq9fq", + "type": "tool-input-delta" + }, + { + "delta": "\"}", + "id": "cityAttractions_pyxssbwnq9fq", + "type": "tool-input-delta" + }, + { + "id": "cityAttractions_pyxssbwnq9fq", + "type": "tool-input-end" + }, + { + "input": "{\"city\":\"San Francisco\"}", + "toolCallId": "cityAttractions_pyxssbwnq9fq", + "toolName": "cityAttractions", + "type": "tool-call" + }, + { + "finishReason": { + "raw": "TOOL_CALL", + "unified": "tool-calls" + }, + "type": "finish", + "usage": { + "inputTokens": { + "noCache": 1549, + "total": 1549 + }, + "outputTokens": { + "text": 95, + "total": 95 + }, + "raw": { + "billed_units": { + "input_tokens": 119, + "output_tokens": 44 + }, + "cached_tokens": 1504, + "tokens": { + "input_tokens": 1549, + "output_tokens": 95 + } + } + } + } +] diff --git a/aimux-providers/tests/cohere_embedding_test.rs b/aimux-providers/tests/cohere_embedding_test.rs index 5631903ab..bb660e698 100644 --- a/aimux-providers/tests/cohere_embedding_test.rs +++ b/aimux-providers/tests/cohere_embedding_test.rs @@ -49,200 +49,204 @@ fn default_options(values: Vec) -> EmbeddingCallOptions { } } -/// TS: "should extract embedding" -#[tokio::test] -async fn should_extract_embedding() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/embed")) - .respond_with(ResponseTemplate::new(200).set_body_json(embedding_response_body())) - .mount(&server) - .await; - - let provider = create_cohere(CohereProviderSettings { - api_key: Some("test-api-key".to_string()), - base_url: Some(server.uri()), - ..Default::default() - }) - .unwrap(); - let model = provider.embedding_model("embed-english-v3.0").unwrap(); - - let result = model - .do_embed(&default_options(test_values())) - .await - .expect("should succeed"); - assert_eq!(result.embeddings.len(), 2); - assert_eq!( - result.embeddings[0], - vec![0.03302002, 0.020904541, -0.019744873, -0.0625, 0.04437256] - ); - assert_eq!( - result.embeddings[1], - vec![ - -0.04660034, - 0.00037765503, - -0.061157227, - -0.08239746, - -0.010360718 - ] - ); -} - -/// TS: "should expose the raw response" -#[tokio::test] -async fn should_expose_raw_response() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/embed")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("test-header", "test-value") - .set_body_json(embedding_response_body()), - ) - .mount(&server) - .await; - - let provider = create_cohere(CohereProviderSettings { - api_key: Some("test-api-key".to_string()), - base_url: Some(server.uri()), - ..Default::default() - }) - .unwrap(); - let model = provider.embedding_model("embed-english-v3.0").unwrap(); - - let result = model - .do_embed(&default_options(test_values())) - .await - .expect("should succeed"); - let headers = result - .response - .as_ref() - .and_then(|r| r.headers.as_ref()) - .expect("headers present"); - assert_eq!( - headers.get("test-header").map(std::string::String::as_str), - Some("test-value") - ); - assert_eq!( - result.response.unwrap().body, - Some(embedding_response_body()) - ); -} - -/// TS: "should extract usage" -#[tokio::test] -async fn should_extract_usage() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/embed")) - .respond_with(ResponseTemplate::new(200).set_body_json(embedding_response_body())) - .mount(&server) - .await; - - let provider = create_cohere(CohereProviderSettings { - api_key: Some("test-api-key".to_string()), - base_url: Some(server.uri()), - ..Default::default() - }) - .unwrap(); - let model = provider.embedding_model("embed-english-v3.0").unwrap(); - - let result = model - .do_embed(&default_options(test_values())) - .await - .expect("should succeed"); - let usage = result.usage.expect("usage should be present"); - assert_eq!(usage.tokens, 10); -} +mod do_embed { + use super::*; + + /// upstream: @ai-sdk/cohere src/cohere-embedding-model.test.ts › doEmbed › should extract embedding + #[tokio::test] + async fn should_extract_embedding() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/embed")) + .respond_with(ResponseTemplate::new(200).set_body_json(embedding_response_body())) + .mount(&server) + .await; + + let provider = create_cohere(CohereProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("embed-english-v3.0").unwrap(); + + let result = model + .do_embed(&default_options(test_values())) + .await + .expect("should succeed"); + assert_eq!(result.embeddings.len(), 2); + assert_eq!( + result.embeddings[0], + vec![0.03302002, 0.020904541, -0.019744873, -0.0625, 0.04437256] + ); + assert_eq!( + result.embeddings[1], + vec![ + -0.04660034, + 0.00037765503, + -0.061157227, + -0.08239746, + -0.010360718 + ] + ); + } -/// TS: "should pass the model and the values" -#[tokio::test] -async fn should_pass_model_and_values() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/embed")) - .respond_with(ResponseTemplate::new(200).set_body_json(embedding_response_body())) - .mount(&server) - .await; - - let provider = create_cohere(CohereProviderSettings { - api_key: Some("test-api-key".to_string()), - base_url: Some(server.uri()), - ..Default::default() - }) - .unwrap(); - let model = provider.embedding_model("embed-english-v3.0").unwrap(); - - let _ = model - .do_embed(&default_options(test_values())) - .await - .expect("should succeed"); - - let requests = server.received_requests().await.expect("requests recorded"); - assert_eq!(requests.len(), 1); - let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - assert_eq!(body["model"], "embed-english-v3.0"); - assert_eq!(body["embedding_types"], json!(["float"])); - assert_eq!(body["input_type"], "search_query"); - assert_eq!( - body["texts"], - json!(["sunny day at the beach", "rainy day in the city"]) - ); -} + /// partial: `should expose the raw response`; full port in src/cohere/tests/embedding_model.rs. + #[tokio::test] + async fn should_expose_raw_response() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/embed")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("test-header", "test-value") + .set_body_json(embedding_response_body()), + ) + .mount(&server) + .await; + + let provider = create_cohere(CohereProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("embed-english-v3.0").unwrap(); + + let result = model + .do_embed(&default_options(test_values())) + .await + .expect("should succeed"); + let headers = result + .response + .as_ref() + .and_then(|r| r.headers.as_ref()) + .expect("headers present"); + assert_eq!( + headers.get("test-header").map(std::string::String::as_str), + Some("test-value") + ); + assert_eq!( + result.response.unwrap().body, + Some(embedding_response_body()) + ); + } -/// TS: "should pass headers" -#[tokio::test] -async fn should_pass_headers() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/embed")) - .respond_with(ResponseTemplate::new(200).set_body_json(embedding_response_body())) - .mount(&server) - .await; - - let provider = create_cohere(CohereProviderSettings { - api_key: Some("test-api-key".to_string()), - base_url: Some(server.uri()), - headers: Some(HashMap::from([( - "Custom-Provider-Header".to_string(), - Some("provider-header-value".to_string()), - )])), - ..Default::default() - }) - .unwrap(); - let model = provider.embedding_model("embed-english-v3.0").unwrap(); + /// upstream: @ai-sdk/cohere src/cohere-embedding-model.test.ts › doEmbed › should extract usage + #[tokio::test] + async fn should_extract_usage() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/embed")) + .respond_with(ResponseTemplate::new(200).set_body_json(embedding_response_body())) + .mount(&server) + .await; + + let provider = create_cohere(CohereProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("embed-english-v3.0").unwrap(); + + let result = model + .do_embed(&default_options(test_values())) + .await + .expect("should succeed"); + let usage = result.usage.expect("usage should be present"); + assert_eq!(usage.tokens, 10); + } - let mut request_headers = HashMap::new(); - request_headers.insert( - "Custom-Request-Header".to_string(), - "request-header-value".to_string(), - ); + /// upstream: @ai-sdk/cohere src/cohere-embedding-model.test.ts › doEmbed › should pass the model and the values + #[tokio::test] + async fn should_pass_model_and_values() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/embed")) + .respond_with(ResponseTemplate::new(200).set_body_json(embedding_response_body())) + .mount(&server) + .await; + + let provider = create_cohere(CohereProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("embed-english-v3.0").unwrap(); + + let _ = model + .do_embed(&default_options(test_values())) + .await + .expect("should succeed"); + + let requests = server.received_requests().await.expect("requests recorded"); + assert_eq!(requests.len(), 1); + let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); + assert_eq!(body["model"], "embed-english-v3.0"); + assert_eq!(body["embedding_types"], json!(["float"])); + assert_eq!(body["input_type"], "search_query"); + assert_eq!( + body["texts"], + json!(["sunny day at the beach", "rainy day in the city"]) + ); + } - let options = EmbeddingCallOptions { - values: test_values(), - abort_signal: None, - provider_options: None, - headers: Some(request_headers), - max_retries: None, - timeout: None, - }; - - let _ = model.do_embed(&options).await.expect("should succeed"); - - let requests = server.received_requests().await.expect("requests recorded"); - assert_eq!(requests.len(), 1); - let h = &requests[0].headers; - assert_eq!( - h.get("custom-provider-header") - .and_then(|v| v.to_str().ok()), - Some("provider-header-value") - ); - assert_eq!( - h.get("authorization").and_then(|v| v.to_str().ok()), - Some("Bearer test-api-key") - ); - assert_eq!( - h.get("custom-request-header").and_then(|v| v.to_str().ok()), - Some("request-header-value") - ); + /// upstream: @ai-sdk/cohere src/cohere-embedding-model.test.ts › doEmbed › should pass headers + #[tokio::test] + async fn should_pass_headers() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/embed")) + .respond_with(ResponseTemplate::new(200).set_body_json(embedding_response_body())) + .mount(&server) + .await; + + let provider = create_cohere(CohereProviderSettings { + api_key: Some("test-api-key".to_string()), + base_url: Some(server.uri()), + headers: Some(HashMap::from([( + "Custom-Provider-Header".to_string(), + Some("provider-header-value".to_string()), + )])), + ..Default::default() + }) + .unwrap(); + let model = provider.embedding_model("embed-english-v3.0").unwrap(); + + let mut request_headers = HashMap::new(); + request_headers.insert( + "Custom-Request-Header".to_string(), + "request-header-value".to_string(), + ); + + let options = EmbeddingCallOptions { + values: test_values(), + abort_signal: None, + provider_options: None, + headers: Some(request_headers), + max_retries: None, + timeout: None, + }; + + let _ = model.do_embed(&options).await.expect("should succeed"); + + let requests = server.received_requests().await.expect("requests recorded"); + assert_eq!(requests.len(), 1); + let h = &requests[0].headers; + assert_eq!( + h.get("custom-provider-header") + .and_then(|v| v.to_str().ok()), + Some("provider-header-value") + ); + assert_eq!( + h.get("authorization").and_then(|v| v.to_str().ok()), + Some("Bearer test-api-key") + ); + assert_eq!( + h.get("custom-request-header").and_then(|v| v.to_str().ok()), + Some("request-header-value") + ); + } } diff --git a/aimux-providers/tests/cohere_upstream_e2e_test.rs b/aimux-providers/tests/cohere_upstream_e2e_test.rs index f38671a0c..a71e905c6 100644 --- a/aimux-providers/tests/cohere_upstream_e2e_test.rs +++ b/aimux-providers/tests/cohere_upstream_e2e_test.rs @@ -213,173 +213,201 @@ fn reranking_fixture() -> Value { .unwrap() } -/// TS: "should extract text response" (cohere/src/cohere-chat-language-model.test.ts) -#[tokio::test] -async fn generate_text() { - let fixture = cohere_text(); - let fetch = MockFetch::new(vec![Canned::json(&fixture)]); - let result = provider(&fetch) - .chat("command-r-plus") - .do_generate(&options()) - .await - .unwrap(); - assert_request(&fetch, "chat", &request(false)); - assert_eq!(result.request.unwrap().body, Some(request(false))); - assert_eq!( - result.content, - vec![GenerateContent::Text { - text: "The capital of France is Paris.".into(), - provider_metadata: None, - }] - ); - assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); - assert_eq!(result.finish_reason.raw.as_deref(), Some("COMPLETE")); - assert_usage(&result.usage, fixture["usage"].clone()); - assert!(result.warnings.is_empty()); - assert!(result.provider_metadata.is_none()); - let response = result.response.unwrap(); - assert_eq!(response.body, Some(fixture)); - assert_eq!(response.id, None); - assert_eq!(response.timestamp, None); - assert_eq!(response.model_id, None); -} +mod do_generate { + use super::*; -/// TS: "should stream text deltas" (cohere/src/cohere-chat-language-model.test.ts) -#[tokio::test] -async fn stream_text() { - let chunks = cohere_text_chunks(); - let body = chunks - .iter() - .map(|chunk| { - format!( - "event: {}\ndata: {chunk}\n\n", - chunk["type"].as_str().unwrap() - ) - }) - .collect(); - let fetch = MockFetch::new(vec![sse(body)]); - let result = provider(&fetch) - .chat("command-r-plus") - .do_stream(&options()) - .await - .unwrap(); - assert_request(&fetch, "chat", &request(true)); - assert_eq!(result.request.unwrap().body, Some(request(true))); - let parts: Vec<_> = result.stream.map(|part| part.unwrap()).collect().await; - assert_eq!(parts.len(), 12); - assert!(matches!(&parts[0], StreamPart::StreamStart { warnings } if warnings.is_empty())); - assert!( - matches!(&parts[1], StreamPart::ResponseMetadata(meta) if meta.id.as_deref() == Some("321d178c-2c12-44d3-ae42-2f5510f6b1cc")) - ); - assert!(matches!(&parts[2], StreamPart::TextStart { id, .. } if id == "0")); - let deltas: Vec<_> = parts - .iter() - .filter_map(|part| match part { - StreamPart::TextDelta { id, delta, .. } => { - assert_eq!(id, "0"); - Some(delta.as_str()) - } - _ => None, - }) - .collect(); - assert_eq!( - deltas, - vec!["The", " capital", " of", " France", " is", " Paris", "."] - ); - assert!(matches!(&parts[10], StreamPart::TextEnd { id, .. } if id == "0")); - match &parts[11] { - StreamPart::Finish { - finish_reason, - usage, - provider_metadata, - } => { - assert_eq!(finish_reason.unified, FinishReasonUnified::Stop); - assert_eq!(finish_reason.raw.as_deref(), Some("COMPLETE")); - assert_usage(usage, chunks.last().unwrap()["delta"]["usage"].clone()); - assert!(provider_metadata.is_none()); + mod text { + use super::*; + + /// partial: `should extract text response`; full port in src/cohere/tests/chat_language_model.rs. + #[tokio::test] + async fn generate_text() { + let fixture = cohere_text(); + let fetch = MockFetch::new(vec![Canned::json(&fixture)]); + let result = provider(&fetch) + .chat("command-r-plus") + .do_generate(&options()) + .await + .unwrap(); + assert_request(&fetch, "chat", &request(false)); + assert_eq!(result.request.unwrap().body, Some(request(false))); + assert_eq!( + result.content, + vec![GenerateContent::Text { + text: "The capital of France is Paris.".into(), + provider_metadata: None, + }] + ); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(result.finish_reason.raw.as_deref(), Some("COMPLETE")); + assert_usage(&result.usage, fixture["usage"].clone()); + assert!(result.warnings.is_empty()); + assert!(result.provider_metadata.is_none()); + let response = result.response.unwrap(); + assert_eq!(response.body, Some(fixture)); + assert_eq!(response.id, None); + assert_eq!(response.timestamp, None); + assert_eq!(response.model_id, None); } - part => panic!("expected finish, got {part:?}"), } -} -/// TS: "should extract tool calls" (cohere/src/cohere-chat-language-model.test.ts) -#[tokio::test] -async fn generate_tool_calls() { - let fixture = cohere_tool_call(); - let fetch = MockFetch::new(vec![Canned::json(&fixture)]); - let schema = json!({"type":"object","properties":{"value":{"type":"string"}},"required":["value"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"}); - let mut opts = options(); - opts.tools = Some(vec![Tool::Function(FunctionTool::new( - "test-tool", - schema.clone(), - ))]); - let result = provider(&fetch) - .chat("command-r-plus") - .do_generate(&opts) - .await - .unwrap(); - let mut expected = request(false); - expected["tools"] = - json!([{"type":"function","function":{"name":"test-tool","parameters":schema}}]); - assert_request(&fetch, "chat", &expected); - assert_eq!(result.content.len(), 2); - for (content, wire) in result - .content - .iter() - .zip(fixture["message"]["tool_calls"].as_array().unwrap()) - { - match content { - GenerateContent::ToolCall(call) => { - assert_eq!(call.tool_call_id, wire["id"]); - assert_eq!(call.tool_name, wire["function"]["name"]); - assert_eq!(call.input, wire["function"]["arguments"]); - assert!(call.provider_metadata.is_none()); + mod tool_call { + use super::*; + + /// partial: `should extract tool calls`; full port in src/cohere/tests/chat_language_model.rs. + #[tokio::test] + async fn generate_tool_calls() { + let fixture = cohere_tool_call(); + let fetch = MockFetch::new(vec![Canned::json(&fixture)]); + let schema = json!({"type":"object","properties":{"value":{"type":"string"}},"required":["value"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"}); + let mut opts = options(); + opts.tools = Some(vec![Tool::Function(FunctionTool::new( + "test-tool", + schema.clone(), + ))]); + let result = provider(&fetch) + .chat("command-r-plus") + .do_generate(&opts) + .await + .unwrap(); + let mut expected = request(false); + expected["tools"] = + json!([{"type":"function","function":{"name":"test-tool","parameters":schema}}]); + assert_request(&fetch, "chat", &expected); + assert_eq!(result.content.len(), 2); + for (content, wire) in result + .content + .iter() + .zip(fixture["message"]["tool_calls"].as_array().unwrap()) + { + match content { + GenerateContent::ToolCall(call) => { + assert_eq!(call.tool_call_id, wire["id"]); + assert_eq!(call.tool_name, wire["function"]["name"]); + assert_eq!(call.input, wire["function"]["arguments"]); + assert!(call.provider_metadata.is_none()); + } + other => panic!("expected tool call, got {other:?}"), + } } - other => panic!("expected tool call, got {other:?}"), + assert_eq!(result.finish_reason.unified, FinishReasonUnified::ToolCalls); + assert_eq!(result.finish_reason.raw.as_deref(), Some("TOOL_CALL")); + assert_eq!(result.usage.input_tokens.total, Some(1549)); + assert_eq!(result.usage.output_tokens.total, Some(103)); + assert_eq!(result.usage.raw, fixture["usage"].as_object().cloned()); } } - assert_eq!(result.finish_reason.unified, FinishReasonUnified::ToolCalls); - assert_eq!(result.finish_reason.raw.as_deref(), Some("TOOL_CALL")); - assert_eq!(result.usage.input_tokens.total, Some(1549)); - assert_eq!(result.usage.output_tokens.total, Some(103)); - assert_eq!(result.usage.raw, fixture["usage"].as_object().cloned()); } -/// TS: "should handle unparsable stream parts" (cohere/src/cohere-chat-language-model.test.ts) -#[tokio::test] -async fn malformed_stream_response() { - let fetch = MockFetch::new(vec![sse( - "event: foo-message\ndata: {unparsable}\n\n".into() - )]); - let result = provider(&fetch) - .chat("command-r-plus") - .do_stream(&options()) - .await - .unwrap(); - assert_request(&fetch, "chat", &request(true)); - let parts: Vec<_> = result.stream.map(|part| part.unwrap()).collect().await; - assert_eq!(parts.len(), 3); - assert!(matches!(&parts[0], StreamPart::StreamStart { warnings } if warnings.is_empty())); - assert!(matches!(&parts[1], StreamPart::Error { .. })); - match &parts[2] { - StreamPart::Finish { - finish_reason, - usage, - provider_metadata, - } => { - assert_eq!(finish_reason.unified, FinishReasonUnified::Error); - assert_eq!(finish_reason.raw, None); +mod do_stream { + use super::*; + + mod text { + use super::*; + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doStream › text › should stream text deltas + #[tokio::test] + async fn stream_text() { + let chunks = cohere_text_chunks(); + let body = chunks + .iter() + .map(|chunk| { + format!( + "event: {}\ndata: {chunk}\n\n", + chunk["type"].as_str().unwrap() + ) + }) + .collect(); + let fetch = MockFetch::new(vec![sse(body)]); + let result = provider(&fetch) + .chat("command-r-plus") + .do_stream(&options()) + .await + .unwrap(); + assert_request(&fetch, "chat", &request(true)); + assert_eq!(result.request.unwrap().body, Some(request(true))); + let parts: Vec<_> = result.stream.map(|part| part.unwrap()).collect().await; + assert_eq!(parts.len(), 12); + assert!( + matches!(&parts[0], StreamPart::StreamStart { warnings } if warnings.is_empty()) + ); + assert!( + matches!(&parts[1], StreamPart::ResponseMetadata(meta) if meta.id.as_deref() == Some("321d178c-2c12-44d3-ae42-2f5510f6b1cc")) + ); + assert!(matches!(&parts[2], StreamPart::TextStart { id, .. } if id == "0")); + let deltas: Vec<_> = parts + .iter() + .filter_map(|part| match part { + StreamPart::TextDelta { id, delta, .. } => { + assert_eq!(id, "0"); + Some(delta.as_str()) + } + _ => None, + }) + .collect(); assert_eq!( - serde_json::to_value(usage).unwrap(), - serde_json::to_value(Usage::default()).unwrap() + deltas, + vec!["The", " capital", " of", " France", " is", " Paris", "."] + ); + assert!(matches!(&parts[10], StreamPart::TextEnd { id, .. } if id == "0")); + match &parts[11] { + StreamPart::Finish { + finish_reason, + usage, + provider_metadata, + } => { + assert_eq!(finish_reason.unified, FinishReasonUnified::Stop); + assert_eq!(finish_reason.raw.as_deref(), Some("COMPLETE")); + assert_usage(usage, chunks.last().unwrap()["delta"]["usage"].clone()); + assert!(provider_metadata.is_none()); + } + part => panic!("expected finish, got {part:?}"), + } + } + } + + mod error_handling { + use super::*; + + /// upstream: @ai-sdk/cohere src/cohere-chat-language-model.test.ts › doStream › error handling › should handle unparsable stream parts + #[tokio::test] + async fn malformed_stream_response() { + let fetch = MockFetch::new(vec![sse( + "event: foo-message\ndata: {unparsable}\n\n".into() + )]); + let result = provider(&fetch) + .chat("command-r-plus") + .do_stream(&options()) + .await + .unwrap(); + assert_request(&fetch, "chat", &request(true)); + let parts: Vec<_> = result.stream.map(|part| part.unwrap()).collect().await; + assert_eq!(parts.len(), 3); + assert!( + matches!(&parts[0], StreamPart::StreamStart { warnings } if warnings.is_empty()) ); - assert!(provider_metadata.is_none()); + assert!(matches!(&parts[1], StreamPart::Error { .. })); + match &parts[2] { + StreamPart::Finish { + finish_reason, + usage, + provider_metadata, + } => { + assert_eq!(finish_reason.unified, FinishReasonUnified::Error); + assert_eq!(finish_reason.raw, None); + assert_eq!( + serde_json::to_value(usage).unwrap(), + serde_json::to_value(Usage::default()).unwrap() + ); + assert!(provider_metadata.is_none()); + } + part => panic!("expected finish, got {part:?}"), + } } - part => panic!("expected finish, got {part:?}"), } } -/// TS: "should extract embedding" (cohere/src/cohere-embedding-model.test.ts) +/// Embed request body, float vectors, usage and raw response of one call; the upstream case itself is in cohere_embedding_test.rs. #[tokio::test] async fn embed_values() { let fixture = embedding_fixture(); @@ -467,19 +495,31 @@ async fn rerank(object_documents: bool) { assert_eq!(response.body, Some(fixture)); } -/// TS: "should send request with stringified json documents" (cohere/src/reranking/cohere-reranking-model.test.ts) -#[tokio::test] -async fn rerank_object_documents() { - rerank(true).await; -} +mod do_rerank { + use super::*; -/// TS: "should send request with text documents" (cohere/src/reranking/cohere-reranking-model.test.ts) -#[tokio::test] -async fn rerank_text_documents() { - rerank(false).await; + mod json_documents { + use super::*; + + /// upstream: @ai-sdk/cohere src/reranking/cohere-reranking-model.test.ts › doRerank › json documents › should send request with stringified json documents + #[tokio::test] + async fn rerank_object_documents() { + rerank(true).await; + } + } + + mod text_documents { + use super::*; + + /// upstream: @ai-sdk/cohere src/reranking/cohere-reranking-model.test.ts › doRerank › text documents › should send request with text documents + #[tokio::test] + async fn rerank_text_documents() { + rerank(false).await; + } + } } -/// TS: `supportedUrls` (cohere-chat-language-model.ts:51-53): image URLs go +/// aimux test of `supportedUrls` (@ai-sdk/cohere@4.0.52 src/cohere-chat-language-model.ts:51-53): image URLs go /// to the API as they are; other URLs are downloaded by the operation. #[tokio::test] async fn chat_supports_http_image_urls() {