litellm-rs 0.6.0

//! Main Bedrock Provider Implementation
//!
//! Contains the BedrockProvider struct and its LLMProvider trait implementation.

use futures::Stream;
use serde_json::Value;
use std::collections::HashMap;
use std::pin::Pin;
use tracing::debug;

use super::client::BedrockClient;
use super::config::BedrockConfig;
use super::error::BedrockErrorMapper;
use super::model_config::BedrockApiType;
use super::model_id::is_runtime_resolved_invoke_model_id;
use super::transformation;
use super::utils::{CostCalculator, validate_region};
use super::{get_model_config_for_model_id, parse_bedrock_model_id};
use crate::core::traits::provider::ProviderConfig as _;

use crate::core::providers::unified_provider::ProviderError;
use crate::core::traits::error_mapper::trait_def::ErrorMapper;
use crate::core::traits::provider::llm_provider::trait_definition::LLMProvider;
use crate::core::types::{
    chat::ChatMessage,
    chat::ChatRequest,
    context::RequestContext,
    embedding::EmbeddingRequest,
    health::HealthStatus,
    message::MessageContent,
    message::MessageRole,
    model::ModelInfo,
    model::ProviderCapability,
    responses::{ChatChunk, ChatResponse, EmbeddingResponse},
};

/// Static capabilities for Bedrock provider
pub(super) const BEDROCK_CAPABILITIES: &[ProviderCapability] = &[
    ProviderCapability::ChatCompletion,
    ProviderCapability::ChatCompletionStream,
    ProviderCapability::FunctionCalling,
    ProviderCapability::Embeddings,
];

fn streaming_operation_for_api_type(api_type: &BedrockApiType) -> &'static str {
    match api_type {
        BedrockApiType::Converse | BedrockApiType::ConverseStream => "converse-stream",
        BedrockApiType::Invoke | BedrockApiType::InvokeStream => "invoke-with-response-stream",
    }
}

fn converse_streaming_request_body(request: &ChatRequest) -> Result<Value, ProviderError> {
    let converse_request = super::chat::converse::transform_to_converse(request)?;
    serde_json::to_value(converse_request)
        .map_err(|e| ProviderError::serialization("bedrock", e.to_string()))
}

/// AWS Bedrock provider implementation
#[derive(Debug, Clone)]
pub struct BedrockProvider {
    client: BedrockClient,
    models: Vec<ModelInfo>,
}

impl BedrockProvider {
    /// Create a new Bedrock provider instance
    pub async fn new(config: BedrockConfig) -> Result<Self, ProviderError> {
        // Validate configuration
        config
            .validate()
            .map_err(|e| ProviderError::configuration("bedrock", e))?;

        // Validate AWS region
        validate_region(&config.aws_region)?;

        // Create Bedrock client
        let client = BedrockClient::new(config)?;

        // Define supported models using cost calculator data
        let mut models = Vec::new();
        let available_models = CostCalculator::get_all_models();

        for model_id in available_models {
            if let Some(pricing) = CostCalculator::get_core_model_pricing(model_id)
                && let Ok(model_config) = get_model_config_for_model_id(model_id)
            {
                models.push(ModelInfo {
                    id: model_id.to_string(),
                    name: format!(
                        "{} (Bedrock)",
                        model_id.split('.').next_back().unwrap_or(model_id)
                    ),
                    provider: "bedrock".to_string(),
                    max_context_length: model_config.max_context_length,
                    max_output_length: model_config.max_output_length,
                    supports_streaming: model_config.supports_streaming,
                    supports_tools: model_config.supports_function_calling,
                    supports_multimodal: model_config.supports_multimodal,
                    input_cost_per_1k_tokens: Some(pricing.input_cost_per_1k_tokens),
                    output_cost_per_1k_tokens: Some(pricing.output_cost_per_1k_tokens),
                    currency: pricing.currency,
                    capabilities: vec![],
                    created_at: None,
                    updated_at: None,
                    metadata: HashMap::new(),
                });
            }
        }

        Ok(Self { client, models })
    }

    /// Generate images using Bedrock image models
    pub async fn generate_image(
        &self,
        request: &crate::core::types::image::ImageGenerationRequest,
    ) -> Result<crate::core::types::responses::ImageGenerationResponse, ProviderError> {
        super::images::execute_image_generation(&self.client, request).await
    }

