Public/kotaemon
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix Yaml datetime format (#79)

Browse Source
This commit is contained in:
trducng (john) 2024-05-20 17:36:14 +07:00 committed by GitHub
parent 5ca3c25404
commit 04e602161b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 35 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
libs/ktem/ktem/embeddings/ui.py
View File

@ -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__"]

7
libs/ktem/ktem/index/ui.py
View File

@ -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:

7
libs/ktem/ktem/llms/ui.py
View File

@ -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:

23
libs/ktem/ktem/utils/file.py Normal file
View File

@ -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")