feat: add test connection for edit spec (#239)

Cinnamon · Sep 8, 2024 · dbb6bb2 · dbb6bb2
1 parent fa881d4
commit dbb6bb2
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 8 deletions.
diff --git a/libs/ktem/ktem/embeddings/ui.py b/libs/ktem/ktem/embeddings/ui.py
@@ -5,6 +5,7 @@
 import yaml
 from ktem.app import BasePage
 from ktem.utils.file import YAMLNoDateSafeLoader
+from theflow.utils.modules import deserialize
 
 from .manager import embedding_models_manager
 
@@ -237,7 +238,7 @@ def on_register_events(self):
 
         self.btn_test_connection.click(
             self.check_connection,
-            inputs=[self.selected_emb_name],
+            inputs=[self.selected_emb_name, self.edit_spec],
             outputs=[self.connection_logs],
         )
 
@@ -330,14 +331,20 @@ def on_btn_delete_click(self):
 
         return btn_delete, btn_delete_yes, btn_delete_no
 
-    def check_connection(self, selected_emb_name):
+    def check_connection(self, selected_emb_name, selected_spec):
         log_content: str = ""
-
         try:
             log_content += f"- Testing model: {selected_emb_name}<br>"
             yield log_content
 
-            emb = embedding_models_manager.get(selected_emb_name)
+            # Parse content & init model
+            info = deepcopy(embedding_models_manager.info()[selected_emb_name])
+
+            # Parse content & create dummy embedding
+            spec = yaml.load(selected_spec, Loader=YAMLNoDateSafeLoader)
+            info["spec"].update(spec)
+
+            emb = deserialize(info["spec"], safe=False)
 
             if emb is None:
                 raise Exception(f"Can not found model: {selected_emb_name}")

diff --git a/libs/ktem/ktem/llms/ui.py b/libs/ktem/ktem/llms/ui.py
@@ -5,6 +5,7 @@
 import yaml
 from ktem.app import BasePage
 from ktem.utils.file import YAMLNoDateSafeLoader
+from theflow.utils.modules import deserialize
 
 from .manager import llms
 
@@ -59,7 +60,9 @@ def on_building_ui(self):
                                     self.connection_logs = gr.HTML("Logs")
 
                                 with gr.Column(scale=1):
-                                    self.btn_test_connection = gr.Button("Test")
+                                    self.btn_test_connection = gr.Button(
+                                        "Test",
+                                    )
 
                         with gr.Row(visible=False) as self._selected_panel_btn:
                             with gr.Column():
@@ -233,7 +236,7 @@ def on_register_events(self):
 
         self.btn_test_connection.click(
             self.check_connection,
-            inputs=[self.selected_llm_name],
+            inputs=[self.selected_llm_name, self.edit_spec],
             outputs=[self.connection_logs],
         )
 
@@ -326,14 +329,21 @@ def on_btn_delete_click(self):
 
         return btn_delete, btn_delete_yes, btn_delete_no
 
-    def check_connection(self, selected_llm_name: str):
+    def check_connection(self, selected_llm_name: str, selected_spec):
         log_content: str = ""
 
         try:
             log_content += f"- Testing model: {selected_llm_name}<br>"
             yield log_content
 
-            llm = llms.get(key=selected_llm_name, default=None)
+            # Parse content & init model
+            info = deepcopy(llms.info()[selected_llm_name])
+
+            # Parse content & create dummy embedding
+            spec = yaml.load(selected_spec, Loader=YAMLNoDateSafeLoader)
+            info["spec"].update(spec)
+
+            llm = deserialize(info["spec"], safe=False)
 
             if llm is None:
                 raise Exception(f"Can not found model: {selected_llm_name}")