modify: 统一插件风格

2026-08-14 16:04:39 +08:00
parent 56dd4f0ebb
commit 4fa066432d
10 changed files with 516 additions and 436 deletions
@@ -15,7 +15,7 @@

 ## 配置

-在 `config/workers/doubao/config.toml` 和 `config/workers/config.toml` 中设置：
+在 `config/workers/doubao/config.toml` 中设置：

 ```toml
 [workers.config]
@@ -64,9 +64,11 @@ model = "doubao-seed-evolving"                              # 模型名称

 ## 实现细节

- 使用 `httpx` 异步 HTTP 客户端调用火山引擎 API
- 支持流式响应（streaming）中的思考模式 token
- API Key 优先从环境变量 `VOLC_ACCESSKEY` 读取
+- 生命周期基于 `Plugin` 基类：`on_init` 按 `is_aicc` 分支创建 Ark（`AsyncArk`）或
+  AICC（`AsyncOpenAI` + `AsyncSecureHttpClient`）客户端并写入 `global_var`；
+  `on_cleanup` 释放客户端引用（无数据库）
+- API Key 优先取配置 `api_key`，其次环境变量 `VOLC_ACCESSKEY`（缺失时插件初始化报错）
+- 支持 AICC 安全通信模式（懒加载 SDK，避免非 AICC 环境导入失败）
 - 超时时间默认 30 秒，在 `config.toml` 的 `defaults.timeout` 中配置

 ---
@@ -1,95 +1,106 @@
 """
 豆包视觉理解 Worker 插件
+
+基于 Plugin 基类统一生命周期：on_init 中按配置创建 Ark / AICC 客户端，
+on_cleanup 释放客户端引用（无数据库）。
 """
 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 chongming_worker.plugin import Plugin
 from .global_var import (
    set_ark_client, set_aicc_client, set_is_aicc,
    cleanup_all,
 )
-from .router import *
+from .router import *  # noqa: F401,F403  触发 @rpc_method 注册

-logger = logging.getLogger("plugins.doubao")
+logger = logging.getLogger(__name__)

-# 自动绑定 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))
+class DoubaoPlugin(Plugin):
+    """豆包插件 — 统一生命周期。"""

-    is_aicc = config.get("is_aicc", False)
-    set_is_aicc(is_aicc)
+    type = "doubao"
+    db_name = None      # 无数据库
+    auto_bind = True    # 首次初始化时自动 auto_bind_models("doubao")

-    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")
+    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))

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

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

-        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")
+        logger.info("豆包插件初始化完成")

-        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)
+    async def on_cleanup(self):
+        """cleanup 时调用：释放客户端引用。"""
+        cleanup_all()
+        logger.info("[doubao] 清理完成")

-        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)
+# 暴露给 app lifespan 的模块级入口（bound method，签名与旧写法完全一致）
+plugin = DoubaoPlugin()
+create_workers = plugin.create_workers
+cleanup = plugin.cleanup