76 lines
2.4 KiB
Python
76 lines
2.4 KiB
Python
import os
|
||
import json
|
||
from openai import OpenAI
|
||
|
||
|
||
def deepseek(user_query, need_extra=False):
|
||
"""
|
||
调用 deepseek 接口生成回答。
|
||
|
||
参数：
|
||
user_query (str): 用户输入的查询内容。
|
||
need_extra (bool): 是否额外返回 usage 信息（completion_tokens），默认为 False。
|
||
|
||
返回：
|
||
如果 need_extra 为 True，则返回一个元组 (full_response, completion_tokens)；
|
||
否则仅返回 full_response。
|
||
"""
|
||
client = OpenAI(
|
||
# 如果没有配置环境变量，请直接将 api_key 替换为实际的 API Key，例如：api_key="sk-xxx"
|
||
api_key=os.getenv("DASHSCOPE_API_KEY"),
|
||
base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
|
||
)
|
||
|
||
# 调用接口，开启流式输出并包含 usage 信息
|
||
completion = client.chat.completions.create(
|
||
model="deepseek-r1-distill-qwen-32b",
|
||
temperature=0.5,
|
||
messages=[
|
||
{'role': 'user', 'content': user_query}
|
||
],
|
||
stream=True,
|
||
stream_options={"include_usage": True}
|
||
)
|
||
|
||
full_response = ""
|
||
completion_tokens = 0
|
||
|
||
for chunk in completion:
|
||
# 解析 chunk 数据：优先使用 to_dict 方法，否则采用 model_dump_json 进行转换
|
||
if hasattr(chunk, 'to_dict'):
|
||
chunk_data = chunk.to_dict()
|
||
else:
|
||
chunk_data = json.loads(chunk.model_dump_json())
|
||
|
||
# 处理 usage 信息
|
||
usage = chunk_data.get('usage')
|
||
if usage is not None:
|
||
completion_tokens = usage.get('completion_tokens', 0)
|
||
|
||
# 处理 choices 信息
|
||
choices = chunk_data.get('choices', [])
|
||
if choices:
|
||
choice = choices[0]
|
||
delta = choice.get('delta', {})
|
||
content = delta.get('content', '')
|
||
if content:
|
||
full_response += content
|
||
# 如需实时输出可取消注释下面一行
|
||
print(content, end='', flush=True)
|
||
# 若需要处理完成原因，可在此处增加相应逻辑
|
||
if choice.get('finish_reason'):
|
||
pass
|
||
|
||
# 根据 need_extra 决定返回值
|
||
if need_extra:
|
||
return full_response, completion_tokens
|
||
else:
|
||
return full_response
|
||
|
||
|
||
# 示例调用
|
||
if __name__ == "__main__":
|
||
query = "1+1等于几？请用json格式回答，键名为'答案',键值为你的回答"
|
||
result = deepseek(query, need_extra=True)
|
||
print("\n完整内容为：", result)
|