加入在线搜索功能

ChuanhuChatbot.py

@@ -70,8 +70,8 @@ with gr.Blocks(css=customCSS,) as demo:
                                     interactive=True, label="Top-p (nucleus sampling)",)
                     temperature = gr.Slider(minimum=-0, maximum=5.0, value=1.0,
                                             step=0.1, interactive=True, label="Temperature",)
-                
-                
+
+
                 with gr.Accordion(label="加载Prompt模板", open=False):
                     with gr.Column():
                         with gr.Row():
@@ -101,9 +101,10 @@ with gr.Blocks(css=customCSS,) as demo:
                                 historyReadBtn = gr.Button("📂 读入对话")
 
                 use_streaming_checkbox = gr.Checkbox(label="实时传输回答", value=True, visible=enable_streaming_option)
-                
+                use_websearch_checkbox = gr.Checkbox(label="使用在线搜索", value=False)
+
+
 
-                
         with gr.Column(scale=5):
             with gr.Row(scale=1):
                 chatbot = gr.Chatbot().style(height=700)  # .style(color_map=("#1D51EE", "#585A5B"))
@@ -118,7 +119,7 @@ with gr.Blocks(css=customCSS,) as demo:
                 retryBtn = gr.Button("🔄 重新生成")
                 delLastBtn = gr.Button("🗑️ 删除最近一条对话")
                 reduceTokenBtn = gr.Button("♻️ 总结对话")
-    
+
 
     gr.HTML("""
     <div style="text-align: center; margin-top: 20px; margin-bottom: 20px;">
@@ -126,10 +127,10 @@ with gr.Blocks(css=customCSS,) as demo:
     gr.Markdown(description)
 
 
-    user_input.submit(predict, [keyTxt, systemPromptTxt, history, user_input, chatbot, token_count, top_p, temperature, use_streaming_checkbox, model_select_dropdown], [chatbot, history, status_display, token_count], show_progress=True)
+    user_input.submit(predict, [keyTxt, systemPromptTxt, history, user_input, chatbot, token_count, top_p, temperature, use_streaming_checkbox, model_select_dropdown, use_websearch_checkbox], [chatbot, history, status_display, token_count], show_progress=True)
     user_input.submit(reset_textbox, [], [user_input])
 
-    submitBtn.click(predict, [keyTxt, systemPromptTxt, history, user_input, chatbot, token_count, top_p, temperature, use_streaming_checkbox, model_select_dropdown], [chatbot, history, status_display, token_count], show_progress=True)
+    submitBtn.click(predict, [keyTxt, systemPromptTxt, history, user_input, chatbot, token_count, top_p, temperature, use_streaming_checkbox, model_select_dropdown, use_websearch_checkbox], [chatbot, history, status_display, token_count], show_progress=True)
     submitBtn.click(reset_textbox, [], [user_input])
 
     emptyBtn.click(reset_state, outputs=[chatbot, history, token_count, status_display], show_progress=True)

presets.py

@@ -40,6 +40,14 @@ pre code {
 
 summarize_prompt = "你是谁？我们刚才聊了什么？" # 总结对话时的 prompt
 MODELS = ["gpt-3.5-turbo", "gpt-3.5-turbo-0301", "gpt-4","gpt-4-0314", "gpt-4-32k", "gpt-4-32k-0314"] # 可选的模型
+websearch_prompt = """Web search results:
+
+{web_results}
+Current date: {current_date}
+
+Instructions: Using the provided web search results, write a comprehensive reply to the given query. Make sure to cite results using [[number](URL)] notation after the reference. If the provided search results refer to multiple subjects with the same name, write separate answers for each subject.
+Query: {query}
+Reply in 中文"""
 
 # 错误信息
 standard_error_msg = "☹️发生了错误：" # 错误信息的标准前缀

requirements.txt

@@ -5,3 +5,4 @@ tiktoken
 socksio
 tqdm
 colorama
+duckduckgo_search

utils.py

@@ -16,6 +16,8 @@ from presets import *
 import tiktoken
 from tqdm import tqdm
 import colorama
+from duckduckgo_search import ddg
+import datetime
 
 # logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s")
 
@@ -224,8 +226,17 @@ def predict_all(openai_api_key, system_prompt, history, inputs, chatbot, all_tok
     return chatbot, history, status_text, all_token_counts
 
 
-def predict(openai_api_key, system_prompt, history, inputs, chatbot, all_token_counts, top_p, temperature, stream=False, selected_model = MODELS[0], should_check_token_count = True):  # repetition_penalty, top_k
+def predict(openai_api_key, system_prompt, history, inputs, chatbot, all_token_counts, top_p, temperature, stream=False, selected_model = MODELS[0], use_websearch_checkbox = False, should_check_token_count = True):  # repetition_penalty, top_k
     logging.info("输入为：" +colorama.Fore.BLUE + f"{inputs}" + colorama.Style.RESET_ALL)
+    if use_websearch_checkbox:
+        results = ddg(inputs, max_results=3)
+        web_results = []
+        for idx, result in enumerate(results):
+            logging.info(f"搜索结果{idx + 1}：{result}")
+            web_results.append(f'[{idx+1}]"{result["body"]}"\nURL: {result["href"]}')
+        web_results = "\n\n".join(web_results)
+        today = datetime.datetime.today().strftime("%Y-%m-%d")
+        inputs = websearch_prompt.replace("{current_date}", today).replace("{query}", inputs).replace("{web_results}", web_results)
     if len(openai_api_key) != 51:
         status_text = standard_error_msg + no_apikey_msg
         logging.info(status_text)