bedrock closeout: aws-embed replaces EOL'd aws-cohere; mistral + embeddings live coverage

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",

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::<i64>().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<i64> {
+    headers
+        .get(name)
+        .and_then(|v| v.to_str().ok())
+        .and_then(|s| s.parse::<i64>().ok())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;