mentedb_embedding/
bedrock_provider.rs

1//! AWS Bedrock embedding provider (Amazon Titan) over the InvokeModel REST API.
2//!
3//! Deliberately dependency light: a synchronous `ureq` call plus a small
4//! hand-rolled SigV4 signer, so the engine gains native Bedrock support without
5//! pulling in the full async AWS SDK. This matches the lean, no-heavy-deps
6//! philosophy of the rest of the engine (see the `http` provider).
7//!
8//! Credentials come from the standard AWS environment variables
9//! (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, and `AWS_SESSION_TOKEN` for
10//! temporary/SSO credentials). Full credential-chain resolution (profiles,
11//! instance metadata) is intentionally out of scope for this provider; export
12//! credentials or use the platform's SDK-backed path for those.
13//!
14//! ```no_run
15//! use mentedb_embedding::bedrock_provider::{BedrockEmbeddingConfig, BedrockEmbeddingProvider};
16//! use mentedb_embedding::provider::EmbeddingProvider;
17//!
18//! let cfg = BedrockEmbeddingConfig::titan_v2("us-east-1")?; // reads creds from env
19//! let provider = BedrockEmbeddingProvider::new(cfg)?;
20//! let vector = provider.embed("the sky is blue")?;
21//! # Ok::<(), mentedb_core::error::MenteError>(())
22//! ```
23
24use hmac::{Hmac, Mac};
25use sha2::{Digest, Sha256};
26
27use mentedb_core::error::{MenteError, MenteResult};
28
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/// AWS credentials for signing, loaded from the standard environment variables.
77#[derive(Clone)]
78pub struct AwsCredentials {
79    pub access_key_id: String,
80    pub secret_access_key: String,
81    pub session_token: Option<String>,
82}
83
84impl AwsCredentials {
85    /// Load from `AWS_ACCESS_KEY_ID` / `AWS_SECRET_ACCESS_KEY` /
86    /// `AWS_SESSION_TOKEN` (the last is required for SSO/temporary credentials).
87    pub fn from_env() -> MenteResult<Self> {
88        let access_key_id = std::env::var("AWS_ACCESS_KEY_ID")
89            .map_err(|_| MenteError::Storage("AWS_ACCESS_KEY_ID not set".to_string()))?;
90        let secret_access_key = std::env::var("AWS_SECRET_ACCESS_KEY")
91            .map_err(|_| MenteError::Storage("AWS_SECRET_ACCESS_KEY not set".to_string()))?;
92        Ok(Self {
93            access_key_id,
94            secret_access_key,
95            session_token: std::env::var("AWS_SESSION_TOKEN").ok(),
96        })
97    }
98}
99
100/// Configuration for the Bedrock Titan embedding provider.
101#[derive(Clone)]
102pub struct BedrockEmbeddingConfig {
103    pub region: String,
104    pub model_id: String,
105    pub dimensions: usize,
106    pub credentials: AwsCredentials,
107}
108
109impl BedrockEmbeddingConfig {
110    /// Amazon Titan Text Embeddings V2 (1024 dimensions), credentials from env.
111    pub fn titan_v2(region: impl Into<String>) -> MenteResult<Self> {
112        Ok(Self {
113            region: region.into(),
114            model_id: "amazon.titan-embed-text-v2:0".to_string(),
115            dimensions: 1024,
116            credentials: AwsCredentials::from_env()?,
117        })
118    }
119
120    /// Amazon Titan Text Embeddings V1 (1536 dimensions), credentials from env.
121    pub fn titan_v1(region: impl Into<String>) -> MenteResult<Self> {
122        Ok(Self {
123            region: region.into(),
124            model_id: "amazon.titan-embed-text-v1".to_string(),
125            dimensions: 1536,
126            credentials: AwsCredentials::from_env()?,
127        })
128    }
129}
130
131/// Synchronous Bedrock (Titan) embedding provider.
132pub struct BedrockEmbeddingProvider {
133    config: BedrockEmbeddingConfig,
134    agent: ureq::Agent,
135}
136
137impl BedrockEmbeddingProvider {
138    pub fn new(config: BedrockEmbeddingConfig) -> MenteResult<Self> {
139        Ok(Self {
140            config,
141            agent: ureq::Agent::new_with_defaults(),
142        })
143    }
144
145    fn invoke(&self, text: &str) -> MenteResult<Vec<f32>> {
146        let cfg = &self.config;
147        let host = format!("bedrock-runtime.{}.amazonaws.com", cfg.region);
148        // The request is sent with the raw model id in the path, while the SigV4
149        // canonical URI uses the percent-encoded form. AWS re-encodes the path it
150        // receives the same way, so signing the encoded form while sending the raw
151        // form matches, which is what the AWS SDKs do for model ids containing ':'.
152        let canonical_uri = format!("/model/{}/invoke", uri_encode_segment(&cfg.model_id));
153        let url = format!("https://{host}/model/{}/invoke", cfg.model_id);
154
155        let body = serde_json::to_vec(&serde_json::json!({ "inputText": text }))
156            .map_err(|e| MenteError::Serialization(e.to_string()))?;
157        let payload_hash = sha256_hex(&body);
158
159        let now = chrono::Utc::now();
160        let amzdate = now.format("%Y%m%dT%H%M%SZ").to_string();
161        let datestamp = now.format("%Y%m%d").to_string();
162
163        // Canonical (sorted, lowercase) headers that we sign and send.
164        let mut headers: Vec<(String, String)> = vec![
165            ("host".to_string(), host.clone()),
166            ("x-amz-content-sha256".to_string(), payload_hash.clone()),
167            ("x-amz-date".to_string(), amzdate.clone()),
168        ];
169        if let Some(token) = &cfg.credentials.session_token {
170            headers.push(("x-amz-security-token".to_string(), token.clone()));
171        }
172        headers.sort_by(|a, b| a.0.cmp(&b.0));
173        let canonical_headers: String = headers.iter().map(|(k, v)| format!("{k}:{v}\n")).collect();
174        let signed_headers = headers
175            .iter()
176            .map(|(k, _)| k.as_str())
177            .collect::<Vec<_>>()
178            .join(";");
179
180        let canonical_request = format!(
181            "POST\n{canonical_uri}\n\n{canonical_headers}\n{signed_headers}\n{payload_hash}"
182        );
183        let scope = format!("{datestamp}/{}/{SERVICE}/aws4_request", cfg.region);
184        let string_to_sign = format!(
185            "AWS4-HMAC-SHA256\n{amzdate}\n{scope}\n{}",
186            sha256_hex(canonical_request.as_bytes())
187        );
188        let key = signing_key(
189            &cfg.credentials.secret_access_key,
190            &datestamp,
191            &cfg.region,
192            SERVICE,
193        );
194        let signature = hex(&hmac_sha256(&key, string_to_sign.as_bytes()));
195        let authorization = format!(
196            "AWS4-HMAC-SHA256 Credential={}/{scope}, SignedHeaders={signed_headers}, Signature={signature}",
197            cfg.credentials.access_key_id
198        );
199
200        let mut req = self
201            .agent
202            .post(&url)
203            .header("content-type", "application/json")
204            .header("accept", "application/json")
205            .header("x-amz-date", &amzdate)
206            .header("x-amz-content-sha256", &payload_hash)
207            .header("authorization", &authorization);
208        if let Some(token) = &cfg.credentials.session_token {
209            req = req.header("x-amz-security-token", token);
210        }
211
212        let mut resp = req
213            .send(body.as_slice())
214            .map_err(|e| MenteError::Storage(format!("bedrock invoke failed: {e}")))?;
215        let parsed: TitanResponse = resp
216            .body_mut()
217            .read_json()
218            .map_err(|e| MenteError::Storage(format!("bedrock response parse failed: {e}")))?;
219        Ok(parsed.embedding)
220    }
221}
222
223#[derive(serde::Deserialize)]
224struct TitanResponse {
225    embedding: Vec<f32>,
226}
227
228impl EmbeddingProvider for BedrockEmbeddingProvider {
229    fn embed(&self, text: &str) -> MenteResult<Vec<f32>> {
230        self.invoke(text)
231    }
232
233    fn embed_batch(&self, texts: &[&str]) -> MenteResult<Vec<Vec<f32>>> {
234        // Titan InvokeModel embeds one input per call; batch by looping.
235        texts.iter().map(|t| self.invoke(t)).collect()
236    }
237
238    fn dimensions(&self) -> usize {
239        self.config.dimensions
240    }
241
242    fn model_name(&self) -> &str {
243        &self.config.model_id
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250
251    /// Verify the SigV4 signing chain against AWS's published worked example
252    /// (docs: "Examples of the complete Signature Version 4 signing process").
253    /// This proves the HMAC key derivation and final signature are correct
254    /// without needing a live AWS call.
255    #[test]
256    fn sigv4_matches_aws_reference_vector() {
257        let secret = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY";
258        let datestamp = "20150830";
259        let region = "us-east-1";
260        let service = "iam";
261        let string_to_sign = "AWS4-HMAC-SHA256\n20150830T123600Z\n\
262             20150830/us-east-1/iam/aws4_request\n\
263             f536975d06c0309214f805bb90ccff089219ecd68b2577efef23edd43b7e1a59";
264        let key = signing_key(secret, datestamp, region, service);
265        let signature = hex(&hmac_sha256(&key, string_to_sign.as_bytes()));
266        assert_eq!(
267            signature,
268            "5d672d79c15b13162d9279b0855cfba6789a8edb4c82c400e06b5924a6f2b5d7"
269        );
270    }
271
272    #[test]
273    fn model_id_path_segment_is_percent_encoded() {
274        // The ':' in a Bedrock model id must be encoded so the request URL and
275        // the signed canonical URI match.
276        assert_eq!(
277            uri_encode_segment("amazon.titan-embed-text-v2:0"),
278            "amazon.titan-embed-text-v2%3A0"
279        );
280    }
281
282    #[test]
283    fn sha256_hex_of_empty_is_known_constant() {
284        assert_eq!(
285            sha256_hex(b""),
286            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
287        );
288    }
289
290    /// Live end-to-end call against real Bedrock. Ignored by default (needs AWS
291    /// credentials in the environment). Run with:
292    ///   AWS_REGION=us-east-1 cargo test -p mentedb-embedding --features bedrock \
293    ///     live_titan_embed -- --ignored --nocapture
294    #[test]
295    #[ignore]
296    fn live_titan_embed() {
297        let region = std::env::var("AWS_REGION").unwrap_or_else(|_| "us-east-1".to_string());
298        let cfg = BedrockEmbeddingConfig::titan_v2(region).expect("credentials from env");
299        let provider = BedrockEmbeddingProvider::new(cfg).expect("provider");
300        let v = provider.embed("the sky is blue").expect("live embed");
301        assert_eq!(v.len(), 1024);
302        assert!(v.iter().any(|&x| x != 0.0), "embedding should be non-zero");
303    }
304}
mentedb_embedding/bedrock_provider.rs

mentedb_embedding/
bedrock_provider.rs
