perf(tools): cache the FunctionTool declaration between LLM requests#6481
Open
dexhunter wants to merge 1 commit into
Open
perf(tools): cache the FunctionTool declaration between LLM requests#6481dexhunter wants to merge 1 commit into
dexhunter wants to merge 1 commit into
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description of change
FunctionTool._get_declaration()rebuilds the declaration from the function signature every time it is called, andLlmRequest.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 andpydantic.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, at8addc447):_get_declaration()—get_temperature_get_declaration()—list_devices_get_declaration()— a tool taking/returning a pydantic modelappend_tools()with all five toolsCall frequency, counted by instrumenting the real
Runner.run_asyncloop rather than estimated: exactly one call per tool per LLM request — 150 calls over 30 LLM requests with 5 tools registered. An agent loop withktool-call rounds pays itk+1times per user turn.This is a synchronous
defreached fromasync def process_llm_request, so it occupies the event loop for that whole time. The quantity it limits is the request throughput ofadk 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:tools=[FunctionTool(f), ...]— beforetools=[FunctionTool(f), ...]— afterBaseToolsetholding tools, no prefix — beforeBaseToolsetholding tools, no prefix — afterBaseToolsetwithtool_name_prefix— beforeBaseToolsetwithtool_name_prefix— aftertools=[plain_function, ...]— beforetools=[plain_function, ...]— afterThe last row is the honest limitation. When
toolscontains a bare callable,_convert_tool_union_to_toolsinagents/llm_agent.pybuilds a brand-newFunctionTool(func=tool_union)on every request, so the cache is cold every time and this change does nothing. ABaseToolinstance 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, whereBaseToolsetkeeps 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/samplesregisters bare callables, so the samples themselves will not move. HandcraftedGoogleToolsubclasses (tools/google_tool.py) areFunctionTools supplied by a toolset and do take the improving path. I could not instantiateBigQueryToolset/SpannerToolset/BigtableToolsetto count their tools empirically, because theirgoogle-cloud-*dependencies were not installed, so I am not putting a number on those.MCPTooland the OpenAPI/REST tools override_get_declarationand 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_declarationappends a note todeclaration.descriptionTransferToAgentTool._get_declarationwrites an enum intoparameters_json_schema['properties']['agent_name']BaseToolsetrewritesdeclaration.namewhen applying a tool name prefixHanding 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,
LongRunningFunctionToolappends 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, becauseFunctionDeclaration.model_validatekeeps a reference to the nestedparameters_json_schemadict rather than copying it, soTransferToAgentToolwrites 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:funcand_ignore_paramsare plain attributes, and_api_variantresolvesGOOGLE_GENAI_USE_ENTERPRISE/GOOGLE_GENAI_USE_VERTEXAIfrom the environment on every access. The variant matters —build_function_declarationdropsresponse_json_schemafor anything other thanVERTEX_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 thatmainalready satisfies and that this change must not break — they are regression guards for the mutation sites above, and they pass onmaintoo.No test was weakened, skipped or marked xfail.
ruff,isort,pyinkandcodespellare clean on both touched files, andmypyreports no error onfunction_tool.pythat is not already present onmain.Manual E2E evidence
Runner setup: an
LlmAgentnamedhome_automation_agentwith the five tools ofcontributing/samples/evaluation/home_automation_agent, driven throughRunner.run_asyncfor 15 user turns against a scripted in-processBaseLlmthat 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_schemais wrapped with a counter, and CPU istime.process_time_ns().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_inforather than silently skipping the prefix path.There is a note at
plugins/bigquery_agent_analytics_plugin.py:366-369saying 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.