modify: aibox_sophon_adapter plugins

doubao/__init__.py

@@ -2,10 +2,14 @@
 豆包视觉理解 Worker 插件
 """
 import os
+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, cleanup_ark_client
+from .global_var import (
+    set_ark_client, set_aicc_client, set_is_aicc,
+    cleanup_all,
+)
 from .router import *
 
 # 自动绑定 gen_models 生成的 Pydantic 模型
@@ -15,25 +19,42 @@ async def create_workers(session, config, count=1, worker_manager=None):
     """
     工厂函数：由 app 生命周期调用，创建 Worker 协程。
     """
-    # 从配置或环境变量获取 API Key
+    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:
         raise ValueError("缺少 API_KEY，请在配置中设置 api_key 或环境变量 VOLC_ACCESSKEY")
 
-    base_url = config.get("base_url", "https://ark.cn-beijing.volces.com/api/v3")
     model = config.get("model", "doubao-seed-evolving")
-
-    # 将模型名存入环境变量，供 router 使用
     os.environ["DOUBAO_MODEL"] = model
 
-    # 创建全局客户端
-    client = AsyncArk(base_url=base_url, api_key=api_key)
-    set_ark_client(client)
+    if is_aicc:
+        # AICC 模式：懒加载，避免非 AICC 环境导入失败
+        from bytedance.volcengine_aicc_sdk.secure_http_client import (
+            AICCConfig, AsyncSecureHttpClient,
+        )
+
+        base_url = config.get("base_url", "https://ark.cn-beijing.volces.com/api/v3")
+        os.environ["DOUBAO_BASE_URL"] = base_url
+
+        aicc_config = AICCConfig(ra_url=base_url, auth_token=api_key)
+        client = AsyncOpenAI(
+            api_key=api_key,
+            base_url=base_url,
+            http_client=AsyncSecureHttpClient(aicc_config=aicc_config),
+        )
+        set_aicc_client(client)
+    else:
+        # Ark 模式：使用 AsyncArk
+        base_url = config.get("base_url", "https://ark.cn-beijing.volces.com/api/v3")
+        client = AsyncArk(base_url=base_url, api_key=api_key)
+        set_ark_client(client)
 
     # 委托给通用运行时创建 Worker 任务
     return _create_workers(session, config, count, worker_manager)
 
 async def cleanup(worker_item: str):
     """清理资源"""
-    cleanup_ark_client()
+    cleanup_all()
     print(f"{worker_item} 清理完毕")