Skip to content

perf(tools): cache the FunctionTool declaration between LLM requests#6481

Open
dexhunter wants to merge 1 commit into
google:mainfrom
dexhunter:perf/cache-function-tool-declaration
Open

perf(tools): cache the FunctionTool declaration between LLM requests#6481
dexhunter wants to merge 1 commit into
google:mainfrom
dexhunter:perf/cache-function-tool-declaration

Conversation

@dexhunter

Copy link
Copy Markdown
Contributor

Description of change

FunctionTool._get_declaration() rebuilds the declaration from the function signature every time it is called, and LlmRequest.append_tools() calls it once per tool on every LLM request. Nothing about that result changes between requests unless the function, the ignored parameters or the API variant change.

The build is not cheap. It runs pydantic.create_model() + model_json_schema() for the parameters and pydantic.TypeAdapter(...).json_schema() for the return type.

Absolute cost, measured on the five tools of contributing/samples/evaluation/home_automation_agent (CPU time, median of 300 reps, Python 3.12.13, pydantic 2.13.4, google-genai 2.14.0, at 8addc447):

what CPU
one _get_declaration()get_temperature 0.53 ms
one _get_declaration()list_devices 0.71 ms
one _get_declaration() — a tool taking/returning a pydantic model 1.10 ms
append_tools() with all five tools 3.48 ms per LLM request
the same, with this change 0.26 ms per LLM request

Call frequency, counted by instrumenting the real Runner.run_async loop rather than estimated: exactly one call per tool per LLM request — 150 calls over 30 LLM requests with 5 tools registered. An agent loop with k tool-call rounds pays it k+1 times per user turn.

This is a synchronous def reached from async def process_llm_request, so it occupies the event loop for that whole time. The quantity it limits is the request throughput of adk api_server, not the latency a single user sees behind a provider call.

Where this helps, and where it does not

This is the part worth reading before reviewing the diff.

Measured end to end through Runner.run_async, 15 user turns / 30 LLM requests, same interpreter and same environment for both sides:

how tools are registered schema builds per LLM request framework CPU per LLM request
tools=[FunctionTool(f), ...] — before 5.0 6.41 ms
tools=[FunctionTool(f), ...] — after 0.167 2.93 ms
a BaseToolset holding tools, no prefix — before 5.0 6.07 ms
a BaseToolset holding tools, no prefix — after 0.167 2.85 ms
a BaseToolset with tool_name_prefix — before 5.0 6.20 ms
a BaseToolset with tool_name_prefix — after 0.167 2.84 ms
tools=[plain_function, ...] — before 5.0 6.28 ms
tools=[plain_function, ...] — after 5.0 6.36 ms — no change

The last row is the honest limitation. When tools contains a bare callable, _convert_tool_union_to_tools in agents/llm_agent.py builds a brand-new FunctionTool(func=tool_union) on every request, so the cache is cold every time and this change does nothing. A BaseTool instance on the other hand is returned as-is and is therefore long-lived, which is why the toolset rows do improve — including the prefixed path, where BaseToolset keeps its shallow copies in _cached_prefixed_tools.

I deliberately did not also memoise the wrapper in _convert_tool_union_to_tools, which is what would extend this to bare callables. It changes tool object identity and lifetime across requests, which is a separate concern from this one and belongs in its own change if you want it.

Every sample agent in contributing/samples registers bare callables, so the samples themselves will not move. Handcrafted GoogleTool subclasses (tools/google_tool.py) are FunctionTools supplied by a toolset and do take the improving path. I could not instantiate BigQueryToolset/SpannerToolset/BigtableToolset to count their tools empirically, because their google-cloud-* dependencies were not installed, so I am not putting a number on those.

MCPTool and the OpenAPI/REST tools override _get_declaration and reuse a supplied schema, so they never touched this path and are unaffected either way.

Why a copy is returned on every call

Callers are allowed to mutate the declaration they receive, and three places in this repository do:

  • LongRunningFunctionTool._get_declaration appends a note to declaration.description
  • TransferToAgentTool._get_declaration writes an enum into parameters_json_schema['properties']['agent_name']
  • BaseToolset rewrites declaration.name when applying a tool name prefix

Handing back a shared cached object is measurably faster (0.18 ms rather than 0.26 ms for the five tools) and is wrong. With a shared object, LongRunningFunctionTool appends its note again on every request: I measured its description growing 195 → 328 → 461 → 594 characters over four requests, and that description is sent to the model. Caching the built dict and re-validating it per call is also unsafe, because FunctionDeclaration.model_validate keeps a reference to the nested parameters_json_schema dict rather than copying it, so TransferToAgentTool writes straight into the cache.

