Merge branch 'dev_fastchat' of github.com:chatchat-space/langchain-ChatGLM into dev_fastchat

server/llm_api.py

@@ -13,7 +13,6 @@ openai_api_port = 8888
 base_url = "http://127.0.0.1:{}"
 queue = Queue()
 sys.modules['fastchat.constants.LOGDIR'] = LOG_PATH
-import parser
 
 
 def set_httpx_timeout(timeout=60.0):

webui.py

@@ -1,3 +1,9 @@
+# 运行方式：
+# 1. 安装必要的包：pip install streamlit-option-menu streamlit-chatbox>=1.1.3
+# 2. 运行本机fastchat服务：python server\llm_api.py 或者 运行对应的sh文件
+# 3. 运行API服务器：python server/api.py。如果使用api = ApiRequest(no_remote_api=True)，该步可以跳过。
+# 4. 运行WEB UI：streamlit run webui.py --server.port 7860
+
 import streamlit as st
 from webui_pages.utils import *
 from streamlit_option_menu import option_menu

webui_pages/dialogue/dialogue.py

@@ -2,37 +2,85 @@ import streamlit as st
 from webui_pages.utils import *
 from streamlit_chatbox import *
 
-chat_box = ChatBox()
-
 
 def dialogue_page(api: ApiRequest):
+    chat_box = ChatBox(
+        greetings=[
+            f"欢迎使用 [`Langchain-Chatglm`](https://github.com/chatchat-space/langchain-chatglm) ! 当前使用模型`{LLM_MODEL}`, 您可以开始提问了.",
+        ]
+    )
+
     with st.sidebar:
+        def on_mode_change():
+            mode = st.session_state.dialogue_mode
+            text = f"已切换到 {mode} 模式。"
+            if mode == "知识库问答":
+                cur_kb = st.session_state.get("selected_kb")
+                if cur_kb:
+                    text = f"{text} 当前知识库： `{cur_kb}`。"
+            chat_box.ai_say(text, not_render=True)
+
         dialogue_mode = st.radio("请选择对话模式",
                                  ["LLM 对话",
                                   "知识库问答",
-                                  "Bing 搜索问答"])
+                                  "Bing 搜索问答",
-        history_len = st.slider("历史对话轮数：", 1, 10, 1)
+                                  "Duck 搜索问答",
+                                ],
+                                on_change=on_mode_change,
+                                key="dialogue_mode",
+                                )
+        history_len = st.slider("历史对话轮数：", 1, 10, 1, disabled=True)
+        # todo: support history len
+        if st.button("清除历史对话"):
+            chat_box.reset_history()
+
+        def on_kb_change():
+            chat_box.ai_say(f"已加载知识库： {st.session_state.selected_kb}", not_render=True)
+
         if dialogue_mode == "知识库问答":
-            selected_kb = st.selectbox("请选择知识库：", get_kb_list())
+            kb_list = api.list_knowledge_bases()
-            with st.expander(f"{selected_kb} 中已存储文件"):
+            selected_kb = st.selectbox(
-                st.write(get_kb_files(selected_kb))
+                "请选择知识库：",
+                kb_list,
+                on_change=on_kb_change,
+                key="selected_kb",
+            )
+            top_k = st.slider("匹配知识条数：", 1, 20, 3, disabled=True)
+            score_threshold = st.slider("知识匹配分数阈值：", 0, 1000, 0, disabled=True)
+            chunk_content = st.checkbox("关联上下文", False, disabled=True)
+            chunk_size = st.slider("关联长度：", 0, 500, 250, disabled=True)
 
     # Display chat messages from history on app rerun
     chat_box.output_messages()
 
     if prompt := st.chat_input("请输入对话内容，换行请使用Ctrl+Enter"):
         chat_box.user_say(prompt)
