From 5396c93cdf92fa3f4b5d62b56ccfce2235e362ca Mon Sep 17 00:00:00 2001 From: Sarmad Qadri Date: Tue, 30 Jan 2024 04:13:48 -0500 Subject: [PATCH] [2/n] Add to_string method to AIConfigRuntime Allow saving an AIConfigRuntime to a string by creating a `to_string` method that encapsulates the core business logic of the `save` method. This is needed for VSCode extension (top of stack), which sends the state of the AIConfigRuntime to the editor as a string for management (the IDE is responsible for managing the document, and for saving it to disk). Test Plan: Tested out getting started ipynb with: save JSON: ``` config = AIConfigRuntime.load('travel.aiconfig.json') config.save('travel2.aiconfig.json') ``` save YAML: ``` config.save('travel2.aiconfig.yaml', mode="yaml") ``` to JSON string: ``` config_str = config.to_string() print(config_str) ``` to YAML string: ``` config_yaml = config.to_string(mode="yaml") print(config_yaml) ``` For further sanity made sure that the generated strings could be loaded back as an AIConfig and `config.run` worked: ``` config = config.load_yaml(config_yaml) config.run(...) ``` --- python/src/aiconfig/Config.py | 79 ++++++++++++++++++++--------------- typescript/lib/config.ts | 48 ++++++++++++++------- 2 files changed, 78 insertions(+), 49 deletions(-) diff --git a/python/src/aiconfig/Config.py b/python/src/aiconfig/Config.py index 40ce861bb..2003d4721 100644 --- a/python/src/aiconfig/Config.py +++ b/python/src/aiconfig/Config.py @@ -455,17 +455,6 @@ def save( config_filepath (str, optional): The file path to the JSON or YAML configuration file. Defaults to "aiconfig.json" or "aiconfig.yaml", depending on the mode. """ - # Decide if we want to serialize as YAML or JSON - - # AIConfig json should only contain the core data fields. These are auxiliary fields that should not be persisted - exclude_options = { - "prompt_index": True, - "file_path": True, - "callback_manager": True, - } - - if not include_outputs: - exclude_options["prompts"] = {"__all__": {"outputs"}} default_filepath = ( "aiconfig.yaml" if mode == "yaml" else "aiconfig.json" @@ -481,30 +470,54 @@ def save( # Default to JSON mode = "json" + config_string = self.to_string(include_outputs, mode) + with open(config_filepath, "w") as file: - # Serialize the AIConfig to JSON - json_data = self.model_dump( - mode="json", - exclude=exclude_options, - exclude_none=True, + file.write(config_string) + + def to_string( + self, + include_outputs: bool = True, + mode: Literal["json", "yaml"] = "json", + ) -> str: + """ + Returns the well-formatted string representing the AIConfig object. + Note that this method will return the string that would be saved as a .aiconfig file using the save() method. + To get the raw string representation of the AIConfig object, use the __str__() method. + """ + # AIConfig json should only contain the core data fields. These are auxiliary fields that should not be persisted + exclude_options = { + "prompt_index": True, + "file_path": True, + "callback_manager": True, + } + + if not include_outputs: + exclude_options["prompts"] = {"__all__": {"outputs"}} + + # Serialize the AIConfig to JSON + json_data = self.model_dump( + mode="json", + exclude=exclude_options, + exclude_none=True, + ) + + if json_data.get("$schema", None) is None: + # Set the schema if it is not set + json_data["$schema"] = "https://json.schemastore.org/aiconfig-1.0" + + if mode == "yaml": + # Save AIConfig JSON as YAML string + return yaml.dump( + json_data, + indent=2, + ) + else: + # Save AIConfig as JSON string, with the schema specified + return json.dumps( + json_data, + indent=2, ) - if mode == "yaml": - # Save AIConfig JSON as YAML to the file - yaml.dump( - json_data, - file, - indent=2, - ) - else: - # Save AIConfig as JSON to the file, with the schema specified - json_data[ - "$schema" - ] = "https://json.schemastore.org/aiconfig-1.0" - json.dump( - json_data, - file, - indent=2, - ) def get_output_text( self, prompt: str | Prompt, output: Optional[dict] = None diff --git a/typescript/lib/config.ts b/typescript/lib/config.ts index 7289b8218..6d39eebdd 100644 --- a/typescript/lib/config.ts +++ b/typescript/lib/config.ts @@ -233,12 +233,40 @@ export class AIConfigRuntime implements AIConfig { * Saves this AIConfig to a file. * @param filePath The path to the file to save to. * @param saveOptions Options that determine how to save the AIConfig to the file. + * @param mode Whether to save the AIConfig as JSON or YAML. If unspecified, the file extension will be used to determine the mode. */ public save( filePath?: string, saveOptions?: SaveOptions, mode?: "json" | "yaml" ) { + const defaultFilePath = mode === "yaml" ? "aiconfig.yaml" : "aiconfig.json"; + if (!filePath) { + filePath = this.filePath ?? defaultFilePath; + } + + if (mode == null) { + if (isYamlExt(filePath)) { + mode = "yaml"; + } else { + // Default to JSON + mode = "json"; + } + } + + const aiConfigString = this.toString(saveOptions, mode); + fs.writeFileSync(filePath, aiConfigString); + } + + /** + * Returns this AIConfig as a string. + * + * Note that this doesn't return the full string representation of the AIConfig object, + * but rather the string representation of the AIConfig object that can be persisted to a file. + * @param saveOptions Options that determine how to serialize the AIConfig to string. + * @param mode Whether to save the AIConfig as JSON or YAML. Defaults to JSON. + */ + public toString(saveOptions?: SaveOptions, mode: "json" | "yaml" = "json") { const keysToOmit = ["filePath", "callbackManager"] as const; try { @@ -254,19 +282,9 @@ export class AIConfigRuntime implements AIConfig { aiConfigObj.prompts = prompts; } - const defaultFilePath = - mode === "yaml" ? "aiconfig.yaml" : "aiconfig.json"; - if (!filePath) { - filePath = this.filePath ?? defaultFilePath; - } - - if (mode == null) { - if (isYamlExt(filePath)) { - mode = "yaml"; - } else { - // Default to JSON - mode = "json"; - } + if (aiConfigObj["$schema"] == null) { + // Add the $schema property to the JSON object before saving it. In the future this can respect the version specified in the AIConfig. + aiConfigObj["$schema"] = "https://json.schemastore.org/aiconfig-1.0"; } // TODO: saqadri - make sure that the object satisfies the AIConfig schema @@ -274,12 +292,10 @@ export class AIConfigRuntime implements AIConfig { if (mode === "yaml") { aiConfigString = yaml.dump(aiConfigObj, { indent: 2 }); } else { - // Add the $schema property to the JSON object before saving it. In the future this can respect the version specified in the AIConfig. - aiConfigObj["$schema"] = "https://json.schemastore.org/aiconfig-1.0"; aiConfigString = JSON.stringify(aiConfigObj, null, 2); } - fs.writeFileSync(filePath, aiConfigString); + return aiConfigString; } catch (error) { console.error(error); throw error;