automatically init vector store info to database when run api.py & webui.py

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
 
 

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,
                       },

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)