57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
import openai
|
||
import os
|
||
from dotenv import load_dotenv, find_dotenv
|
||
|
||
|
||
# 如果你设置的是全局的环境变量,这行代码则没有任何作用。
|
||
_ = load_dotenv(find_dotenv())
|
||
|
||
# 获取环境变量 OPENAI_API_KEY
|
||
openai.api_key = os.environ['OPENAI_API_KEY']
|
||
|
||
# 一个封装 OpenAI 接口的函数,参数为 Prompt,返回对应结果
|
||
|
||
|
||
def get_completion(prompt,
|
||
model="gpt-3.5-turbo"
|
||
):
|
||
'''
|
||
prompt: 对应的提示词
|
||
model: 调用的模型,默认为 gpt-3.5-turbo(ChatGPT)。你也可以选择其他模型。
|
||
https://platform.openai.com/docs/models/overview
|
||
'''
|
||
|
||
messages = [{"role": "user", "content": prompt}]
|
||
|
||
# 调用 OpenAI 的 ChatCompletion 接口
|
||
response = openai.ChatCompletion.create(
|
||
model=model,
|
||
messages=messages,
|
||
temperature=0
|
||
)
|
||
|
||
return response.choices[0].message["content"]
|
||
|
||
|
||
def get_completion_from_messages(messages,
|
||
model="gpt-3.5-turbo",
|
||
temperature=0,
|
||
max_tokens=500):
|
||
'''
|
||
prompt: 对应的提示词
|
||
model: 调用的模型,默认为 gpt-3.5-turbo(ChatGPT)。你也可以选择其他模型。
|
||
https://platform.openai.com/docs/models/overview
|
||
temperature: 模型输出的随机程度。默认为0,表示输出将非常确定。增加温度会使输出更随机。
|
||
max_tokens: 定模型输出的最大的 token 数。
|
||
'''
|
||
|
||
# 调用 OpenAI 的 ChatCompletion 接口
|
||
response = openai.ChatCompletion.create(
|
||
model=model,
|
||
messages=messages,
|
||
temperature=temperature,
|
||
max_tokens=max_tokens
|
||
)
|
||
|
||
return response.choices[0].message["content"]
|