feat: Add AWS Bedrock LLM Support

This commit adds support for AWS Bedrock for LLM parsing.

The implementation follows the approach of other LLM providers and uses the `BEDROCK_API_KEY` and `BEDROCK_REGION` environment variables for authentication.

This resolves issue #1162.

Author: belloibrahv

python/cocoindex/llm.py

@@ -14,6 +14,7 @@ class LlmApiType(Enum):
     OPEN_ROUTER = "OpenRouter"
     VOYAGE = "Voyage"
     VLLM = "Vllm"
+    BEDROCK = "Bedrock"
 
 
 @dataclass

python/cocoindex/tests/test_engine_value.py

@@ -1064,6 +1064,25 @@ def test_full_roundtrip_vector_numeric_types() -> None:
         validate_full_roundtrip(value_u64, Vector[np.uint64, Literal[3]])
 
 
+def test_llm_api_type_bedrock() -> None:
+    """Test that LlmApiType.BEDROCK is available and works."""
+    from cocoindex.llm import LlmApiType, LlmSpec
+
+    # Test enum availability
+    assert hasattr(LlmApiType, "BEDROCK")
+    assert LlmApiType.BEDROCK.value == "Bedrock"
+
+    # Test LlmSpec creation with Bedrock
+    spec = LlmSpec(
+        api_type=LlmApiType.BEDROCK, model="us.anthropic.claude-3-5-haiku-20241022-v1:0"
+    )
+
+    assert spec.api_type == LlmApiType.BEDROCK
+    assert spec.model == "us.anthropic.claude-3-5-haiku-20241022-v1:0"
+    assert spec.address is None
+    assert spec.api_config is None
+
+
 def test_full_roundtrip_vector_of_vector() -> None:
     """Test full roundtrip for vector of vector."""
     value_f32 = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32)

src/llm/bedrock.rs

