codefuse-chatbot/examples/webui/document.py

378 lines
15 KiB
Python
Raw Permalink Normal View History

2023-09-28 10:58:58 +08:00
import streamlit as st
import os
import time
import traceback
from typing import Literal, Dict, Tuple
from st_aggrid import AgGrid, JsCode
from st_aggrid.grid_options_builder import GridOptionsBuilder
import pandas as pd
from .utils import *
from muagent.utils.path_utils import *
from muagent.service.service_factory import get_kb_details, get_kb_doc_details
from muagent.orm import table_init
2023-09-28 10:58:58 +08:00
from configs.model_config import (
KB_ROOT_PATH, kbs_config, DEFAULT_VS_TYPE, WEB_CRAWL_PATH,
EMBEDDING_DEVICE, EMBEDDING_ENGINE, EMBEDDING_MODEL, embedding_model_dict,
llm_model_dict
)
2023-09-28 10:58:58 +08:00
# SENTENCE_SIZE = 100
cell_renderer = JsCode("""function(params) {if(params.value==true){return ''}else{return '×'}}""")
def config_aggrid(
df: pd.DataFrame,
columns: Dict[Tuple[str, str], Dict] = {},
selection_mode: Literal["single", "multiple", "disabled"] = "single",
use_checkbox: bool = False,
) -> GridOptionsBuilder:
gb = GridOptionsBuilder.from_dataframe(df)
gb.configure_column("No", width=40)
for (col, header), kw in columns.items():
gb.configure_column(col, header, wrapHeaderText=True, **kw)
gb.configure_selection(
selection_mode=selection_mode,
use_checkbox=use_checkbox,
# pre_selected_rows=st.session_state.get("selected_rows", [0]),
)
return gb
def file_exists(kb: str, selected_rows: List) -> Tuple[str, str]:
'''
check whether a doc file exists in local knowledge base folder.
return the file's name and path if it exists.
'''
if selected_rows:
file_name = selected_rows[0]["file_name"]
file_path = get_file_path(kb, file_name, KB_ROOT_PATH)
2023-09-28 10:58:58 +08:00
if os.path.isfile(file_path):
return file_name, file_path
return "", ""
def knowledge_page(
api: ApiRequest,
embedding_model_dict: dict = embedding_model_dict,
kbs_config: dict = kbs_config,
embedding_model: str = EMBEDDING_MODEL,
default_vs_type: str = DEFAULT_VS_TYPE,
web_crawl_path: str = WEB_CRAWL_PATH
):
2023-09-28 10:58:58 +08:00
# 判断表是否存在并进行初始化
table_init()
try:
kb_list = {x["kb_name"]: x for x in get_kb_details(KB_ROOT_PATH)}
2023-09-28 10:58:58 +08:00
except Exception as e:
st.error("获取知识库信息错误,请检查是否已按照 `README.md` 中 `4 知识库初始化与迁移` 步骤完成初始化或迁移,或是否为数据库连接错误。")
st.stop()
kb_names = list(kb_list.keys())
if "selected_kb_name" in st.session_state and st.session_state["selected_kb_name"] in kb_names:
selected_kb_index = kb_names.index(st.session_state["selected_kb_name"])
else:
selected_kb_index = 0
def format_selected_kb(kb_name: str) -> str:
if kb := kb_list.get(kb_name):
return f"{kb_name} ({kb['vs_type']} @ {kb['embed_model']})"
else:
return kb_name
selected_kb = st.selectbox(
"请选择或新建知识库:",
kb_names + ["新建知识库"],
format_func=format_selected_kb,
index=selected_kb_index
)
if selected_kb == "新建知识库":
with st.form("新建知识库"):
kb_name = st.text_input(
"新建知识库名称",
placeholder="新知识库名称,不支持中文命名",
key="kb_name",
)
cols = st.columns(2)
vs_types = list(kbs_config.keys())
vs_type = cols[0].selectbox(
"向量库类型",
vs_types,
index=vs_types.index(default_vs_type),
2023-09-28 10:58:58 +08:00
key="vs_type",
)
embed_models = list(embedding_model_dict.keys())
embed_model = cols[1].selectbox(
"Embedding 模型",
embed_models,
index=embed_models.index(embedding_model),
2023-09-28 10:58:58 +08:00
key="embed_model",
)
submit_create_kb = st.form_submit_button(
"新建",
# disabled=not bool(kb_name),
use_container_width=True,
)
if submit_create_kb:
if not kb_name or not kb_name.strip():
st.error(f"知识库名称不能为空!")
elif kb_name in kb_list:
st.error(f"名为 {kb_name} 的知识库已经存在!")
else:
ret = api.create_knowledge_base(
knowledge_base_name=kb_name,
vector_store_type=vs_type,
embed_model=embed_model,
embed_engine=EMBEDDING_ENGINE,
embedding_device= EMBEDDING_DEVICE,
embed_model_path=embedding_model_dict[embed_model],
api_key=llm_model_dict[LLM_MODEL]["api_key"],
api_base_url=llm_model_dict[LLM_MODEL]["api_base_url"],
2023-09-28 10:58:58 +08:00
)
st.toast(ret.get("msg", " "))
st.session_state["selected_kb_name"] = kb_name
st.experimental_rerun()
elif selected_kb:
kb = selected_kb
# 上传文件
# sentence_size = st.slider("文本入库分句长度限制", 1, 1000, SENTENCE_SIZE, disabled=True)
files = st.file_uploader("上传知识文件",
[i for ls in LOADER2EXT_DICT.values() for i in ls],
accept_multiple_files=True,
)
if st.button(
"添加文件到知识库",
# help="请先上传文件,再点击添加",
# use_container_width=True,
disabled=len(files) == 0,
):
data = [{"file": f, "knowledge_base_name": kb, "not_refresh_vs_cache": True, "embed_model": EMBEDDING_MODEL,
"embed_model_path": embedding_model_dict[EMBEDDING_MODEL],
"model_device": EMBEDDING_DEVICE,
"embed_engine": EMBEDDING_ENGINE,
"api_key": llm_model_dict[LLM_MODEL]["api_key"],
"api_base_url": llm_model_dict[LLM_MODEL]["api_base_url"],
}
for f in files]
data[-1]["not_refresh_vs_cache"]=False
for k in data:
pass
ret = api.upload_kb_doc(**k)
if msg := check_success_msg(ret):
st.toast(msg, icon="")
elif msg := check_error_msg(ret):
st.toast(msg, icon="")
st.session_state.files = []
2023-09-28 10:58:58 +08:00
base_url = st.text_input(
"待获取内容的URL地址",
placeholder="请填写正确可打开的URL地址",
key="base_url",
)
if st.button(
"添加URL内容到知识库",
disabled= base_url is None or base_url=="",
):
filename = base_url.replace("https://", " ").\
replace("http://", " ").replace("/", " ").\
replace("?", " ").replace("=", " ").replace(".", " ").strip()
html_name = "_".join(filename.split(" ",) + ["html.jsonl"])
text_name = "_".join(filename.split(" ",) + ["text.jsonl"])
html_path = os.path.join(web_crawl_path, html_name,)
text_path = os.path.join(web_crawl_path, text_name,)
2023-09-28 10:58:58 +08:00
# if not os.path.exists(text_dir) or :
st.toast(base_url)
st.toast(html_path)
st.toast(text_path)
res = api.web_crawl(
base_url=base_url,
html_dir=html_path,
text_dir=text_path,
do_dfs = False,
reptile_lib="requests",
method="get",
time_sleep=2,
)
if res["status"] == 200:
st.toast(res["response"], icon="")
data = [{"file": text_path, "filename": text_name, "knowledge_base_name": kb, "not_refresh_vs_cache": False,
"embed_model": EMBEDDING_MODEL,
"embed_model_path": embedding_model_dict[EMBEDDING_MODEL],
"model_device": EMBEDDING_DEVICE,
"embed_engine": EMBEDDING_ENGINE,
"api_key": llm_model_dict[LLM_MODEL]["api_key"],
"api_base_url": llm_model_dict[LLM_MODEL]["api_base_url"],}]
2023-09-28 10:58:58 +08:00
for k in data:
ret = api.upload_kb_doc(**k)
logger.info(ret)
if msg := check_success_msg(ret):
st.toast(msg, icon="")
elif msg := check_error_msg(ret):
st.toast(msg, icon="")
st.session_state.files = []
else:
st.toast(res["response"], icon="")
if os.path.exists(html_path):
os.remove(html_path)
st.divider()
# 知识库详情
# st.info("请选择文件,点击按钮进行操作。")
doc_details = pd.DataFrame(get_kb_doc_details(kb, KB_ROOT_PATH))
2023-09-28 10:58:58 +08:00
if not len(doc_details):
st.info(f"知识库 `{kb}` 中暂无文件")
else:
st.write(f"知识库 `{kb}` 中已有文件:")
st.info("知识库中包含源文件与向量库,请从下表中选择文件后操作")
doc_details.drop(columns=["kb_name"], inplace=True)
doc_details = doc_details[[
"No", "file_name", "document_loader", "text_splitter", "in_folder", "in_db",
]]
# doc_details["in_folder"] = doc_details["in_folder"].replace(True, "✓").replace(False, "×")
# doc_details["in_db"] = doc_details["in_db"].replace(True, "✓").replace(False, "×")
gb = config_aggrid(
doc_details,
{
("No", "序号"): {},
("file_name", "文档名称"): {},
# ("file_ext", "文档类型"): {},
# ("file_version", "文档版本"): {},
("document_loader", "文档加载器"): {},
("text_splitter", "分词器"): {},
# ("create_time", "创建时间"): {},
("in_folder", "源文件"): {"cellRenderer": cell_renderer},
("in_db", "向量库"): {"cellRenderer": cell_renderer},
},
"multiple",
)
doc_grid = AgGrid(
doc_details,
gb.build(),
columns_auto_size_mode="FIT_CONTENTS",
theme="alpine",
custom_css={
"#gridToolBar": {"display": "none"},
},
allow_unsafe_jscode=True
)
selected_rows = doc_grid.get("selected_rows", [])
cols = st.columns(4)
file_name, file_path = file_exists(kb, selected_rows)
if file_path:
with open(file_path, "rb") as fp:
cols[0].download_button(
"下载选中文档",
fp,
file_name=file_name,
use_container_width=True, )
else:
cols[0].download_button(
"下载选中文档",
"",
disabled=True,
use_container_width=True, )
st.write()
# 将文件分词并加载到向量库中
if cols[1].button(
"重新添加至向量库" if selected_rows and (pd.DataFrame(selected_rows)["in_db"]).any() else "添加至向量库",
disabled=not file_exists(kb, selected_rows)[0],
use_container_width=True,
):
for row in selected_rows:
api.update_kb_doc(kb, row["file_name"],
embed_engine=EMBEDDING_ENGINE,embed_model=EMBEDDING_MODEL,
embed_model_path=embedding_model_dict[EMBEDDING_MODEL],
model_device=EMBEDDING_DEVICE,
api_key=llm_model_dict[LLM_MODEL]["api_key"],
api_base_url=llm_model_dict[LLM_MODEL]["api_base_url"],
)
2023-09-28 10:58:58 +08:00
st.experimental_rerun()
# 将文件从向量库中删除,但不删除文件本身。
if cols[2].button(
"从向量库删除",
disabled=not (selected_rows and selected_rows[0]["in_db"]),
use_container_width=True,
):
for row in selected_rows:
api.delete_kb_doc(kb, row["file_name"],
embed_engine=EMBEDDING_ENGINE,embed_model=EMBEDDING_MODEL,
embed_model_path=embedding_model_dict[EMBEDDING_MODEL],
model_device=EMBEDDING_DEVICE,
api_key=llm_model_dict[LLM_MODEL]["api_key"],
api_base_url=llm_model_dict[LLM_MODEL]["api_base_url"],)
2023-09-28 10:58:58 +08:00
st.experimental_rerun()
if cols[3].button(
"从知识库中删除",
type="primary",
use_container_width=True,
):
for row in selected_rows:
ret = api.delete_kb_doc(kb, row["file_name"], True,
embed_engine=EMBEDDING_ENGINE,embed_model=EMBEDDING_MODEL,
embed_model_path=embedding_model_dict[EMBEDDING_MODEL],
model_device=EMBEDDING_DEVICE,
api_key=llm_model_dict[LLM_MODEL]["api_key"],
api_base_url=llm_model_dict[LLM_MODEL]["api_base_url"],)
2023-09-28 10:58:58 +08:00
st.toast(ret.get("msg", " "))
st.experimental_rerun()
st.divider()
cols = st.columns(3)
# todo: freezed
if cols[0].button(
"依据源文件重建向量库",
# help="无需上传文件通过其它方式将文档拷贝到对应知识库content目录下点击本按钮即可重建知识库。",
use_container_width=True,
type="primary",
):
with st.spinner("向量库重构中,请耐心等待,勿刷新或关闭页面。"):
empty = st.empty()
empty.progress(0.0, "")
for d in api.recreate_vector_store(
kb, vs_type=default_vs_type, embed_model=embedding_model, embedding_device=EMBEDDING_DEVICE,
embed_model_path=embedding_model_dict[embedding_model], embed_engine=EMBEDDING_ENGINE,
api_key=llm_model_dict[LLM_MODEL]["api_key"],
api_base_url=llm_model_dict[LLM_MODEL]["api_base_url"],
):
2023-09-28 10:58:58 +08:00
if msg := check_error_msg(d):
st.toast(msg)
else:
empty.progress(d["finished"] / d["total"], f"正在处理: {d['doc']}")
st.experimental_rerun()
if cols[2].button(
"删除知识库",
use_container_width=True,
):
ret = api.delete_knowledge_base(kb,)
2023-09-28 10:58:58 +08:00
st.toast(ret.get("msg", " "))
time.sleep(1)
st.experimental_rerun()