So the cache stores the expensive build and each call gets a private copy.deepcopy. The 0.08 ms that copy costs relative to the shared-object version is the price of not changing behaviour, and I would rather pay it.

The key is (func, ignore_params, api variant) because all three can change on a live tool: func and _ignore_params are plain attributes, and _api_variant resolves GOOGLE_GENAI_USE_ENTERPRISE/GOOGLE_GENAI_USE_VERTEXAI from the environment on every access. The variant matters — build_function_declaration drops response_json_schema for anything other than VERTEX_AI.

Testing plan

New file tests/unittests/tools/test_function_tool_declaration_cache.py, 11 tests. One of them (test_declaration_is_built_once_for_repeated_calls) is the new behaviour and fails without this change. The other ten are invariants that main already satisfies and that this change must not break — they are regression guards for the mutation sites above, and they pass on main too.

$ pytest tests/unittests/tools/test_function_tool_declaration_cache.py -q
11 passed, 5 warnings in 1.30s

$ pytest tests/unittests/tools -q
1659 passed, 299 warnings in 51.60s

$ pytest tests/unittests/models tests/unittests/agents tests/unittests/flows -q
1992 passed, 2 xfailed, 587 warnings in 36.70s

No test was weakened, skipped or marked xfail. ruff, isort, pyink and codespell are clean on both touched files, and mypy reports no error on function_tool.py that is not already present on main.

Manual E2E evidence

Runner setup: an LlmAgent named home_automation_agent with the five tools of contributing/samples/evaluation/home_automation_agent, driven through Runner.run_async for 15 user turns against a scripted in-process BaseLlm that answers one tool call and then one final message, so the reported CPU is framework only and no network time is included. build_function_declaration_with_json_schema is wrapped with a counter, and CPU is time.process_time_ns().

# tools registered as long-lived FunctionTool instances
before: {"llm_requests": 30, "schema_builds": 150, "builds_per_llm_request": 5.0,   "framework_cpu_ms_per_llm_request": 6.4930}
after:  {"llm_requests": 30, "schema_builds": 5,   "builds_per_llm_request": 0.167, "framework_cpu_ms_per_llm_request": 2.9522}

# the same agent with a BaseToolset using tool_name_prefix="demo"
before: {"llm_requests": 30, "schema_builds": 150, "builds_per_llm_request": 5.0,   "framework_cpu_ms_per_llm_request": 6.2036}
after:  {"llm_requests": 30, "schema_builds": 5,   "builds_per_llm_request": 0.167, "framework_cpu_ms_per_llm_request": 2.8437}

# tools registered as bare callables - unchanged, as explained above
before: {"llm_requests": 30, "schema_builds": 150, "builds_per_llm_request": 5.0,   "framework_cpu_ms_per_llm_request": 6.2778}
after:  {"llm_requests": 30, "schema_builds": 150, "builds_per_llm_request": 5.0,   "framework_cpu_ms_per_llm_request": 6.3598}

Both sides were run in the same interpreter, same virtualenv and same worktree, toggling only this diff, alternating and repeated. The declaration name, description and parameter schema were verified identical before and after for every tool, and the prefixed toolset run was checked to actually produce demo_get_device_info rather than silently skipping the prefix path.

There is a note at plugins/bigquery_agent_analytics_plugin.py:366-369 saying this rebuild "repeats work the framework already did when assembling the request" and to "revisit with a cache if it shows up on the hot path". This is that revisit; if you would rather leave it alone for now, say so and I will close this.

Supplementary autoresearch

Fourteen structurally distinct approaches to this were implemented and measured against the same workload before settling on the one in this diff, including the ones that were faster but unshippable and the ones that barely helped: https://dashboard.weco.ai/share/0ivwDzA07OLt8bLf4rj0X6xgrAt3KRdm. That is an external record of the exploration, published for transparency; the diff here is written by hand.

FunctionTool._get_declaration() rebuilds the declaration from the function
signature on every call, and LlmRequest.append_tools() calls it once per tool
on every LLM request. The build runs pydantic create_model() plus
model_json_schema() for the parameters and a TypeAdapter().json_schema() for
the return type, which costs 0.53-0.71 ms of CPU per tool for the five tools of
contributing/samples/evaluation/home_automation_agent.

Cache the built declaration per tool, keyed on the function, the ignored
parameters and the API variant, since all three can change during the life of
a tool. Callers mutate the declaration they are given - LongRunningFunctionTool
appends to its description, TransferToAgentTool writes an enum into its
parameter schema and BaseToolset rewrites its name when prefixing - so each
call still returns a private copy and the cached value stays pristine.
@adk-bot adk-bot added the tools [Component] This issue is related to tools label Jul 25, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

tools [Component] This issue is related to tools

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants