Merge pull request zenml-io#31 from sklarfox/feature/openai-api-key

Feature: Register OpenAI API key

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": {

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,
+};

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);
+  }
+};

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;
+}

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,
   ];
 
   /**