@@ -0,0 +1,185 @@
+use crate::prelude::*;
+use base64::prelude::*;
+
+use crate::llm::{
+    LlmGenerateRequest, LlmGenerateResponse, LlmGenerationClient, OutputFormat,
+    ToJsonSchemaOptions, detect_image_mime_type,
+};
+use anyhow::Context;
+use urlencoding::encode;
+
+pub struct Client {
+    api_key: String,
+    region: String,
+    client: reqwest::Client,
+}
+
+impl Client {
+    pub async fn new(address: Option<String>) -> Result<Self> {
+        if address.is_some() {
+            api_bail!("Bedrock doesn't support custom API address");
+        }
+
+        let api_key = match std::env::var("BEDROCK_API_KEY") {
+            Ok(val) => val,
+            Err(_) => api_bail!("BEDROCK_API_KEY environment variable must be set"),
+        };
+
+        // Default to us-east-1 if no region specified
+        let region = std::env::var("BEDROCK_REGION").unwrap_or_else(|_| "us-east-1".to_string());
+
+        Ok(Self {
+            api_key,
+            region,
+            client: reqwest::Client::new(),
+        })
+    }
+}
+
+#[async_trait]
+impl LlmGenerationClient for Client {
+    async fn generate<'req>(
+        &self,
+        request: LlmGenerateRequest<'req>,
+    ) -> Result<LlmGenerateResponse> {
+        let mut user_content_parts: Vec<serde_json::Value> = Vec::new();
+
+        // Add image part if present
+        if let Some(image_bytes) = &request.image {
+            let base64_image = BASE64_STANDARD.encode(image_bytes.as_ref());
+            let mime_type = detect_image_mime_type(image_bytes.as_ref())?;
+            user_content_parts.push(serde_json::json!({
+                "image": {
+                    "format": mime_type.split('/').nth(1).unwrap_or("png"),
+                    "source": {
+                        "bytes": base64_image,
+                    }
+                }
+            }));
+        }
+
+        // Add text part
+        user_content_parts.push(serde_json::json!({
+            "text": request.user_prompt
+        }));
+
+        let messages = vec![serde_json::json!({
+            "role": "user",
+            "content": user_content_parts
+        })];
+
+        let mut payload = serde_json::json!({
+            "messages": messages,
+            "inferenceConfig": {
+                "maxTokens": 4096
+            }
+        });
+
+        // Add system prompt if present
+        if let Some(system) = request.system_prompt {
+            payload["system"] = serde_json::json!([{
+                "text": system
+            }]);
+        }
+
+        // Handle structured output using tool schema
+        if let Some(OutputFormat::JsonSchema { schema, name }) = request.output_format.as_ref() {
+            let schema_json = serde_json::to_value(schema)?;
+            payload["toolConfig"] = serde_json::json!({
+                "tools": [{
+                    "toolSpec": {
+                        "name": name,
+                        "description": format!("Extract structured data according to the schema"),
+                        "inputSchema": {
+                            "json": schema_json
+                        }
+                    }
+                }]
+            });
+        }
+
+        // Construct the Bedrock Runtime API URL
+        let url = format!(
+            "https://bedrock-runtime.{}.amazonaws.com/model/{}/converse",
+            self.region, request.model
+        );
+
+        let encoded_api_key = encode(&self.api_key);
+
+        let resp = retryable::run(
+            || async {
+                self.client
+                    .post(&url)
+                    .header(
+                        "Authorization",
+                        format!("Bearer {}", encoded_api_key.as_ref()),
+                    )
+                    .header("Content-Type", "application/json")
+                    .json(&payload)
+                    .send()
+                    .await?
+                    .error_for_status()
+            },
+            &retryable::HEAVY_LOADED_OPTIONS,
+        )
+        .await
+        .context("Bedrock API error")?;
+
+        let resp_json: serde_json::Value = resp.json().await.context("Invalid JSON")?;
+
+        // Check for errors in the response
+        if let Some(error) = resp_json.get("error") {
+            bail!("Bedrock API error: {:?}", error);
+        }
+
+        // Debug print full response (uncomment for debugging)
+        // println!("Bedrock API full response: {resp_json:?}");
+
+        // Extract the response content
+        let output = &resp_json["output"];
+        let message = &output["message"];
+        let content = &message["content"];
+
+        let text = if let Some(content_array) = content.as_array() {
+            // Look for tool use first (structured output)
+            let mut extracted_json: Option<serde_json::Value> = None;
+            for item in content_array {
+                if let Some(tool_use) = item.get("toolUse") {
+                    if let Some(input) = tool_use.get("input") {
+                        extracted_json = Some(input.clone());
+                        break;
+                    }
+                }
+            }
+
+            if let Some(json) = extracted_json {
+                // Return the structured output as JSON
+                serde_json::to_string(&json)?
+            } else {
+                // Fall back to text content
+                let mut text_parts = Vec::new();
+                for item in content_array {
+                    if let Some(text) = item.get("text") {
+                        if let Some(text_str) = text.as_str() {
+                            text_parts.push(text_str);
+                        }
+                    }
+                }
+                text_parts.join("")
+            }
+        } else {
+            return Err(anyhow::anyhow!("No content found in Bedrock response"));
+        };
+
+        Ok(LlmGenerateResponse { text })
+    }
+
+    fn json_schema_options(&self) -> ToJsonSchemaOptions {
+        ToJsonSchemaOptions {
+            fields_always_required: false,
+            supports_format: false,
+            extract_descriptions: false,
+            top_level_must_be_object: true,
+        }
+    }
+}

src/llm/mod.rs

@@ -18,6 +18,7 @@ pub enum LlmApiType {
     Voyage,
     Vllm,
     VertexAi,
+    Bedrock,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -106,6 +107,7 @@ pub trait LlmEmbeddingClient: Send + Sync {
 }
 
 mod anthropic;
+mod bedrock;
 mod gemini;
 mod litellm;
 mod ollama;
@@ -134,6 +136,9 @@ pub async fn new_llm_generation_client(
         LlmApiType::Anthropic => {
             Box::new(anthropic::Client::new(address).await?) as Box<dyn LlmGenerationClient>
         }
+        LlmApiType::Bedrock => {
+            Box::new(bedrock::Client::new(address).await?) as Box<dyn LlmGenerationClient>
+        }
         LlmApiType::LiteLlm => {
             Box::new(litellm::Client::new_litellm(address).await?) as Box<dyn LlmGenerationClient>
         }
@@ -169,7 +174,11 @@ pub async fn new_llm_embedding_client(
         }
         LlmApiType::VertexAi => Box::new(gemini::VertexAiClient::new(address, api_config).await?)
             as Box<dyn LlmEmbeddingClient>,
-        LlmApiType::OpenRouter | LlmApiType::LiteLlm | LlmApiType::Vllm | LlmApiType::Anthropic => {
+        LlmApiType::OpenRouter
+        | LlmApiType::LiteLlm
+        | LlmApiType::Vllm
+        | LlmApiType::Anthropic
+        | LlmApiType::Bedrock => {
             api_bail!("Embedding is not supported for API type {:?}", api_type)
         }
     };