diff --git a/crates/consts/src/protocol.rs b/crates/consts/src/protocol.rs index 1205ec5..e6c033e 100644 --- a/crates/consts/src/protocol.rs +++ b/crates/consts/src/protocol.rs @@ -45,8 +45,8 @@ pub enum Protocol { AwsAnthropic, /// Any Bedrock model through the Converse API (Messages engine, transcoded) AwsConverse, - /// Cohere Command on AWS Bedrock (SigV4) - AwsCohere, + /// Titan/Cohere embeddings on AWS Bedrock (InvokeModel) + AwsEmbed, /// Llama on AWS Bedrock (SigV4) AwsLlama, /// Alibaba DashScope native (input.messages/parameters/output.choices) @@ -77,7 +77,7 @@ impl Protocol { Protocol::MinimaxV1, Protocol::AwsAnthropic, Protocol::AwsConverse, - Protocol::AwsCohere, + Protocol::AwsEmbed, Protocol::AwsLlama, Protocol::Dashscope, Protocol::Moderations, @@ -104,7 +104,7 @@ impl Protocol { Protocol::MinimaxV1 => "minimax-v1", Protocol::AwsAnthropic => "aws-anthropic", Protocol::AwsConverse => "aws-converse", - Protocol::AwsCohere => "aws-cohere", + Protocol::AwsEmbed => "aws-embed", Protocol::AwsLlama => "aws-llama", Protocol::Dashscope => "dashscope", Protocol::Moderations => "moderations", diff --git a/crates/engines/src/bedrock.rs b/crates/engines/src/bedrock.rs index 7d4fd92..d7aba68 100644 --- a/crates/engines/src/bedrock.rs +++ b/crates/engines/src/bedrock.rs @@ -208,24 +208,31 @@ pub(crate) async fn bedrock_invoke( Ok((status, v, headers)) } +/// Billed input tokens, `fallback` when the header is absent — an embed reply +/// carries the input header alone, so the pairwise form never matches it. +pub(crate) fn bedrock_input_tokens(headers: &HeaderMap, fallback: i64) -> i64 { + token_header(headers, "x-amzn-bedrock-input-token-count").unwrap_or(fallback) +} + /// Bedrock stamps every InvokeModel reply with the billed counts while only some /// family bodies carry them, so the headers win when present. pub(crate) fn bedrock_header_usage(headers: &HeaderMap, body: (i64, i64)) -> (i64, i64) { - let count = |name: &str| { - headers - .get(name) - .and_then(|v| v.to_str().ok()) - .and_then(|s| s.parse::().ok()) - }; match ( - count("x-amzn-bedrock-input-token-count"), - count("x-amzn-bedrock-output-token-count"), + token_header(headers, "x-amzn-bedrock-input-token-count"), + token_header(headers, "x-amzn-bedrock-output-token-count"), ) { (Some(input), Some(output)) => (input, output), _ => body, } } +fn token_header(headers: &HeaderMap, name: &str) -> Option { + headers + .get(name) + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.parse::().ok()) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/engines/src/bespoke.rs b/crates/engines/src/bespoke.rs index 9ea40ea..b9b9ad1 100644 --- a/crates/engines/src/bespoke.rs +++ b/crates/engines/src/bespoke.rs @@ -3,12 +3,14 @@ //! response shape (the mock answers in the same shapes). AWS engines compute a //! real SigV4 Authorization header. -use gw_models::{GResult, GatewayResponse}; +use gw_models::{GResult, GatewayError, GatewayResponse, TypedParams}; use gw_protocol::object; use serde_json::{Map, Value, json}; use crate::base::{Base, base_engine}; -use crate::bedrock::{bedrock_header_usage, bedrock_invoke, bedrock_stream, invocation_metrics}; +use crate::bedrock::{ + bedrock_header_usage, bedrock_input_tokens, bedrock_invoke, bedrock_stream, invocation_metrics, +}; use crate::engine::{EngineOutcome, ModelEngine, StreamChunk, reject_minimax_error}; use crate::transport::Headers; @@ -108,109 +110,69 @@ impl ModelEngine for MinimaxV1Engine { } } -base_engine!(CohereEngine); - -impl CohereEngine { - fn build_body(&mut self) -> Value { - // cohere's system slot is `preamble` (system turns are filtered above) - let system = self.base.system_text(); - let mut history = simple_turns(&mut self.base, ("CHATBOT", "USER"), ("role", "message")); - let message = history - .pop() - .map(|mut last| last["message"].take()) - .unwrap_or(Value::String(String::new())); - let mut body = json!({}); - body["message"] = message; - body["chat_history"] = Value::Array(history); - if !system.is_empty() { - body["preamble"] = system.into(); - } - if let Some(p) = self.base.chat_params() - && let Some(mt) = p.max_tokens - { - body["max_tokens"] = json!(mt); - } - body - } -} +base_engine!(AwsEmbedEngine); #[async_trait::async_trait] -impl ModelEngine for CohereEngine { - /// Bedrock Cohere Command: `{message, chat_history}` → `{text, finish_reason}` - /// (legacy `generations[0]`), billed counts in the Bedrock headers. +impl ModelEngine for AwsEmbedEngine { + /// Bedrock embeddings over InvokeModel, answered in the OpenAI list shape: + /// Titan takes exactly one `{inputText}`, Cohere batches `{texts}`; usage + /// from the `x-amzn-bedrock-*` header, the body count as fallback. async fn run(&mut self) -> GResult { let model = self.base.model_name()?.to_owned(); - let body = self.build_body(); - if self.base.request.stream { - let mut resp = GatewayResponse { - model, - is_messages_protocol: true, - ..Default::default() - }; - let mut full = String::new(); - let (status, r) = bedrock_stream(&mut self.base, body, |v| { - let mut chunks = Vec::new(); - if let Some(t) = v["text"].as_str() - && !t.is_empty() - && v["event_type"] != "stream-end" - { - full.push_str(t); - chunks.push(StreamChunk { - delta: t.to_owned(), - ..Default::default() - }); - } - if let Some(fr) = v["finish_reason"].as_str() { - resp.finish_reason = cohere_finish_reason(Some(fr)); - chunks.push(StreamChunk { - finish_reason: Some(resp.finish_reason.clone()), - ..Default::default() - }); - } - invocation_metrics(&v, &mut resp); - Ok(chunks) - }) - .await?; - resp.message = full; - resp.total_tokens = resp.prompt_tokens.saturating_add(resp.completion_tokens); - resp.raw_usage = Some( - json!({"input_tokens": resp.prompt_tokens, "output_tokens": resp.completion_tokens}), - ); - return Ok(EngineOutcome::from_pump(resp, status, r)); - } - let (status, mut v, headers) = bedrock_invoke(&mut self.base, &model, body).await?; - let message = crate::engine::take_string(&mut v, "/text") - .or_else(|| crate::engine::take_string(&mut v, "/generations/0/text")) - .unwrap_or_default(); - let finish_reason = cohere_finish_reason( - v["finish_reason"] - .as_str() - .or_else(|| v["generations"][0]["finish_reason"].as_str()), - ); - let meta = &v["meta"]; - let body_count = |key: &str| match crate::engine::tok(&meta["billed_units"][key]) { - 0 => crate::engine::tok(&meta["tokens"][key]), - n => n, + let (texts, dimensions) = match self.base.take_typed() { + Some(TypedParams::Embeddings(p)) if !p.input.is_empty() => (p.input, p.dimensions), + _ => { + return Err(GatewayError::bad_request( + "aws-embed serves embeddings input only", + )); + } }; - let (input, output) = bedrock_header_usage( - &headers, - (body_count("input_tokens"), body_count("output_tokens")), - ); + let mut data = Vec::with_capacity(texts.len()); + let (status, prompt_tokens) = if model.starts_with("cohere.") { + // bedrock's cohere embed requires an input_type; the gateway embeds for retrieval storage + let mut body = json!({"input_type": "search_document"}); + body["texts"] = Value::Array(texts.into_iter().map(Value::String).collect()); + let (st, mut v, headers) = bedrock_invoke(&mut self.base, &model, body).await?; + if let Value::Array(rows) = v["embeddings"].take() { + data.extend(rows.into_iter().enumerate().map(embedding_row)); + } + (st, bedrock_input_tokens(&headers, 0)) + } else { + let [text] = <[String; 1]>::try_from(texts).map_err(|_| { + GatewayError::bad_request("titan embeddings require exactly one input") + })?; + let mut body = json!({}); + body["inputText"] = Value::String(text); + if let Some(d) = dimensions { + body["dimensions"] = json!(d); + } + let (st, mut v, headers) = bedrock_invoke(&mut self.base, &model, body).await?; + let body_count = crate::engine::tok(&v["inputTextTokenCount"]); + data.push(embedding_row((0, v["embedding"].take()))); + (st, bedrock_input_tokens(&headers, body_count)) + }; + let mut v2 = json!({"object": "list", "model": model, + "usage": {"prompt_tokens": prompt_tokens, "total_tokens": prompt_tokens}}); + v2["data"] = Value::Array(data); let resp = GatewayResponse { - message, model, - finish_reason, - prompt_tokens: input, - completion_tokens: output, - total_tokens: input.saturating_add(output), - raw_usage: Some(json!({"input_tokens": input, "output_tokens": output})), - is_messages_protocol: true, // anthropic's usage fields align with cohere's input/output + prompt_tokens, + total_tokens: prompt_tokens, + raw_usage: Some(v2["usage"].clone()), + response_v2: Some(v2), + finish_reason: "stop".to_owned(), ..Default::default() }; Ok(EngineOutcome::with_status(resp, status)) } } +fn embedding_row((index, embedding): (usize, Value)) -> Value { + let mut row = json!({"object": "embedding", "index": index}); + row["embedding"] = embedding; + row +} + /// The non-system turns as `{role_key: ai|user, content_key: text}` objects, /// moved out of the request (the vendors' flat two-role wires). fn simple_turns( @@ -264,14 +226,6 @@ fn llama_prompt(model: &str, messages: &[gw_models::ChatMsg]) -> String { prompt } -fn cohere_finish_reason(vendor: Option<&str>) -> String { - match vendor { - None | Some("COMPLETE") => "stop".to_owned(), - Some("MAX_TOKENS") => "length".to_owned(), - Some(other) => other.to_lowercase(), - } -} - base_engine!(LlamaEngine); #[async_trait::async_trait] @@ -297,14 +251,14 @@ impl ModelEngine for LlamaEngine { ..Default::default() }; let mut full = String::new(); - let (status, r) = bedrock_stream(&mut self.base, body, |v| { + let (status, r) = bedrock_stream(&mut self.base, body, |mut v| { let mut chunks = Vec::new(); - if let Some(t) = v["generation"].as_str() + if let Some(Value::String(t)) = v.get_mut("generation").map(Value::take) && !t.is_empty() { - full.push_str(t); + full.push_str(&t); chunks.push(StreamChunk { - delta: t.to_owned(), + delta: t, ..Default::default() }); } @@ -542,7 +496,7 @@ mod tests { use std::sync::Arc; use gw_consts::Protocol; - use gw_models::{ChatMsg, GatewayRequest, ModelParamV2}; + use gw_models::{ChatMsg, EmbeddingParams, GatewayRequest, ModelParamV2}; use super::*; use crate::transport::{ @@ -610,48 +564,61 @@ mod tests { assert!(out.response.total_tokens > 0); } + fn embed_req(model: &str, inputs: &[&str], dimensions: Option) -> GatewayRequest { + let mut r = req(Protocol::AwsEmbed, model); + r.model_param_v2.as_mut().unwrap().typed = Some(TypedParams::Embeddings(EmbeddingParams { + input: inputs.iter().map(|s| (*s).to_owned()).collect(), + dimensions, + })); + r + } + #[tokio::test] - async fn cohere_wire_shape() { - let mut e = CohereEngine::new(req(Protocol::AwsCohere, "cohere.command-r-v1:0"), t()); - let out = e.run().await.unwrap(); - assert!( - out.response - .message - .contains("[mock-cohere] you said: hello bespoke") - ); - assert!(out.response.prompt_tokens > 0 && out.response.completion_tokens > 0); + async fn titan_embeddings_invoke_once_and_bill_the_body_count() { + let r = embed_req("amazon.titan-embed-text-v2:0", &["hello world"], Some(256)); + let out = AwsEmbedEngine::new(r, t()).run().await.unwrap(); + let v = out.response.response_v2.unwrap(); + assert_eq!(v["data"].as_array().unwrap().len(), 1); + assert_eq!(v["data"][0]["index"], 0); + assert!(v["data"][0]["embedding"].is_array()); + assert_eq!(out.response.prompt_tokens, 2); + assert_eq!(out.response.completion_tokens, 0); + assert_eq!(v["usage"]["total_tokens"], 2); } #[tokio::test] - async fn bedrock_headers_bill_command_r_and_the_legacy_command_shape_parses() { - let mut e = CohereEngine::new( - req(Protocol::AwsCohere, "cohere.command-r-v1:0"), - Arc::new(BedrockReply( - r#"{"response_id":"r","text":"hi there","finish_reason":"COMPLETE"}"#, - 57, - 9, - )), - ); - let out = e.run().await.unwrap(); - assert_eq!(out.response.message, "hi there"); - assert_eq!(out.response.finish_reason, "stop"); - assert_eq!( - (out.response.prompt_tokens, out.response.completion_tokens), - (57, 9) - ); + async fn cohere_embeddings_batch_once_and_bill_the_bedrock_headers() { + let r = embed_req("cohere.embed-english-v3", &["a", "b", "c"], None); + let out = AwsEmbedEngine::new(r, t()).run().await.unwrap(); + let v = out.response.response_v2.unwrap(); + assert_eq!(v["data"].as_array().unwrap().len(), 3); + assert_eq!(out.response.prompt_tokens, 12); + assert_eq!(v["usage"]["prompt_tokens"], 12); + } - let mut e = CohereEngine::new( - req(Protocol::AwsCohere, "cohere.command-text-v14"), - Arc::new(BedrockReply( - r#"{"generations":[{"id":"g","text":"legacy hi","finish_reason":"MAX_TOKENS"}],"prompt":""}"#, - 12, - 4, - )), + #[tokio::test] + async fn titan_rejects_a_batch_the_vendor_cannot_serve() { + let r = embed_req("amazon.titan-embed-text-v2:0", &["one", "two"], None); + assert!(AwsEmbedEngine::new(r, t()).run().await.is_err()); + } + + #[tokio::test] + async fn aws_embed_rejects_a_chat_request() { + let out = AwsEmbedEngine::new(req(Protocol::AwsEmbed, "amazon.titan-embed-text-v2:0"), t()) + .run() + .await; + assert!(out.is_err()); + } + + #[tokio::test] + async fn bedrock_headers_bill_the_invoke() { + let mut e = AwsEmbedEngine::new( + embed_req("cohere.embed-english-v3", &["only text"], None), + Arc::new(BedrockReply(r#"{"embeddings":[[0.5,0.25]]}"#, 57, 0)), ); let out = e.run().await.unwrap(); - assert_eq!(out.response.message, "legacy hi"); - assert_eq!(out.response.finish_reason, "length"); - assert_eq!(out.response.total_tokens, 16); + assert_eq!(out.response.prompt_tokens, 57); + assert_eq!(out.response.total_tokens, 57); let mut e = LlamaEngine::new( req(Protocol::AwsLlama, "meta.llama3-1-8b-instruct-v1:0"), @@ -681,19 +648,6 @@ mod tests { assert_eq!(out.response.finish_reason, "stop"); assert!(out.chunks.iter().any(|c| c.finish_reason.is_some())); assert!(out.response.total_tokens > 0); - - let mut r = req(Protocol::AwsCohere, "cohere.command-r-v1:0"); - r.stream = true; - let out = CohereEngine::new(r, t()).run().await.unwrap(); - assert!( - out.response - .message - .contains("[mock-cohere] you said: hello bespoke"), - "{:?}", - out.response - ); - assert_eq!(out.response.finish_reason, "stop"); - assert!(out.response.prompt_tokens > 0 && out.response.completion_tokens > 0); } #[tokio::test] diff --git a/crates/engines/src/factory.rs b/crates/engines/src/factory.rs index 28e2c99..aca70e9 100644 --- a/crates/engines/src/factory.rs +++ b/crates/engines/src/factory.rs @@ -6,7 +6,7 @@ use gw_consts::{ErrCode, Protocol}; use gw_models::{GResult, GatewayError, GatewayRequest}; -use crate::bespoke::{CohereEngine, DashScopeEngine, ErnieEngine, LlamaEngine, MinimaxV1Engine}; +use crate::bespoke::{AwsEmbedEngine, DashScopeEngine, ErnieEngine, LlamaEngine, MinimaxV1Engine}; use crate::claude_engine::ClaudeEngine; use crate::engine::ModelEngine; use crate::families::{ @@ -44,7 +44,7 @@ pub fn get_engine( Protocol::Passthrough => Box::new(PassthroughEngine::new(request, transport)), Protocol::Ernie => Box::new(ErnieEngine::new(request, transport)), Protocol::MinimaxV1 => Box::new(MinimaxV1Engine::new(request, transport)), - Protocol::AwsCohere => Box::new(CohereEngine::new(request, transport)), + Protocol::AwsEmbed => Box::new(AwsEmbedEngine::new(request, transport)), Protocol::AwsLlama => Box::new(LlamaEngine::new(request, transport)), Protocol::Dashscope => Box::new(DashScopeEngine::new(request, transport)), Protocol::Realtime => { diff --git a/crates/engines/src/lib.rs b/crates/engines/src/lib.rs index 36dd1e9..bf2c672 100644 --- a/crates/engines/src/lib.rs +++ b/crates/engines/src/lib.rs @@ -26,7 +26,7 @@ pub mod usage_extract; mod base; -pub use bespoke::{CohereEngine, DashScopeEngine, ErnieEngine, LlamaEngine, MinimaxV1Engine}; +pub use bespoke::{AwsEmbedEngine, DashScopeEngine, ErnieEngine, LlamaEngine, MinimaxV1Engine}; pub use claude_engine::{ClaudeEngine, anthropic_native_chunks}; pub use engine::{EngineOutcome, ModelEngine, StreamChunk}; pub use factory::get_engine; diff --git a/crates/engines/src/transport.rs b/crates/engines/src/transport.rs index b41c7f9..4b770ed 100644 --- a/crates/engines/src/transport.rs +++ b/crates/engines/src/transport.rs @@ -425,7 +425,7 @@ impl MockTransport { } let reply = match req.protocol { Protocol::AwsAnthropic => self.anthropic_reply(req)?, - Protocol::AwsCohere => self.cohere_reply(req)?, + Protocol::AwsEmbed => self.embed_reply(req)?, Protocol::AwsLlama => self.llama_reply(req)?, _ => return Err(GatewayError::internal("mock bedrock protocol")), }; @@ -562,16 +562,20 @@ impl MockTransport { }) } - fn cohere_reply(&self, req: &UpstreamRequest) -> GResult { - let body = Self::parse(&req.body, "cohere")?; - let user = body["message"].as_str().unwrap_or_default(); - let reply = format!("[mock-cohere] you said: {user}"); - Self::ok_json(json!({ - "response_id": "cohere-mock", "generation_id": "gen-mock", - "text": reply, "finish_reason": "COMPLETE", - "meta": {"tokens": {"input_tokens": Self::tokens(user) + 3, - "output_tokens": Self::tokens(&reply)}} - })) + fn embed_reply(&self, req: &UpstreamRequest) -> GResult { + let body = Self::parse(&req.body, "embed")?; + if let Some(text) = body["inputText"].as_str() { + return Self::ok_json(json!({"embedding": [0.1, 0.2, 0.3], + "inputTextTokenCount": Self::tokens(text)})); + } + let n = body["texts"].as_array().map_or(0, Vec::len); + let mut reply = Self::ok_json(json!({"embeddings": vec![vec![0.1, 0.2, 0.3]; n]}))?; + // real embed replies carry ONLY the input-count header, never the output one + reply.headers.insert( + "x-amzn-bedrock-input-token-count", + reqwest::header::HeaderValue::from(n as u64 * 4), + ); + Ok(reply) } fn llama_reply(&self, req: &UpstreamRequest) -> GResult { @@ -1035,8 +1039,6 @@ impl Transport for MockTransport { self.ernie_reply(&req) } else if u.contains("minimax") { self.minimax_reply(&req) - } else if u.contains("cohere") { - self.cohere_reply(&req) } else if u.contains("meta.llama") { self.llama_reply(&req) } else if u.contains("/messages") { diff --git a/crates/engines/tests/request_construction.rs b/crates/engines/tests/request_construction.rs index 50d77f4..7c8d830 100644 --- a/crates/engines/tests/request_construction.rs +++ b/crates/engines/tests/request_construction.rs @@ -10,7 +10,7 @@ use async_trait::async_trait; use gw_consts::Protocol; use gw_engines::transport::{Transport, UpstreamBody, UpstreamRequest, UpstreamResponse}; use gw_engines::{ - AudioEngine, AudioKind, ClaudeEngine, CohereEngine, CompletionsEngine, DashScopeEngine, + AudioEngine, AudioKind, AwsEmbedEngine, ClaudeEngine, CompletionsEngine, DashScopeEngine, EmbeddingsEngine, ErnieEngine, ImageEngine, LlamaEngine, MinimaxV1Engine, ModelEngine, OpenAiEngine, ResponsesEngine, SearchEngine, VertexEngine, VideoEngine, }; @@ -338,10 +338,12 @@ async fn go_live_seam_aws_sigv4_uses_real_credentials() { std::env::set_var("GW_TEST_AWS_SK", "realsecretkeyvalue"); } - let t = RecordingTransport::new( - r#"{"text":"ok","meta":{"tokens":{"input_tokens":1,"output_tokens":1}}}"#, - ); - let mut req = chat_req(Protocol::AwsCohere, "cohere.command-r"); + let t = RecordingTransport::new(r#"{"embeddings":[[0.5]]}"#); + let mut req = chat_req(Protocol::AwsEmbed, "cohere.embed-english-v3"); + req.model_param_v2.as_mut().unwrap().typed = Some(TypedParams::Embeddings(EmbeddingParams { + input: vec!["hello".into()], + dimensions: None, + })); req.account = Some(std::sync::Arc::new(Account { name: "live-bedrock".into(), endpoint: "https://bedrock-runtime.eu-west-1.amazonaws.com".into(), @@ -349,7 +351,7 @@ async fn go_live_seam_aws_sigv4_uses_real_credentials() { secret_key_env: "GW_TEST_AWS_SK".into(), ..Default::default() })); - let _ = CohereEngine::new(req, t.clone()).run().await.unwrap(); + let _ = AwsEmbedEngine::new(req, t.clone()).run().await.unwrap(); assert!( t.url() .starts_with("https://bedrock-runtime.eu-west-1.amazonaws.com/model/"), @@ -364,7 +366,7 @@ async fn go_live_seam_aws_sigv4_uses_real_credentials() { "SigV4 must sign with the real access key in the endpoint's region, got: {auth}" ); assert!( - t.url().ends_with("/model/cohere.command-r/invoke"), + t.url().ends_with("/model/cohere.embed-english-v3/invoke"), "url: {}", t.url() ); @@ -545,27 +547,6 @@ async fn system_prompt_reaches_every_bespoke_wire() { .all(|m| m["sender_type"] != "USER" || m["text"] != "be brief"), "system must not be downgraded to a USER turn: {mb}" ); - - let cohere = RecordingTransport::new( - r#"{"text":"ok","finish_reason":"COMPLETE","meta":{"tokens":{"input_tokens":1,"output_tokens":1}}}"#, - ); - CohereEngine::new( - chat_req(Protocol::AwsCohere, "cohere.command-r-v1:0"), - cohere.clone(), - ) - .run() - .await - .unwrap(); - let cb = cohere.body_json(); - assert_eq!(cb["preamble"], "be brief", "cohere system slot"); - assert!( - cb["chat_history"] - .as_array() - .unwrap() - .iter() - .all(|m| m["message"] != "be brief"), - "system must not leak into chat_history: {cb}" - ); } #[tokio::test] @@ -628,20 +609,17 @@ async fn minimax_v1_request_shape() { } #[tokio::test] -async fn cohere_request_shape_with_sigv4() { - let t = RecordingTransport::new( - r#"{"text":"ok","finish_reason":"COMPLETE","meta":{"tokens":{"input_tokens":1,"output_tokens":1}}}"#, - ); - let _ = CohereEngine::new( - chat_req(Protocol::AwsCohere, "cohere.command-r-v1:0"), - t.clone(), - ) - .run() - .await - .unwrap(); +async fn embed_request_shape_with_sigv4() { + let t = RecordingTransport::new(r#"{"embeddings":[[0.5],[0.25]]}"#); + let mut req = chat_req(Protocol::AwsEmbed, "cohere.embed-english-v3"); + req.model_param_v2.as_mut().unwrap().typed = Some(TypedParams::Embeddings(EmbeddingParams { + input: vec!["first".into(), "second".into()], + dimensions: None, + })); + let _ = AwsEmbedEngine::new(req, t.clone()).run().await.unwrap(); let b = t.body_json(); - assert_eq!(b["message"], "hello"); - assert!(b["chat_history"].is_array()); + assert_eq!(b["texts"], serde_json::json!(["first", "second"])); + assert_eq!(b["input_type"], "search_document"); let auth = t.header("authorization").expect("SigV4 auth header"); assert!( auth.starts_with("AWS4-HMAC-SHA256 Credential="), diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 99bb9ae..24a59af 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -1909,7 +1909,9 @@ async fn admin_config_get(State(s): State, headers: HeaderMap) -> Resp }; match store.load_latest().await { Ok(Some((version, yaml))) => { - Json(json!({ "version": version, "yaml": yaml })).into_response() + let mut out = json!({ "version": version }); + out["yaml"] = Value::String(yaml); + Json(out).into_response() } Ok(None) => error_response(404, "config store is empty"), Err(e) => gateway_error(e), @@ -2165,7 +2167,9 @@ async fn admin_models_status(State(s): State, scope: AdminScope) -> Re }) }) .collect(); - Json(json!({ "models": rows })).into_response() + let mut out = json!({}); + out["models"] = Value::Array(rows); + Json(out).into_response() } /// GET /admin/usage — ledger rollup by (tenant, requested model). A tenant @@ -2294,13 +2298,9 @@ async fn admin_usage_series( "vendor_cost_micros": if redact_vendor { 0 } else { totals.vendor_cost_micros }, })); } - Json(json!({ - "bucket": bucket_name, - "since": since, - "until": until, - "series": series, - })) - .into_response() + let mut out = json!({"bucket": bucket_name, "since": since, "until": until}); + out["series"] = Value::Array(series); + Json(out).into_response() } /// A CSV field, RFC-4180 quoted and neutralized against spreadsheet formula diff --git a/docs/providers.md b/docs/providers.md index 02c7856..51bcb92 100644 --- a/docs/providers.md +++ b/docs/providers.md @@ -76,7 +76,7 @@ usage); the rest are marked non-streaming below and always answer buffered: | `ernie` | Baidu Ernie (Wenxin) | `https://aip.baidubce.com` | a `bce-v3/…` key goes as `Bearer`, a legacy token as the `access_token` query param (non-streaming); Qianfan's OpenAI-compatible `https://qianfan.baidubce.com/v2` also works as `kind: openai` + `endpoint` | | `aws-anthropic` | Anthropic Claude on AWS Bedrock | `https://bedrock-runtime..amazonaws.com` | SigV4 (see below); model name = the Bedrock model id (`anthropic.claude-…`, `us.anthropic.claude-…`); the full Messages engine (system, tools, thinking dialects by generation, prompt-cache breakpoints, signed reasoning) on the InvokeModel wire — `anthropic_version` in the body, model and streaming in the path; streams via InvokeModelWithResponseStream (EventStream frames decoded into the same event sequence) | | `aws-converse` | any model on AWS Bedrock via the Converse API | `https://bedrock-runtime..amazonaws.com` | SigV4 or API key (see below); model name = the Bedrock model id or inference profile (`eu.amazon.nova-micro-v1:0`, `us.meta.llama3-3-70b-instruct-v1:0`, `mistral.pixtral-large-2502-v1:0`, `anthropic.claude-…`); the Messages engine transcoded to Converse — system, tools + tool results, images, thinking replay, prompt-cache points — and back (buffered and `converse-stream`); Claude reasoning knobs ride in `additionalModelRequestFields`, other passthrough extras too | -| `aws-cohere` | Cohere Command on AWS Bedrock | `https://bedrock-runtime..amazonaws.com` | SigV4 (see below); model name = the Bedrock model id; Command R (`{message, chat_history, preamble}` → `text`) and the legacy Command `generations[]` shape; usage from Bedrock's `x-amzn-bedrock-*-token-count` headers, or `amazon-bedrock-invocationMetrics` on a stream | +| `aws-embed` | Titan / Cohere embeddings on AWS Bedrock | `https://bedrock-runtime..amazonaws.com` | SigV4 or API-key Bearer (see below); model name = the Bedrock model id; Titan `{inputText}` → `{embedding, inputTextTokenCount}` takes exactly one input per call (a batch is refused with 400; `dimensions` forwarded when the client sends it), Cohere `{texts, input_type: search_document}` → `{embeddings}` in one call; answered in the OpenAI `/v1/embeddings` list shape, usage from Bedrock's `x-amzn-bedrock-*-token-count` headers | | `aws-llama` | Meta Llama on AWS Bedrock | `https://bedrock-runtime..amazonaws.com` | SigV4 (see below); model name = the Bedrock model id or inference profile (`meta.llama3-8b-instruct-v1:0`, `us.meta.llama3-3-70b-instruct-v1:0`, `us.meta.llama4-scout-17b-instruct-v1:0`); the conversation is rendered into the Llama 3 (or Llama 4) chat template; usage from the token-count headers / invocation metrics, else the body counts | | `minimax-v1` | MiniMax legacy v1 (`abab*`) | `https://api.minimax.chat` | `Bearer` (non-streaming); kept for existing accounts — the vendor has retired it for new ones; new integrations should use MiniMax's OpenAI-/Anthropic-compatible endpoints | @@ -148,10 +148,10 @@ the native `/v1/messages` stream); the SigV4 path is verified up to an accepted signature. Vendor limits seen on that wire: Bedrock's Llama answers a tool request as JSON text rather than a `toolUse` block, Qwen rejects `stopSequences`, non-Claude models reject `strict` (the gateway forwards it -only to Claude). Cohere Command R on Bedrock is verified against the -[ministack](https://github.com/ministackorg/ministack) emulator's -family-faithful InvokeModel replies, EventStream framing and token-count -headers only — AWS marks Command R / R+ legacy and gates them per account. +only to Claude). Cohere Command on Bedrock is not carried: AWS answered live +probes with end-of-life for the legacy Command models and legacy-gates +Command R per account, so the former `aws-cohere` chat protocol was removed. +Titan v2 and Cohere v3 embeddings are live-verified through `aws-embed`. `GW_TRANSPORT` overrides transport routing: unset (or any value other than `mock`/`http`) routes `mock://` sentinel URLs in-process and real URLs over @@ -191,7 +191,7 @@ console session it was minted in. accounts: - {name: bedrock, provider: aws, endpoint: "https://bedrock-runtime.us-east-1.amazonaws.com", api_key_env: AWS_ACCESS_KEY_ID, secret_key_env: AWS_SECRET_ACCESS_KEY, - protocols: ["aws-anthropic", "aws-llama", "aws-cohere"]} + protocols: ["aws-anthropic", "aws-llama", "aws-embed"]} - {name: bedrock-eu, provider: aws, endpoint: "https://bedrock-runtime.eu-north-1.amazonaws.com", api_key_env: AWS_BEARER_TOKEN_BEDROCK, protocols: ["aws-anthropic", "aws-converse"]} models: diff --git a/scripts/live-matrix/live.yaml b/scripts/live-matrix/live.yaml index f1aed63..e94fdc2 100644 --- a/scripts/live-matrix/live.yaml +++ b/scripts/live-matrix/live.yaml @@ -104,6 +104,9 @@ models: - {name: "us.anthropic.claude-haiku-4-5-20251001-v1:0", protocol: aws-anthropic, input_price_per_1k_micros: 1000, output_price_per_1k_micros: 5000, token_rate: {read_cache: 0.1, write_cache: 1.25}} - {name: "us.anthropic.claude-sonnet-4-5-20250929-v1:0", protocol: aws-converse, input_price_per_1k_micros: 3000, output_price_per_1k_micros: 15000, token_rate: {read_cache: 0.1, write_cache: 1.25}, prompt_cache: true} - {name: "us.amazon.nova-lite-v1:0", protocol: aws-converse, input_price_per_1k_micros: 60, output_price_per_1k_micros: 240} + - {name: "mistral.mistral-large-2402-v1:0", protocol: aws-converse, input_price_per_1k_micros: 4000, output_price_per_1k_micros: 12000} + - {name: "amazon.titan-embed-text-v2:0", protocol: aws-embed, input_price_per_1k_micros: 20, output_price_per_1k_micros: 0} + - {name: "cohere.embed-english-v3", protocol: aws-embed, input_price_per_1k_micros: 100, output_price_per_1k_micros: 0} - {name: "us.deepseek.r1-v1:0", protocol: aws-converse, input_price_per_1k_micros: 1350, output_price_per_1k_micros: 5400} - {name: "us.meta.llama3-1-8b-instruct-v1:0", protocol: aws-llama, input_price_per_1k_micros: 220, output_price_per_1k_micros: 220} @@ -125,7 +128,7 @@ accounts: protocols: ["search"] - {name: cohere, provider: cohere, endpoint: https://api.cohere.com, api_key_env: COHERE_API_KEY, protocols: ["rerank"], cost_unit_price_micros: 1000} - {name: jina, provider: jina, endpoint: https://api.jina.ai, api_key_env: JINA_API_KEY, protocols: ["rerank"]} - - {name: bedrock, provider: aws, endpoint: https://bedrock-runtime.us-east-1.amazonaws.com, api_key_env: AWS_BEARER_TOKEN_BEDROCK, protocols: ["aws-anthropic", "aws-converse", "aws-llama"], timeout_seconds: 120} + - {name: bedrock, provider: aws, endpoint: https://bedrock-runtime.us-east-1.amazonaws.com, api_key_env: AWS_BEARER_TOKEN_BEDROCK, protocols: ["aws-anthropic", "aws-converse", "aws-llama", "aws-embed"], timeout_seconds: 120} security: dlp_redact: false diff --git a/scripts/live-matrix/live_matrix.py b/scripts/live-matrix/live_matrix.py index b8a4a4f..27c5f2f 100644 --- a/scripts/live-matrix/live_matrix.py +++ b/scripts/live-matrix/live_matrix.py @@ -333,10 +333,11 @@ def case_messages( record(name + " [thinking present]", thinking_blocks > 0, note) -def case_embeddings(gw: Gateway, model: str) -> None: +def case_embeddings(gw: Gateway, model: str, inputs: list[str] | None = None) -> None: + """Titan takes exactly one input per invoke; every other embedder batches.""" name = f"{model} embeddings" before, _ = gw.ledger() - st, txt = gw.call("/v1/embeddings", {"model": model, "input": ["gateway live test", "second input"]}) + st, txt = gw.call("/v1/embeddings", {"model": model, "input": inputs or ["gateway live test", "second input"]}) if st != 200: record(name, False, f"HTTP {st}: {txt[:200]}") return @@ -1157,6 +1158,10 @@ def run_group(gw: Gateway, group: str) -> None: case_thinking_replay(gw, sonnet, native=True) case_chat(gw, "us.amazon.nova-lite-v1:0", "converse chat") case_chat(gw, "us.amazon.nova-lite-v1:0", "converse chat", stream=True) + case_chat(gw, "mistral.mistral-large-2402-v1:0", "converse chat") + case_chat(gw, "mistral.mistral-large-2402-v1:0", "converse chat", stream=True) + case_embeddings(gw, "amazon.titan-embed-text-v2:0", ["gateway live test"]) + case_embeddings(gw, "cohere.embed-english-v3") case_chat(gw, "us.deepseek.r1-v1:0", "converse reasoning", prompt=prime, expect_reasoning=True) case_chat(gw, "us.meta.llama3-1-8b-instruct-v1:0", "aws-llama") case_chat(gw, "us.meta.llama3-1-8b-instruct-v1:0", "aws-llama", stream=True)