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)
|