From 4c7d64bc5fdfeb33c63524ddd61173c655ab28e3 Mon Sep 17 00:00:00 2001 From: HugoCasa Date: Tue, 12 Dec 2023 07:57:18 +0100 Subject: [PATCH] feat: oauth flow openai key (#2833) * feat: oauth flow openai key * fix: rawvalue insert --- ...a8bb2a9d69b7b6e7ed6b0d0064fa2ce9c2e31.json | 20 ++ ...6789887e47db9695fc26b1dfaedd735e1e60.json} | 4 +- backend/windmill-api/src/openai.rs | 299 ++++++++++++++---- frontend/src/lib/components/copilot/lib.ts | 8 +- .../src/lib/components/instanceSettings.ts | 10 + .../(logged)/workspace_settings/+page.svelte | 16 +- 6 files changed, 281 insertions(+), 76 deletions(-) create mode 100644 backend/.sqlx/query-0cc3618495d5d024b2a173c58a3a8bb2a9d69b7b6e7ed6b0d0064fa2ce9c2e31.json rename backend/.sqlx/{query-500393aa7c3f8c96ec8d92869c5c4744eebad8e9e31336c26cf51e678c3556f8.json => query-acdaa5151f8f7f37bb8c8c5a7d146789887e47db9695fc26b1dfaedd735e1e60.json} (61%) diff --git a/backend/.sqlx/query-0cc3618495d5d024b2a173c58a3a8bb2a9d69b7b6e7ed6b0d0064fa2ce9c2e31.json b/backend/.sqlx/query-0cc3618495d5d024b2a173c58a3a8bb2a9d69b7b6e7ed6b0d0064fa2ce9c2e31.json new file mode 100644 index 0000000000..79e36950f6 --- /dev/null +++ b/backend/.sqlx/query-0cc3618495d5d024b2a173c58a3a8bb2a9d69b7b6e7ed6b0d0064fa2ce9c2e31.json @@ -0,0 +1,20 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT value\n FROM global_settings\n WHERE name = 'openai_azure_base_path'", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "value", + "type_info": "Jsonb" + } + ], + "parameters": { + "Left": [] + }, + "nullable": [ + false + ] + }, + "hash": "0cc3618495d5d024b2a173c58a3a8bb2a9d69b7b6e7ed6b0d0064fa2ce9c2e31" +} diff --git a/backend/.sqlx/query-500393aa7c3f8c96ec8d92869c5c4744eebad8e9e31336c26cf51e678c3556f8.json b/backend/.sqlx/query-acdaa5151f8f7f37bb8c8c5a7d146789887e47db9695fc26b1dfaedd735e1e60.json similarity index 61% rename from backend/.sqlx/query-500393aa7c3f8c96ec8d92869c5c4744eebad8e9e31336c26cf51e678c3556f8.json rename to backend/.sqlx/query-acdaa5151f8f7f37bb8c8c5a7d146789887e47db9695fc26b1dfaedd735e1e60.json index 8b6e9d9ade..90b77d6cb7 100644 --- a/backend/.sqlx/query-500393aa7c3f8c96ec8d92869c5c4744eebad8e9e31336c26cf51e678c3556f8.json +++ b/backend/.sqlx/query-acdaa5151f8f7f37bb8c8c5a7d146789887e47db9695fc26b1dfaedd735e1e60.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "SELECT value\n FROM resource\n WHERE path = $1 AND workspace_id = $2", + "query": "SELECT value\n FROM resource\n WHERE path = $1 AND workspace_id = $2", "describe": { "columns": [ { @@ -19,5 +19,5 @@ true ] }, - "hash": "500393aa7c3f8c96ec8d92869c5c4744eebad8e9e31336c26cf51e678c3556f8" + "hash": "acdaa5151f8f7f37bb8c8c5a7d146789887e47db9695fc26b1dfaedd735e1e60" } diff --git a/backend/windmill-api/src/openai.rs b/backend/windmill-api/src/openai.rs index 7bfa926ee5..8d76b19341 100644 --- a/backend/windmill-api/src/openai.rs +++ b/backend/windmill-api/src/openai.rs @@ -1,3 +1,5 @@ +use std::{collections::HashMap, sync::Arc}; + use crate::{ db::{ApiAuthed, DB}, variables::build_crypt, @@ -6,13 +8,14 @@ use crate::{ use axum::{ body::{Bytes, StreamBody}, - extract::{Extension, Path}, + extract::{Extension, Path, Query}, response::IntoResponse, routing::post, Router, }; use magic_crypt::MagicCryptTrait; -use serde_json::json; +use serde_json::value::RawValue; +use tokio::sync::RwLock; use windmill_audit::{audit_log, ActionKind}; use windmill_common::error::{to_anyhow, Error}; @@ -30,22 +33,30 @@ struct OpenaiResource { organization_id: Option, } -fn create_openai_json_error(msg: String) -> Error { - Error::OpenAIError( - serde_json::to_string(&json!({ - "error": { - "message": msg - } - })) - .unwrap(), - ) +#[derive(Deserialize)] +struct OpenaiClientCredentialsOauthResource { + client_id: String, + client_secret: String, + token_url: String, + user: Option, +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum OpenaiConfig { + Resource(OpenaiResource), + ClientCredentialsOauthResource(OpenaiClientCredentialsOauthResource), } struct Variable { value: String, is_secret: bool, } -async fn get_variable(path: String, db: &DB, w_id: &String) -> Result { +async fn get_variable_or_self(path: String, db: &DB, w_id: &String) -> Result { + if !path.starts_with("$var:") { + return Ok(path); + } + let path = path.strip_prefix("$var:").unwrap().to_string(); let mut tx = db.begin().await?; let mut variable = sqlx::query_as!( Variable, @@ -71,63 +82,218 @@ lazy_static::lazy_static! { pub static ref OPENAI_AZURE_BASE_PATH: Option = std::env::var("OPENAI_AZURE_BASE_PATH").ok(); } +#[derive(Deserialize)] +struct OpenaiCredentials { + access_token: String, +} +async fn get_openai_key_using_credentials_flow( + mut resource: OpenaiClientCredentialsOauthResource, + db: &DB, + w_id: &String, +) -> Result { + 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(¶ms) + .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::().await.map_err(|err| { + Error::InternalErr(format!( + "Failed to parse OpenAI credentials from credentials flow: {}", + err + )) + })?; + Ok(response.access_token) +} + +struct OpenaiKeyCache { + api_key: String, + organization_id: Option, + azure_base_path: Option, + user: Option, + expires_at: std::time::Instant, +} + +impl OpenaiKeyCache { + fn new( + api_key: String, + organization_id: Option, + azure_base_path: Option, + expires_at: std::time::Instant, + user: Option, + ) -> Self { + Self { api_key, organization_id, azure_base_path, expires_at, user } + } + fn is_expired(&self) -> bool { + self.expires_at < std::time::Instant::now() + } +} + +lazy_static::lazy_static! { + static ref OPENAI_KEY_CACHE: Arc>> = Arc::new(RwLock::new(HashMap::new())); +} + +#[derive(Deserialize)] +struct ProxyQueryParams { + no_cache: Option, +} async fn proxy( authed: ApiAuthed, Extension(db): Extension, Path((w_id, openai_path)): Path<(String, String)>, - body: Bytes, + Query(query_params): Query, + mut body: Bytes, ) -> impl IntoResponse { - let mut tx = db.begin().await?; - let openai_resource_path = sqlx::query_scalar!( - "SELECT openai_resource_path FROM workspace_settings WHERE workspace_id = $1", - &w_id - ) - .fetch_one(&mut *tx) - .await?; - tx.commit().await?; + let mut cache = OPENAI_KEY_CACHE.write().await; + let workspace_cache = 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.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(create_openai_json_error( - "OpenAI resource not configured".to_string(), - )); + 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::; + 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::(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(), + ); + 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> = 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 openai_resource_path = openai_resource_path.unwrap(); - - tx = db.begin().await?; - let resource = sqlx::query_scalar!( - "SELECT value - FROM resource - WHERE path = $1 AND workspace_id = $2", - &openai_resource_path, - &w_id - ) - .fetch_optional(&mut *tx) - .await? - .ok_or_else(|| { - create_openai_json_error(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(create_openai_json_error( - "OpenAI resource missing value".to_string(), - )); - } - - let mut resource: OpenaiResource = serde_json::from_value(resource.unwrap()) - .map_err(|e| Error::InternalErr(format!("validating openai resource {e}")))?; - - if resource.api_key.starts_with("$var:") { - let openai_api_key_path = resource.api_key.strip_prefix("$var:").unwrap().to_string(); - resource.api_key = get_variable(openai_api_key_path, &db, &w_id).await?; - } - - let base_url = if let Some(base_url) = &*OPENAI_AZURE_BASE_PATH { + let base_url = if let Some(base_url) = azure_base_path { base_url } else { - "https://api.openai.com/v1" + "https://api.openai.com/v1".to_string() }; let url = format!("{}/{}", base_url, openai_path); @@ -138,29 +304,24 @@ async fn proxy( if base_url != "https://api.openai.com/v1" { request = request - .header("api-key", resource.api_key) + .header("api-key", api_key) .query(&[("api-version", "2023-05-15")]) } else { - request = request.header("authorization", format!("Bearer {}", resource.api_key)) + request = request.header("authorization", format!("Bearer {}", api_key)) } - if let Some(mut org_id) = resource.organization_id { - tracing::info!("org_id: {:?}", org_id); - if org_id.starts_with("$var:") { - let openai_organisation_path = org_id.strip_prefix("$var:").unwrap().to_string(); - org_id = get_variable(openai_organisation_path, &db, &w_id).await?; - } + if let Some(org_id) = organization_id { request = request.header("OpenAI-Organization", org_id); } let response = request.send().await.map_err(to_anyhow)?; - tx = db.begin().await?; + let mut tx = db.begin().await?; audit_log( &mut *tx, &authed.username, "openai.request", - ActionKind::Update, + ActionKind::Execute, &w_id, Some(&authed.email), Some([("openai_path", &format!("{:?}", openai_path)[..])].into()), diff --git a/frontend/src/lib/components/copilot/lib.ts b/frontend/src/lib/components/copilot/lib.ts index 514b024eb1..85c7b365e6 100644 --- a/frontend/src/lib/components/copilot/lib.ts +++ b/frontend/src/lib/components/copilot/lib.ts @@ -73,7 +73,7 @@ export async function testKey({ } ) } else { - await getNonStreamingCompletion(messages, abortController) + await getNonStreamingCompletion(messages, abortController, undefined, true) } } @@ -233,7 +233,8 @@ const PROMPTS_CONFIGS = { export async function getNonStreamingCompletion( messages: ChatCompletionMessageParam[], abortController: AbortController, - model: string = 'gpt-4-1106-preview' + model: string = 'gpt-4-1106-preview', + noCache?: boolean ) { const openaiClient = workspacedOpenai.getClient() const completion = await openaiClient.chat.completions.create( @@ -244,6 +245,9 @@ export async function getNonStreamingCompletion( model }, { + query: { + no_cache: noCache + }, signal: abortController.signal } ) diff --git a/frontend/src/lib/components/instanceSettings.ts b/frontend/src/lib/components/instanceSettings.ts index 3e1e9fd040..34827e3e56 100644 --- a/frontend/src/lib/components/instanceSettings.ts +++ b/frontend/src/lib/components/instanceSettings.ts @@ -98,6 +98,16 @@ export const settings: Record = { fieldType: 'boolean', storage: 'setting', ee_only: 'No workaround around this' + }, + { + label: 'Azure OpenAI base path', + description: + 'All Windmill AI features will run on the specified deployed model. Format: https://{your-resource-name}.openai.azure.com/openai/deployments/{deployment-id}', + key: 'openai_azure_base_path', + fieldType: 'text', + storage: 'setting', + ee_only: + 'You can still set this setting by using OPENAI_AZURE_BASE_PATH as env variable to the server containers' } ], SMTP: [ diff --git a/frontend/src/routes/(root)/(logged)/workspace_settings/+page.svelte b/frontend/src/routes/(root)/(logged)/workspace_settings/+page.svelte index 35f7a3b131..d64a07e8c1 100644 --- a/frontend/src/routes/(root)/(logged)/workspace_settings/+page.svelte +++ b/frontend/src/routes/(root)/(logged)/workspace_settings/+page.svelte @@ -20,7 +20,8 @@ Script, WorkspaceService, HelpersService, - JobService + JobService, + ResourceService } from '$lib/gen' import { enterpriseLicense, @@ -72,6 +73,7 @@ | 'webhook' | 'deploy_to' | 'error_handler') ?? 'users' + let usingOpenaiClientCredentialsOauth = false // function getDropDownItems(username: string): DropdownItem[] { // return [ @@ -248,6 +250,12 @@ ? settings.large_file_storage?.s3_resource_path?.replace('$res:', '') : undefined gitSyncResourcePath = settings.git_sync?.git_repo_resource_path?.replace('$res:', '') + + // check openai_client_credentials_oauth + const resourceTypes = await ResourceService.listResourceTypeNames({ + workspace: $workspaceStore! + }) + usingOpenaiClientCredentialsOauth = resourceTypes.includes('openai_client_credentials_oauth') } $: { @@ -649,9 +657,11 @@
- {#key openaiResourceInitialPath} + {#key [openaiResourceInitialPath, usingOpenaiClientCredentialsOauth]} { editCopilotConfig(ev.detail)