-        chat_box.ai_say("正在思考...")
+        if dialogue_mode == "LLM 对话":
-        # with api.chat_fastchat([{"role": "user", "content": "prompt"}], stream=streaming) as r:
+            chat_box.ai_say("正在思考...")
-        # todo: support history len
+            text = ""
-        text = ""
+            r = api.chat_chat(prompt, no_remote_api=True)
-        r = api.chat_chat(prompt, no_remote_api=True)
+            for t in r:
-        for t in r:
+                text += t
-            text += t
+                chat_box.update_msg(text)
-            chat_box.update_msg(text)
+            chat_box.update_msg(text, streaming=False) # 更新最终的字符串，去除光标
-        chat_box.update_msg(text, streaming=False)
+        elif dialogue_mode == "知识库问答":
-        # with  api.chat_chat(prompt) as r:
+            chat_box.ai_say(f"正在查询知识库： `{selected_kb}` ...")
-        #     for t in r.iter_text(None):
+            text = ""
-        #         text += t
+            for t in api.knowledge_base_chat(prompt, selected_kb):
-        #         chat_box.update_msg(text)
+                text += t
-        #     chat_box.update_msg(text, streaming=False)
+                chat_box.update_msg(text)
+            chat_box.update_msg(text, streaming=False)
+        elif dialogue_mode == "Bing 搜索问答":
+            chat_box.ai_say("正在执行Bing搜索...")
+            text = ""
+            for t in api.bing_search_chat(prompt):
+                text += t
+                chat_box.update_msg(text)
+            chat_box.update_msg(text, streaming=False)
+        elif dialogue_mode == "Duck 搜索问答":
+            chat_box.ai_say("正在执行Duckduck搜索...")
+            text = ""
+            for t in api.duckduckgo_search_chat(prompt):
+                text += t
+                chat_box.update_msg(text)
+            chat_box.update_msg(text, streaming=False)

webui_pages/utils.py

@@ -1,5 +1,6 @@
 # 该文件包含webui通用工具，可以被不同的webui使用
 
+import tempfile
 from typing import *
 from pathlib import Path
 import os
@@ -12,6 +13,7 @@ import httpx
 import asyncio
 from server.chat.openai_chat import OpenAiChatMsgIn
 from fastapi.responses import StreamingResponse
+import contextlib
 
 
 def set_httpx_timeout(timeout=60.0):
@@ -87,9 +89,11 @@ class ApiRequest:
         self,
         base_url: str = "http://127.0.0.1:7861",
         timeout: float = 60.0,
+        no_remote_api: bool = False,   # call api view function directly
     ):
         self.base_url = base_url
         self.timeout = timeout
