LongArticlesJob/applications/api/deep_seek_api_by_volcanoengine.py

"""
@author: luojunhui
"""
from typing import Any, Dict, List, Optional

from volcenginesdkarkruntime import Ark

from applications.utils.common import safe_json_parse
from config import deep_seek_model
from config import deep_seek_default_model
from config import deep_seek_api_key_byte_dance


def fetch_deepseek_completion(
    model: str,
    prompt: str,
    output_type: str = "text",
) -> Optional[Dict | List]:
    input_messages = [{"role": "user", "content": prompt}]
    kwargs: dict[str, Any] = {
        "model": deep_seek_model.get(model, deep_seek_default_model),
        "messages": input_messages,
    }

    client = Ark(
        base_url="https://ark.cn-beijing.volces.com/api/v3",
        api_key=deep_seek_api_key_byte_dance
    )

    try:
        response = client.chat.completions.create(**kwargs)
        output_text = response.choices[0].message.content

        if output_type == "text":
            return output_text
        elif output_type == "json":
            return safe_json_parse(output_text)
        else:
            raise ValueError(f"Invalid output_type: {output_type}")

    except Exception as e:
        print(f"[ERROR] fetch_deepseek_completion (volcengine) failed: {e}")
        return None