From 0da44d31a1ee9aaf05d92a6b93dadf1ac5d897b0 Mon Sep 17 00:00:00 2001
From: zy123 <646228430@qq.com>
Date: Thu, 19 Dec 2024 12:25:29 +0800
Subject: [PATCH] =?UTF-8?q?12.19=20=E4=BF=AE=E5=A4=8D=E8=B1=86=E5=8C=85?=
 =?UTF-8?q?=E6=A8=A1=E5=9E=8B=E4=BD=BF=E7=94=A8=E7=9A=84bug?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 flask_app/general/doubao.py | 28 ++++++++++++++++++----------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/flask_app/general/doubao.py b/flask_app/general/doubao.py
index 119c782..b0c6b22 100644
--- a/flask_app/general/doubao.py
+++ b/flask_app/general/doubao.py
@@ -193,23 +193,30 @@ def doubao_model(full_user_query):
     doubao_api_key = os.getenv("DOUBAO_API_KEY")
 
     # 定义主模型和备用模型
-    models = [
-        "ep-20241119121710-425g6",  # 豆包Pro 32k模型
-        "ep-20241119121743-xt6wg"  # 128k模型
-    ]
+    models = {
+        "pro_32k": "ep-20241119121710-425g6",  # 豆包Pro 32k模型
+        "pro_128k": "ep-20241119121743-xt6wg"  # 128k模型
+    }
+
+    # 判断用户查询字符串的长度
+    if len(full_user_query) > 32000:
+        selected_model = models["pro_128k"]  # 如果长度超过32k，直接使用128k模型
+    else:
+        selected_model = models["pro_32k"]  # 默认使用32k模型
+
     # 请求头
     headers = {
         "Content-Type": "application/json",
         "Authorization": "Bearer " + doubao_api_key
     }
+
     max_retries = 1  # 最大重试次数
     attempt = 0
-    while attempt <= max_retries and attempt < len(models):
-        # 设置当前使用的模型
-        model_name = models[attempt]
+
+    while attempt <= max_retries:
         # 请求数据
         data = {
-            "model": model_name,
+            "model": selected_model,
             "messages": [
                 {
                     "role": "user",
@@ -225,11 +232,12 @@ def doubao_model(full_user_query):
             return response.json()["choices"][0]["message"]["content"]
         except requests.exceptions.RequestException as e:
             attempt += 1
-            if attempt > max_retries or attempt >= len(models):
+            if attempt > max_retries:
                 print(f"请求失败，尝试了 {attempt} 次。错误信息：{e}")
                 return None  # 或者根据需要返回其他默认值
             else:
-                print(f"请求出错：{e}。正在尝试第 {attempt} 次重试，使用备用模型...")
+                print(f"请求出错：{e}。正在尝试第 {attempt} 次重试，继续使用模型 {selected_model}...")
+
 
 def generate_full_user_query(file_path, prompt_template):
     """