llm/providers/bedrock/
provider.rs

1use super::mappers::{default_cache_point, map_messages, map_tools};
2use super::streaming::process_bedrock_stream;
3use crate::provider::{LlmResponseStream, ProviderFactory, StreamingModelProvider, get_context_window};
4use crate::{Context, LlmError, ProviderAuthMode, ProviderConnectionConfig, Result};
5use aws_config::Region;
6use aws_sdk_bedrockruntime::config::{BehaviorVersion, Credentials};
7use aws_sdk_bedrockruntime::error::SdkError;
8use aws_sdk_bedrockruntime::operation::converse_stream::ConverseStreamError;
9use aws_sdk_bedrockruntime::primitives::event_stream::EventReceiver;
10use aws_sdk_bedrockruntime::types::error::ConverseStreamOutputError;
11use aws_sdk_bedrockruntime::types::{ConverseStreamOutput, InferenceConfiguration};
12use aws_sdk_bedrockruntime::{Client, Config};
13use futures::StreamExt;
14use tracing::{error, info};
15
16const DEFAULT_MODEL: &str = "anthropic.claude-sonnet-4-5-20250929-v1:0";
17const DEFAULT_MAX_TOKENS: i32 = 16_384;
18const DEFAULT_REGION: &str = "us-east-1";
19
20/// AWS credentials for explicit authentication with Bedrock.
21#[derive(Clone)]
22pub struct AwsCredentials {
23    pub access_key_id: String,
24    pub secret_access_key: String,
25    pub session_token: Option<String>,
26}
27
28#[derive(Clone)]
29pub struct BedrockProvider {
30    client: Client,
31    model: String,
32    inference_profile_arn: Option<String>,
33}
34
35impl BedrockProvider {
36    /// Create a provider using the default AWS credential chain
37    /// (env vars, `~/.aws/credentials`, IAM roles, SSO).
38    pub async fn new() -> Self {
39        Self::new_with_connection(ProviderConnectionConfig::default()).await
40    }
41
42    pub async fn new_with_connection(connection: ProviderConnectionConfig) -> Self {
43        let client = if connection.auth_mode == ProviderAuthMode::None {
44            build_no_auth_client(connection.base_url.as_deref(), region_from_env().as_deref())
45        } else {
46            let mut loader = aws_config::defaults(BehaviorVersion::latest());
47            if let Some(url) = &connection.base_url {
48                loader = loader.endpoint_url(url.clone());
49            }
50            let config = loader.load().await;
51            Client::new(&config)
52        };
53
54        Self { client, model: DEFAULT_MODEL.to_string(), inference_profile_arn: connection.inference_profile_arn }
55    }
56
57    /// Create a provider from explicit configuration without async credential discovery.
58    pub fn from_config(credentials: Option<AwsCredentials>, region: Option<&str>) -> Self {
59        let client = build_client(credentials, region);
60
61        Self { client, model: DEFAULT_MODEL.to_string(), inference_profile_arn: None }
62    }
63
64    pub fn with_model(mut self, model: &str) -> Self {
65        self.model = model.to_string();
66        self
67    }
68
69    pub fn with_inference_profile_arn(mut self, arn: impl Into<String>) -> Self {
70        self.inference_profile_arn = Some(arn.into());
71        self
72    }
73
74    fn request_model_id(&self) -> &str {
75        self.inference_profile_arn.as_deref().unwrap_or(&self.model)
76    }
77
78    async fn send_converse_stream(
79        &self,
80        context: &Context,
81    ) -> Result<EventReceiver<ConverseStreamOutput, ConverseStreamOutputError>> {
82        let cache_point =
83            self.model().is_some_and(|m| m.supports_prompt_caching()).then(default_cache_point).transpose()?;
84        let (system_blocks, messages) = map_messages(context.messages(), cache_point.as_ref())?;
85        let settings = context.model_settings();
86        let max_tokens = settings.max_tokens.and_then(|m| i32::try_from(m).ok()).unwrap_or(DEFAULT_MAX_TOKENS);
87        let mut inference_config = InferenceConfiguration::builder().max_tokens(max_tokens);
88
89        if let Some(temp) = settings.temperature {
90            inference_config = inference_config.temperature(temp);
91        }
92
93        if let Some(top_p) = settings.top_p {
94            inference_config = inference_config.top_p(top_p);
95        }
96
97        let inference_config = inference_config.build();
98
99        let mut request = self
100            .client
101            .converse_stream()
102            .model_id(self.request_model_id())
103            .set_messages(Some(messages))
104            .inference_config(inference_config);
105
106        if !system_blocks.is_empty() {
107            request = request.set_system(Some(system_blocks));
108        }
109
110        if !context.tools().is_empty() {
111            let tool_config = map_tools(context.tools(), cache_point.as_ref())?;
112            request = request.tool_config(tool_config);
113        }
114
115        if let Some(arn) = self.inference_profile_arn.as_deref() {
116            info!(model = %self.model, inference_profile_arn = %arn, "Sending Bedrock converse_stream request");
117        } else {
118            info!(model = %self.model, "Sending Bedrock converse_stream request");
119        }
120
121        let response = request.send().await.map_err(|e| {
122            error!(model = %self.model, error = ?e, "Bedrock API error");
123            LlmError::from(e)
124        })?;
125
126        Ok(response.stream)
127    }
128}
129
130impl ProviderFactory for BedrockProvider {
131    async fn from_env() -> Result<Self> {
132        Ok(Self::new().await)
133    }
134
135    async fn from_env_with_connection(connection: ProviderConnectionConfig) -> Result<Self> {
136        Ok(Self::new_with_connection(connection).await)
137    }
138
139    fn with_model(self, model: &str) -> Self {
140        self.with_model(model)
141    }
142}
143
144impl StreamingModelProvider for BedrockProvider {
145    fn model(&self) -> Option<crate::LlmModel> {
146        format!("bedrock:{}", self.model).parse().ok()
147    }
148
149    fn context_window(&self) -> Option<u32> {
150        get_context_window("bedrock", &self.model)
151    }
152
153    fn stream_response(&self, context: &Context) -> LlmResponseStream {
154        let provider = self.clone();
155        let context = context.clone();
156
157        Box::pin(async_stream::stream! {
158            match provider.send_converse_stream(&context).await {
159                Ok(receiver) => {
160                    let mut stream = Box::pin(process_bedrock_stream(receiver));
161                    while let Some(result) = stream.next().await {
162                        yield result;
163                    }
164                }
165                Err(e) => {
166                    yield Err(e);
167                }
168            }
169        })
170    }
171
172    fn display_name(&self) -> String {
173        format!("Bedrock ({})", self.model)
174    }
175}
176
177impl From<SdkError<ConverseStreamError>> for LlmError {
178    fn from(e: SdkError<ConverseStreamError>) -> Self {
179        let message = format!("Bedrock API error: {e}");
180        match e {
181            SdkError::TimeoutError(_) => LlmError::Timeout(message),
182            SdkError::DispatchFailure(_) => LlmError::Network(message),
183            SdkError::ResponseError(_) => LlmError::ServerError { status: None, message },
184            SdkError::ServiceError(svc) => {
185                let inner = svc.err();
186                if inner.is_throttling_exception() {
187                    LlmError::RateLimited(message)
188                } else if inner.is_service_unavailable_exception()
189                    || inner.is_internal_server_exception()
190                    || inner.is_model_stream_error_exception()
191                {
192                    LlmError::ServerError { status: None, message }
193                } else {
194                    LlmError::ApiError(message)
195                }
196            }
197            _ => LlmError::ApiError(message),
198        }
199    }
200}
201
202fn build_client(credentials: Option<AwsCredentials>, region: Option<&str>) -> Client {
203    let mut config = Config::builder().behavior_version(BehaviorVersion::latest());
204
205    if let Some(creds) = credentials {
206        config = config.credentials_provider(Credentials::new(
207            creds.access_key_id,
208            creds.secret_access_key,
209            creds.session_token,
210            None,
211            "aether-bedrock-provider",
212        ));
213    }
214
215    config = config.region(Region::new(region.unwrap_or(DEFAULT_REGION).to_string()));
216
217    Client::from_conf(config.build())
218}
219
220fn build_no_auth_client(base_url: Option<&str>, region: Option<&str>) -> Client {
221    let mut config = Config::builder()
222        .behavior_version(BehaviorVersion::latest())
223        .allow_no_auth()
224        .region(Region::new(region.unwrap_or(DEFAULT_REGION).to_string()));
225
226    if let Some(url) = base_url {
227        config = config.endpoint_url(url);
228    }
229
230    Client::from_conf(config.build())
231}
232
233fn region_from_env() -> Option<String> {
234    ["AWS_REGION", "AWS_DEFAULT_REGION"].into_iter().find_map(|name| match std::env::var(name) {
235        Ok(value) if !value.is_empty() => Some(value),
236        _ => None,
237    })
238}
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243    use crate::types::IsoString;
244    use crate::{ChatMessage, ContentBlock};
245    use axum::Router;
246    use axum::body::Body;
247    use axum::extract::State;
248    use axum::http::{HeaderMap, Method, Request, StatusCode};
249    use axum::response::IntoResponse;
250    use axum::routing::any;
251    use std::sync::Arc;
252    use tokio::net::TcpListener;
253    use tokio::sync::{Mutex, oneshot};
254
255    fn inference_profile_arn(model: &str) -> String {
256        format!("arn:aws:bedrock:us-west-2:000000000000:inference-profile/{model}")
257    }
258
259    fn application_inference_profile_arn() -> &'static str {
260        "arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/000000000000"
261    }
262
263    fn test_provider() -> BedrockProvider {
264        BedrockProvider::from_config(None, None)
265    }
266
267    #[test]
268    fn test_display_name() {
269        assert_eq!(test_provider().display_name(), "Bedrock (anthropic.claude-sonnet-4-5-20250929-v1:0)");
270    }
271
272    #[test]
273    fn test_with_model() {
274        let provider = test_provider().with_model("anthropic.claude-opus-4-20250514-v1:0");
275        assert_eq!(provider.display_name(), "Bedrock (anthropic.claude-opus-4-20250514-v1:0)");
276    }
277
278    #[test]
279    fn test_default_values() {
280        let provider = test_provider();
281        assert_eq!(provider.model, "anthropic.claude-sonnet-4-5-20250929-v1:0");
282    }
283
284    #[tokio::test]
285    async fn auth_none_sends_unsigned_request_to_custom_endpoint() {
286        let endpoint = FakeBedrockEndpoint::start().await;
287        let provider = BedrockProvider::new_with_connection(ProviderConnectionConfig {
288            base_url: Some(endpoint.url.clone()),
289            auth_mode: ProviderAuthMode::None,
290            ..Default::default()
291        })
292        .await;
293
294        let context = Context::new(
295            vec![ChatMessage::User { content: vec![ContentBlock::text("hello")], timestamp: IsoString::now() }],
296            vec![],
297        );
298
299        let result = provider.send_converse_stream(&context).await;
300        let request = endpoint.request.await.expect("fake Bedrock endpoint received no request");
301
302        assert!(result.is_err());
303        assert_eq!(request.method, Method::POST);
304        assert!(request.path.starts_with("/model/"), "{}", request.path);
305        assert!(!request.headers.contains_key("authorization"), "request was signed: {:?}", request.headers);
306        assert!(
307            !request.headers.contains_key("x-amz-security-token"),
308            "request included session token: {:?}",
309            request.headers
310        );
311    }
312
313    #[test]
314    fn test_from_config_with_credentials() {
315        let credentials = AwsCredentials {
316            access_key_id: "AKIAIOSFODNN7EXAMPLE".to_string(),
317            secret_access_key: "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".to_string(),
318            session_token: None,
319        };
320
321        let provider = BedrockProvider::from_config(Some(credentials), None);
322        assert_eq!(provider.model, DEFAULT_MODEL);
323    }
324
325    #[test]
326    fn test_from_config_with_credentials_and_region() {
327        let credentials = AwsCredentials {
328            access_key_id: "AKIAIOSFODNN7EXAMPLE".to_string(),
329            secret_access_key: "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".to_string(),
330            session_token: Some("FwoGZXIvYXdzEBYaD...".to_string()),
331        };
332
333        let provider = BedrockProvider::from_config(Some(credentials), Some("us-west-2"))
334            .with_model("anthropic.claude-opus-4-20250514-v1:0");
335
336        assert_eq!(provider.model, "anthropic.claude-opus-4-20250514-v1:0");
337    }
338
339    #[test]
340    fn test_from_config_with_region_only() {
341        let provider = BedrockProvider::from_config(None, Some("eu-west-1"));
342        assert_eq!(provider.model, DEFAULT_MODEL);
343    }
344
345    #[test]
346    fn catalog_foundation_id_resolves_context_window() {
347        let provider = test_provider().with_model("anthropic.claude-sonnet-4-5-20250929-v1:0");
348        assert!(provider.context_window().is_some());
349        assert_eq!(provider.model().unwrap().to_string(), "bedrock:anthropic.claude-sonnet-4-5-20250929-v1:0");
350    }
351
352    #[test]
353    fn cross_region_profile_id_in_catalog_resolves() {
354        let provider = test_provider().with_model("us.anthropic.claude-opus-4-6-v1");
355        assert!(provider.context_window().is_some());
356    }
357
358    #[test]
359    fn unknown_cross_region_profile_id_falls_through_to_profile() {
360        let id = "us.anthropic.claude-future-model-v99:0";
361        let provider = test_provider().with_model(id);
362        assert_eq!(provider.context_window(), None);
363        assert_eq!(provider.model().unwrap().to_string(), format!("bedrock:{id}"));
364        assert_eq!(provider.display_name(), format!("Bedrock ({id})"));
365    }
366
367    #[tokio::test]
368    async fn separate_inference_profile_arn_is_used_as_request_model_id() {
369        let endpoint = FakeBedrockEndpoint::start().await;
370        let provider = BedrockProvider::new_with_connection(ProviderConnectionConfig {
371            base_url: Some(endpoint.url.clone()),
372            auth_mode: ProviderAuthMode::None,
373            inference_profile_arn: Some(application_inference_profile_arn().to_string()),
374        })
375        .await
376        .with_model("anthropic.claude-sonnet-4-5-20250929-v1:0");
377        let context = Context::new(
378            vec![ChatMessage::User { content: vec![ContentBlock::text("hello")], timestamp: IsoString::now() }],
379            vec![],
380        );
381
382        let result = provider.send_converse_stream(&context).await;
383        let request = endpoint.request.await.expect("fake Bedrock endpoint received no request");
384
385        assert!(result.is_err());
386        assert!(
387            request.path.contains("arn%3Aaws%3Abedrock%3Aus-west-2%3A000000000000%3Aapplication-inference-profile"),
388            "{}",
389            request.path
390        );
391        assert_eq!(provider.context_window(), Some(200_000));
392        assert_eq!(provider.model().unwrap().to_string(), "bedrock:anthropic.claude-sonnet-4-5-20250929-v1:0");
393    }
394
395    #[test]
396    fn with_inference_profile_arn_keeps_canonical_model_identity() {
397        let arn = inference_profile_arn("us.anthropic.claude-sonnet-4-5-20250929-v1:0");
398        let provider =
399            test_provider().with_model("anthropic.claude-sonnet-4-5-20250929-v1:0").with_inference_profile_arn(&arn);
400
401        assert_eq!(provider.request_model_id(), arn);
402        assert_eq!(provider.context_window(), Some(200_000));
403        assert_eq!(provider.model().unwrap().to_string(), "bedrock:anthropic.claude-sonnet-4-5-20250929-v1:0");
404    }
405
406    #[test]
407    fn prompt_caching_support_comes_from_canonical_model() {
408        let cached = test_provider().with_model("anthropic.claude-sonnet-4-5-20250929-v1:0");
409        assert!(cached.model().unwrap().supports_prompt_caching());
410
411        let unknown_profile = test_provider().with_model("us.anthropic.claude-future-model-v99:0");
412        assert!(!unknown_profile.model().unwrap().supports_prompt_caching());
413    }
414
415    struct FakeBedrockEndpoint {
416        url: String,
417        request: oneshot::Receiver<CapturedRequest>,
418    }
419
420    struct CapturedRequest {
421        method: Method,
422        path: String,
423        headers: HeaderMap,
424    }
425
426    #[derive(Clone)]
427    struct FakeBedrockState {
428        request_tx: Arc<Mutex<Option<oneshot::Sender<CapturedRequest>>>>,
429        shutdown_tx: Arc<Mutex<Option<oneshot::Sender<()>>>>,
430    }
431
432    impl FakeBedrockEndpoint {
433        async fn start() -> Self {
434            let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind fake Bedrock endpoint");
435            let url = format!("http://{}", listener.local_addr().expect("fake Bedrock endpoint address"));
436            let (request_tx, request) = oneshot::channel();
437            let (shutdown_tx, shutdown) = oneshot::channel();
438            let state = FakeBedrockState {
439                request_tx: Arc::new(Mutex::new(Some(request_tx))),
440                shutdown_tx: Arc::new(Mutex::new(Some(shutdown_tx))),
441            };
442
443            let app = Router::new().fallback(any(capture_bedrock_request)).with_state(state);
444            tokio::spawn(async move {
445                axum::serve(listener, app)
446                    .with_graceful_shutdown(async {
447                        let _ = shutdown.await;
448                    })
449                    .await
450                    .expect("serve fake Bedrock endpoint");
451            });
452
453            Self { url, request }
454        }
455    }
456
457    async fn capture_bedrock_request(
458        State(state): State<FakeBedrockState>,
459        request: Request<Body>,
460    ) -> impl IntoResponse {
461        let (parts, _) = request.into_parts();
462        if let Some(tx) = state.request_tx.lock().await.take() {
463            let _ = tx.send(CapturedRequest {
464                method: parts.method,
465                path: parts.uri.path().to_string(),
466                headers: parts.headers,
467            });
468        }
469        if let Some(tx) = state.shutdown_tx.lock().await.take() {
470            let _ = tx.send(());
471        }
472        (StatusCode::FORBIDDEN, "{}")
473    }
474}
llm/providers/bedrock/provider.rs

llm/providers/bedrock/
provider.rs
