diff options
| -rw-r--r-- | config.example.yaml | 1 | ||||
| -rw-r--r-- | src/client/bedrock.rs | 10 |
2 files changed, 11 insertions, 0 deletions
diff --git a/config.example.yaml b/config.example.yaml index 9537b0e..7c367d7 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -233,6 +233,7 @@ clients: access_key_id: xxx secret_access_key: xxx region: xxx + session_token: xxx # Optional, only needed for temporary credentials # See https://developers.cloudflare.com/workers-ai/ - type: openai-compatible diff --git a/src/client/bedrock.rs b/src/client/bedrock.rs index 6ccefe0..88d57d9 100644 --- a/src/client/bedrock.rs +++ b/src/client/bedrock.rs @@ -19,6 +19,7 @@ pub struct BedrockConfig { pub access_key_id: Option<String>, pub secret_access_key: Option<String>, pub region: Option<String>, + pub session_token: Option<String>, #[serde(default)] pub models: Vec<ModelData>, pub patch: Option<RequestPatch>, @@ -29,6 +30,7 @@ impl BedrockClient { config_get_fn!(access_key_id, get_access_key_id); config_get_fn!(secret_access_key, get_secret_access_key); config_get_fn!(region, get_region); + config_get_fn!(session_token, get_session_token); pub const PROMPTS: [PromptAction<'static>; 3] = [ ("access_key_id", "AWS Access Key ID", None), @@ -44,6 +46,7 @@ impl BedrockClient { let access_key_id = self.get_access_key_id()?; let secret_access_key = self.get_secret_access_key()?; let region = self.get_region()?; + let session_token = self.get_session_token().ok(); let host = format!("bedrock-runtime.{region}.amazonaws.com"); let model_name = &self.model.real_name(); @@ -70,6 +73,7 @@ impl BedrockClient { access_key_id, secret_access_key, region, + session_token, }, AwsRequest { method: Method::POST, @@ -93,6 +97,7 @@ impl BedrockClient { let access_key_id = self.get_access_key_id()?; let secret_access_key = self.get_secret_access_key()?; let region = self.get_region()?; + let session_token = self.get_session_token().ok(); let host = format!("bedrock-runtime.{region}.amazonaws.com"); let uri = format!("/model/{}/invoke", self.model.real_name()); @@ -121,6 +126,7 @@ impl BedrockClient { access_key_id, secret_access_key, region, + session_token, }, AwsRequest { method: Method::POST, @@ -527,6 +533,7 @@ struct AwsCredentials { access_key_id: String, secret_access_key: String, region: String, + session_token: Option<String>, } #[derive(Debug)] @@ -563,6 +570,9 @@ fn aws_fetch( let date_stamp = amz_date[0..8].to_string(); headers.insert("host".into(), host.clone()); headers.insert("x-amz-date".into(), amz_date.clone()); + if let Some(token) = credentials.session_token.clone() { + headers.insert("x-amz-security-token".into(), token); + } let canonical_headers = headers .iter() |