39 lines
1.4 KiB
Python
39 lines
1.4 KiB
Python
|
"""
|
||
|
|
豆包视觉理解 Worker 插件
|
||
|
|
"""
|
||
|
|
import os
|
||
|
|
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 .router import *
|
||
|
|
|
||
|
|
# 自动绑定 gen_models 生成的 Pydantic 模型
|
||
|
|
auto_bind_models("doubao")
|
||
|
|
|
||
|
|
async def create_workers(session, config, count=1, worker_manager=None):
|
||
|
|
"""
|
||
|
|
工厂函数：由 app 生命周期调用，创建 Worker 协程。
|
||
|
|
"""
|
||
|
|
# 从配置或环境变量获取 API Key
|
||
|
|
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)
|
||
|
|
|
||
|
|
# 委托给通用运行时创建 Worker 任务
|
||
|
|
return _create_workers(session, config, count, worker_manager)
|
||
|
|
|
||
|
|
async def cleanup(worker_item: str):
|
||
|
|
"""清理资源"""
|
||
|
|
cleanup_ark_client()
|
||
|
|
print(f"{worker_item} 清理完毕")
|