+        self.no_remote_api = no_remote_api
 
     def _parse_url(self, url: str) -> str:
         if (not url.startswith("http")
@@ -112,7 +116,7 @@ class ApiRequest:
         kwargs.setdefault("timeout", self.timeout)
         while retry > 0:
             try:
-                return httpx.get(url, params, **kwargs)
+                return httpx.get(url, params=params, **kwargs)
             except:
                 retry -= 1
 
@@ -128,7 +132,7 @@ class ApiRequest:
         async with httpx.AsyncClient() as client:
             while retry > 0:
                 try:
-                    return await client.get(url, params, **kwargs)
+                    return await client.get(url, params=params, **kwargs)
                 except:
                     retry -= 1
 
@@ -170,7 +174,7 @@ class ApiRequest:
                 except:
                     retry -= 1
 
-    def _stream2generator(self, response: StreamingResponse):
+    def _fastapi_stream2generator(self, response: StreamingResponse):
         '''
         将api.py中视图函数返回的StreamingResponse转化为同步生成器
         '''
@@ -180,6 +184,16 @@ class ApiRequest:
             loop = asyncio.new_event_loop()
         return iter_over_async(response.body_iterator, loop)
 
+    def _httpx_stream2generator(self,response: contextlib._GeneratorContextManager):
+        '''
+        将httpx.stream返回的GeneratorContextManager转化为普通生成器
+        '''
+        with response as r:
+            for chunk in r.iter_text(None):
+                yield chunk
+
+    # 对话相关操作
+
     def chat_fastchat(
         self,
         messages: List[Dict],
@@ -187,12 +201,14 @@ class ApiRequest:
         model: str = LLM_MODEL,
         temperature: float = 0.7,
         max_tokens: int = 1024, # todo:根据message内容自动计算max_tokens
-        no_remote_api=False,  # all api view function directly
+        no_remote_api: bool = None,
         **kwargs: Any,
     ):
         '''
         对应api.py/chat/fastchat接口
         '''
+        if no_remote_api is None:
+            no_remote_api = self.no_remote_api
         msg = OpenAiChatMsgIn(**{
             "messages": messages,
             "stream": stream,
@@ -205,7 +221,7 @@ class ApiRequest:
         if no_remote_api:
             from server.chat.openai_chat import openai_chat
             response = openai_chat(msg)
-            return self._stream2generator(response)
+            return self._fastapi_stream2generator(response)
         else:
             data = msg.dict(exclude_unset=True, exclude_none=True)
             response = self.post(
@@ -213,35 +229,260 @@ class ApiRequest:
                 json=data,
                 stream=stream,
             )
-            return response
+            return self._httpx_stream2generator(response)
 
     def chat_chat(
         self,
         query: str,
-        no_remote_api: bool = False,
+        no_remote_api: bool = None,
     ):
         '''
         对应api.py/chat/chat接口
         '''
+        if no_remote_api is None:
+            no_remote_api = self.no_remote_api
+
         if no_remote_api:
             from server.chat.chat import chat
             response = chat(query)
-            return self._stream2generator(response)
+            return self._fastapi_stream2generator(response)
         else:
             response = self.post("/chat/chat", json=f"{query}", stream=True)
-            return response
+            return self._httpx_stream2generator(response)
+
+    def knowledge_base_chat(
+        self,
+        query: str,
+        knowledge_base_name: str,
+        no_remote_api: bool = None,
+    ):
+        '''
+        对应api.py/chat/knowledge_base_chat接口
+        '''
+        if no_remote_api is None:
+            no_remote_api = self.no_remote_api
+
+        if no_remote_api:
+            from server.chat.knowledge_base_chat import knowledge_base_chat
+            response = knowledge_base_chat(query, knowledge_base_name)
+            return self._fastapi_stream2generator(response)
+        else:
+            response = self.post(
+                "/chat/knowledge_base_chat",
+                json={"query": query, "knowledge_base_name": knowledge_base_name},
+                stream=True,
+            )
+            return self._httpx_stream2generator(response)
+
+    def duckduckgo_search_chat(
+        self,
+        query: str,
+        no_remote_api: bool = None,
+    ):
+        '''
+        对应api.py/chat/duckduckgo_search_chat接口
+        '''
+        if no_remote_api is None:
+            no_remote_api = self.no_remote_api
+
+        if no_remote_api:
+            from server.chat.duckduckgo_search_chat import duckduckgo_search_chat
+            response = duckduckgo_search_chat(query)
+            return self._fastapi_stream2generator(response)
+        else:
+            response = self.post(
+                "/chat/duckduckgo_search_chat",
+                json=f"{query}",
+                stream=True,
+            )
+            return self._httpx_stream2generator(response)
+
+    def bing_search_chat(
+        self,
+        query: str,
+        no_remote_api: bool = None,
+    ):
+        '''
+        对应api.py/chat/bing_search_chat接口
+        '''
+        if no_remote_api is None:
+            no_remote_api = self.no_remote_api
+
+        if no_remote_api:
+            from server.chat.bing_search_chat import bing_search_chat
+            response = bing_search_chat(query)
+            return self._fastapi_stream2generator(response)
+        else:
+            response = self.post(
+                "/chat/bing_search_chat",
+                json=f"{query}",
+                stream=True,
+            )
+            return self._httpx_stream2generator(response)
+
+    # 知识库相关操作
+
+    def list_knowledge_bases(
+        self,
+        no_remote_api: bool = None,
+    ):
+        '''
+        对应api.py/knowledge_base/list_knowledge_bases接口
+        '''
+        if no_remote_api is None:
+            no_remote_api = self.no_remote_api
+
+        if no_remote_api:
+            from server.knowledge_base.kb_api import list_kbs
+            response = run_async(list_kbs())
+            return response.data
+        else:
+            response = self.get("/knowledge_base/list_knowledge_bases")
+            return response.json().get("data")
+
+    def create_knowledge_base(
+        self,
+        knowledge_base_name: str,
+        no_remote_api: bool = None,
+    ):
+        '''
+        对应api.py/knowledge_base/create_knowledge_base接口
+        '''
+        if no_remote_api is None:
+            no_remote_api = self.no_remote_api
+
+        if no_remote_api:
+            from server.knowledge_base.kb_api import create_kb
+            response = run_async(create_kb(knowledge_base_name))
+            return response.dict()
+        else:
+            response = self.post(
+                "/knowledge_base/create_knowledge_base",
+                json={"knowledge_base_name": knowledge_base_name},
+            )
+            return response.json()
+
+    def delete_knowledge_base(
+        self,
+        knowledge_base_name: str,
+        no_remote_api: bool = None,
+    ):
+        '''
+        对应api.py/knowledge_base/delete_knowledge_base接口
+        '''
+        if no_remote_api is None:
+            no_remote_api = self.no_remote_api
+
+        if no_remote_api:
+            from server.knowledge_base.kb_api import delete_kb
+            response = run_async(delete_kb(knowledge_base_name))
+            return response.dict()
+        else:
+            response = self.delete(
+                "/knowledge_base/delete_knowledge_base",
+                json={"knowledge_base_name": knowledge_base_name},
+            )
+            return response.json()
+
+    def list_kb_docs(
+        self,
+        knowledge_base_name: str,
+        no_remote_api: bool = None,
+    ):
+        '''
+        对应api.py/knowledge_base/list_docs接口
+        '''
+        if no_remote_api is None:
+            no_remote_api = self.no_remote_api
+
+        if no_remote_api:
+            from server.knowledge_base.kb_doc_api import list_docs
+            response = run_async(list_docs(knowledge_base_name))
+            return response.data
+        else:
+            response = self.get(
+                "/knowledge_base/list_docs",
+                params={"knowledge_base_name": knowledge_base_name}
+            )
+            return response.json().get("data")
+
+    def upload_kb_doc(
+        self,
+        file: Union[str, Path],
+        knowledge_base_name: str,
+        no_remote_api: bool = None,
+    ):
+        '''
+        对应api.py/knowledge_base/upload_docs接口
+        '''
+        if no_remote_api is None:
+            no_remote_api = self.no_remote_api
+
+        file = Path(file).absolute()
+        filename = file.name
+
+        if no_remote_api:
+            from server.knowledge_base.kb_doc_api import upload_doc
+            from fastapi import UploadFile
+            from tempfile import SpooledTemporaryFile
+
+            temp_file = SpooledTemporaryFile(max_size=10 * 1024 * 1024)
+            with file.open("rb") as fp:
+                temp_file.write(fp.read())
+            response = run_async(upload_doc(
+                UploadFile(temp_file, filename=filename),
+                knowledge_base_name,
+            ))
+            return response.dict()
+        else:
+            response = self.post(
+                "/knowledge_base/upload_doc",
+                data={"knowledge_base_name": knowledge_base_name},
+                files={"file": (filename, file.open("rb"))},
+            )
+            return response.json()
+
+    def delete_kb_doc(
+        self,
+        knowledge_base_name: str,
+        doc_name: str,
+        no_remote_api: bool = None,
+    ):
+        '''
+        对应api.py/knowledge_base/delete_doc接口
+        '''
+        if no_remote_api is None:
+            no_remote_api = self.no_remote_api
+
+        if no_remote_api:
+            from server.knowledge_base.kb_doc_api import delete_doc
+            response = run_async(delete_doc(knowledge_base_name, doc_name))
+            return response.dict()
+        else:
+            response = self.delete(
+                "/knowledge_base/delete_doc",
+                json={"knowledge_base_name": knowledge_base_name, "doc_name": doc_name},
+            )
+            return response.json()
 
 
 if __name__ == "__main__":
     api = ApiRequest()
+
     # print(api.chat_fastchat(
     #     messages=[{"role": "user", "content": "hello"}]
     # ))
 
-    with api.chat_chat("你好") as r:
+    # with api.chat_chat("你好") as r:
-        for t in r.iter_text(None):
+    #     for t in r.iter_text(None):
-            print(t)
+    #         print(t)
 
-    r = api.chat_chat("你好", no_remote_api=True)
+    # r = api.chat_chat("你好", no_remote_api=True)
-    for t in r:
+    # for t in r:
-        print(t)
+    #     print(t)
+
+    # r = api.duckduckgo_search_chat("室温超导最新研究进展", no_remote_api=True)
+    # for t in r:
+    #     print(t)
+
+    print(api.list_knowledge_bases())