codefuse-chatbot/examples/webui/document.py

378 lines
15 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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
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
)
# 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)
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
):
# 判断表是否存在并进行初始化
table_init()
try:
kb_list = {x["kb_name"]: x for x in get_kb_details(KB_ROOT_PATH)}
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),
key="vs_type",
)
embed_models = list(embedding_model_dict.keys())
embed_model = cols[1].selectbox(
"Embedding 模型",
embed_models,
index=embed_models.index(embedding_model),
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"],
)
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 = []
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,)
# 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"],}]
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))
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"],
)
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"],)
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"],)
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"],
):
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,)
st.toast(ret.get("msg", " "))
time.sleep(1)
st.experimental_rerun()