Feature: Register OpenAI API key

package.json

@@ -297,6 +297,12 @@
         "title": "Restart LSP Server",
         "icon": "$(debug-restart)",
         "category": "ZenML Environment"
+      },
+      {
+        "command": "zenml.registerOpenAIAPIKey",
+        "title": "Register OpenAI API Key",
+        "icon": "$(add)",
+        "category": "ZenML Secrets"
       }
     ],
     "viewsContainers": {

src/commands/secrets/cmds.ts

@@ -0,0 +1,44 @@
+// 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.');
+};
+
+export const secretsCommands = {
+  registerOpenAIAPIKey,
+};

src/commands/secrets/registry.ts

@@ -0,0 +1,46 @@
+// 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)
+      ),
+    ];
+
+    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);
+  }
+};

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

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