Update bin_public/app/chat_func.py

bin_public/app/chat_func.py

@@ -1,13 +1,12 @@
 # -*- coding:utf-8 -*-
 from __future__ import annotations
 
-import requests
 import urllib3
 
 from tqdm import tqdm
 from duckduckgo_search import ddg
 from llama_func import *
-from bin_public.utils.tools import *
+from bin_public.utils.utils import *
 from bin_public.utils.utils_db import *
 # logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s")
 
@@ -94,7 +93,8 @@ def stream_predict(
     top_p,
     temperature,
     selected_model,
-    fake_input=None
+    fake_input=None,
+    display_append=""
 ):
     def get_return_value():
         return chatbot, history, status_text, all_token_counts
@@ -106,9 +106,9 @@ def stream_predict(
     history.append(construct_user(inputs))
     history.append(construct_assistant(""))
     if fake_input:
-        chatbot.append((parse_text(fake_input), ""))
+        chatbot.append((fake_input, ""))
     else:
-        chatbot.append((parse_text(inputs), ""))
+        chatbot.append((inputs, ""))
     user_token_count = 0
     if len(all_token_counts) == 0:
         system_prompt_token_count = count_token(construct_system(system_prompt))
@@ -183,7 +183,7 @@ def stream_predict(
                     yield get_return_value()
                     break
                 history[-1] = construct_assistant(partial_words)
-                chatbot[-1] = (chatbot[-1][0], parse_text(partial_words))
+                chatbot[-1] = (chatbot[-1][0], partial_words+display_append)
                 all_token_counts[-1] += 1
                 yield get_return_value()
 
@@ -198,15 +198,16 @@ def predict_all(
     top_p,
     temperature,
     selected_model,
-    fake_input=None
+    fake_input=None,
+    display_append=""
 ):
     logging.info("一次性回答模式")
     history.append(construct_user(inputs))
     history.append(construct_assistant(""))
     if fake_input:
-        chatbot.append((parse_text(fake_input), ""))
+        chatbot.append((fake_input, ""))
     else:
-        chatbot.append((parse_text(inputs), ""))
+        chatbot.append((inputs, ""))
     all_token_counts.append(count_token(construct_user(inputs)))
     try:
         response = get_response(
@@ -232,7 +233,7 @@ def predict_all(
     response = json.loads(response.text)
     content = response["choices"][0]["message"]["content"]
     history[-1] = construct_assistant(content)
-    chatbot[-1] = (chatbot[-1][0], parse_text(content))
+    chatbot[-1] = (chatbot[-1][0], content+display_append)
     total_token_count = response["usage"]["total_tokens"]
     all_token_counts[-1] = total_token_count - sum(all_token_counts)
     status_text = construct_token_message(total_token_count)
@@ -287,7 +288,7 @@ def predict(
     if len(openai_api_key) != 51:
         status_text = standard_error_msg + no_apikey_msg
         logging.info(status_text)
-        chatbot.append((parse_text(inputs), ""))
+        chatbot.append(inputs, "")
         if len(history) == 0:
             history.append(construct_user(inputs))
             history.append("")
@@ -341,12 +342,22 @@ def predict(
         holo_query_insert_chat_message(invite_code, inputs, history[-1]['content'], token, history)
 
     if use_websearch:
-        response = history[-1]['content']
-        response += "\n\n" + "\n".join(link_references)
-        logging.info(f"Added link references.")
-        logging.info(response)
-        chatbot[-1] = (parse_text(old_inputs), response)
-        yield chatbot, history, status_text, all_token_counts
+        search_results = ddg(inputs, max_results=5)
+        old_inputs = inputs
+        web_results = []
+        for idx, result in enumerate(search_results):
+            logging.info(f"搜索结果{idx + 1}：{result}")
+            domain_name = urllib3.util.parse_url(result["href"]).host
+            web_results.append(f'[{idx + 1}]"{result["body"]}"\nURL: {result["href"]}')
+            link_references.append(f"{idx + 1}. [{domain_name}]({result['href']})\n")
+        link_references = "\n\n" + "".join(link_references)
+        inputs = (
+            replace_today(WEBSEARCH_PTOMPT_TEMPLATE)
+                .replace("{query}", inputs)
+                .replace("{web_results}", "\n\n".join(web_results))
+        )
+    else:
+        link_references = ""
 
     if stream:
         max_token = max_token_streaming