use std::pin::Pin;
use std::sync::Arc;
use std::time::SystemTime;
use async_trait::async_trait;
use aws_credential_types::Credentials;
use aws_credential_types::provider::ProvideCredentials;
use aws_sigv4::http_request::{SignableBody, SignableRequest, SigningSettings, sign};
use aws_sigv4::sign::v4;
use aws_smithy_eventstream::frame::{DecodedFrame, MessageFrameDecoder};
use futures::{Stream, StreamExt};
fn url_encode_model_id(model_id: &str) -> String {
super::config::encode_path_segment(model_id)
}
use super::config::BedrockProviderConfig;
use super::convert::{format_request, format_response, format_stream_event};
use super::error::parse_bedrock_error;
use super::stream::parse_event_stream_frame;
use super::types::{BedrockConverseResponse, BedrockModelListResponse};
use crate::error::KovaError;
use crate::models::{
ConversationMessage, InferenceConfig, ModelInfo, ModelResponse, StreamEvent, ToolDefinition,
};
use crate::provider::LlmProvider;
use crate::provider::http::map_request_error;
pub struct BedrockProvider {
client: reqwest::Client,
config: BedrockProviderConfig,
credentials_provider: Arc<dyn ProvideCredentials>,
}
impl BedrockProvider {
pub async fn new(config: BedrockProviderConfig) -> Result<Self, KovaError> {
let credentials_provider: Arc<dyn ProvideCredentials> = if let (
Some(access_key_id),
Some(secret_access_key),
) =
(&config.access_key_id, &config.secret_access_key)
{
let creds = Credentials::from_keys(
access_key_id.clone(),
secret_access_key.clone(),
config.session_token.clone(),
);
Arc::new(creds)
} else {
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
if let Some(profile) = &config.profile {
loader = loader.profile_name(profile);
}
let sdk_config = loader.load().await;
let provider = sdk_config.credentials_provider().ok_or_else(|| {
KovaError::Provider {
message: "Failed to resolve AWS credentials: no credentials provider found in the default chain".to_string(),
status_code: None,
}
})?;
Arc::from(provider)
};
let client = reqwest::Client::builder()
.timeout(config.timeout)
.build()
.map_err(|e| KovaError::Provider {
message: format!("Failed to build HTTP client: {e}"),
status_code: None,
})?;
Ok(Self {
client,
config,
credentials_provider,
})
}
async fn sign_request(
&self,
method: &str,
url: &str,
body: &[u8],
) -> Result<Vec<(String, String)>, KovaError> {
let creds = self
.credentials_provider
.provide_credentials()
.await
.map_err(|e| KovaError::Provider {
message: format!("Failed to resolve AWS credentials: {e}"),
status_code: None,
})?;
let identity = creds.into();
let signing_settings = SigningSettings::default();
let signing_params = v4::SigningParams::builder()
.identity(&identity)
.region(&self.config.region)
.name("bedrock")
.time(SystemTime::now())
.settings(signing_settings)
.build()
.map_err(|e| KovaError::Provider {
message: format!("Failed to build SigV4 signing params: {e}"),
status_code: None,
})?
.into();
let headers = [("content-type", "application/json")];
let signable_request =
SignableRequest::new(method, url, headers.into_iter(), SignableBody::Bytes(body))
.map_err(|e| KovaError::Provider {
message: format!("Failed to create signable request: {e}"),
status_code: None,
})?;
let (signing_instructions, _signature) = sign(signable_request, &signing_params)
.map_err(|e| KovaError::Provider {
message: format!("SigV4 signing failed: {e}"),
status_code: None,
})?
.into_parts();
Ok(signing_instructions
.headers()
.map(|(name, value)| (name.to_string(), value.to_string()))
.collect())
}
fn resolve_base_url(&self, service: &str) -> String {
self.config
.endpoint_url
.clone()
.unwrap_or_else(|| format!("https://{}.{}.amazonaws.com", service, self.config.region))
}
}
#[async_trait]
impl LlmProvider for BedrockProvider {
async fn chat_completion(
&self,
messages: &[ConversationMessage],
tools: &[ToolDefinition],
config: &InferenceConfig,
) -> Result<ModelResponse, KovaError> {
let span = tracing::info_span!(
"llm.chat_completion",
provider = "bedrock",
model = %self.config.model_id,
otel.status_code = tracing::field::Empty,
llm.input_tokens = tracing::field::Empty,
llm.output_tokens = tracing::field::Empty,
llm.stop_reason = tracing::field::Empty,
);
let _guard = span.enter();
let base = self.resolve_base_url("bedrock-runtime");
let encoded_model = url_encode_model_id(&self.config.model_id);
let url = format!("{}/model/{}/converse", base, encoded_model);
let bedrock_request = format_request(
messages,
tools,
config,
self.config.additional_model_request_fields.clone(),
);
let body = serde_json::to_vec(&bedrock_request).map_err(|e| KovaError::Provider {
message: format!("Failed to serialize request: {e}"),
status_code: None,
})?;
tracing::debug!(body = %String::from_utf8_lossy(&body), "Bedrock request body");
let signed_headers = self.sign_request("POST", &url, &body).await?;
let mut req = self
.client
.post(&url)
.header("content-type", "application/json")
.body(body);
for (name, value) in &signed_headers {
req = req.header(name.as_str(), value.as_str());
}
let response = req.send().await.map_err(|e| {
let err = map_request_error(e, self.config.timeout);
tracing::Span::current().record("otel.status_code", "ERROR");
tracing::warn!(error = %err, "Bedrock chat_completion request failed");
err
})?;
let status = response.status();
if !status.is_success() {
let body = response.text().await.unwrap_or_default();
let err = parse_bedrock_error(status.as_u16(), &body);
tracing::Span::current().record("otel.status_code", "ERROR");
tracing::warn!(error = %err, "Bedrock provider returned error");
return Err(err);
}
let response_text = response.text().await.map_err(|e| {
let err = KovaError::Provider {
message: format!("Failed to read response body: {e}"),
status_code: None,
};
tracing::Span::current().record("otel.status_code", "ERROR");
err
})?;
tracing::debug!(body = %response_text, "Bedrock response body");
let bedrock_response: BedrockConverseResponse = serde_json::from_str(&response_text)
.map_err(|e| {
let err = KovaError::Provider {
message: format!("Failed to deserialize response: {e}"),
status_code: None,
};
tracing::Span::current().record("otel.status_code", "ERROR");
err
})?;
tracing::Span::current().record("llm.input_tokens", bedrock_response.usage.input_tokens);
tracing::Span::current().record("llm.output_tokens", bedrock_response.usage.output_tokens);
tracing::Span::current().record("llm.stop_reason", bedrock_response.stop_reason.as_str());
tracing::info!("Bedrock chat completion succeeded");
format_response(bedrock_response)
}
async fn chat_completion_stream(
&self,
messages: &[ConversationMessage],
tools: &[ToolDefinition],
config: &InferenceConfig,
) -> Result<Pin<Box<dyn Stream<Item = Result<StreamEvent, KovaError>> + Send>>, KovaError> {
let span = tracing::info_span!(
"llm.chat_completion_stream",
provider = "bedrock",
model = %self.config.model_id,
otel.status_code = tracing::field::Empty,
);
let _guard = span.enter();
let base = self.resolve_base_url("bedrock-runtime");
let encoded_model = url_encode_model_id(&self.config.model_id);
let url = format!("{}/model/{}/converse-stream", base, encoded_model);
let bedrock_request = format_request(
messages,
tools,
config,
self.config.additional_model_request_fields.clone(),
);
let body = serde_json::to_vec(&bedrock_request).map_err(|e| KovaError::Provider {
message: format!("Failed to serialize request: {e}"),
status_code: None,
})?;
tracing::debug!(body = %String::from_utf8_lossy(&body), "Bedrock stream request body");
let signed_headers = self.sign_request("POST", &url, &body).await?;
let mut req = self
.client
.post(&url)
.header("content-type", "application/json")
.body(body);
for (name, value) in &signed_headers {
req = req.header(name.as_str(), value.as_str());
}
let response = req.send().await.map_err(|e| {
let err = map_request_error(e, self.config.timeout);
tracing::Span::current().record("otel.status_code", "ERROR");
tracing::warn!(error = %err, "Bedrock stream request failed");
err
})?;
let status = response.status();
if !status.is_success() {
let body = response.text().await.unwrap_or_default();
let err = parse_bedrock_error(status.as_u16(), &body);
tracing::Span::current().record("otel.status_code", "ERROR");
tracing::warn!(error = %err, "Bedrock stream provider returned error");
return Err(err);
}
let byte_stream = response.bytes_stream();
let stream = futures::stream::try_unfold(
(
byte_stream,
MessageFrameDecoder::new(),
bytes::BytesMut::new(),
),
|(mut byte_stream, mut decoder, mut buffer)| async move {
loop {
match decoder.decode_frame(&mut buffer) {
Ok(DecodedFrame::Complete(message)) => {
match parse_event_stream_frame(&message) {
Ok(Some(bedrock_event)) => {
if let Some(stream_event) = format_stream_event(bedrock_event) {
return Ok(Some((
stream_event,
(byte_stream, decoder, buffer),
)));
}
continue;
}
Ok(None) => continue,
Err(e) => return Err(e),
}
}
Ok(DecodedFrame::Incomplete) => {}
Err(e) => {
return Err(KovaError::Stream(format!(
"Failed to decode event stream frame: {e}"
)));
}
}
match byte_stream.next().await {
Some(Ok(chunk)) => buffer.extend_from_slice(&chunk),
Some(Err(e)) => {
return Err(KovaError::Stream(format!("Stream connection error: {e}")));
}
None => return Ok(None),
}
}
},
);
Ok(Box::pin(stream))
}
async fn list_models(&self) -> Result<Vec<ModelInfo>, KovaError> {
let span = tracing::info_span!(
"llm.list_models",
provider = "bedrock",
otel.status_code = tracing::field::Empty,
);
let _guard = span.enter();
let base = self.resolve_base_url("bedrock");
let url = format!("{}/foundation-models", base);
let signed_headers = self.sign_request("GET", &url, &[]).await?;
let mut req = self.client.get(&url);
for (name, value) in &signed_headers {
req = req.header(name.as_str(), value.as_str());
}
let response = req.send().await.map_err(|e| {
let err = map_request_error(e, self.config.timeout);
tracing::Span::current().record("otel.status_code", "ERROR");
tracing::warn!(error = %err, "Bedrock list_models request failed");
err
})?;
let status = response.status();
if !status.is_success() {
let body = response.text().await.unwrap_or_default();
let err = parse_bedrock_error(status.as_u16(), &body);
tracing::Span::current().record("otel.status_code", "ERROR");
tracing::warn!(error = %err, "Bedrock list_models provider returned error");
return Err(err);
}
let list_response: BedrockModelListResponse = response.json().await.map_err(|e| {
let err = KovaError::Provider {
message: format!("Failed to deserialize model list response: {e}"),
status_code: None,
};
tracing::Span::current().record("otel.status_code", "ERROR");
err
})?;
Ok(list_response
.model_summaries
.into_iter()
.map(|s| ModelInfo {
id: s.model_id,
object: "model".to_string(),
created: 0,
owned_by: s.provider_name,
})
.collect())
}
}
#[cfg(test)]
#[path = "provider_tests.rs"]
mod tests;