starred/windmill
1
0
Fork 0
mirror of https://github.com/windmill-labs/windmill.git synced 2026-10-29 00:03:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: move bedrock proxy handling to windmill-ai (#9309)

Browse Source 
* refactor: move bedrock proxy handling to windmill-ai

* docs: track ai refactor follow-ups
This commit is contained in:
centdix
2026-06-08 11:35:45 +02:00
committed by tristantr
parent 528a7afcd0
commit 2df3dd09dd
9 changed files with 923 additions and 998 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/windmill-api/src/ai.rs
+39 -89
View File
@@ -1,5 +1,3 @@
#[cfg(feature = "bedrock")]
use crate::bedrock;
use crate::db::{ApiAuthed, DB};
use crate::utils::check_scopes;
@@ -20,6 +18,10 @@ use windmill_ai::ai_cache::current_instance_ai_config_revision;
use windmill_ai::ai_providers::{
empty_string_as_none, AIPlatform, AIProvider, ProviderConfig, ProviderModel,
};
#[cfg(feature = "bedrock")]
use windmill_ai::providers::bedrock::{
handle_bedrock_proxy, BedrockProxyResponse, BedrockProxyResponseBody,
};
use windmill_ai::providers::{
create_proxy_query_builder,
google_ai::{
@@ -540,6 +542,18 @@ fn google_ai_proxy_response_to_body(
(response.status_code, response.headers, body)
}
#[cfg(feature = "bedrock")]
fn bedrock_proxy_response_to_body(
response: BedrockProxyResponse,
) -> (http::StatusCode, HeaderMap, axum::body::Body) {
let body = match response.body {
BedrockProxyResponseBody::Fixed(body) => axum::body::Body::from(body),
BedrockProxyResponseBody::Stream(stream) => axum::body::Body::from_stream(stream),
};
(response.status_code, response.headers, body)
}
pub(crate) fn inject_keepalives<S>(
upstream: S,
interval: Duration,
@@ -885,95 +899,31 @@ async fn proxy(
// 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!(proxy_mode, ProxyExecutionMode::NativeAwsBedrock)
&& 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)
};
if matches!(proxy_mode, ProxyExecutionMode::NativeAwsBedrock) {
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?;
// For Bedrock requests, use the SDK-based approach
if matches!(proxy_mode, ProxyExecutionMode::NativeAwsBedrock) {
let region = request_config
.region
.as_deref()
.unwrap_or(windmill_ai::ai_providers::USE_ENV_REGION);
let credentials = request_config.into_provider_credentials(provider.clone());
let response = handle_bedrock_proxy(&ProxyBuildArgs {
method: &method,
path: &ai_path,
headers: &headers,
body: &body,
credentials: &credentials,
})
.await?;
// 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(),
request_config.aws_session_token.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(),
request_config.aws_session_token.as_deref(),
region,
)
.await;
}
}
// 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(),
request_config.aws_session_token.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(),
request_config.aws_session_token.as_deref(),
region,
)
.await;
}
}
}
return Ok(bedrock_proxy_response_to_body(response));
}
// When bedrock feature is disabled, return error for Bedrock provider
backend/windmill-api/src/bedrock.rs
-873
View File
@@ -1,873 +0,0 @@
//! 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 serde::Deserialize;
use windmill_ai::ai_bedrock::build_tool_config;
use windmill_ai::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,
};
use windmill_ai::ai_types::{
OpenAIFunction, OpenAIMessage, OpenAIToolCall, ToolDef, ToolDefFunction,
};
use windmill_common::error::{Error, Result};
// ============================================================================
// Shared Request Types for SDK-Based Handlers
// ============================================================================
/// OpenAI-format request body for Bedrock SDK handlers
#[derive(Deserialize, Debug)]
struct OpenAIRequest {
messages: Vec<OpenAIMessage>,
#[serde(default)]
tools: Option<Vec<OpenAIToolDef>>,
#[serde(default)]
tool_choice: Option<serde_json::Value>,
#[serde(default)]
max_tokens: Option<i32>,
#[serde(default)]
temperature: Option<f32>,
}
#[derive(Deserialize, Debug)]
struct OpenAIToolDef {
#[serde(default)]
#[allow(dead_code)]
r#type: Option<String>,
function: OpenAIToolFunction,
}
#[derive(Deserialize, Debug)]
struct OpenAIToolFunction {
name: String,
#[serde(default)]
description: Option<String>,
#[serde(default)]
parameters: Option<serde_json::Value>,
}
// ============================================================================
// 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,
session_token: Option<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>,
aws_session_token: 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(),
session_token: aws_session_token
.filter(|token| !token.is_empty())
.map(str::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>,
aws_session_token: Option<&str>,
region: &str,
) -> Result<BedrockClient> {
match determine_auth_config(
api_key,
aws_access_key_id,
aws_secret_access_key,
aws_session_token,
) {
BedrockAuthConfig::BearerToken(key) => BedrockClient::from_bearer_token(key, region).await,
BedrockAuthConfig::IamCredentials { access_key_id, secret_access_key, session_token } => {
BedrockClient::from_credentials(access_key_id, secret_access_key, session_token, 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>,
enable_prompt_caching: bool,
) -> Result<Option<aws_sdk_bedrockruntime::types::ToolConfiguration>> {
if let Some(tools) = tools {
let tool_defs: Vec<ToolDef> = tools
.iter()
.map(|t| ToolDef {
r#type: "function".to_string(),
function: 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);
build_tool_config(Some(&tool_defs), force_tool_use, enable_prompt_caching)
} 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>,
aws_session_token: Option<&str>,
region: &str,
) -> Result<aws_sdk_bedrock::Client> {
use aws_config::BehaviorVersion;
use windmill_ai::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,
aws_session_token,
) {
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, session_token } => {
let credentials = aws_credential_types::Credentials::new(
access_key_id,
secret_access_key,
session_token,
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>,
aws_session_token: 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,
aws_session_token,
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<serde_json::Value> = 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::<Vec<_>>(),
"outputModalities": m.output_modalities().iter().map(|o| o.as_str()).collect::<Vec<_>>(),
"responseStreamingSupported": m.response_streaming_supported(),
"inferenceTypesSupported": m.inference_types_supported().iter().map(|i| i.as_str()).collect::<Vec<_>>(),
})
})
.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>,
aws_session_token: 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,
aws_session_token,
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<serde_json::Value> = 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>,
aws_session_token: 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,
aws_session_token,
region,
)
.await?;
// Convert messages using shared conversion
let enable_prompt_caching =
windmill_ai::ai_bedrock::bedrock_model_supports_prompt_caching(model);
let (bedrock_messages, system_prompts) =
windmill_ai::ai_bedrock::openai_messages_to_bedrock(
&openai_req.messages,
enable_prompt_caching,
)?;
// Build inference configuration
let inference_config = windmill_ai::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(),
enable_prompt_caching,
)?;
// Build the SDK request
let mut request_builder = bedrock_client
.client()
.converse_stream()
.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 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");
// 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),
))
}
/// 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,
) -> impl futures::Stream<Item = std::result::Result<bytes::Bytes, std::io::Error>> + Send {
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
struct StreamState {
id: String,
model: String,
created: u64,
tool_calls: HashMap<usize, (String, String, String)>, // 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>,
aws_session_token: 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,
aws_session_token,
region,
)
.await?;
// Convert messages using shared conversion
let enable_prompt_caching =
windmill_ai::ai_bedrock::bedrock_model_supports_prompt_caching(model);
let (bedrock_messages, system_prompts) =
windmill_ai::ai_bedrock::openai_messages_to_bedrock(
&openai_req.messages,
enable_prompt_caching,
)?;
// Build inference configuration
let inference_config = windmill_ai::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(),
enable_prompt_caching,
)?;
// 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",
"tool_use" => "tool_calls",
"stop_sequence" => "stop",
"guardrail_intervened" | "content_filtered" => "content_filter",
_ => "stop",
};
// Extract message content
let mut text_content = String::new();
let mut tool_calls: Vec<OpenAIToolCall> = Vec::new();
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: 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,
});
}
_ => {}
}
}
}
}
// Build the message
let message = if !tool_calls.is_empty() {
serde_json::json!({
"role": "assistant",
"content": if text_content.is_empty() { serde_json::Value::Null } else { serde_json::Value::String(text_content) },
"tool_calls": tool_calls
})
} else {
serde_json::json!({
"role": "assistant",
"content": text_content
})
};
// Extract usage information
let usage = if let Some(usage_data) = response.usage() {
serde_json::json!({
"prompt_tokens": usage_data.input_tokens(),
"completion_tokens": usage_data.output_tokens(),
"total_tokens": usage_data.total_tokens()
})
} else {
serde_json::json!({
"prompt_tokens": 0,
"completion_tokens": 0,
"total_tokens": 0
})
};
// Build OpenAI-format response
let openai_resp = serde_json::json!({
"id": id,
"object": "chat.completion",
"created": created,
"model": model,
"choices": [{
"index": 0,
"message": message,
"finish_reason": finish_reason
}],
"usage": usage
});
let response_body = serde_json::to_vec(&openai_resp)
.map_err(|e| Error::internal_err(format!("Failed to serialize OpenAI response: {}", e)))?;
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),
))
}
/// 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));
}
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,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn determine_auth_config_prioritizes_bearer_token() {
let config = determine_auth_config(
Some("bearer-token"),
Some("AKIA123"),
Some("secret"),
Some("session-token"),
);
match config {
BedrockAuthConfig::BearerToken(token) => assert_eq!(token, "bearer-token"),
_ => panic!("expected bearer token auth config"),
}
}
#[test]
fn determine_auth_config_uses_iam_with_optional_session_token() {
let config =
determine_auth_config(None, Some("AKIA123"), Some("secret"), Some("session-token"));
match config {
BedrockAuthConfig::IamCredentials {
access_key_id,
secret_access_key,
session_token,
} => {
assert_eq!(access_key_id, "AKIA123");
assert_eq!(secret_access_key, "secret");
assert_eq!(session_token.as_deref(), Some("session-token"));
}
_ => panic!("expected IAM auth config"),
}
}
#[test]
fn determine_auth_config_treats_empty_session_token_as_none() {
let config = determine_auth_config(None, Some("AKIA123"), Some("secret"), Some(""));
match config {
BedrockAuthConfig::IamCredentials { session_token, .. } => {
assert!(session_token.is_none());
}
_ => panic!("expected IAM auth config"),
}
}
#[test]
fn determine_auth_config_falls_back_to_environment() {
let config = determine_auth_config(None, Some("AKIA123"), None, Some("session-token"));
assert!(matches!(config, BedrockAuthConfig::Environment));
}
}
backend/windmill-api/src/lib.rs
-2
View File
@@ -73,8 +73,6 @@ 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;
mod db;