74 lines
2.4 KiB
Python
74 lines
2.4 KiB
Python
# pip install openai python-dotenv
|
||
import os
|
||
import json
|
||
from openai import OpenAI
|
||
from dotenv import load_dotenv
|
||
from httpx import Client, Proxy
|
||
|
||
load_dotenv()
|
||
|
||
|
||
class KIMI:
|
||
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.getenv(
|
||
"KIMI_API_KEY", os.getenv("KIMI_API_KEY")
|
||
)
|
||
if not self.key:
|
||
raise RuntimeError("找不到 KIMI_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://api.moonshot.cn/v1", # 已去掉多余空格
|
||
http_client=httpx_client,
|
||
)
|
||
|
||
def chat(self, question_json: dict) -> str:
|
||
resp = self.client.chat.completions.create(
|
||
model="kimi-k2-turbo-preview",
|
||
messages=[
|
||
{"role": "system", "content": self.system_prompt},
|
||
{
|
||
"role": "user",
|
||
"content": json.dumps(question_json, ensure_ascii=False),
|
||
},
|
||
],
|
||
temperature=0,
|
||
)
|
||
return resp.choices[0].message.content or ""
|
||
|
||
|
||
# 本地单条测试
|
||
if __name__ == "__main__":
|
||
question = {
|
||
"id": "25594265",
|
||
"index": "1",
|
||
"type": "单选",
|
||
"title": "下列哪个属于AI视觉识别应用？",
|
||
"chooies": [
|
||
{"num": "A", "txt": "音乐合成"},
|
||
{"num": "B", "txt": "植物识别"},
|
||
{"num": "C", "txt": "文案创作"},
|
||
{"num": "D", "txt": "机器翻译"},
|
||
],
|
||
}
|
||
print(KIMI().chat(question)) # 期望输出：B
|