diff --git a/server/api.py b/server/api.py index 21c9c1f..ac92a9e 100644 --- a/server/api.py +++ b/server/api.py @@ -104,6 +104,11 @@ def create_app(): tags=["Knowledge Base Management"], summary="根据content中文档重建向量库,流式输出处理进度。" )(recreate_vector_store) + + # init local vector store info to database + from webui_pages.utils import init_vs_database + init_vs_database() + return app diff --git a/webui.py b/webui.py index c883b10..513379d 100644 --- a/webui.py +++ b/webui.py @@ -13,34 +13,15 @@ from webui_pages import * api = ApiRequest(base_url="http://127.0.0.1:7861", no_remote_api=False) if __name__ == "__main__": - st.set_page_config("langchain-chatglm WebUI", layout="wide") + # init local vector store info to database + init_vs_database() + + st.set_page_config("langchain-chatglm WebUI") if not chat_box.chat_inited: st.toast(f"欢迎使用 [`Langchain-Chatglm`](https://github.com/chatchat-space/langchain-chatglm) ! \n\n当前使用模型`{LLM_MODEL}`, 您可以开始提问了.") st.toast(" ") - # pages = {"对话1": {"icon": "chat", - # "func": dialogue_page, - # }, - # "对话2": {"icon": "chat", - # "func": dialogue_page, - # }, - # "对话3": {"icon": "chat", - # "func": dialogue_page, - # }, - # "新建对话": {"icon": "plus-circle", - # "func": dialogue_page, - # }, - # "---": {"icon": None, - # "func": None}, - # "知识库管理": {"icon": "database-fill-gear", - # "func": knowledge_base_page, - # }, - # "模型配置": {"icon": "gear", - # "func": model_config_page, - # } - # } - pages = {"对话": {"icon": "chat", "func": dialogue_page, }, diff --git a/webui_pages/utils.py b/webui_pages/utils.py index 19c375f..b4c2feb 100644 --- a/webui_pages/utils.py +++ b/webui_pages/utils.py @@ -614,11 +614,38 @@ def get_kb_doc_details(api: ApiRequest, kb: str) -> pd.DataFrame: return df -if __name__ == "__main__": +def init_vs_database(recreate_vs: bool = False): + ''' + init local vector store info to database + ''' from server.db.base import Base, engine + from server.db.repository.knowledge_base_repository import add_kb_to_db, kb_exists + from server.db.repository.knowledge_file_repository import add_doc_to_db + from server.knowledge_base.utils import KnowledgeFile + Base.metadata.create_all(bind=engine) + if recreate_vs: + api = ApiRequest(no_remote_api=True) + for kb in list_kbs_from_folder(): + for t in api.recreate_vector_store(kb): + print(t) + else: # add vs info to db only + for kb in list_kbs_from_folder(): + if not kb_exists(kb): + add_kb_to_db(kb, "faiss", EMBEDDING_MODEL) + for doc in list_docs_from_folder(kb): + try: + kb_file = KnowledgeFile(doc, kb) + add_doc_to_db(kb_file) + except Exception as e: + print(e) + + +if __name__ == "__main__": api = ApiRequest(no_remote_api=True) + # init vector store database + init_vs_database() # print(api.chat_fastchat( # messages=[{"role": "user", "content": "hello"}] @@ -637,8 +664,3 @@ if __name__ == "__main__": # print(t) # print(api.list_knowledge_bases()) - - # recreate all vector store - for kb in list_kbs_from_folder(): - for t in api.recreate_vector_store(kb): - print(t)