feat: Support mistral anthropic for ai (#4692)

* wip: openai proxy and other ai proxy integration * fixing migration script * wip: support different ai provider in front, fix proxy openai * wip: adding frontend ai provider * updated copilot types * wip: working on anthropic integration * done AI proxy front * adding new type and support for anthropic * updating gitignore * adding streaming response * added streaming prompt * push lib/gen * wip: fixing anthropic * anthropic fully supported * fix backend missing var error and fully support stream event for anthropic * remove gen directory * fixing openapi file * add support for mistral, and update create workspace components * remove deref.json * remove package-json * openapi * fix ui enable code * added utility function for init workspace ai provider * fix workspace switch bug * update anthropic property and fixed frontend error * fix workspace settings * update error message and fix typo migration file * chore: update openapi file * fix dev file * add .sqlx * all * update sqlx --------- Co-authored-by: dieriba <t.dieriba@gmail.com> Co-authored-by: Ruben Fiszel <ruben@windmill.dev> Co-authored-by: Ruben Fiszel <ruben@rubenfiszel.com>
2026-10-28 00:02:42 +00:00 · 2024-11-15 11:29:52 -05:00
parent 979baab7d3
commit f8c9237f8a
51 changed files with 1326 additions and 667 deletions
@@ -0,0 +1,507 @@
+use crate::{
+    db::{ApiAuthed, DB},
+    variables::decrypt,
+};
+use anthropic::AnthropicCache;
+use axum::{
+    body::Bytes,
+    extract::{Path, Query},
+    response::IntoResponse,
+    routing::post,
+    Extension, Router,
+};
+use lazy_static::lazy_static;
+use mistral::MistralCache;
+use openai::OpenaiCache;
+use quick_cache::sync::Cache;
+use reqwest::{Client, RequestBuilder};
+use serde::{Deserialize, Deserializer};
+use windmill_audit::audit_ee::audit_log;
+use windmill_audit::ActionKind;
+use windmill_common::error::{to_anyhow, Result};
+use windmill_common::variables::build_crypt;
+
+use windmill_common::error::Error;
+
+use serde_json::value::{RawValue, Value};
+use std::collections::HashMap;
+
+lazy_static::lazy_static! {
+    static ref HTTP_CLIENT: Client = reqwest::ClientBuilder::new()
+        .timeout(std::time::Duration::from_secs(60 * 5))
+        .user_agent("windmill/beta")
+        .build().unwrap();
+}
+
+trait AiRequest {
+    fn prepare_request(self, path: &str, body: Bytes) -> Result<RequestBuilder>;
+}
+
+mod openai {
+    use super::*;
+
+    use super::{get_variable_or_self, KeyCache};
+
+    const API_VERSION: &str = "2023-05-15";
+
+    #[derive(Deserialize, Debug)]
+    struct OpenaiResource {
+        api_key: String,
+        organization_id: Option<String>,
+    }
+
+    #[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 }
+        }
+    }
+
+    const BASE_URL: &str = "https://api.openai.com/v1";
+    impl AiRequest for OpenaiCache {
+        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::InternalErr(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::InternalErr(format!("Failed to parse user: {}", e)))?,
+                );
+
+                body = serde_json::to_vec(&json_body)
+                    .map_err(|e| {
+                        Error::InternalErr(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,
+        db: &DB,
+        w_id: &str,
+    ) -> Result<String> {
+        resource.client_id = get_variable_or_self(resource.client_id, db, w_id).await?;
+        resource.client_secret = get_variable_or_self(resource.client_secret, db, w_id).await?;
+        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");
+        let response = HTTP_CLIENT
+            .post(resource.token_url)
+            .form(&params)
+            .basic_auth(resource.client_id, Some(resource.client_secret))
+            .send()
+            .await
+            .map_err(|err| {
+                Error::InternalErr(format!(
+                    "Failed to get OpenAI credentials using credentials flow: {}",
+                    err
+                ))
+            })?;
+        let response = response.json::<OpenaiCredentials>().await.map_err(|err| {
+            Error::InternalErr(format!(
+                "Failed to parse OpenAI credentials 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::InternalErr(format!("validating openai resource {e:#}")))?;
+
+        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::InternalErr(format!("validating openai azure base path {e:#}"))
+                })?,
+            )
+        } else {
+            OPENAI_AZURE_BASE_PATH.clone()
+        };
+
+        let workspace_cache = OpenaiCache::new(
+            resource.api_key.clone(),
+            resource.organization_id.clone(),
+            azure_base_path.clone(),
+            user.clone(),
+        );
+        Ok(KeyCache::Openai(workspace_cache))
+    }
+}
+
+mod anthropic {
+
+    use super::*;
+
+    #[derive(Clone, Deserialize, Debug)]
+    pub struct AnthropicCache {
+        #[serde(rename = "apiKey")]
+        pub api_key: String,
+    }
+
+    const API_VERSION: &str = "2023-06-01";
+
+    impl AnthropicCache {
+        pub fn new(api_key: String) -> Self {
+            Self { api_key }
+        }
+    }
+
+    const BASE_URL: &str = "https://api.anthropic.com";
+    impl AiRequest for AnthropicCache {
+        fn prepare_request(self, anthropic_path: &str, body: Bytes) -> Result<RequestBuilder> {
+            let AnthropicCache { api_key } = self;
+            let url = format!("{}/{}", BASE_URL, anthropic_path);
+            let request = HTTP_CLIENT
+                .post(url)
+                .header("x-api-key", api_key)
+                .header("anthropic-version", API_VERSION)
+                .header("content-type", "application/json")
+                .body(body);
+            Ok(request)
+        }
+    }
+
+    pub async fn get_cached_value(db: &DB, w_id: &str, resource: Value) -> Result<KeyCache> {
+        let mut resource: AnthropicCache = serde_json::from_value(resource)
+            .map_err(|e| Error::InternalErr(format!("validating anthropic resource {e:#}")))?;
+        resource.api_key = get_variable_or_self(resource.api_key, db, w_id).await?;
+        let workspace_cache = AnthropicCache::new(resource.api_key);
+        Ok(KeyCache::Anthropic(workspace_cache))
+    }
+}
+
+mod mistral {
+    use super::*;
+    #[derive(Deserialize, Clone, Debug)]
+    pub struct MistralCache {
+        #[serde(rename = "apiKey")]
+        pub api_key: String,
+    }
+
+    impl MistralCache {
+        pub fn new(api_key: String) -> Self {
+            Self { api_key }
+        }
+    }
+
+    const BASE_URL: &str = "https://api.mistral.ai";
+    impl AiRequest for MistralCache {
+        fn prepare_request(self, mistral_path: &str, body: Bytes) -> Result<RequestBuilder> {
+            let MistralCache { api_key } = self;
+
+            let url = format!("{}/{}", BASE_URL, mistral_path);
+            let request = HTTP_CLIENT
+                .post(url)
+                .header("content-type", "application/json")
+                .header("Accept", "application/json")
+                .header("authorization", format!("Bearer {}", api_key))
+                .body(body);
+            Ok(request)
+        }
+    }
+
+    pub async fn get_cached_value(db: &DB, w_id: &str, resource: Value) -> Result<KeyCache> {
+        let mut resource: MistralCache = serde_json::from_value(resource)
+            .map_err(|e| Error::InternalErr(format!("validating mistral resource {e:#}")))?;
+        resource.api_key = get_variable_or_self(resource.api_key, db, w_id).await?;
+
+        let workspace_cache = MistralCache::new(resource.api_key);
+        Ok(KeyCache::Mistral(workspace_cache))
+    }
+}
+
+#[derive(Clone, Debug)]
+pub enum KeyCache {
+    Openai(OpenaiCache),
+    Anthropic(AnthropicCache),
+    Mistral(MistralCache),
+}
+
+#[derive(Clone, Debug)]
+pub struct AiCache {
+    pub path: String,
+    pub cached_key: KeyCache,
+    pub expires_at: std::time::Instant,
+}
+
+impl AiCache {
+    pub fn new(path: String, cached_key: KeyCache) -> Self {
+        Self {
+            path,
+            cached_key,
+            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, AiCache> = Cache::new(500);
+}
+
+struct Variable {
+    value: String,
+    is_secret: bool,
+}
+
+#[derive(Deserialize, Debug)]
+struct ProxyQueryParams {
+    no_cache: Option<bool>,
+}
+
+async fn get_variable_or_self(path: String, db: &DB, w_id: &str) -> Result<String> {
+    if !path.starts_with("$var:") {
+        return Ok(path);
+    }
+    let path = path.strip_prefix("$var:").unwrap().to_string();
+    let mut variable = sqlx::query_as!(
+        Variable,
+        "SELECT value, is_secret
+        FROM variable
+        WHERE path = $1 AND workspace_id = $2",
+        &path,
+        &w_id
+    )
+    .fetch_one(db)
+    .await?;
+    if variable.is_secret {
+        let mc = build_crypt(db, w_id).await?;
+        variable.value = decrypt(&mc, variable.value)?;
+    }
+    Ok(variable.value)
+}
+
+#[derive(Deserialize, Debug)]
+pub struct AiResource {
+    pub path: String,
+    #[serde(deserialize_with = "check_if_valid_ai_provider")]
+    pub provider: String,
+}
+
+fn check_if_valid_ai_provider<'de, D>(provider: D) -> std::result::Result<String, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    let provider = String::deserialize(provider)?;
+    match provider.as_str() {
+        "anthropic" | "openai" | "mistral" => Ok(provider),
+        _ => Err(serde::de::Error::custom(
+            "Only the following Ai providers are supported: openai, anthropic and mistral"
+                .to_string(),
+        )),
+    }
+}
+
+pub fn workspaced_service() -> Router {
+    let router = Router::new().route("/proxy/*ai", post(proxy));
+
+    router
+}
+
+async fn proxy(
+    authed: ApiAuthed,
+    Extension(db): Extension<DB>,
+    Path((w_id, ai_path)): Path<(String, String)>,
+    Query(query_params): Query<ProxyQueryParams>,
+    body: Bytes,
+) -> impl IntoResponse {
+    let workspace_cache = AI_KEY_CACHE.get(&w_id);
+    let ai_cache = match workspace_cache {
+        Some(cache) if !cache.is_expired() && !query_params.no_cache.unwrap_or(false) => {
+            cache.cached_key
+        }
+        _ => {
+            let ai_resource = sqlx::query_scalar!(
+                "SELECT ai_resource FROM workspace_settings WHERE workspace_id = $1",
+                &w_id
+            )
+            .fetch_one(&db)
+            .await?;
+
+            if ai_resource.is_none() {
+                return Err(Error::InternalErr("AI resource not configured".to_string()));
+            }
+
+            let ai_resource = serde_json::from_value::<AiResource>(ai_resource.unwrap())
+                .map_err(|e| Error::BadRequest(e.to_string()))?;
+            let ai_resource_path = ai_resource.path;
+
+            let resource = sqlx::query_scalar!(
+                "SELECT value
+                FROM resource
+                WHERE path = $1 AND workspace_id = $2",
+                &ai_resource_path,
+                &w_id
+            )
+            .fetch_optional(&db)
+            .await?
+            .ok_or_else(|| {
+                Error::InternalErr(format!(
+                    "Could not find the {} resource at path {ai_resource_path}, update the resource path in the workspace settings", ai_resource.provider
+                ))
+            })?;
+
+            if resource.is_none() {
+                return Err(Error::InternalErr(format!(
+                    "{} resource missing value",
+                    ai_resource.provider
+                )));
+            }
+
+            let resource = resource.unwrap();
+
+            let ai_cache = match ai_resource.provider.as_str() {
+                "openai" => openai::get_cached_value(&db, &w_id, resource).await,
+                "anthropic" => anthropic::get_cached_value(&db, &w_id, resource).await,
+                "mistral" => mistral::get_cached_value(&db, &w_id, resource).await,
+                provider => {
+                    return Err(Error::BadRequest(format!("{} is not supported", provider)))
+                }
+            };
+            let ai_cache = ai_cache?;
+            AI_KEY_CACHE.insert(
+                w_id.clone(),
+                AiCache::new(ai_resource_path, ai_cache.clone()),
+            );
+            ai_cache
+        }
+    };
+    let (path, request) = match ai_cache {
+        KeyCache::Openai(cached) => ("openai_path", cached.prepare_request(&ai_path, body)),
+        KeyCache::Anthropic(cached) => ("anthropic_path", cached.prepare_request(&ai_path, body)),
+        KeyCache::Mistral(cached) => ("mistral_path", cached.prepare_request(&ai_path, body)),
+    };
+
+    let response = request?.send().await.map_err(to_anyhow)?;
+
+    let mut tx = db.begin().await?;
+
+    audit_log(
+        &mut *tx,
+        &authed,
+        "ai.request",
+        ActionKind::Execute,
+        &w_id,
+        Some(&authed.email),
+        Some([(path, &format!("{:?}", ai_path)[..])].into()),
+    )
+    .await?;
+    tx.commit().await?;
+
+    if response.error_for_status_ref().is_err() {
+        let err_msg = response.text().await.unwrap_or("".to_string());
+        return Err(Error::AiError(err_msg));
+    }
+
+    let status_code = response.status();
+    let headers = response.headers().clone();
+    let stream = response.bytes_stream();
+    Ok((status_code, headers, axum::body::Body::from_stream(stream)))
+}
@@ -62,13 +62,14 @@ mod http_triggers;
 mod indexer_ee;
 mod inputs;
 mod integration;
+mod ai;
+
 #[cfg(feature = "parquet")]
 mod job_helpers_ee;
 pub mod job_metrics;
 pub mod jobs;
 pub mod oauth2_ee;
 mod oidc_ee;
-mod openai;
 mod raw_apps;
 mod resources;
 mod saml_ee;
@@ -94,6 +95,7 @@ mod workspaces_ee;

 pub const DEFAULT_BODY_LIMIT: usize = 2097152 * 100; // 200MB

+
 lazy_static::lazy_static! {

    pub static ref REQUEST_SIZE_LIMIT: Arc<RwLock<usize>> = Arc::new(RwLock::new(DEFAULT_BODY_LIMIT));
@@ -278,7 +280,7 @@ pub async fn run_server(
                        .nest("/job_helpers", job_helpers_service)
                        .nest("/jobs", jobs::workspaced_service())
                        .nest("/oauth", oauth2_ee::workspaced_service())
-                        .nest("/openai", openai::workspaced_service())
+                        .nest("/ai", ai::workspaced_service())
                        .nest("/raw_apps", raw_apps::workspaced_service())
                        .nest("/resources", resources::workspaced_service())
                        .nest("/schedules", schedule::workspaced_service())
@@ -1,344 +0,0 @@
-use std::collections::HashMap;
-
-use crate::db::{ApiAuthed, DB};
-
-use axum::{
-    body::Bytes,
-    extract::{Extension, Path, Query},
-    response::IntoResponse,
-    routing::post,
-    Router,
-};
-use quick_cache::sync::Cache;
-use reqwest::Client;
-use serde_json::value::RawValue;
-use windmill_audit::audit_ee::audit_log;
-use windmill_audit::ActionKind;
-use windmill_common::{
-    error::{to_anyhow, Error},
-    variables::build_crypt,
-};
-
-use crate::variables::decrypt;
-use serde::Deserialize;
-
-lazy_static::lazy_static! {
-    static ref HTTP_CLIENT: Client = reqwest::ClientBuilder::new()
-        .timeout(std::time::Duration::from_secs(60 * 5))
-        .user_agent("windmill/beta")
-        .build().unwrap();
-}
-
-pub fn workspaced_service() -> Router {
-    let router = Router::new().route("/proxy/*openai_path", post(proxy));
-
-    router
-}
-
-#[derive(Deserialize)]
-struct OpenaiResource {
-    api_key: String,
-    organization_id: Option<String>,
-}
-
-#[derive(Deserialize)]
-struct OpenaiClientCredentialsOauthResource {
-    client_id: String,
-    client_secret: String,
-    token_url: String,
-    user: Option<String>,
-}
-
-#[derive(Deserialize)]
-#[serde(untagged)]
-enum OpenaiConfig {
-    Resource(OpenaiResource),
-    ClientCredentialsOauthResource(OpenaiClientCredentialsOauthResource),
-}
-
-struct Variable {
-    value: String,
-    is_secret: bool,
-}
-async fn get_variable_or_self(path: String, db: &DB, w_id: &String) -> Result<String, Error> {
-    if !path.starts_with("$var:") {
-        return Ok(path);
-    }
-    let path = path.strip_prefix("$var:").unwrap().to_string();
-    let mut variable = sqlx::query_as!(
-        Variable,
-        "SELECT value, is_secret
-        FROM variable
-        WHERE path = $1 AND workspace_id = $2",
-        &path,
-        &w_id
-    )
-    .fetch_one(db)
-    .await?;
-    if variable.is_secret {
-        let mc = build_crypt(&db, &w_id).await?;
-        variable.value = decrypt(&mc, variable.value)?;
-    }
-    Ok(variable.value)
-}
-
-lazy_static::lazy_static! {
-    pub static ref OPENAI_AZURE_BASE_PATH: Option<String> = std::env::var("OPENAI_AZURE_BASE_PATH").ok();
-    static ref OPENAI_KEY_CACHE: Cache<String, OpenaiKeyCache> = Cache::new(500);
-}
-
-#[derive(Deserialize)]
-struct OpenaiCredentials {
-    access_token: String,
-}
-async fn get_openai_key_using_credentials_flow(
-    mut resource: OpenaiClientCredentialsOauthResource,
-    db: &DB,
-    w_id: &String,
-) -> Result<String, Error> {
-    resource.client_id = get_variable_or_self(resource.client_id, &db, &w_id).await?;
-    resource.client_secret = get_variable_or_self(resource.client_secret, &db, &w_id).await?;
-    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");
-    let response = HTTP_CLIENT
-        .post(resource.token_url)
-        .form(&params)
-        .basic_auth(resource.client_id, Some(resource.client_secret))
-        .send()
-        .await
-        .map_err(|err| {
-            Error::InternalErr(format!(
-                "Failed to get OpenAI credentials using credentials flow: {}",
-                err
-            ))
-        })?;
-    let response = response.json::<OpenaiCredentials>().await.map_err(|err| {
-        Error::InternalErr(format!(
-            "Failed to parse OpenAI credentials from credentials flow: {}",
-            err
-        ))
-    })?;
-    Ok(response.access_token)
-}
-
-#[derive(Clone)]
-struct OpenaiKeyCache {
-    api_key: String,
-    organization_id: Option<String>,
-    azure_base_path: Option<String>,
-    user: Option<String>,
-    expires_at: std::time::Instant,
-}
-
-impl OpenaiKeyCache {
-    fn new(
-        api_key: String,
-        organization_id: Option<String>,
-        azure_base_path: Option<String>,
-        expires_at: std::time::Instant,
-        user: Option<String>,
-    ) -> Self {
-        Self { api_key, organization_id, azure_base_path, expires_at, user }
-    }
-    fn is_expired(&self) -> bool {
-        self.expires_at < std::time::Instant::now()
-    }
-}
-
-#[derive(Deserialize)]
-struct ProxyQueryParams {
-    no_cache: Option<bool>,
-}
-async fn proxy(
-    authed: ApiAuthed,
-    Extension(db): Extension<DB>,
-    Path((w_id, openai_path)): Path<(String, String)>,
-    Query(query_params): Query<ProxyQueryParams>,
-    mut body: Bytes,
-) -> impl IntoResponse {
-    let workspace_cache = OPENAI_KEY_CACHE.get(&w_id);
-    let (api_key, organization_id, azure_base_path, user) = if query_params
-        .no_cache
-        .unwrap_or(false)
-        || workspace_cache.is_none()
-        || workspace_cache.clone().unwrap().is_expired()
-    {
-        let openai_resource_path = sqlx::query_scalar!(
-            "SELECT openai_resource_path FROM workspace_settings WHERE workspace_id = $1",
-            &w_id
-        )
-        .fetch_one(&db)
-        .await?;
-
-        if openai_resource_path.is_none() {
-            return Err(Error::InternalErr(
-                "OpenAI resource not configured".to_string(),
-            ));
-        }
-
-        let openai_resource_path = openai_resource_path.unwrap();
-
-        let resource = sqlx::query_scalar!(
-            "SELECT value
-            FROM resource
-            WHERE path = $1 AND workspace_id = $2",
-            &openai_resource_path,
-            &w_id
-        )
-        .fetch_optional(&db)
-        .await?
-        .ok_or_else(|| {
-            Error::InternalErr(format!(
-                "Could not find the OpenAI resource at path {openai_resource_path}, update the resource path in the workspace settings"
-            ))
-        })?;
-
-        if resource.is_none() {
-            return Err(Error::InternalErr(
-                "OpenAI resource missing value".to_string(),
-            ));
-        }
-
-        let config: OpenaiConfig = serde_json::from_value(resource.unwrap())
-            .map_err(|e| Error::InternalErr(format!("validating openai resource {e:#}")))?;
-
-        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 resource.organization_id.is_some() {
-            resource.organization_id =
-                Some(get_variable_or_self(resource.organization_id.unwrap(), &db, &w_id).await?);
-        }
-
-        if user.is_some() {
-            user = Some(get_variable_or_self(user.unwrap(), &db, &w_id).await?);
-        }
-
-        let expires_at = std::time::Instant::now() + std::time::Duration::from_secs(60);
-
-        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::InternalErr(format!("validating openai azure base path {e:#}"))
-                })?,
-            )
-        } else {
-            OPENAI_AZURE_BASE_PATH.clone()
-        };
-
-        let workspace_cache = OpenaiKeyCache::new(
-            resource.api_key.clone(),
-            resource.organization_id.clone(),
-            azure_base_path.clone(),
-            expires_at,
-            user.clone(),
-        );
-        OPENAI_KEY_CACHE.insert(w_id.clone(), workspace_cache);
-        (
-            resource.api_key,
-            resource.organization_id,
-            azure_base_path,
-            user,
-        )
-    } else {
-        tracing::debug!("Using cached OpenAI key");
-        let workspace_cache = workspace_cache.unwrap();
-        (
-            workspace_cache.api_key.clone(),
-            workspace_cache.organization_id.clone(),
-            workspace_cache.azure_base_path.clone(),
-            workspace_cache.user.clone(),
-        )
-    };
-
-    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::InternalErr(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::InternalErr(format!("Failed to parse user: {}", e)))?,
-        );
-
-        body = serde_json::to_vec(&json_body)
-            .map_err(|e| Error::InternalErr(format!("Failed to reserialize request body: {}", e)))?
-            .into();
-    }
-
-    let base_url = if let Some(base_url) = azure_base_path {
-        base_url
-    } else {
-        "https://api.openai.com/v1".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 != "https://api.openai.com/v1" {
-        request = request
-            .header("api-key", api_key)
-            .query(&[("api-version", "2023-05-15")])
-    } else {
-        request = request.header("authorization", format!("Bearer {}", api_key))
-    }
-
-    if let Some(org_id) = organization_id {
-        request = request.header("OpenAI-Organization", org_id);
-    }
-
-    let response = request.send().await.map_err(to_anyhow)?;
-
-    let mut tx = db.begin().await?;
-    audit_log(
-        &mut *tx,
-        &authed,
-        "openai.request",
-        ActionKind::Execute,
-        &w_id,
-        Some(&authed.email),
-        Some([("openai_path", &format!("{:?}", openai_path)[..])].into()),
-    )
-    .await?;
-    tx.commit().await?;
-
-    if response.error_for_status_ref().is_err() {
-        return Err(Error::OpenAIError(
-            response.text().await.unwrap_or("".to_string()),
-        ));
-    }
-
-    let status_code = response.status();
-    let headers = response.headers().clone();
-    let stream = response.bytes_stream();
-
-    Ok((status_code, headers, axum::body::Body::from_stream(stream)))
-}
@@ -8,6 +8,7 @@

 use std::collections::HashMap;

