Add knowledge base save and load feature

ChuanhuChatbot.py

@@ -32,7 +32,6 @@
 def create_new_model():
  return get_model(model_name=MODELS[DEFAULT_MODEL], access_key=my_api_key)[0]
 
-
 with gr.Blocks(theme=small_and_beautiful_theme) as demo:
  user_name = gr.Textbox("", visible=False)
  # 激活/logout路由
@@ -261,6 +260,14 @@ def create_new_model():
  "使用在线搜索"), value=False, elem_classes="switch-checkbox", elem_id="gr-websearch-cb", visible=False)
  index_files = gr.Files(label=i18n(
  "上传"), type="file", file_types=[".pdf", ".docx", ".pptx", ".epub", ".xlsx", ".txt", "text", "image"], elem_id="upload-index-file")
+ with gr.Row():
+ knowledge_base_name = gr.Dropdown(
+ label=i18n("选择或输入知识库名字"), choices=load_knowledge_base_names(), multiselect=True, visible=True, elem_id="gr-kb-select-dropdown")
+ with gr.Row():
+ knowledge_base_refresh_btn = gr.Button(
+ i18n("🔄"))
+ save_index_file_btn = gr.Button(
+ i18n("保存知识库"), elem_id="save-index-file-btn")
  two_column = gr.Checkbox(label=i18n(
  "双栏pdf"), value=advance_docs["pdf"].get("two_column", False))
  summarize_btn = gr.Button(i18n("总结"))
@@ -594,6 +601,9 @@ def create_greeting(request: gr.Request):
 
  index_files.upload(handle_file_upload, [current_model, index_files, chatbot, language_select_dropdown], [
  index_files, chatbot, status_display])
+ knowledge_base_name.select(load_index_file, [current_model, knowledge_base_name], [index_files])
+ knowledge_base_refresh_btn.click(lambda: gr.Dropdown.update(choices=load_knowledge_base_names()), [], [knowledge_base_name])
+ save_index_file_btn.click(save_index_file, [current_model, index_files, knowledge_base_name], [])
  summarize_btn.click(handle_summarize_index, [
  current_model, index_files, chatbot, language_select_dropdown], [chatbot, status_display])
 

locale/en_US.json

@@ -116,6 +116,8 @@
  "刷新状态": "Refresh Status",
  "剩余配额不足，[进一步了解](https://github.com/GaiZhenbiao/ChuanhuChatGPT/wiki/%E5%B8%B8%E8%A7%81%E9%97%AE%E9%A2%98#you-exceeded-your-current-quota-please-check-your-plan-and-billing-details)": "Insufficient remaining quota, [learn more](https://github.com/GaiZhenbiao/ChuanhuChatGPT/wiki/%E5%B8%B8%E8%A7%81%E9%97%AE%E9%A2%98#you-exceeded-your-current-quota-please-check-your-plan-and-billing-details)",
  "加载Prompt模板": "Load Prompt Template",
+ "选择或输入知识库名字": "Select or enter knowledge base",
+ "保存知识库": "Save knowledge base",
  "单轮对话": "Single-turn",
  "历史记录（JSON）": "History file (JSON)",
  "参数": "Parameters",

locale/ja_JP.json

@@ -128,6 +128,8 @@
  "由Bilibili [土川虎虎虎](https://space.bilibili.com/29125536)、[明昭MZhao](https://space.bilibili.com/24807452) 和 [Keldos](https://github.com/Keldos-Li) 开发<br />访问川虎Chat的 [GitHub项目](https://github.com/GaiZhenbiao/ChuanhuChatGPT) 下载最新版脚本": "開発：Bilibili [土川虎虎虎](https://space.bilibili.com/29125536) と [明昭MZhao](https://space.bilibili.com/24807452) と [Keldos](https://github.com/Keldos-Li)\n\n最新コードは川虎Chatのサイトへ [GitHubプロジェクト](https://github.com/GaiZhenbiao/ChuanhuChatGPT)",
  "由于下面的原因，Google 拒绝返回 Gemini 的回答：\\n\\n": "GoogleがGeminiの回答を返信しない理由：",
  "知识库": "ファイル収納庫",
+ "选择或输入知识库名字": "ナレッジベース名を選択または入力",
+ "保存知识库": "ナレッジベースを保存",
  "知识库文件": "ナレッジベースファイル",
  "立即重启": "今すぐ再起動",
  "第一条提问": "最初の質問",

locale/ko_KR.json

@@ -129,6 +129,8 @@
  "由于下面的原因，Google 拒绝返回 Gemini 的回答：\\n\\n": "Google이 Gemini의 답변을 반환하는 것을 거부하는 이유에는 다음과 같은 이유가 있습니다:",
  "知识库": "knowledge base",
  "知识库文件": "knowledge base 파일",
+ "选择或输入知识库名字": "지식 라이브러리 이름 선택 또는 입력",
+ "保存知识库": "지식 라이브러리 저장",
  "立即重启": "지금 재시작",
  "第一条提问": "첫 번째 질문",
  "索引已保存至本地!": "로컬에 인덱스가 저장되었습니다!",

