From dbb6bb275f89bdb9e897035e36ec49db32fe201b Mon Sep 17 00:00:00 2001
From: kan_cin <kan@cinnamon.is>
Date: Sun, 8 Sep 2024 10:55:13 +0700
Subject: [PATCH] feat: add test connection for edit spec (#239)

---
 libs/ktem/ktem/embeddings/ui.py | 15 +++++++++++----
 libs/ktem/ktem/llms/ui.py       | 18 ++++++++++++++----
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/libs/ktem/ktem/embeddings/ui.py b/libs/ktem/ktem/embeddings/ui.py
index 1f119d8..c464005 100644
--- a/libs/ktem/ktem/embeddings/ui.py
+++ b/libs/ktem/ktem/embeddings/ui.py
@@ -5,6 +5,7 @@ import pandas as pd
 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 @@ class EmbeddingManagement(BasePage):
 
         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 @@ class EmbeddingManagement(BasePage):
 
         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
index 7296c06..e7cd6a9 100644
--- a/libs/ktem/ktem/llms/ui.py
+++ b/libs/ktem/ktem/llms/ui.py
@@ -5,6 +5,7 @@ import pandas as pd
 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 @@ class LLMManagement(BasePage):
                                     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 @@ class LLMManagement(BasePage):
 
         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 @@ class LLMManagement(BasePage):
 
         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}")