zy123/zbparse
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

12.19 修复豆包模型使用的bug

Browse Source
This commit is contained in:
zy123 2024-12-19 12:25:29 +08:00
parent f0ed0efc6d
commit 0da44d31a1
1 changed files with 18 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
flask_app/general/doubao.py
View File

@ -193,23 +193,30 @@ def doubao_model(full_user_query):
doubao_api_key = os.getenv("DOUBAO_API_KEY") doubao_api_key = os.getenv("DOUBAO_API_KEY")
# 定义主模型和备用模型 # 定义主模型和备用模型
models = [ models = {
"ep-20241119121710-425g6", # 豆包Pro 32k模型 "pro_32k": "ep-20241119121710-425g6", # 豆包Pro 32k模型
"ep-20241119121743-xt6wg" # 128k模型 "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 = { headers = {
"Content-Type": "application/json", "Content-Type": "application/json",
"Authorization": "Bearer " + doubao_api_key "Authorization": "Bearer " + doubao_api_key
} }
max_retries = 1 # 最大重试次数 max_retries = 1 # 最大重试次数
attempt = 0 attempt = 0
while attempt <= max_retries and attempt < len(models):
# 设置当前使用的模型 while attempt <= max_retries:
model_name = models[attempt]
# 请求数据 # 请求数据
data = { data = {
"model": model_name, "model": selected_model,
"messages": [ "messages": [
{ {
"role": "user", "role": "user",
@ -225,11 +232,12 @@ def doubao_model(full_user_query):
return response.json()["choices"][0]["message"]["content"] return response.json()["choices"][0]["message"]["content"]
except requests.exceptions.RequestException as e: except requests.exceptions.RequestException as e:
attempt += 1 attempt += 1
if attempt > max_retries or attempt >= len(models): if attempt > max_retries:
print(f"请求失败，尝试了 {attempt} 次。错误信息：{e}") print(f"请求失败，尝试了 {attempt} 次。错误信息：{e}")
return None # 或者根据需要返回其他默认值 return None # 或者根据需要返回其他默认值
else: else:
print(f"请求出错：{e}。正在尝试第 {attempt} 次重试，使用备用模型...") print(f"请求出错：{e}。正在尝试第 {attempt} 次重试，继续使用模型 {selected_model}...")
def generate_full_user_query(file_path, prompt_template): def generate_full_user_query(file_path, prompt_template):
""" """