rag_server/applications/api/deepseek.py

"""
@author: luojunhui
@description: deepseek 官方api (async版)
"""

from __future__ import annotations

import json
from typing import Dict, List, Optional
from openai import AsyncOpenAI

from applications.config import DEEPSEEK_MODEL
from applications.config import DEEPSEEK_API_KEY


async def fetch_deepseek_completion(
    model: str,
    prompt: str,
    output_type: str = "text",
    tool_calls: bool = False,
    tools: List[Dict] = None,
) -> Optional[Dict | List]:
    messages = [{"role": "user", "content": prompt}]
    kwargs = {
        "model": DEEPSEEK_MODEL.get(model, "deepseek-chat"),
        "messages": messages,
    }

    # add tool calls
    if tool_calls and tools:
        kwargs["tools"] = tools
        kwargs["tool_choice"] = "auto"

    client = AsyncOpenAI(api_key=DEEPSEEK_API_KEY, base_url="https://api.deepseek.com")

    if output_type == "json":
        kwargs["response_format"] = {"type": "json_object"}

    try:
        response = await client.chat.completions.create(**kwargs)
        choice = response.choices[0]

        if output_type == "text":
            return choice.message.content  # 只返回文本
        elif output_type == "json":
            return json.loads(choice.message.content)
        else:
            raise ValueError(f"Invalid output_type: {output_type}")

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


__all__ = ["fetch_deepseek_completion"]