chongming-plugins/doubao/__init__.py

"""
豆包视觉理解 Worker 插件

基于 Plugin 基类统一生命周期：on_init 中按配置创建 Ark / AICC 客户端，
on_cleanup 释放客户端引用（无数据库）。
"""
import os
import logging
from openai import AsyncOpenAI
from volcenginesdkarkruntime import AsyncArk

from chongming_worker.plugin import Plugin
from .global_var import (
    set_ark_client, set_aicc_client, set_is_aicc,
    cleanup_all,
)
from .router import *  # noqa: F401,F403  触发 @rpc_method 注册

logger = logging.getLogger(__name__)


class DoubaoPlugin(Plugin):
    """豆包插件 — 统一生命周期。"""

    type = "doubao"
    db_name = None      # 无数据库
    auto_bind = True    # 首次初始化时自动 auto_bind_models("doubao")

    async def on_init(self):
        """首次 create_workers 时调用一次：初始化豆包客户端。"""
        config = self.config
        logger.info("豆包插件初始化开始 (count=%s, aicc=%s)",
                    config.get("count", 1), config.get("is_aicc", False))

        is_aicc = config.get("is_aicc", False)
        set_is_aicc(is_aicc)

        api_key = config.get("api_key") or os.environ.get("VOLC_ACCESSKEY")
        if not api_key:
            logger.error("豆包插件初始化失败: 缺少 API_KEY")
            raise ValueError("缺少 API_KEY，请在配置中设置 api_key 或环境变量 VOLC_ACCESSKEY")

        model = config.get("model", "doubao-seed-evolving")
        os.environ["DOUBAO_MODEL"] = model
        logger.info("豆包模型: %s", model)

        if is_aicc:
            logger.info("豆包使用 AICC 安全通信模式")
            # AICC 模式：懒加载，避免非 AICC 环境导入失败
            try:
                from bytedance.volcengine_aicc_sdk.secure_http_client import (
                    AICCConfig, AsyncSecureHttpClient,
                )
            except ImportError as e:
                logger.error("AICC SDK 导入失败: %s，请确认已安装 bytedance_volcengine_aicc_sdk", e)
                raise

            base_url = config.get("base_url", "https://ark.cn-beijing.volces.com/api/v3")
            # ra_url 为远程证明（安全令牌）端点地址，可能与推理 base_url 不同
            # 若不单独配置则默认使用 base_url 的域名 + /v1/security/token
            ra_url = config.get("ra_url", base_url)
            ak = config.get("ak", "")
            sk = config.get("sk", "")
            ra_service_name = config.get("ra_service_name", "AICC.ConfidentialInference")
            ra_policy_id = config.get("ra_policy_id", "router_policy")

            os.environ["DOUBAO_BASE_URL"] = base_url
            logger.info("豆包 AICC base_url: %s, ra_url: %s, ra_service_name: %s, ra_policy_id: %s",
                        base_url, ra_url, ra_service_name, ra_policy_id)

            aicc_config = AICCConfig(
                ra_url=ra_url,
                auth_token=api_key,
                ak=ak,
                sk=sk,
                aicc_llm_service_name=ra_service_name,
                policy_id=ra_policy_id,
            )
            client = AsyncOpenAI(
                api_key=api_key,
                base_url=base_url,
                http_client=AsyncSecureHttpClient(aicc_config=aicc_config), # type: ignore
            )
            set_aicc_client(client)
            logger.info("豆包 AICC 客户端已创建")
        else:
            logger.info("豆包使用标准 Ark SDK 模式")
            # Ark 模式：使用 AsyncArk
            base_url = config.get("base_url", "https://ark.cn-beijing.volces.com/api/v3")
            os.environ["DOUBAO_BASE_URL"] = base_url
            client = AsyncArk(base_url=base_url, api_key=api_key)
            set_ark_client(client)
            logger.info("豆包 Ark 客户端已创建 (base_url=%s)", base_url)

        logger.info("豆包插件初始化完成")

    async def on_cleanup(self):
        """cleanup 时调用：释放客户端引用。"""
        cleanup_all()
        logger.info("[doubao] 清理完成")


# 暴露给 app lifespan 的模块级入口（bound method，签名与旧写法完全一致）
plugin = DoubaoPlugin()
create_workers = plugin.create_workers
cleanup = plugin.cleanup