fix(aiagent): fix endpoint for azure (#6633)

* fix ai agent for azure

* cleaning

* add parentheses

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<String> {
-        #[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)]

backend/windmill-common/src/ai_providers.rs

@@ -10,6 +10,9 @@ lazy_static::lazy_static! {
     static ref OPENAI_AZURE_BASE_PATH: Option<String> = 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<String> {
+        #[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 {

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 {

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<ResponseFormat>,
 }
 
-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))]
+        }
     }
 }

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))]

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<dyn QueryBui
     match provider.kind {
         AIProvider::GoogleAI => 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
     }
 }

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()