google / google/adk-python

Backport `perf: cache the FunctionTool declaration` (57f3af24) to the 1.x line

Aperta
#6,833 1 commento 0 reazioni 1 assegnatario Rivendicata da @sanketpatil06 Vedi su GitHub
needs review tools
Lingua principale
Python
Stelle
21.5k
Fork
4k
Merge medio
1g 14h
PR unite (30g)
37

Descrizione

### 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.

Guida per i contributori

Apri la guida per i contributori

Valutazione

Questa issue non è ancora stata valutata.

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.