locale/ru_RU.json

@@ -149,6 +149,8 @@
  "由Bilibili [土川虎虎虎](https://space.bilibili.com/29125536)、[明昭MZhao](https://space.bilibili.com/24807452) 和 [Keldos](https://github.com/Keldos-Li) 开发<br />访问川虎Chat的 [GitHub项目](https://github.com/GaiZhenbiao/ChuanhuChatGPT) 下载最新版脚本": "Разработано [土川虎虎虎](https://space.bilibili.com/29125536), [明昭MZhao](https://space.bilibili.com/24807452) и [Keldos](https://github.com/Keldos-Li).<br />посетите [GitHub Project](https://github.com/GaiZhenbiao/ChuanhuChatGPT) чата Chuanhu, чтобы загрузить последнюю версию скрипта",
  "知识库": "База знаний",
  "知识库文件": "Файл базы знаний",
+ "选择或输入知识库名字": "Выберите или введите имя базы знаний",
+ "保存知识库": "Сохранить базу знаний",
  "立即重启": "Перезапустить сейчас",
  "第一条提问": "Первый вопрос",
  "索引构建完成": "Индексирование завершено.",

locale/sv_SE.json

@@ -129,6 +129,8 @@
  "由于下面的原因，Google 拒绝返回 Gemini 的回答：\\n\\n": "Google vägrar att returnera Gemini's svar av följande skäl:",
  "知识库": "kunskapsbank",
  "知识库文件": "kunskapsbankfil",
+ "选择或输入知识库名字": "Välj eller skriv in kunskapsbankens namn",
+ "保存知识库": "Spara kunskapsbank",
  "立即重启": "Starta om nu",
  "第一条提问": "Första frågan",
  "索引已保存至本地!": "Index har sparats lokalt!",

locale/vi_VN.json

@@ -149,6 +149,8 @@
  "由Bilibili [土川虎虎虎](https://space.bilibili.com/29125536)、[明昭MZhao](https://space.bilibili.com/24807452) 和 [Keldos](https://github.com/Keldos-Li) 开发<br />访问川虎Chat的 [GitHub项目](https://github.com/GaiZhenbiao/ChuanhuChatGPT) 下载最新版脚本": "Phát triển bởi Bilibili [土川虎虎虎](https://space.bilibili.com/29125536), [明昭MZhao](https://space.bilibili.com/24807452) và [Keldos](https://github.com/Keldos-Li)\n\nTải mã nguồn mới nhất từ [GitHub](https://github.com/GaiZhenbiao/ChuanhuChatGPT)",
  "知识库": "Cơ sở kiến thức",
  "知识库文件": "Tệp cơ sở kiến thức",
+ "选择或输入知识库名字": "Chọn hoặc nhập tên cơ sở kiến thức",
+ "保存知识库": "Lưu cơ sở kiến thức",
  "立即重启": "Khởi động lại ngay",
  "第一条提问": "Theo câu hỏi đầu tiên",
  "索引构建完成": "Xây dựng chỉ mục hoàn tất",

modules/models/base_model.py

@@ -399,6 +399,21 @@ def handle_file_upload(self, files, chatbot, language):
  status = i18n("索引构建失败！") + str(e)
  return gr.Files.update(), chatbot, status
 
+ def load_index_file(self, knowledge_base_name):
+ return [os.path.join("./files/knowledge_base", name) for name in knowledge_base_name]
+
+ def save_index_file(self, files, knowledge_base_name):
+ if files:
+ knowledge_base_path = f"./files/knowledge_base/{knowledge_base_name}"
+ Path(knowledge_base_path).mkdir(parents=True, exist_ok=True)
+ for file in files:
+ try:
+ shutil.copy(file.name, str(Path(knowledge_base_path) / Path(file.name).name))
+ except FileExistsError:
+ pass
+ gr.Info(i18n("索引文件保存成功！"))
+ return
+
  def summarize_index(self, files, chatbot, language):
  status = gr.Markdown.update()
  if files:

modules/utils.py

@@ -147,6 +147,14 @@ def set_single_turn(current_model, *args):
 def handle_file_upload(current_model, *args):
  return current_model.handle_file_upload(*args)
 
+def save_index_file(current_model, *args):
+ return current_model.save_index_file(*args)
+
+def load_index_file(current_model, *args):
+ return current_model.load_index_file(*args)
+
+def load_knowledge_base_names():
+ return [x.name for x in Path("./files/knowledge_base/").glob("*")]
 
 def handle_summarize_index(current_model, *args):
  return current_model.summarize_index(*args)
@@ -1411,4 +1419,4 @@ def setup_wizard():
 def reboot_chuanhu():
  import sys
  print(colorama.Back.GREEN + i18n("正在尝试重启...") + colorama.Style.RESET_ALL)
- os.execl(sys.executable, sys.executable, *sys.argv)
+ os.execl(sys.executable, sys.executable, *sys.argv)