feat: added endpoint_url to amazon bedrock (#2555)

Added support to endpoint_url when calling amazon bedrock. This is
needed to support private endpoints behind VPCs on aws.

---------

Signed-off-by: Roey Ben Chaim <roey.benhaim@gmail.com>
Co-authored-by: Sam Lijin <sam@boundaryml.com>

Author: roeybc

engine/baml-lib/llm-client/src/clients/aws_bedrock.rs

@@ -22,6 +22,7 @@ pub struct UnresolvedAwsBedrock<Meta> {
     secret_access_key: Option<StringOr>,
     session_token: Option<StringOr>,
     profile: Option<StringOr>,
+    endpoint_url: Option<StringOr>,
     role_selection: UnresolvedRolesSelection,
     allowed_role_metadata: UnresolvedAllowedRoleMetadata,
     supported_request_modes: SupportedRequestModes,
@@ -78,6 +79,7 @@ pub struct ResolvedAwsBedrock {
     pub secret_access_key: Option<ApiKeyWithProvenance>,
     pub session_token: Option<String>,
     pub profile: Option<String>,
+    pub endpoint_url: Option<String>,
     pub inference_config: Option<InferenceConfiguration>,
     role_selection: RolesSelection,
     pub allowed_role_metadata: AllowedRoleMetadata,
@@ -95,6 +97,7 @@ impl std::fmt::Debug for ResolvedAwsBedrock {
             .field("secret_access_key", &"<no-repr-available>")
             .field("session_token", &self.session_token)
             .field("profile", &self.profile)
+            .field("endpoint_url", &self.endpoint_url)
             .field("inference_config", &"<no-repr-available>")
             .field("role_selection", &self.role_selection)
             .field("allowed_role_metadata", &self.allowed_role_metadata)
@@ -121,6 +124,12 @@ impl ResolvedAwsBedrock {
                 serde_json::Value::String(region.clone()),
             );
         }
+        if let Some(endpoint_url) = &self.endpoint_url {
+            options.insert(
+                "endpoint_url".to_string(),
+                serde_json::Value::String(endpoint_url.clone()),
+            );
+        }
         options
     }
 
@@ -162,6 +171,7 @@ impl<Meta: Clone> UnresolvedAwsBedrock<Meta> {
             secret_access_key: self.secret_access_key.clone(),
             session_token: self.session_token.clone(),
             profile: self.profile.clone(),
+            endpoint_url: self.endpoint_url.clone(),
             role_selection: self.role_selection.clone(),
             allowed_role_metadata: self.allowed_role_metadata.clone(),
             supported_request_modes: self.supported_request_modes.clone(),
@@ -215,6 +225,11 @@ impl<Meta: Clone> UnresolvedAwsBedrock<Meta> {
             }
         }
 
+        match self.endpoint_url.as_ref() {
+            Some(endpoint_url) => env_vars.extend(endpoint_url.required_env_vars()),
+            None => {}
+        }
+
         env_vars.extend(self.role_selection.required_env_vars());
         env_vars.extend(self.allowed_role_metadata.required_env_vars());
         env_vars.extend(self.supported_request_modes.required_env_vars());
@@ -357,6 +372,18 @@ impl<Meta: Clone> UnresolvedAwsBedrock<Meta> {
             },
         };
 
+        let endpoint_url = match self.endpoint_url.as_ref() {
+            Some(endpoint_url) => {
+                let url = endpoint_url.resolve(ctx)?;
+                if url.is_empty() {
+                    None
+                } else {
+                    Some(url)
+                }
+            }
+            None => None,
+        };
+
         #[cfg(target_arch = "wasm32")]
         {
             if region.is_none() {
@@ -377,6 +404,7 @@ impl<Meta: Clone> UnresolvedAwsBedrock<Meta> {
             secret_access_key,
             session_token,
             profile,
+            endpoint_url,
             role_selection,
             allowed_role_metadata: self.allowed_role_metadata.resolve(ctx)?,
             supported_request_modes: self.supported_request_modes.clone(),
@@ -430,6 +458,9 @@ impl<Meta: Clone> UnresolvedAwsBedrock<Meta> {
         let profile = properties
             .ensure_string("profile", false)
             .map(|(_, v, _)| v.clone());
+        let endpoint_url = properties
+            .ensure_string("endpoint_url", false)
+            .map(|(_, v, _)| v.clone());
 
         let role_selection = properties.ensure_roles_selection();
         let allowed_metadata = properties.ensure_allowed_metadata();
@@ -511,6 +542,7 @@ impl<Meta: Clone> UnresolvedAwsBedrock<Meta> {
             secret_access_key,
             session_token,
             profile,
+            endpoint_url,
             role_selection,
             allowed_role_metadata: allowed_metadata,
             supported_request_modes,

engine/baml-runtime/src/internal/llm_client/primitive/aws/aws_client.rs

@@ -592,7 +592,7 @@ impl AwsClient {
         let config = loader.load().await;
         let http_client = custom_http_client::client()?;
 
-        let bedrock_config = aws_sdk_bedrockruntime::config::Builder::from(&config)
+        let mut bedrock_config = aws_sdk_bedrockruntime::config::Builder::from(&config)
             // To support HTTPS_PROXY https://github.com/awslabs/aws-sdk-rust/issues/169
             .http_client(http_client)
             // Adding a custom http client (above) breaks the stalled stream protection for some reason. If a bedrock request takes longer than 5s (the default grace period, it makes it error out), so we disable it.
@@ -601,9 +601,14 @@ impl AwsClient {
                 call_stack,
                 http_request_id.clone(),
                 &self.properties,
-            ))
-            .build();
-        Ok(BedrockRuntimeClient::from_conf(bedrock_config))
+            ));
+
+        // Set endpoint_url if specified
+        if let Some(endpoint_url) = self.properties.endpoint_url.as_ref() {
+            bedrock_config = bedrock_config.endpoint_url(endpoint_url);
+        }
+
+        Ok(BedrockRuntimeClient::from_conf(bedrock_config.build()))
     }
 
     async fn chat_anyhow(&self, response: &ConverseOutput) -> Result<String> {

fern/03-reference/baml/clients/providers/aws-bedrock.mdx

@@ -479,6 +479,10 @@ You can use this to modify the `region`, `access_key_id`, `secret_access_key`, a
   AWS profile name from credentials file. **Default: `AWS_PROFILE` environment variable**
 </ParamField>
 
+<ParamField path="endpoint_url" type="string">
+  AWS endpoint URL. Useful for using a VPC endpoint.
+</ParamField>
+
 <Markdown src="/snippets/role-selection.mdx" />
 <Markdown src="/snippets/allowed-role-metadata-basic.mdx" />
 <Markdown src="/snippets/supports-streaming.mdx" />

integ-tests/baml_src/clients.baml

@@ -227,9 +227,10 @@ client<llm> AwsBedrock {
     }
     // max_tokens 100000
     // max_completion_tokens 100000
-    // model "anthropic.claude-3-5-sonnet-20240620-v1:0"
+    model "anthropic.claude-3-5-sonnet-20240620-v1:0"
     // model_id "anthropic.claude-3-haiku-20240307-v1:0"
-    model "arn:aws:bedrock:us-east-1:404337120808:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0"
+    //model "arn:aws:bedrock:us-east-1:404337120808:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0"
+    endpoint_url "https://bedrock-runtime.us-east-1.amazonaws.com/"
     // region "us-east-1"
     // access_key_id env.AWS_ACCESS_KEY_ID
     // secret_access_key env.AWS_SECRET_ACCESS_KEY
@@ -253,6 +254,14 @@ client<llm> AwsBedrockInvalidRegion {
   }
 }
 
+client<llm> AwsBedrockInvalidEndpoint {
+  provider aws-bedrock
+  options {
+    model "anthropic.claude-3-5-sonnet-20240620-v1:0"
+    endpoint_url "https://bedrock-runtime.us-doenotexist-3.amazonaws.com"
+  }
+}
+
 client<llm> AwsBedrockInvalidAccessKey {
   provider aws-bedrock
   options {

integ-tests/baml_src/generators.baml

@@ -49,5 +49,5 @@ generator lang_go {
   output_dir "../go"
   version "0.211.0"
   client_package_name "example.com/integ-tests"
-  on_generate "mise x -- gofmt -w . && mise x -- goimports -w . && mise x -- go mod tidy"
+  on_generate "gofmt -w . && goimports -w . && go mod tidy"
 }

integ-tests/baml_src/test-files/providers/aws.baml

@@ -37,6 +37,13 @@ function TestAwsInvalidRegion(input: string) -> string {
   "#
 }
 
+function TestAwsInvalidEndpoint(input: string) -> string {
+  client AwsBedrockInvalidEndpoint
+  prompt #"
+    Write a nice short story about {{ input }}. Keep it to 15 words or less.
+  "#
+}
+
 function TestAwsInvalidAccessKey(input: string) -> string {
   client AwsBedrockInvalidAccessKey
   prompt #"