From 3e4c2b056c9c2ff78514e26ccb217fb3360c5a67 Mon Sep 17 00:00:00 2001
From: 505030475 <505030475@qq.com>
Date: Tue, 30 May 2023 19:55:38 +0800
Subject: [PATCH] knowledge base

---
 crazy_functional.py                | 12 ++++++++++++
 crazy_functions/Langchain知识库.py | 23 +++++++++++++++--------
 crazy_functions/crazy_utils.py     |  3 +++
 3 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/crazy_functional.py b/crazy_functional.py
index 3b295ac..9ce8ee1 100644
--- a/crazy_functional.py
+++ b/crazy_functional.py
@@ -301,9 +301,21 @@ def get_crazy_functions():
             "构建知识库（请先上传文件素材）": {
                 "Color": "stop",
                 "AsButton": False,
+                "AdvancedArgs": True,
+                "ArgsReminder": "待注入的知识库名称id, 默认为default",
                 "Function": HotReload(知识库问答)
             }
         })
+        from crazy_functions.Langchain知识库 import 读取知识库作答
+        function_plugins.update({
+            "知识库问答": {
+                "Color": "stop",
+                "AsButton": False,
+                "AdvancedArgs": True,
+                "ArgsReminder": "待提取的知识库名称id, 默认为default, 您需要首先调用构建知识库",
+                "Function": HotReload(读取知识库作答)
+            }
+        })
     except:
         print('Load function plugin failed')
 
diff --git a/crazy_functions/Langchain知识库.py b/crazy_functions/Langchain知识库.py
index 0bdb7f5..4fee382 100644
--- a/crazy_functions/Langchain知识库.py
+++ b/crazy_functions/Langchain知识库.py
@@ -29,9 +29,14 @@ def 知识库问答(txt, llm_kwargs, plugin_kwargs, chatbot, history, system_pro
         )
         yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
         return
+    
+    # < --------------------读取参数--------------- >
+    if ("advanced_arg" in plugin_kwargs) and (plugin_kwargs["advanced_arg"] == ""): plugin_kwargs.pop("advanced_arg")
+    kai_id = plugin_kwargs.get("advanced_arg", 'default')
+
     # < --------------------读取文件--------------- >
     file_manifest = []
-    spl = ["doc", "docx", "email", "epub", "html", "image", "json", "md", "msg", "odt", "pdf", "ppt", "pptx", "rtf", "text"]
+    spl = ["txt", "doc", "docx", "email", "epub", "html", "image", "json", "md", "msg", "odt", "pdf", "ppt", "pptx", "rtf", "text"]
     for sp in spl:
         _, file_manifest_tmp, _ = get_files_from_everything(txt, type=f'.{sp}')
         file_manifest += file_manifest_tmp
@@ -55,13 +60,15 @@ def 知识库问答(txt, llm_kwargs, plugin_kwargs, chatbot, history, system_pro
     print('Establishing knowledge archive ...')
     with ProxyNetworkActivate():    # 临时地激活代理网络
         kai = knowledge_archive_interface()
-        kai.feed_archive(file_manifest=file_manifest, id="default")
-
-    chatbot.append(['知识库构建成功', "正在将知识库存储至cookie中"])
-    yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
-    chatbot._cookies['langchain_plugin_embedding'] = kai.get_current_archive_id()
-    chatbot._cookies['lock_plugin'] = 'crazy_functions.Langchain知识库->读取知识库作答'
-    chatbot.append(['完成', "“根据知识库作答”函数插件已经接管问答系统, 提问吧! 但注意, 您接下来不能再使用其他插件了，刷新页面即可以退出知识库问答模式。"])
+        kai.feed_archive(file_manifest=file_manifest, id=kai_id)
+    kai_files = kai.get_loaded_file()
+    kai_files = '<br/>'.join(kai_files)
+    # chatbot.append(['知识库构建成功', "正在将知识库存储至cookie中"])
+    # yield from update_ui(chatbot=chatbot, history=history) # 刷新界面
+    # chatbot._cookies['langchain_plugin_embedding'] = kai.get_current_archive_id()
+    # chatbot._cookies['lock_plugin'] = 'crazy_functions.Langchain知识库->读取知识库作答'
+    # chatbot.append(['完成', "“根据知识库作答”函数插件已经接管问答系统, 提问吧! 但注意, 您接下来不能再使用其他插件了，刷新页面即可以退出知识库问答模式。"])
+    chatbot.append(['构建完成', f"当前知识库内的文件：\n\n---\n\n{kai_files}\n\n---\n\n请切换至“知识库问答”插件进行知识库访问, 或者使用此插件继续上传更多文件。"])
     yield from update_ui(chatbot=chatbot, history=history) # 刷新界面 # 由于请求gpt需要一段时间，我们先及时地做一次界面更新
 
 @CatchException
diff --git a/crazy_functions/crazy_utils.py b/crazy_functions/crazy_utils.py
index 3416d91..a447d22 100644
--- a/crazy_functions/crazy_utils.py
+++ b/crazy_functions/crazy_utils.py
@@ -661,6 +661,9 @@ class knowledge_archive_interface():
 
     def get_current_archive_id(self):
         return self.current_id
+    
+    def get_loaded_file(self):
+        return self.qa_handle.get_loaded_file()
 
     def answer_with_archive_by_id(self, txt, id):
         self.threadLock.acquire()