starred/windmill
1
0
Fork 0
mirror of https://github.com/windmill-labs/windmill.git synced 2026-08-19 08:01:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(ai): native bedrock compatibility (#7668)

Browse Source 
* check endpoint

* use flag for cred check on worker

* use sdk for auth chat side

* cleaning

* cleaning

* also handle GET

* Add session token support in BedrockClient::from_credentials

Users with temporary STS credentials need session tokens for AWS
authentication. This adds the optional session_token parameter to
support these use cases.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Cache AWS SDK config loading to avoid repeated I/O

aws_config::load_defaults() performs environment variable lookups and
potentially file I/O on each call. Cache the result in a static OnceCell
to improve performance on high-frequency Bedrock requests.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Add debouncing to Bedrock credential check button

Prevent rapid clicks from spawning multiple concurrent flow preview
jobs for the worker credential check.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Remove orphaned SigV4 signing comment

The manual SigV4 signing code was removed in favor of the AWS SDK,
but this documentation comment was left behind.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* Remove verbose tracing info and change-related comments

- Remove per-request auth method info logs (too verbose for production)
- Simplify from_env log to single debug-level message
- Remove comments describing future changes rather than current behavior

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* cleaning

* Fix Bedrock feature flag compilation warnings

Gate AWS-specific imports, struct fields, and methods behind the bedrock
feature flag to eliminate dead code warnings when building without the
bedrock feature enabled.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* bedrock flag

* cleaning

* Move BedrockCredentialsCheck from edit drawer to add drawer

The component was incorrectly shown when editing resources instead of
when adding new ones. Moved it from ResourceEditor.svelte to
ApiConnectForm.svelte.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* cleaning

* cleaning

* make aws-config optional