    /// Access the Agents client
    pub fn agents_client(&self) -> super::agents::AgentClient<'_> {
        super::agents::AgentClient::new(&self.client)
    }

    /// Access the Knowledge Bases client
    pub fn knowledge_bases_client(&self) -> super::knowledge_bases::KnowledgeBaseClient<'_> {
        super::knowledge_bases::KnowledgeBaseClient::new(&self.client)
    }

    /// Access the Batch processing client
    pub fn batch_client(&self) -> super::batch::BatchClient<'_> {
        super::batch::BatchClient::new(&self.client)
    }

    /// Access the Guardrails client
    pub fn guardrails_client(&self) -> super::guardrails::GuardrailClient<'_> {
        super::guardrails::GuardrailClient::new(&self.client)
    }

    /// Check if model is an embedding model
    pub(super) fn is_embedding_model(&self, model: &str) -> bool {
        model.contains("embed")
    }

    /// Create a test provider for unit testing
    #[cfg(test)]
    pub(super) fn new_for_test(client: BedrockClient, models: Vec<ModelInfo>) -> Self {
        Self { client, models }
    }

    /// Convert messages to a single prompt string for models that require it
    pub(super) fn messages_to_prompt(
        &self,
        messages: &[ChatMessage],
    ) -> Result<String, ProviderError> {
        let mut prompt = String::new();

        for message in messages {
            let content = match &message.content {
                Some(MessageContent::Text(text)) => text.clone(),
                Some(MessageContent::Parts(parts)) => {
                    // Extract text from parts
                    parts
                        .iter()
                        .filter_map(|part| {
                            if let crate::core::types::content::ContentPart::Text { text } = part {
                                Some(text.clone())
                            } else {
                                None
                            }
                        })
                        .collect::<Vec<_>>()
                        .join(" ")
                }
                None => continue,
            };

            match message.role {
                MessageRole::System | MessageRole::Developer => {
                    prompt.push_str(&format!("System: {}\n\n", content))
                }
                MessageRole::User => prompt.push_str(&format!("Human: {}\n\n", content)),
                MessageRole::Assistant => prompt.push_str(&format!("Assistant: {}\n\n", content)),
                MessageRole::Function | MessageRole::Tool => {
                    prompt.push_str(&format!("Tool: {}\n\n", content));
                }
            }
        }

        // Add Assistant prompt at the end for completion
        prompt.push_str("Assistant:");

        Ok(prompt)
    }
}

