"""
@author: luojunhui
"""
from openai import OpenAI

from config import deep_seek_model
from config import deep_seek_default_model
from config import deep_seek_api_key_byte_dance


def deep_seek_api(model, prompt):
    """
    deep_seek方法
    """
    client = OpenAI(
        api_key=deep_seek_api_key_byte_dance,
        base_url="https://ark.cn-beijing.volces.com/api/v3",
        # 深度推理模型耗费时间会较长，建议您设置一个较长的超时时间，推荐为30分钟
        timeout=1800,
        )
    response = client.chat.completions.create(
        model=deep_seek_model.get(model, deep_seek_default_model),
        messages=[
            {"role": "user", "content": prompt}
        ]
    )
    return response.choices[0].message.content