starred/windmill
1
0
Fork 0
mirror of https://github.com/windmill-labs/windmill.git synced 2026-08-25 00:01:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: oauth flow openai key (#2833)

Browse Source 
* feat: oauth flow openai key

* fix: rawvalue insert
This commit is contained in:
HugoCasa
2023-12-12 07:57:18 +01:00
committed by GitHub
parent 350c8e843f
commit 4c7d64bc5f
6 changed files with 281 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/.sqlx/query-0cc3618495d5d024b2a173c58a3a8bb2a9d69b7b6e7ed6b0d0064fa2ce9c2e31.json
Generated
+20
View File
@@ -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"
}
backend/.sqlx/query-500393aa7c3f8c96ec8d92869c5c4744eebad8e9e31336c26cf51e678c3556f8.json → backend/.sqlx/query-acdaa5151f8f7f37bb8c8c5a7d146789887e47db9695fc26b1dfaedd735e1e60.json
Generated
+2 -2
View File
@@ -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"
}
backend/windmill-api/src/openai.rs
+230 -69
View File
@@ -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<String>,
}
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<String>,
}
#[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<String, Error> {
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 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<String> = 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<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)
}
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()
}
}
lazy_static::lazy_static! {
static ref OPENAI_KEY_CACHE: Arc<RwLock<HashMap<String, OpenaiKeyCache>>> = Arc::new(RwLock::new(HashMap::new()));
}
#[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)>,
body: Bytes,
Query(query_params): Query<ProxyQueryParams>,
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::<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(),
);
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 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()),
frontend/src/lib/components/copilot/lib.ts
+6 -2
View File
@@ -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
}
)
frontend/src/lib/components/instanceSettings.ts
+10
View File
@@ -98,6 +98,16 @@ export const settings: Record<string, Setting[]> = {
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: [
frontend/src/routes/(root)/(logged)/workspace_settings/+page.svelte
+13 -3
View File
@@ -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 @@
</Alert>
</div>
<div class="mt-5 flex gap-1">
{#key openaiResourceInitialPath}
{#key [openaiResourceInitialPath, usingOpenaiClientCredentialsOauth]}
<ResourcePicker
resourceType="openai"
resourceType={usingOpenaiClientCredentialsOauth
? 'openai_client_credentials_oauth'
: 'openai'}
initialValue={openaiResourceInitialPath}
on:change={(ev) => {
editCopilotConfig(ev.detail)