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)