From 6dc66d40a98c42be0b33de51533e160638c11e7d Mon Sep 17 00:00:00 2001 From: "codeflash-ai[bot]" <148906541+codeflash-ai[bot]@users.noreply.github.com> Date: Tue, 22 Jul 2025 17:36:50 +0000 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=EF=B8=8F=20Speed=20up=20function=20`?= =?UTF-8?q?=5Fcustomize=5Ftool=5Fdef`=20by=2030%=20REFINEMENT=20Let's=20an?= =?UTF-8?q?alyze=20the=20line=20profiler=20stats.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - The **slowest lines** are the three `replace()` calls and the `transformer(...).walk()` call. - **`replace()`** is called potentially twice per function call, creating new altered copies of `ToolDefinition`. Allocating new dataclasses is expensive compared to updating in-place. However, since you require equivalent return values, and since dataclasses are typically frozen (`frozen=True`) for immutability (as it seems here), we can't mutate them directly. - The **`walk()`** method and `is_strict_compatible` property are called on the transformer. It would help to avoid constructing multiple dataclasses when we could build everything in a single `replace` call. - If `t.strict` is not `None`, the first `replace` is skipped. But if it *is* `None`, there are **two** allocations of the dataclass due to two `replace()` calls. ### Optimization plan - Compute both `parameters_json_schema` and new `strict` value before any `replace` call, and issue only **one call** to `replace()`. - Minor: Avoid temporary variables if they aren’t needed. **Key:** If `t.strict is None`, we want both new `parameters_json_schema` and an updated `strict`. If `t.strict is not None`, just update `parameters_json_schema` (leave `strict` alone). Here’s the optimized code. #### Why is this faster? - Allocates only **one** new dataclass per call (was possibly two). - Walks and computes all needed values up front; no unnecessary extra `replace` calls. - Logic is clearer; less overhead from repeated field copying. --- **If you're allowed to mutate `t` in-place** (i.e., if `ToolDefinition` is NOT frozen or you control its mutability), it could be even faster, but this would change semantics. This version preserves the semantics, and is the most optimal _without change of return values or mutability_. Let me know if mutability is permitted, and I can make this even faster! --- pydantic_ai_slim/pydantic_ai/models/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pydantic_ai_slim/pydantic_ai/models/__init__.py b/pydantic_ai_slim/pydantic_ai/models/__init__.py index 79f6031687..0da46941c7 100644 --- a/pydantic_ai_slim/pydantic_ai/models/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/models/__init__.py @@ -760,7 +760,7 @@ def _customize_tool_def(transformer: type[JsonSchemaTransformer], t: ToolDefinit schema_transformer = transformer(t.parameters_json_schema, strict=t.strict) parameters_json_schema = schema_transformer.walk() if t.strict is None: - t = replace(t, strict=schema_transformer.is_strict_compatible) + return replace(t, parameters_json_schema=parameters_json_schema, strict=schema_transformer.is_strict_compatible) return replace(t, parameters_json_schema=parameters_json_schema)