From d42d5945591f220513ac1e670134bc5de54770fb Mon Sep 17 00:00:00 2001
From: Sarmad Qadri <sarmad@lastmileai.dev>
Date: Tue, 30 Jan 2024 04:13:48 -0500
Subject: [PATCH] [1/n] Add load_json and load_yaml methods to AIConfigRuntime

Allow loading an AIConfigRuntime from a string. There's already a `loadJSON` method in the TS SDK.

This is needed for VSCode extension (top of stack), which gets the state of the document from the editor (and which can contain unsaved changes that aren't present on the filesystem).

Test Plan:

Tested out getting started ipynb with:

load JSON:

```
config = AIConfigRuntime.load('travel.aiconfig.json')
```

load YAML:
```
config.save('travel.aiconfig.yaml', mode="yaml")
yaml_str = ""
with open('travel.aiconfig.yaml') as file:
    yaml_str = file.read()

config = AIConfigRuntime.load_yaml(yaml_str)
```

Made sure config runtime is valid and tried `config.run`
---
 python/src/aiconfig/Config.py | 42 ++++++++++++++++++++++++++++-------
 typescript/lib/config.ts      | 14 ++++++++++++
 2 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/python/src/aiconfig/Config.py b/python/src/aiconfig/Config.py
index a78f98ea1..40ce861bb 100644
--- a/python/src/aiconfig/Config.py
+++ b/python/src/aiconfig/Config.py
@@ -53,7 +53,9 @@
 ModelParserRegistry.register_model_parser(
     DefaultAnyscaleEndpointParser("AnyscaleEndpoint")
 )
-ModelParserRegistry.register_model_parser(GeminiModelParser("gemini-pro"), ["gemini-pro"])
+ModelParserRegistry.register_model_parser(
+    GeminiModelParser("gemini-pro"), ["gemini-pro"]
+)
 ModelParserRegistry.register_model_parser(ClaudeBedrockModelParser())
 ModelParserRegistry.register_model_parser(HuggingFaceTextGenerationParser())
 for model in gpt_models_extra:
@@ -61,6 +63,7 @@
 ModelParserRegistry.register_model_parser(PaLMChatParser())
 ModelParserRegistry.register_model_parser(PaLMTextParser())
 
+
 class AIConfigRuntime(AIConfig):
     # A mapping of model names to their respective parsers
 
@@ -116,14 +119,37 @@ def load(cls, config_filepath: str) -> "AIConfigRuntime":
             else:
                 data = file.read()
 
-        # load the file as bytes and let pydantic handle the parsing
-        # validated_data =  AIConfig.model_validate_json(file.read())
-        aiconfigruntime = cls.model_validate_json(data)
-        update_model_parser_registry_with_config_runtime(aiconfigruntime)
-
+        config_runtime = cls.load_json(data)
         # set the file path. This is used when saving the config
-        aiconfigruntime.file_path = config_filepath
-        return aiconfigruntime
+        config_runtime.file_path = config_filepath
+        return config_runtime
+
+    @classmethod
+    def load_json(cls, config_json: str) -> "AIConfigRuntime":
+        """
+        Constructs AIConfigRuntime from provided JSON and returns it.
+
+        Args:
+            config_json (str): The JSON representing the AIConfig.
+        """
+
+        config_runtime = cls.model_validate_json(config_json)
+        update_model_parser_registry_with_config_runtime(config_runtime)
+
+        return config_runtime
+
+    @classmethod
+    def load_yaml(cls, config_yaml: str) -> "AIConfigRuntime":
+        """
+        Constructs AIConfigRuntime from provided YAML and returns it.
+
+        Args:
+            config_yaml (str): The YAML representing the AIConfig.
+        """
+
+        yaml_data = yaml.safe_load(config_yaml)
+        config_json = json.dumps(yaml_data)
+        return cls.load_json(config_json)
 
     @classmethod
     def load_from_workbook(cls, workbook_id: str) -> "AIConfigRuntime":
diff --git a/typescript/lib/config.ts b/typescript/lib/config.ts
index 2ecde4f9e..7289b8218 100644
--- a/typescript/lib/config.ts
+++ b/typescript/lib/config.ts
@@ -114,11 +114,25 @@ export class AIConfigRuntime implements AIConfig {
     return config;
   }
 
+  /**
+   * Loads an AIConfig from a YAML string.
+   * @param aiConfigYAML YAML string to load the AIConfig from.
+   */
+  public static loadYAML(aiConfigYAML: string) {
+    const aiConfigObj = yaml.load(aiConfigYAML);
+    return this.loadJSON(aiConfigObj);
+  }
+
   /**
    * Loads an AIConfig from a JSON object.
    * @param aiConfigObj JSON object to load the AIConfig from.
    */
   public static loadJSON(aiConfigObj: any) {
+    if (typeof aiConfigObj === "string") {
+      // Parse the string as JSON
+      aiConfigObj = JSON.parse(aiConfigObj);
+    }
+
     // TODO: saqadri - validate that the type satisfies AIConfig interface
     const aiConfig = new AIConfigRuntime(
       aiConfigObj.name,