mentedb_embedding/
bedrock_provider.rs

1//! AWS Bedrock embedding provider (Amazon Titan and Cohere Embed) over the
2//! InvokeModel REST API.
3//!
4//! Deliberately dependency light: a synchronous `ureq` call plus a small
5//! hand-rolled SigV4 signer, so the engine gains native Bedrock support without
6//! pulling in the full async AWS SDK. This matches the lean, no-heavy-deps
7//! philosophy of the rest of the engine (see the `http` provider).
8//!
9//! Credentials come from [`AwsCredentialSource`]: static keys from the
10//! environment, or the ECS task role / EKS Pod Identity endpoint, refreshed
11//! before they expire.
12//!
13//! ```no_run
14//! use mentedb_embedding::bedrock_provider::{BedrockEmbeddingConfig, BedrockEmbeddingProvider};
15//! use mentedb_embedding::provider::EmbeddingProvider;
16//!
17//! let cfg = BedrockEmbeddingConfig::titan_v2("us-east-1")?; // reads creds from env
18//! let provider = BedrockEmbeddingProvider::new(cfg)?;
19//! let vector = provider.embed("the sky is blue")?;
20//! # Ok::<(), mentedb_core::error::MenteError>(())
21//! ```
22
23use hmac::{Hmac, Mac};
24use sha2::{Digest, Sha256};
25
26use mentedb_core::error::{MenteError, MenteResult};
27
28use crate::aws_credentials::AwsCredentialSource;
29use crate::provider::EmbeddingProvider;
30
31type HmacSha256 = Hmac<Sha256>;
32
33const SERVICE: &str = "bedrock";
34
35fn hex(bytes: &[u8]) -> String {
36    let mut s = String::with_capacity(bytes.len() * 2);
37    for b in bytes {
38        s.push_str(&format!("{b:02x}"));
39    }
40    s
41}
42
43fn sha256_hex(data: &[u8]) -> String {
44    hex(&Sha256::digest(data))
45}
46
47fn hmac_sha256(key: &[u8], data: &[u8]) -> Vec<u8> {
48    let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts any key length");
49    mac.update(data);
50    mac.finalize().into_bytes().to_vec()
51}
52
53/// SigV4 signing key: HMAC chain over date, region, service, "aws4_request".
54fn signing_key(secret: &str, datestamp: &str, region: &str, service: &str) -> Vec<u8> {
55    let k_date = hmac_sha256(format!("AWS4{secret}").as_bytes(), datestamp.as_bytes());
56    let k_region = hmac_sha256(&k_date, region.as_bytes());
57    let k_service = hmac_sha256(&k_region, service.as_bytes());
58    hmac_sha256(&k_service, b"aws4_request")
59}
60
61/// URI-encode a single path segment per SigV4 rules (unreserved chars pass
62/// through; everything else, including the `:` in a model id, is percent
63/// encoded). The request URL and the signed canonical URI must match exactly.
64fn uri_encode_segment(s: &str) -> String {
65    let mut out = String::with_capacity(s.len());
66    for b in s.bytes() {
67        if b.is_ascii_alphanumeric() || matches!(b, b'-' | b'.' | b'_' | b'~') {
68            out.push(b as char);
69        } else {
70            out.push_str(&format!("%{b:02X}"));
71        }
72    }
73    out
74}
75
76/// Most texts Cohere Embed accepts in one InvokeModel call.
77const COHERE_MAX_BATCH: usize = 96;
78
79/// Request and response shape of a Bedrock embedding model.
80#[derive(Clone, Copy, Debug, PartialEq, Eq)]
81pub enum BedrockEmbeddingFamily {
82    /// Amazon Titan: one text per call.
83    Titan,
84    /// Cohere Embed: up to 96 texts per call.
85    Cohere,
86}
87
88impl BedrockEmbeddingFamily {
89    /// Family for a model id, by its provider prefix.
90    pub fn for_model(model_id: &str) -> MenteResult<Self> {
91        let bare = model_id.rsplit('/').next().unwrap_or(model_id);
92        let bare = bare
93            .strip_prefix("us.")
94            .or_else(|| bare.strip_prefix("eu."))
95            .or_else(|| bare.strip_prefix("apac."))
96            .or_else(|| bare.strip_prefix("global."))
97            .unwrap_or(bare);
98        if bare.starts_with("amazon.titan-embed") {
99            Ok(Self::Titan)
100        } else if bare.starts_with("cohere.embed") {
101            Ok(Self::Cohere)
102        } else {
103            Err(MenteError::Storage(format!(
104                "unsupported Bedrock embedding model '{model_id}' (expected amazon.titan-embed-* \
105                 or cohere.embed-*)"
106            )))
107        }
108    }
109}
110
111/// Configuration for the Bedrock embedding provider.
112#[derive(Clone)]
113pub struct BedrockEmbeddingConfig {
114    pub region: String,
115    pub model_id: String,
116    pub family: BedrockEmbeddingFamily,
117    pub dimensions: usize,
118    pub credentials: AwsCredentialSource,
119    /// Whole call timeout for one embedding request. An embed sits on the
120    /// write and search paths, so a stalled call must fail fast rather than
121    /// hold a request open. Default 3 seconds.
122    pub timeout: std::time::Duration,
123}
124
125/// Default whole call timeout for a Bedrock embedding request.
126pub const DEFAULT_EMBED_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(3);
127
128impl BedrockEmbeddingConfig {
129    /// Any supported model. `dimensions` must match what the model returns.
130    pub fn new(
131        region: impl Into<String>,
132        model_id: impl Into<String>,
133        dimensions: usize,
134        credentials: AwsCredentialSource,
135    ) -> MenteResult<Self> {
136        let model_id = model_id.into();
137        Ok(Self {
138            region: region.into(),
139            family: BedrockEmbeddingFamily::for_model(&model_id)?,
140            model_id,
141            dimensions,
142            credentials,
143            timeout: DEFAULT_EMBED_TIMEOUT,
144        })
145    }
146
147    /// Use `timeout` as the whole call timeout (zero keeps the default).
148    pub fn with_timeout(mut self, timeout: std::time::Duration) -> Self {
149        if !timeout.is_zero() {
150            self.timeout = timeout;
151        }
152        self
153    }
154
155    /// A known model with its native dimensions, credentials from the
156    /// environment or the container role.
157    pub fn for_model(region: impl Into<String>, model_id: &str) -> MenteResult<Self> {
158        let dimensions = match model_id {
159            "amazon.titan-embed-text-v2:0" => 1024,
160            "amazon.titan-embed-text-v1" => 1536,
161            "cohere.embed-english-v3" | "cohere.embed-multilingual-v3" => 1024,
162            other => {
163                return Err(MenteError::Storage(format!(
164                    "unknown dimensions for Bedrock embedding model '{other}'; use \
165                     BedrockEmbeddingConfig::new with explicit dimensions"
166                )));
167            }
168        };
169        Self::new(
170            region,
171            model_id,
172            dimensions,
173            AwsCredentialSource::from_env()?,
174        )
175    }
176
177    /// Amazon Titan Text Embeddings V2 (1024 dimensions).
178    pub fn titan_v2(region: impl Into<String>) -> MenteResult<Self> {
179        Self::for_model(region, "amazon.titan-embed-text-v2:0")
180    }
181
182    /// Amazon Titan Text Embeddings V1 (1536 dimensions).
183    pub fn titan_v1(region: impl Into<String>) -> MenteResult<Self> {
184        Self::for_model(region, "amazon.titan-embed-text-v1")
185    }
186
187    /// Cohere Embed English v3 (1024 dimensions).
188    pub fn cohere_english_v3(region: impl Into<String>) -> MenteResult<Self> {
189        Self::for_model(region, "cohere.embed-english-v3")
190    }
191}
192
193/// The HTTP agent embedding calls go through, bounded by `timeout` per call.
194fn http_agent(timeout: std::time::Duration) -> ureq::Agent {
195    ureq::Agent::config_builder()
196        .timeout_global(Some(timeout))
197        .build()
198        .into()
199}
200
201/// Synchronous Bedrock embedding provider.
202pub struct BedrockEmbeddingProvider {
203    config: BedrockEmbeddingConfig,
204    agent: ureq::Agent,
205}
206
207impl BedrockEmbeddingProvider {
208    pub fn new(config: BedrockEmbeddingConfig) -> MenteResult<Self> {
209        let agent = http_agent(config.timeout);
210        Ok(Self { config, agent })
211    }
212
213    /// Embed up to one call's worth of texts.
214    fn invoke(&self, texts: &[&str]) -> MenteResult<Vec<Vec<f32>>> {
215        let body = match self.config.family {
216            BedrockEmbeddingFamily::Titan => serde_json::json!({ "inputText": texts[0] }),
217            // Stored memories and queries are embedded the same way, so the
218            // document input type is used for both, as on the cloud platform.
219            BedrockEmbeddingFamily::Cohere => serde_json::json!({
220                "texts": texts,
221                "input_type": "search_document",
222            }),
223        };
224        let value = self.post(&body)?;
225        let vectors = match self.config.family {
226            BedrockEmbeddingFamily::Titan => {
227                let r: TitanResponse = serde_json::from_value(value).map_err(|e| {
228                    MenteError::Storage(format!("bedrock response parse failed: {e}"))
229                })?;
230                vec![r.embedding]
231            }
232            BedrockEmbeddingFamily::Cohere => {
233                let r: CohereResponse = serde_json::from_value(value).map_err(|e| {
234                    MenteError::Storage(format!("bedrock response parse failed: {e}"))
235                })?;
236                r.embeddings
237            }
238        };
239        if vectors.len() != texts.len() {
240            return Err(MenteError::Storage(format!(
241                "bedrock returned {} embeddings for {} texts",
242                vectors.len(),
243                texts.len()
244            )));
245        }
246        Ok(vectors)
247    }
248
249    /// Sign and send one InvokeModel request.
250    fn post(&self, body: &serde_json::Value) -> MenteResult<serde_json::Value> {
251        let cfg = &self.config;
252        let creds = cfg.credentials.credentials()?;
253        let host = format!("bedrock-runtime.{}.amazonaws.com", cfg.region);
254        // The request is sent with the raw model id in the path, while the SigV4
255        // canonical URI uses the percent-encoded form. AWS re-encodes the path it
256        // receives the same way, so signing the encoded form while sending the raw
257        // form matches, which is what the AWS SDKs do for model ids containing ':'.
258        let canonical_uri = format!("/model/{}/invoke", uri_encode_segment(&cfg.model_id));
259        let url = format!("https://{host}/model/{}/invoke", cfg.model_id);
260
261        let body =
262            serde_json::to_vec(body).map_err(|e| MenteError::Serialization(e.to_string()))?;
263        let payload_hash = sha256_hex(&body);
264
265        let now = chrono::Utc::now();
266        let amzdate = now.format("%Y%m%dT%H%M%SZ").to_string();
267        let datestamp = now.format("%Y%m%d").to_string();
268
269        // Canonical (sorted, lowercase) headers that we sign and send.
270        let mut headers: Vec<(String, String)> = vec![
271            ("host".to_string(), host.clone()),
272            ("x-amz-content-sha256".to_string(), payload_hash.clone()),
273            ("x-amz-date".to_string(), amzdate.clone()),
274        ];
275        if let Some(token) = &creds.session_token {
276            headers.push(("x-amz-security-token".to_string(), token.clone()));
277        }
278        headers.sort_by(|a, b| a.0.cmp(&b.0));
279        let canonical_headers: String = headers.iter().map(|(k, v)| format!("{k}:{v}\n")).collect();
280        let signed_headers = headers
281            .iter()
282            .map(|(k, _)| k.as_str())
283            .collect::<Vec<_>>()
284            .join(";");
285
286        let canonical_request = format!(
287            "POST\n{canonical_uri}\n\n{canonical_headers}\n{signed_headers}\n{payload_hash}"
288        );
289        let scope = format!("{datestamp}/{}/{SERVICE}/aws4_request", cfg.region);
290        let string_to_sign = format!(
291            "AWS4-HMAC-SHA256\n{amzdate}\n{scope}\n{}",
292            sha256_hex(canonical_request.as_bytes())
293        );
294        let key = signing_key(&creds.secret_access_key, &datestamp, &cfg.region, SERVICE);
295        let signature = hex(&hmac_sha256(&key, string_to_sign.as_bytes()));
296        let authorization = format!(
297            "AWS4-HMAC-SHA256 Credential={}/{scope}, SignedHeaders={signed_headers}, Signature={signature}",
298            creds.access_key_id
299        );
300
301        let mut req = self
302            .agent
303            .post(&url)
304            .header("content-type", "application/json")
305            .header("accept", "application/json")
306            .header("x-amz-date", &amzdate)
307            .header("x-amz-content-sha256", &payload_hash)
308            .header("authorization", &authorization);
309        if let Some(token) = &creds.session_token {
310            req = req.header("x-amz-security-token", token);
311        }
312
313        let mut resp = req
314            .send(body.as_slice())
315            .map_err(|e| MenteError::Storage(format!("bedrock invoke failed: {e}")))?;
316        resp.body_mut()
317            .read_json()
318            .map_err(|e| MenteError::Storage(format!("bedrock response parse failed: {e}")))
319    }
320}
321
322#[derive(serde::Deserialize)]
323struct TitanResponse {
324    embedding: Vec<f32>,
325}
326
327#[derive(serde::Deserialize)]
328struct CohereResponse {
329    embeddings: Vec<Vec<f32>>,
330}
331
332impl EmbeddingProvider for BedrockEmbeddingProvider {
333    fn embed(&self, text: &str) -> MenteResult<Vec<f32>> {
334        self.invoke(&[text])?
335            .pop()
336            .ok_or_else(|| MenteError::Storage("bedrock returned no embedding".to_string()))
337    }
338
339    fn embed_batch(&self, texts: &[&str]) -> MenteResult<Vec<Vec<f32>>> {
340        let per_call = match self.config.family {
341            BedrockEmbeddingFamily::Titan => 1,
342            BedrockEmbeddingFamily::Cohere => COHERE_MAX_BATCH,
343        };
344        let mut out = Vec::with_capacity(texts.len());
345        for chunk in texts.chunks(per_call) {
346            out.extend(self.invoke(chunk)?);
347        }
348        Ok(out)
349    }
350
351    fn dimensions(&self) -> usize {
352        self.config.dimensions
353    }
354
355    fn model_name(&self) -> &str {
356        &self.config.model_id
357    }
358}
359
360#[cfg(test)]
361mod tests {
362    use super::*;
363
364    #[test]
365    fn embed_calls_give_up_at_the_timeout() {
366        // A server that accepts and never answers, like a stalled endpoint.
367        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
368        let port = listener.local_addr().unwrap().port();
369        let _hold = std::thread::spawn(move || {
370            let conns: Vec<_> = listener.incoming().take(1).collect();
371            std::thread::sleep(std::time::Duration::from_secs(10));
372            drop(conns);
373        });
374        let agent = http_agent(std::time::Duration::from_millis(300));
375        let started = std::time::Instant::now();
376        let result = agent.get(&format!("http://127.0.0.1:{port}/")).call();
377        assert!(result.is_err());
378        assert!(
379            started.elapsed() < std::time::Duration::from_secs(3),
380            "took {:?}",
381            started.elapsed()
382        );
383        assert_eq!(DEFAULT_EMBED_TIMEOUT, std::time::Duration::from_secs(3));
384    }
385
386    /// Verify the SigV4 signing chain against AWS's published worked example
387    /// (docs: "Examples of the complete Signature Version 4 signing process").
388    /// This proves the HMAC key derivation and final signature are correct
389    /// without needing a live AWS call.
390    #[test]
391    fn sigv4_matches_aws_reference_vector() {
392        let secret = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY";
393        let datestamp = "20150830";
394        let region = "us-east-1";
395        let service = "iam";
396        let string_to_sign = "AWS4-HMAC-SHA256\n20150830T123600Z\n\
397             20150830/us-east-1/iam/aws4_request\n\
398             f536975d06c0309214f805bb90ccff089219ecd68b2577efef23edd43b7e1a59";
399        let key = signing_key(secret, datestamp, region, service);
400        let signature = hex(&hmac_sha256(&key, string_to_sign.as_bytes()));
401        assert_eq!(
402            signature,
403            "5d672d79c15b13162d9279b0855cfba6789a8edb4c82c400e06b5924a6f2b5d7"
404        );
405    }
406
407    #[test]
408    fn model_id_path_segment_is_percent_encoded() {
409        // The ':' in a Bedrock model id must be encoded so the request URL and
410        // the signed canonical URI match.
411        assert_eq!(
412            uri_encode_segment("amazon.titan-embed-text-v2:0"),
413            "amazon.titan-embed-text-v2%3A0"
414        );
415    }
416
417    #[test]
418    fn model_family_follows_the_provider_prefix() {
419        use BedrockEmbeddingFamily::*;
420        assert_eq!(
421            BedrockEmbeddingFamily::for_model("amazon.titan-embed-text-v2:0").unwrap(),
422            Titan
423        );
424        assert_eq!(
425            BedrockEmbeddingFamily::for_model("cohere.embed-english-v3").unwrap(),
426            Cohere
427        );
428        assert_eq!(
429            BedrockEmbeddingFamily::for_model("us.cohere.embed-english-v3").unwrap(),
430            Cohere
431        );
432        assert!(BedrockEmbeddingFamily::for_model("anthropic.claude-haiku-5-5").is_err());
433    }
434
435    #[test]
436    fn sha256_hex_of_empty_is_known_constant() {
437        assert_eq!(
438            sha256_hex(b""),
439            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
440        );
441    }
442
443    /// Live end-to-end call against real Bedrock. Ignored by default (needs AWS
444    /// credentials in the environment). Run with:
445    ///   AWS_REGION=us-east-1 cargo test -p mentedb-embedding --features bedrock \
446    ///     live_titan_embed -- --ignored --nocapture
447    #[test]
448    #[ignore]
449    fn live_titan_embed() {
450        let region = std::env::var("AWS_REGION").unwrap_or_else(|_| "us-east-1".to_string());
451        let cfg = BedrockEmbeddingConfig::titan_v2(region).expect("credentials from env");
452        let provider = BedrockEmbeddingProvider::new(cfg).expect("provider");
453        let v = provider.embed("the sky is blue").expect("live embed");
454        assert_eq!(v.len(), 1024);
455        assert!(v.iter().any(|&x| x != 0.0), "embedding should be non-zero");
456    }
457
458    /// Live Cohere batch call. Ignored by default, like `live_titan_embed`.
459    #[test]
460    #[ignore]
461    fn live_cohere_embed_batch() {
462        let region = std::env::var("AWS_REGION").unwrap_or_else(|_| "us-east-1".to_string());
463        let cfg = BedrockEmbeddingConfig::cohere_english_v3(region).expect("credentials");
464        let provider = BedrockEmbeddingProvider::new(cfg).expect("provider");
465        let texts: Vec<String> = (0..100).map(|i| format!("memory number {i}")).collect();
466        let refs: Vec<&str> = texts.iter().map(|s| s.as_str()).collect();
467        let v = provider.embed_batch(&refs).expect("live batch");
468        assert_eq!(v.len(), 100, "100 texts span two calls of at most 96");
469        assert!(v.iter().all(|e| e.len() == 1024));
470    }
471}
mentedb_embedding/bedrock_provider.rs

mentedb_embedding/
bedrock_provider.rs
