yujunjun/kotaemon
1
0
Fork 0
mirror of https://github.com/Cinnamon/kotaemon.git synced 2025-06-26 23:19:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add test connection feature (#166)

Browse Source 
* feat: add test connection feature

* fix: typo
This commit is contained in:
kan_cin 2024-09-01 08:22:36 +07:00 committed by GitHub
parent c1e8c37e5e
commit 041d229282
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 112 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

58
libs/ktem/ktem/embeddings/ui.py
View File

@ -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}<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)

56
libs/ktem/ktem/llms/ui.py
View File

@ -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}<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)