From 66905d39c48693a7b8568b3644e62cfcc108b03c Mon Sep 17 00:00:00 2001 From: "Duc Nguyen (john)" Date: Fri, 12 Apr 2024 15:41:09 +0700 Subject: [PATCH] Allow adding, updating and deleting indices (#24) * Allow adding indices * Allow deleting indices * Allow updating the indices * When there are multiple indices, group them below Indices tab * Update elem classes --- libs/ktem/ktem/assets/css/main.css | 4 +- libs/ktem/ktem/index/file/ui.py | 1 + libs/ktem/ktem/index/manager.py | 103 ++++++++++++++++---- libs/ktem/ktem/index/ui.py | 145 +++++++++++++++++++++++++++-- libs/ktem/ktem/main.py | 36 +++++-- 5 files changed, 252 insertions(+), 37 deletions(-) diff --git a/libs/ktem/ktem/assets/css/main.css b/libs/ktem/ktem/assets/css/main.css index 4757121..025baba 100644 --- a/libs/ktem/ktem/assets/css/main.css +++ b/libs/ktem/ktem/assets/css/main.css @@ -53,6 +53,7 @@ button.selected { } #chat-tab, +#indices-tab, #settings-tab, #help-tab, #resources-tab, @@ -98,7 +99,8 @@ button.selected { .setting-answer-mode-description { margin: 5px 5px 2px !important; } -*/ mark { + +mark { background-color: #1496bb; } diff --git a/libs/ktem/ktem/index/file/ui.py b/libs/ktem/ktem/index/file/ui.py index 10683f6..c166123 100644 --- a/libs/ktem/ktem/index/file/ui.py +++ b/libs/ktem/ktem/index/file/ui.py @@ -111,6 +111,7 @@ class FileIndexPage(BasePage): file_types=self._supported_file_types, file_count="multiple", container=True, + show_label=False, ) msg = self.upload_instruction() diff --git a/libs/ktem/ktem/index/manager.py b/libs/ktem/ktem/index/manager.py index 40b8a34..35d17cf 100644 --- a/libs/ktem/ktem/index/manager.py +++ b/libs/ktem/ktem/index/manager.py @@ -23,17 +23,14 @@ class IndexManager: def __init__(self, app): self._app = app self._indices = [] - self._index_types = {} + self._index_types: dict[str, Type[BaseIndex]] = {} - def add_index_type(self, cls: Type[BaseIndex]): - """Register index type to the system""" - self._index_types[cls.__name__] = cls - - def list_index_types(self) -> dict: + @property + def index_types(self) -> dict: """List the index_type of the index""" return self._index_types - def build_index(self, name: str, config: dict, index_type: str, id=None): + def build_index(self, name: str, config: dict, index_type: str): """Build the index Building the index simply means recording the index information into the @@ -49,22 +46,48 @@ class IndexManager: Returns: BaseIndex: the index object """ - index_cls = import_dotted_string(index_type, safe=False) - index = index_cls(app=self._app, id=id, name=name, config=config) - index.on_create() with Session(engine) as sess: - index_entry = Index( - id=index.id, name=index.name, config=index.config, index_type=index_type - ) - sess.add(index_entry) + entry = Index(name=name, config=config, index_type=index_type) + sess.add(entry) sess.commit() - sess.refresh(index_entry) + sess.refresh(entry) - index.id = index_entry.id + try: + # build the index + index_cls = import_dotted_string(index_type, safe=False) + index = index_cls(app=self._app, id=entry.id, name=name, config=config) + index.on_create() + except Exception as e: + sess.delete(entry) + sess.commit() + raise ValueError(f'Cannot create index "{name}": {e}') return index + def update_index(self, id: int, name: str, config: dict): + """Update the index information + + Args: + id: the id of the index + name: the new name of the index + config: the new config of the index + """ + with Session(engine) as sess: + entry = sess.get(Index, id) + if entry is None: + raise ValueError(f"Index with id {id} does not exist") + + entry.name = name + entry.config = config + sess.commit() + + for index in self._indices: + if index.id == id: + index.name = name + index.config = config + break + def start_index(self, id: int, name: str, config: dict, index_type: str): """Start the index @@ -81,6 +104,50 @@ class IndexManager: self._indices.append(index) return index + def delete_index(self, id: int): + """Delete the index from the database""" + index: Optional[BaseIndex] = None + for _ in self._indices: + if _.id == id: + index = _ + break + + if index is None: + raise ValueError( + "Index does not exist. If you have already removed the index, " + "please restart to reflect the changes." + ) + + try: + # clean up + index.on_delete() + + # remove from database + with Session(engine) as sess: + item = sess.query(Index).filter_by(id=id).first() + sess.delete(item) + sess.commit() + + new_indices = [_ for _ in self._indices if _.id != id] + self._indices = new_indices + except Exception as e: + raise ValueError(f"Cannot delete index {index.name}: {e}") + + def load_index_types(self): + """Load the supported index types""" + self._index_types = {} + + # built-in index types + from .file.index import FileIndex + + for index in [FileIndex]: + self._index_types[f"{index.__module__}.{index.__qualname__}"] = index + + # developer-defined custom index types + for index_str in settings.KH_INDEX_TYPES: + cls: Type[BaseIndex] = import_dotted_string(index_str, safe=False) + self._index_types[f"{cls.__module__}.{cls.__qualname__}"] = cls + def exists(self, id: Optional[int] = None, name: Optional[str] = None) -> bool: """Check if the index exists @@ -107,9 +174,7 @@ class IndexManager: Load the index from database """ - for index in settings.KH_INDEX_TYPES: - index_cls = import_dotted_string(index, safe=False) - self.add_index_type(index_cls) + self.load_index_types() for index in settings.KH_INDICES: if not self.exists(index["id"]): diff --git a/libs/ktem/ktem/index/ui.py b/libs/ktem/ktem/index/ui.py index 9c7d4dc..6f553ef 100644 --- a/libs/ktem/ktem/index/ui.py +++ b/libs/ktem/ktem/index/ui.py @@ -3,6 +3,8 @@ import pandas as pd import yaml from ktem.app import BasePage +from .manager import IndexManager + def format_description(cls): user_settings = cls.get_admin_settings() @@ -17,7 +19,7 @@ def format_description(cls): class IndexManagement(BasePage): def __init__(self, app): self._app = app - self.manager = app.index_manager + self.manager: IndexManager = app.index_manager self.spec_desc_default = ( "# Spec description\n\nSelect an index to view the spec description." ) @@ -38,16 +40,15 @@ class IndexManagement(BasePage): label="Index name", ) self.edit_spec = gr.Textbox( - label="Specification", - info="Specification of the Index in YAML format", + label="Index config", + info="Admin configuration of the Index in YAML format", lines=10, ) gr.Markdown( - "IMPORTANT: Changing or deleting the name or " - "specification of the index will require restarting " - "the system. Some settings will require rebuilding " - "the index." + "IMPORTANT: Changing or deleting the index will require " + "restarting the system. Some config settings will require " + "rebuilding the index for the index to work properly." ) with gr.Row(): self.btn_edit_save = gr.Button( @@ -68,6 +69,27 @@ class IndexManagement(BasePage): with gr.Column(): self.edit_spec_desc = gr.Markdown("# Spec description") + with gr.Tab(label="Add"): + with gr.Row(): + with gr.Column(scale=2): + self.name = gr.Textbox( + label="Index name", + info="Must be unique and non-empty.", + ) + self.index_type = gr.Dropdown(label="Index type") + self.spec = gr.Textbox( + label="Specification", + info="Specification of the index in YAML format.", + ) + gr.Markdown( + "Note: " + "After creating index, please restart the app" + ) + self.btn_new = gr.Button("Add", variant="primary") + + with gr.Column(scale=3): + self.spec_desc = gr.Markdown(self.spec_desc_default) + def _on_app_created(self): """Called when the app is created""" self._app.app.load( @@ -75,8 +97,34 @@ class IndexManagement(BasePage): inputs=None, outputs=[self.index_list], ) + self._app.app.load( + lambda: gr.update( + choices=[ + (key.split(".")[-1], key) for key in self.manager.index_types.keys() + ] + ), + outputs=[self.index_type], + ) def on_register_events(self): + self.index_type.select( + self.on_index_type_change, + inputs=[self.index_type], + outputs=[self.spec, self.spec_desc], + ) + self.btn_new.click( + self.create_index, + inputs=[self.name, self.index_type, self.spec], + outputs=None, + ).success(self.list_indices, inputs=None, outputs=[self.index_list]).success( + lambda: ("", None, "", self.spec_desc_default), + outputs=[ + self.name, + self.index_type, + self.spec, + self.spec_desc, + ], + ) self.index_list.select( self.select_index, inputs=self.index_list, @@ -85,7 +133,7 @@ class IndexManagement(BasePage): ) self.selected_index_id.change( - self.on_change_selected_index, + self.on_selected_index_change, inputs=[self.selected_index_id], outputs=[ self._selected_panel, @@ -112,6 +160,16 @@ class IndexManagement(BasePage): ], show_progress="hidden", ) + self.btn_delete_yes.click( + self.delete_index, + inputs=[self.selected_index_id], + outputs=[self.selected_index_id], + show_progress="hidden", + ).then( + self.list_indices, + inputs=None, + outputs=[self.index_list], + ) self.btn_delete_no.click( lambda: ( gr.update(visible=True), @@ -128,11 +186,57 @@ class IndexManagement(BasePage): ], show_progress="hidden", ) + self.btn_edit_save.click( + self.update_index, + inputs=[ + self.selected_index_id, + self.edit_name, + self.edit_spec, + ], + show_progress="hidden", + ).then( + self.list_indices, + inputs=None, + outputs=[self.index_list], + ) self.btn_close.click( lambda: -1, outputs=[self.selected_index_id], ) + def on_index_type_change(self, index_type: str): + """Update the spec description and pre-fill the default values + + Args: + index_type: the name of the index type, this is usually the class name + + Returns: + A tuple of the default spec and the description + """ + index_type_cls = self.manager.index_types[index_type] + required: dict = { + key: value.get("value", None) + for key, value in index_type_cls.get_admin_settings().items() + } + + return yaml.dump(required, sort_keys=False), format_description(index_type_cls) + + def create_index(self, name: str, index_type: str, config: str): + """Create the index + + Args: + name: the name of the index + index_type: the type of the index + config: the expected config of the index + """ + try: + self.manager.build_index( + name=name, config=yaml.safe_load(config), index_type=index_type + ) + gr.Info(f'Create index "{name}" successfully. Please restart the app!') + except Exception as e: + raise gr.Error(f"Failed to create Embedding model {name}: {e}") + def list_indices(self): """List the indices constructed by the user""" items = [] @@ -163,7 +267,12 @@ class IndexManagement(BasePage): return int(index_list["ID"][ev.index[0]]) - def on_change_selected_index(self, selected_index_id: int): + def on_selected_index_change(self, selected_index_id: int): + """Show the relevant index as user selects it on the UI + + Args: + selected_index_id: the id of the selected index + """ if selected_index_id == -1: _selected_panel = gr.update(visible=False) edit_spec = gr.update(value="") @@ -182,3 +291,21 @@ class IndexManagement(BasePage): edit_spec_desc, edit_name, ) + + def update_index(self, selected_index_id: int, name: str, config: str): + try: + spec = yaml.safe_load(config) + self.manager.update_index(selected_index_id, name, spec) + gr.Info(f'Update index "{name}" successfully. Please restart the app!') + except Exception as e: + raise gr.Error(f'Failed to save index "{name}": {e}') + + def delete_index(self, selected_index_id): + try: + self.manager.delete_index(selected_index_id) + gr.Info("Delete index successfully. Please restart the app!") + except Exception as e: + gr.Warning(f"Fail to delete index: {e}") + return selected_index_id + + return -1 diff --git a/libs/ktem/ktem/main.py b/libs/ktem/ktem/main.py index 1a6908d..6182e39 100644 --- a/libs/ktem/ktem/main.py +++ b/libs/ktem/ktem/main.py @@ -41,16 +41,36 @@ class App(BaseApp): ) as self._tabs["chat-tab"]: self.chat_page = ChatPage(self) - for index in self.index_manager.indices: + if len(self.index_manager.indices) == 1: + for index in self.index_manager.indices: + with gr.Tab( + f"{index.name} Index", + elem_id="indices-tab", + elem_classes=[ + "fill-main-area-height", + "scrollable", + "indices-tab", + ], + id="indices-tab", + visible=not self.f_user_management, + ) as self._tabs[f"{index.id}-tab"]: + page = index.get_index_page_ui() + setattr(self, f"_index_{index.id}", page) + elif len(self.index_manager.indices) > 1: with gr.Tab( - f"{index.name} Index", - elem_id=f"{index.id}-tab", - elem_classes="indices-tab", - id=f"{index.id}-tab", + "Indices", + elem_id="indices-tab", + elem_classes=["fill-main-area-height", "scrollable", "indices-tab"], + id="indices-tab", visible=not self.f_user_management, - ) as self._tabs[f"{index.id}-tab"]: - page = index.get_index_page_ui() - setattr(self, f"_index_{index.id}", page) + ) as self._tabs["indices-tab"]: + for index in self.index_manager.indices: + with gr.Tab( + f"{index.name}", + elem_id=f"{index.id}-tab", + ) as self._tabs[f"{index.id}-tab"]: + page = index.get_index_page_ui() + setattr(self, f"_index_{index.id}", page) with gr.Tab( "Resources",