67 lines
2.1 KiB
Python
67 lines
2.1 KiB
Python
# Please install OpenAI SDK first: `pip3 install openai`
|
||
import os
|
||
import json
|
||
from openai import OpenAI
|
||
from dotenv import load_dotenv
|
||
|
||
load_dotenv()
|
||
|
||
|
||
class DeepSeek:
|
||
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("DEEPSEEK_API_KEY")
|
||
if not self.key:
|
||
raise Exception("找不到 DEEPSEEK_API_KEY 请设置环境变量在运行代码")
|
||
self.client = OpenAI(
|
||
api_key=self.key,
|
||
base_url="https://api.deepseek.com",
|
||
)
|
||
|
||
def chat(self, question_json: dict) -> str:
|
||
response = self.client.chat.completions.create(
|
||
# model="deepseek-chat",
|
||
# model="deepseek-reasoner",
|
||
# model="deepseek-v4-pro",
|
||
model="deepseek-v4-flash",
|
||
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": "机器翻译"},
|
||
],
|
||
}
|
||
DeepSeek().chat(question)
|