From 041d22928212ab012461e32cc8f21a9f65387964 Mon Sep 17 00:00:00 2001 From: kan_cin Date: Sun, 1 Sep 2024 08:22:36 +0700 Subject: [PATCH] feat: add test connection feature (#166) * feat: add test connection feature * fix: typo --- libs/ktem/ktem/embeddings/ui.py | 58 ++++++++++++++++++++++++++++++++- libs/ktem/ktem/llms/ui.py | 56 ++++++++++++++++++++++++++++++- 2 files changed, 112 insertions(+), 2 deletions(-) diff --git a/libs/ktem/ktem/embeddings/ui.py b/libs/ktem/ktem/embeddings/ui.py index 1b2e549..1f119d8 100644 --- a/libs/ktem/ktem/embeddings/ui.py +++ b/libs/ktem/ktem/embeddings/ui.py @@ -52,6 +52,18 @@ class EmbeddingManagement(BasePage): lines=10, ) + with gr.Accordion( + label="Test connection", visible=False, open=False + ) as self._check_connection_panel: + with gr.Row(): + with gr.Column(scale=4): + self.connection_logs = gr.HTML( + "Logs", + ) + + with gr.Column(scale=1): + self.btn_test_connection = gr.Button("Test") + with gr.Row(visible=False) as self._selected_panel_btn: with gr.Column(): self.btn_edit_save = gr.Button( @@ -174,9 +186,11 @@ class EmbeddingManagement(BasePage): self.edit_spec, self.edit_spec_desc, self.edit_default, + self._check_connection_panel, ], show_progress="hidden", - ) + ).success(lambda: gr.update(value=""), outputs=[self.connection_logs]) + self.btn_delete.click( self.on_btn_delete_click, inputs=[], @@ -221,6 +235,12 @@ class EmbeddingManagement(BasePage): outputs=[self.selected_emb_name], ) + self.btn_test_connection.click( + self.check_connection, + inputs=[self.selected_emb_name], + outputs=[self.connection_logs], + ) + def create_emb(self, name, choices, spec, default): try: spec = yaml.load(spec, Loader=YAMLNoDateSafeLoader) @@ -266,6 +286,7 @@ class EmbeddingManagement(BasePage): def on_selected_emb_change(self, selected_emb_name): if selected_emb_name == "": + _check_connection_panel = gr.update(visible=False) _selected_panel = gr.update(visible=False) _selected_panel_btn = gr.update(visible=False) btn_delete = gr.update(visible=True) @@ -275,6 +296,7 @@ class EmbeddingManagement(BasePage): edit_spec_desc = gr.update(value="") edit_default = gr.update(value=False) else: + _check_connection_panel = gr.update(visible=True) _selected_panel = gr.update(visible=True) _selected_panel_btn = gr.update(visible=True) btn_delete = gr.update(visible=True) @@ -298,6 +320,7 @@ class EmbeddingManagement(BasePage): edit_spec, edit_spec_desc, edit_default, + _check_connection_panel, ) def on_btn_delete_click(self): @@ -307,6 +330,39 @@ class EmbeddingManagement(BasePage): return btn_delete, btn_delete_yes, btn_delete_no + def check_connection(self, selected_emb_name): + log_content: str = "" + + try: + log_content += f"- Testing model: {selected_emb_name}
" + yield log_content + + emb = embedding_models_manager.get(selected_emb_name) + + if emb is None: + raise Exception(f"Can not found model: {selected_emb_name}") + + log_content += "- Sending a message `Hi`
" + yield log_content + _ = emb("Hi") + + log_content += ( + "- Connection success. " + "
" + ) + yield log_content + + gr.Info(f"Embedding {selected_emb_name} connect successfully") + except Exception as e: + print(e) + log_content += ( + f"- Connection failed. " + f"Got error:\n {str(e)}" + ) + yield log_content + + return log_content + def save_emb(self, selected_emb_name, default, spec): try: spec = yaml.load(spec, Loader=YAMLNoDateSafeLoader) diff --git a/libs/ktem/ktem/llms/ui.py b/libs/ktem/ktem/llms/ui.py index 116982e..7296c06 100644 --- a/libs/ktem/ktem/llms/ui.py +++ b/libs/ktem/ktem/llms/ui.py @@ -51,6 +51,16 @@ class LLMManagement(BasePage): lines=10, ) + with gr.Accordion( + label="Test connection", visible=False, open=False + ) as self._check_connection_panel: + with gr.Row(): + with gr.Column(scale=4): + self.connection_logs = gr.HTML("Logs") + + with gr.Column(scale=1): + self.btn_test_connection = gr.Button("Test") + with gr.Row(visible=False) as self._selected_panel_btn: with gr.Column(): self.btn_edit_save = gr.Button( @@ -171,9 +181,12 @@ class LLMManagement(BasePage): self.edit_spec, self.edit_spec_desc, self.edit_default, + # check connection panel + self._check_connection_panel, ], show_progress="hidden", - ) + ).success(lambda: gr.update(value=""), outputs=[self.connection_logs]) + self.btn_delete.click( self.on_btn_delete_click, inputs=[], @@ -218,6 +231,12 @@ class LLMManagement(BasePage): outputs=[self.selected_llm_name], ) + self.btn_test_connection.click( + self.check_connection, + inputs=[self.selected_llm_name], + outputs=[self.connection_logs], + ) + def create_llm(self, name, choices, spec, default): try: spec = yaml.load(spec, Loader=YAMLNoDateSafeLoader) @@ -263,6 +282,7 @@ class LLMManagement(BasePage): def on_selected_llm_change(self, selected_llm_name): if selected_llm_name == "": + _check_connection_panel = gr.update(visible=False) _selected_panel = gr.update(visible=False) _selected_panel_btn = gr.update(visible=False) btn_delete = gr.update(visible=True) @@ -272,6 +292,7 @@ class LLMManagement(BasePage): edit_spec_desc = gr.update(value="") edit_default = gr.update(value=False) else: + _check_connection_panel = gr.update(visible=True) _selected_panel = gr.update(visible=True) _selected_panel_btn = gr.update(visible=True) btn_delete = gr.update(visible=True) @@ -295,6 +316,7 @@ class LLMManagement(BasePage): edit_spec, edit_spec_desc, edit_default, + _check_connection_panel, ) def on_btn_delete_click(self): @@ -304,6 +326,38 @@ class LLMManagement(BasePage): return btn_delete, btn_delete_yes, btn_delete_no + def check_connection(self, selected_llm_name: str): + log_content: str = "" + + try: + log_content += f"- Testing model: {selected_llm_name}
" + yield log_content + + llm = llms.get(key=selected_llm_name, default=None) + + if llm is None: + raise Exception(f"Can not found model: {selected_llm_name}") + + log_content += "- Sending a message `Hi`
" + yield log_content + respond = llm("Hi") + + log_content += ( + f"- Connection success. " + f"Got response:\n {respond}
" + ) + yield log_content + + gr.Info(f"LLM {selected_llm_name} connect successfully") + except Exception as e: + log_content += ( + f"- Connection failed. " + f"Got error:\n {e}" + ) + yield log_content + + return log_content + def save_llm(self, selected_llm_name, default, spec): try: spec = yaml.load(spec, Loader=YAMLNoDateSafeLoader)