import streamlit as st from webui_pages.utils import * # import streamlit_antd_components as sac from st_aggrid import AgGrid from st_aggrid.grid_options_builder import GridOptionsBuilder import pandas as pd from server.knowledge_base.utils import get_file_path, list_kbs_from_folder, list_docs_from_folder from server.knowledge_base.kb_service.base import KBServiceFactory from server.db.repository.knowledge_base_repository import get_kb_detail from server.db.repository.knowledge_file_repository import get_file_detail # from streamlit_chatbox import * from typing import Literal, Dict SENTENCE_SIZE = 100 def get_kb_details(api: ApiRequest) -> pd.DataFrame: kbs_in_folder = list_kbs_from_folder() kbs_in_db = api.list_knowledge_bases() result = {} for kb in kbs_in_folder: result[kb] = { "kb_name": kb, "vs_type": "", "embed_model": "", "file_count": 0, "create_time": None, "in_folder": True, "in_db": False, } for kb in kbs_in_db: kb_detail = get_kb_detail(kb) if kb_detail: kb_detail["in_db"] = True if kb in result: result[kb].update(kb_detail) else: kb_detail["in_folder"] = False result[kb] = kb_detail df = pd.DataFrame(result.values(), columns=[ "kb_name", "vs_type", "embed_model", "file_count", "create_time", "in_folder", "in_db", ]) df.insert(0, "No", range(1, len(df) + 1)) return df def get_kb_doc_details(api: ApiRequest, kb: str) -> pd.DataFrame: docs_in_folder = list_docs_from_folder(kb) docs_in_db = api.list_kb_docs(kb) result = {} for doc in docs_in_folder: result[doc] = { "kb_name": kb, "file_name": doc, "file_ext": os.path.splitext(doc)[-1], "file_version": 0, "document_loader": "", "text_splitter": "", "create_time": None, "in_folder": True, "in_db": False, } for doc in docs_in_db: doc_detail = get_file_detail(kb, doc) if doc_detail: doc_detail["in_db"] = True if doc in result: result[doc].update(doc_detail) else: doc_detail["in_folder"] = False result[doc] = doc_detail df = pd.DataFrame(result.values(), columns=[ "kb_name", "file_name", "file_ext", "file_version", "document_loader", "text_splitter", "create_time", "in_folder", "in_db", ]) df.insert(0, "No", range(1, len(df) + 1)) return df def config_aggrid( df: pd.DataFrame, titles: Dict[str, str] = {}, selection_mode: Literal["single", "multiple", "disabled"] = "single", use_checkbox: bool = False, ) -> GridOptionsBuilder: gb = GridOptionsBuilder.from_dataframe(df) gb.configure_column("No", width=50) for k, v in titles.items(): gb.configure_column(k, v, maxWidth=100, wrapHeaderText=True) gb.configure_selection(selection_mode, use_checkbox, pre_selected_rows=[0]) return gb # kb_box = ChatBox(session_key="kb_messages") def knowledge_base_page(api: ApiRequest): api = ApiRequest(base_url="http://127.0.0.1:7861", no_remote_api=True) kb_details = get_kb_details(api) kb_list = list(kb_details.kb_name) cols = st.columns([2, 1, 1]) new_kb_name = cols[0].text_input( "新知识库名称", placeholder="新知识库名称", label_visibility="collapsed", key="new_kb_name", ) if cols[1].button("新建", disabled=not bool(new_kb_name)) and new_kb_name: if new_kb_name in kb_list: st.error(f"名为 {new_kb_name} 的知识库已经存在！") else: ret = api.create_knowledge_base(new_kb_name) st.toast(ret["msg"]) st.experimental_rerun() if cols[2].button("删除", disabled=not bool(new_kb_name)) and new_kb_name: if new_kb_name in kb_list: ret = api.delete_knowledge_base(new_kb_name) st.toast(ret["msg"]) st.experimental_rerun() else: st.error(f"名为 {new_kb_name} 的知识库不存在！") st.write("知识库：") if kb_list: gb = config_aggrid( kb_details, { "kb_name": "知识库名称", "vs_type": "知识库类型", "embed_model": "嵌入模型", "file_count": "文档数量", "create_time": "创建时间", "in_folder": "存在于文件夹", "in_db": "存在于数据库", } ) kb_grid = AgGrid(kb_details, gb.build()) if kb_grid.selected_rows: kb = kb_grid.selected_rows[0]["kb_name"] with st.sidebar: sentence_size = st.slider("文本入库分句长度限制", 1, 1000, SENTENCE_SIZE, disabled=True) files = st.file_uploader("上传知识文件", ["docx", "txt", "md", "csv", "xlsx", "pdf"], accept_multiple_files=True, ) if st.button( "添加文件到知识库", help="请先上传文件，再点击添加", use_container_width=True, disabled=len(files)==0, ): for f in files: ret = api.upload_kb_doc(f, kb) if ret["code"] == 200: st.toast(ret["msg"], icon="✔") else: st.toast(ret["msg"], icon="❌") st.session_state.files = [] if st.button( "重建知识库", help="无需上传文件，通过其它方式将文档拷贝到对应知识库content目录下，点击本按钮即可重建知识库。", use_container_width=True, disabled=True, ): progress = st.progress(0.0, "") for d in api.recreate_vector_store(kb): progress.progress(d["finished"] / d["t]otal"], f"正在处理： {d['doc']}") # 知识库详情 st.subheader(f"知识库 {kb} 详情") doc_details = get_kb_doc_details(api, kb) doc_details.drop(columns=["kb_name"], inplace=True) gb = config_aggrid( doc_details, { "file_name": "文档名称", "file_ext": "文档类型", "file_version": "文档版本", "document_loader": "文档加载器", "text_splitter": "分词器", "create_time": "创建时间", "in_folder": "存在于文件夹", "in_db": "存在于数据库", }, "multiple", ) doc_grid = AgGrid(doc_details, gb.build()) cols = st.columns(3) selected_rows = doc_grid.get("selected_rows", []) cols = st.columns(4) if selected_rows: file_name = selected_rows[0]["file_name"] file_path = get_file_path(kb, file_name) with open(file_path, "rb") as fp: cols[0].download_button("下载选中文档", fp, file_name=file_name) else: cols[0].download_button("下载选中文档", "", disabled=True) if cols[1].button("入库", disabled=len(selected_rows)==0): for row in selected_rows: api.update_kb_doc(kb, row["file_name"]) st.experimental_rerun() if cols[2].button("出库", disabled=len(selected_rows)==0): for row in selected_rows: api.delete_kb_doc(kb, row["file_name"]) st.experimental_rerun() if cols[3].button("删除选中文档！", type="primary"): for row in selected_rows: ret = api.delete_kb_doc(kb, row["file_name"], True) st.toast(ret["msg"]) st.experimental_rerun() st.write("本文档包含以下知识条目：(待定内容)")