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

refactor: unify AI provider credentials (#9317)

Browse Source 
* refactor: use provider credentials for worker builders

* refactor: resolve api proxy credentials directly

* fix: lazy load frontend eval modes
This commit is contained in:
centdix
2026-05-26 05:51:32 +00:00
committed by GitHub
parent 8bf7fd2c92
commit 4be930f585
7 changed files with 272 additions and 326 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ai_evals/adapters/frontend/benchmarkRunner.ts
+15 -11
View File
@@ -10,10 +10,6 @@ import { runSuite } from "../../core/runSuite";
import type { BenchmarkRunResult, ModeRunner } from "../../core/types";
import { resolveWindmillBackendSettings } from "../../core/windmillBackendSettings";
import { emitFrontendBenchmarkProgress } from "./progress";
import { createAppModeRunner } from "../../modes/app";
import { createFlowModeRunner } from "../../modes/flow";
import { createGlobalModeRunner } from "../../modes/global";
import { createScriptModeRunner } from "../../modes/script";
import { DEFAULT_JUDGE_MODEL } from "../../core/judge";
export type FrontendBenchmarkMode = "flow" | "app" | "script" | "global";
@@ -40,7 +36,7 @@ export async function runFrontendBenchmarkFromEnv(): Promise<BenchmarkRunResult>
const backendSettings = resolveWindmillBackendSettings();
const selectedCases = await loadSelectedCases(mode, caseIds);
const modeRunner = getModeRunner(
const modeRunner = await getModeRunner(
mode,
getFrontendEvalModel(model),
backendValidation,
@@ -69,25 +65,33 @@ export async function runFrontendBenchmarkFromEnv(): Promise<BenchmarkRunResult>
});
}
function getModeRunner(
async function getModeRunner(
mode: FrontendBenchmarkMode,
model: ReturnType<typeof getFrontendEvalModel>,
backendValidation: ReturnType<typeof resolveBackendValidationSettings>,
backendSettings: ReturnType<typeof resolveWindmillBackendSettings>,
): ModeRunner<any, any, any> {
): Promise<ModeRunner<any, any, any>> {
switch (mode) {
case "flow":
case "flow": {
const { createFlowModeRunner } = await import("../../modes/flow");
return createFlowModeRunner(model, backendValidation, backendSettings);
case "app":
}
case "app": {
const { createAppModeRunner } = await import("../../modes/app");
return createAppModeRunner(model, backendSettings);
case "script":
}
case "script": {
const { createScriptModeRunner } = await import("../../modes/script");
return createScriptModeRunner(
model,
backendValidation,
backendSettings,
);
case "global":
}
case "global": {
const { createGlobalModeRunner } = await import("../../modes/global");
return createGlobalModeRunner(model, backendSettings);
}
}
}
backend/windmill-ai/src/providers/mod.rs
+3 -23
View File
@@ -6,35 +6,15 @@ pub mod openai;
pub mod openrouter;
pub mod other;
use crate::{
ai_providers::AIProvider, proxy::ProviderCredentials, query_builder::QueryBuilder,
types::ProviderWithResource,
};
use crate::{ai_providers::AIProvider, proxy::ProviderCredentials, query_builder::QueryBuilder};
use self::{
anthropic::AnthropicQueryBuilder, google_ai::GoogleAIQueryBuilder, openai::OpenAIQueryBuilder,
openrouter::OpenRouterQueryBuilder, other::OtherQueryBuilder,
};
/// Factory function to create the appropriate query builder for a provider.
pub fn create_query_builder(provider: &ProviderWithResource) -> Box<dyn QueryBuilder> {
match provider.kind {
AIProvider::GoogleAI => {
Box::new(GoogleAIQueryBuilder::new(provider.get_platform().clone()))
}
AIProvider::OpenAI => Box::new(OpenAIQueryBuilder::new(provider.kind.clone())),
AIProvider::Anthropic => Box::new(AnthropicQueryBuilder::new(
provider.kind.clone(),
provider.get_platform().clone(),
provider.get_enable_1m_context(),
)),
AIProvider::OpenRouter => Box::new(OpenRouterQueryBuilder::new()),
_ => Box::new(OtherQueryBuilder::new(provider.kind.clone())),
}
}
/// Factory function to create the appropriate query builder from resolved proxy credentials.
pub fn create_proxy_query_builder(credentials: &ProviderCredentials) -> Box<dyn QueryBuilder> {
/// Factory function to create the appropriate query builder from resolved credentials.
pub fn create_query_builder(credentials: &ProviderCredentials) -> Box<dyn QueryBuilder> {
match credentials.provider {
AIProvider::GoogleAI => Box::new(GoogleAIQueryBuilder::new(credentials.platform.clone())),
AIProvider::OpenAI => Box::new(OpenAIQueryBuilder::new(credentials.provider.clone())),
backend/windmill-ai/src/proxy.rs
+4 -3
View File
@@ -7,10 +7,11 @@ use windmill_common::error::{Error, Result};
use crate::ai_providers::{AIPlatform, AIProvider};
use crate::utils::AI_HTTP_HEADERS;
/// Resolved provider credentials and proxy-specific context.
/// Resolved provider credentials shared by API proxy and worker execution.
///
/// This is intentionally separate from the worker's `ProviderWithResource`: API
/// proxy credentials are already resolved from workspace or instance resources.
/// Raw API resources and worker agent payloads convert into this shape at their
/// execution boundaries. Request-specific state such as the selected model stays
/// outside this type.
#[derive(Clone, Debug)]
pub struct ProviderCredentials {
pub provider: AIProvider,
backend/windmill-ai/src/types.rs
+29
View File
@@ -18,6 +18,7 @@ pub struct McpToolSource {
use crate::{
ai_google::sanitize_schema_for_google,
ai_providers::{empty_string_as_none, AIProvider},
proxy::ProviderCredentials,
};
use windmill_common::{db::DB, error::Error, flow_status::AgentAction, flows::FlowModule};
use windmill_parser::Typ;
@@ -222,6 +223,34 @@ impl ProviderWithResource {
.await
}
/// Convert worker agent provider input into resolved runtime credentials.
///
/// Callers must only pass resources that were already authorized for the
/// current job/workspace; this helper does not perform access checks.
pub async fn to_provider_credentials(&self, db: &DB) -> Result<ProviderCredentials, Error> {
let base_url = if self.kind == AIProvider::AWSBedrock {
String::new()
} else {
self.get_base_url(db).await?
};
Ok(ProviderCredentials {
provider: self.kind.clone(),
base_url,
api_key: self.resource.api_key.clone(),
access_token: None,
organization_id: None,
user: None,
region: self.resource.region.clone(),
aws_access_key_id: self.resource.aws_access_key_id.clone(),
aws_secret_access_key: self.resource.aws_secret_access_key.clone(),
aws_session_token: self.resource.aws_session_token.clone(),
platform: self.resource.platform.clone(),
enable_1m_context: self.resource.enable_1m_context,
custom_headers: self.resource.headers.clone(),
})
}
#[cfg(feature = "bedrock")]
pub fn get_region(&self) -> Option<&str> {
self.resource.region.as_deref()
backend/windmill-api/src/ai.rs
+137 -262
View File
@@ -23,7 +23,7 @@ use windmill_ai::providers::bedrock::{
handle_bedrock_proxy, BedrockProxyResponse, BedrockProxyResponseBody,
};
use windmill_ai::providers::{
create_proxy_query_builder,
create_query_builder,
google_ai::{
handle_google_ai_chat_proxy, handle_google_ai_models_proxy, GoogleAIProxyResponse,
GoogleAIProxyResponseBody,
@@ -114,7 +114,7 @@ lazy_static::lazy_static! {
static ref OPENAI_AZURE_BASE_PATH: Option<String> = std::env::var("OPENAI_AZURE_BASE_PATH").ok();
pub static ref AI_REQUEST_CACHE: Cache<(String, AIProvider), ExpiringAIRequestConfig> = Cache::new(500);
pub static ref AI_REQUEST_CACHE: Cache<(String, AIProvider), ExpiringProviderCredentials> = Cache::new(500);
}
@@ -181,26 +181,6 @@ enum AIResource {
Standard(AIStandardResource),
}
#[derive(Deserialize, Clone, Debug)]
struct AIRequestConfig {
pub base_url: String,
pub api_key: Option<String>,
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>,
#[allow(dead_code)]
pub aws_session_token: Option<String>,
pub platform: AIPlatform,
pub enable_1m_context: bool,
pub custom_headers: HashMap<String, String>,
}
/// Resolve a `$var:` reference. When `user_db`/`authed` are provided the query
/// goes through an RLS-scoped connection so the caller can only read variables
/// they are authorised to access. Without auth context the raw pool is used
@@ -218,196 +198,142 @@ async fn resolve_var(
}
}
impl AIRequestConfig {
pub async fn new(
provider: &AIProvider,
db: &DB,
w_id: &str,
resource: AIResource,
authed: Option<&ApiAuthed>,
) -> Result<Self> {
// When authed is provided, resolve $var: references through RLS so that
// users can only read variables they have permission to access.
let user_db = authed.map(|_| UserDB::new(db.clone()));
async fn resolve_provider_credentials(
provider: &AIProvider,
db: &DB,
w_id: &str,
resource: AIResource,
authed: Option<&ApiAuthed>,
) -> Result<ProviderCredentials> {
// When authed is provided, resolve $var: references through RLS so that
// users can only read variables they have permission to access.
let user_db = authed.map(|_| UserDB::new(db.clone()));
let (
api_key,
access_token,
organization_id,
base_url,
user,
region,
aws_access_key_id,
aws_secret_access_key,
aws_session_token,
platform,
enable_1m_context,
custom_headers,
) = match resource {
AIResource::Standard(resource) => {
let region = resource.region.clone();
let platform = resource.platform.clone();
let enable_1m_context = resource.enable_1m_context;
let custom_headers = resource.headers.clone();
// Skip get_base_url for Bedrock - it uses SDK directly, not HTTP
let base_url = if matches!(provider, AIProvider::AWSBedrock) {
String::new()
} else {
provider.get_base_url(resource.base_url, db).await?
};
let api_key = if let Some(api_key) = resource.api_key {
Some(resolve_var(api_key, db, w_id, user_db.as_ref(), authed).await?)
} else {
None
};
let organization_id = if let Some(organization_id) = resource.organization_id {
Some(resolve_var(organization_id, db, w_id, user_db.as_ref(), authed).await?)
} else {
None
};
let aws_access_key_id = if let Some(access_key_id) = resource.aws_access_key_id {
Some(resolve_var(access_key_id, db, w_id, user_db.as_ref(), authed).await?)
} else {
None
};
let aws_secret_access_key = if let Some(secret_access_key) =
resource.aws_secret_access_key
{
match resource {
AIResource::Standard(resource) => {
// Skip get_base_url for Bedrock - it uses SDK directly, not HTTP
let base_url = if matches!(provider, AIProvider::AWSBedrock) {
String::new()
} else {
provider.get_base_url(resource.base_url, db).await?
};
let api_key = if let Some(api_key) = resource.api_key {
Some(resolve_var(api_key, db, w_id, user_db.as_ref(), authed).await?)
} else {
None
};
let organization_id = if let Some(organization_id) = resource.organization_id {
Some(resolve_var(organization_id, db, w_id, user_db.as_ref(), authed).await?)
} else {
None
};
let aws_access_key_id = if let Some(access_key_id) = resource.aws_access_key_id {
Some(resolve_var(access_key_id, db, w_id, user_db.as_ref(), authed).await?)
} else {
None
};
let aws_secret_access_key =
if let Some(secret_access_key) = resource.aws_secret_access_key {
Some(resolve_var(secret_access_key, db, w_id, user_db.as_ref(), authed).await?)
} else {
None
};
let aws_session_token = if let Some(session_token) = resource.aws_session_token {
Some(resolve_var(session_token, db, w_id, user_db.as_ref(), authed).await?)
} else {
None
};
let aws_session_token = if let Some(session_token) = resource.aws_session_token {
Some(resolve_var(session_token, db, w_id, user_db.as_ref(), authed).await?)
} else {
None
};
(
api_key,
None,
organization_id,
base_url,
None,
region,
aws_access_key_id,
aws_secret_access_key,
aws_session_token,
platform,
enable_1m_context,
custom_headers,
)
}
AIResource::OAuth(resource) => {
let user = if let Some(user) = resource.user.clone() {
Some(resolve_var(user, db, w_id, user_db.as_ref(), authed).await?)
} else {
None
};
let token =
Self::get_token_using_oauth(resource, db, w_id, user_db.as_ref(), authed)
.await?;
let base_url = provider.get_base_url(None, db).await?;
Ok(ProviderCredentials {
provider: provider.clone(),
base_url,
api_key,
access_token: None,
organization_id,
user: None,
region: resource.region,
aws_access_key_id,
aws_secret_access_key,
aws_session_token,
platform: resource.platform,
enable_1m_context: resource.enable_1m_context,
custom_headers: resource.headers,
})
}
AIResource::OAuth(resource) => {
let user = if let Some(user) = resource.user.clone() {
Some(resolve_var(user, db, w_id, user_db.as_ref(), authed).await?)
} else {
None
};
let token = get_token_using_oauth(resource, db, w_id, user_db.as_ref(), authed).await?;
let base_url = provider.get_base_url(None, db).await?;
(
None,
Some(token),
None,
base_url,
user,
None,
None,
None,
None,
AIPlatform::Standard,
false,
HashMap::new(),
)
}
};
Ok(Self {
base_url,
organization_id,
api_key,
access_token,
user,
region,
aws_access_key_id,
aws_secret_access_key,
aws_session_token,
platform,
enable_1m_context,
custom_headers,
})
}
async fn get_token_using_oauth(
mut resource: AIOAuthResource,
db: &DB,
w_id: &str,
user_db: Option<&UserDB>,
authed: Option<&ApiAuthed>,
) -> Result<String> {
resource.client_id = resolve_var(resource.client_id, db, w_id, user_db, authed).await?;
resource.client_secret =
resolve_var(resource.client_secret, db, w_id, user_db, authed).await?;
resource.token_url = resolve_var(resource.token_url, db, w_id, user_db, authed).await?;
let mut params = HashMap::new();
params.insert("grant_type", "client_credentials");
params.insert("scope", "https://cognitiveservices.azure.com/.default");
let response = HTTP_CLIENT
.post(resource.token_url)
.form(&params)
.basic_auth(resource.client_id, Some(resource.client_secret))
.send()
.await
.and_then(|r| r.error_for_status())
.map_err(|err| {
Error::internal_err(format!(
"Failed to get access token using credentials flow: {}",
err
))
})?;
let response = response.json::<OAuthTokens>().await.map_err(|err| {
Error::internal_err(format!(
"Failed to parse access token from credentials flow: {}",
err
))
})?;
Ok(response.access_token)
}
fn into_provider_credentials(self, provider: AIProvider) -> ProviderCredentials {
ProviderCredentials {
provider,
base_url: self.base_url,
api_key: self.api_key,
access_token: self.access_token,
organization_id: self.organization_id,
user: self.user,
region: self.region,
aws_access_key_id: self.aws_access_key_id,
aws_secret_access_key: self.aws_secret_access_key,
aws_session_token: self.aws_session_token,
platform: self.platform,
enable_1m_context: self.enable_1m_context,
custom_headers: self.custom_headers,
Ok(ProviderCredentials {
provider: provider.clone(),
base_url,
api_key: None,
access_token: Some(token),
organization_id: None,
user,
region: None,
aws_access_key_id: None,
aws_secret_access_key: None,
aws_session_token: None,
platform: AIPlatform::Standard,
enable_1m_context: false,
custom_headers: HashMap::new(),
})
}
}
}
async fn get_token_using_oauth(
mut resource: AIOAuthResource,
db: &DB,
w_id: &str,
user_db: Option<&UserDB>,
authed: Option<&ApiAuthed>,
) -> Result<String> {
resource.client_id = resolve_var(resource.client_id, db, w_id, user_db, authed).await?;
resource.client_secret = resolve_var(resource.client_secret, db, w_id, user_db, authed).await?;
resource.token_url = resolve_var(resource.token_url, db, w_id, user_db, authed).await?;
let mut params = HashMap::new();
params.insert("grant_type", "client_credentials");
params.insert("scope", "https://cognitiveservices.azure.com/.default");
let response = HTTP_CLIENT
.post(resource.token_url)
.form(&params)
.basic_auth(resource.client_id, Some(resource.client_secret))
.send()
.await
.and_then(|r| r.error_for_status())
.map_err(|err| {
Error::internal_err(format!(
"Failed to get access token using credentials flow: {}",
err
))
})?;
let response = response.json::<OAuthTokens>().await.map_err(|err| {
Error::internal_err(format!(
"Failed to parse access token from credentials flow: {}",
err
))
})?;
Ok(response.access_token)
}
#[derive(Clone, Debug)]
pub struct ExpiringAIRequestConfig {
config: AIRequestConfig,
pub struct ExpiringProviderCredentials {
credentials: ProviderCredentials,
expires_at: std::time::Instant,
instance_ai_config_revision: Option<u64>,
}
impl ExpiringAIRequestConfig {
fn new(config: AIRequestConfig, instance_ai_config_revision: Option<u64>) -> Self {
impl ExpiringProviderCredentials {
fn new(credentials: ProviderCredentials, instance_ai_config_revision: Option<u64>) -> Self {
Self {
config,
credentials,
expires_at: std::time::Instant::now() + std::time::Duration::from_secs(60),
instance_ai_config_revision,
}
@@ -622,7 +548,7 @@ async fn global_proxy(
enable_1m_context: false,
custom_headers: HashMap::new(),
};
let query_builder = create_proxy_query_builder(&credentials);
let query_builder = create_query_builder(&credentials);
let proxy_request = query_builder.build_proxy_request(&ProxyBuildArgs {
method: &method,
path: &ai_path,
@@ -715,9 +641,9 @@ async fn proxy(
check_scopes(&authed, || format!("resources:read:{}", resource_path))?;
}
let request_config = match workspace_cache {
let credentials = match workspace_cache {
Some(request_cache) if !request_cache.is_expired() && forced_resource_path.is_none() => {
request_cache.config
request_cache.credentials
}
_ => {
let (resource_path, save_to_cache, resource_workspace, instance_ai_config_revision) =
@@ -825,7 +751,7 @@ async fn proxy(
} else {
None
};
let request_config = AIRequestConfig::new(
let credentials = resolve_provider_credentials(
&provider,
&db,
&resource_workspace,
@@ -836,13 +762,13 @@ async fn proxy(
if save_to_cache {
AI_REQUEST_CACHE.insert(
(w_id.clone(), provider.clone()),
ExpiringAIRequestConfig::new(
request_config.clone(),
ExpiringProviderCredentials::new(
credentials.clone(),
instance_ai_config_revision,
),
);
}
request_config
credentials
}
};
@@ -876,7 +802,6 @@ async fn proxy(
.await?;
tx.commit().await?;
let credentials = request_config.into_provider_credentials(provider.clone());
let proxy_args = ProxyBuildArgs {
method: &method,
path: &ai_path,
@@ -913,7 +838,6 @@ async fn proxy(
.await?;
tx.commit().await?;
let credentials = request_config.into_provider_credentials(provider.clone());
let response = handle_bedrock_proxy(&ProxyBuildArgs {
method: &method,
path: &ai_path,
@@ -936,8 +860,7 @@ async fn proxy(
let request = match proxy_mode {
ProxyExecutionMode::HttpForward => {
let credentials = request_config.into_provider_credentials(provider.clone());
let query_builder = create_proxy_query_builder(&credentials);
let query_builder = create_query_builder(&credentials);
let proxy_request = query_builder.build_proxy_request(&ProxyBuildArgs {
method: &method,
path: &ai_path,
@@ -1003,8 +926,9 @@ mod tests {
static TEST_LOCK: LazyLock<Mutex<()>> = LazyLock::new(|| Mutex::new(()));
fn sample_request_config() -> AIRequestConfig {
AIRequestConfig {
fn sample_provider_credentials() -> ProviderCredentials {
ProviderCredentials {
provider: AIProvider::OpenAI,
base_url: "https://example.com".to_string(),
api_key: None,
access_token: None,
@@ -1020,70 +944,21 @@ mod tests {
}
}
#[test]
fn maps_request_config_to_provider_credentials() {
let mut custom_headers = HashMap::new();
custom_headers.insert("X-Test".to_string(), "yes".to_string());
let config = AIRequestConfig {
base_url: "https://example.com".to_string(),
api_key: Some("api-key".to_string()),
access_token: Some("access-token".to_string()),
organization_id: Some("org-id".to_string()),
user: Some("user-id".to_string()),
region: Some("us-east-1".to_string()),
aws_access_key_id: Some("aws-access-key".to_string()),
aws_secret_access_key: Some("aws-secret-key".to_string()),
aws_session_token: Some("aws-session-token".to_string()),
platform: AIPlatform::GoogleVertexAi,
enable_1m_context: true,
custom_headers,
};
let credentials = config.into_provider_credentials(AIProvider::Anthropic);
assert_eq!(credentials.provider, AIProvider::Anthropic);
assert_eq!(credentials.base_url, "https://example.com");
assert_eq!(credentials.api_key.as_deref(), Some("api-key"));
assert_eq!(credentials.access_token.as_deref(), Some("access-token"));
assert_eq!(credentials.organization_id.as_deref(), Some("org-id"));
assert_eq!(credentials.user.as_deref(), Some("user-id"));
assert_eq!(credentials.region.as_deref(), Some("us-east-1"));
assert_eq!(
credentials.aws_access_key_id.as_deref(),
Some("aws-access-key")
);
assert_eq!(
credentials.aws_secret_access_key.as_deref(),
Some("aws-secret-key")
);
assert_eq!(
credentials.aws_session_token.as_deref(),
Some("aws-session-token")
);
assert_eq!(credentials.platform, AIPlatform::GoogleVertexAi);
assert!(credentials.enable_1m_context);
assert_eq!(
credentials.custom_headers.get("X-Test").map(String::as_str),
Some("yes")
);
}
#[test]
fn invalidates_all_cached_providers_for_workspace() {
let _guard = TEST_LOCK.lock().unwrap();
AI_REQUEST_CACHE.clear();
AI_REQUEST_CACHE.insert(
("workspace-a".to_string(), AIProvider::OpenAI),
ExpiringAIRequestConfig::new(sample_request_config(), None),
ExpiringProviderCredentials::new(sample_provider_credentials(), None),
);
AI_REQUEST_CACHE.insert(
("workspace-a".to_string(), AIProvider::Anthropic),
ExpiringAIRequestConfig::new(sample_request_config(), None),
ExpiringProviderCredentials::new(sample_provider_credentials(), None),
);
AI_REQUEST_CACHE.insert(
("workspace-b".to_string(), AIProvider::OpenAI),
ExpiringAIRequestConfig::new(sample_request_config(), None),
ExpiringProviderCredentials::new(sample_provider_credentials(), None),
);
invalidate_ai_request_cache_for_workspace("workspace-a");
@@ -1104,8 +979,8 @@ mod tests {
let _guard = TEST_LOCK.lock().unwrap();
AI_REQUEST_CACHE.clear();
let cached = ExpiringAIRequestConfig::new(
sample_request_config(),
let cached = ExpiringProviderCredentials::new(
sample_provider_credentials(),
Some(current_instance_ai_config_revision()),
);
assert!(!cached.is_expired());
backend/windmill-worker/src/ai_executor.rs
+14 -18
View File
@@ -586,16 +586,12 @@ pub async fn run_agent(
tool_abort_handles: ToolAbortHandles,
) -> error::Result<Box<RawValue>> {
let output_type = args.output_type.as_ref().unwrap_or(&OutputType::Text);
// Skip get_base_url for Bedrock - it uses SDK directly, not HTTP
let base_url = if args.provider.kind == AIProvider::AWSBedrock {
String::new()
} else {
args.provider.get_base_url(db).await?
};
let api_key = args.provider.get_api_key().unwrap_or("");
let credentials = args.provider.to_provider_credentials(db).await?;
let base_url = &credentials.base_url;
let api_key = credentials.api_key.as_deref().unwrap_or("");
// Create the query builder for the provider
let query_builder = create_query_builder(&args.provider);
let query_builder = create_query_builder(&credentials);
// Initialize messages
let mut messages =
@@ -859,12 +855,12 @@ pub async fn run_agent(
}
// Handle AWS Bedrock provider specially using the official SDK
let parsed = if args.provider.kind == AIProvider::AWSBedrock {
let parsed = if credentials.provider == AIProvider::AWSBedrock {
#[cfg(feature = "bedrock")]
{
let region = args
.provider
.get_region()
let region = credentials
.region
.as_deref()
.unwrap_or(windmill_ai::ai_providers::USE_ENV_REGION);
// Use Bedrock SDK via dedicated query builder
windmill_ai::providers::bedrock::BedrockQueryBuilder::default()
@@ -880,9 +876,9 @@ pub async fn run_agent(
client,
&job.workspace_id,
structured_output_tool_name.as_deref(),
args.provider.get_aws_access_key_id(),
args.provider.get_aws_secret_access_key(),
args.provider.get_aws_session_token(),
credentials.aws_access_key_id.as_deref(),
credentials.aws_secret_access_key.as_deref(),
credentials.aws_session_token.as_deref(),
)
.await?
}
@@ -913,14 +909,14 @@ pub async fn run_agent(
.await?;
let endpoint =
query_builder.get_endpoint(&base_url, args.provider.get_model(), output_type);
let auth_headers = query_builder.get_auth_headers(api_key, &base_url, output_type);
query_builder.get_endpoint(base_url, args.provider.get_model(), output_type);
let auth_headers = query_builder.get_auth_headers(api_key, base_url, output_type);
let timeout = resolve_job_timeout(conn, &job.workspace_id, job.id, job.timeout)
.await
.0;
let resource_headers = args.provider.get_headers();
let resource_headers = &credentials.custom_headers;
// Helper to build HTTP request with headers
let build_http_request = |body: String| {
docs/windmill-ai-refactor-plan.md
+70 -9
View File
@@ -5,7 +5,7 @@
AI provider logic is currently split across three crates with duplicate code:
- **windmill-common** — base types (`ai_types`, `ai_providers`, `ai_google`, `ai_bedrock`, `ai_cache`)
- **windmill-api** — chat proxy routes (`ai.rs`), audit logging, caching, and DB-backed credential resolution through `AIRequestConfig`
- **windmill-api** — chat proxy routes (`ai.rs`), audit logging, caching, and DB-backed credential resolution into `ProviderCredentials`
- **windmill-worker** — agent execution (`ai/` module) with `QueryBuilder` trait, SSE parsers, provider implementations
The goal: a single `windmill-ai` crate with all AI provider logic. Worker agent execution uses `QueryBuilder`; the API proxy uses `QueryBuilder::build_proxy_request` for HTTP-forwarding providers and native proxy handlers for providers that need response conversion or SDK execution.
@@ -25,7 +25,7 @@ windmill-common does **NOT** re-export from windmill-ai (would be circular). All
## Reviewer Note: Keep API Proxy Unification Split
The crate boundary, shared utilities, SSE parsers, image handling, worker provider implementations, and provider-specific API proxy transformations are now in `windmill-ai`. The remaining duplication is credential shape and resolution: `windmill-api` still resolves DB-backed proxy credentials through `AIRequestConfig`, while worker agent execution still receives `ProviderWithResource`.
The crate boundary, shared utilities, SSE parsers, image handling, worker provider implementations, provider-specific API proxy transformations, and resolved runtime credential shape are now in `windmill-ai`. Raw API resources and worker agent provider payloads remain separate input/deserialization shapes and convert into `ProviderCredentials` at execution boundaries.
Do not jump directly from the current state to full proxy and credential unification in one PR. The API proxy combines request transformation, endpoint selection, auth headers, custom headers, OAuth user injection, Azure URL handling, Anthropic Vertex handling, Bedrock SDK calls, and SSE keepalive behavior. Split the work by risk:
- Introduce shared proxy request and credential types first.
@@ -62,7 +62,7 @@ Out of scope:
Validation:
- `cargo test -p windmill-ai proxy`
- `cargo test -p windmill-api maps_request_config_to_provider_credentials`
- `cargo test -p windmill-api invalidates_all_cached_providers_for_workspace`
- `cargo check -p windmill-ai -p windmill-api`
- `cargo check -p windmill-ai -p windmill-api --features bedrock`
@@ -98,14 +98,14 @@ Out of scope:
Validation:
- `cargo test -p windmill-ai google_ai`
- `cargo test -p windmill-ai proxy`
- `cargo test -p windmill-api maps_request_config_to_provider_credentials`
- `cargo test -p windmill-api invalidates_all_cached_providers_for_workspace`
- `cargo test -p windmill-ai anthropic`
Follow-up status: Bedrock native proxy handling has since moved into
`windmill-ai`, and the API-local `windmill-api/src/bedrock.rs` module has been
removed.
## Current Phase PR: Bedrock Native Proxy Migration
## Completed Phase: Bedrock Native Proxy Migration ✅
Goal: move the remaining native-provider API proxy execution out of
`windmill-api` and into `windmill-ai`, while leaving API-owned routing,
@@ -157,6 +157,61 @@ visible for later hardening work.
keepalive wrapper so it works for both `reqwest::Error` streams and Bedrock's
SDK-backed `std::io::Error` streams.
## Completed Phase: Credential Unification Phase 1 ✅
Goal: make `ProviderCredentials` the shared resolved runtime credential shape
without overloading it with raw resource input or model-selection state.
`AIRequestConfig` and `ProviderWithResource` are not equivalent concepts:
`AIRequestConfig` is API-side resolved state after DB, variable, OAuth, and
resource handling, while `ProviderWithResource` is worker-side raw agent input
that also carries the selected model. Keep raw/deserialization types separate and
convert them into `ProviderCredentials` at execution boundaries.
Suggested PR title: `refactor(ai): use provider credentials for worker builders`.
Scope:
- Add a worker-side conversion from `ProviderWithResource` to
`ProviderCredentials`.
- Keep `model` outside `ProviderCredentials`; it remains agent request data.
- Keep `ProviderWithResource` as the backward-compatible deserialization type for
existing agent payloads.
- Use `ProviderCredentials` for worker query-builder creation.
- Collapse `create_query_builder` and `create_proxy_query_builder` into one
`create_query_builder(&ProviderCredentials)` factory.
Out of scope:
- Do not remove API-local `AIRequestConfig` yet.
- Do not change API request-cache behavior.
- Do not change worker agent payload shape or serialized field names.
Validation:
- `cargo check -p windmill-ai -p windmill-api -p windmill-worker`
- `cargo check -p windmill-ai -p windmill-api -p windmill-worker --features bedrock`
## Completed Phase: Credential Unification Phase 2 ✅
Goal: remove the API-local resolved credential wrapper after worker execution
already uses the shared shape.
Suggested PR title: `refactor(ai): resolve api proxy credentials directly`.
Scope:
- Change API credential resolution to return `ProviderCredentials` directly.
- Replace `ExpiringAIRequestConfig` with an expiring `ProviderCredentials`
cache entry.
- Remove `AIRequestConfig::into_provider_credentials`.
- Delete `AIRequestConfig` entirely if no API-only behavior remains.
Out of scope:
- Do not merge raw worker resource input into `ProviderCredentials`.
- Do not put model selection into `ProviderCredentials`.
Validation:
- `cargo check -p windmill-ai -p windmill-api -p windmill-worker`
- `cargo check -p windmill-ai -p windmill-api -p windmill-worker --features bedrock`
- `cargo test -p windmill-api invalidates_all_cached_providers_for_workspace`
## Step-by-Step Plan
Each step produces a compiling, working backend.
@@ -313,18 +368,24 @@ pub struct ProxyRequest {
- `supports_native_fim`, `transform_fim_to_chat_completions` — moved to windmill-ai
**Keep** in API:
- `AIRequestConfig::new` credential resolution until it is refactored to produce `ProviderCredentials`
- credential resolution from DB, workspace settings, instance settings, variables, and OAuth into `ProviderCredentials`
- HTTP routes, audit logging, request caching
- `inject_keepalives`, `is_sse_response` helpers
- `AIConfig`, `ExpiringAIRequestConfig` caching types
- `AIConfig`, `ExpiringProviderCredentials` caching types
---
### Step 9: Unify credential resolution
Merge `AIRequestConfig` (API-side) and `ProviderWithResource` (worker-side) into a single credential shape in windmill-ai.
Make `ProviderCredentials` the single resolved runtime credential shape in
windmill-ai, while keeping raw API and worker input/deserialization types at
their boundaries.
Both currently carry: api_key, base_url, region, platform, custom_headers, AWS credentials. The API's `AIRequestConfig::new` resolves credentials from DB (workspace/instance settings). The worker's `ProviderWithResource` gets credentials from the flow module definition.
The API's `resolve_provider_credentials` resolves credentials from DB, workspace
or instance settings, variables, and OAuth. The worker's `ProviderWithResource`
gets raw credentials from the flow module definition and also carries the
selected model. Convert both paths into `ProviderCredentials`; do not make
`ProviderCredentials` carry raw resource state or the model.
Extend `windmill_ai::proxy::ProviderCredentials` as needed so both can produce it:
```rust