ouyang/AstrBot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #58 from O2022/patch-2

Browse Source 
feat: 实现OpenAI画图
This commit is contained in:
Soulter
2023-03-29 17:21:49 +08:00
committed by GitHub
parent bfa8f137de 37e9373561
commit 134a8e233a
1 changed files with 59 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

83
cores/openai/core.py
View File

@@ -31,35 +31,70 @@ class ChatGPT:
self.chatGPT_configs = chatGPT_configs self.chatGPT_configs = chatGPT_configs
self.openai_configs = cfg self.openai_configs = cfg
def chat(self, req, image_mode = False): def chat(self, req, image_mode = False, img_num = 1, img_size = "1024x1024"):
# ChatGPT API 2023/3/2 # ChatGPT API 2023/3/2
# messages = [{"role": "user", "content": prompt}] # messages = [{"role": "user", "content": prompt}]
try: if not image_mode:
response = openai.ChatCompletion.create( try:
messages=req,
**self.chatGPT_configs
)
except Exception as e:
print(e)
if 'You exceeded' in str(e) or 'Billing hard limit has been reached' in str(e) or 'No API key provided' in str(e) or 'Incorrect API key provided' in str(e):
print("[System] 当前Key已超额或者不正常,正在切换")
self.key_stat[openai.api_key]['exceed'] = True
self.save_key_record()
response, is_switched = self.handle_switch_key(req)
if not is_switched:
# 所有Key都超额或不正常
raise e
else:
response = openai.ChatCompletion.create( response = openai.ChatCompletion.create(
messages=req, messages=req,
**self.chatGPT_configs **self.chatGPT_configs
) )
self.key_stat[openai.api_key]['used'] += response['usage']['total_tokens'] except Exception as e:
self.save_key_record() print(e)
print("[ChatGPT] "+str(response["choices"][0]["message"]["content"])) if 'You exceeded' in str(e) or 'Billing hard limit has been reached' in str(e) or 'No API key provided' in str(e) or 'Incorrect API key provided' in str(e):
return str(response["choices"][0]["message"]["content"]).strip(), response['usage']['total_tokens'] print("[System] 当前Key已超额或者不正常,正在切换")
self.key_stat[openai.api_key]['exceed'] = True
self.save_key_record()
response, is_switched = self.handle_switch_key(req)
if not is_switched:
# 所有Key都超额或不正常
raise e
else:
response = openai.ChatCompletion.create(
messages=req,
**self.chatGPT_configs
)
self.key_stat[openai.api_key]['used'] += response['usage']['total_tokens']
self.save_key_record()
print("[ChatGPT] "+str(response["choices"][0]["message"]["content"]))
return str(response["choices"][0]["message"]["content"]).strip(), response['usage']['total_tokens']
else:
try:
# print("test1")
response = openai.Image.create(
prompt=req[0]['content'],
n=img_num,
size=img_size
)
# print("test2")
image_url = []
for i in range(img_num):
image_url.append(response['data'][i]['url'])
print(image_url)
except Exception as e:
print(e)
if 'You exceeded' in str(e) or 'Billing hard limit has been reached' in str(
e) or 'No API key provided' in str(e) or 'Incorrect API key provided' in str(e):
print("[System] 当前Key已超额或者不正常,正在切换")
self.key_stat[openai.api_key]['exceed'] = True
self.save_key_record()
response, is_switched = self.handle_switch_key(req)
if not is_switched:
# 所有Key都超额或不正常
raise e
else:
response = openai.Image.create(
prompt=req[0]['content'],
n=img_num,
size=img_size
)
image_url = []
for i in range(img_num):
image_url.append(response['data'][i]['url'])
return image_url
def handle_switch_key(self, req): def handle_switch_key(self, req):
# messages = [{"role": "user", "content": prompt}] # messages = [{"role": "user", "content": prompt}]
while True: while True:
@@ -153,4 +188,4 @@ class ChatGPT:
pass pass
if openai.api_key == None: if openai.api_key == None:
self.handle_switch_key("") self.handle_switch_key("")
self.save_key_record() self.save_key_record()