75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
# Please install OpenAI SDK first: `pip3 install openai`
|
||
import os
|
||
import json
|
||
from openai import OpenAI
|
||
from dotenv import load_dotenv
|
||
from httpx import Client, Proxy
|
||
|
||
load_dotenv()
|
||
|
||
|
||
class DOUBAO:
|
||
system_prompt = (
|
||
"You are a professional QA answer assistant.\n"
|
||
"The user will provide a JSON object containing a question.\n"
|
||
"Your job:\n"
|
||
"1. Read the 'title' and the 'chooies'.\n"
|
||
"2. Determine the correct answer(s).\n"
|
||
"3. Output ONLY the 'num' of the correct choices.\n"
|
||
"4. If multiple answers exist, join them with English commas, e.g., 'A,C,D'.\n"
|
||
"5. If you cannot determine the answer confidently, return an empty string.\n"
|
||
"6. Do not output explanations or any extra text."
|
||
)
|
||
|
||
def __init__(self) -> None:
|
||
self.key = os.environ.get("DOUBAO_API_KEY")
|
||
if not self.key:
|
||
raise Exception("找不到 DOUBAO_API_KEY 请设置环境变量在运行代码")
|
||
|
||
httpx_client = Client(
|
||
# proxy=Proxy("socks5://127.0.0.1:9000"), # 传入SOCKS代理地址
|
||
timeout=30.0, # 可选：设置超时时间
|
||
# verify=False,
|
||
)
|
||
|
||
self.client = OpenAI(
|
||
api_key=self.key,
|
||
base_url="https://ark.cn-beijing.volces.com/api/v3",
|
||
http_client=httpx_client,
|
||
)
|
||
|
||
def chat(self, question_json: dict) -> str:
|
||
response = self.client.chat.completions.create(
|
||
# model="deepseek-chat",
|
||
model="doubao-seed-code-preview-251028",
|
||
# model="deepseek-reasoner",
|
||
messages=[
|
||
{"role": "system", "content": self.system_prompt},
|
||
{"role": "user", "content": json.dumps(question_json)},
|
||
],
|
||
stream=False,
|
||
)
|
||
|
||
# print(response.choices[0].message.content)
|
||
return (
|
||
response.choices[0].message.content
|
||
if response.choices[0].message.content
|
||
else ""
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
question = {
|
||
"id": "25594265",
|
||
"index": "1",
|
||
"type": "单选",
|
||
"title": "下列哪个属于AI视觉识别应用？",
|
||
"chooies": [
|
||
{"num": "A", "txt": "音乐合成"},
|
||
{"num": "B", "txt": "植物识别"},
|
||
{"num": "C", "txt": "文案创作"},
|
||
{"num": "D", "txt": "机器翻译"},
|
||
],
|
||
}
|
||
DOUBAO().chat(question)
|