+use crate::ai::{AiResource, AI_KEY_CACHE};
 use crate::db::ApiAuthed;
 use crate::users_ee::send_email_if_possible;
 use crate::utils::get_instance_username_or_create_pending;
@@ -168,7 +169,7 @@ pub struct WorkspaceSettings {
    pub plan: Option<String>,
    pub webhook: Option<String>,
    pub deploy_to: Option<String>,
-    pub openai_resource_path: Option<String>,
+    pub ai_resource: Option<serde_json::Value>,
    pub code_completion_enabled: bool,
    pub error_handler: Option<String>,
    pub error_handler_extra_args: Option<serde_json::Value>,
@@ -234,7 +235,7 @@ struct EditWebhook {

 #[derive(Deserialize)]
 struct EditCopilotConfig {
-    openai_resource_path: Option<String>,
+    ai_resource: Option<serde_json::Value>,
    code_completion_enabled: bool,
 }

@@ -645,23 +646,33 @@ async fn edit_copilot_config(

    let mut tx = db.begin().await?;

-    if let Some(openai_resource_path) = &eo.openai_resource_path {
+    if let Some(ai_resource) = &eo.ai_resource {
+        let path = serde_json::from_value::<AiResource>(ai_resource.clone())
+            .map_err(|e| Error::BadRequest(e.to_string()))?
+            .path;
        sqlx::query!(
-            "UPDATE workspace_settings SET openai_resource_path = $1, code_completion_enabled = $2 WHERE workspace_id = $3",
-            openai_resource_path,
+            "UPDATE workspace_settings SET ai_resource = $1, code_completion_enabled = $2 WHERE workspace_id = $3",
+            ai_resource,
            eo.code_completion_enabled,
            &w_id
        )
        .execute(&mut *tx)
        .await?;
+
+        if let Some(cached) = AI_KEY_CACHE.get(&w_id) {
+            if cached.path != path {
+                AI_KEY_CACHE.remove(&w_id);
+            }
+        }
    } else {
        sqlx::query!(
-            "UPDATE workspace_settings SET openai_resource_path = NULL, code_completion_enabled = $1 WHERE workspace_id = $2",
+            "UPDATE workspace_settings SET ai_resource = NULL, code_completion_enabled = $1 WHERE workspace_id = $2",
            eo.code_completion_enabled,
            &w_id,
        )
        .execute(&mut *tx)
        .await?;
+        AI_KEY_CACHE.remove(&w_id);
    }
    audit_log(
        &mut *tx,
@@ -672,10 +683,7 @@ async fn edit_copilot_config(
        Some(&authed.email),
        Some(
            [
-                (
-                    "openai_resource_path",
-                    &format!("{:?}", eo.openai_resource_path)[..],
-                ),
+                ("ai_resource", &format!("{:?}", eo.ai_resource)[..]),
                (
                    "code_completion_enabled",
                    &format!("{:?}", eo.code_completion_enabled)[..],
@@ -692,7 +700,8 @@ async fn edit_copilot_config(

 #[derive(Serialize)]
 struct CopilotInfo {
-    pub exists_openai_resource_path: bool,
+    pub ai_provider: String,
+    pub exists_ai_resource: bool,
    pub code_completion_enabled: bool,
 }
 async fn get_copilot_info(
@@ -701,16 +710,32 @@ async fn get_copilot_info(
 ) -> JsonResult<CopilotInfo> {
    let mut tx = db.begin().await?;
    let record = sqlx::query!(
-        "SELECT openai_resource_path, code_completion_enabled FROM workspace_settings WHERE workspace_id = $1",
+        "SELECT ai_resource, code_completion_enabled FROM workspace_settings WHERE workspace_id = $1",
        &w_id
    )
    .fetch_one(&mut *tx)
    .await
-    .map_err(|e| Error::InternalErr(format!("getting openai_resource_path and code_completion_enabled: {e:#}")))?;
+    .map_err(|e| Error::InternalErr(format!("getting ai_resource and code_completion_enabled: {e:#}")))?;
    tx.commit().await?;

+    let (ai_provider, exists_ai_resource) = if let Some(ai_resource) = record.ai_resource {
+        let ai_resource = serde_json::from_value::<AiResource>(ai_resource);
+        let exist = ai_resource.is_ok();
+        (
+            if exist {
+                ai_resource.unwrap().provider
+            } else {
+                "".to_string()
+            },
+            exist,
+        )
+    } else {
+        ("".to_string(), false)
+    };
+
    Ok(Json(CopilotInfo {
-        exists_openai_resource_path: record.openai_resource_path.is_some(),
+        ai_provider,
+        exists_ai_resource,
        code_completion_enabled: record.code_completion_enabled,
    }))
 }
@@ -2223,7 +2248,7 @@ struct SimplifiedSettings {
    error_handler: Option<String>,
    error_handler_extra_args: Option<Value>,
    error_handler_muted_on_cancel: bool,
-    openai_resource_path: Option<String>,
+    ai_resource: Option<serde_json::Value>,
    code_completion_enabled: bool,
    large_file_storage: Option<Value>,
    git_sync: Option<Value>,
@@ -2588,7 +2613,7 @@ async fn tarball_workspace(
                webhook, 
                deploy_to, 
                error_handler, 
-                openai_resource_path, 
+                ai_resource, 
                code_completion_enabled, 
                error_handler_extra_args, 
                error_handler_muted_on_cancel,