Fix Yaml datetime format (#79)

libs/ktem/ktem/embeddings/ui.py

@@ -4,6 +4,7 @@ import gradio as gr
 import pandas as pd
 import yaml
 from ktem.app import BasePage
+from ktem.utils.file import YAMLNoDateSafeLoader
 
 from .manager import embedding_models_manager
 
@@ -222,7 +223,7 @@ class EmbeddingManagement(BasePage):
 
     def create_emb(self, name, choices, spec, default):
         try:
-            spec = yaml.safe_load(spec)
+            spec = yaml.load(spec, Loader=YAMLNoDateSafeLoader)
             spec["__type__"] = (
                 embedding_models_manager.vendors()[choices].__module__
                 + "."
@@ -308,7 +309,7 @@ class EmbeddingManagement(BasePage):
 
     def save_emb(self, selected_emb_name, default, spec):
         try:
-            spec = yaml.safe_load(spec)
+            spec = yaml.load(spec, Loader=YAMLNoDateSafeLoader)
             spec["__type__"] = embedding_models_manager.info()[selected_emb_name][
                 "spec"
             ]["__type__"]

libs/ktem/ktem/index/ui.py

@@ -2,6 +2,7 @@ import gradio as gr
 import pandas as pd
 import yaml
 from ktem.app import BasePage
+from ktem.utils.file import YAMLNoDateSafeLoader
 
 from .manager import IndexManager
 
@@ -231,7 +232,9 @@ class IndexManagement(BasePage):
         """
         try:
             self.manager.build_index(
-                name=name, config=yaml.safe_load(config), index_type=index_type
+                name=name,
+                config=yaml.load(config, Loader=YAMLNoDateSafeLoader),
+                index_type=index_type,
             )
             gr.Info(f'Create index "{name}" successfully. Please restart the app!')
         except Exception as e:
@@ -294,7 +297,7 @@ class IndexManagement(BasePage):
 
     def update_index(self, selected_index_id: int, name: str, config: str):
         try:
-            spec = yaml.safe_load(config)
+            spec = yaml.load(config, Loader=YAMLNoDateSafeLoader)
             self.manager.update_index(selected_index_id, name, spec)
             gr.Info(f'Update index "{name}" successfully. Please restart the app!')
         except Exception as e:

libs/ktem/ktem/llms/ui.py

@@ -4,6 +4,7 @@ import gradio as gr
 import pandas as pd
 import yaml
 from ktem.app import BasePage
+from ktem.utils.file import YAMLNoDateSafeLoader
 
 from .manager import llms
 
@@ -219,7 +220,7 @@ class LLMManagement(BasePage):
 
     def create_llm(self, name, choices, spec, default):
         try:
-            spec = yaml.safe_load(spec)
+            spec = yaml.load(spec, Loader=YAMLNoDateSafeLoader)
             spec["__type__"] = (
                 llms.vendors()[choices].__module__
                 + "."
@@ -305,12 +306,12 @@ class LLMManagement(BasePage):
 
     def save_llm(self, selected_llm_name, default, spec):
         try:
-            spec = yaml.safe_load(spec)
+            spec = yaml.load(spec, Loader=YAMLNoDateSafeLoader)
             spec["__type__"] = llms.info()[selected_llm_name]["spec"]["__type__"]
             llms.update(selected_llm_name, spec=spec, default=default)
             gr.Info(f"LLM {selected_llm_name} saved successfully")
         except Exception as e:
-            gr.Error(f"Failed to save LLM {selected_llm_name}: {e}")
+            raise gr.Error(f"Failed to save LLM {selected_llm_name}: {e}")
 
     def delete_llm(self, selected_llm_name):
         try:

libs/ktem/ktem/utils/file.py

@@ -0,0 +1,23 @@
+import yaml
+
+
+class YAMLNoDateSafeLoader(yaml.SafeLoader):
+    """Load datetime as strings, not dates"""
+
+    @classmethod
+    def remove_implicit_resolver(cls, tag_to_remove):
+        """Remove implicit resolvers for a particular tag
+
+        Args:
+            tag_to_remove (str): YAML tag to remove
+        """
+        if "yaml_implicit_resolvers" not in cls.__dict__:
+            cls.yaml_implicit_resolvers = cls.yaml_implicit_resolvers.copy()
+
+        for first_letter, mappings in cls.yaml_implicit_resolvers.items():
+            cls.yaml_implicit_resolvers[first_letter] = [
+                (tag, regexp) for tag, regexp in mappings if tag != tag_to_remove
+            ]
+
+
+YAMLNoDateSafeLoader.remove_implicit_resolver("tag:yaml.org,2002:timestamp")