DataOps-platform/app/core/llm/deepseek_client.py

"""
DeepSeek API helpers (OpenAI-compatible SDK).
"""

from __future__ import annotations

import os
from typing import Any

from flask import current_app, has_app_context
from openai import OpenAI

DEEPSEEK_DEFAULT_BASE_URL = "https://api.deepseek.com"


def _clean_secret(value: str | None) -> str:
    if not value:
        return ""
    return str(value).strip().strip("\r\n\t")


def get_llm_api_key() -> str:
    """Resolve API key: DEEPSEEK_API_KEY first, then LLM_API_KEY / app config."""
    candidates = []
    if has_app_context():
        candidates.extend(
            [
                current_app.config.get("DEEPSEEK_API_KEY"),
                current_app.config.get("LLM_API_KEY"),
            ]
        )
    candidates.extend([os.environ.get("DEEPSEEK_API_KEY"), os.environ.get("LLM_API_KEY")])

    for candidate in candidates:
        cleaned = _clean_secret(candidate)
        if cleaned and cleaned not in {
            "replace-with-your-deepseek-api-key",
            "your-api-key",
        }:
            return cleaned
    return ""


def normalize_llm_base_url(raw: str | None = None) -> str:
    """Normalize DeepSeek OpenAI-compatible base URL for SDK usage."""
    if raw is None:
        if has_app_context():
            raw = str(current_app.config.get("LLM_BASE_URL") or "")
        if not raw:
            raw = os.environ.get("LLM_BASE_URL", DEEPSEEK_DEFAULT_BASE_URL)

    url = _clean_secret(raw) or DEEPSEEK_DEFAULT_BASE_URL
    url = url.rstrip("/")
    # OpenAI SDK 会自动追加 /v1；若 env 已带 /v1，去掉以避免重复
    if url.endswith("/v1"):
        url = url[:-3]
    return url or DEEPSEEK_DEFAULT_BASE_URL


def get_llm_base_url() -> str:
    return normalize_llm_base_url()


def get_llm_chat_completions_url() -> str:
    """Return the HTTP endpoint for raw requests.post() callers."""
    return f"{get_llm_base_url()}/v1/chat/completions"


def get_llm_model() -> str:
    raw = ""
    if has_app_context():
        raw = str(current_app.config.get("LLM_MODEL_NAME") or "")
    if not raw:
        raw = os.environ.get("LLM_MODEL_NAME", "deepseek-chat")
    return _clean_secret(raw) or "deepseek-chat"


def create_llm_client() -> OpenAI:
    api_key = get_llm_api_key()
    if not api_key:
        raise ValueError(
            "DeepSeek API Key 未配置，请在 /etc/dataops-platform/dataops.env 中设置 DEEPSEEK_API_KEY"
        )
    return OpenAI(api_key=api_key, base_url=get_llm_base_url())


def chat_completions_create(
    client: OpenAI,
    *,
    messages: list[dict[str, str]],
    temperature: float = 0.7,
    max_tokens: int = 1024,
    use_thinking: bool = False,
    **kwargs: Any,
):
    """Create a chat completion; optional DeepSeek thinking mode for complex tasks."""
    create_kwargs: dict[str, Any] = {
        "model": get_llm_model(),
        "messages": messages,
        "stream": False,
        "temperature": temperature,
        "max_tokens": max_tokens,
        **kwargs,
    }
    if use_thinking:
        create_kwargs["reasoning_effort"] = current_app.config.get(
            "LLM_REASONING_EFFORT", "high"
        )
        create_kwargs["extra_body"] = {"thinking": {"type": "enabled"}}
    return client.chat.completions.create(**create_kwargs)