chat-ai-api/app/core/ark_client.py

from openai import OpenAI
from volcenginesdkarkruntime import Ark
from ..db.ai_config import get_config_by_app_name

_client_cache: dict[str, OpenAI] = {}
_ark_client_cache: dict[str, Ark] = {}


def get_ark_client(app_name: str = "com.yunxiangshengtai") -> Ark:
    if app_name in _ark_client_cache:
        return _ark_client_cache[app_name]

    config = get_config_by_app_name(app_name)
    if not config:
        raise ValueError(f"未找到appName '{app_name}' 的配置")

    client = Ark(
        base_url=config["baseUrl"],
        api_key=config["apiKey"]
    )
    _ark_client_cache[app_name] = client
    return client

def get_client(app_name: str = "com.yunxiangshengtai") -> OpenAI:
    if app_name in _client_cache:
        return _client_cache[app_name]

    config = get_config_by_app_name(app_name)
    if not config:
        raise ValueError(f"未找到appName '{app_name}' 的配置")

    client = OpenAI(
        base_url=config["baseUrl"],
        api_key=config["apiKey"],
        default_headers={"ark-beta-knowledge-search": "true"}
    )
    _client_cache[app_name] = client
    return client