chongming-plugins/doubao/__init__.py

"""
豆包视觉理解 Worker 插件
"""
import os
import logging
from openai import AsyncOpenAI
from volcenginesdkarkruntime import AsyncArk
from chongming_worker import create_workers as _create_workers
from chongming_worker import auto_bind_models
from .global_var import (
    set_ark_client, set_aicc_client, set_is_aicc,
    cleanup_all,
)
from .router import *

logger = logging.getLogger("plugins.doubao")

# 自动绑定 gen_models 生成的 Pydantic 模型
auto_bind_models("doubao")

async def create_workers(session, config, count=1, worker_manager=None):
    """
    工厂函数：由 app 生命周期调用，创建 Worker 协程。
    """
    logger.info("豆包插件初始化开始 (count=%d, aicc=%s)", count, 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),
        )
        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)

    # 委托给通用运行时创建 Worker 任务
    logger.info("豆包插件初始化完成，创建 %d 个 Worker 副本", count)
    return _create_workers(session, config, count, worker_manager)

async def cleanup(worker_item: str):
    """清理资源"""
    logger.info("豆包插件清理: %s", worker_item)
    cleanup_all()
    logger.info("豆包插件清理完毕: %s", worker_item)