|
from ..presets import * |
|
from ..utils import * |
|
|
|
from .base_model import BaseLLMModel |
|
|
|
|
|
class ERNIE_Client(BaseLLMModel): |
|
def __init__(self, model_name, api_key, secret_key) -> None: |
|
super().__init__(model_name=model_name) |
|
self.api_key = api_key |
|
self.api_secret = secret_key |
|
if None in [self.api_secret, self.api_key]: |
|
raise Exception("请在配置文件或者环境变量中设置文心一言的API Key 和 Secret Key") |
|
|
|
if self.model_name == "ERNIE-Bot-turbo": |
|
self.ERNIE_url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?access_token=" |
|
elif self.model_name == "ERNIE-Bot": |
|
self.ERNIE_url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token=" |
|
elif self.model_name == "ERNIE-Bot-4": |
|
self.ERNIE_url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=" |
|
|
|
def get_access_token(self): |
|
""" |
|
使用 AK，SK 生成鉴权签名（Access Token） |
|
:return: access_token，或是None(如果错误) |
|
""" |
|
url = "https://aip.baidubce.com/oauth/2.0/token?client_id=" + self.api_key + "&client_secret=" + self.api_secret + "&grant_type=client_credentials" |
|
|
|
payload = json.dumps("") |
|
headers = { |
|
'Content-Type': 'application/json', |
|
'Accept': 'application/json' |
|
} |
|
|
|
response = requests.request("POST", url, headers=headers, data=payload) |
|
|
|
return response.json()["access_token"] |
|
def get_answer_stream_iter(self): |
|
url = self.ERNIE_url + self.get_access_token() |
|
system_prompt = self.system_prompt |
|
history = self.history |
|
if system_prompt is not None: |
|
history = [construct_system(system_prompt), *history] |
|
|
|
|
|
history = [i for i in history if i["role"] != "system"] |
|
|
|
payload = json.dumps({ |
|
"messages":history, |
|
"stream": True |
|
}) |
|
headers = { |
|
'Content-Type': 'application/json' |
|
} |
|
|
|
response = requests.request("POST", url, headers=headers, data=payload, stream=True) |
|
|
|
if response.status_code == 200: |
|
partial_text = "" |
|
for line in response.iter_lines(): |
|
if len(line) == 0: |
|
continue |
|
line = json.loads(line[5:]) |
|
partial_text += line['result'] |
|
yield partial_text |
|
else: |
|
yield STANDARD_ERROR_MSG + GENERAL_ERROR_MSG |
|
|
|
|
|
def get_answer_at_once(self): |
|
url = self.ERNIE_url + self.get_access_token() |
|
system_prompt = self.system_prompt |
|
history = self.history |
|
if system_prompt is not None: |
|
history = [construct_system(system_prompt), *history] |
|
|
|
|
|
history = [i for i in history if i["role"] != "system"] |
|
|
|
payload = json.dumps({ |
|
"messages": history, |
|
"stream": True |
|
}) |
|
headers = { |
|
'Content-Type': 'application/json' |
|
} |
|
|
|
response = requests.request("POST", url, headers=headers, data=payload, stream=True) |
|
|
|
if response.status_code == 200: |
|
|
|
return str(response.json()["result"]),len(response.json()["result"]) |
|
else: |
|
return "获取资源错误", 0 |
|
|
|
|
|
|
