update blackbox chroma and chat

This commit is contained in:
ACBBZ
2024-05-28 08:23:10 +00:00
parent 347eea2f9c
commit ca34962376
4 changed files with 113 additions and 48 deletions

View File

@ -82,14 +82,52 @@ class Chat(Blackbox):
if user_model_key is None or user_model_key.isspace() or user_model_key == "":
user_model_key = "YOUR_API_KEY"
# 文心格式和openai的不一样需要单独处理
if re.search(r"ernie", user_model_name):
# key = "24.22873ef3acf61fb343812681e4df251a.2592000.1719453781.282335-46723715" 没充钱只有ernie-speed-128k能用
key = user_model_key
if re.search(r"ernie-speed-128k", user_model_name):
url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-speed-128k?access_token=" + key
elif re.search(r"ernie-3.5-8k", user_model_name):
url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token=" + key
elif re.search(r"ernie-4.0-8k", user_model_name):
url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=" + key
payload = json.dumps({
"system": prompt_template,
"messages": user_context + [
{
"role": "user",
"content": user_question
}
],
"temperature": user_temperature,
"top_p": user_top_p,
"stop": [str(user_stop)],
"max_output_tokens": user_max_tokens
})
headers = {
'Content-Type': 'application/json'
}
response = requests.request("POST", url, headers=headers, data=payload)
return response.json()["result"]
# gpt-4 gpt-3.5-turbo
if re.search(r"gpt", user_model_name):
elif re.search(r"gpt", user_model_name):
url = 'https://api.openai.com/v1/completions'
key = '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'
# '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'
key = user_model_key
# 自定义model
else:
url = user_model_url
key = user_model_key
prompt_template = [
{"role": "system", "content": user_template},
]