---------

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Co-authored-by: Ruben Fiszel <ruben@windmill.dev>
This commit is contained in:
centdix
2026-01-26 12:04:14 +01:00
committed by GitHub
parent 5de8ecacdf
commit 2553d987d9
22 changed files with 2123 additions and 1367 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/build_windows_worker_.yml
+1 -1
View File
@@ -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"
.github/workflows/docker-image.yml
+2 -2
View File
@@ -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 }}
.github/workflows/publish_windows_worker.yml
+1 -1
View File
@@ -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"
backend/Cargo.lock
Generated
+31 -3
View File
@@ -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",
backend/Cargo.toml
+4 -2
View File
@@ -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"
backend/windmill-api/Cargo.toml
+5 -1
View File
@@ -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 }
backend/windmill-api/src/ai.rs
+141 -111
View File
@@ -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<String>,
pub organization_id: Option<String>,
pub user: Option<String>,
#[allow(dead_code)]
pub region: Option<String>,
#[allow(dead_code)]
pub aws_access_key_id: Option<String>,
#[allow(dead_code)]
pub aws_secret_access_key: Option<String>,
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<String>,
) -> Result<Json<windmill_common::ai_bedrock::BedrockCredentialsCheck>> {
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)))
}
backend/windmill-api/src/bedrock.rs
+671 -574
View File
File diff suppressed because it is too large Load Diff
backend/windmill-api/src/lib.rs
+1
View File
@@ -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;
backend/windmill-common/Cargo.toml
+3
View File
@@ -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
backend/windmill-common/src/ai_bedrock.rs
+693
View File
@@ -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<aws_config::SdkConfig> =
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<String>,
pub region: Option<String>,
pub error: Option<String>,
}
/// 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<Self, Error> {
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<String>,
region: &str,
) -> Result<Self, Error> {
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<Self, Error> {
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<E, R>(error: &aws_sdk_bedrockruntime::error::SdkError<E, R>) -> 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<Message>, Vec<SystemContentBlock>), Error> {
let mut bedrock_messages = Vec::new();
let mut system_prompts = Vec::new();
let mut pending_tool_results: Vec<ContentBlock> = 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<u8>), Error> {
if !url.starts_with("data:") {
return Err(Error::internal_err("Image URL must be a data URL"));
}
// Parse data:image/png;base64,<data>
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<Option<ContentBlock>, 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<Message, Error> {
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<ContentBlock, Error> {
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<ContentBlock, Error> {
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::<serde_json::Value>(&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<Vec<Tool>, 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<f32>,
max_tokens: Option<i32>,
) -> Option<InferenceConfiguration> {
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<String> {
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<StreamingToolCall> {
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<String> {
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<StreamingToolCall>) -> Vec<OpenAIToolCall> {
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<Option<aws_sdk_bedrockruntime::types::ToolConfiguration>, 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)
}
}
backend/windmill-common/src/ai_providers.rs
+26 -8
View File
@@ -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()
))
}
}
}
}
backend/windmill-common/src/ai_types.rs
+116
View File
@@ -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<String>,
}
#[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<ContentPart>),
}
#[derive(Serialize, Clone, Debug)]
pub struct ToolDefFunction {
pub name: String,
pub description: Option<String>,
pub parameters: Box<RawValue>,
}
#[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<String>,
}
/// 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<GoogleExtraContent>,
}
#[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<ExtraContent>,
}
/// 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<OpenAIContent>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<OpenAIToolCall>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
/// Worker-specific field for tracking agent actions in flows.
/// Never serialized; defaults to None when deserializing.
#[serde(skip_serializing, default)]
pub agent_action: Option<AgentAction>,
#[serde(skip_serializing_if = "Option::is_none")]
pub annotations: Option<Vec<UrlCitation>>,
}
backend/windmill-common/src/lib.rs
+3
View File
@@ -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;
backend/windmill-worker/Cargo.toml
+2 -1
View File
@@ -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
backend/windmill-worker/src/ai/providers/bedrock.rs
+57 -540
View File
@@ -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<Self, Error> {
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<Self, Error> {
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<E, R>(error: &aws_sdk_bedrockruntime::error::SdkError<E, R>) -> 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<Message>, Vec<SystemContentBlock>), 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<u8>), Error> {
if !url.starts_with("data:") {
return Err(Error::internal_err("Image URL must be a data URL"));
}
// Parse data:image/png;base64,<data>
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<Option<ContentBlock>, 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<Message, Error> {
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<ContentBlock, Error> {
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<Message, Error> {
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::<serde_json::Value>(&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<Vec<Tool>, 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<f32>,
max_tokens: Option<i32>,
) -> Option<InferenceConfiguration> {
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<String> {
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<StreamingToolCall> {
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<String> {
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<StreamingToolCall>) -> Vec<OpenAIToolCall> {
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<ParsedResponse, Error> {
// 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<Option<aws_sdk_bedrockruntime::types::ToolConfiguration>, 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<aws_sdk_bedrockruntime::types::ToolConfiguration>,
stream_event_processor: Option<StreamEventProcessor>,
) -> Result<ParsedResponse, Error> {
// 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<String, StreamingToolCall> = HashMap::new();
let mut current_tool_use_id: Option<String> = 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) = &current_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)));
}
backend/windmill-worker/src/ai/providers/mod.rs
+1
View File
@@ -1,4 +1,5 @@
pub mod anthropic;
#[cfg(feature = "bedrock")]
pub mod bedrock;
pub mod google_ai;
pub mod openai;
backend/windmill-worker/src/ai/providers/openai.rs
+3 -29
View File
@@ -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<String>,
}
/// 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<GoogleExtraContent>,
}
#[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<ExtraContent>,
}
pub use windmill_common::ai_types::{
ExtraContent, OpenAIFunction, OpenAIToolCall,
};
// Responses API structures
#[derive(Deserialize)]
backend/windmill-worker/src/ai/types.rs
+14 -66
View File
@@ -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<String>,
}
#[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<ContentPart>),
}
#[derive(Deserialize, Serialize, Clone, Default, Debug)]
pub struct OpenAIMessage {
pub role: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub content: Option<OpenAIContent>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<OpenAIToolCall>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
#[serde(skip_serializing)]
pub agent_action: Option<AgentAction>,
#[serde(skip_serializing_if = "Option::is_none")]
pub annotations: Option<Vec<UrlCitation>>,
}
// 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<bool>,
}
#[derive(Serialize, Clone, Debug)]
pub struct ToolDefFunction {
pub name: String,
pub description: Option<String>,
pub parameters: Box<RawValue>,
}
#[derive(Serialize, Clone, Debug)]
pub struct ToolDef {
pub r#type: String,
pub function: ToolDefFunction,
}
#[derive(Serialize, Clone, Debug)]
pub struct Tool {
pub module: Option<FlowModule>,
@@ -159,6 +97,8 @@ struct AIAgentArgsRaw {
memory: Option<Memory>,
// Legacy field for backward compatibility
messages_context_length: Option<usize>,
#[serde(default)]
credentials_check: Option<bool>,
}
#[derive(Debug, Deserialize)]
@@ -175,6 +115,7 @@ pub struct AIAgentArgs {
pub streaming: Option<bool>,
pub max_iterations: Option<usize>,
pub memory: Option<Memory>,
pub credentials_check: bool,
}
impl From<AIAgentArgsRaw> for AIAgentArgs {
@@ -206,6 +147,7 @@ impl From<AIAgentArgsRaw> 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<String>,
#[allow(dead_code)]
pub region: Option<String>,
#[allow(dead_code)]
#[serde(alias = "awsAccessKeyId")]
pub aws_access_key_id: Option<String>,
#[allow(dead_code)]
#[serde(alias = "awsSecretAccessKey")]
pub aws_secret_access_key: Option<String>,
/// 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()
}
backend/windmill-worker/src/ai_executor.rs
+84 -28
View File
@@ -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<Box<RawValue>, 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::<AIAgentArgs>(&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<Box<RawValue>, 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()))
}
frontend/src/lib/components/ApiConnectForm.svelte
+4
View File
@@ -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)}
/>
</div>
{#if resourceType?.includes('bedrock')}
<BedrockCredentialsCheck />
{/if}
{:else}
<p class="text-primary font-normal text-xs mb-4"
>No corresponding resource type found in your workspace for {resourceType}. Define the value in
frontend/src/lib/components/BedrockCredentialsCheck.svelte
+260
View File
@@ -0,0 +1,260 @@
<script lang="ts">
import { JobService } from '$lib/gen'
import { workspaceStore } from '$lib/stores'
import { tryEvery } from '$lib/utils'
import { Check, LoaderCircle, Server, X, Cpu } from 'lucide-svelte'
import Button from './common/button/Button.svelte'
interface CredentialsCheckResult {
available: boolean
access_key_id_prefix?: string | null
region?: string | null
source?: string | null
error?: string | null
}
type CheckStatus = 'idle' | 'loading' | 'success' | 'error'
let apiStatus: CheckStatus = $state('idle')
let apiResult: CredentialsCheckResult | null = $state(null)
let workerStatus: CheckStatus = $state('idle')
let workerResult: (CredentialsCheckResult & { worker?: string }) | null = $state(null)
let isChecking = $state(false)
async function checkApiCredentials() {
apiStatus = 'loading'
apiResult = null
try {
const response = await fetch(`/api/w/${$workspaceStore}/ai/check_bedrock_credentials`)
if (!response.ok) {
throw new Error(`HTTP error: ${response.status}`)
}
apiResult = await response.json()
apiStatus = apiResult?.available ? 'success' : 'error'
} catch (err) {
apiResult = {
available: false,
error: err instanceof Error ? err.message : 'Unknown error'
}
apiStatus = 'error'
}
}
async function checkWorkerCredentials() {
workerStatus = 'loading'
workerResult = null
try {
// Create minimal flow with AI agent dry_run step
const flowValue = {
modules: [
{
id: 'a',
value: {
type: 'aiagent' as const,
input_transforms: {
provider: {
type: 'static' as const,
value: {
kind: 'aws_bedrock',
resource: {
api_key: 'dry_run_placeholder',
region: 'us-east-1'
},
model: 'dry_run_placeholder'
}
},
user_message: { type: 'static' as const, value: 'dry_run_placeholder' },
output_type: { type: 'static' as const, value: 'text' },
credentials_check: { type: 'static' as const, value: true }
},
tools: []
}
}
]
}
const job = await JobService.runFlowPreview({
workspace: $workspaceStore!,
requestBody: {
value: flowValue,
args: {}
}
})
tryEvery({
tryCode: async () => {
const testResult = await JobService.getCompletedJob({
workspace: $workspaceStore!,
id: job
})
if (testResult.success && testResult.result) {
const result = testResult.result as {
credentials_check?: boolean
credentials?: CredentialsCheckResult
}
if (result?.credentials_check && result?.credentials) {
workerResult = {
...result.credentials,
source: 'worker_process',
worker: testResult.worker ?? undefined
}
workerStatus = workerResult.available ? 'success' : 'error'
} else {
workerResult = {
available: false,
error: 'Unexpected response format',
worker: testResult.worker ?? undefined
}
workerStatus = 'error'
}
} else {
workerResult = {
available: false,
error:
(testResult.result as Record<string, any>)?.['error']?.['message'] ?? 'Job failed',
worker: testResult.worker ?? undefined
}
workerStatus = 'error'
}
},
timeoutCode: async () => {
workerResult = {
available: false,
error: 'Timeout: job did not complete within 10s'
}
workerStatus = 'error'
try {
await JobService.cancelQueuedJob({
workspace: $workspaceStore!,
id: job,
requestBody: {
reason: 'Timeout checking Bedrock credentials'
}
})
} catch (err) {
console.error(err)
}
},
interval: 500,
timeout: 10000
})
} catch (err) {
workerResult = {
available: false,
error: err instanceof Error ? err.message : 'Unknown error'
}
workerStatus = 'error'
}
}
async function checkBoth() {
if (isChecking) return
isChecking = true
try {
await Promise.all([checkApiCredentials(), checkWorkerCredentials()])
} finally {
isChecking = false
}
}
</script>
<div class="flex flex-col gap-3 p-3 border rounded-md bg-surface-secondary">
<div class="flex items-center justify-between">
<h4 class="text-sm font-semibold">AWS Environment Credentials Check</h4>
<Button size="xs" variant="border" on:click={checkBoth} disabled={isChecking}>
{#if isChecking}
<LoaderCircle class="animate-spin mr-1.5 h-3.5 w-3.5" />
{/if}
Check Credentials
</Button>
</div>
<p class="text-xs text-secondary">
Check if AWS credentials are available from the environment. If available, you do not need to
configure credentials manually.
</p>
<div class="grid grid-cols-2 gap-3">
<!-- API Server Check -->
<div class="flex flex-col gap-1.5 p-2 border rounded bg-surface">
<div class="flex items-center gap-1.5 text-xs font-medium">
<Server class="h-3.5 w-3.5" />
<span>API Server</span>
{#if apiStatus === 'loading'}
<LoaderCircle class="animate-spin h-3.5 w-3.5 ml-auto text-blue-500" />
{:else if apiStatus === 'success'}
<Check class="h-3.5 w-3.5 ml-auto text-green-500" />
{:else if apiStatus === 'error'}
<X class="h-3.5 w-3.5 ml-auto text-red-500" />
{/if}
</div>
{#if apiResult}
<div class="text-xs">
{#if apiResult.available}
<div class="text-green-600 dark:text-green-400">
Available: {apiResult.access_key_id_prefix}
</div>
{#if apiResult.region}
<div class="text-secondary">Region: {apiResult.region}</div>
{/if}
{:else}
<div class="text-red-600 dark:text-red-400 break-words">
{apiResult.error ?? 'Not available'}
</div>
{/if}
</div>
{:else if apiStatus === 'idle'}
<div class="text-xs text-tertiary">Click "Check Credentials" to test</div>
{/if}
</div>
<!-- Worker Check -->
<div class="flex flex-col gap-1.5 p-2 border rounded bg-surface">
<div class="flex items-center gap-1.5 text-xs font-medium">
<Cpu class="h-3.5 w-3.5" />
<span>Worker</span>
{#if workerStatus === 'loading'}
<LoaderCircle class="animate-spin h-3.5 w-3.5 ml-auto text-blue-500" />
{:else if workerStatus === 'success'}
<Check class="h-3.5 w-3.5 ml-auto text-green-500" />
{:else if workerStatus === 'error'}
<X class="h-3.5 w-3.5 ml-auto text-red-500" />
{/if}
</div>
{#if workerResult}
<div class="text-xs">
{#if workerResult.available}
<div class="text-green-600 dark:text-green-400">
Available: {workerResult.access_key_id_prefix}
</div>
{#if workerResult.region}
<div class="text-secondary">Region: {workerResult.region}</div>
{/if}
{#if workerResult.source}
<div class="text-secondary">Source: {workerResult.source}</div>
{/if}
{#if workerResult.worker}
<div class="text-secondary">Worker: {workerResult.worker}</div>
{/if}
{:else}
<div class="text-red-600 dark:text-red-400 break-words">
{workerResult.error ?? 'Not available'}
</div>
{#if workerResult.worker}
<div class="text-secondary">Worker: {workerResult.worker}</div>
{/if}
{/if}
</div>
{:else if workerStatus === 'idle'}
<div class="text-xs text-tertiary">Click "Check Credentials" to test</div>
{/if}
</div>
</div>
</div>