WordsStory / LLMutils.py
redstoneleo's picture
Upload 12 files
54bc341 verified
import json,re,os
from openai import OpenAI
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"), # Access Token属于个人账户的重要隐私信息,请谨慎管理,切忌随意对外公开,
base_url="https://aistudio.baidu.com/llm/lmapi/v3") # aistudio 大模型 api 服务域名
MODEL_NAME='ernie-4.5-turbo-128k'
def chat(messages,MODEL_NAME):
response = client.chat.completions.create(
model=MODEL_NAME,
messages=messages,
stream=False,
temperature=0,
response_format={"type": "json_object"},
)
response_text = response.choices[0].message.content
try:
json_output = json.loads(response_text)
except json.decoder.JSONDecodeError as e:
print(response_text)
match = re.search(r"\{.*?\}", response_text, re.DOTALL)
if match:
json_output = json.loads(match.group(0))
return json_output