diff --git a/python/src/aiconfig/editor/server/server_utils.py b/python/src/aiconfig/editor/server/server_utils.py index 0788f45da..753a069c9 100644 --- a/python/src/aiconfig/editor/server/server_utils.py +++ b/python/src/aiconfig/editor/server/server_utils.py @@ -10,7 +10,7 @@ from textwrap import dedent from threading import Event from types import ModuleType -from typing import Any, Callable, NewType, Type, TypeVar, cast +from typing import Any, Callable, NewType, Type, TypeVar, cast, Optional import lastmile_utils.lib.core.api as core_utils import result @@ -66,6 +66,7 @@ class EditServerConfig(core_utils.Record): log_level: str | int = "INFO" server_mode: ServerMode = ServerMode.PROD parsers_module_path: str = "aiconfig_model_registry.py" + env_file_path: Optional[str] = None @field_validator("server_mode", mode="before") def convert_to_mode( @@ -84,6 +85,7 @@ class StartServerConfig(core_utils.Record): log_level: str | int = "INFO" server_mode: ServerMode = ServerMode.PROD parsers_module_path: str = "aiconfig_model_registry.py" + env_file_path: Optional[str] = None @field_validator("server_mode", mode="before") def convert_to_mode( @@ -294,13 +296,12 @@ def init_server_state( aiconfigrc_path: str, ) -> Result[None, str]: LOGGER.info("Initializing server state for 'edit' command") - # TODO: saqadri - load specific .env file if specified - dotenv.load_dotenv() _load_user_parser_module_if_exists( initialization_settings.parsers_module_path ) state = get_server_state(app) state.aiconfigrc_path = aiconfigrc_path + state.env_file_path = initialization_settings.env_file_path if isinstance(initialization_settings, StartServerConfig): # The aiconfig will be loaded later, when the editor sends the payload. diff --git a/vscode-extension/package.json b/vscode-extension/package.json index b2d13c4a8..09bdb06b9 100644 --- a/vscode-extension/package.json +++ b/vscode-extension/package.json @@ -128,6 +128,11 @@ "type": "string", "default": "", "description": "Version of the AIConfig Editor extension last time it was activated. We use this value to check if the extension has been updated to prompt users to reload VS Code" + }, + "vscode-aiconfig.env_file_path": { + "type": "string", + "default": "", + "description": "File path to your .env file containing your API keys" } } }, diff --git a/vscode-extension/src/constants.ts b/vscode-extension/src/constants.ts index a2033eb9b..0741b4b25 100644 --- a/vscode-extension/src/constants.ts +++ b/vscode-extension/src/constants.ts @@ -3,3 +3,5 @@ export const PYTHON_INTERPRETER_CACHE_KEY_NAME = "pythonInterpreter"; // Used for prompting user to reload VS Code window on update, and ensuring // that walkthrough gets shown on first install export const VERSION_KEY_NAME = "version"; + +export const ENV_FILE_PATH = "env_file_path"; \ No newline at end of file diff --git a/vscode-extension/src/editor_server/editorServer.ts b/vscode-extension/src/editor_server/editorServer.ts index c548d9898..0ae6710b2 100644 --- a/vscode-extension/src/editor_server/editorServer.ts +++ b/vscode-extension/src/editor_server/editorServer.ts @@ -3,6 +3,7 @@ import { EXTENSION_NAME } from "../util"; import { getPythonPath } from "../utilities/pythonSetupUtils"; import { ChildProcessWithoutNullStreams, spawn } from "child_process"; import serverPortManager from "./serverPortManager"; +import { ENV_FILE_PATH } from "../constants"; export enum EditorServerState { Starting = "Starting", @@ -73,6 +74,8 @@ export class EditorServer { const modelRegistryPathArgs = modelRegistryPath ? ["--parsers-module-path", modelRegistryPath] : []; + const filePath = config.get(ENV_FILE_PATH) as string; + const envFilePathArgs = filePath ? ["--env-file-path", filePath] : []; this.port = await serverPortManager.getPort(); @@ -89,6 +92,7 @@ export class EditorServer { "--server-port", this.port.toString(), ...modelRegistryPathArgs, + ...envFilePathArgs, ], { cwd: this.cwd, diff --git a/vscode-extension/src/extension.ts b/vscode-extension/src/extension.ts index 84602c42e..e3a577d3d 100644 --- a/vscode-extension/src/extension.ts +++ b/vscode-extension/src/extension.ts @@ -27,6 +27,7 @@ import { setupEnvironmentVariables, getConfigurationTarget, getModeFromDocument, + updateServerEnv } from "./util"; import { initialize, @@ -34,6 +35,7 @@ import { savePythonInterpreterToCache, } from "./utilities/pythonSetupUtils"; import { performVersionInstallAndUpdateActionsIfNeeded } from "./utilities/versionUpdateUtils"; +import { ENV_FILE_PATH } from "./constants"; // This method is called when your extension is activated // Your extension is activated the very first time a command is executed @@ -208,6 +210,25 @@ export async function activate(context: vscode.ExtensionContext) { } ) ); + + // Handle changes to the .env path + context.subscriptions.push(vscode.workspace.onDidChangeConfiguration( async (event) => { + if (event.affectsConfiguration("vscode-aiconfig" + "." + ENV_FILE_PATH)) { + // Get new env + const envPath = vscode.workspace.getConfiguration("vscode-aiconfig").get(ENV_FILE_PATH) as string; + console.log(`New .env path set: ${envPath}`); + // set env on all AIConfig Servers + const editors: Array = Array.from( + aiconfigEditorManager.getRegisteredEditors() + ); + if (editors.length > 0) { + editors.forEach(async (editor) => { + await updateServerEnv(editor.editorServer.url, envPath); + }); + + } + } + })); } // This method is called when your extension is deactivated diff --git a/vscode-extension/src/util.ts b/vscode-extension/src/util.ts index dd6ee61f3..9d9ad0e4e 100644 --- a/vscode-extension/src/util.ts +++ b/vscode-extension/src/util.ts @@ -8,6 +8,8 @@ import fs from "fs"; import path from "path"; import os from "os"; +import { ENV_FILE_PATH } from "./constants"; + export const EXTENSION_NAME = "vscode-aiconfig"; export const COMMANDS = { INIT: `${EXTENSION_NAME}.init`, @@ -45,6 +47,8 @@ export const EDITOR_SERVER_ROUTE_TABLE = { urlJoin(hostUrl, EDITOR_SERVER_API_ENDPOINT, "/load_content"), LOAD_MODEL_PARSER_MODULE: (hostUrl: string) => urlJoin(hostUrl, EDITOR_SERVER_API_ENDPOINT, "/load_model_parser_module"), + SET_ENV_FILE_PATH: (hostUrl: string) => + urlJoin(hostUrl, EDITOR_SERVER_API_ENDPOINT, "/set_env_file_path"), }; export async function isServerReady(serverUrl: string) { @@ -98,6 +102,15 @@ export async function updateServerState( }); } +export async function updateServerEnv(serverUrl: string, filePath: string) { + return await ufetch.post( + EDITOR_SERVER_ROUTE_TABLE.SET_ENV_FILE_PATH(serverUrl), + { + [ENV_FILE_PATH]: filePath, + } + ); +} + // Figure out what kind of AIConfig this is that we are loading export function getModeFromDocument( document: vscode.TextDocument @@ -381,6 +394,12 @@ export async function setupEnvironmentVariables( "Please define your environment variables." ); } + + // Update Server Env FLow + // Set the .env file path in the settings + // vscode Extension has a listener for changes defined at activation. + const config = vscode.workspace.getConfiguration(EXTENSION_NAME); + await config.update(ENV_FILE_PATH, envPath, getConfigurationTarget()); } /**