Backport `perf: cache the FunctionTool declaration` (57f3af24) to the 1.x line
- Vorherrschende Sprache
- Python
- Sterne
- 21.5k
- Forks
- 4k
- Ø Merge
- 1 T. 14 Std.
- Gemergte PRs (30 T.)
- 37
Beschreibung
### Is your feature request related to a specific problem?
`FunctionTool._get_declaration()` rebuilds the declaration from the function signature on every call
(`inspect.signature` -> pydantic `create_model` -> JSON schema generation). On the 1.x line there is
no caching, so the same result is rebuilt every time.
This was fixed on `main` by `perf: cache the FunctionTool declaration`
(https://github.com/google/adk-python/commit/57f3af24), shipped in 2.6.0, but it has not been
backported to 1.x. I checked the raw `src/google/adk/tools/function_tool.py` at each tag: no cache in
`v1.39.0` (current 1.x latest), `lru_cache` present in `v2.6.0`.
It matters more than it looks because `trace_call_llm` is called from inside the streaming loop
(`src/google/adk/flows/llm_flows/base_llm_flow.py`: span opened at `:1596`, `async for llm_response
in agen:` at `:1668`, `trace_call_llm(...)` at `:1669`), so it runs once per streamed chunk. Any
OpenTelemetry instrumentation that walks `llm_request.tools_dict` and reads each tool's declaration
ends up calling `_get_declaration()` chunks x tools times per turn. With a few hundred chunks and a
handful of tools that is thousands of rebuilds per answer, all synchronous CPU on the event loop,
which delays every other coroutine in the process.
### Describe the Solution You'd Like
Backport 57f3af24 to the 1.x line. It is self-contained: a module-level `lru_cache` keyed on the
function plus the parameters that affect the schema, with `model_copy(deep=True)` on return so
callers still get an independent object.
### Impact on your work
Not blocking, there are workarounds on the caller side. Filing it because 1.x users have no way to
get the fix short of a major-version upgrade.
### Willingness to contribute
Happy to, though I understand the repo imports from an internal source of truth, so an issue seemed
more useful than a PR.
### Describe Alternatives You've Considered
Upgrading to 2.x, which is the obvious answer but is a major-version migration for a small isolated
change.
### Additional Context
Measurements on google-adk 1.39.0, `timeit` `number=20 repeat=7 min`, Linux x86_64, CPython 3.12.3:
| annotation style | params | schema bytes | `_get_declaration()` | per param |
|---|---|---|---|---|
| `Annotated` + `Field` | 5 / 15 / 27 | 815 / 2250 / 3990 B | 1.962 / 5.486 / 9.651 ms | ~357-392 us |
| `Annotated` + `Field` + `Literal` | 5 / 15 / 27 | 733 / 2018 / 3590 B | 2.216 / 6.299 / 11.311 ms | ~419-443 us |
| plain type hints | 5 / 15 / 27 | 296 / 691 / 1171 B | 0.317 / 0.796 / 1.377 ms | ~51-63 us |
`same content: True same object: False` - identical result, rebuilt every call.
Cost is linear in the number of parameters, and a parameter carrying
`Annotated[..., pydantic.Field(...)]` costs roughly 8x a plainly annotated one. Schema size explains
little: the 3990 B row is faster than the 3590 B one. So the tools hit hardest are the ones written
the way the docs encourage, with many parameters each carrying `Field(description=...)` and often a
`Literal` enum.
Extrapolating the 27-parameter tool: 4 tools x 300 chunks is 1200 calls, roughly 12 s of CPU;
7 tools x 300 chunks is 2100 calls, roughly 20 s.
Reproduction, self-contained, generated tools only:
repro.py
```python
"""Cost of FunctionTool._get_declaration() by annotation style. Self-contained."""
import platform
import timeit
from typing import Annotated, Literal, Optional
import google.adk
from google.adk.tools import FunctionTool
from pydantic import Field
def make_annotated_tool(n: int) -> FunctionTool:
params = ", ".join(
f"p{i}: Annotated[Optional[str], Field(description='Parameter number {i} "
f"used to filter results by attribute {i}.')] = None"
for i in range(n)
)
ns = {"Annotated": Annotated, "Optional": Optional, "Field": Field}
exec(f"def annotated_tool({params}) -> dict:\n '''doc'''\n return {{}}\n", ns)
return FunctionTool(func=ns["annotated_tool"])
def make_literal_tool(n: int) -> FunctionTool:
params = ", ".join(
f"p{i}: Annotated[Optional[Literal['a{i}','b{i}','c{i}']], "
f"Field(description='Enum parameter {i}.')] = None"
for i in range(n)
)
ns = {"Annotated": Annotated, "Optional": Optional, "Literal": Literal, "Field": Field}
exec(f"def literal_tool({params}) -> dict:\n '''doc'''\n return {{}}\n", ns)
return FunctionTool(func=ns["literal_tool"])
def make_plain_tool(n: int) -> FunctionTool:
params = ", ".join(f"p{i}: Optional[str] = None" for i in range(n))
ns = {"Optional": Optional}
exec(f"def plain_tool({params}) -> dict:\n '''doc'''\n return {{}}\n", ns)
return FunctionTool(func=ns["plain_tool"])
def ms_per_call(tool: FunctionTool, number: int = 20, repeat: int = 7) -> float:
tool._get_declaration()
return min(timeit.repeat(tool._get_declaration, number=number, repeat=repeat)) / number * 1000
print("python ", platform.python_version(), platform.machine())
print("google-adk", getattr(google.adk, "__version__", "?"))
print()
print(f"{'annotation style':32s} {'params':>6s} {'bytes':>7s} {'ms/call':>9s} {'us/param':>9s}")
for label, factory in (
("Annotated + Field", make_annotated_tool),
("Annotated + Field + Literal", make_literal_tool),
("plain type hints", make_plain_tool),
):
for n in (5, 15, 27):
t = factory(n)
d = t._get_declaration()
ms = ms_per_call(t)
print(
f"{label:32s} {n:6d} {len(d.model_dump_json(exclude_none=True)):6d}B "
f"{ms:9.3f} {ms / n * 1000:9.0f}"
)
t = make_literal_tool(27)
a, b = t._get_declaration(), t._get_declaration()
print()
print("same content:", a.model_dump_json() == b.model_dump_json(), " same object:", a is b)
```
Related: #4233 reports the same class of problem from the serialization side.
I have also opened https://github.com/Arize-ai/openinference/pull/3591 against `Arize-ai/openinference`, which stops that instrumentation from
re-deriving request-side attributes on every chunk. That removes most of the multiplier from the
instrumentation side, but the rebuild cost on 1.x remains for any other caller.
Beitragsleitfaden
Bewertung
Dieses Issue wurde noch nicht bewertet.