impl LLMProvider for BedrockProvider {
    fn name(&self) -> &'static str {
        "bedrock"
    }

    fn capabilities(&self) -> &'static [ProviderCapability] {
        BEDROCK_CAPABILITIES
    }

    fn models(&self) -> &[ModelInfo] {
        &self.models
    }

    fn supports_model(&self, model: &str) -> bool {
        let parsed = parse_bedrock_model_id(model);
        if parsed.user_selector.starts_with("bedrock/")
            || parsed.execution_model_id.starts_with("arn:")
        {
            return get_model_config_for_model_id(model).is_ok();
        }

        parsed
            .metadata_lookup_ids
            .iter()
            .any(|lookup_id| super::model_config::get_model_config(lookup_id).is_ok())
    }

    fn get_supported_openai_params(&self, _model: &str) -> &'static [&'static str] {
        &[
            "temperature",
            "top_p",
            "max_tokens",
            "stream",
            "stop",
            "tools",
            "tool_choice",
        ]
    }

    async fn map_openai_params(
        &self,
        params: HashMap<String, Value>,
        _model: &str,
    ) -> Result<HashMap<String, Value>, ProviderError> {
        // Bedrock has some differences from OpenAI format
        let mut mapped = HashMap::new();

        for (key, value) in params {
            match key.as_str() {
                // Map OpenAI parameters to Bedrock format
                "max_tokens" => mapped.insert("max_tokens_to_sample".to_string(), value),
                "temperature" | "top_p" | "stream" | "stop" => mapped.insert(key, value),
                // Skip unsupported parameters
                _ => None,
            };
        }

        Ok(mapped)
    }

    async fn transform_request(
        &self,
        request: ChatRequest,
        _context: RequestContext,
    ) -> Result<Value, ProviderError> {
        if is_runtime_resolved_invoke_model_id(&request.model) {
            return super::chat::transformations::transform_runtime_invoke_request(&request);
        }

        transformation::transform_chat_request(
            &request.model,
            &request.messages,
            request.max_tokens,
            request.temperature,
            request.top_p,
            |msgs| self.messages_to_prompt(msgs),
        )
    }

    async fn transform_response(
        &self,
        raw_response: &[u8],
        model: &str,
        _request_id: &str,
    ) -> Result<ChatResponse, ProviderError> {
        transformation::transform_chat_response(raw_response, model)
    }

    async fn chat_completion(
        &self,
        request: ChatRequest,
        _context: RequestContext,
    ) -> Result<ChatResponse, ProviderError> {
        debug!("Bedrock chat request: model={}", request.model);

        // Check if it's an embedding model
        if self.is_embedding_model(&request.model) {
            return Err(ProviderError::invalid_request(
                "bedrock",
                "Use embeddings endpoint for embedding models".to_string(),
            ));
        }

        // Use the chat module's routing logic
        let response_value = super::chat::route_chat_request(&self.client, &request).await?;

        // Convert the response to bytes for transform_response
        let response_bytes = serde_json::to_vec(&response_value)
            .map_err(|e| ProviderError::serialization("bedrock", e.to_string()))?;

        self.transform_response(&response_bytes, &request.model, "bedrock-request")
            .await
    }

    async fn chat_completion_stream(
        &self,
        request: ChatRequest,
        context: RequestContext,
    ) -> Result<Pin<Box<dyn Stream<Item = Result<ChatChunk, ProviderError>> + Send>>, ProviderError>
    {
        debug!("Bedrock streaming chat request: model={}", request.model);

        // Check if it's an embedding model
        if self.is_embedding_model(&request.model) {
            return Err(ProviderError::invalid_request(
                "bedrock",
                "Use embeddings endpoint for embedding models".to_string(),
            ));
        }

        // Get model configuration
        let model_config = get_model_config_for_model_id(&request.model)?;

        if !model_config.supports_streaming {
            return Err(ProviderError::not_supported(
                "bedrock",
                format!("Model {} does not support streaming", request.model),
            ));
        }

        // Use streaming endpoint
        let operation = streaming_operation_for_api_type(&model_config.api_type);
        let body = match &model_config.api_type {
            BedrockApiType::Converse | BedrockApiType::ConverseStream => {
                converse_streaming_request_body(&request)?
            }
            BedrockApiType::Invoke | BedrockApiType::InvokeStream => {
                self.transform_request(request.clone(), context).await?
            }
        };

        // Send streaming request
        let execution_model_id = parse_bedrock_model_id(&request.model).execution_model_id;
        let response = self
            .client
            .send_streaming_request(&execution_model_id, operation, &body)
            .await?;

        // Create BedrockStream
        let stream = super::streaming::BedrockStream::new(
            response.bytes_stream(),
            model_config.family.clone(),
            model_config.api_type.clone(),
            request.model.clone(),
        );

        Ok(Box::pin(stream))
    }

    async fn embeddings(
        &self,
        request: EmbeddingRequest,
        _context: RequestContext,
    ) -> Result<EmbeddingResponse, ProviderError> {
        debug!("Bedrock embedding request: model={}", request.model);

        // Use the embeddings module
        super::embeddings::execute_embedding(&self.client, &request).await
    }

    async fn health_check(&self) -> HealthStatus {
        match self.client.health_check().await {
            Ok(is_healthy) => {
                if is_healthy {
                    HealthStatus::Healthy
                } else {
                    HealthStatus::Unhealthy
                }
            }
            Err(error) => {
                debug!(error = %error, "Bedrock health check failed");
                HealthStatus::Unhealthy
            }
        }
    }

    fn get_error_mapper(&self) -> Box<dyn ErrorMapper<ProviderError>> {
        Box::new(BedrockErrorMapper)
    }

    async fn calculate_cost(
        &self,
        model: &str,
        input_tokens: u32,
        output_tokens: u32,
    ) -> Result<f64, ProviderError> {
        CostCalculator::calculate_cost(model, input_tokens, output_tokens)
            .ok_or_else(|| ProviderError::model_not_found("bedrock", model.to_string()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn streaming_operation_promotes_non_stream_api_types() {
        assert_eq!(
            streaming_operation_for_api_type(&BedrockApiType::Converse),
            "converse-stream"
        );
        assert_eq!(
            streaming_operation_for_api_type(&BedrockApiType::Invoke),
            "invoke-with-response-stream"
        );
    }

    #[test]
    fn catalog_streaming_converse_model_uses_converse_stream_operation() {
        let config = get_model_config_for_model_id("anthropic.claude-opus-4-6-v1:0").unwrap();
        assert_eq!(config.api_type, BedrockApiType::Converse);
        assert!(config.supports_streaming);
        assert_eq!(
            streaming_operation_for_api_type(&config.api_type),
            "converse-stream"
        );
    }

    #[test]
    fn converse_streaming_request_uses_converse_body_shape() {
        let mut request = ChatRequest::new("anthropic.claude-opus-4-6-v1:0")
            .add_system_message("Use concise answers.")
            .add_user_message("hello");
        request.max_tokens = Some(64);
        request.temperature = Some(0.2);
        request.stream = true;

        let body = converse_streaming_request_body(&request).unwrap();

        assert_eq!(body["system"][0]["text"], "Use concise answers.");
        assert_eq!(body["messages"][0]["role"], "user");
        assert_eq!(body["messages"][0]["content"][0]["text"], "hello");
        assert_eq!(body["inferenceConfig"]["maxTokens"], 64);
        assert!(body.get("max_tokens").is_none());
    }

    #[test]
    fn supports_model_does_not_capture_plain_non_bedrock_ids() {
        let config = BedrockConfig {
            aws_access_key_id: "AKIATEST123456789012".to_string(),
            aws_secret_access_key: "test-secret-key".to_string(),
            aws_session_token: None,
            aws_region: "us-east-1".to_string(),
            timeout_seconds: 30,
            max_retries: 3,
            endpoint_access: Default::default(),
        };
        let client = BedrockClient::new(config)
            .unwrap_or_else(|err| panic!("test Bedrock client should build: {err}"));
        let provider = BedrockProvider::new_for_test(client, vec![]);

        assert!(provider.supports_model("anthropic.claude-3-sonnet-20240229"));
        assert!(provider.supports_model("bedrock/my-team-profile"));
        assert!(!provider.supports_model("my-team-profile"));
        assert!(!provider.supports_model("gpt-4o"));
    }
}