From 04e602161b1a40b02c13ac6d4b4880b69ea8d286 Mon Sep 17 00:00:00 2001 From: "trducng (john)" Date: Mon, 20 May 2024 17:36:14 +0700 Subject: [PATCH] Fix Yaml datetime format (#79) --- libs/ktem/ktem/embeddings/ui.py | 5 +++-- libs/ktem/ktem/index/ui.py | 7 +++++-- libs/ktem/ktem/llms/ui.py | 7 ++++--- libs/ktem/ktem/utils/file.py | 23 +++++++++++++++++++++++ 4 files changed, 35 insertions(+), 7 deletions(-) create mode 100644 libs/ktem/ktem/utils/file.py diff --git a/libs/ktem/ktem/embeddings/ui.py b/libs/ktem/ktem/embeddings/ui.py index c5fee3f..f8db9f1 100644 --- a/libs/ktem/ktem/embeddings/ui.py +++ b/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__"] diff --git a/libs/ktem/ktem/index/ui.py b/libs/ktem/ktem/index/ui.py index 6f553ef..9d6b7b2 100644 --- a/libs/ktem/ktem/index/ui.py +++ b/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: diff --git a/libs/ktem/ktem/llms/ui.py b/libs/ktem/ktem/llms/ui.py index 02d13a2..8b6f4ee 100644 --- a/libs/ktem/ktem/llms/ui.py +++ b/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: diff --git a/libs/ktem/ktem/utils/file.py b/libs/ktem/ktem/utils/file.py new file mode 100644 index 0000000..7ab7b36 --- /dev/null +++ b/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")