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): """