diff --git a/backend/windmill-api/src/ai.rs b/backend/windmill-api/src/ai.rs index abde227214..3c3ba77bb5 100644 --- a/backend/windmill-api/src/ai.rs +++ b/backend/windmill-api/src/ai.rs @@ -10,8 +10,8 @@ use reqwest::{Client, RequestBuilder}; use serde::{Deserialize, Serialize}; use serde_json::value::RawValue; use std::collections::HashMap; -use windmill_common::ai_providers::{AIProvider, ProviderConfig, ProviderModel}; use windmill_audit::{audit_oss::audit_log, ActionKind}; +use windmill_common::ai_providers::{AIProvider, ProviderConfig, ProviderModel, AZURE_API_VERSION}; use windmill_common::error::{to_anyhow, Error, Result}; lazy_static::lazy_static! { @@ -25,9 +25,6 @@ lazy_static::lazy_static! { pub static ref AI_REQUEST_CACHE: Cache<(String, AIProvider), ExpiringAIRequestConfig> = Cache::new(500); } -const AZURE_API_VERSION: &str = "2025-04-01-preview"; -const OPENAI_BASE_URL: &str = "https://api.openai.com/v1"; - #[derive(Deserialize, Debug)] struct AIOAuthResource { client_id: String, @@ -154,21 +151,13 @@ impl AIRequestConfig { let base_url = self.base_url.trim_end_matches('/'); - let is_azure = matches!(provider, AIProvider::OpenAI) && base_url != OPENAI_BASE_URL - || matches!(provider, AIProvider::AzureOpenAI); + let is_azure = provider.is_azure_openai(base_url); let is_anthropic = matches!(provider, AIProvider::Anthropic); let is_anthropic_sdk = headers.get("X-Anthropic-SDK").is_some(); let url = if is_azure && method != Method::GET { - if base_url.ends_with("/deployments") { - let model = Self::get_azure_model(&body)?; - format!("{}/{}/{}", base_url, model, path) - } else if base_url.ends_with("/openai") { - let model = Self::get_azure_model(&body)?; - format!("{}/deployments/{}/{}", base_url, model, path) - } else { - format!("{}/{}", base_url, path) - } + let model = AIProvider::extract_model_from_body(&body)?; + AIProvider::build_azure_openai_url(base_url, &model, path) } else if is_anthropic_sdk { let truncated_base_url = base_url.trim_end_matches("/v1"); format!("{}/{}", truncated_base_url, path) @@ -233,18 +222,6 @@ impl AIRequestConfig { .map_err(|e| Error::internal_err(format!("Failed to reserialize request body: {}", e)))? .into()) } - - fn get_azure_model(body: &Bytes) -> Result { - #[derive(Deserialize, Debug)] - struct AzureModel { - model: String, - } - - let azure_model: AzureModel = serde_json::from_slice(body) - .map_err(|e| Error::internal_err(format!("Failed to parse request body: {}", e)))?; - - Ok(azure_model.model) - } } #[derive(Clone, Debug)] diff --git a/backend/windmill-common/src/ai_providers.rs b/backend/windmill-common/src/ai_providers.rs index 577f613f21..46d5d687c3 100644 --- a/backend/windmill-common/src/ai_providers.rs +++ b/backend/windmill-common/src/ai_providers.rs @@ -10,6 +10,9 @@ lazy_static::lazy_static! { static ref OPENAI_AZURE_BASE_PATH: Option = std::env::var("OPENAI_AZURE_BASE_PATH").ok(); } +pub const AZURE_API_VERSION: &str = "2025-04-01-preview"; +pub const OPENAI_BASE_URL: &str = "https://api.openai.com/v1"; + #[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Hash, Clone)] #[serde(rename_all = "lowercase")] pub enum AIProvider { @@ -78,6 +81,39 @@ impl AIProvider { pub fn is_anthropic(&self) -> bool { matches!(self, AIProvider::Anthropic) } + + /// Check if this provider/URL combination represents Azure OpenAI + pub fn is_azure_openai(&self, base_url: &str) -> bool { + (matches!(self, AIProvider::OpenAI) && base_url != OPENAI_BASE_URL) + || matches!(self, AIProvider::AzureOpenAI) + } + + /// Build Azure OpenAI URL with deployment model path + pub fn build_azure_openai_url(base_url: &str, model: &str, path: &str) -> String { + let base_url = base_url.trim_end_matches('/'); + + if base_url.ends_with("/deployments") { + format!("{}/{}/{}", base_url, model, path) + } else if base_url.ends_with("/openai") { + format!("{}/deployments/{}/{}", base_url, model, path) + } else { + format!("{}/{}", base_url, path) + } + } + + /// Extract model from request body (needed for Azure deployments) + pub fn extract_model_from_body(body: &[u8]) -> Result { + #[derive(serde::Deserialize)] + struct ModelRequest { + model: String, + } + + let model_request: ModelRequest = serde_json::from_slice(body).map_err(|e| { + Error::internal_err(format!("Failed to parse request body for model: {}", e)) + })?; + + Ok(model_request.model) + } } impl TryFrom<&str> for AIProvider { diff --git a/backend/windmill-worker/src/ai/providers/google_ai.rs b/backend/windmill-worker/src/ai/providers/google_ai.rs index 9b08d40e46..b5a27a5f99 100644 --- a/backend/windmill-worker/src/ai/providers/google_ai.rs +++ b/backend/windmill-worker/src/ai/providers/google_ai.rs @@ -1,7 +1,7 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; use serde_json; -use windmill_common::{client::AuthedClient, error::Error}; +use windmill_common::{ai_providers::AIProvider, client::AuthedClient, error::Error}; use crate::ai::{ image_handler::download_and_encode_s3_image, @@ -166,7 +166,7 @@ impl QueryBuilder for GoogleAIQueryBuilder { match args.output_type { OutputType::Text => { // For text output, use OpenAI-compatible format - let openai_builder = super::openai::OpenAIQueryBuilder::new(); + let openai_builder = super::openai::OpenAIQueryBuilder::new(AIProvider::GoogleAI); openai_builder .build_request(args, client, workspace_id) .await @@ -180,7 +180,7 @@ impl QueryBuilder for GoogleAIQueryBuilder { // For chat completions (text), use OpenAI parser if url.contains("/chat/completions") { - let openai_builder = super::openai::OpenAIQueryBuilder::new(); + let openai_builder = super::openai::OpenAIQueryBuilder::new(AIProvider::GoogleAI); return openai_builder.parse_response(response).await; } @@ -259,6 +259,7 @@ impl QueryBuilder for GoogleAIQueryBuilder { fn get_auth_headers( &self, api_key: &str, + _base_url: &str, output_type: &OutputType, ) -> Vec<(&'static str, String)> { match output_type { diff --git a/backend/windmill-worker/src/ai/providers/openai.rs b/backend/windmill-worker/src/ai/providers/openai.rs index 26bc804aed..421c5761f6 100644 --- a/backend/windmill-worker/src/ai/providers/openai.rs +++ b/backend/windmill-worker/src/ai/providers/openai.rs @@ -1,7 +1,7 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; use serde_json; -use windmill_common::{client::AuthedClient, error::Error}; +use windmill_common::{ai_providers::AIProvider, client::AuthedClient, error::Error}; use crate::ai::{ image_handler::download_and_encode_s3_image, @@ -92,11 +92,13 @@ pub struct OpenAIRequest<'a> { pub response_format: Option, } -pub struct OpenAIQueryBuilder; +pub struct OpenAIQueryBuilder { + provider_kind: AIProvider, +} impl OpenAIQueryBuilder { - pub fn new() -> Self { - Self + pub fn new(provider_kind: AIProvider) -> Self { + Self { provider_kind } } pub async fn prepare_messages_for_api( @@ -332,18 +334,29 @@ impl QueryBuilder for OpenAIQueryBuilder { } } - fn get_endpoint(&self, base_url: &str, _model: &str, output_type: &OutputType) -> String { - match output_type { - OutputType::Text => format!("{}/chat/completions", base_url), - OutputType::Image => format!("{}/responses", base_url), + fn get_endpoint(&self, base_url: &str, model: &str, output_type: &OutputType) -> String { + let path = match output_type { + OutputType::Text => "chat/completions", + OutputType::Image => "responses", + }; + + if self.provider_kind.is_azure_openai(base_url) { + AIProvider::build_azure_openai_url(base_url, model, path) + } else { + format!("{}/{}", base_url, path) } } fn get_auth_headers( &self, api_key: &str, + base_url: &str, _output_type: &OutputType, ) -> Vec<(&'static str, String)> { - vec![("Authorization", format!("Bearer {}", api_key))] + if self.provider_kind.is_azure_openai(base_url) { + vec![("api-key", api_key.to_string())] + } else { + vec![("Authorization", format!("Bearer {}", api_key))] + } } } diff --git a/backend/windmill-worker/src/ai/providers/openrouter.rs b/backend/windmill-worker/src/ai/providers/openrouter.rs index 9fc6c6f051..3d5dd0314a 100644 --- a/backend/windmill-worker/src/ai/providers/openrouter.rs +++ b/backend/windmill-worker/src/ai/providers/openrouter.rs @@ -1,7 +1,7 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; use serde_json; -use windmill_common::{client::AuthedClient, error::Error}; +use windmill_common::{ai_providers::AIProvider, client::AuthedClient, error::Error}; use crate::ai::{ providers::openai::{OpenAIQueryBuilder, OpenAIResponse}, @@ -59,7 +59,7 @@ pub struct OpenRouterQueryBuilder { impl OpenRouterQueryBuilder { pub fn new() -> Self { - Self { openai_builder: OpenAIQueryBuilder::new() } + Self { openai_builder: OpenAIQueryBuilder::new(AIProvider::OpenRouter) } } } @@ -195,6 +195,7 @@ impl QueryBuilder for OpenRouterQueryBuilder { fn get_auth_headers( &self, api_key: &str, + _base_url: &str, _output_type: &OutputType, ) -> Vec<(&'static str, String)> { vec![("Authorization", format!("Bearer {}", api_key))] diff --git a/backend/windmill-worker/src/ai/query_builder.rs b/backend/windmill-worker/src/ai/query_builder.rs index 07b9c96e4e..6b043726bc 100644 --- a/backend/windmill-worker/src/ai/query_builder.rs +++ b/backend/windmill-worker/src/ai/query_builder.rs @@ -54,6 +54,7 @@ pub trait QueryBuilder: Send + Sync { fn get_auth_headers( &self, api_key: &str, + base_url: &str, output_type: &OutputType, ) -> Vec<(&'static str, String)>; } @@ -65,6 +66,6 @@ pub fn create_query_builder(provider: &ProviderWithResource) -> Box Box::new(GoogleAIQueryBuilder::new()), AIProvider::OpenRouter => Box::new(OpenRouterQueryBuilder::new()), - _ => Box::new(OpenAIQueryBuilder::new()), // Use OpenAI as default for all other providers + _ => Box::new(OpenAIQueryBuilder::new(provider.kind.clone())), // Pass provider kind for Azure handling } } diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index ddc2cc480f..9d5ac8aa66 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -6,6 +6,7 @@ use ulid; use uuid::Uuid; use windmill_common::{ ai_providers::AIProvider, + ai_providers::AZURE_API_VERSION, cache, client::AuthedClient, db::DB, @@ -526,7 +527,7 @@ pub async fn run_agent( let endpoint = query_builder.get_endpoint(&base_url, args.provider.get_model(), output_type); - let auth_headers = query_builder.get_auth_headers(api_key, output_type); + let auth_headers = query_builder.get_auth_headers(api_key, &base_url, output_type); let mut request = HTTP_CLIENT .post(&endpoint) @@ -538,6 +539,10 @@ pub async fn run_agent( request = request.header(header_name, header_value); } + if args.provider.kind.is_azure_openai(&base_url) { + request = request.query(&[("api-version", AZURE_API_VERSION)]) + } + let resp = request .body(request_body) .send()