diff --git a/package.json b/package.json index 2c9051e5..6071c824 100644 --- a/package.json +++ b/package.json @@ -297,6 +297,12 @@ "title": "Restart LSP Server", "icon": "$(debug-restart)", "category": "ZenML Environment" + }, + { + "command": "zenml.registerLLMAPIKey", + "title": "Register LLM API Key", + "icon": "$(add)", + "category": "ZenML Secrets" } ], "viewsContainers": { diff --git a/src/commands/secrets/cmds.ts b/src/commands/secrets/cmds.ts new file mode 100644 index 00000000..e7819c65 --- /dev/null +++ b/src/commands/secrets/cmds.ts @@ -0,0 +1,62 @@ +// 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'; + +const registerLLMAPIKey = async (context: ExtensionContext) => { + const options: vscode.QuickPickItem[] = [ + { label: 'Anthropic' }, + { label: 'Gemini' }, + { label: 'OpenAI' }, + ]; + + 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 undefined; + } + + const model = selectedOption.label; + const secretKey = `zenml.${model.toLowerCase()}.key`; + + let apiKey = await context.secrets.get(secretKey); + + if (apiKey) { + apiKey = await vscode.window.showInputBox({ + prompt: `${model} API Key already exists, enter a new value to update.`, + password: true, + }); + } else { + apiKey = await vscode.window.showInputBox({ + prompt: `Please enter your ${model} API key`, + password: true, + }); + } + + if (apiKey === undefined) { + vscode.window.showWarningMessage('API key input was canceled.'); + return; + } + + await context.secrets.store(secretKey, apiKey); + vscode.window.showInformationMessage(`${model} API key stored successfully.`); +}; + +export const secretsCommands = { + registerLLMAPIKey, +}; diff --git a/src/commands/secrets/registry.ts b/src/commands/secrets/registry.ts new file mode 100644 index 00000000..af861a6b --- /dev/null +++ b/src/commands/secrets/registry.ts @@ -0,0 +1,43 @@ +// 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 { secretsCommands } from './cmds'; +import { registerCommand } from '../../common/vscodeapi'; +import { ZenExtension } from '../../services/ZenExtension'; +import { ExtensionContext, commands } from 'vscode'; + +/** + * 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. + */ +export const registerSecretsCommands = (context: ExtensionContext) => { + try { + const registeredCommands = [ + registerCommand( + 'zenml.registerLLMAPIKey', + async () => await secretsCommands.registerLLMAPIKey(context) + ), + ]; + + registeredCommands.forEach(cmd => { + context.subscriptions.push(cmd); + ZenExtension.commandDisposables.push(cmd); + }); + + commands.executeCommand('setContext', 'secretsCommandsRegistered', true); + } catch (error) { + console.error('Error registering secrets commands:', error); + commands.executeCommand('setContext', 'secretsCommandsRegistered', false); + } +}; diff --git a/src/common/vscodeapi.ts b/src/common/vscodeapi.ts index ff2463bd..5cdeaf58 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 undefined; + } + + return secret; +} 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, ]; /**