Update app.py

update the chat mode with memory.

app.py

@@ -18,6 +18,19 @@ from g4f.Provider import (
     GetGpt
 )
 import os
+import json
+import pandas as pd
+
+from models_for_langchain.model import CustomLLM
+from langchain.memory import ConversationBufferWindowMemory, ConversationTokenBufferMemory
+from langchain import LLMChain, PromptTemplate
+from langchain.prompts import (
+    ChatPromptTemplate,
+    PromptTemplate,
+    SystemMessagePromptTemplate,
+    AIMessagePromptTemplate,
+    HumanMessagePromptTemplate,
+)
 
 provider_dict = {
     'Ails': Ails,
@@ -37,92 +50,113 @@ provider_dict = {
     'GetGpt': GetGpt
 }
 
+prompt_set_list = {}
+for prompt_file in os.listdir("prompt_set"):
+    key = prompt_file
+    if '.csv' in key:
+        df = pd.read_csv("prompt_set/" + prompt_file)
+        prompt_dict = dict(zip(df['act'], df['prompt']))
+    else:
+        with open("prompt_set/" + prompt_file, encoding='utf-8') as f:
+            ds = json.load(f)
+        prompt_dict = {item["act"]: item["prompt"] for item in ds}
+    prompt_set_list[key] = prompt_dict
+
 with gr.Blocks() as demo:
-  chatbot = gr.Chatbot([[None, None]], label='AI')
-  msg = gr.Textbox(value="", label='')
-  with gr.Row():
-    clear = gr.Button("清空对话")
-    remember = gr.Button("记住最近一次对话")
-  system_msg = gr.Textbox(value="你是一名助手，可以解答问题。", label='系统提示')
-  with gr.Row():
-    model_name = gr.Dropdown(['gpt-3.5-turbo', 'gpt-4'], value='gpt-3.5-turbo', label='模型')
-    provider_name  =  gr.Dropdown(provider_dict.keys(), value='GetGpt', label='提供者')
+    llm = CustomLLM()
+
+    template = """
+    Chat with human based on following instructions:
+    ```
+    {system_instruction}
+    ```
+    The following is a conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.
+    {{chat_history}}
+    Human: {{human_input}}
+    Chatbot:"""
+
+    memory = ConversationBufferWindowMemory(k=10, memory_key="chat_history")
+
+    chatbot = gr.Chatbot([[None, None]], label='AI')
+    msg = gr.Textbox(value="", label='请输入：')
+    with gr.Row():
+        clear = gr.Button("清空对话", scale=2)
+        chat_mode = gr.Checkbox(value=True, label='聊天模式', interactive=True, scale=1)
+    system_msg = gr.Textbox(value="你是一名助手，可以解答问题。", label='系统提示')
+    with gr.Row():
+        default_prompt_set = "1 中文提示词.json"
+        prompt_set_name = gr.Dropdown(prompt_set_list.keys(), value=default_prompt_set, label='提示词集合')
+        prompt_name = gr.Dropdown(prompt_set_list[default_prompt_set].keys(), label='提示词')
+    with gr.Row():
+        model_name = gr.Dropdown(['gpt-3.5-turbo', 'gpt-4'], value='gpt-3.5-turbo', label='模型')
+        provider_name  =  gr.Dropdown(provider_dict.keys(), value='GetGpt', label='提供者')
+
+    def change_prompt_set(prompt_set_name):
+        return gr.Dropdown.update(choices=list(prompt_set_list[prompt_set_name].keys()))
+
+    def change_prompt(prompt_set_name, prompt_name):
+        return gr.update(value=prompt_set_list[prompt_set_name][prompt_name])
+
+    def user(user_message, history):
+        return gr.update(value="", interactive=False), history + [[user_message, None]]
 
-  memory_msg = gr.Textbox(value="", label='记忆')
+    def bot(history, model_name, provider_name, system_msg, chat_mode):
+        history[-1][1] = ''
+        if len(system_msg)>3000:
+            system_msg = system_msg[:2000] + system_msg[-1000:]
 
-  def user(user_message, history):
-      return gr.update(value="", interactive=False), history + [[user_message, None]]
+        if chat_mode:
+            global template, memory
+            llm.model_name = model_name
+            llm.provider_name = provider_name
+            prompt = PromptTemplate(
+                                input_variables=["chat_history", "human_input"], template=template.format(system_instruction=system_msg)
+                            )
+            llm_chain = LLMChain(
+                                llm=llm,
+                                prompt=prompt,
+                                verbose=False,
+                                memory=memory,
+                            )
+            bot_msg = llm_chain.run(history[-1][0])
+            for c in bot_msg:
+                history[-1][1] += c
+            yield history
+        else:
+            prompt = """
+            请你仔细阅读以下提示，然后针对用户的话进行回答。
+            提示：
+            ```
+            {}
+            ```
+            用户最新的话：
+            ```
+            {}
+            ```
+            请回答：
+            """
+            bot_msg = g4f.ChatCompletion.create(model=model_name, 
+                                            provider=provider_dict[provider_name], 
+                                            messages=[{"role": "user", 
+                                                        "content": prompt.format(system_msg,
+                                                                                history[-1][0])}],
+                                            stream=True)
+            for c in bot_msg:
+                history[-1][1] += c
+            yield history
 
-  def bot(history, model_name, provider_name, system_msg, memory_msg):
-      history[-1][1] = ''
-      if len(system_msg)>3000:
-         system_msg = system_msg[:1500] + system_msg[-1500:]
-      prompt = """
-      请你仔细阅读以下提示，然后针对用户的话进行回答。
-      提示：
-      ```
-      {}
-      ```
-      历史对话中的关键信息：
-      ```
-      {}
-      ```
-      用户最新的话：
-      ```
-      {}
-      ```
-      请回答：
-      """
-      bot_msg = g4f.ChatCompletion.create(model=model_name, 
-                                          provider=provider_dict[provider_name], 
-                                          messages=[{"role": "user", 
-                                                     "content": prompt.format(system_msg,
-                                                                              memory_msg,
-                                                                              history[-1][0])}],
-                                          stream=True)
-      for c in bot_msg:
-          history[-1][1] += c
-          yield history
+    def empty_chat():
+        global memory
+        memory = ConversationBufferWindowMemory(k=10, memory_key="chat_history")
+        return None
+    response = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot, [chatbot, model_name, provider_name, system_msg, chat_mode], chatbot
+    )
+    prompt_set_name.select(change_prompt_set, prompt_set_name, prompt_name)
+    prompt_name.select(change_prompt, [prompt_set_name, prompt_name], system_msg)
 
