Skip to content

Fix LLM call limits for non-ADK runtimes - #1038

Open
FirstayZheng wants to merge 1 commit into
volcengine:mainfrom
FirstayZheng:fix/non-adk-max-llm-calls
Open

Fix LLM call limits for non-ADK runtimes#1038
FirstayZheng wants to merge 1 commit into
volcengine:mainfrom
FirstayZheng:fix/non-adk-max-llm-calls

Conversation

@FirstayZheng

Copy link
Copy Markdown
Contributor

背景

RunConfig(max_llm_calls=...)MODEL_AGENT_MAX_LLM_CALLS 依赖 ADK 原生 BaseLlmFlow._call_llm_async 在真正模型调用前执行 invocation_context.increment_llm_call_count()。Codex 和 PiAgent runtime 绕过了 ADK 原生 LLM flow,因此此前这些 runtime 不会消耗 LLM 调用预算,导致 max_llm_calls 实际失效。

修复内容

  • Codex runtime:

    • ResponsesShim 的 turn context 中新增 before_model_call 回调。
    • 每次执行 litellm.aresponses(...) 前调用该回调。
    • 在 Codex runtime 注册 turn 时传入 ctx.increment_llm_call_count
    • 同一个 Codex turn 内如果 tool loop 触发多次 backend model call,会逐次计数;超限时会在下一次 backend call 前抛出 LlmCallsLimitExceededError,避免实际请求模型。
    • shim 内记录 turn-scoped 异常,并由 runtime 侧重新抛出,保留 ADK 原生限流异常语义。
  • PiAgent runtime:

    • 在调用 client.prompt(prompt) 前执行一次 ctx.increment_llm_call_count()
    • 注释和测试中明确这是 per prompt 的保底计数;Pi 子进程内部多轮模型调用目前 veADK 无法精确观测。
  • 测试:

    • 新增 Codex tool loop 超限用例,验证第二次 backend call 前触发限流且 backend 未执行。
    • 新增 Codex 普通单次调用计数用例。
    • 新增 PiAgent per prompt 计数用例。
    • 新增 PiAgent 超限时不进入 client.prompt(prompt) 用例。
    • 更新 Codex SDK contract test 的 fake shim/context 以覆盖新接口。

验证

已执行:

.venv/bin/pre-commit run --files veadk/runtime/codex/proxy.py veadk/runtime/codex/runtime.py veadk/runtime/piagent/runtime.py tests/runtime/codex/test_codex_runtime.py tests/runtime/codex/test_codex_runtime_sdk.py tests/runtime/piagent/test_piagent_runtime.py
.venv/bin/python -m pytest tests/runtime/codex/test_codex_runtime.py tests/runtime/codex/test_codex_runtime_sdk.py tests/runtime/piagent/test_piagent_runtime.py

结果:86 passed

另跑过全量 pytest,当前仓库有 4 个与本次 runtime 改动无关的既有/非相关失败:两个 generated-agent debug 用例在全量顺序下 504 但单跑通过;两个 Studio RBAC media audit log 用例单跑仍因 caplog 抓不到项目自定义 stdout logger 失败。

手动集成验证:

  • Codex:第一次 backend call 后进入 tool loop,第二次 backend call 前触发 LlmCallsLimitExceededErrorbackend_calls=1
  • PiAgent:超限时在 client.prompt(prompt) 前抛出 LlmCallsLimitExceededError,fake Pi 未收到 prompt。

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant