TOOL_CALL_MAX_RETRY = 10
TOOL_CALL_TIMEOUT = 20.0 # 单次工具调用最大时长(秒
rate_limiter = AsyncLimiter(max_rate=20, time_period=1)
LOCAL_TEST_MODEL = "doubao-seed-2-0-mini-260428"
"description": "搜索的关键词,1~100个字符(过长会截断),不支持多词搜索",
"description": "搜索结果数量",
name="code_sandbox_grader_single_grader",
description="code_sandbox_grader",
instance=PluginInstance.CPU2MEM4,
proxy: RolloutInferenceProxy,
sample: ChatCompletionSample,
) -> Optional[RolloutResult]:
训练时,每个样本会调用 n_sample 次rollout函数,每次调用rollout函数当前仅支持返回一条轨迹(wrap_client_inplace的client最后一次调用上下文)
# NOTE: 由于search mcp的封装,无法修改tool定义,使用直接调用api的方式
# from arkitect.core.component.tool.tool_pool import ToolPool, build_tool_pool
# from arkitect.core.component.tool.builder import build_mcp_clients_from_config
# mcp_clients, cleanup = build_mcp_clients_from_config(CONFIG_FILE_PATH)
# tool_pool = build_tool_pool(list(mcp_clients.values()))
# tools = [tool.model_dump() for tool in await tool_pool.list_tools()]
schema = search_tool["function"]["parameters"]
search_api_key = os.getenv("SEARCH_API_KEY")
assert search_api_key, "env SEARCH_API_KEY not set"
async with httpx.AsyncClient(
base_url="https://open.feedcoopapi.com",
"Authorization": f"Bearer {search_api_key}",
"Content-Type": "application/json",
req = sample.model_dump()
messages = req.pop("messages")
client = AsyncArk(base_url=proxy.url, api_key=proxy.jwt_token)
wrap_inplace_trace(client)
wrap_async_client_inplace(client, proxy=proxy)
completion = await client.chat.completions.create(
# NOTE: 为了避免训练效果受影响,被wrap后的client无法在此处指定采样参数或thinking type, 需要在样本或任务超参中指定
messages.append(completion.choices[0].message.model_dump())
if completion.choices[0].finish_reason != "tool_calls":
tool_calls = completion.choices[0].message.tool_calls
for tool_call in tool_calls or []:
name = tool_call.function.name
params = json.loads(tool_call.function.arguments)
assert name == "web_search", f"tool name not web_search: {name}"
jsonschema.validate(instance=params, schema=schema)
jsonschema.ValidationError,
# NOTE: extra字段用于传递额外的信息给reward函数,此处通知reward需要惩罚为-1分
logger.error(f"tool call format check failed: {repr(e)}")
status=PluginStatus.SUCCESS,
res = await exec_tool_call(
tool_client=tool_client,
"tool_call_id": tool_call.id,
status=PluginStatus.SUCCESS,
jitter=backoff.full_jitter, # full jitter
process_inputs=lambda d: {key: d["kwargs"][key] for key in ["name", "arguments"]},
async def exec_tool_call(
tool_client: httpx.AsyncClient,
) -> str | list[ChatCompletionContentPartParam]:
# from arkitect.core.component.tool.utils import convert_to_chat_completion_content_part_param
# if isinstance(tool_client, ToolPool):
# res = await tool_client.execute_tool(name, arguments)
# content = convert_to_chat_completion_content_part_param(res)
async with rate_limiter:
resp = await asyncio.wait_for(
"search_api/web_search",
json={**arguments, "SearchType": "web"},
timeout=TOOL_CALL_TIMEOUT,
return json.dumps(data, ensure_ascii=False)
# NOTE: 可以选择一些错误返回RETRY进行rollout重试,重试达到上线后导致任务失败,默认exception为Failure会重试后拉挂任务
raise RolloutRetryException(
f"web_search tool call got resp without Result: {data}"
except asyncio.TimeoutError:
raise RolloutRetryException(
f"tool call timeout after {TOOL_CALL_TIMEOUT}s"
except httpx.HTTPStatusError as e:
raise RolloutRetryException(
f"tool call HTTPStatusError: {e.response.status_code} {e.response.text}"