Public/kotaemon
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Allow adding, updating and deleting indices (#24)

Browse Source 
* Allow adding indices

* Allow deleting indices

* Allow updating the indices

* When there are multiple indices, group them below Indices tab

* Update elem classes
This commit is contained in:
Duc Nguyen (john) 2024-04-12 15:41:09 +07:00 committed by GitHub
parent 4efe9c02a8
commit 66905d39c4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 252 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
libs/ktem/ktem/assets/css/main.css
View File

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

1
libs/ktem/ktem/index/file/ui.py
View File

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

103
libs/ktem/ktem/index/manager.py
View File

@ -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"]):

145
libs/ktem/ktem/index/ui.py
View File

@ -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(
"<mark>Note</mark>: "
"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

36
libs/ktem/ktem/main.py
View File

@ -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",