107 lines
4.3 KiB
Python
107 lines
4.3 KiB
Python
"""
|
|
豆包视觉理解 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
|