diff --git a/.github/workflows/build_windows_worker_.yml b/.github/workflows/build_windows_worker_.yml index c61f607e17..79cde5003e 100644 --- a/.github/workflows/build_windows_worker_.yml +++ b/.github/workflows/build_windows_worker_.yml @@ -57,7 +57,7 @@ jobs: $env:OPENSSL_DIR="${Env:VCPKG_INSTALLATION_ROOT}\installed\x64-windows-static" mkdir frontend/build && cd backend New-Item -Path . -Name "windmill-api/openapi-deref.yaml" -ItemType "File" -Force - cargo build --release --features=enterprise,stripe,embedding,parquet,prometheus,openidconnect,cloud,jemalloc,tantivy,license,http_trigger,zip,oauth2,kafka,nats,sqs_trigger,postgres_trigger,gcp_trigger,mqtt_trigger,websocket,smtp,static_frontend,all_languages_windows,mcp,private + cargo build --release --features=enterprise,stripe,embedding,parquet,prometheus,openidconnect,cloud,jemalloc,tantivy,license,http_trigger,zip,oauth2,kafka,nats,sqs_trigger,postgres_trigger,gcp_trigger,mqtt_trigger,websocket,smtp,static_frontend,all_languages_windows,mcp,bedrock,private - name: Rename binary with corresponding architecture run: | Rename-Item -Path ".\backend\target\release\windmill.exe" -NewName "windmill-ee.exe" diff --git a/.github/workflows/docker-image.yml b/.github/workflows/docker-image.yml index 49e878f13c..cad751582f 100644 --- a/.github/workflows/docker-image.yml +++ b/.github/workflows/docker-image.yml @@ -97,7 +97,7 @@ jobs: platforms: linux/amd64,linux/arm64 push: true build-args: | - features=embedding,parquet,openidconnect,jemalloc,license,http_trigger,zip,oauth2,dind,postgres_trigger,mqtt_trigger,websocket,smtp,static_frontend,agent_worker_server,all_languages,deno_core,mcp,private + features=embedding,parquet,openidconnect,jemalloc,license,http_trigger,zip,oauth2,dind,postgres_trigger,mqtt_trigger,websocket,smtp,static_frontend,agent_worker_server,all_languages,deno_core,mcp,bedrock,private tags: | ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ env.DEV_SHA }} ${{ steps.meta-public.outputs.tags }} @@ -159,7 +159,7 @@ jobs: platforms: linux/amd64,linux/arm64 push: true build-args: | - features=enterprise,enterprise_saml,stripe,embedding,parquet,prometheus,openidconnect,cloud,jemalloc,agent_worker_server,tantivy,license,http_trigger,zip,oauth2,kafka,sqs_trigger,nats,otel,dind,postgres_trigger,mqtt_trigger,gcp_trigger,websocket,smtp,static_frontend,all_languages,private,deno_core,mcp + features=enterprise,enterprise_saml,stripe,embedding,parquet,prometheus,openidconnect,cloud,jemalloc,agent_worker_server,tantivy,license,http_trigger,zip,oauth2,kafka,sqs_trigger,nats,otel,dind,postgres_trigger,mqtt_trigger,gcp_trigger,websocket,smtp,static_frontend,all_languages,private,deno_core,mcp,bedrock tags: | ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}-ee:${{ env.DEV_SHA }} ${{ steps.meta-ee-public.outputs.tags }} diff --git a/.github/workflows/publish_windows_worker.yml b/.github/workflows/publish_windows_worker.yml index 95545c14a1..059f299331 100644 --- a/.github/workflows/publish_windows_worker.yml +++ b/.github/workflows/publish_windows_worker.yml @@ -59,7 +59,7 @@ jobs: $env:OPENSSL_DIR="${Env:VCPKG_INSTALLATION_ROOT}\installed\x64-windows-static" mkdir frontend/build && cd backend New-Item -Path . -Name "windmill-api/openapi-deref.yaml" -ItemType "File" -Force - cargo build --release --features=enterprise,stripe,embedding,parquet,prometheus,openidconnect,cloud,jemalloc,tantivy,license,http_trigger,zip,oauth2,kafka,sqs_trigger,nats,postgres_trigger,mqtt_trigger,gcp_trigger,websocket,smtp,static_frontend,all_languages_windows,mcp,private + cargo build --release --features=enterprise,stripe,embedding,parquet,prometheus,openidconnect,cloud,jemalloc,tantivy,license,http_trigger,zip,oauth2,kafka,sqs_trigger,nats,postgres_trigger,mqtt_trigger,gcp_trigger,websocket,smtp,static_frontend,all_languages_windows,mcp,bedrock,private - name: Rename binary with corresponding architecture run: | Rename-Item -Path ".\backend\target\release\windmill.exe" -NewName "windmill-ee.exe" diff --git a/backend/Cargo.lock b/backend/Cargo.lock index d1a6c5733d..8ac22c2cba 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -906,10 +906,33 @@ dependencies = [ ] [[package]] -name = "aws-sdk-bedrockruntime" -version = "1.113.0" +name = "aws-sdk-bedrock" +version = "1.129.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d5d2b8f081b9e8ff455b8dd7387b6b02263c3dac73172d188d2b523ff1e775e9" +checksum = "2d009c7cbb8332c805be2be32c0def27cb7bd4804d9333f5cf204abfd96b6d25" +dependencies = [ + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand", + "http 0.2.12", + "regex-lite", + "tracing", +] + +[[package]] +name = "aws-sdk-bedrockruntime" +version = "1.122.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c14680affe00d9bda89ddeb7a41004ac130165aef6678078da58df22dec3ff1" dependencies = [ "aws-credential-types", "aws-runtime", @@ -918,6 +941,7 @@ dependencies = [ "aws-smithy-eventstream", "aws-smithy-http", "aws-smithy-json", + "aws-smithy-observability", "aws-smithy-runtime", "aws-smithy-runtime-api", "aws-smithy-types", @@ -15454,12 +15478,15 @@ dependencies = [ "async_zip", "aws-config", "aws-credential-types", + "aws-sdk-bedrock", + "aws-sdk-bedrockruntime", "aws-sdk-config", "aws-sdk-sqs", "aws-sdk-sso", "aws-sdk-ssooidc", "aws-sdk-sts", "aws-sigv4", + "aws-smithy-types", "axum 0.7.9", "backon", "base32", @@ -15618,6 +15645,7 @@ dependencies = [ "async-trait", "aws-config", "aws-credential-types", + "aws-sdk-bedrockruntime", "aws-sdk-rds", "aws-sdk-sts", "aws-smithy-types", diff --git a/backend/Cargo.toml b/backend/Cargo.toml index 15f412acf2..4cbacd713e 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -77,6 +77,7 @@ websocket = ["windmill-api/websocket"] http_trigger = ["windmill-api/http_trigger"] postgres_trigger = ["windmill-api/postgres_trigger"] mcp = ["windmill-api/mcp", "windmill-worker/mcp"] +bedrock = ["windmill-api/bedrock", "windmill-worker/bedrock", "windmill-common/bedrock"] mqtt_trigger = ["windmill-api/mqtt_trigger"] sqs_trigger = ["windmill-api/sqs_trigger", "windmill-common/aws_auth", "windmill-api/openidconnect"] gcp_trigger = ["windmill-api/gcp_trigger"] @@ -106,7 +107,7 @@ all_languages_windows = ["python", "deno_core", "rust", "mysql", "oracledb", "du all_sqlx_features = ["all_languages", "enterprise", "enterprise_saml", "embedding", "parquet", "prometheus", "flow_testing", "openidconnect", "cloud", "jemalloc", "tantivy", "sqlx", "kafka", "nats", "otel", "dind", "websocket", "http_trigger", "postgres_trigger", "mcp", "mqtt_trigger", "sqs_trigger", "gcp_trigger", "smtp", "stripe", - "license", "oauth2", "zip", "static_frontend", "scoped_cache", "agent_worker_server"] + "license", "oauth2", "zip", "static_frontend", "scoped_cache", "agent_worker_server", "bedrock"] [patch.crates-io] object_store = { git = "https://github.com/apache/arrow-rs-object-store", rev = "36752c975d4f29e20b57c91f81a10872dcd48ae7" } @@ -394,7 +395,8 @@ datafusion = "47.0.0" object_store = { git = "https://github.com/apache/arrow-rs-object-store", rev = "36752c975d4f29e20b57c91f81a10872dcd48ae7", features = ["aws", "azure", "gcp"] } openidconnect = { version = "4.0.0-rc.1" } aws-config = "^1" -aws-sdk-bedrockruntime = "=1.113.0" +aws-sdk-bedrock = "1.129.0" +aws-sdk-bedrockruntime = "=1.122.0" aws-credential-types = "^1" aws-smithy-types = "^1" aws-sdk-sqs = "=1.77.0" diff --git a/backend/windmill-api/Cargo.toml b/backend/windmill-api/Cargo.toml index ae4372ac24..ae98191b09 100644 --- a/backend/windmill-api/Cargo.toml +++ b/backend/windmill-api/Cargo.toml @@ -32,11 +32,12 @@ http_trigger = ["dep:matchit", "dep:thiserror", "dep:sha1", "dep:constant_time_e static_frontend = ["dep:rust-embed"] postgres_trigger = ["dep:rust-postgres", "dep:pg_escape", "dep:byteorder", "dep:thiserror", "dep:rust_decimal", "dep:rust-postgres-native-tls"] mqtt_trigger = ["dep:thiserror", "dep:rumqttc"] -sqs_trigger = ["dep:aws-sdk-sqs", "dep:aws-sdk-sts", "dep:aws-sdk-sso", "dep:aws-sdk-ssooidc", "dep:thiserror", "dep:aws-config", "dep:backon"] +sqs_trigger = ["dep:aws-sdk-sqs", "dep:aws-sdk-sts", "dep:aws-sdk-sso", "dep:aws-sdk-ssooidc", "dep:thiserror", "dep:backon", "dep:aws-config"] deno_core = ["dep:deno_core", "dep:deno_error"] gcp_trigger = ["dep:thiserror", "dep:google-cloud-pubsub", "dep:google-cloud-googleapis", "dep:tonic"] cloud = ["windmill-common/cloud"] mcp = ["dep:windmill-mcp", "windmill-mcp/server", "windmill-mcp/auth"] +bedrock = ["dep:aws-sdk-bedrock", "dep:aws-sdk-bedrockruntime", "windmill-common/bedrock", "dep:aws-config"] python = [] [dependencies] @@ -151,6 +152,9 @@ aws-sigv4.workspace = true aws-sdk-config.workspace = true aws-config = { workspace = true, optional = true } aws-credential-types.workspace = true +aws-sdk-bedrock = { workspace = true, optional = true } +aws-sdk-bedrockruntime = { workspace = true, optional = true } +aws-smithy-types.workspace = true async-trait.workspace = true google-cloud-pubsub = { workspace = true, optional = true } google-cloud-googleapis = { workspace = true , optional = true } diff --git a/backend/windmill-api/src/ai.rs b/backend/windmill-api/src/ai.rs index 84f03c12b8..224167514d 100644 --- a/backend/windmill-api/src/ai.rs +++ b/backend/windmill-api/src/ai.rs @@ -1,7 +1,14 @@ +#[cfg(feature = "bedrock")] use crate::bedrock; use crate::db::{ApiAuthed, DB}; -use axum::{body::Bytes, extract::Path, response::IntoResponse, routing::post, Extension, Router}; +#[cfg(feature = "bedrock")] +use axum::routing::get; +use axum::{ + body::Bytes, extract::Path, response::IntoResponse, routing::post, Extension, Router, +}; +#[cfg(feature = "bedrock")] +use axum::Json; use http::{HeaderMap, Method}; use quick_cache::sync::Cache; use reqwest::{Client, RequestBuilder}; @@ -170,8 +177,11 @@ struct AIRequestConfig { pub access_token: Option, pub organization_id: Option, pub user: Option, + #[allow(dead_code)] pub region: Option, + #[allow(dead_code)] pub aws_access_key_id: Option, + #[allow(dead_code)] pub aws_secret_access_key: Option, pub platform: AnthropicPlatform, } @@ -323,7 +333,6 @@ impl AIRequestConfig { let is_anthropic = matches!(provider, AIProvider::Anthropic); let is_anthropic_vertex = is_anthropic && self.platform == AnthropicPlatform::GoogleVertexAi; let is_anthropic_sdk = headers.get("X-Anthropic-SDK").is_some(); - let is_bedrock = matches!(provider, AIProvider::AWSBedrock); let is_google_ai = matches!(provider, AIProvider::GoogleAI); // GoogleAI uses OpenAI-compatible endpoint in the proxy (for the chat), but not for the ai agent @@ -334,27 +343,8 @@ impl AIRequestConfig { }; let base_url = base_url.as_str(); - // Check if using IAM credentials for Bedrock (instead of bearer token) - let use_iam_auth = - is_bedrock && self.aws_access_key_id.is_some() && self.aws_secret_access_key.is_some(); - - // Handle AWS Bedrock transformation - let (url, body) = if is_bedrock && method != Method::GET { - let (model, transformed_body, is_streaming) = - bedrock::transform_openai_to_bedrock(&body)?; - let endpoint = if is_streaming { - "converse-stream" - } else { - "converse" - }; - let bedrock_url = format!("{}/model/{}/{}", base_url, model, endpoint); - (bedrock_url, transformed_body) - } else if is_bedrock && (path == "foundation-models" || path == "inference-profiles") { - // AWS Bedrock foundation-models and inference-profiles endpoints use different base URL (without -runtime) - let bedrock_base_url = base_url.replace("bedrock-runtime.", "bedrock."); - let bedrock_url = format!("{}/{}", bedrock_base_url, path); - (bedrock_url, body) - } else if is_anthropic_vertex && method != Method::GET { + // Build URL based on provider + let (url, body) = if is_anthropic_vertex && method != Method::GET { let (model, transformed_body) = transform_anthropic_for_vertex(&body)?; let vertex_url = format!("{}/{}:streamRawPredict", base_url, model); (vertex_url, transformed_body) @@ -387,40 +377,21 @@ impl AIRequestConfig { } } - // For Bedrock with IAM credentials, sign the request using SigV4 - if use_iam_auth { - let region = self.region.as_deref().ok_or_else(|| { - Error::internal_err("AWS region must be set for IAM authentication with Bedrock") - })?; - let signed_headers = bedrock::sign_bedrock_request( - method.as_str(), - &url, - &body, - self.aws_access_key_id.as_ref().unwrap(), - self.aws_secret_access_key.as_ref().unwrap(), - region, - )?; + // Add authentication headers + if let Some(api_key) = self.api_key { + if is_azure { + request = request.header("api-key", api_key.clone()) + } else { + request = request.header("authorization", format!("Bearer {}", api_key.clone())) + } + // For standard Anthropic API, also add X-API-Key header (but not for Vertex AI) + if is_anthropic && !is_anthropic_vertex { + request = request.header("X-API-Key", api_key); + } + } - for (header_name, header_value) in signed_headers { - request = request.header(header_name, header_value); - } - } else { - // For non-IAM auth, use bearer token or API key - if let Some(api_key) = self.api_key.clone() { - if is_azure { - request = request.header("api-key", api_key.clone()) - } else { - request = request.header("authorization", format!("Bearer {}", api_key.clone())) - } - // For standard Anthropic API, also add X-API-Key header - if is_anthropic && !is_anthropic_vertex { - request = request.header("X-API-Key", api_key); - } - } - - if let Some(access_token) = self.access_token { - request = request.header("authorization", format!("Bearer {}", access_token)) - } + if let Some(access_token) = self.access_token { + request = request.header("authorization", format!("Bearer {}", access_token)) } request = request.body(body); @@ -565,7 +536,22 @@ pub fn global_service() -> Router { } pub fn workspaced_service() -> Router { - Router::new().route("/proxy/*ai", post(proxy).get(proxy)) + let router = Router::new().route("/proxy/*ai", post(proxy).get(proxy)); + + #[cfg(feature = "bedrock")] + let router = router.route("/check_bedrock_credentials", get(check_bedrock_credentials)); + + router +} + +/// Check if AWS Bedrock credentials are available from environment variables. +#[cfg(feature = "bedrock")] +async fn check_bedrock_credentials( + _authed: ApiAuthed, + Path(_w_id): Path, +) -> Result> { + let response = windmill_common::ai_bedrock::check_env_credentials().await; + Ok(Json(response)) } async fn global_proxy( @@ -735,20 +721,101 @@ async fn proxy( ai_path = chat_path; } - // Extract model and streaming flag for Bedrock transformation (only for POST requests) - let (model, is_streaming) = - if matches!(provider, AIProvider::AWSBedrock) && method == Method::POST { - #[derive(Deserialize, Debug)] - struct BedrockRequest { - model: String, - stream: bool, + // Handle Bedrock-specific logic when the feature is enabled + #[cfg(feature = "bedrock")] + { + // Extract model and streaming flag for Bedrock transformation (only for POST requests) + let (model, is_streaming) = + if matches!(provider, AIProvider::AWSBedrock) && method == Method::POST { + #[derive(Deserialize, Debug)] + struct BedrockRequest { + model: String, + #[serde(default)] + stream: bool, + } + let parsed: BedrockRequest = serde_json::from_slice(&body) + .map_err(|e| Error::internal_err(format!("Failed to parse request body: {}", e)))?; + (Some(parsed.model), parsed.stream) + } else { + (None, false) + }; + + // For Bedrock requests, use the SDK-based approach + if matches!(provider, AIProvider::AWSBedrock) { + let region = request_config + .region + .as_deref() + .ok_or_else(|| Error::internal_err("AWS region must be set for Bedrock"))?; + + // Audit log before making the SDK request + let mut tx = db.begin().await?; + audit_log( + &mut *tx, + &authed, + "ai.request", + ActionKind::Execute, + &w_id, + Some(&authed.email), + Some([("ai_config_path", &format!("{:?}", ai_path)[..])].into()), + ) + .await?; + tx.commit().await?; + + // Handle GET requests for control plane operations + if method == Method::GET { + if ai_path == "foundation-models" { + return bedrock::list_foundation_models( + request_config.api_key.as_deref(), + request_config.aws_access_key_id.as_deref(), + request_config.aws_secret_access_key.as_deref(), + region, + ) + .await; + } else if ai_path == "inference-profiles" { + return bedrock::list_inference_profiles( + request_config.api_key.as_deref(), + request_config.aws_access_key_id.as_deref(), + request_config.aws_secret_access_key.as_deref(), + region, + ) + .await; + } } - let parsed: BedrockRequest = serde_json::from_slice(&body) - .map_err(|e| Error::internal_err(format!("Failed to parse request body: {}", e)))?; - (Some(parsed.model), parsed.stream) - } else { - (None, false) - }; + + // Handle POST requests for inference + if method == Method::POST && model.is_some() { + if is_streaming { + return bedrock::handle_bedrock_sdk_streaming( + model.as_ref().unwrap(), + &body, + request_config.api_key.as_deref(), + request_config.aws_access_key_id.as_deref(), + request_config.aws_secret_access_key.as_deref(), + region, + ) + .await; + } else { + return bedrock::handle_bedrock_sdk_non_streaming( + model.as_ref().unwrap(), + &body, + request_config.api_key.as_deref(), + request_config.aws_access_key_id.as_deref(), + request_config.aws_secret_access_key.as_deref(), + region, + ) + .await; + } + } + } + } + + // When bedrock feature is disabled, return error for Bedrock provider + #[cfg(not(feature = "bedrock"))] + if matches!(provider, AIProvider::AWSBedrock) { + return Err(Error::BadRequest( + "AWS Bedrock support is not enabled. Build with 'bedrock' feature.".to_string(), + )); + } let request = request_config.prepare_request(&provider, &ai_path, method, headers, body)?; @@ -773,45 +840,8 @@ async fn proxy( return Err(Error::AIError(err_msg)); } - // Transform Bedrock responses back to OpenAI format - if matches!(provider, AIProvider::AWSBedrock) && model.is_some() { - if is_streaming { - // Transform streaming response - use http::StatusCode; - - let mut response_headers = HeaderMap::new(); - response_headers.insert("content-type", "text/event-stream".parse().unwrap()); - response_headers.insert("cache-control", "no-cache".parse().unwrap()); - response_headers.insert("connection", "keep-alive".parse().unwrap()); - - let stream = response.bytes_stream(); - let transformed_stream = - bedrock::transform_bedrock_stream_to_openai(stream, model.unwrap()); - - Ok(( - StatusCode::OK, - response_headers, - axum::body::Body::from_stream(transformed_stream), - )) - } else { - // Transform non-streaming response - let transformed_body = - bedrock::transform_bedrock_to_openai(response, model.unwrap()).await?; - - let mut response_headers = HeaderMap::new(); - response_headers.insert("content-type", "application/json".parse().unwrap()); - - Ok(( - http::StatusCode::OK, - response_headers, - axum::body::Body::from(transformed_body), - )) - } - } else { - // Pass through for other providers - let status_code = response.status(); - let headers = response.headers().clone(); - let stream = response.bytes_stream(); - Ok((status_code, headers, axum::body::Body::from_stream(stream))) - } + let status_code = response.status(); + let headers = response.headers().clone(); + let stream = response.bytes_stream(); + Ok((status_code, headers, axum::body::Body::from_stream(stream))) } diff --git a/backend/windmill-api/src/bedrock.rs b/backend/windmill-api/src/bedrock.rs index 81ee7056b6..8965c8c377 100644 --- a/backend/windmill-api/src/bedrock.rs +++ b/backend/windmill-api/src/bedrock.rs @@ -1,318 +1,648 @@ +//! AWS Bedrock SDK-based operations for the AI chat proxy. +//! +//! This module provides SDK-based request handling for Bedrock: +//! +//! ## Inference (Runtime SDK): +//! - `handle_bedrock_sdk_streaming`: Uses BedrockClient for streaming requests +//! - `handle_bedrock_sdk_non_streaming`: Uses BedrockClient for non-streaming requests +//! - `sdk_stream_to_sse`: Converts SDK ConverseStream events to SSE format +//! +//! ## Control Plane (Bedrock SDK): +//! - `list_foundation_models`: Lists available foundation models +//! - `list_inference_profiles`: Lists inference profiles +//! +//! Shared AWS SDK code is available in `windmill_common::ai_bedrock`, including: +//! - `BedrockClient`: SDK wrapper with bearer token and IAM auth +//! - Stream event parsing functions +//! - Helper utilities + use axum::body::Bytes; -use aws_sigv4::http_request::{sign, SignableBody, SignableRequest, SigningSettings}; -use aws_sigv4::sign::v4; -use bytes; -use futures; -use std::time::SystemTime; -use uuid; +use serde::Deserialize; +use windmill_common::ai_bedrock::{ + bedrock_stream_event_is_block_stop, bedrock_stream_event_to_text, + bedrock_stream_event_to_tool_delta, bedrock_stream_event_to_tool_start, format_bedrock_error, + BedrockClient, OpenAIMessage, OpenAIToolCall, +}; use windmill_common::error::{Error, Result}; -/// Sign a request for AWS Bedrock using SigV4 -/// -/// Returns a vector of (header_name, header_value) tuples to add to the request -pub fn sign_bedrock_request( - method: &str, - uri: &str, - body: &[u8], - access_key_id: &str, - secret_access_key: &str, - region: &str, -) -> Result> { - let identity = aws_credential_types::Credentials::new( - access_key_id, - secret_access_key, - None, // session token - None, // expiration - "windmill", - ) - .into(); +// ============================================================================ +// Shared Request Types for SDK-Based Handlers +// ============================================================================ - let signing_settings = SigningSettings::default(); - let signing_params = v4::SigningParams::builder() - .identity(&identity) - .region(region) - .name("bedrock") - .time(SystemTime::now()) - .settings(signing_settings) - .build() - .map_err(|e| Error::internal_err(format!("Failed to build signing params: {}", e)))?; - - // Parse the URI to extract path and query - let parsed_uri: http::Uri = uri - .parse() - .map_err(|e| Error::internal_err(format!("Failed to parse URI: {}", e)))?; - - let path_and_query = parsed_uri - .path_and_query() - .map(|pq| pq.as_str()) - .unwrap_or("/"); - - let signable_request = SignableRequest::new( - method, - path_and_query, - std::iter::once(("host", parsed_uri.host().unwrap_or(""))), - SignableBody::Bytes(body), - ) - .map_err(|e| Error::internal_err(format!("Failed to create signable request: {}", e)))?; - - let (signing_instructions, _signature) = sign(signable_request, &signing_params.into()) - .map_err(|e| Error::internal_err(format!("Failed to sign request: {}", e)))? - .into_parts(); - - // Collect the headers to add - let mut headers = Vec::new(); - for (name, value) in signing_instructions.headers() { - headers.push((name.to_string(), value.to_string())); - } - - Ok(headers) +/// OpenAI-format request body for Bedrock SDK handlers +#[derive(Deserialize, Debug)] +struct OpenAIRequest { + messages: Vec, + #[serde(default)] + tools: Option>, + #[serde(default)] + tool_choice: Option, + #[serde(default)] + max_tokens: Option, + #[serde(default)] + temperature: Option, } -/// Transform OpenAI format request to AWS Bedrock Converse format -/// Returns: (model_id, transformed_body, is_streaming) -pub fn transform_openai_to_bedrock(body: &[u8]) -> Result<(String, Bytes, bool)> { - use serde_json::Value; +#[derive(Deserialize, Debug)] +struct OpenAIToolDef { + #[serde(default)] + #[allow(dead_code)] + r#type: Option, + function: OpenAIToolFunction, +} - // Parse the OpenAI request - let openai_req: Value = serde_json::from_slice(body) +#[derive(Deserialize, Debug)] +struct OpenAIToolFunction { + name: String, + #[serde(default)] + description: Option, + #[serde(default)] + parameters: Option, +} + +// ============================================================================ +// Shared Helper Functions for SDK-Based Handlers +// ============================================================================ + +/// Authentication configuration for Bedrock clients +enum BedrockAuthConfig { + BearerToken(String), + IamCredentials { + access_key_id: String, + secret_access_key: String, + }, + Environment, +} + +/// Determine auth configuration with priority: bearer token → IAM credentials → environment +fn determine_auth_config( + api_key: Option<&str>, + aws_access_key_id: Option<&str>, + aws_secret_access_key: Option<&str>, +) -> BedrockAuthConfig { + if let Some(key) = api_key.filter(|k| !k.is_empty()) { + BedrockAuthConfig::BearerToken(key.to_string()) + } else if let (Some(access_key_id), Some(secret_access_key)) = ( + aws_access_key_id.filter(|s| !s.is_empty()), + aws_secret_access_key.filter(|s| !s.is_empty()), + ) { + BedrockAuthConfig::IamCredentials { + access_key_id: access_key_id.to_string(), + secret_access_key: secret_access_key.to_string(), + } + } else { + BedrockAuthConfig::Environment + } +} + +/// Create a BedrockClient with auth priority: bearer token → IAM credentials → environment +async fn create_bedrock_client( + api_key: Option<&str>, + aws_access_key_id: Option<&str>, + aws_secret_access_key: Option<&str>, + region: &str, +) -> Result { + match determine_auth_config(api_key, aws_access_key_id, aws_secret_access_key) { + BedrockAuthConfig::BearerToken(key) => { + BedrockClient::from_bearer_token(key, region).await + } + BedrockAuthConfig::IamCredentials { + access_key_id, + secret_access_key, + } => { + BedrockClient::from_credentials(access_key_id, secret_access_key, None, region).await + } + BedrockAuthConfig::Environment => BedrockClient::from_env(region).await, + } +} + +/// Convert OpenAIToolDef array to tool configuration for Bedrock SDK +fn build_tool_config_from_request( + tools: Option<&[OpenAIToolDef]>, + tool_choice: Option<&serde_json::Value>, +) -> Result> { + if let Some(tools) = tools { + let tool_defs: Vec = tools + .iter() + .map(|t| windmill_common::ai_bedrock::ToolDef { + r#type: "function".to_string(), + function: windmill_common::ai_bedrock::ToolDefFunction { + name: t.function.name.clone(), + description: t.function.description.clone(), + parameters: Box::from( + serde_json::value::RawValue::from_string( + serde_json::to_string( + &t.function + .parameters + .clone() + .unwrap_or(serde_json::json!({})), + ) + .unwrap_or_default(), + ) + .unwrap_or_else(|_| { + serde_json::value::RawValue::from_string("{}".to_string()).unwrap() + }), + ), + }, + }) + .collect(); + + // Determine if we should force tool use based on tool_choice + let force_tool_use = tool_choice + .map(|tc| tc == "required" || tc.as_str() == Some("required")) + .unwrap_or(false); + + windmill_common::ai_bedrock::build_tool_config(Some(&tool_defs), force_tool_use) + } else { + Ok(None) + } +} + +// ============================================================================ +// Control Plane Operations (using aws-sdk-bedrock) +// ============================================================================ + +/// Create a Bedrock control plane client with auth priority: bearer token → IAM credentials → environment +async fn create_bedrock_control_client( + api_key: Option<&str>, + aws_access_key_id: Option<&str>, + aws_secret_access_key: Option<&str>, + region: &str, +) -> Result { + use aws_config::BehaviorVersion; + use windmill_common::ai_bedrock::BearerTokenProvider; + + let region_provider = aws_sdk_bedrock::config::Region::new(region.to_string()); + + match determine_auth_config(api_key, aws_access_key_id, aws_secret_access_key) { + BedrockAuthConfig::BearerToken(key) => { + let config = aws_sdk_bedrock::config::Builder::new() + .region(region_provider) + .behavior_version(BehaviorVersion::latest()) + .token_provider(BearerTokenProvider::new(key)) + .build(); + Ok(aws_sdk_bedrock::Client::from_conf(config)) + } + BedrockAuthConfig::IamCredentials { + access_key_id, + secret_access_key, + } => { + let credentials = aws_credential_types::Credentials::new( + access_key_id, + secret_access_key, + None, + None, + "windmill", + ); + let config = aws_sdk_bedrock::config::Builder::new() + .region(region_provider) + .behavior_version(BehaviorVersion::latest()) + .credentials_provider(credentials) + .build(); + Ok(aws_sdk_bedrock::Client::from_conf(config)) + } + BedrockAuthConfig::Environment => { + let config = aws_config::defaults(BehaviorVersion::latest()) + .region(region_provider) + .load() + .await; + Ok(aws_sdk_bedrock::Client::new(&config)) + } + } +} + +/// List foundation models using the Bedrock SDK +pub async fn list_foundation_models( + api_key: Option<&str>, + aws_access_key_id: Option<&str>, + aws_secret_access_key: Option<&str>, + region: &str, +) -> Result<(http::StatusCode, http::HeaderMap, axum::body::Body)> { + let client = + create_bedrock_control_client(api_key, aws_access_key_id, aws_secret_access_key, region) + .await?; + + let response = client + .list_foundation_models() + .send() + .await + .map_err(|e| Error::internal_err(format!("Failed to list foundation models: {}", e)))?; + + // Convert to JSON response + let models: Vec = response + .model_summaries() + .iter() + .map(|m| { + serde_json::json!({ + "modelId": m.model_id(), + "modelName": m.model_name(), + "providerName": m.provider_name(), + "modelArn": m.model_arn(), + "inputModalities": m.input_modalities().iter().map(|i| i.as_str()).collect::>(), + "outputModalities": m.output_modalities().iter().map(|o| o.as_str()).collect::>(), + "responseStreamingSupported": m.response_streaming_supported(), + "inferenceTypesSupported": m.inference_types_supported().iter().map(|i| i.as_str()).collect::>(), + }) + }) + .collect(); + + let body = serde_json::json!({ "modelSummaries": models }); + let body_bytes = serde_json::to_vec(&body) + .map_err(|e| Error::internal_err(format!("Failed to serialize response: {}", e)))?; + + let mut headers = http::HeaderMap::new(); + headers.insert("content-type", "application/json".parse().unwrap()); + + Ok(( + http::StatusCode::OK, + headers, + axum::body::Body::from(body_bytes), + )) +} + +/// List inference profiles using the Bedrock SDK +pub async fn list_inference_profiles( + api_key: Option<&str>, + aws_access_key_id: Option<&str>, + aws_secret_access_key: Option<&str>, + region: &str, +) -> Result<(http::StatusCode, http::HeaderMap, axum::body::Body)> { + let client = + create_bedrock_control_client(api_key, aws_access_key_id, aws_secret_access_key, region) + .await?; + + let response = + client.list_inference_profiles().send().await.map_err(|e| { + Error::internal_err(format!("Failed to list inference profiles: {}", e)) + })?; + + // Convert to JSON response + let profiles: Vec = response + .inference_profile_summaries() + .iter() + .map(|p| { + serde_json::json!({ + "inferenceProfileId": p.inference_profile_id(), + "inferenceProfileName": p.inference_profile_name(), + "inferenceProfileArn": p.inference_profile_arn(), + "description": p.description(), + "status": p.status().as_str(), + "type": p.r#type().as_str(), + }) + }) + .collect(); + + let body = serde_json::json!({ "inferenceProfileSummaries": profiles }); + let body_bytes = serde_json::to_vec(&body) + .map_err(|e| Error::internal_err(format!("Failed to serialize response: {}", e)))?; + + let mut headers = http::HeaderMap::new(); + headers.insert("content-type", "application/json".parse().unwrap()); + + Ok(( + http::StatusCode::OK, + headers, + axum::body::Body::from(body_bytes), + )) +} + +// ============================================================================ +// Inference Operations (using aws-sdk-bedrockruntime) +// ============================================================================ + +/// Handle Bedrock streaming request using the AWS SDK. +/// +/// This function uses the shared BedrockClient to make streaming requests +/// and converts the SDK stream events to SSE format for the proxy response. +/// +/// Auth priority: bearer token → IAM credentials → environment credentials +pub async fn handle_bedrock_sdk_streaming( + model: &str, + body: &Bytes, + api_key: Option<&str>, + aws_access_key_id: Option<&str>, + aws_secret_access_key: Option<&str>, + region: &str, +) -> Result<(http::StatusCode, http::HeaderMap, axum::body::Body)> { + let openai_req: OpenAIRequest = serde_json::from_slice(body) .map_err(|e| Error::internal_err(format!("Failed to parse OpenAI request: {}", e)))?; - // Extract model and streaming flag - let model = openai_req["model"] - .as_str() - .ok_or_else(|| Error::BadRequest("Missing 'model' field in request".to_string()))? - .to_string(); + // Create Bedrock client using shared helper + let bedrock_client = + create_bedrock_client(api_key, aws_access_key_id, aws_secret_access_key, region).await?; - let is_streaming = openai_req["stream"].as_bool().unwrap_or(false); + // Convert messages using shared conversion + let (bedrock_messages, system_prompts) = + windmill_common::ai_bedrock::openai_messages_to_bedrock(&openai_req.messages)?; - // Build Bedrock request - let mut bedrock_req = serde_json::json!({}); + // Build inference configuration + let inference_config = windmill_common::ai_bedrock::create_inference_config( + openai_req.temperature, + openai_req.max_tokens, + ); - // Transform messages - if let Some(messages) = openai_req["messages"].as_array() { - let mut system_messages = Vec::new(); - let mut conversation_messages = Vec::new(); + // Convert tools using shared helper + let tool_config = build_tool_config_from_request( + openai_req.tools.as_deref(), + openai_req.tool_choice.as_ref(), + )?; - for msg in messages { - let role = msg["role"].as_str().unwrap_or(""); + // Build the SDK request + let mut request_builder = bedrock_client + .client() + .converse_stream() + .model_id(model) + .set_messages(Some(bedrock_messages)); - match role { - "system" => { - // Extract system messages to separate array - if let Some(content) = msg["content"].as_str() { - system_messages.push(serde_json::json!({"text": content})); - } - } - "user" | "assistant" => { - // Normalize content to array format - let mut content = if let Some(text) = msg["content"].as_str() { - // Simple string → array of content blocks - vec![serde_json::json!({"text": text})] - } else if let Some(content_array) = msg["content"].as_array() { - // Already an array - transform each item - content_array - .iter() - .filter_map(|item| { - if let Some(text) = item["text"].as_str() { - Some(serde_json::json!({"text": text})) - } else if item["type"].as_str() == Some("text") { - Some(serde_json::json!({"text": item["text"]})) - } else if item["type"].as_str() == Some("image_url") { - // Transform image_url format if needed - // For now, pass through - may need more sophisticated handling - Some(item.clone()) - } else { - None - } - }) - .collect() - } else { - vec![] - }; - - // Handle tool_calls for assistant messages (OpenAI → Bedrock toolUse) - if role == "assistant" { - if let Some(tool_calls) = msg["tool_calls"].as_array() { - for tool_call in tool_calls { - if tool_call["type"].as_str() == Some("function") { - let tool_use_id = tool_call["id"].as_str().unwrap_or(""); - let function_name = - tool_call["function"]["name"].as_str().unwrap_or(""); - let arguments_str = - tool_call["function"]["arguments"].as_str().unwrap_or("{}"); - - // Parse arguments JSON string to object - let input = serde_json::from_str::(arguments_str) - .map_err(|e| { - Error::internal_err(format!( - "Failed to parse tool call arguments: {}", - e - )) - })?; - - content.push(serde_json::json!({ - "toolUse": { - "toolUseId": tool_use_id, - "name": function_name, - "input": input - } - })); - } - } - } - } - - // Only add message if it has content - if !content.is_empty() { - conversation_messages.push(serde_json::json!({ - "role": role, - "content": content - })); - } - } - "tool" => { - // Transform tool response to Bedrock format - let tool_call_id = msg["tool_call_id"].as_str().unwrap_or(""); - let content = msg["content"].as_str().unwrap_or(""); - - // Try to parse content as JSON - // Bedrock requires json field to be an object, not a primitive or array - let tool_result_content = - if let Ok(json_content) = serde_json::from_str::(content) { - if json_content.is_object() { - vec![serde_json::json!({"json": json_content})] - } else { - // Wrap primitives and arrays in an object - vec![serde_json::json!({"json": {"result": json_content}})] - } - } else { - vec![serde_json::json!({"text": content})] - }; - - conversation_messages.push(serde_json::json!({ - "role": "user", - "content": [{ - "toolResult": { - "toolUseId": tool_call_id, - "content": tool_result_content - } - }] - })); - } - _ => {} - } - } - - if !system_messages.is_empty() { - bedrock_req["system"] = Value::Array(system_messages); - } - bedrock_req["messages"] = Value::Array(conversation_messages); + if !system_prompts.is_empty() { + request_builder = request_builder.set_system(Some(system_prompts)); } - // Transform inference parameters - let mut inference_config = serde_json::json!({}); - if let Some(max_tokens) = openai_req["max_tokens"].as_i64() { - inference_config["maxTokens"] = Value::Number(max_tokens.into()); - } - if let Some(temperature) = openai_req["temperature"].as_f64() { - inference_config["temperature"] = serde_json::json!(temperature); - } - if let Some(top_p) = openai_req["top_p"].as_f64() { - inference_config["topP"] = serde_json::json!(top_p); - } - if let Some(stop) = openai_req["stop"].as_array() { - let stop_sequences: Vec = stop - .iter() - .filter_map(|s| s.as_str().map(|s| s.to_string())) - .collect(); - if !stop_sequences.is_empty() { - inference_config["stopSequences"] = - Value::Array(stop_sequences.into_iter().map(Value::String).collect()); - } - } - if !inference_config.as_object().unwrap().is_empty() { - bedrock_req["inferenceConfig"] = inference_config; + if let Some(config) = inference_config { + request_builder = request_builder.inference_config(config); } - // Transform tools if present - if let Some(tools) = openai_req["tools"].as_array() { - let mut bedrock_tools = Vec::new(); - - for tool in tools { - if tool["type"].as_str() == Some("function") { - if let Some(function) = tool["function"].as_object() { - bedrock_tools.push(serde_json::json!({ - "toolSpec": { - "name": function.get("name"), - "description": function.get("description") - .and_then(|v| v.as_str()) - .filter(|s| !s.is_empty()) - .unwrap_or("Tool function"), - "inputSchema": { - "json": function.get("parameters") - } - } - })); - } - } - } - - if !bedrock_tools.is_empty() { - let mut tool_config = serde_json::json!({ - "tools": bedrock_tools - }); - - // Transform tool_choice - if let Some(tool_choice) = openai_req.get("tool_choice") { - if tool_choice == "auto" { - tool_config["toolChoice"] = serde_json::json!({"auto": {}}); - } else if tool_choice == "required" { - tool_config["toolChoice"] = serde_json::json!({"any": {}}); - } else if let Some(obj) = tool_choice.as_object() { - if obj.get("type").and_then(|v| v.as_str()) == Some("function") { - if let Some(function) = obj.get("function").and_then(|v| v.as_object()) { - if let Some(name) = function.get("name").and_then(|v| v.as_str()) { - tool_config["toolChoice"] = serde_json::json!({ - "tool": {"name": name} - }); - } - } - } - } - } - - bedrock_req["toolConfig"] = tool_config; - } + if let Some(config) = tool_config { + request_builder = request_builder.set_tool_config(Some(config)); } - let transformed_body = serde_json::to_vec(&bedrock_req) - .map_err(|e| Error::internal_err(format!("Failed to serialize Bedrock request: {}", e)))? - .into(); + // Send the request and get the stream + tracing::debug!("Bedrock SDK streaming: sending converse_stream request"); + let stream_output = request_builder.send().await.map_err(|e| { + let error_msg = format!("Bedrock SDK streaming error: {}", format_bedrock_error(&e)); + tracing::error!("Bedrock SDK streaming failed: {}", error_msg); + Error::internal_err(error_msg) + })?; + tracing::debug!("Bedrock SDK streaming: stream established successfully"); - Ok((model, transformed_body, is_streaming)) + // Convert SDK stream to SSE (pass the inner stream, not the full output) + let sse_stream = sdk_stream_to_sse(stream_output.stream, model.to_string()); + + // Build response headers + let mut response_headers = http::HeaderMap::new(); + response_headers.insert("content-type", "text/event-stream".parse().unwrap()); + response_headers.insert("cache-control", "no-cache".parse().unwrap()); + response_headers.insert("connection", "keep-alive".parse().unwrap()); + + Ok(( + http::StatusCode::OK, + response_headers, + axum::body::Body::from_stream(sse_stream), + )) } -/// Transform AWS Bedrock Converse response to OpenAI format -pub async fn transform_bedrock_to_openai( - response: reqwest::Response, +/// Convert AWS SDK ConverseStream events to SSE format. +/// +/// Uses shared stream parsing functions from windmill_common::ai_bedrock +/// to extract text deltas and tool calls from the SDK stream events. +pub fn sdk_stream_to_sse( + stream: aws_sdk_bedrockruntime::primitives::event_stream::EventReceiver< + aws_sdk_bedrockruntime::types::ConverseStreamOutput, + aws_sdk_bedrockruntime::types::error::ConverseStreamOutputError, + >, model: String, -) -> Result { - use serde_json::Value; +) -> impl futures::Stream> + Send { + use std::collections::HashMap; - let bedrock_resp: Value = response - .json() - .await - .map_err(|e| Error::internal_err(format!("Failed to parse Bedrock response: {}", e)))?; - - // Generate unique ID and timestamp let id = format!("chatcmpl-{}", uuid::Uuid::new_v4().simple()); let created = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap() .as_secs(); - // Extract stop reason and map to finish_reason - let stop_reason = bedrock_resp["stopReason"].as_str().unwrap_or("end_turn"); + // State to track partial tool calls + struct StreamState { + id: String, + model: String, + created: u64, + tool_calls: HashMap, // index -> (id, name, args) + current_tool_index: usize, + } + + let state = std::sync::Arc::new(tokio::sync::Mutex::new(StreamState { + id: id.clone(), + model: model.clone(), + created, + tool_calls: HashMap::new(), + current_tool_index: 0, + })); + + async_stream::stream! { + let mut stream = stream; + let state = state.clone(); + + loop { + match stream.recv().await { + Ok(Some(event)) => { + let mut state = state.lock().await; + + // Handle tool use start + if let Some(tool_call) = bedrock_stream_event_to_tool_start(&event) { + let index = state.current_tool_index; + state.tool_calls.insert( + index, + (tool_call.id.clone(), tool_call.name.clone(), String::new()), + ); + + // Send initial tool call chunk + let chunk = serde_json::json!({ + "id": state.id, + "object": "chat.completion.chunk", + "created": state.created, + "model": state.model, + "choices": [{ + "index": 0, + "delta": { + "tool_calls": [{ + "index": index, + "id": tool_call.id, + "type": "function", + "function": { + "name": tool_call.name, + "arguments": "" + } + }] + }, + "finish_reason": serde_json::Value::Null + }] + }); + + yield Ok(bytes::Bytes::from(format!("data: {}\n\n", chunk))); + } + + // Handle text delta + if let Some(text) = bedrock_stream_event_to_text(&event) { + let chunk = serde_json::json!({ + "id": state.id, + "object": "chat.completion.chunk", + "created": state.created, + "model": state.model, + "choices": [{ + "index": 0, + "delta": { + "content": text + }, + "finish_reason": serde_json::Value::Null + }] + }); + + yield Ok(bytes::Bytes::from(format!("data: {}\n\n", chunk))); + } + + // Handle tool use input delta + if let Some(input_delta) = bedrock_stream_event_to_tool_delta(&event) { + let index = state.current_tool_index; + if let Some((_id, _name, ref mut args)) = state.tool_calls.get_mut(&index) { + args.push_str(&input_delta); + + let chunk = serde_json::json!({ + "id": state.id, + "object": "chat.completion.chunk", + "created": state.created, + "model": state.model, + "choices": [{ + "index": 0, + "delta": { + "tool_calls": [{ + "index": index, + "function": { + "arguments": input_delta + } + }] + }, + "finish_reason": serde_json::Value::Null + }] + }); + + yield Ok(bytes::Bytes::from(format!("data: {}\n\n", chunk))); + } + } + + // Handle content block stop + if bedrock_stream_event_is_block_stop(&event) { + state.current_tool_index += 1; + } + + // Handle message stop + if let aws_sdk_bedrockruntime::types::ConverseStreamOutput::MessageStop(stop) = &event { + let stop_reason = stop.stop_reason().as_str(); + let finish_reason = match stop_reason { + "end_turn" => "stop", + "max_tokens" => "length", + "tool_use" => "tool_calls", + "stop_sequence" => "stop", + "guardrail_intervened" | "content_filtered" => "content_filter", + _ => "stop", + }; + + let chunk = serde_json::json!({ + "id": state.id, + "object": "chat.completion.chunk", + "created": state.created, + "model": state.model, + "choices": [{ + "index": 0, + "delta": {}, + "finish_reason": finish_reason + }] + }); + + yield Ok(bytes::Bytes::from(format!("data: {}\n\n", chunk))); + } + } + Ok(None) => break, + Err(e) => { + yield Err(std::io::Error::new( + std::io::ErrorKind::Other, + e.to_string(), + )); + break; + } + } + } + + // Send [DONE] at the end + yield Ok(bytes::Bytes::from("data: [DONE]\n\n")); + } +} + +/// Handle non-streaming Bedrock request using the AWS SDK. +/// +/// Auth priority: bearer token → IAM credentials → environment credentials +pub async fn handle_bedrock_sdk_non_streaming( + model: &str, + body: &Bytes, + api_key: Option<&str>, + aws_access_key_id: Option<&str>, + aws_secret_access_key: Option<&str>, + region: &str, +) -> Result<(http::StatusCode, http::HeaderMap, axum::body::Body)> { + let openai_req: OpenAIRequest = serde_json::from_slice(body) + .map_err(|e| Error::internal_err(format!("Failed to parse OpenAI request: {}", e)))?; + + // Create Bedrock client using shared helper + let bedrock_client = + create_bedrock_client(api_key, aws_access_key_id, aws_secret_access_key, region).await?; + + // Convert messages using shared conversion + let (bedrock_messages, system_prompts) = + windmill_common::ai_bedrock::openai_messages_to_bedrock(&openai_req.messages)?; + + // Build inference configuration + let inference_config = windmill_common::ai_bedrock::create_inference_config( + openai_req.temperature, + openai_req.max_tokens, + ); + + // Convert tools using shared helper + let tool_config = build_tool_config_from_request( + openai_req.tools.as_deref(), + openai_req.tool_choice.as_ref(), + )?; + + // Build the SDK request (non-streaming) + let mut request_builder = bedrock_client + .client() + .converse() + .model_id(model) + .set_messages(Some(bedrock_messages)); + + if !system_prompts.is_empty() { + request_builder = request_builder.set_system(Some(system_prompts)); + } + + if let Some(config) = inference_config { + request_builder = request_builder.inference_config(config); + } + + if let Some(config) = tool_config { + request_builder = request_builder.set_tool_config(Some(config)); + } + + // Send the request + tracing::debug!("Bedrock SDK non-streaming: sending converse request"); + let response = request_builder.send().await.map_err(|e| { + let error_msg = format!( + "Bedrock SDK non-streaming error: {}", + format_bedrock_error(&e) + ); + tracing::error!("Bedrock SDK non-streaming failed: {}", error_msg); + Error::internal_err(error_msg) + })?; + tracing::debug!( + "Bedrock SDK non-streaming: response received, stop_reason={}", + response.stop_reason().as_str() + ); + + // Convert response to OpenAI format + let id = format!("chatcmpl-{}", uuid::Uuid::new_v4().simple()); + let created = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + + // Extract stop reason + let stop_reason = response.stop_reason().as_str(); let finish_reason = match stop_reason { "end_turn" => "stop", "max_tokens" => "length", @@ -323,28 +653,31 @@ pub async fn transform_bedrock_to_openai( }; // Extract message content - let message_content = &bedrock_resp["output"]["message"]["content"]; let mut text_content = String::new(); - let mut tool_calls = Vec::new(); + let mut tool_calls: Vec = Vec::new(); - if let Some(content_array) = message_content.as_array() { - for (_index, block) in content_array.iter().enumerate() { - if let Some(text) = block["text"].as_str() { - text_content.push_str(text); - } else if let Some(tool_use) = block.get("toolUse") { - // Transform tool use to OpenAI tool_calls format - let tool_call_id = tool_use["toolUseId"].as_str().unwrap_or(""); - let name = tool_use["name"].as_str().unwrap_or(""); - let input = &tool_use["input"]; - - tool_calls.push(serde_json::json!({ - "id": tool_call_id, - "type": "function", - "function": { - "name": name, - "arguments": serde_json::to_string(input).unwrap_or_default() + if let Some(output) = response.output() { + if let aws_sdk_bedrockruntime::types::ConverseOutput::Message(message) = output { + for block in message.content() { + match block { + aws_sdk_bedrockruntime::types::ContentBlock::Text(text) => { + text_content.push_str(text); } - })); + aws_sdk_bedrockruntime::types::ContentBlock::ToolUse(tool_use) => { + // Convert Document back to JSON string + let input_json = document_to_json(tool_use.input()); + tool_calls.push(OpenAIToolCall { + id: tool_use.tool_use_id().to_string(), + function: windmill_common::ai_bedrock::OpenAIFunction { + name: tool_use.name().to_string(), + arguments: serde_json::to_string(&input_json).unwrap_or_default(), + }, + r#type: "function".to_string(), + extra_content: None, + }); + } + _ => {} + } } } } @@ -353,7 +686,7 @@ pub async fn transform_bedrock_to_openai( let message = if !tool_calls.is_empty() { serde_json::json!({ "role": "assistant", - "content": if text_content.is_empty() { Value::Null } else { Value::String(text_content) }, + "content": if text_content.is_empty() { serde_json::Value::Null } else { serde_json::Value::String(text_content) }, "tool_calls": tool_calls }) } else { @@ -364,11 +697,11 @@ pub async fn transform_bedrock_to_openai( }; // Extract usage information - let usage = if let Some(usage_data) = bedrock_resp.get("usage") { + let usage = if let Some(usage_data) = response.usage() { serde_json::json!({ - "prompt_tokens": usage_data["inputTokens"].as_i64().unwrap_or(0), - "completion_tokens": usage_data["outputTokens"].as_i64().unwrap_or(0), - "total_tokens": usage_data["totalTokens"].as_i64().unwrap_or(0) + "prompt_tokens": usage_data.input_tokens(), + "completion_tokens": usage_data.output_tokens(), + "total_tokens": usage_data.total_tokens() }) } else { serde_json::json!({ @@ -393,274 +726,38 @@ pub async fn transform_bedrock_to_openai( }); let response_body = serde_json::to_vec(&openai_resp) - .map_err(|e| Error::internal_err(format!("Failed to serialize OpenAI response: {}", e)))? - .into(); + .map_err(|e| Error::internal_err(format!("Failed to serialize OpenAI response: {}", e)))?; - Ok(response_body) + let mut response_headers = http::HeaderMap::new(); + response_headers.insert("content-type", "application/json".parse().unwrap()); + + Ok(( + http::StatusCode::OK, + response_headers, + axum::body::Body::from(response_body), + )) } -/// Transform AWS Bedrock streaming response to OpenAI SSE format -/// Bedrock uses AWS event stream binary format, not SSE -pub fn transform_bedrock_stream_to_openai( - stream: impl futures::Stream> - + Send - + 'static, - model: String, -) -> impl futures::Stream> + Send { - use futures::stream::StreamExt; - use serde_json::Value; - use std::collections::HashMap; - - let id = format!("chatcmpl-{}", uuid::Uuid::new_v4().simple()); - let created = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_secs(); - - // State to track partial tool calls and binary buffer - struct StreamState { - id: String, - model: String, - created: u64, - tool_calls: HashMap, // index -> (id, name, args) - buffer: Vec, // Binary buffer for AWS event stream - } - - let state = std::sync::Arc::new(tokio::sync::Mutex::new(StreamState { - id: id.clone(), - model: model.clone(), - created, - tool_calls: HashMap::new(), - buffer: Vec::new(), - })); - - stream - .then(move |chunk_result| { - let state = state.clone(); - async move { - match chunk_result { - Ok(chunk) => { - let mut state = state.lock().await; - state.buffer.extend_from_slice(&chunk); - - let mut events = Vec::new(); - - // Parse AWS event stream messages from buffer - loop { - // Need at least 12 bytes for prelude (8) + prelude CRC (4) - if state.buffer.len() < 12 { - break; - } - - // Read prelude: total_length (4 bytes) + headers_length (4 bytes) - let total_length = u32::from_be_bytes([ - state.buffer[0], - state.buffer[1], - state.buffer[2], - state.buffer[3], - ]) as usize; - - // Check if we have the complete message - if state.buffer.len() < total_length { - break; - } - - let headers_length = u32::from_be_bytes([ - state.buffer[4], - state.buffer[5], - state.buffer[6], - state.buffer[7], - ]) as usize; - - // Skip prelude CRC (4 bytes after prelude) - let headers_start = 12; - let payload_start = headers_start + headers_length; - let payload_end = total_length - 4; // Exclude message CRC - - // Parse headers to extract event type - let mut event_type = None; - let mut pos = headers_start; - while pos < payload_start { - if pos + 1 > state.buffer.len() { - break; - } - let name_len = state.buffer[pos] as usize; - pos += 1; - - if pos + name_len > state.buffer.len() { - break; - } - let name = String::from_utf8_lossy(&state.buffer[pos..pos + name_len]).to_string(); - pos += name_len; - - if pos + 3 > state.buffer.len() { - break; - } - let value_type = state.buffer[pos]; - pos += 1; - let value_len = u16::from_be_bytes([state.buffer[pos], state.buffer[pos + 1]]) as usize; - pos += 2; - - if pos + value_len > state.buffer.len() { - break; - } - - if value_type == 7 && name == ":event-type" { - event_type = Some(String::from_utf8_lossy(&state.buffer[pos..pos + value_len]).to_string()); - } - pos += value_len; - } - - // Extract JSON payload (copy to avoid borrow issues) - let payload = state.buffer[payload_start..payload_end].to_vec(); - - // Remove processed message from buffer - state.buffer.drain(0..total_length); - - // Process the event - if let Some(evt_type) = event_type { - if let Ok(payload_str) = std::str::from_utf8(&payload) { - if let Ok(parsed_data) = serde_json::from_str::(payload_str) { - // Transform based on event type - match evt_type.as_str() { - "messageStart" => { - // No output for messageStart - } - "contentBlockStart" => { - let index = parsed_data["contentBlockIndex"].as_u64().unwrap_or(0) as usize; - - if let Some(tool_use) = parsed_data["start"].get("toolUse") { - let tool_id = tool_use["toolUseId"].as_str().unwrap_or("").to_string(); - let name = tool_use["name"].as_str().unwrap_or("").to_string(); - - state.tool_calls.insert(index, (tool_id.clone(), name.clone(), String::new())); - - // Send initial tool call chunk - let chunk = serde_json::json!({ - "id": state.id, - "object": "chat.completion.chunk", - "created": state.created, - "model": state.model, - "choices": [{ - "index": 0, - "delta": { - "tool_calls": [{ - "index": index, - "id": tool_id, - "type": "function", - "function": { - "name": name, - "arguments": "" - } - }] - }, - "finish_reason": Value::Null - }] - }); - - events.push(Ok(bytes::Bytes::from(format!("data: {}\n\n", chunk)))); - } - } - "contentBlockDelta" => { - let index = parsed_data["contentBlockIndex"].as_u64().unwrap_or(0) as usize; - - if let Some(text) = parsed_data["delta"]["text"].as_str() { - // Text content delta - let chunk = serde_json::json!({ - "id": state.id, - "object": "chat.completion.chunk", - "created": state.created, - "model": state.model, - "choices": [{ - "index": 0, - "delta": { - "content": text - }, - "finish_reason": Value::Null - }] - }); - - events.push(Ok(bytes::Bytes::from(format!("data: {}\n\n", chunk)))); - } else if let Some(tool_use_input) = parsed_data["delta"]["toolUse"]["input"].as_str() { - // Tool use arguments delta - if let Some((_tool_id, _name, ref mut args)) = state.tool_calls.get_mut(&index) { - args.push_str(tool_use_input); - - let chunk = serde_json::json!({ - "id": state.id, - "object": "chat.completion.chunk", - "created": state.created, - "model": state.model, - "choices": [{ - "index": 0, - "delta": { - "tool_calls": [{ - "index": index, - "function": { - "arguments": tool_use_input - } - }] - }, - "finish_reason": Value::Null - }] - }); - - events.push(Ok(bytes::Bytes::from(format!("data: {}\n\n", chunk)))); - } - } - } - "contentBlockStop" => { - // No output needed - } - "messageStop" => { - let stop_reason = parsed_data["stopReason"].as_str().unwrap_or("end_turn"); - let finish_reason = match stop_reason { - "end_turn" => "stop", - "max_tokens" => "length", - "tool_use" => "tool_calls", - "stop_sequence" => "stop", - "guardrail_intervened" | "content_filtered" => "content_filter", - _ => "stop", - }; - - let chunk = serde_json::json!({ - "id": state.id, - "object": "chat.completion.chunk", - "created": state.created, - "model": state.model, - "choices": [{ - "index": 0, - "delta": {}, - "finish_reason": finish_reason - }] - }); - - events.push(Ok(bytes::Bytes::from(format!("data: {}\n\n", chunk)))); - } - "metadata" => { - // Could include usage info here if needed - } - _ => {} - } - } - } - } - } // end loop - - events - } - Err(e) => { - vec![Err(std::io::Error::new( - std::io::ErrorKind::Other, - e.to_string(), - ))] - } - } +/// Convert AWS Smithy Document to serde_json::Value +fn document_to_json(doc: &aws_smithy_types::Document) -> serde_json::Value { + match doc { + aws_smithy_types::Document::Object(map) => { + let mut json_map = serde_json::Map::new(); + for (k, v) in map { + json_map.insert(k.clone(), document_to_json(v)); } - }) - .flat_map(|events| futures::stream::iter(events)) - .chain(futures::stream::iter(vec![ - // Send [DONE] at the end - Ok(bytes::Bytes::from("data: [DONE]\n\n")) - ])) + serde_json::Value::Object(json_map) + } + aws_smithy_types::Document::Array(arr) => { + serde_json::Value::Array(arr.iter().map(document_to_json).collect()) + } + aws_smithy_types::Document::Number(num) => match num { + aws_smithy_types::Number::PosInt(n) => serde_json::Value::Number((*n).into()), + aws_smithy_types::Number::NegInt(n) => serde_json::Value::Number((*n).into()), + aws_smithy_types::Number::Float(f) => serde_json::json!(*f), + }, + aws_smithy_types::Document::String(s) => serde_json::Value::String(s.clone()), + aws_smithy_types::Document::Bool(b) => serde_json::Value::Bool(*b), + aws_smithy_types::Document::Null => serde_json::Value::Null, + } } diff --git a/backend/windmill-api/src/lib.rs b/backend/windmill-api/src/lib.rs index 7c625ac634..db561b7ea6 100644 --- a/backend/windmill-api/src/lib.rs +++ b/backend/windmill-api/src/lib.rs @@ -81,6 +81,7 @@ pub mod auth; #[cfg(all(feature = "private", feature = "parquet"))] pub mod azure_proxy_ee; mod azure_proxy_oss; +#[cfg(feature = "bedrock")] mod bedrock; mod capture; mod concurrency_groups; diff --git a/backend/windmill-common/Cargo.toml b/backend/windmill-common/Cargo.toml index 8f15435285..bdd84663f6 100644 --- a/backend/windmill-common/Cargo.toml +++ b/backend/windmill-common/Cargo.toml @@ -20,6 +20,8 @@ smtp = ["dep:mail-send"] scoped_cache = [] cloud = [] openidconnect = ["dep:openidconnect"] +bedrock = ["dep:aws-sdk-bedrockruntime"] + [lib] name = "windmill_common" path = "src/lib.rs" @@ -65,6 +67,7 @@ aws-config.workspace = true aws-sdk-sts = { workspace = true, optional = true } aws-credential-types.workspace = true aws-smithy-types.workspace = true +aws-sdk-bedrockruntime = { workspace = true, optional = true } base64.workspace = true bitflags.workspace = true once_cell.workspace = true diff --git a/backend/windmill-common/src/ai_bedrock.rs b/backend/windmill-common/src/ai_bedrock.rs new file mode 100644 index 0000000000..b7fcdc3aa7 --- /dev/null +++ b/backend/windmill-common/src/ai_bedrock.rs @@ -0,0 +1,693 @@ +//! Shared AWS Bedrock SDK code for AI chat proxy and worker. +//! +//! This module provides: +//! - BedrockClient: SDK wrapper with bearer token and IAM credentials auth +//! - Message/tool conversion: OpenAI format <-> Bedrock Converse API format +//! - Stream event parsing: Extract text/tool deltas from Bedrock stream events +//! +//! Used by both windmill-api (chat proxy) and windmill-worker (AI agent). + +use aws_config::BehaviorVersion; +use aws_credential_types::provider::token::ProvideToken; +use aws_credential_types::provider::ProvideCredentials; +use aws_sdk_bedrockruntime::types::{ + ContentBlock, ConversationRole, ConverseStreamOutput, ImageBlock, ImageFormat, ImageSource, + InferenceConfiguration, Message, SystemContentBlock, Tool, ToolInputSchema, ToolSpecification, +}; +use aws_sdk_bedrockruntime::Client as BedrockRuntimeClient; +use serde::{Deserialize, Serialize}; + +use crate::error::Error; + +// Re-export types from ai_types +pub use crate::ai_types::{ + ContentPart, ExtraContent, ImageUrlData, OpenAIContent, OpenAIFunction, OpenAIMessage, + OpenAIToolCall, ToolDef, ToolDefFunction, UrlCitation, +}; + +// ============================================================================ +// Cached AWS SDK Config +// ============================================================================ + +/// Cached AWS SDK config loaded from environment +/// Avoids repeated I/O for environment variable lookups and file reads +static AWS_SDK_CONFIG: tokio::sync::OnceCell = + tokio::sync::OnceCell::const_new(); + +/// Get or initialize the cached AWS SDK config +async fn get_aws_sdk_config() -> &'static aws_config::SdkConfig { + AWS_SDK_CONFIG + .get_or_init(|| async { aws_config::load_defaults(BehaviorVersion::latest()).await }) + .await +} + +// ============================================================================ +// Bedrock Client +// ============================================================================ + +/// Result of checking AWS Bedrock credentials availability +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct BedrockCredentialsCheck { + pub available: bool, + pub access_key_id_prefix: Option, + pub region: Option, + pub error: Option, +} + +/// Check if AWS credentials are available from the environment +pub async fn check_env_credentials() -> BedrockCredentialsCheck { + let config = get_aws_sdk_config().await; + + if let Some(creds_provider) = config.credentials_provider() { + match creds_provider.provide_credentials().await { + Ok(creds) => { + let access_key_id = creds.access_key_id(); + let prefix = if access_key_id.len() >= 8 { + format!("{}...", &access_key_id[..8]) + } else { + access_key_id.to_string() + }; + + BedrockCredentialsCheck { + available: true, + access_key_id_prefix: Some(prefix), + region: config.region().map(|r| r.to_string()), + error: None, + } + } + Err(e) => BedrockCredentialsCheck { + available: false, + access_key_id_prefix: None, + region: None, + error: Some(format!("Failed to retrieve credentials: {}", e)), + }, + } + } else { + BedrockCredentialsCheck { + available: false, + access_key_id_prefix: None, + region: None, + error: Some("No credentials provider configured".to_string()), + } + } +} + +/// Constants for commonly used strings to avoid allocations +pub const FUNCTION_TYPE: &str = "function"; + +#[derive(Debug, Clone)] +pub struct BearerTokenProvider { + token: String, +} + +impl BearerTokenProvider { + pub fn new(token: String) -> Self { + Self { token } + } +} + +impl ProvideToken for BearerTokenProvider { + fn provide_token<'a>(&'a self) -> aws_credential_types::provider::future::ProvideToken<'a> + where + Self: 'a, + { + aws_credential_types::provider::future::ProvideToken::ready(Ok( + aws_credential_types::Token::new(self.token.clone(), None), + )) + } +} + +pub struct BedrockClient { + client: BedrockRuntimeClient, +} + +impl BedrockClient { + pub async fn from_bearer_token(bearer_token: String, region: &str) -> Result { + let config = aws_sdk_bedrockruntime::config::Builder::new() + .region(aws_config::Region::new(region.to_string())) + .behavior_version(BehaviorVersion::latest()) + .token_provider(BearerTokenProvider::new(bearer_token)) + .build(); + + Ok(Self { client: BedrockRuntimeClient::from_conf(config) }) + } + + pub async fn from_credentials( + access_key_id: String, + secret_access_key: String, + session_token: Option, + region: &str, + ) -> Result { + let credentials = aws_credential_types::Credentials::new( + access_key_id, + secret_access_key, + session_token, + None, // expiration + "windmill", + ); + + let config = aws_sdk_bedrockruntime::config::Builder::new() + .region(aws_config::Region::new(region.to_string())) + .behavior_version(BehaviorVersion::latest()) + .credentials_provider(credentials) + .build(); + + Ok(Self { client: BedrockRuntimeClient::from_conf(config) }) + } + + pub async fn from_env(region: &str) -> Result { + let config = get_aws_sdk_config().await; + + // Verify that credentials are actually available + if let Some(creds_provider) = config.credentials_provider() { + match creds_provider.provide_credentials().await { + Ok(creds) => { + tracing::debug!( + "Bedrock: using env credentials, access_key={}...", + &creds.access_key_id().get(..8).unwrap_or("N/A"), + ); + } + Err(e) => { + return Err(Error::internal_err(format!( + "AWS credentials not available from environment: {}", + e + ))); + } + } + } else { + return Err(Error::internal_err( + "No AWS credentials provider configured in environment".to_string(), + )); + } + + // Build client, only override region if explicitly provided + let mut builder = aws_sdk_bedrockruntime::config::Builder::from(config); + if !region.is_empty() { + builder = builder.region(aws_config::Region::new(region.to_string())); + } + let bedrock_config = builder.build(); + + let client = aws_sdk_bedrockruntime::Client::from_conf(bedrock_config); + Ok(Self { client }) + } + + pub fn client(&self) -> &BedrockRuntimeClient { + &self.client + } +} + +// ============================================================================ +// Error Formatting +// ============================================================================ + +/// Format AWS SDK errors with detailed information +pub fn format_bedrock_error(error: &aws_sdk_bedrockruntime::error::SdkError) -> String +where + E: std::fmt::Debug + std::fmt::Display, + R: std::fmt::Debug, +{ + use aws_sdk_bedrockruntime::error::SdkError; + + match error { + SdkError::ServiceError(err) => { + format!("Service error: {} (details: {:?})", err.err(), err) + } + SdkError::ConstructionFailure(err) => { + format!("Request construction failed: {:?}", err) + } + SdkError::DispatchFailure(err) => { + format!("Request dispatch failed: {:?}", err) + } + SdkError::ResponseError(err) => { + format!("Response error: {:?}", err) + } + SdkError::TimeoutError(err) => { + format!("Request timeout: {:?}", err) + } + _ => format!("{:?}", error), + } +} + +// ============================================================================ +// Type Conversion Utilities +// ============================================================================ + +/// Convert serde_json::Value to AWS Smithy Document +pub fn json_to_document(value: serde_json::Value) -> aws_smithy_types::Document { + use aws_smithy_types::Document; + use serde_json::Value; + + match value { + Value::Object(map) => { + let mut doc_map = std::collections::HashMap::new(); + for (k, v) in map { + doc_map.insert(k, json_to_document(v)); + } + Document::Object(doc_map) + } + Value::Array(arr) => Document::Array(arr.into_iter().map(json_to_document).collect()), + Value::Number(num) => { + if let Some(i) = num.as_i64() { + Document::Number(aws_smithy_types::Number::PosInt(i as u64)) + } else if let Some(f) = num.as_f64() { + Document::Number(aws_smithy_types::Number::Float(f)) + } else { + Document::Number(aws_smithy_types::Number::PosInt(0)) + } + } + Value::String(s) => Document::String(s), + Value::Bool(b) => Document::Bool(b), + Value::Null => Document::Null, + } +} + +// ============================================================================ +// Message Conversion (OpenAI -> Bedrock) +// ============================================================================ + +/// Convert OpenAI-style messages to Bedrock format +/// +/// Separates system messages from conversation messages as required by Bedrock API. +/// +/// Important: Bedrock requires messages to alternate between user and assistant roles. +/// When an assistant message has tool_use blocks, the next user message must contain +/// ALL corresponding tool_result blocks. This function groups consecutive tool messages +/// into a single user message. +/// +/// # Returns +/// Tuple of (conversation_messages, system_prompts) +pub fn openai_messages_to_bedrock( + messages: &[OpenAIMessage], +) -> Result<(Vec, Vec), Error> { + let mut bedrock_messages = Vec::new(); + let mut system_prompts = Vec::new(); + let mut pending_tool_results: Vec = Vec::new(); + + for msg in messages { + match msg.role.as_str() { + "system" => { + // Extract system messages separately + if let Some(content) = &msg.content { + let text = content_to_text(content); + if !text.is_empty() { + system_prompts.push(SystemContentBlock::Text(text)); + } + } + } + "user" | "assistant" => { + // Before adding a user/assistant message, flush any pending tool results + if !pending_tool_results.is_empty() { + let tool_result_message = Message::builder() + .role(ConversationRole::User) + .set_content(Some(pending_tool_results.drain(..).collect())) + .build() + .map_err(|e| { + Error::internal_err(format!( + "Failed to build tool results message: {}", + e + )) + })?; + bedrock_messages.push(tool_result_message); + } + bedrock_messages.push(convert_message(msg)?); + } + "tool" => { + // Accumulate tool results - they will be flushed as a single message + // when we encounter a non-tool message or at the end + let tool_result = convert_tool_result_content(msg)?; + pending_tool_results.push(tool_result); + } + _ => { + return Err(Error::BadRequest(format!("Unsupported role: {}", msg.role))); + } + } + } + + // Flush any remaining tool results at the end + if !pending_tool_results.is_empty() { + let tool_result_message = Message::builder() + .role(ConversationRole::User) + .set_content(Some(pending_tool_results)) + .build() + .map_err(|e| { + Error::internal_err(format!("Failed to build tool results message: {}", e)) + })?; + bedrock_messages.push(tool_result_message); + } + + Ok((bedrock_messages, system_prompts)) +} + +/// Helper to extract text from OpenAIContent (ignoring images) +/// +/// This is public so it can be reused by the worker module. +pub fn content_to_text(content: &OpenAIContent) -> String { + match content { + OpenAIContent::Text(text) => text.to_string(), + OpenAIContent::Parts(parts) => { + // Extract only text parts and join them + let text_parts: Vec<&str> = parts + .iter() + .filter_map(|part| match part { + ContentPart::Text { text } => Some(text.as_str()), + _ => None, + }) + .collect(); + text_parts.join(" ") + } + } +} + +/// Parse image data URL and extract format and base64 data +fn parse_image_data_url(url: &str) -> Result<(ImageFormat, Vec), Error> { + if !url.starts_with("data:") { + return Err(Error::internal_err("Image URL must be a data URL")); + } + + // Parse data:image/png;base64, + let base64_start = url + .find("base64,") + .ok_or_else(|| Error::internal_err("Invalid data URL format"))?; + + let base64_data = &url[base64_start + 7..]; + let mime_type = url + .split(';') + .next() + .and_then(|s| s.strip_prefix("data:")) + .unwrap_or("image/png"); + + // Extract format from MIME type (e.g., "image/png" -> "png") + let format_str = mime_type + .rsplit_once('/') + .map(|(_, format)| format) + .unwrap_or("png"); + + // Map to ImageFormat enum + let format = match format_str { + "png" => ImageFormat::Png, + "jpeg" | "jpg" => ImageFormat::Jpeg, + "gif" => ImageFormat::Gif, + "webp" => ImageFormat::Webp, + _ => ImageFormat::Png, // Default to PNG + }; + + // Decode base64 + let bytes = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, base64_data) + .map_err(|e| Error::internal_err(format!("Failed to decode base64 image: {}", e)))?; + + Ok((format, bytes)) +} + +/// Convert a ContentPart to Bedrock ContentBlock +fn content_part_to_block(part: &ContentPart) -> Result, Error> { + match part { + ContentPart::Text { text } => { + if text.is_empty() { + Ok(None) + } else { + Ok(Some(ContentBlock::Text(text.clone()))) + } + } + ContentPart::ImageUrl { image_url } => { + let (format, bytes) = parse_image_data_url(&image_url.url)?; + + let image_source = ImageSource::Bytes(bytes.into()); + let image_block = ImageBlock::builder() + .format(format) + .source(image_source) + .build() + .map_err(|e| Error::internal_err(format!("Failed to build image block: {}", e)))?; + + Ok(Some(ContentBlock::Image(image_block))) + } + ContentPart::S3Object { .. } => { + // S3Objects should be converted to ImageUrl before calling this function + Ok(None) + } + } +} + +/// Convert a single OpenAI message to Bedrock Message +fn convert_message(msg: &OpenAIMessage) -> Result { + let role = match msg.role.as_str() { + "user" => ConversationRole::User, + "assistant" => ConversationRole::Assistant, + _ => { + return Err(Error::internal_err(format!("Unsupported role: {}", msg.role))); + } + }; + + let mut content_blocks = Vec::new(); + + // Handle content (text and/or images) + if let Some(content) = &msg.content { + match content { + OpenAIContent::Text(text) => { + if !text.is_empty() { + content_blocks.push(ContentBlock::Text(text.clone())); + } + } + OpenAIContent::Parts(parts) => { + for part in parts { + if let Some(block) = content_part_to_block(part)? { + content_blocks.push(block); + } + } + } + } + } + + // Handle tool calls (for assistant messages) + if let Some(tool_calls) = &msg.tool_calls { + for tc in tool_calls { + content_blocks.push(convert_tool_call_to_content(tc)?); + } + } + + // Bedrock requires at least one content block + if content_blocks.is_empty() { + content_blocks.push(ContentBlock::Text(String::new())); + } + + Message::builder() + .role(role) + .set_content(Some(content_blocks)) + .build() + .map_err(|e| Error::internal_err(format!("Failed to build message: {}", e))) +} + +/// Convert OpenAI tool call to Bedrock ToolUse content block +fn convert_tool_call_to_content(tool_call: &OpenAIToolCall) -> Result { + let input = json_to_document( + serde_json::from_str(&tool_call.function.arguments) + .unwrap_or_else(|_| serde_json::json!({})), + ); + Ok(ContentBlock::ToolUse( + aws_sdk_bedrockruntime::types::ToolUseBlock::builder() + .tool_use_id(&tool_call.id) + .name(&tool_call.function.name) + .input(input) + .build() + .map_err(|e| Error::internal_err(format!("Failed to build tool use: {}", e)))?, + )) +} + +/// Convert tool result message to Bedrock ToolResult ContentBlock +/// +/// Returns just the ContentBlock (not a full Message) so multiple tool results +/// can be combined into a single user message. +fn convert_tool_result_content(msg: &OpenAIMessage) -> Result { + let tool_call_id = msg + .tool_call_id + .as_ref() + .ok_or_else(|| Error::internal_err("Tool message missing tool_call_id"))?; + + let content_str = msg + .content + .as_ref() + .map(|c| content_to_text(c)) + .unwrap_or_default(); + + // Try to parse as JSON, otherwise use text + let tool_result_content = + if let Ok(json_val) = serde_json::from_str::(&content_str) { + if json_val.is_object() { + vec![aws_sdk_bedrockruntime::types::ToolResultContentBlock::Json( + json_to_document(json_val), + )] + } else { + // Wrap primitives and arrays in an object + vec![aws_sdk_bedrockruntime::types::ToolResultContentBlock::Json( + json_to_document(serde_json::json!({"result": json_val})), + )] + } + } else { + vec![aws_sdk_bedrockruntime::types::ToolResultContentBlock::Text( + content_str.to_string(), + )] + }; + + Ok(ContentBlock::ToolResult( + aws_sdk_bedrockruntime::types::ToolResultBlock::builder() + .tool_use_id(tool_call_id) + .set_content(Some(tool_result_content)) + .build() + .map_err(|e| Error::internal_err(format!("Failed to build tool result: {}", e)))?, + )) +} + +// ============================================================================ +// Tool Conversion (OpenAI -> Bedrock) +// ============================================================================ + +/// Convert OpenAI tool definitions to Bedrock format +pub fn openai_tools_to_bedrock(tools: &[ToolDef]) -> Result, Error> { + tools + .iter() + .map(|tool_def| { + let spec = &tool_def.function; + + // Convert parameters (RawValue) to Document via serde_json::Value + let param_value: serde_json::Value = serde_json::from_str(spec.parameters.get()) + .map_err(|e| Error::internal_err(format!("Invalid tool schema: {}", e)))?; + let input_schema = ToolInputSchema::Json(json_to_document(param_value)); + + let tool_spec = ToolSpecification::builder() + .name(&spec.name) + .set_description(spec.description.clone()) + .input_schema(input_schema) + .build() + .map_err(|e| Error::internal_err(format!("Failed to build tool spec: {}", e)))?; + + Ok(Tool::ToolSpec(tool_spec)) + }) + .collect() +} + +// ============================================================================ +// Inference Configuration +// ============================================================================ + +/// Create inference configuration from parameters +pub fn create_inference_config( + temperature: Option, + max_tokens: Option, +) -> Option { + if temperature.is_none() && max_tokens.is_none() { + return None; + } + + let mut builder = InferenceConfiguration::builder(); + + if let Some(temp) = temperature { + builder = builder.temperature(temp); + } + + if let Some(max_tok) = max_tokens { + builder = builder.max_tokens(max_tok); + } + + Some(builder.build()) +} + +// ============================================================================ +// Stream Event Parsing +// ============================================================================ + +/// Extract text delta from Bedrock stream event +pub fn bedrock_stream_event_to_text(event: &ConverseStreamOutput) -> Option { + match event { + ConverseStreamOutput::ContentBlockDelta(delta) => delta + .delta() + .and_then(|d| d.as_text().ok()) + .map(|s| s.to_string()), + _ => None, + } +} + +/// Represents a streaming tool call being accumulated +#[derive(Debug, Clone)] +pub struct StreamingToolCall { + pub id: String, + pub name: String, + pub arguments: String, +} + +/// Extract tool use start event from stream +pub fn bedrock_stream_event_to_tool_start( + event: &ConverseStreamOutput, +) -> Option { + match event { + ConverseStreamOutput::ContentBlockStart(start) => { + if let Some(tool_use) = start.start().and_then(|s| s.as_tool_use().ok()) { + Some(StreamingToolCall { + id: tool_use.tool_use_id().to_string(), + name: tool_use.name().to_string(), + arguments: String::new(), + }) + } else { + None + } + } + _ => None, + } +} + +/// Extract tool use input delta from stream +pub fn bedrock_stream_event_to_tool_delta(event: &ConverseStreamOutput) -> Option { + match event { + ConverseStreamOutput::ContentBlockDelta(delta) => delta + .delta() + .and_then(|d| d.as_tool_use().ok()) + .map(|tool_use| tool_use.input().to_string()), + _ => None, + } +} + +/// Check if stream event indicates content block stop +pub fn bedrock_stream_event_is_block_stop(event: &ConverseStreamOutput) -> bool { + matches!(event, ConverseStreamOutput::ContentBlockStop(_)) +} + +/// Convert accumulated streaming tool calls to OpenAI format +pub fn streaming_tool_calls_to_openai(tool_calls: Vec) -> Vec { + tool_calls + .into_iter() + .map(|tc| OpenAIToolCall { + id: tc.id, + function: OpenAIFunction { name: tc.name, arguments: tc.arguments }, + r#type: FUNCTION_TYPE.to_string(), + extra_content: None, // Bedrock doesn't use thought signatures + }) + .collect() +} + +// ============================================================================ +// Tool Configuration Builder +// ============================================================================ + +/// Build tool configuration with optional ToolChoice for structured output +pub fn build_tool_config( + tools: Option<&[ToolDef]>, + force_tool_use: bool, +) -> Result, Error> { + if let Some(tools) = tools { + let bedrock_tools = openai_tools_to_bedrock(tools)?; + let mut tool_config_builder = aws_sdk_bedrockruntime::types::ToolConfiguration::builder() + .set_tools(Some(bedrock_tools)); + + // For structured output, force the model to use the tool + if force_tool_use { + tool_config_builder = + tool_config_builder.tool_choice(aws_sdk_bedrockruntime::types::ToolChoice::Any( + aws_sdk_bedrockruntime::types::AnyToolChoice::builder().build(), + )); + } + + Ok(Some(tool_config_builder.build().map_err(|e| { + Error::internal_err(format!("Failed to build tool configuration: {}", e)) + })?)) + } else { + Ok(None) + } +} diff --git a/backend/windmill-common/src/ai_providers.rs b/backend/windmill-common/src/ai_providers.rs index 34153195ac..48b2ef8281 100644 --- a/backend/windmill-common/src/ai_providers.rs +++ b/backend/windmill-common/src/ai_providers.rs @@ -74,14 +74,32 @@ impl AIProvider { AIProvider::TogetherAI => Ok("https://api.together.xyz/v1".to_string()), AIProvider::Anthropic => Ok("https://api.anthropic.com/v1".to_string()), AIProvider::Mistral => Ok("https://api.mistral.ai/v1".to_string()), - AIProvider::AWSBedrock => Ok(format!( - "https://bedrock-runtime.{}.amazonaws.com", - region.unwrap_or_else(|| "us-east-1".to_string()) - )), - AIProvider::CustomAI | AIProvider::AzureOpenAI => Err(Error::BadRequest(format!( - "{:?} provider requires a base URL in the resource", - self - ))), + p @ (AIProvider::CustomAI | AIProvider::AzureOpenAI) => { + if let Some(base_url) = resource_base_url { + Ok(base_url) + } else { + Err(Error::BadRequest(format!( + "{:?} provider requires a base URL in the resource", + p + ))) + } + } + AIProvider::AWSBedrock => { + #[cfg(feature = "bedrock")] + { + Ok(format!( + "https://bedrock-runtime.{}.amazonaws.com", + region.unwrap_or_else(|| "us-east-1".to_string()) + )) + } + #[cfg(not(feature = "bedrock"))] + { + let _ = region; + Err(Error::BadRequest( + "AWS Bedrock support is not enabled. Build with 'bedrock' feature.".to_string() + )) + } + } } } diff --git a/backend/windmill-common/src/ai_types.rs b/backend/windmill-common/src/ai_types.rs new file mode 100644 index 0000000000..2d72190f5b --- /dev/null +++ b/backend/windmill-common/src/ai_types.rs @@ -0,0 +1,116 @@ +//! Shared AI types used across all providers. +//! +//! This module contains common types for OpenAI-compatible message formats +//! that are used by all AI providers, not just Bedrock. + +use serde::{Deserialize, Serialize}; +use serde_json::value::RawValue; + +use crate::flow_status::AgentAction; +use crate::s3_helpers::S3Object; + +// ============================================================================ +// Shared Types for OpenAI-compatible message format +// ============================================================================ + +/// URL citation annotation for web search results +#[derive(Deserialize, Serialize, Clone, Debug)] +pub struct UrlCitation { + pub start_index: usize, + pub end_index: usize, + pub url: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub title: Option, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ContentPart { + Text { + text: String, + }, + #[serde(rename = "image_url")] + ImageUrl { + image_url: ImageUrlData, + }, + #[serde(rename = "s3_object")] + S3Object { + s3_object: S3Object, + }, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct ImageUrlData { + pub url: String, // data:image/png;base64,... or https://... +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(untagged)] +pub enum OpenAIContent { + Text(String), + Parts(Vec), +} + +#[derive(Serialize, Clone, Debug)] +pub struct ToolDefFunction { + pub name: String, + pub description: Option, + pub parameters: Box, +} + +#[derive(Serialize, Clone, Debug)] +pub struct ToolDef { + pub r#type: String, + pub function: ToolDefFunction, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub struct OpenAIFunction { + pub name: String, + pub arguments: String, +} + +/// Google-specific extra content for thought signatures (Gemini 3 Pro / 2.5) +#[derive(Deserialize, Serialize, Clone, Debug, Default)] +pub struct GoogleExtraContent { + #[serde(skip_serializing_if = "Option::is_none")] + pub thought_signature: Option, +} + +/// Extra content for provider-specific metadata (e.g., Google thought signatures) +#[derive(Deserialize, Serialize, Clone, Debug, Default)] +pub struct ExtraContent { + #[serde(skip_serializing_if = "Option::is_none")] + pub google: Option, +} + +#[derive(Deserialize, Serialize, Clone, Debug)] +pub struct OpenAIToolCall { + pub id: String, + pub function: OpenAIFunction, + pub r#type: String, + /// Extra content for provider-specific metadata (e.g., Google Gemini thought signatures) + #[serde(skip_serializing_if = "Option::is_none")] + pub extra_content: Option, +} + +/// OpenAI-compatible message format used across all AI providers. +/// +/// The `agent_action` field is used by the worker for flow-specific tracking +/// and is never serialized to JSON (skip_serializing, default). +#[derive(Deserialize, Serialize, Clone, Default, Debug)] +pub struct OpenAIMessage { + pub role: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub content: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_calls: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_call_id: Option, + /// Worker-specific field for tracking agent actions in flows. + /// Never serialized; defaults to None when deserializing. + #[serde(skip_serializing, default)] + pub agent_action: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub annotations: Option>, +} diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index 90896608ef..e107fa564d 100644 --- a/backend/windmill-common/src/lib.rs +++ b/backend/windmill-common/src/lib.rs @@ -28,6 +28,9 @@ use sqlx::{Acquire, Postgres}; pub mod agent_workers; pub mod ai_providers; +pub mod ai_types; +#[cfg(feature = "bedrock")] +pub mod ai_bedrock; pub mod apps; pub mod audit; pub mod assets; diff --git a/backend/windmill-worker/Cargo.toml b/backend/windmill-worker/Cargo.toml index 1dbb702381..5790d16b05 100644 --- a/backend/windmill-worker/Cargo.toml +++ b/backend/windmill-worker/Cargo.toml @@ -36,6 +36,7 @@ nu = ["dep:windmill-parser-nu"] java = ["dep:windmill-parser-java"] ruby = ["dep:windmill-parser-ruby"] duckdb = ["dep:libloading"] +bedrock = ["dep:aws-sdk-bedrockruntime", "windmill-common/bedrock"] [dependencies] windmill-queue.workspace = true @@ -59,7 +60,7 @@ windmill-parser-sql.workspace = true windmill-parser-graphql.workspace = true windmill-parser-php = { workspace = true, optional = true } windmill-git-sync.workspace = true -aws-sdk-bedrockruntime.workspace = true +aws-sdk-bedrockruntime = { workspace = true, optional = true } aws-config.workspace = true aws-credential-types.workspace = true aws-smithy-types.workspace = true diff --git a/backend/windmill-worker/src/ai/providers/bedrock.rs b/backend/windmill-worker/src/ai/providers/bedrock.rs index 95068f5c01..f4df4bc393 100644 --- a/backend/windmill-worker/src/ai/providers/bedrock.rs +++ b/backend/windmill-worker/src/ai/providers/bedrock.rs @@ -1,500 +1,32 @@ +//! AWS Bedrock provider for the AI agent. +//! +//! Uses shared SDK code from windmill_common::ai_bedrock for: +//! - BedrockClient (SDK wrapper with auth) +//! - Message conversion (OpenAI format -> Bedrock format) +//! - Stream event parsing +//! - Helper utilities + use crate::ai::{ image_handler::prepare_messages_for_api, - providers::openai::{OpenAIFunction, OpenAIToolCall}, query_builder::{ParsedResponse, StreamEventProcessor}, types::StreamingEvent, - types::{ContentPart, OpenAIContent, OpenAIMessage, ToolDef}, + types::{OpenAIMessage, ToolDef}, }; -use aws_config::BehaviorVersion; -use aws_credential_types::provider::token::ProvideToken; -use aws_sdk_bedrockruntime::types::{ - ContentBlock, ConversationRole, ConverseStreamOutput, ImageBlock, ImageFormat, ImageSource, - InferenceConfiguration, Message, SystemContentBlock, Tool, ToolInputSchema, ToolSpecification, -}; -use aws_sdk_bedrockruntime::Client as BedrockRuntimeClient; use std::collections::HashMap; use windmill_common::{client::AuthedClient, error::Error}; -/// Constants for commonly used strings to avoid allocations -const FUNCTION_TYPE: &str = "function"; +// Re-export from shared module for use by other parts of the worker +pub use windmill_common::ai_bedrock::{check_env_credentials, BedrockClient}; +use windmill_common::ai_bedrock::{ + bedrock_stream_event_is_block_stop, bedrock_stream_event_to_text, + bedrock_stream_event_to_tool_delta, bedrock_stream_event_to_tool_start, build_tool_config, + create_inference_config, format_bedrock_error, openai_messages_to_bedrock, + streaming_tool_calls_to_openai, StreamingToolCall, +}; -#[derive(Debug, Clone)] -pub struct BearerTokenProvider { - token: String, -} - -impl BearerTokenProvider { - pub fn new(token: String) -> Self { - Self { token } - } -} - -impl ProvideToken for BearerTokenProvider { - fn provide_token<'a>(&'a self) -> aws_credential_types::provider::future::ProvideToken<'a> - where - Self: 'a, - { - aws_credential_types::provider::future::ProvideToken::ready(Ok( - aws_credential_types::Token::new(self.token.clone(), None), - )) - } -} - -pub struct BedrockClient { - client: BedrockRuntimeClient, -} - -impl BedrockClient { - pub async fn from_bearer_token(bearer_token: String, region: &str) -> Result { - let config = aws_sdk_bedrockruntime::config::Builder::new() - .region(aws_config::Region::new(region.to_string())) - .behavior_version(BehaviorVersion::latest()) - .token_provider(BearerTokenProvider::new(bearer_token)) - .build(); - - Ok(Self { client: BedrockRuntimeClient::from_conf(config) }) - } - - pub async fn from_credentials( - access_key_id: String, - secret_access_key: String, - region: &str, - ) -> Result { - let credentials = aws_credential_types::Credentials::new( - access_key_id, - secret_access_key, - None, // session token - None, // expiration - "windmill", - ); - - let config = aws_sdk_bedrockruntime::config::Builder::new() - .region(aws_config::Region::new(region.to_string())) - .behavior_version(BehaviorVersion::latest()) - .credentials_provider(credentials) - .build(); - - Ok(Self { client: BedrockRuntimeClient::from_conf(config) }) - } - - pub fn client(&self) -> &BedrockRuntimeClient { - &self.client - } -} - -/// Format AWS SDK errors with detailed information -fn format_bedrock_error(error: &aws_sdk_bedrockruntime::error::SdkError) -> String -where - E: std::fmt::Debug + std::fmt::Display, - R: std::fmt::Debug, -{ - use aws_sdk_bedrockruntime::error::SdkError; - - match error { - SdkError::ServiceError(err) => { - // Include both the display and debug representations for maximum detail - format!("Service error: {} (details: {:?})", err.err(), err) - } - SdkError::ConstructionFailure(err) => { - format!("Request construction failed: {:?}", err) - } - SdkError::DispatchFailure(err) => { - format!("Request dispatch failed: {:?}", err) - } - SdkError::ResponseError(err) => { - format!("Response error: {:?}", err) - } - SdkError::TimeoutError(err) => { - format!("Request timeout: {:?}", err) - } - _ => format!("{:?}", error), - } -} - -/// Convert serde_json::Value to AWS Smithy Document -fn json_to_document(value: serde_json::Value) -> aws_smithy_types::Document { - use aws_smithy_types::Document; - use serde_json::Value; - - match value { - Value::Object(map) => { - let mut doc_map = std::collections::HashMap::new(); - for (k, v) in map { - doc_map.insert(k, json_to_document(v)); - } - Document::Object(doc_map) - } - Value::Array(arr) => Document::Array(arr.into_iter().map(json_to_document).collect()), - Value::Number(num) => { - if let Some(i) = num.as_i64() { - Document::Number(aws_smithy_types::Number::PosInt(i as u64)) - } else if let Some(f) = num.as_f64() { - Document::Number(aws_smithy_types::Number::Float(f)) - } else { - Document::Number(aws_smithy_types::Number::PosInt(0)) - } - } - Value::String(s) => Document::String(s), - Value::Bool(b) => Document::Bool(b), - Value::Null => Document::Null, - } -} - -/// Convert OpenAI-style messages to Bedrock format -/// -/// Separates system messages from conversation messages as required by Bedrock API. -/// -/// # Returns -/// Tuple of (conversation_messages, system_prompts) -pub fn openai_messages_to_bedrock( - messages: &[OpenAIMessage], -) -> Result<(Vec, Vec), Error> { - let mut bedrock_messages = Vec::new(); - let mut system_prompts = Vec::new(); - - for msg in messages { - match msg.role.as_str() { - "system" => { - // Extract system messages separately - if let Some(ref content) = msg.content { - let text = content_to_text(content); - if !text.is_empty() { - system_prompts.push(SystemContentBlock::Text(text)); - } - } - } - "user" | "assistant" => { - bedrock_messages.push(convert_message(msg)?); - } - "tool" => { - // Tool results are handled as user messages with ToolResult content - bedrock_messages.push(convert_tool_message(msg)?); - } - _ => { - return Err(Error::BadRequest(format!("Unsupported role: {}", msg.role))); - } - } - } - - Ok((bedrock_messages, system_prompts)) -} - -/// Helper to extract text from OpenAIContent (ignoring images) -fn content_to_text(content: &OpenAIContent) -> String { - match content { - OpenAIContent::Text(text) => text.to_string(), - OpenAIContent::Parts(parts) => { - // Extract only text parts and join them - let text_parts: Vec<&str> = parts - .iter() - .filter_map(|part| match part { - ContentPart::Text { text } => Some(text.as_str()), - _ => None, - }) - .collect(); - text_parts.join(" ") - } - } -} - -/// Parse image data URL and extract format and base64 data -fn parse_image_data_url(url: &str) -> Result<(ImageFormat, Vec), Error> { - if !url.starts_with("data:") { - return Err(Error::internal_err("Image URL must be a data URL")); - } - - // Parse data:image/png;base64, - let base64_start = url - .find("base64,") - .ok_or_else(|| Error::internal_err("Invalid data URL format"))?; - - let base64_data = &url[base64_start + 7..]; - let mime_type = url - .split(';') - .next() - .and_then(|s| s.strip_prefix("data:")) - .unwrap_or("image/png"); - - // Extract format from MIME type (e.g., "image/png" -> "png") - let format_str = mime_type - .rsplit_once('/') - .map(|(_, format)| format) - .unwrap_or("png"); - - // Map to ImageFormat enum - let format = match format_str { - "png" => ImageFormat::Png, - "jpeg" | "jpg" => ImageFormat::Jpeg, - "gif" => ImageFormat::Gif, - "webp" => ImageFormat::Webp, - _ => ImageFormat::Png, // Default to PNG - }; - - // Decode base64 - let bytes = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, base64_data) - .map_err(|e| Error::internal_err(format!("Failed to decode base64 image: {}", e)))?; - - Ok((format, bytes)) -} - -/// Convert a ContentPart to Bedrock ContentBlock -fn content_part_to_block(part: &ContentPart) -> Result, Error> { - match part { - ContentPart::Text { text } => { - if text.is_empty() { - Ok(None) - } else { - Ok(Some(ContentBlock::Text(text.clone()))) - } - } - ContentPart::ImageUrl { image_url } => { - let (format, bytes) = parse_image_data_url(&image_url.url)?; - - let image_source = ImageSource::Bytes(bytes.into()); - let image_block = ImageBlock::builder() - .format(format) - .source(image_source) - .build() - .map_err(|e| Error::internal_err(format!("Failed to build image block: {}", e)))?; - - Ok(Some(ContentBlock::Image(image_block))) - } - ContentPart::S3Object { .. } => { - // S3Objects are already converted to ImageUrl by prepare_messages_for_api - // If we somehow get here, skip it - Ok(None) - } - } -} - -/// Convert a single OpenAI message to Bedrock Message -fn convert_message(msg: &OpenAIMessage) -> Result { - let role = match msg.role.as_str() { - "user" => ConversationRole::User, - "assistant" => ConversationRole::Assistant, - _ => { - return Err(Error::internal_err(format!( - "Unsupported role: {}", - msg.role - ))); - } - }; - - let mut content_blocks = Vec::new(); - - // Handle content (text and/or images) - if let Some(ref content) = msg.content { - match content { - OpenAIContent::Text(text) => { - if !text.is_empty() { - content_blocks.push(ContentBlock::Text(text.clone())); - } - } - OpenAIContent::Parts(parts) => { - for part in parts { - if let Some(block) = content_part_to_block(part)? { - content_blocks.push(block); - } - } - } - } - } - - // Handle tool calls (for assistant messages) - if let Some(ref tool_calls) = msg.tool_calls { - for tc in tool_calls { - content_blocks.push(convert_tool_call_to_content(tc)?); - } - } - - // Bedrock requires at least one content block - if content_blocks.is_empty() { - content_blocks.push(ContentBlock::Text(String::new())); - } - - Message::builder() - .role(role) - .set_content(Some(content_blocks)) - .build() - .map_err(|e| Error::internal_err(format!("Failed to build message: {}", e))) -} - -/// Convert OpenAI tool call to Bedrock ToolUse content block -fn convert_tool_call_to_content(tool_call: &OpenAIToolCall) -> Result { - let input = json_to_document( - serde_json::from_str(&tool_call.function.arguments) - .unwrap_or_else(|_| serde_json::json!({})), - ); - Ok(ContentBlock::ToolUse( - aws_sdk_bedrockruntime::types::ToolUseBlock::builder() - .tool_use_id(&tool_call.id) - .name(&tool_call.function.name) - .input(input) - .build() - .map_err(|e| Error::internal_err(format!("Failed to build tool use: {}", e)))?, - )) -} - -/// Convert tool result message to Bedrock format -fn convert_tool_message(msg: &OpenAIMessage) -> Result { - let tool_call_id = msg - .tool_call_id - .as_ref() - .ok_or_else(|| Error::internal_err("Tool message missing tool_call_id"))?; - - let content_str = msg - .content - .as_ref() - .map(|c| content_to_text(c)) - .unwrap_or_default(); - - // Try to parse as JSON, otherwise use text - let tool_result_content = - if let Ok(json_val) = serde_json::from_str::(&content_str) { - if json_val.is_object() { - vec![aws_sdk_bedrockruntime::types::ToolResultContentBlock::Json( - json_to_document(json_val), - )] - } else { - // Wrap primitives and arrays in an object - vec![aws_sdk_bedrockruntime::types::ToolResultContentBlock::Json( - json_to_document(serde_json::json!({"result": json_val})), - )] - } - } else { - vec![aws_sdk_bedrockruntime::types::ToolResultContentBlock::Text( - content_str.to_string(), - )] - }; - - let tool_result = ContentBlock::ToolResult( - aws_sdk_bedrockruntime::types::ToolResultBlock::builder() - .tool_use_id(tool_call_id) - .set_content(Some(tool_result_content)) - .build() - .map_err(|e| Error::internal_err(format!("Failed to build tool result: {}", e)))?, - ); - - Message::builder() - .role(ConversationRole::User) - .content(tool_result) - .build() - .map_err(|e| Error::internal_err(format!("Failed to build tool result message: {}", e))) -} - -/// Convert OpenAI tool definitions to Bedrock format -pub fn openai_tools_to_bedrock(tools: &[ToolDef]) -> Result, Error> { - tools - .iter() - .map(|tool_def| { - let spec = &tool_def.function; - - // Convert parameters (RawValue) to Document via serde_json::Value - let param_value: serde_json::Value = serde_json::from_str(spec.parameters.get()) - .map_err(|e| Error::internal_err(format!("Invalid tool schema: {}", e)))?; - let input_schema = ToolInputSchema::Json(json_to_document(param_value)); - - let tool_spec = ToolSpecification::builder() - .name(&spec.name) - .set_description(spec.description.clone()) - .input_schema(input_schema) - .build() - .map_err(|e| Error::internal_err(format!("Failed to build tool spec: {}", e)))?; - - Ok(Tool::ToolSpec(tool_spec)) - }) - .collect() -} - -/// Create inference configuration from parameters -pub fn create_inference_config( - temperature: Option, - max_tokens: Option, -) -> Option { - if temperature.is_none() && max_tokens.is_none() { - return None; - } - - let mut builder = InferenceConfiguration::builder(); - - if let Some(temp) = temperature { - builder = builder.temperature(temp); - } - - if let Some(max_tok) = max_tokens { - builder = builder.max_tokens(max_tok); - } - - Some(builder.build()) -} - -/// Extract text delta from Bedrock stream event -pub fn bedrock_stream_event_to_text(event: &ConverseStreamOutput) -> Option { - match event { - ConverseStreamOutput::ContentBlockDelta(delta) => delta - .delta() - .and_then(|d| d.as_text().ok()) - .map(|s| s.to_string()), - _ => None, - } -} - -/// Represents a streaming tool call being accumulated -#[derive(Debug, Clone)] -pub struct StreamingToolCall { - pub id: String, - pub name: String, - pub arguments: String, -} - -/// Extract tool use start event from stream -pub fn bedrock_stream_event_to_tool_start( - event: &ConverseStreamOutput, -) -> Option { - match event { - ConverseStreamOutput::ContentBlockStart(start) => { - if let Some(tool_use) = start.start().and_then(|s| s.as_tool_use().ok()) { - Some(StreamingToolCall { - id: tool_use.tool_use_id().to_string(), - name: tool_use.name().to_string(), - arguments: String::new(), - }) - } else { - None - } - } - _ => None, - } -} - -/// Extract tool use input delta from stream -pub fn bedrock_stream_event_to_tool_delta(event: &ConverseStreamOutput) -> Option { - match event { - ConverseStreamOutput::ContentBlockDelta(delta) => delta - .delta() - .and_then(|d| d.as_tool_use().ok()) - .map(|tool_use| tool_use.input().to_string()), - _ => None, - } -} - -/// Check if stream event indicates content block stop -pub fn bedrock_stream_event_is_block_stop(event: &ConverseStreamOutput) -> bool { - matches!(event, ConverseStreamOutput::ContentBlockStop(_)) -} - -/// Convert accumulated streaming tool calls to OpenAI format -pub fn streaming_tool_calls_to_openai(tool_calls: Vec) -> Vec { - tool_calls - .into_iter() - .map(|tc| OpenAIToolCall { - id: tc.id, - function: OpenAIFunction { name: tc.name, arguments: tc.arguments }, - r#type: FUNCTION_TYPE.to_string(), - extra_content: None, // Bedrock doesn't use thought signatures - }) - .collect() -} +// ============================================================================ +// Query Builder (Worker-specific orchestration) +// ============================================================================ #[derive(Default)] pub struct BedrockQueryBuilder; @@ -517,17 +49,21 @@ impl BedrockQueryBuilder { aws_access_key_id: Option<&str>, aws_secret_access_key: Option<&str>, ) -> Result { - // Create Bedrock client - use IAM credentials if provided, otherwise fall back to bearer token - let bedrock_client = match (aws_access_key_id, aws_secret_access_key) { - (Some(access_key_id), Some(secret_access_key)) => { - BedrockClient::from_credentials( - access_key_id.to_string(), - secret_access_key.to_string(), - region, - ) - .await? - } - _ => BedrockClient::from_bearer_token(api_key.to_string(), region).await?, + let bedrock_client = if !api_key.is_empty() { + BedrockClient::from_bearer_token(api_key.to_string(), region).await? + } else if let (Some(access_key_id), Some(secret_access_key)) = ( + aws_access_key_id.filter(|s| !s.is_empty()), + aws_secret_access_key.filter(|s| !s.is_empty()), + ) { + BedrockClient::from_credentials( + access_key_id.to_string(), + secret_access_key.to_string(), + None, + region, + ) + .await? + } else { + BedrockClient::from_env(region).await? }; // Prepare messages: convert S3Objects to ImageUrls by downloading from S3 @@ -536,11 +72,11 @@ impl BedrockQueryBuilder { // Convert messages to Bedrock format (separates system prompts) let (bedrock_messages, system_prompts) = openai_messages_to_bedrock(&prepared_messages)?; - // Build inference configuration + // Build inference configuration using shared helper let inference_config = create_inference_config(temperature, max_tokens.map(|t| t as i32)); // Build tool configuration with optional ToolChoice - let tool_config = self.build_tool_config(tools, structured_output_tool_name.is_some())?; + let tool_config = build_tool_config(tools, structured_output_tool_name.is_some())?; self.execute_converse_stream( &bedrock_client, @@ -554,36 +90,7 @@ impl BedrockQueryBuilder { .await } - /// Build tool configuration with optional ToolChoice for structured output - fn build_tool_config( - &self, - tools: Option<&[ToolDef]>, - force_tool_use: bool, - ) -> Result, Error> { - if let Some(tools) = tools { - let bedrock_tools = openai_tools_to_bedrock(tools)?; - let mut tool_config_builder = - aws_sdk_bedrockruntime::types::ToolConfiguration::builder() - .set_tools(Some(bedrock_tools)); - - // For structured output, force the model to use the tool - if force_tool_use { - tool_config_builder = tool_config_builder.tool_choice( - aws_sdk_bedrockruntime::types::ToolChoice::Any( - aws_sdk_bedrockruntime::types::AnyToolChoice::builder().build(), - ), - ); - } - - Ok(Some(tool_config_builder.build().map_err(|e| { - Error::internal_err(format!("Failed to build tool configuration: {}", e)) - })?)) - } else { - Ok(None) - } - } - - /// Execute streaming Bedrock request + /// Execute streaming Bedrock request using shared stream parsing functions async fn execute_converse_stream( &self, bedrock_client: &BedrockClient, @@ -594,7 +101,13 @@ impl BedrockQueryBuilder { tool_config: Option, stream_event_processor: Option, ) -> Result { - // Build streaming request + tracing::debug!( + "Worker Bedrock: executing converse_stream, messages={}, system_prompts={}, has_tools={}", + bedrock_messages.len(), + system_prompts.len(), + tool_config.is_some() + ); + let mut request_builder = bedrock_client .client() .converse_stream() @@ -613,46 +126,50 @@ impl BedrockQueryBuilder { request_builder = request_builder.set_tool_config(Some(config)); } - // Execute streaming request let mut stream = request_builder .send() .await .map_err(|e| { let error_msg = format!("Bedrock streaming API error: {}", format_bedrock_error(&e)); + tracing::error!("Worker Bedrock: {}", error_msg); Error::internal_err(error_msg) })? .stream; + tracing::debug!("Worker Bedrock: stream established, processing events"); + let mut accumulated_text = String::new(); let mut events_str = String::new(); let mut accumulated_tool_calls: HashMap = HashMap::new(); let mut current_tool_use_id: Option = None; - // Process stream events + // Process stream events using shared parsing functions loop { match stream.recv().await { Ok(Some(event)) => { - // Handle tool use start + // Handle tool use start using shared parser if let Some(tool_call) = bedrock_stream_event_to_tool_start(&event) { current_tool_use_id = Some(tool_call.id.clone()); accumulated_tool_calls.insert(tool_call.id.clone(), tool_call); } - // Handle text delta + // Handle text delta using shared parser if let Some(text_delta) = bedrock_stream_event_to_text(&event) { accumulated_text.push_str(&text_delta); if let Some(processor) = stream_event_processor.as_ref() { processor .send( - StreamingEvent::TokenDelta { content: text_delta }, + StreamingEvent::TokenDelta { + content: text_delta, + }, &mut events_str, ) .await?; } } - // Handle tool use input delta + // Handle tool use input delta using shared parser if let Some(input_delta) = bedrock_stream_event_to_tool_delta(&event) { if let Some(tool_id) = ¤t_tool_use_id { if let Some(tool_call) = accumulated_tool_calls.get_mut(tool_id) { @@ -661,12 +178,12 @@ impl BedrockQueryBuilder { } } - // Handle content block stop + // Handle content block stop using shared parser if bedrock_stream_event_is_block_stop(&event) { current_tool_use_id = None; } } - Ok(None) => break, // Stream ended + Ok(None) => break, Err(e) => { return Err(Error::internal_err(format!("Bedrock stream error: {}", e))); } diff --git a/backend/windmill-worker/src/ai/providers/mod.rs b/backend/windmill-worker/src/ai/providers/mod.rs index e35d01d5b8..f558c28289 100644 --- a/backend/windmill-worker/src/ai/providers/mod.rs +++ b/backend/windmill-worker/src/ai/providers/mod.rs @@ -1,4 +1,5 @@ pub mod anthropic; +#[cfg(feature = "bedrock")] pub mod bedrock; pub mod google_ai; pub mod openai; diff --git a/backend/windmill-worker/src/ai/providers/openai.rs b/backend/windmill-worker/src/ai/providers/openai.rs index ae33ce5121..e27d752871 100644 --- a/backend/windmill-worker/src/ai/providers/openai.rs +++ b/backend/windmill-worker/src/ai/providers/openai.rs @@ -11,35 +11,9 @@ use crate::ai::{ utils::extract_text_content, }; -#[derive(Deserialize, Serialize, Clone, Debug)] -pub struct OpenAIFunction { - pub name: String, - pub arguments: String, -} - -/// Google-specific extra content for thought signatures (Gemini 3 Pro / 2.5) -#[derive(Deserialize, Serialize, Clone, Debug, Default)] -pub struct GoogleExtraContent { - #[serde(skip_serializing_if = "Option::is_none")] - pub thought_signature: Option, -} - -/// Extra content for provider-specific metadata (e.g., Google thought signatures) -#[derive(Deserialize, Serialize, Clone, Debug, Default)] -pub struct ExtraContent { - #[serde(skip_serializing_if = "Option::is_none")] - pub google: Option, -} - -#[derive(Deserialize, Serialize, Clone, Debug)] -pub struct OpenAIToolCall { - pub id: String, - pub function: OpenAIFunction, - pub r#type: String, - /// Extra content for provider-specific metadata (e.g., Google Gemini thought signatures) - #[serde(skip_serializing_if = "Option::is_none")] - pub extra_content: Option, -} +pub use windmill_common::ai_types::{ + ExtraContent, OpenAIFunction, OpenAIToolCall, +}; // Responses API structures #[derive(Deserialize)] diff --git a/backend/windmill-worker/src/ai/types.rs b/backend/windmill-worker/src/ai/types.rs index 1a638a5e76..bd5d1e9691 100644 --- a/backend/windmill-worker/src/ai/types.rs +++ b/backend/windmill-worker/src/ai/types.rs @@ -1,4 +1,3 @@ -use crate::ai::providers::openai::OpenAIToolCall; use serde::{Deserialize, Serialize}; use serde_json::value::RawValue; use std::collections::HashMap; @@ -20,58 +19,10 @@ use windmill_common::{ }; use windmill_parser::Typ; -/// URL citation annotation for web search results -#[derive(Deserialize, Serialize, Clone, Debug)] -pub struct UrlCitation { - pub start_index: usize, - pub end_index: usize, - pub url: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub title: Option, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(tag = "type", rename_all = "snake_case")] -pub enum ContentPart { - Text { - text: String, - }, - #[serde(rename = "image_url")] - ImageUrl { - image_url: ImageUrlData, - }, - #[serde(rename = "s3_object")] - S3Object { - s3_object: S3Object, - }, -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -pub struct ImageUrlData { - pub url: String, // data:image/png;base64,... or https://... -} - -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(untagged)] -pub enum OpenAIContent { - Text(String), - Parts(Vec), -} - -#[derive(Deserialize, Serialize, Clone, Default, Debug)] -pub struct OpenAIMessage { - pub role: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub content: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub tool_calls: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub tool_call_id: Option, - #[serde(skip_serializing)] - pub agent_action: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub annotations: Option>, -} +// Re-export shared types from windmill_common::ai_types +pub use windmill_common::ai_types::{ + ContentPart, ImageUrlData, OpenAIContent, OpenAIMessage, ToolDef, ToolDefFunction, UrlCitation, +}; /// same as OpenAIMessage but with agent_action field included in the serialization #[derive(Serialize)] @@ -96,19 +47,6 @@ pub struct JsonSchemaFormat { pub strict: Option, } -#[derive(Serialize, Clone, Debug)] -pub struct ToolDefFunction { - pub name: String, - pub description: Option, - pub parameters: Box, -} - -#[derive(Serialize, Clone, Debug)] -pub struct ToolDef { - pub r#type: String, - pub function: ToolDefFunction, -} - #[derive(Serialize, Clone, Debug)] pub struct Tool { pub module: Option, @@ -159,6 +97,8 @@ struct AIAgentArgsRaw { memory: Option, // Legacy field for backward compatibility messages_context_length: Option, + #[serde(default)] + credentials_check: Option, } #[derive(Debug, Deserialize)] @@ -175,6 +115,7 @@ pub struct AIAgentArgs { pub streaming: Option, pub max_iterations: Option, pub memory: Option, + pub credentials_check: bool, } impl From for AIAgentArgs { @@ -206,6 +147,7 @@ impl From for AIAgentArgs { streaming: raw.streaming, max_iterations: raw.max_iterations, memory, + credentials_check: raw.credentials_check.unwrap_or(false), } } } @@ -224,9 +166,12 @@ pub struct ProviderResource { pub api_key: String, #[serde(alias = "baseUrl")] pub base_url: Option, + #[allow(dead_code)] pub region: Option, + #[allow(dead_code)] #[serde(alias = "awsAccessKeyId")] pub aws_access_key_id: Option, + #[allow(dead_code)] #[serde(alias = "awsSecretAccessKey")] pub aws_secret_access_key: Option, /// Platform for Anthropic API (standard or google_vertex_ai) @@ -260,14 +205,17 @@ impl ProviderWithResource { .await } + #[cfg(feature = "bedrock")] pub fn get_region(&self) -> Option<&str> { self.resource.region.as_deref() } + #[cfg(feature = "bedrock")] pub fn get_aws_access_key_id(&self) -> Option<&str> { self.resource.aws_access_key_id.as_deref() } + #[cfg(feature = "bedrock")] pub fn get_aws_secret_access_key(&self) -> Option<&str> { self.resource.aws_secret_access_key.as_deref() } diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index f8c51e5675..8069aa1773 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -1,3 +1,5 @@ +#[cfg(feature = "bedrock")] +use crate::ai::providers::bedrock::check_env_credentials; use crate::ai::tools::{execute_tool_calls, ToolExecutionContext}; use crate::ai::utils::{ add_message_to_conversation, any_tool_needs_previous_result, cleanup_mcp_clients, @@ -103,9 +105,18 @@ pub async fn handle_ai_agent_job( killpill_rx: &mut tokio::sync::broadcast::Receiver<()>, has_stream: &mut bool, ) -> Result, Error> { - let args = build_args_map(job, client, conn).await?; + // build_args_map returns None if no $res:/$var: transforms needed, in which case use original args + let args = match build_args_map(job, client, conn).await? { + Some(transformed) => transformed, + None => job.args.as_ref().map(|a| a.0.clone()).unwrap_or_default(), + }; let args = serde_json::from_str::(&serde_json::to_string(&args)?)?; + // Handle dry_run mode - check credentials without making API calls + if args.credentials_check { + return handle_credentials_check(&args.provider).await; + } + let Some(flow_step_id) = &job.flow_step_id else { return Err(Error::internal_err( "AI agent job has no flow step id".to_string(), @@ -315,7 +326,8 @@ pub async fn handle_ai_agent_job( let mut tools = tools; let mcp_clients = if !mcp_configs.is_empty() { - let (clients, mcp_tools) = load_mcp_tools(db, &job.workspace_id, mcp_configs, &client.token).await?; + let (clients, mcp_tools) = + load_mcp_tools(db, &job.workspace_id, mcp_configs, &client.token).await?; tools.extend(mcp_tools); clients } else { @@ -400,7 +412,6 @@ pub async fn run_agent( let output_type = args.output_type.as_ref().unwrap_or(&OutputType::Text); let base_url = args.provider.get_base_url(db).await?; let api_key = args.provider.get_api_key(); - let region = args.provider.get_region(); // Create the query builder for the provider let query_builder = create_query_builder(&args.provider); @@ -645,34 +656,43 @@ pub async fn run_agent( break; } - // Special handling for AWS Bedrock using the official SDK + // Handle AWS Bedrock provider specially using the official SDK let parsed = if args.provider.kind == AIProvider::AWSBedrock { - let Some(region) = region else { + #[cfg(feature = "bedrock")] + { + let region = args.provider.get_region(); + let Some(region) = region else { + return Err(Error::internal_err( + "AWS Bedrock region is required".to_string(), + )); + }; + // Use Bedrock SDK via dedicated query builder + crate::ai::providers::bedrock::BedrockQueryBuilder::default() + .execute_request( + &messages, + tool_defs.as_deref(), + args.provider.get_model(), + args.temperature, + args.max_completion_tokens, + api_key, + region, + stream_event_processor.clone(), + client, + &job.workspace_id, + structured_output_tool_name.as_deref(), + args.provider.get_aws_access_key_id(), + args.provider.get_aws_secret_access_key(), + ) + .await? + } + #[cfg(not(feature = "bedrock"))] + { return Err(Error::internal_err( - "AWS Bedrock region is required".to_string(), + "AWS Bedrock support is not enabled. Build with 'bedrock' feature.".to_string(), )); - }; - // Use Bedrock SDK via dedicated query builder - // Always use streaming for text output - crate::ai::providers::bedrock::BedrockQueryBuilder::default() - .execute_request( - &messages, - tool_defs.as_deref(), - args.provider.get_model(), - args.temperature, - args.max_completion_tokens, - api_key, - region, - stream_event_processor.clone(), - client, - &job.workspace_id, - structured_output_tool_name.as_deref(), - args.provider.get_aws_access_key_id(), - args.provider.get_aws_secret_access_key(), - ) - .await? + } } else { - // For non-Bedrock providers, use HTTP client + // For all other providers, use the HTTP client approach let build_args = BuildRequestArgs { messages: &messages, tools: tool_defs.as_deref(), @@ -687,7 +707,6 @@ pub async fn run_agent( has_websearch, }; - // Always use streaming for text output let request_body = query_builder .build_request(&build_args, client, &job.workspace_id) .await?; @@ -1041,3 +1060,40 @@ pub async fn run_agent( }, })) } + +/// Handle credentials check mode - check credentials without making API calls +async fn handle_credentials_check(provider: &ProviderWithResource) -> Result, Error> { + let result = match &provider.kind { + #[cfg(feature = "bedrock")] + AIProvider::AWSBedrock => { + let check = check_env_credentials().await; + serde_json::json!({ + "credentials_check": true, + "provider": "aws_bedrock", + "credentials": { + "available": check.available, + "access_key_id_prefix": check.access_key_id_prefix, + "region": check.region, + "error": check.error + } + }) + } + #[cfg(not(feature = "bedrock"))] + AIProvider::AWSBedrock => { + serde_json::json!({ + "credentials_check": true, + "provider": "aws_bedrock", + "error": "AWS Bedrock support is not enabled. Build with 'bedrock' feature." + }) + } + other => { + serde_json::json!({ + "credentials_check": true, + "provider": format!("{:?}", other), + "message": "Credentials check not implemented for this provider" + }) + } + }; + + serde_json::value::to_raw_value(&result).map_err(|e| Error::internal_err(e.to_string())) +} diff --git a/frontend/src/lib/components/ApiConnectForm.svelte b/frontend/src/lib/components/ApiConnectForm.svelte index 73da863251..05bb15b8c6 100644 --- a/frontend/src/lib/components/ApiConnectForm.svelte +++ b/frontend/src/lib/components/ApiConnectForm.svelte @@ -13,6 +13,7 @@ import { untrack } from 'svelte' import { base } from '$lib/base' import GitHubAppIntegration from './GitHubAppIntegration.svelte' + import BedrockCredentialsCheck from './BedrockCredentialsCheck.svelte' interface Props { resourceType: string @@ -210,6 +211,9 @@ onDescriptionUpdate={(newDescription) => (description = newDescription)} /> + {#if resourceType?.includes('bedrock')} + + {/if} {:else}

No corresponding resource type found in your workspace for {resourceType}. Define the value in diff --git a/frontend/src/lib/components/BedrockCredentialsCheck.svelte b/frontend/src/lib/components/BedrockCredentialsCheck.svelte new file mode 100644 index 0000000000..989367ef23 --- /dev/null +++ b/frontend/src/lib/components/BedrockCredentialsCheck.svelte @@ -0,0 +1,260 @@ + + +

+
+

AWS Environment Credentials Check

+ +
+ +

+ Check if AWS credentials are available from the environment. If available, you do not need to + configure credentials manually. +

+ +
+ +
+
+ + API Server + {#if apiStatus === 'loading'} + + {:else if apiStatus === 'success'} + + {:else if apiStatus === 'error'} + + {/if} +
+ + {#if apiResult} +
+ {#if apiResult.available} +
+ Available: {apiResult.access_key_id_prefix} +
+ {#if apiResult.region} +
Region: {apiResult.region}
+ {/if} + {:else} +
+ {apiResult.error ?? 'Not available'} +
+ {/if} +
+ {:else if apiStatus === 'idle'} +
Click "Check Credentials" to test
+ {/if} +
+ + +
+
+ + Worker + {#if workerStatus === 'loading'} + + {:else if workerStatus === 'success'} + + {:else if workerStatus === 'error'} + + {/if} +
+ + {#if workerResult} +
+ {#if workerResult.available} +
+ Available: {workerResult.access_key_id_prefix} +
+ {#if workerResult.region} +
Region: {workerResult.region}
+ {/if} + {#if workerResult.source} +
Source: {workerResult.source}
+ {/if} + {#if workerResult.worker} +
Worker: {workerResult.worker}
+ {/if} + {:else} +
+ {workerResult.error ?? 'Not available'} +
+ {#if workerResult.worker} +
Worker: {workerResult.worker}
+ {/if} + {/if} +
+ {:else if workerStatus === 'idle'} +
Click "Check Credentials" to test
+ {/if} +
+
+