-  def update_memory(memory_msg, history, model_name, provider_name):
-     if history==None:
-        return gr.update(value=memory_msg, interactive=False)
-     prompt = """
-     请你记录对话内容，尤其是用户提供的各种关键信息。已知的对话内容总结如下：
-     ```
-     {}
-     ```
-     最新的对话内容如下：
-     ```
-     用户:{}
-     AI:{}
-     ```
-     总结格式如下：
-     ```
-     关键信息总结：XXX
-     ```
-     请对已知对话和最新对话的关键信息进行总结，以应对后续的聊天：
-     """
-     memory = g4f.ChatCompletion.create(model=model_name,
-                                            provider=provider_dict[provider_name],
-                                            messages=[
-                                               {
-                                                  'role':'user',
-                                                  'content': prompt.format(memory_msg,
-                                                                           history[-1][0],
-                                                                           history[-1][1])
-                                               }
-                                            ])
-     return gr.update(value=memory, interactive=False)
-  
-  response = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-      bot, [chatbot, model_name, provider_name, system_msg, memory_msg], chatbot
-  )
-  response.then(lambda: gr.update(interactive=True), None, [msg], queue=False)
-  remember.click(
-     update_memory, [memory_msg, chatbot, model_name, provider_name], memory_msg
-  )
-  clear.click(lambda: (None, None), None, [chatbot, memory_msg], queue=False)
+    response.then(lambda: gr.update(interactive=True), None, [msg], queue=False)
+    clear.click(empty_chat, None, [chatbot], queue=False)
 
 demo.title = "AI Chat"
 demo.queue()