feat: make azure a standalone AI provider (#5558)

* feat: make azure a standalone AI provider

* oups

* nit

* fix: openai/azure oauth

* nit

* nits
This commit is contained in:
HugoCasa
2025-04-03 21:26:56 +02:00
committed by GitHub
parent efaaebd299
commit e8d3bc402d
16 changed files with 549 additions and 613 deletions
@@ -1,23 +0,0 @@
{
"db_name": "PostgreSQL",
"query": "SELECT value\n FROM resource\n WHERE path = $1 AND workspace_id = $2",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "value",
"type_info": "Jsonb"
}
],
"parameters": {
"Left": [
"Text",
"Text"
]
},
"nullable": [
true
]
},
"hash": "0a9dd1addaf48eeb46eed59abb6daf9819d07b08cf7ca442ea7ec78a9b2b63b2"
}
@@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT value as \"value: sqlx::types::Json<Box<RawValue>>\" FROM resource WHERE path = $1 AND workspace_id = $2",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "value: sqlx::types::Json<Box<RawValue>>",
"type_info": "Jsonb"
}
],
"parameters": {
"Left": [
"Text",
"Text"
]
},
"nullable": [
true
]
},
"hash": "30483ae46f6d0452126eb2cd07fc4d960961cc6ee61cf065113b7a48f97caecc"
}
@@ -1,23 +0,0 @@
{
"db_name": "PostgreSQL",
"query": "SELECT value FROM resource WHERE path = $1 AND workspace_id = $2",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "value",
"type_info": "Jsonb"
}
],
"parameters": {
"Left": [
"Text",
"Text"
]
},
"nullable": [
true
]
},
"hash": "63c16a4277983aaed0aed54972923919cee3cc444725ac6b7906922554bae800"
}
@@ -1,6 +1,6 @@
{
"db_name": "PostgreSQL",
"query": "SELECT value\n FROM global_settings\n WHERE name = 'openai_azure_base_path'",
"query": "SELECT value\n FROM global_settings\n WHERE name = 'openai_azure_base_path'",
"describe": {
"columns": [
{
@@ -16,5 +16,5 @@
false
]
},
"hash": "0cc3618495d5d024b2a173c58a3a8bb2a9d69b7b6e7ed6b0d0064fa2ce9c2e31"
"hash": "a21a16064b51580a8f5c2505cb0c701281dbfa94e40994fdd1cadc86a26c294e"
}
@@ -15,7 +15,7 @@
]
},
"nullable": [
true
null
]
},
"hash": "ddf2eccb78a310ed00c7d8b9c3f05d394a7cbcf0038c72a78add5c7b02ef5927"
+1 -1
View File
@@ -12497,7 +12497,7 @@ components:
AIProvider:
type: string
enum: [openai, anthropic, mistral, deepseek, googleai, groq, openrouter, togetherai, customai]
enum: [openai, azure_openai, anthropic, mistral, deepseek, googleai, groq, openrouter, togetherai, customai]
AIProviderModel:
type: object
+223 -325
View File
@@ -3,16 +3,12 @@ use crate::{
variables::get_variable_or_self,
};
use anyhow::Context;
use axum::{body::Bytes, extract::Path, response::IntoResponse, routing::post, Extension, Router};
use http::HeaderMap;
use lazy_static::lazy_static;
use openai::OpenaiCache;
use openai_api_compatible::OpenaiApiCompatibleCache;
use quick_cache::sync::Cache;
use reqwest::{Client, RequestBuilder};
use serde::{Deserialize, Serialize};
use serde_json::value::{RawValue, Value};
use serde_json::value::RawValue;
use std::collections::HashMap;
use windmill_audit::{audit_ee::audit_log, ActionKind};
use windmill_common::error::{to_anyhow, Error, Result};
@@ -22,175 +18,94 @@ lazy_static::lazy_static! {
.timeout(std::time::Duration::from_secs(60 * 5))
.user_agent("windmill/beta")
.build().unwrap();
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);
}
mod openai_api_compatible {
use super::*;
const AZURE_API_VERSION: &str = "2024-10-21";
const OPENAI_BASE_URL: &str = "https://api.openai.com/v1";
#[derive(Deserialize, Clone, Debug)]
pub struct OpenaiApiCompatibleCache {
pub base_url: String,
pub api_key: Option<String>,
}
#[derive(Deserialize, Debug)]
struct AIOAuthResource {
client_id: String,
client_secret: String,
token_url: String,
user: Option<String>,
}
impl OpenaiApiCompatibleCache {
pub fn prepare_request(self, path: &str, body: Bytes) -> Result<RequestBuilder> {
let url = format!("{}/{}", self.base_url, path);
#[derive(Deserialize, Debug)]
struct AIStandardResource {
#[serde(alias = "baseUrl")]
base_url: Option<String>,
#[serde(alias = "apiKey")]
api_key: Option<String>,
organization_id: Option<String>,
}
let mut request = HTTP_CLIENT
.post(url)
.header("content-type", "application/json")
.body(body);
#[derive(Deserialize, Debug)]
struct OAuthTokens {
access_token: String,
}
if let Some(api_key) = self.api_key {
request = request.header("Authorization", format!("Bearer {}", api_key));
}
#[derive(Deserialize, Debug)]
#[serde(untagged)]
enum AIResource {
OAuth(AIOAuthResource),
Standard(AIStandardResource),
}
Ok(request)
}
}
#[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>,
}
pub async fn get_cached_value(
impl AIRequestConfig {
pub async fn new(
provider: &AIProvider,
db: &DB,
w_id: &str,
resource: Value,
base_url: Option<String>,
) -> Result<KeyCache> {
let mut resource: OpenaiApiCompatibleCache = if let Some(base_url) = base_url {
let api_key = match resource {
Value::Object(mut obj) => obj
.remove("api_key")
.map(|v| serde_json::from_value::<String>(v.clone()).ok())
.flatten()
.or_else(|| {
obj.remove("apiKey")
.map(|v| serde_json::from_value::<String>(v.clone()).ok())
.flatten()
}),
_ => None,
};
OpenaiApiCompatibleCache { base_url, api_key }
} else {
serde_json::from_value(resource).with_context(|| "validating custom AI resource")?
resource: AIResource,
) -> Result<Self> {
let (api_key, access_token, organization_id, base_url, user) = match resource {
AIResource::Standard(resource) => {
let base_url = provider.get_base_url(resource.base_url, db).await?;
let api_key = if let Some(api_key) = resource.api_key {
Some(get_variable_or_self(api_key, db, w_id).await?)
} else {
None
};
let organization_id = if let Some(organization_id) = resource.organization_id {
Some(get_variable_or_self(organization_id, db, w_id).await?)
} else {
None
};
(api_key, None, organization_id, base_url, None)
}
AIResource::OAuth(resource) => {
let user = if let Some(user) = resource.user.clone() {
Some(get_variable_or_self(user, db, w_id).await?)
} else {
None
};
let token = Self::get_token_using_oauth(resource, db, w_id).await?;
let base_url = provider.get_base_url(None, db).await?;
(None, Some(token), None, base_url, user)
}
};
if let Some(api_key) = resource.api_key {
resource.api_key = Some(get_variable_or_self(api_key, db, w_id).await?);
}
Ok(KeyCache::OpenaiApiCompatible(resource))
}
}
mod openai {
use super::*;
const API_VERSION: &str = "2024-10-21";
#[derive(Deserialize, Debug)]
struct OpenaiResource {
api_key: String,
organization_id: Option<String>,
Ok(Self { base_url, organization_id, api_key, access_token, user })
}
#[derive(Deserialize, Debug)]
struct OpenaiClientCredentialsOauthResource {
client_id: String,
client_secret: String,
token_url: String,
user: Option<String>,
}
#[derive(Deserialize, Debug)]
#[serde(untagged, rename_all = "snake_case")]
enum OpenaiConfig {
Resource(OpenaiResource),
ClientCredentialsOauthResource(OpenaiClientCredentialsOauthResource),
}
lazy_static::lazy_static! {
pub static ref OPENAI_AZURE_BASE_PATH: Option<String> = std::env::var("OPENAI_AZURE_BASE_PATH").ok();
}
#[derive(Deserialize, Debug)]
struct OpenaiCredentials {
access_token: String,
}
#[derive(Clone, Debug, Deserialize)]
pub struct OpenaiCache {
api_key: String,
organization_id: Option<String>,
azure_base_path: Option<String>,
user: Option<String>,
}
impl OpenaiCache {
pub fn new(
api_key: String,
organization_id: Option<String>,
azure_base_path: Option<String>,
user: Option<String>,
) -> Self {
Self { api_key, organization_id, azure_base_path, user }
}
}
pub const BASE_URL: &str = "https://api.openai.com/v1";
impl OpenaiCache {
pub fn prepare_request(self, openai_path: &str, mut body: Bytes) -> Result<RequestBuilder> {
let OpenaiCache { api_key, azure_base_path, organization_id, user } = self;
if user.is_some() {
tracing::debug!("Adding user to request body");
let mut json_body: HashMap<String, Box<RawValue>> = serde_json::from_slice(&body)
.map_err(|e| {
Error::internal_err(format!("Failed to parse request body: {}", e))
})?;
let user_json_string = serde_json::Value::String(user.unwrap()).to_string(); // makes sure to escape characters
json_body.insert(
"user".to_string(),
RawValue::from_string(user_json_string)
.map_err(|e| Error::internal_err(format!("Failed to parse user: {}", e)))?,
);
body = serde_json::to_vec(&json_body)
.map_err(|e| {
Error::internal_err(format!("Failed to reserialize request body: {}", e))
})?
.into();
}
let base_url = if let Some(base_url) = azure_base_path {
base_url
} else {
BASE_URL.to_string()
};
let url = format!("{}/{}", base_url, openai_path);
let mut request = HTTP_CLIENT
.post(url)
.header("content-type", "application/json")
.body(body);
if base_url != BASE_URL {
request = request
.header("api-key", api_key)
.query(&[("api-version", API_VERSION)])
} else {
request = request.header("authorization", format!("Bearer {}", api_key))
}
if let Some(org_id) = organization_id {
request = request.header("OpenAI-Organization", org_id);
}
Ok(request)
}
}
async fn get_openai_key_using_credentials_flow(
mut resource: OpenaiClientCredentialsOauthResource,
async fn get_token_using_oauth(
mut resource: AIOAuthResource,
db: &DB,
w_id: &str,
) -> Result<String> {
@@ -199,115 +114,114 @@ mod openai {
resource.token_url = get_variable_or_self(resource.token_url, db, w_id).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 OpenAI credentials using credentials flow: {}",
"Failed to get access token using credentials flow: {}",
err
))
})?;
let response = response.json::<OpenaiCredentials>().await.map_err(|err| {
let response = response.json::<OAuthTokens>().await.map_err(|err| {
Error::internal_err(format!(
"Failed to parse OpenAI credentials from credentials flow: {}",
"Failed to parse access token from credentials flow: {}",
err
))
})?;
Ok(response.access_token)
}
pub async fn get_cached_value(db: &DB, w_id: &str, resource: Value) -> Result<KeyCache> {
let config = serde_json::from_value(resource)
.map_err(|e| Error::internal_err(format!("validating openai resource {e:#}")))?;
pub fn prepare_request(
self,
provider: &AIProvider,
path: &str,
body: Bytes,
) -> Result<RequestBuilder> {
let url = format!("{}/{}", self.base_url, path);
let mut user = None::<String>;
let mut resource = match config {
OpenaiConfig::Resource(resource) => {
tracing::debug!("Getting OpenAI key from static resource");
resource
}
OpenaiConfig::ClientCredentialsOauthResource(resource) => {
tracing::debug!("Getting OpenAI key with client credentials flow");
user = resource.user.clone();
let token = get_openai_key_using_credentials_flow(resource, db, w_id).await?;
OpenaiResource { api_key: token, organization_id: None }
}
};
resource.api_key = get_variable_or_self(resource.api_key, db, w_id).await?;
if let Some(organization_id) = resource.organization_id {
resource.organization_id = Some(get_variable_or_self(organization_id, db, w_id).await?);
}
if user.is_some() {
user = Some(get_variable_or_self(user.unwrap(), db, w_id).await?);
}
let azure_base_path = sqlx::query_scalar!(
"SELECT value
FROM global_settings
WHERE name = 'openai_azure_base_path'",
)
.fetch_optional(db)
.await?;
let azure_base_path = if let Some(azure_base_path) = azure_base_path {
Some(
serde_json::from_value::<String>(azure_base_path).map_err(|e| {
Error::internal_err(format!("validating openai azure base path {e:#}"))
})?,
)
let body = if let Some(user) = self.user {
Self::add_user_to_body(body, user)?
} else {
OPENAI_AZURE_BASE_PATH.clone()
body
};
let workspace_cache = OpenaiCache::new(
resource.api_key.clone(),
resource.organization_id.clone(),
azure_base_path.clone(),
user.clone(),
let is_azure = matches!(provider, AIProvider::OpenAI) && self.base_url != OPENAI_BASE_URL
|| matches!(provider, AIProvider::AzureOpenAI);
let mut request = HTTP_CLIENT
.post(url)
.header("content-type", "application/json")
.body(body);
if is_azure {
request = request.query(&[("api-version", AZURE_API_VERSION)])
}
if let Some(api_key) = self.api_key {
if is_azure {
request = request.header("api-key", api_key)
} else {
request = request.header("authorization", format!("Bearer {}", api_key))
}
}
if let Some(access_token) = self.access_token {
request = request.header("authorization", format!("Bearer {}", access_token))
}
if let Some(org_id) = self.organization_id {
request = request.header("OpenAI-Organization", org_id);
}
Ok(request)
}
fn add_user_to_body(body: Bytes, user: String) -> Result<Bytes> {
tracing::debug!("Adding user to request body");
let mut json_body: HashMap<String, Box<RawValue>> = serde_json::from_slice(&body)
.map_err(|e| Error::internal_err(format!("Failed to parse request body: {}", e)))?;
let user_json_string = serde_json::Value::String(user).to_string(); // makes sure to escape characters
json_body.insert(
"user".to_string(),
RawValue::from_string(user_json_string)
.map_err(|e| Error::internal_err(format!("Failed to parse user: {}", e)))?,
);
Ok(KeyCache::Openai(workspace_cache))
Ok(serde_json::to_vec(&json_body)
.map_err(|e| Error::internal_err(format!("Failed to reserialize request body: {}", e)))?
.into())
}
}
#[derive(Clone, Debug)]
pub enum KeyCache {
Openai(OpenaiCache),
OpenaiApiCompatible(OpenaiApiCompatibleCache),
pub struct ExpiringAIRequestConfig {
config: AIRequestConfig,
expires_at: std::time::Instant,
}
#[derive(Clone, Debug)]
pub struct AICache {
pub cached_key: KeyCache,
pub expires_at: std::time::Instant,
}
impl AICache {
pub fn new(cached_key: KeyCache) -> Self {
Self {
cached_key,
expires_at: std::time::Instant::now() + std::time::Duration::from_secs(60),
}
impl ExpiringAIRequestConfig {
fn new(config: AIRequestConfig) -> Self {
Self { config, expires_at: std::time::Instant::now() + std::time::Duration::from_secs(60) }
}
fn is_expired(&self) -> bool {
self.expires_at < std::time::Instant::now()
}
}
lazy_static! {
pub static ref AI_KEY_CACHE: Cache<(String, AIProvider), AICache> = Cache::new(500);
}
#[derive(Serialize, Deserialize, Debug, Eq, PartialEq, Hash, Clone)]
#[serde(rename_all = "lowercase")]
pub enum AIProvider {
OpenAI,
#[serde(rename = "azure_openai")]
AzureOpenAI,
Anthropic,
Mistral,
DeepSeek,
@@ -319,19 +233,48 @@ pub enum AIProvider {
}
impl AIProvider {
pub fn get_base_url(&self) -> Result<Option<String>> {
pub async fn get_base_url(&self, resource_base_url: Option<String>, db: &DB) -> Result<String> {
match self {
AIProvider::DeepSeek => Ok(Some("https://api.deepseek.com/v1".to_string())),
AIProvider::GoogleAI => Ok(Some(
"https://generativelanguage.googleapis.com/v1beta/openai".to_string(),
)),
AIProvider::Groq => Ok(Some("https://api.groq.com/openai/v1".to_string())),
AIProvider::OpenRouter => Ok(Some("https://openrouter.ai/api/v1".to_string())),
AIProvider::TogetherAI => Ok(Some("https://api.together.xyz/v1".to_string())),
AIProvider::Anthropic => Ok(Some("https://api.anthropic.com/v1".to_string())),
AIProvider::Mistral => Ok(Some("https://api.mistral.ai/v1".to_string())),
AIProvider::CustomAI => Ok(None),
AIProvider::OpenAI => Ok(Some(openai::BASE_URL.to_string())),
AIProvider::OpenAI => {
let azure_base_path = sqlx::query_scalar!(
"SELECT value
FROM global_settings
WHERE name = 'openai_azure_base_path'",
)
.fetch_optional(db)
.await?;
let azure_base_path = if let Some(azure_base_path) = azure_base_path {
Some(
serde_json::from_value::<String>(azure_base_path).map_err(|e| {
Error::internal_err(format!("validating openai azure base path {e:#}"))
})?,
)
} else {
OPENAI_AZURE_BASE_PATH.clone()
};
Ok(azure_base_path.unwrap_or(OPENAI_BASE_URL.to_string()))
}
AIProvider::DeepSeek => Ok("https://api.deepseek.com/v1".to_string()),
AIProvider::GoogleAI => {
Ok("https://generativelanguage.googleapis.com/v1beta/openai".to_string())
}
AIProvider::Groq => Ok("https://api.groq.com/openai/v1".to_string()),
AIProvider::OpenRouter => Ok("https://openrouter.ai/api/v1".to_string()),
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()),
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
)))
}
}
}
}
}
@@ -339,18 +282,9 @@ impl AIProvider {
impl TryFrom<&str> for AIProvider {
type Error = Error;
fn try_from(s: &str) -> Result<Self> {
match s {
"openai" => Ok(AIProvider::OpenAI),
"anthropic" => Ok(AIProvider::Anthropic),
"mistral" => Ok(AIProvider::Mistral),
"groq" => Ok(AIProvider::Groq),
"openrouter" => Ok(AIProvider::OpenRouter),
"togetherai" => Ok(AIProvider::TogetherAI),
"deepseek" => Ok(AIProvider::DeepSeek),
"googleai" => Ok(AIProvider::GoogleAI),
"customai" => Ok(AIProvider::CustomAI),
_ => Err(Error::BadRequest(format!("Invalid AI provider: {}", s))),
}
let s = serde_json::from_value::<AIProvider>(serde_json::Value::String(s.to_string()))
.map_err(|e| Error::BadRequest(format!("Invalid AI provider: {}", e)))?;
Ok(s)
}
}
@@ -407,11 +341,7 @@ async fn global_proxy(
return Err(Error::BadRequest("API key is required".to_string()));
};
let base_url = provider.get_base_url()?;
let Some(base_url) = base_url else {
return Err(Error::BadRequest("Provider is not supported".to_string()));
};
let base_url = provider.get_base_url(None, &db).await?;
let url = format!("{}/{}", base_url, ai_path);
@@ -439,7 +369,7 @@ async fn global_proxy(
if response.error_for_status_ref().is_err() {
let err_msg = response.text().await.unwrap_or("".to_string());
return Err(Error::AiError(err_msg));
return Err(Error::AIError(err_msg));
}
let status_code = response.status();
@@ -464,32 +394,19 @@ async fn proxy(
None => return Err(Error::BadRequest("Provider is required".to_string())),
};
let workspace_cache = AI_KEY_CACHE.get(&(w_id.clone(), provider.clone()));
let workspace_cache = AI_REQUEST_CACHE.get(&(w_id.clone(), provider.clone()));
let forced_resource_path = headers
.get("X-Resource-Path")
.map(|v| v.to_str().unwrap_or("").to_string());
let ai_cache = match workspace_cache {
Some(cache) if !cache.is_expired() && forced_resource_path.is_none() => cache.cached_key,
let request_config = match workspace_cache {
Some(request_cache) if !request_cache.is_expired() && forced_resource_path.is_none() => {
request_cache.config
}
_ => {
let (resource, ai_provider, save_to_cache) = if let Some(resource_path) =
forced_resource_path
{
// forced resource path, get the resource directly
let resource = sqlx::query_scalar!(
"SELECT value FROM resource WHERE path = $1 AND workspace_id = $2",
&resource_path,
&w_id
)
.fetch_optional(&db)
.await?
.ok_or_else(|| {
Error::NotFound(format!(
"Could not find the resource {}, update the resource path in the workspace settings", resource_path
))
})?;
(resource, provider, false)
let (resource_path, save_to_cache) = if let Some(resource_path) = forced_resource_path {
// forced resource path
(resource_path, false)
} else {
let ai_config = sqlx::query_scalar!(
"SELECT ai_config FROM workspace_settings WHERE workspace_id = $1",
@@ -504,13 +421,13 @@ async fn proxy(
));
}
let ai_config = serde_json::from_value::<AIConfig>(ai_config.unwrap())
let mut ai_config = serde_json::from_value::<AIConfig>(ai_config.unwrap())
.map_err(|e| Error::BadRequest(e.to_string()))?;
let provider_config = ai_config
.providers
.as_ref()
.map(|providers| providers.get(&provider))
.as_mut()
.map(|providers| providers.remove(&provider))
.flatten()
.ok_or_else(|| {
Error::BadRequest(format!("Provider {:?} not configured", provider))
@@ -519,56 +436,37 @@ async fn proxy(
if provider_config.resource_path.is_empty() {
return Err(Error::BadRequest("Resource path is empty".to_string()));
}
let resource = sqlx::query_scalar!(
"SELECT value
FROM resource
WHERE path = $1 AND workspace_id = $2",
&provider_config.resource_path,
&w_id
)
.fetch_optional(&db)
.await?
.ok_or_else(|| {
Error::NotFound(format!(
"Could not find the {:?} resource at path {}, update the resource path in the workspace settings", provider, provider_config.resource_path
))
})?;
(resource, provider, true)
(provider_config.resource_path, true)
};
let Some(resource) = resource else {
return Err(Error::internal_err(format!(
"{:?} resource missing value",
ai_provider
)));
};
let resource= sqlx::query_scalar!(
"SELECT value as \"value: sqlx::types::Json<Box<RawValue>>\" FROM resource WHERE path = $1 AND workspace_id = $2",
&resource_path,
&w_id
)
.fetch_optional(&db)
.await?
.ok_or_else(|| Error::NotFound(format!("Could not find the resource {}, update the resource path in the workspace settings", resource_path)))?
.ok_or_else(|| Error::BadRequest(format!("Empty resource value for {}", resource_path)))?;
let ai_cache = match ai_provider {
AIProvider::OpenAI => openai::get_cached_value(&db, &w_id, resource).await?,
_ => {
openai_api_compatible::get_cached_value(
&db,
&w_id,
resource,
ai_provider.get_base_url()?,
)
.await?
}
};
let resource = serde_json::from_str::<AIResource>(resource.0.get())
.map_err(|e| Error::BadRequest(e.to_string()))?;
let request_config = AIRequestConfig::new(&provider, &db, &w_id, resource).await?;
if save_to_cache {
AI_KEY_CACHE.insert((w_id.clone(), ai_provider), AICache::new(ai_cache.clone()));
AI_REQUEST_CACHE.insert(
(w_id.clone(), provider.clone()),
ExpiringAIRequestConfig::new(request_config.clone()),
);
}
ai_cache
request_config
}
};
let request = match ai_cache {
KeyCache::Openai(cached) => cached.prepare_request(&ai_path, body),
KeyCache::OpenaiApiCompatible(cached) => cached.prepare_request(&ai_path, body),
};
let request = request_config.prepare_request(&provider, &ai_path, body)?;
let response = request?.send().await.map_err(to_anyhow)?;
let response = request.send().await.map_err(to_anyhow)?;
let mut tx = db.begin().await?;
@@ -586,7 +484,7 @@ async fn proxy(
if response.error_for_status_ref().is_err() {
let err_msg = response.text().await.unwrap_or("".to_string());
return Err(Error::AiError(err_msg));
return Err(Error::AIError(err_msg));
}
let status_code = response.status();
+2 -2
View File
@@ -8,7 +8,7 @@
use std::collections::HashMap;
use crate::ai::{AIConfig, AI_KEY_CACHE};
use crate::ai::{AIConfig, AI_REQUEST_CACHE};
use crate::db::ApiAuthed;
use crate::users_ee::send_email_if_possible;
use crate::utils::get_instance_username_or_create_pending;
@@ -704,7 +704,7 @@ async fn edit_copilot_config(
if let Some(ref providers) = ai_config.providers {
for provider in providers.keys() {
AI_KEY_CACHE.remove(&(w_id.clone(), provider.clone()));
AI_REQUEST_CACHE.remove(&(w_id.clone(), provider.clone()));
}
}
+2 -2
View File
@@ -69,7 +69,7 @@ pub enum Error {
#[error("Error: {0:#?}")]
JsonErr(serde_json::Value),
#[error("{0}")]
AiError(String),
AIError(String),
#[error("{0}")]
AlreadyCompleted(String),
#[error("Find python error: {0}")]
@@ -179,7 +179,7 @@ impl IntoResponse for Error {
Self::RequireAdmin(_) => axum::http::StatusCode::FORBIDDEN,
Self::SqlErr { .. }
| Self::BadRequest(_)
| Self::AiError(_)
| Self::AIError(_)
| Self::QuotaExceeded(_) => axum::http::StatusCode::BAD_REQUEST,
_ => axum::http::StatusCode::INTERNAL_SERVER_ERROR,
};
@@ -438,7 +438,7 @@ export class Autocompletor {
endColumn: 10000
})
const returnedCode = await autocompleteRequest(
let returnedCode = await autocompleteRequest(
{
prefix,
modifiablePrefix,
@@ -454,6 +454,8 @@ export class Autocompletor {
return
}
returnedCode = returnedCode.replace('<CURSOR>', '')
const editableCode = model.getValueInRange({
startLineNumber: modifiableStart,
startColumn: 1,
@@ -1,8 +1,8 @@
<script lang="ts">
import { Button } from '$lib/components/common'
import { getAstNode, type HastNode } from 'svelte-exmarkdown'
import { getAstNode } from 'svelte-exmarkdown'
import { editor as meditor } from 'monaco-editor'
import { getContext } from 'svelte'
import { getContext, untrack } from 'svelte'
import { Loader2 } from 'lucide-svelte'
import { initializeVscode } from '$lib/components/vscode'
import type { AIChatContext, ContextElement, DisplayMessage } from './core'
@@ -102,12 +102,13 @@
)
let loading = $state(true)
function shouldStopLoading(astNode: HastNode, replying: boolean) {
if (!replying || $currentReply.length > (astNode.position?.end.offset ?? 0)) {
$effect(() => {
// we only want to trigger when astNode offset is updated not currentReply, otherwise as there is some delay on the offset update, loading would be set to false too early
const completeReply = untrack(() => $currentReply)
if (!$loadingContext || completeReply.length > (astNode.current.position?.end.offset ?? 0)) {
loading = false
}
}
$effect(() => shouldStopLoading(astNode.current, $loadingContext))
})
let diffEl: HTMLDivElement | undefined = $state()
let diffEditor: meditor.IStandaloneDiffEditor | undefined = $state()
@@ -24,6 +24,7 @@ export const SUPPORTED_LANGUAGES = new Set(Object.keys(GEN_CONFIG.prompts))
// need at least one model for each provider except customai
export const AI_DEFAULT_MODELS: Record<AIProvider, string[]> = {
openai: ['gpt-4o', 'gpt-4o-mini'],
azure_openai: ['gpt-4o', 'gpt-4o-mini'],
anthropic: [
'claude-3-7-sonnet-latest',
'claude-3-7-sonnet-latest/thinking',
@@ -81,6 +82,7 @@ const DEFAULT_COMPLETION_CONFIG: ChatCompletionCreateParams = {
export const PROVIDER_COMPLETION_CONFIG_MAP: Record<AIProvider, ChatCompletionCreateParams> = {
openai: DEFAULT_COMPLETION_CONFIG,
azure_openai: DEFAULT_COMPLETION_CONFIG,
groq: DEFAULT_COMPLETION_CONFIG,
openrouter: DEFAULT_COMPLETION_CONFIG,
togetherai: DEFAULT_COMPLETION_CONFIG,
@@ -0,0 +1,235 @@
<script lang="ts">
import { WorkspaceService, type AIConfig, type AIProvider } from '$lib/gen'
import { setCopilotInfo, workspaceStore } from '$lib/stores'
import { sendUserToast } from '$lib/toast'
import MultiSelect from 'svelte-multiselect'
import { AI_DEFAULT_MODELS } from '../copilot/lib'
import TestAiKey from '../copilot/TestAIKey.svelte'
import Description from '../Description.svelte'
import Label from '../Label.svelte'
import ResourcePicker from '../ResourcePicker.svelte'
import Toggle from '../Toggle.svelte'
import ArgEnum from '../ArgEnum.svelte'
import Button from '../common/button/Button.svelte'
const aiProviderLabels: [AIProvider, string][] = [
['openai', 'OpenAI'],
['azure_openai', 'Azure OpenAI'],
['anthropic', 'Anthropic'],
['mistral', 'Mistral'],
['deepseek', 'DeepSeek'],
['googleai', 'Google AI'],
['groq', 'Groq'],
['openrouter', 'OpenRouter'],
['togetherai', 'Together AI'],
['customai', 'Custom AI']
]
let {
aiProviders = $bindable(),
codeCompletionModel = $bindable(),
defaultModel = $bindable(),
usingOpenaiClientCredentialsOauth = $bindable()
}: {
aiProviders: Exclude<AIConfig['providers'], undefined>
codeCompletionModel: string | undefined
defaultModel: string | undefined
usingOpenaiClientCredentialsOauth: boolean
} = $props()
let availableAIModels = $derived(Object.values(aiProviders).flatMap((p) => p.models))
let modelProviderMap = $derived(
Object.fromEntries(
Object.entries(aiProviders).flatMap(([provider, config]) =>
config.models.map((m) => [m, provider as AIProvider])
)
)
)
$effect(() => {
if (Object.keys(aiProviders).length < 1) {
codeCompletionModel = undefined
defaultModel = undefined
}
})
async function editCopilotConfig(): Promise<void> {
if (Object.keys(aiProviders ?? {}).length > 0) {
const code_completion_model = codeCompletionModel
? { model: codeCompletionModel, provider: modelProviderMap[codeCompletionModel] }
: undefined
const default_model = defaultModel
? { model: defaultModel, provider: modelProviderMap[defaultModel] }
: undefined
await WorkspaceService.editCopilotConfig({
workspace: $workspaceStore!,
requestBody: {
providers: aiProviders,
code_completion_model,
default_model
}
})
setCopilotInfo({
providers: aiProviders,
code_completion_model,
default_model
})
} else {
await WorkspaceService.editCopilotConfig({
workspace: $workspaceStore!,
requestBody: {}
})
setCopilotInfo({})
}
sendUserToast(`Copilot settings updated`)
}
</script>
<div class="flex flex-col gap-4 my-8">
<div class="flex flex-col gap-1">
<div class="text-primary text-lg font-semibold"> Windmill AI</div>
<Description link="https://www.windmill.dev/docs/core_concepts/ai_generation">
Windmill AI integrates with your favorite AI providers and models.
</Description>
</div>
</div>
<div class="flex flex-col gap-8">
<div class="flex flex-col gap-2">
<p class="font-semibold">AI Providers</p>
<div class="flex flex-col gap-4">
{#each aiProviderLabels as [provider, label]}
<div class="flex flex-col gap-2">
<Toggle
options={{
right: label
}}
checked={!!aiProviders[provider]}
on:change={(e) => {
if (e.detail) {
aiProviders[provider] = {
resource_path: '',
models:
AI_DEFAULT_MODELS[provider].length > 0 ? [AI_DEFAULT_MODELS[provider][0]] : []
}
if (AI_DEFAULT_MODELS[provider].length > 0 && !defaultModel) {
defaultModel = AI_DEFAULT_MODELS[provider][0]
}
} else {
aiProviders = Object.fromEntries(
Object.entries(aiProviders).filter(([key]) => key !== provider)
)
}
}}
/>
{#if aiProviders[provider]}
<div class="mb-4 flex flex-col gap-2">
<div class="flex flex-row gap-1">
{#key aiProviders[provider].resource_path}
<!-- this can be removed once the parent component moves to runes -->
<!-- svelte-ignore binding_property_non_reactive -->
<ResourcePicker
resourceType={provider === 'openai' && usingOpenaiClientCredentialsOauth
? 'openai_client_credentials_oauth'
: provider}
initialValue={aiProviders[provider].resource_path}
bind:value={aiProviders[provider].resource_path}
on:change={() => {
if (
aiProviders[provider].resource_path &&
aiProviders[provider].models.length === 0 &&
AI_DEFAULT_MODELS[provider].length > 0
) {
aiProviders[provider].models = AI_DEFAULT_MODELS[provider].slice(0, 1)
}
}}
/>
{/key}
<TestAiKey
aiProvider={provider}
resourcePath={aiProviders[provider].resource_path}
model={aiProviders[provider].models[0]}
/>
</div>
<Label label="Enabled models">
<!-- this can be removed once the parent component moves to runes -->
<!-- svelte-ignore binding_property_non_reactive -->
<MultiSelect
options={AI_DEFAULT_MODELS[provider]}
ulOptionsClass={'!bg-surface-secondary'}
allowUserOptions="append"
bind:selected={aiProviders[provider].models}
/>
</Label>
</div>
{/if}
</div>
{/each}
</div>
</div>
{#if Object.keys(aiProviders).length > 0}
<div class="flex flex-col gap-2">
<p class="font-semibold">Settings</p>
<div class="flex flex-col gap-4">
<Label label="Default chat model">
<ArgEnum
enum_={availableAIModels}
bind:value={defaultModel}
disabled={false}
autofocus={false}
defaultValue={undefined}
valid={true}
create={false}
required={false}
/>
</Label>
<div class="flex flex-col gap-2">
<Toggle
on:change={(e) => {
if (e.detail) {
codeCompletionModel = ''
} else {
codeCompletionModel = undefined
}
}}
checked={codeCompletionModel != undefined}
options={{
right: 'Code completion'
}}
/>
{#if codeCompletionModel != undefined}
<Label label="Code completion model">
<ArgEnum
enum_={availableAIModels}
bind:value={codeCompletionModel}
disabled={false}
autofocus={false}
defaultValue={undefined}
valid={true}
create={false}
required={false}
/>
<p class="text-xs">
We highly recommend using Mistral's Codestral model for code completion.
</p>
</Label>
{/if}
</div>
</div>
</div>
{/if}
<Button
wrapperClasses="self-start"
disabled={!Object.values(aiProviders).every((p) => p.resource_path) ||
(codeCompletionModel != undefined && codeCompletionModel.length === 0) ||
(Object.keys(aiProviders).length > 0 && !defaultModel)}
on:click={editCopilotConfig}
>
Save
</Button>
</div>
+25 -5
View File
@@ -102,6 +102,17 @@ export function setCopilotInfo(aiConfig: AIConfig) {
([provider, providerConfig]) =>
providerConfig.models.map((m) => ({ model: m, provider: provider as AIProvider }))
)
copilotSessionModel.update((model) => {
if (
model &&
!aiModels.some((m) => m.model === model.model && m.provider === model.provider)
) {
return undefined
}
return model
})
copilotInfo.set({
enabled: true,
codeCompletionModel: aiConfig.code_completion_model,
@@ -109,6 +120,8 @@ export function setCopilotInfo(aiConfig: AIConfig) {
aiModels: aiModels
})
} else {
copilotSessionModel.set(undefined)
copilotInfo.set({
enabled: false,
codeCompletionModel: undefined,
@@ -139,9 +152,9 @@ const sessionProvider = getLocalSetting(COPILOT_SESSION_PROVIDER_SETTING_NAME)
export const copilotSessionModel = writable<AIProviderModel | undefined>(
sessionModel && sessionProvider
? {
model: sessionModel,
provider: sessionProvider as AIProvider
}
model: sessionModel,
provider: sessionProvider as AIProvider
}
: undefined
)
export const usedTriggerKinds = writable<string[]>([])
@@ -158,10 +171,17 @@ type SQLBaseSchema = {
}
}
export const SQLSchemaLanguages = ['mysql', 'bigquery', 'postgresql', 'snowflake', 'mssql', 'oracledb'] as const
export const SQLSchemaLanguages = [
'mysql',
'bigquery',
'postgresql',
'snowflake',
'mssql',
'oracledb'
] as const
export interface SQLSchema {
lang: typeof SQLSchemaLanguages[number]
lang: (typeof SQLSchemaLanguages)[number]
schema: SQLBaseSchema
publicOnly: boolean | undefined
stringified: string
@@ -265,9 +265,6 @@
<ToggleButton value="anthropic" label="Anthropic" {item} />
<ToggleButton value="mistral" label="Mistral" {item} />
<ToggleButton value="deepseek" label="DeepSeek" {item} />
<ToggleButton value="groq" label="Groq" {item} />
<ToggleButton value="openrouter" label="OpenRouter" {item} />
<ToggleButton value="togetherai" label="Together AI" {item} />
</ToggleButtonGroup>
</div>
</label>
@@ -21,7 +21,6 @@
JobService,
ResourceService,
SettingService,
type AIProvider,
type AIConfig
} from '$lib/gen'
import {
@@ -30,8 +29,7 @@
userStore,
usersWorkspaceStore,
workspaceStore,
isCriticalAlertsUIOpen,
setCopilotInfo
isCriticalAlertsUIOpen
} from '$lib/stores'
import { sendUserToast } from '$lib/toast'
import { emptyString, tryEvery } from '$lib/utils'
@@ -48,7 +46,6 @@
import PremiumInfo from '$lib/components/settings/PremiumInfo.svelte'
import Toggle from '$lib/components/Toggle.svelte'
import TestAIKey from '$lib/components/copilot/TestAIKey.svelte'
import Portal from '$lib/components/Portal.svelte'
import { fade } from 'svelte/transition'
@@ -62,12 +59,9 @@
} from '$lib/workspace_settings'
import { base } from '$lib/base'
import { hubPaths } from '$lib/hub'
import { AI_DEFAULT_MODELS } from '$lib/components/copilot/lib'
import Description from '$lib/components/Description.svelte'
import ConnectionSection from '$lib/components/ConnectionSection.svelte'
import MultiSelect from 'svelte-multiselect'
import Label from '$lib/components/Label.svelte'
import ArgEnum from '$lib/components/ArgEnum.svelte'
import AiSettings from '$lib/components/workspaceSettings/AISettings.svelte'
type GitSyncTypeMap = {
scripts: boolean
@@ -118,31 +112,10 @@
let criticalAlertUIMuted: boolean | undefined = undefined
let initialCriticalAlertUIMuted: boolean | undefined = undefined
const aiProviderLabels: [AIProvider, string][] = [
['openai', 'OpenAI'],
['anthropic', 'Anthropic'],
['mistral', 'Mistral'],
['deepseek', 'DeepSeek'],
['googleai', 'Google AI'],
['groq', 'Groq'],
['openrouter', 'OpenRouter'],
['togetherai', 'Together AI'],
['customai', 'Custom AI']
]
let aiProviders: Exclude<AIConfig['providers'], undefined> = {}
let codeCompletionModel: string | undefined = undefined
let defaultModel: string | undefined = undefined
$: availableAIModels = Object.values(aiProviders).flatMap((p) => p.models)
$: modelProviderMap = Object.fromEntries(
Object.entries(aiProviders).flatMap(([provider, config]) =>
config.models.map((m) => [m, provider as AIProvider])
)
)
$: Object.keys(aiProviders).length < 1 &&
((codeCompletionModel = undefined), (defaultModel = undefined))
let s3ResourceSettings: S3ResourceSettings = {
resourceType: 's3',
resourcePath: undefined,
@@ -239,37 +212,6 @@
}
}
async function editCopilotConfig(): Promise<void> {
if (Object.keys(aiProviders ?? {}).length > 0) {
const code_completion_model = codeCompletionModel
? { model: codeCompletionModel, provider: modelProviderMap[codeCompletionModel] }
: undefined
const default_model = defaultModel
? { model: defaultModel, provider: modelProviderMap[defaultModel] }
: undefined
await WorkspaceService.editCopilotConfig({
workspace: $workspaceStore!,
requestBody: {
providers: aiProviders,
code_completion_model,
default_model
}
})
setCopilotInfo({
providers: aiProviders,
code_completion_model,
default_model
})
} else {
await WorkspaceService.editCopilotConfig({
workspace: $workspaceStore!,
requestBody: {}
})
setCopilotInfo({})
}
sendUserToast(`Copilot settings updated`)
}
async function editWindmillLFSSettings(): Promise<void> {
const large_file_storage = convertFrontendToBackendSetting(s3ResourceSettings)
await WorkspaceService.editLargeFileStorageConfig({
@@ -419,9 +361,12 @@
await loadWorkspaceEncryptionKey()
const timeEnd = new Date().getTime()
sendUserToast('All workspace secrets have been re-encrypted with the new key')
setTimeout(() => {
workspaceReencryptionInProgress = false
}, 1000 - (timeEnd - timeStart))
setTimeout(
() => {
workspaceReencryptionInProgress = false
},
1000 - (timeEnd - timeStart)
)
}
let loadedSettings = false
@@ -461,11 +406,11 @@
errorHandlerSelected = emptyString(errorHandlerScriptPath)
? 'custom'
: errorHandlerScriptPath.startsWith('hub/') &&
errorHandlerScriptPath.endsWith('/workspace-or-schedule-error-handler-slack')
? 'slack'
: errorHandlerScriptPath.endsWith('/workspace-or-schedule-error-handler-teams')
? 'teams'
: 'custom'
errorHandlerScriptPath.endsWith('/workspace-or-schedule-error-handler-slack')
? 'slack'
: errorHandlerScriptPath.endsWith('/workspace-or-schedule-error-handler-teams')
? 'teams'
: 'custom'
}
errorHandlerExtraArgs = settings.error_handler_extra_args ?? {}
workspaceDefaultAppPath = settings.default_app
@@ -476,8 +421,8 @@
gitSyncTestJobs = []
gitSyncSettings = {
include_path:
settings.git_sync.include_path?.length ?? 0 > 0
? settings.git_sync.include_path ?? []
(settings.git_sync.include_path?.length ?? 0 > 0)
? (settings.git_sync.include_path ?? [])
: ['f/**'],
repositories: (settings.git_sync.repositories ?? []).map((settings) => {
gitSyncTestJobs.push({
@@ -541,8 +486,8 @@
if (settings.deploy_ui != undefined && settings.deploy_ui != null) {
deployUiSettings = {
include_path:
settings.deploy_ui.include_path?.length ?? 0 > 0
? settings.deploy_ui.include_path ?? []
(settings.deploy_ui.include_path?.length ?? 0 > 0)
? (settings.deploy_ui.include_path ?? [])
: [],
include_type: {
scripts: (settings.deploy_ui.include_type?.indexOf('script') ?? -1) >= 0,
@@ -1064,153 +1009,12 @@
</div>
</div>
{:else if tab == 'ai'}
<div class="flex flex-col gap-4 my-8">
<div class="flex flex-col gap-1">
<div class="text-primary text-lg font-semibold"> Windmill AI</div>
<Description link="https://www.windmill.dev/docs/core_concepts/ai_generation">
Windmill AI integrates with your favorite AI providers and models.
</Description>
</div>
</div>
<div class="flex flex-col gap-8">
<div class="flex flex-col gap-2">
<p class="font-semibold">AI Providers</p>
<div class="flex flex-col gap-4">
{#each aiProviderLabels as [provider, label]}
<div class="flex flex-col gap-2">
<Toggle
options={{
right: label
}}
checked={!!aiProviders[provider]}
on:change={(e) => {
if (e.detail) {
aiProviders[provider] = {
resource_path: '',
models:
AI_DEFAULT_MODELS[provider].length > 0
? [AI_DEFAULT_MODELS[provider][0]]
: []
}
if (AI_DEFAULT_MODELS[provider].length > 0 && !defaultModel) {
defaultModel = AI_DEFAULT_MODELS[provider][0]
}
} else {
aiProviders = Object.fromEntries(
Object.entries(aiProviders).filter(([key]) => key !== provider)
)
}
}}
/>
{#if aiProviders[provider]}
<div class="mb-4 flex flex-col gap-2">
<div class="flex flex-row gap-1">
{#key aiProviders[provider].resource_path}
<ResourcePicker
resourceType={usingOpenaiClientCredentialsOauth
? 'openai_client_credentials_oauth'
: provider}
initialValue={aiProviders[provider].resource_path}
bind:value={aiProviders[provider].resource_path}
on:change={() => {
if (
aiProviders[provider].resource_path &&
aiProviders[provider].models.length === 0 &&
AI_DEFAULT_MODELS[provider].length > 0
) {
aiProviders[provider].models = AI_DEFAULT_MODELS[provider].slice(0, 1)
}
}}
/>
{/key}
<TestAIKey
aiProvider={provider}
resourcePath={aiProviders[provider].resource_path}
model={aiProviders[provider].models[0]}
/>
</div>
<Label label="Enabled models">
<MultiSelect
options={AI_DEFAULT_MODELS[provider]}
ulOptionsClass={'!bg-surface-secondary'}
allowUserOptions="append"
bind:selected={aiProviders[provider].models}
/>
</Label>
</div>
{/if}
</div>
{/each}
</div>
</div>
{#if Object.keys(aiProviders).length > 0}
<div class="flex flex-col gap-2">
<p class="font-semibold">Settings</p>
<div class="flex flex-col gap-4">
<Label label="Default chat model">
<ArgEnum
enum_={availableAIModels}
bind:value={defaultModel}
disabled={false}
autofocus={false}
defaultValue={undefined}
valid={true}
create={false}
required={false}
/>
</Label>
<div class="flex flex-col gap-2">
<Toggle
on:change={(e) => {
if (e.detail) {
codeCompletionModel = ''
} else {
codeCompletionModel = undefined
}
}}
checked={codeCompletionModel != undefined}
options={{
right: 'Code completion'
}}
/>
{#if codeCompletionModel != undefined}
<Label label="Code completion model">
<ArgEnum
enum_={availableAIModels}
bind:value={codeCompletionModel}
disabled={false}
autofocus={false}
defaultValue={undefined}
valid={true}
create={false}
required={false}
/>
<p class="text-xs">
We highly recommend using Mistral's Codestral model for code completion.
</p>
</Label>
{/if}
</div>
</div>
</div>
{/if}
<Button
wrapperClasses="self-start"
disabled={!Object.values(aiProviders).every((p) => p.resource_path) ||
(codeCompletionModel != undefined && codeCompletionModel.length === 0) ||
(Object.keys(aiProviders).length > 0 && !defaultModel)}
on:click={editCopilotConfig}
>
Save
</Button>
</div>
<AiSettings
{aiProviders}
{codeCompletionModel}
{defaultModel}
{usingOpenaiClientCredentialsOauth}
/>
{:else if tab == 'windmill_lfs'}
<div class="flex flex-col gap-4 my-8">
<div class="flex flex-col gap-1">