Skip to content

Commit

Permalink
feat: add test connection feature (#166)
Browse files Browse the repository at this point in the history 
* feat: add test connection feature

* fix: typo
  • Loading branch information
@phv2312
phv2312 authored Sep 1, 2024
1 parent c1e8c37 commit 041d229
Show file tree
Hide file tree
Showing 2 changed files with 112 additions and 2 deletions.
58 changes: 57 additions & 1 deletion libs/ktem/ktem/embeddings/ui.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,18 @@ def on_building_ui(self):
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(
Expand Down Expand Up @@ -174,9 +186,11 @@ def on_register_events(self):
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=[],
Expand Down Expand Up @@ -221,6 +235,12 @@ def on_register_events(self):
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)
Expand Down Expand Up @@ -266,6 +286,7 @@ def select_emb(self, emb_list, ev: gr.SelectData):

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)
Expand All @@ -275,6 +296,7 @@ def on_selected_emb_change(self, selected_emb_name):
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)
Expand All @@ -298,6 +320,7 @@ def on_selected_emb_change(self, selected_emb_name):
edit_spec,
edit_spec_desc,
edit_default,
_check_connection_panel,
)

def on_btn_delete_click(self):
Expand All @@ -307,6 +330,39 @@ def on_btn_delete_click(self):

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}<br>"
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`<br>"
yield log_content
_ = emb("Hi")

log_content += (
"<mark style='background: yellow; color: red'>- Connection success. "
"</mark><br>"
)
yield log_content

gr.Info(f"Embedding {selected_emb_name} connect successfully")
except Exception as e:
print(e)
log_content += (
f"<mark style='color: yellow; background: red'>- Connection failed. "
f"Got error:\n {str(e)}</mark>"
)
yield log_content

return log_content

def save_emb(self, selected_emb_name, default, spec):
try:
spec = yaml.load(spec, Loader=YAMLNoDateSafeLoader)
Expand Down
56 changes: 55 additions & 1 deletion libs/ktem/ktem/llms/ui.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,16 @@ def on_building_ui(self):
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(
Expand Down Expand Up @@ -171,9 +181,12 @@ def on_register_events(self):
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=[],
Expand Down Expand Up @@ -218,6 +231,12 @@ def on_register_events(self):
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)
Expand Down Expand Up @@ -263,6 +282,7 @@ def select_llm(self, llm_list, ev: gr.SelectData):

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)
Expand All @@ -272,6 +292,7 @@ def on_selected_llm_change(self, selected_llm_name):
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)
Expand All @@ -295,6 +316,7 @@ def on_selected_llm_change(self, selected_llm_name):
edit_spec,
edit_spec_desc,
edit_default,
_check_connection_panel,
)

def on_btn_delete_click(self):
Expand All @@ -304,6 +326,38 @@ def on_btn_delete_click(self):

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}<br>"
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`<br>"
yield log_content
respond = llm("Hi")

log_content += (
f"<mark style='background: yellow; color: red'>- Connection success. "
f"Got response:\n {respond}</mark><br>"
)
yield log_content

gr.Info(f"LLM {selected_llm_name} connect successfully")
except Exception as e:
log_content += (
f"<mark style='color: yellow; background: red'>- Connection failed. "
f"Got error:\n {e}</mark>"
)
yield log_content

return log_content

def save_llm(self, selected_llm_name, default, spec):
try:
spec = yaml.load(spec, Loader=YAMLNoDateSafeLoader)
Expand Down

0 comments on commit 041d229

Please sign in to comment.