From 94b0f73130c10e218bc0ebb9fd3fbb33de832e77 Mon Sep 17 00:00:00 2001 From: Alex Sklar Date: Wed, 21 Aug 2024 12:25:34 -0700 Subject: [PATCH 1/7] feat(secrets): add commands to register, update, and remove OpenAI API keys with the VSCode SecretsStorage API --- package.json | 12 +++++++ src/commands/secrets/cmds.ts | 56 ++++++++++++++++++++++++++++++++ src/commands/secrets/registry.ts | 50 ++++++++++++++++++++++++++++ src/services/ZenExtension.ts | 2 ++ 4 files changed, 120 insertions(+) create mode 100644 src/commands/secrets/cmds.ts create mode 100644 src/commands/secrets/registry.ts diff --git a/package.json b/package.json index 2c9051e5..fe1d0a1c 100644 --- a/package.json +++ b/package.json @@ -297,6 +297,18 @@ "title": "Restart LSP Server", "icon": "$(debug-restart)", "category": "ZenML Environment" + }, + { + "command": "zenml.registerOpenAIAPIKey", + "title": "Register OpenAI API Key", + "icon": "$(add)", + "category": "ZenML Secrets" + }, + { + "command": "zenml.deleteOpenAIAPIKey", + "title": "Delete OpenAI API Key", + "icon": "$(trash)", + "category": "ZenML Secrets" } ], "viewsContainers": { diff --git a/src/commands/secrets/cmds.ts b/src/commands/secrets/cmds.ts new file mode 100644 index 00000000..b68ad9b9 --- /dev/null +++ b/src/commands/secrets/cmds.ts @@ -0,0 +1,56 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied.See the License for the specific language governing +// permissions and limitations under the License. + +import * as vscode from 'vscode'; +import type { ExtensionContext } from 'vscode'; + +// TODO I don't think retrieval of an api key will live in here + +const registerOpenAIAPIKey = async (context: ExtensionContext) => { + let apiKey = await context.secrets.get('OPENAI_API_KEY'); + + if (apiKey) { + apiKey = await vscode.window.showInputBox({ + prompt: 'OpenAI API Key already exists, enter a new value to update.', + password: true, + }); + } else { + apiKey = await vscode.window.showInputBox({ + prompt: 'Please enter your OpenAI API key', + password: true, + }); + } + + if (apiKey === undefined) { + return undefined; + } + + await context.secrets.store('OPENAI_API_KEY', apiKey); + vscode.window.showInformationMessage('OpenAI API key stored successfully.'); +}; + +const deleteOpenAIAPIKey = async (context: ExtensionContext) => { + const apiKey = await context.secrets.get('OPENAI_API_KEY'); + + if (apiKey === undefined) { + vscode.window.showInformationMessage('No OpenAI API key exists.'); + return; + } + await context.secrets.delete('OPENAI_API_KEY'); + vscode.window.showInformationMessage('OpenAI API key successfully removed.'); +}; + +export const secretsCommands = { + registerOpenAIAPIKey, + deleteOpenAIAPIKey, +}; diff --git a/src/commands/secrets/registry.ts b/src/commands/secrets/registry.ts new file mode 100644 index 00000000..951599ec --- /dev/null +++ b/src/commands/secrets/registry.ts @@ -0,0 +1,50 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied.See the License for the specific language governing +// permissions and limitations under the License. + +// TODO the registration of secrets commands should go in here, which is then imported into the ZenExtension file +// In registry, the "register command thing is passed context (i.e. context.secrets) as an arg" + +import { secretsCommands } from './cmds'; +import { registerCommand } from '../../common/vscodeapi'; +import { ZenExtension } from '../../services/ZenExtension'; +import { ExtensionContext, commands } from 'vscode'; + +/** + * Registers pipeline-related commands for the extension. + * + * @param {ExtensionContext} context - The context in which the extension operates, used for registering commands and managing their lifecycle. + */ +export const registerSecretsCommands = (context: ExtensionContext) => { + try { + const registeredCommands = [ + registerCommand( + 'zenml.registerOpenAIAPIKey', + async () => await secretsCommands.registerOpenAIAPIKey(context) + ), + registerCommand( + 'zenml.deleteOpenAIAPIKey', + async () => await secretsCommands.deleteOpenAIAPIKey(context) + ), + ]; + + registeredCommands.forEach(cmd => { + context.subscriptions.push(cmd); + ZenExtension.commandDisposables.push(cmd); + }); + + commands.executeCommand('setContext', 'secretsCommandsRegistered', true); + } catch (error) { + console.error('Error registering pipeline commands:', error); + commands.executeCommand('setContext', 'secretsCommandsRegistered', false); + } +}; diff --git a/src/services/ZenExtension.ts b/src/services/ZenExtension.ts index 43a0f8e7..68e5f6eb 100644 --- a/src/services/ZenExtension.ts +++ b/src/services/ZenExtension.ts @@ -17,6 +17,7 @@ import * as vscode from 'vscode'; import { registerPipelineCommands } from '../commands/pipelines/registry'; import { registerServerCommands } from '../commands/server/registry'; import { registerStackCommands } from '../commands/stack/registry'; +import { registerSecretsCommands } from '../commands/secrets/registry'; import { EXTENSION_ROOT_DIR } from '../common/constants'; import { registerLogger, traceLog, traceVerbose } from '../common/log/logging'; import { @@ -73,6 +74,7 @@ export class ZenExtension { registerStackCommands, registerComponentCommands, registerPipelineCommands, + registerSecretsCommands, ]; /** From 97dc9dd937d9416455c3e14c09cb92f14245f2ae Mon Sep 17 00:00:00 2001 From: Alex Sklar Date: Wed, 21 Aug 2024 12:44:19 -0700 Subject: [PATCH 2/7] feat(secrets): add function to retrieve secrets from VSCode --- src/common/vscodeapi.ts | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/common/vscodeapi.ts b/src/common/vscodeapi.ts index ff2463bd..da826b46 100644 --- a/src/common/vscodeapi.ts +++ b/src/common/vscodeapi.ts @@ -18,6 +18,7 @@ import { ConfigurationScope, Disposable, DocumentSelector, + ExtensionContext, languages, LanguageStatusItem, LogOutputChannel, @@ -69,3 +70,14 @@ export function createLanguageStatusItem( ): LanguageStatusItem { return languages.createLanguageStatusItem(id, selector); } + +export async function getSecret(context: ExtensionContext, key: string) { + const secret = await context.secrets.get(key); + + if (secret === undefined) { + console.error(`The requested secret with key '${key}' does not exist.`); + return; + } + + return secret; +} From 89b25e4df0f924fa778bbefbebf748a9d4b71b12 Mon Sep 17 00:00:00 2001 From: Alex Sklar Date: Thu, 22 Aug 2024 10:58:04 -0700 Subject: [PATCH 3/7] feat: remove unnecessary deleteOpenAIKey function --- package.json | 6 ------ src/commands/secrets/cmds.ts | 12 ------------ src/commands/secrets/registry.ts | 4 ---- 3 files changed, 22 deletions(-) diff --git a/package.json b/package.json index fe1d0a1c..73b7ecfb 100644 --- a/package.json +++ b/package.json @@ -303,12 +303,6 @@ "title": "Register OpenAI API Key", "icon": "$(add)", "category": "ZenML Secrets" - }, - { - "command": "zenml.deleteOpenAIAPIKey", - "title": "Delete OpenAI API Key", - "icon": "$(trash)", - "category": "ZenML Secrets" } ], "viewsContainers": { diff --git a/src/commands/secrets/cmds.ts b/src/commands/secrets/cmds.ts index b68ad9b9..faf0ebf3 100644 --- a/src/commands/secrets/cmds.ts +++ b/src/commands/secrets/cmds.ts @@ -39,18 +39,6 @@ const registerOpenAIAPIKey = async (context: ExtensionContext) => { vscode.window.showInformationMessage('OpenAI API key stored successfully.'); }; -const deleteOpenAIAPIKey = async (context: ExtensionContext) => { - const apiKey = await context.secrets.get('OPENAI_API_KEY'); - - if (apiKey === undefined) { - vscode.window.showInformationMessage('No OpenAI API key exists.'); - return; - } - await context.secrets.delete('OPENAI_API_KEY'); - vscode.window.showInformationMessage('OpenAI API key successfully removed.'); -}; - export const secretsCommands = { registerOpenAIAPIKey, - deleteOpenAIAPIKey, }; diff --git a/src/commands/secrets/registry.ts b/src/commands/secrets/registry.ts index 951599ec..c63ed07f 100644 --- a/src/commands/secrets/registry.ts +++ b/src/commands/secrets/registry.ts @@ -31,10 +31,6 @@ export const registerSecretsCommands = (context: ExtensionContext) => { 'zenml.registerOpenAIAPIKey', async () => await secretsCommands.registerOpenAIAPIKey(context) ), - registerCommand( - 'zenml.deleteOpenAIAPIKey', - async () => await secretsCommands.deleteOpenAIAPIKey(context) - ), ]; registeredCommands.forEach(cmd => { From 40f39ad9f200a2a9094480c5c4bc1fc9165479e2 Mon Sep 17 00:00:00 2001 From: Alex Sklar Date: Fri, 23 Aug 2024 12:56:17 -0700 Subject: [PATCH 4/7] feat: generalize LLM API Key registry command, support for Anthropic, Gemini, OpenAI --- package.json | 4 ++-- src/commands/secrets/cmds.ts | 37 ++++++++++++++++++++++++-------- src/commands/secrets/registry.ts | 4 ++-- 3 files changed, 32 insertions(+), 13 deletions(-) diff --git a/package.json b/package.json index 73b7ecfb..6071c824 100644 --- a/package.json +++ b/package.json @@ -299,8 +299,8 @@ "category": "ZenML Environment" }, { - "command": "zenml.registerOpenAIAPIKey", - "title": "Register OpenAI API Key", + "command": "zenml.registerLLMAPIKey", + "title": "Register LLM API Key", "icon": "$(add)", "category": "ZenML Secrets" } diff --git a/src/commands/secrets/cmds.ts b/src/commands/secrets/cmds.ts index faf0ebf3..e35e06e1 100644 --- a/src/commands/secrets/cmds.ts +++ b/src/commands/secrets/cmds.ts @@ -14,31 +14,50 @@ import * as vscode from 'vscode'; import type { ExtensionContext } from 'vscode'; -// TODO I don't think retrieval of an api key will live in here +const registerLLMAPIKey = async (context: ExtensionContext) => { + const options: vscode.QuickPickItem[] = [ + { label: 'Anthropic' }, + { label: 'Gemini' }, + { label: 'OpenAI' }, + ]; -const registerOpenAIAPIKey = async (context: ExtensionContext) => { - let apiKey = await context.secrets.get('OPENAI_API_KEY'); + const selectedOption = await vscode.window.showQuickPick(options, { + placeHolder: 'Please select an LLM.', + canPickMany: false, + }); + + if (selectedOption === undefined) { + vscode.window.showWarningMessage('API key input was canceled.'); + return; + } + + const model = selectedOption.label; + const secretKey = `${model.toUpperCase()}_API_KEY`; + + let apiKey = await context.secrets.get(secretKey); + console.log(secretKey, apiKey); if (apiKey) { apiKey = await vscode.window.showInputBox({ - prompt: 'OpenAI API Key already exists, enter a new value to update.', + prompt: `${model} API Key already exists, enter a new value to update.`, password: true, }); } else { apiKey = await vscode.window.showInputBox({ - prompt: 'Please enter your OpenAI API key', + prompt: `Please enter your ${model} API key`, password: true, }); } if (apiKey === undefined) { - return undefined; + vscode.window.showWarningMessage('API key input was canceled.'); + return; } - await context.secrets.store('OPENAI_API_KEY', apiKey); - vscode.window.showInformationMessage('OpenAI API key stored successfully.'); + await context.secrets.store(secretKey, apiKey); + vscode.window.showInformationMessage(`${model} API key stored successfully.`); }; export const secretsCommands = { - registerOpenAIAPIKey, + registerLLMAPIKey, }; diff --git a/src/commands/secrets/registry.ts b/src/commands/secrets/registry.ts index c63ed07f..df0b5fce 100644 --- a/src/commands/secrets/registry.ts +++ b/src/commands/secrets/registry.ts @@ -28,8 +28,8 @@ export const registerSecretsCommands = (context: ExtensionContext) => { try { const registeredCommands = [ registerCommand( - 'zenml.registerOpenAIAPIKey', - async () => await secretsCommands.registerOpenAIAPIKey(context) + 'zenml.registerLLMAPIKey', + async () => await secretsCommands.registerLLMAPIKey(context) ), ]; From adbc8c9863231768f308adcb1d5341485a925941 Mon Sep 17 00:00:00 2001 From: Alex Sklar Date: Fri, 23 Aug 2024 12:59:30 -0700 Subject: [PATCH 5/7] chore: remove todo comments --- src/commands/secrets/registry.ts | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/commands/secrets/registry.ts b/src/commands/secrets/registry.ts index df0b5fce..af861a6b 100644 --- a/src/commands/secrets/registry.ts +++ b/src/commands/secrets/registry.ts @@ -11,16 +11,13 @@ // or implied.See the License for the specific language governing // permissions and limitations under the License. -// TODO the registration of secrets commands should go in here, which is then imported into the ZenExtension file -// In registry, the "register command thing is passed context (i.e. context.secrets) as an arg" - import { secretsCommands } from './cmds'; import { registerCommand } from '../../common/vscodeapi'; import { ZenExtension } from '../../services/ZenExtension'; import { ExtensionContext, commands } from 'vscode'; /** - * Registers pipeline-related commands for the extension. + * Registers secrets related commands for the extension. * * @param {ExtensionContext} context - The context in which the extension operates, used for registering commands and managing their lifecycle. */ @@ -40,7 +37,7 @@ export const registerSecretsCommands = (context: ExtensionContext) => { commands.executeCommand('setContext', 'secretsCommandsRegistered', true); } catch (error) { - console.error('Error registering pipeline commands:', error); + console.error('Error registering secrets commands:', error); commands.executeCommand('setContext', 'secretsCommandsRegistered', false); } }; From fc0917b91a771f3ac5925f3ee3a2a17f06d6f604 Mon Sep 17 00:00:00 2001 From: Alex Sklar Date: Fri, 23 Aug 2024 13:04:22 -0700 Subject: [PATCH 6/7] chore: address coderabbitai suggestions --- src/commands/secrets/cmds.ts | 3 +-- src/common/vscodeapi.ts | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/commands/secrets/cmds.ts b/src/commands/secrets/cmds.ts index e35e06e1..99514123 100644 --- a/src/commands/secrets/cmds.ts +++ b/src/commands/secrets/cmds.ts @@ -28,14 +28,13 @@ const registerLLMAPIKey = async (context: ExtensionContext) => { if (selectedOption === undefined) { vscode.window.showWarningMessage('API key input was canceled.'); - return; + return undefined; } const model = selectedOption.label; const secretKey = `${model.toUpperCase()}_API_KEY`; let apiKey = await context.secrets.get(secretKey); - console.log(secretKey, apiKey); if (apiKey) { apiKey = await vscode.window.showInputBox({ diff --git a/src/common/vscodeapi.ts b/src/common/vscodeapi.ts index da826b46..5cdeaf58 100644 --- a/src/common/vscodeapi.ts +++ b/src/common/vscodeapi.ts @@ -76,7 +76,7 @@ export async function getSecret(context: ExtensionContext, key: string) { if (secret === undefined) { console.error(`The requested secret with key '${key}' does not exist.`); - return; + return undefined; } return secret; From 1d525b131cea8a310998217c55077d3a23f7b58b Mon Sep 17 00:00:00 2001 From: Alex Sklar Date: Thu, 29 Aug 2024 11:51:31 -0700 Subject: [PATCH 7/7] refactor: change api key naming convention to `zenml.{provider}.key` --- src/commands/secrets/cmds.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/commands/secrets/cmds.ts b/src/commands/secrets/cmds.ts index 99514123..e7819c65 100644 --- a/src/commands/secrets/cmds.ts +++ b/src/commands/secrets/cmds.ts @@ -32,7 +32,7 @@ const registerLLMAPIKey = async (context: ExtensionContext) => { } const model = selectedOption.label; - const secretKey = `${model.toUpperCase()}_API_KEY`; + const secretKey = `zenml.${model.toLowerCase()}.key`; let apiKey = await context.secrets.get(secretKey);