From 98eff4a24f8fc828f1bc0ea6d3ad46f963929eed Mon Sep 17 00:00:00 2001 From: ColtenOuO Date: Sat, 25 Jul 2026 12:21:22 +0000 Subject: [PATCH] Decode base64 for bytes-typed HookToolset parameters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit JSON has no binary type, so HookToolset advertises a hook method's bytes-typed parameters as strings, and nothing converted them back. A method such as S3Hook.load_bytes received a str and failed deep inside BytesIO with a TypeError that gave no hint about where the bad value came from. What that string was meant to contain was never stated. Models routinely send base64 for binary payloads, and read as UTF-8 that text would have been written to storage verbatim — a successful upload of the wrong bytes. Declaring the encoding makes the contract explicit, and a value that is not valid base64 now fails loudly as a retryable error rather than being stored. The schema and the decoding are driven by one condition so a parameter can never be advertised as base64 without also being decoded. Annotations are resolved per parameter because get_type_hints is all-or-nothing: a single TYPE_CHECKING-only import elsewhere in a signature, as in CloudKMSHook.encrypt, would otherwise turn the conversion off for the whole method with nothing to show for it. --- providers/common/ai/docs/toolsets.rst | 13 ++ .../providers/common/ai/toolsets/hook.py | 111 ++++++++++-- .../unit/common/ai/toolsets/test_hook.py | 171 +++++++++++++++++- 3 files changed, 274 insertions(+), 21 deletions(-) diff --git a/providers/common/ai/docs/toolsets.rst b/providers/common/ai/docs/toolsets.rst index a4b08bf869356..b8ca60ea8695c 100644 --- a/providers/common/ai/docs/toolsets.rst +++ b/providers/common/ai/docs/toolsets.rst @@ -110,6 +110,19 @@ Parameters - ``tool_name_prefix``: Optional prefix prepended to each tool name (e.g. ``"s3_"`` produces ``"s3_list_keys"``). +Binary parameters +^^^^^^^^^^^^^^^^^ + +JSON has no binary type, so a ``bytes``-typed parameter is advertised to the +model as a base64 string (``contentEncoding: base64``, repeated in the +parameter description because most function-calling APIs drop the keyword) and +decoded before the hook method is called. A value that is not valid base64 is +rejected and fed back to the model as a retry rather than stored as the wrong +bytes. + +This applies to ``bytes`` and ``bytes | None`` only. A parameter that already +accepts text, such as ``bytes | str``, is passed through untouched. + ``SQLToolset`` -------------- diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py index 63037e1a4f8f9..3c3999a50b61d 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py @@ -18,12 +18,14 @@ from __future__ import annotations +import base64 import inspect import json import re import types from typing import TYPE_CHECKING, Any, Union, get_args, get_origin, get_type_hints +from pydantic_ai.exceptions import ModelRetry from pydantic_ai.tools import ToolDefinition from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool @@ -47,6 +49,10 @@ bytes: {"type": "string"}, } +_BASE64_PARAM_NOTE = "Provide this value base64-encoded." + +_MAX_UNRESOLVABLE_ANNOTATION_NAMES = 20 + class HookToolset(AbstractToolset[Any]): """ @@ -86,6 +92,7 @@ def __init__( self._allowed_methods = allowed_methods self._tool_name_prefix = tool_name_prefix self._id = f"hook-{type(hook).__name__}" + self._bytes_params: dict[str, frozenset[str]] = {} @property def id(self) -> str: @@ -97,14 +104,23 @@ async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]: method = getattr(self._hook, method_name) tool_name = f"{self._tool_name_prefix}{method_name}" if self._tool_name_prefix else method_name - json_schema = _build_json_schema_from_signature(method) + json_schema, bytes_params = _introspect_signature(method) + self._bytes_params[tool_name] = bytes_params description = _extract_description(method) param_docs = _parse_param_docs(method.__doc__ or "") + properties = json_schema.get("properties", {}) # Enrich parameter descriptions from docstring. for param_name, param_desc in param_docs.items(): - if param_name in json_schema.get("properties", {}): - json_schema["properties"][param_name]["description"] = param_desc + if param_name in properties: + properties[param_name]["description"] = param_desc + + # After the loop above, which would otherwise overwrite the note. + for param_name in bytes_params & properties.keys(): + existing = properties[param_name].get("description") + properties[param_name]["description"] = ( + f"{existing} {_BASE64_PARAM_NOTE}" if existing else _BASE64_PARAM_NOTE + ) # sequential=True because hook methods perform synchronous I/O # (network calls, DB queries) and should not run concurrently. @@ -136,7 +152,14 @@ async def call_tool( ) -> Any: method_name = name.removeprefix(self._tool_name_prefix) if self._tool_name_prefix else name method: Callable[..., Any] = getattr(self._hook, method_name) - result = method(**tool_args) + bytes_params = self._bytes_params.get(name) + if bytes_params is None: + bytes_params = _introspect_signature(method)[1] + # Decoding belongs here rather than in the args validator: validated args + # travel the whole toolset chain, and CachingToolset fingerprints them with + # a plain json.dumps, so bytes upstream of this point would make every + # binary call unverifiable on durable replay. + result = method(**_decode_bytes_args(tool_args, bytes_params)) return _serialize_for_llm(result) @@ -173,17 +196,76 @@ def _python_type_to_json_schema(annotation: Any) -> dict[str, Any]: return dict(schema) if schema else {} -def _build_json_schema_from_signature(method: Callable[..., Any]) -> dict[str, Any]: - """Build a JSON Schema ``object`` from a method's signature and type hints.""" - sig = inspect.signature(method) +def _resolves_to_bytes(annotation: Any) -> bool: + """Whether ``annotation`` is ``bytes`` or ``Optional[bytes]``/``bytes | None``.""" + if annotation is bytes: + return True + origin = get_origin(annotation) + if origin is types.UnionType or origin is Union: + non_none = [a for a in get_args(annotation) if a is not type(None)] + if len(non_none) == 1: + return _resolves_to_bytes(non_none[0]) + return False - try: - hints = get_type_hints(method) - except Exception: - hints = {} + +def _decode_bytes_args(tool_args: dict[str, Any], bytes_params: frozenset[str]) -> dict[str, Any]: + """Decode the base64 strings the model supplied for ``bytes``-typed parameters.""" + if not bytes_params: + return tool_args + + decoded: dict[str, Any] = {} + for key, value in tool_args.items(): + if key not in bytes_params or not isinstance(value, str): + decoded[key] = value + continue + try: + # ``validate=True``: the default discards invalid characters instead of + # erroring, which is the silent corruption this decoding exists to avoid. + decoded[key] = base64.b64decode(value, validate=True) + except ValueError as e: + # ModelRetry is the only exception pydantic-ai feeds back to the model; + # anything else fails the run without giving it a chance to correct. + raise ModelRetry(f"Parameter {key!r} must be base64-encoded binary data.") from e + return decoded + + +def _resolve_annotations(method: Callable[..., Any]) -> dict[str, Any]: + """ + Resolve ``method``'s annotations, tolerating names that are not importable at runtime. + + ``get_type_hints`` is all-or-nothing: a single unresolvable annotation discards + the hints for every parameter, which would silently stop base64 decoding for the + rest of the signature. ``CloudKMSHook.encrypt`` hits this today — its ``bytes`` + parameters sit next to ``retry: Retry | _MethodDefault``, where ``Retry`` is + imported under ``TYPE_CHECKING``. Unresolvable names are substituted with ``Any`` + so the parameters that *can* be resolved still are. + """ + localns: dict[str, Any] = {} + for _ in range(_MAX_UNRESOLVABLE_ANNOTATION_NAMES): + try: + return get_type_hints(method, localns=localns) + except NameError as e: + if not e.name or e.name in localns: + break + localns[e.name] = Any + except TypeError: + break + return {} + + +def _introspect_signature(method: Callable[..., Any]) -> tuple[dict[str, Any], frozenset[str]]: + """ + Build ``method``'s JSON Schema and the names of its ``bytes``-typed parameters. + + Both come from one pass so the schema advertised to the model and the arguments + decoded on the way back can never disagree about which parameters are binary. + """ + sig = inspect.signature(method) + hints = _resolve_annotations(method) properties: dict[str, Any] = {} required: list[str] = [] + bytes_params: set[str] = set() allows_additional_properties = False for name, param in sig.parameters.items(): @@ -198,6 +280,11 @@ def _build_json_schema_from_signature(method: Callable[..., Any]) -> dict[str, A annotation = hints.get(name, param.annotation) prop = _python_type_to_json_schema(annotation) properties[name] = prop + # One condition drives both, so a parameter can never be advertised as + # base64 without also being decoded on the way back. + if _resolves_to_bytes(annotation): + bytes_params.add(name) + prop["contentEncoding"] = "base64" if param.default is inspect.Parameter.empty: required.append(name) @@ -207,7 +294,7 @@ def _build_json_schema_from_signature(method: Callable[..., Any]) -> dict[str, A schema["required"] = required if allows_additional_properties: schema["additionalProperties"] = True - return schema + return schema, frozenset(bytes_params) def _extract_description(method: Callable[..., Any]) -> str: diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py b/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py index ae2d4c6f86754..340f264072eae 100644 --- a/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py +++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py @@ -17,20 +17,28 @@ from __future__ import annotations import asyncio +import base64 +import json +from typing import TYPE_CHECKING from unittest.mock import MagicMock import pytest +from pydantic_ai.exceptions import ModelRetry from pydantic_core import ValidationError from airflow.providers.common.ai.toolsets.hook import ( + _BASE64_PARAM_NOTE, HookToolset, - _build_json_schema_from_signature, _extract_description, + _introspect_signature, _parse_param_docs, _serialize_for_llm, ) from airflow.providers.common.ai.utils.tool_definition import _SUPPORTS_RETURN_SCHEMA +if TYPE_CHECKING: + from decimal import Decimal + class _FakeHook: """Fake hook for testing HookToolset introspection.""" @@ -55,6 +63,22 @@ def request( ) -> dict[str, object]: return {"endpoint": endpoint, "data": data, **kwargs} + def upload_bytes(self, data: bytes, key: str) -> str: + """Upload raw bytes to storage. + + :param data: Content to store. + :param key: Destination key. + """ + return f"uploaded {len(data)} bytes to {key} (type={type(data).__name__})" + + def upload_optional_bytes(self, data: bytes | None = None) -> str: + """Upload optional raw bytes to storage.""" + return f"data type={type(data).__name__}" + + def echo_bytes(self, data: bytes) -> str: + """Return the hex representation of raw bytes.""" + return data.hex() + class TestHookToolsetInit: def test_requires_non_empty_allowed_methods(self): @@ -148,6 +172,29 @@ def test_param_docs_enriched_in_schema(self): assert "description" in props["bucket"] assert "S3 bucket" in props["bucket"]["description"] + def test_bytes_params_declare_base64_encoding(self): + ts = HookToolset(_FakeHook(), allowed_methods=["upload_bytes"]) + tools = asyncio.run(ts.get_tools(ctx=MagicMock())) + + props = tools["upload_bytes"].tool_def.parameters_json_schema["properties"] + assert props["data"] == { + "type": "string", + "contentEncoding": "base64", + "description": f"Content to store. {_BASE64_PARAM_NOTE}", + } + assert props["key"] == {"type": "string", "description": "Destination key."} + + def test_base64_note_survives_a_param_without_docs(self): + ts = HookToolset(_FakeHook(), allowed_methods=["upload_optional_bytes"]) + tools = asyncio.run(ts.get_tools(ctx=MagicMock())) + + props = tools["upload_optional_bytes"].tool_def.parameters_json_schema["properties"] + assert props["data"] == { + "anyOf": [{"type": "string"}, {"type": "null"}], + "contentEncoding": "base64", + "description": _BASE64_PARAM_NOTE, + } + class TestHookToolsetArgsValidator: @pytest.fixture @@ -203,13 +250,119 @@ def test_dispatches_with_prefix(self): ) assert result == "contents of test.txt" + @pytest.mark.parametrize( + ("method_name", "tool_args", "expected"), + [ + pytest.param( + "upload_bytes", + {"data": "aGVsbG8gd29ybGQ=", "key": "greeting.txt"}, + "uploaded 11 bytes to greeting.txt (type=bytes)", + id="bytes", + ), + pytest.param("upload_optional_bytes", {"data": "aGk="}, "data type=bytes", id="optional-bytes"), + pytest.param("upload_optional_bytes", {"data": None}, "data type=NoneType", id="explicit-null"), + pytest.param("upload_optional_bytes", {}, "data type=NoneType", id="omitted"), + ], + ) + def test_decodes_base64_for_bytes_params(self, method_name, tool_args, expected): + ts = HookToolset(_FakeHook(), allowed_methods=[method_name]) + tools = asyncio.run(ts.get_tools(ctx=MagicMock())) + + result = asyncio.run(ts.call_tool(method_name, tool_args, ctx=MagicMock(), tool=tools[method_name])) + assert result == expected + + def test_decoded_bytes_are_byte_exact(self): + ts = HookToolset(_FakeHook(), allowed_methods=["echo_bytes"]) + tools = asyncio.run(ts.get_tools(ctx=MagicMock())) + payload = b"\x89PNG\r\n\x1a\n" + + result = asyncio.run( + ts.call_tool( + "echo_bytes", + {"data": base64.b64encode(payload).decode("ascii")}, + ctx=MagicMock(), + tool=tools["echo_bytes"], + ) + ) + assert result == payload.hex() + + @pytest.mark.parametrize( + "value", ["hello world", "not base64!!", "abc"], ids=["text", "punctuation", "bad-padding"] + ) + def test_undecodable_value_asks_the_model_to_retry(self, value): + ts = HookToolset(_FakeHook(), allowed_methods=["upload_bytes"]) + tools = asyncio.run(ts.get_tools(ctx=MagicMock())) + + with pytest.raises(ModelRetry, match="'data' must be base64-encoded"): + asyncio.run( + ts.call_tool( + "upload_bytes", + {"data": value, "key": "greeting.txt"}, + ctx=MagicMock(), + tool=tools["upload_bytes"], + ) + ) + + +def _takes_bytes(value: bytes): ... + + +def _takes_optional_bytes(value: bytes | None = None): ... + + +def _takes_bytes_or_str(value: bytes | str): ... + + +def _takes_list_of_bytes(value: list[bytes]): ... + + +def _takes_str(value: str): ... + + +def _takes_unannotated(value): ... + + +class TestBytesParamDetection: + @pytest.mark.parametrize( + ("func", "is_bytes_param"), + [ + pytest.param(_takes_bytes, True, id="bytes"), + pytest.param(_takes_optional_bytes, True, id="optional-bytes"), + pytest.param(_takes_bytes_or_str, False, id="bytes-or-str"), + pytest.param(_takes_list_of_bytes, False, id="list-of-bytes"), + pytest.param(_takes_str, False, id="str"), + pytest.param(_takes_unannotated, False, id="unannotated"), + ], + ) + def test_only_bytes_and_optional_bytes_are_decoded(self, func, is_bytes_param): + schema, bytes_params = _introspect_signature(func) + assert ("value" in bytes_params) is is_bytes_param + # A parameter advertised as base64 that is never decoded would hand the + # hook the encoded text — the corruption this decoding exists to prevent. + assert ("base64" in json.dumps(schema["properties"]["value"])) is is_bytes_param + + def test_var_args_are_never_decoded(self): + def fn(*chunks: bytes, **extra: bytes): ... + + schema, bytes_params = _introspect_signature(fn) + assert bytes_params == frozenset() + assert schema["properties"] == {} + + def test_unresolvable_annotation_does_not_disable_sibling_params(self): + # Mirrors CloudKMSHook.encrypt, whose bytes params sit next to a parameter + # annotated with a TYPE_CHECKING-only import. + def fn(data: bytes, precision: Decimal | None = None): ... + + _, bytes_params = _introspect_signature(fn) + assert bytes_params == {"data"} + class TestBuildJsonSchemaFromSignature: def test_basic_types(self): def fn(name: str, count: int, rate: float, active: bool): pass - schema = _build_json_schema_from_signature(fn) + schema, _ = _introspect_signature(fn) assert schema["properties"]["name"] == {"type": "string"} assert schema["properties"]["count"] == {"type": "integer"} assert schema["properties"]["rate"] == {"type": "number"} @@ -220,7 +373,7 @@ def test_optional_params_accept_null(self): def fn(name: str, prefix: str | None = None): pass - schema = _build_json_schema_from_signature(fn) + schema, _ = _introspect_signature(fn) assert schema["required"] == ["name"] assert schema["properties"]["prefix"] == {"anyOf": [{"type": "string"}, {"type": "null"}]} @@ -228,28 +381,28 @@ def test_union_types(self): def fn(data: dict[str, object] | str): pass - schema = _build_json_schema_from_signature(fn) + schema, _ = _introspect_signature(fn) assert schema["properties"]["data"] == {"anyOf": [{"type": "object"}, {"type": "string"}]} def test_list_type(self): def fn(items: list[str]): pass - schema = _build_json_schema_from_signature(fn) + schema, _ = _introspect_signature(fn) assert schema["properties"]["items"] == {"type": "array", "items": {"type": "string"}} def test_no_annotation_is_untyped(self): def fn(x): pass - schema = _build_json_schema_from_signature(fn) + schema, _ = _introspect_signature(fn) assert schema["properties"]["x"] == {} def test_kwargs_allow_additional_properties(self): def fn(x: int, **kwargs): pass - schema = _build_json_schema_from_signature(fn) + schema, _ = _introspect_signature(fn) assert schema["additionalProperties"] is True def test_skips_self_and_cls(self): @@ -257,14 +410,14 @@ class Foo: def method(self, x: int): pass - schema = _build_json_schema_from_signature(Foo().method) + schema, _ = _introspect_signature(Foo().method) assert "self" not in schema["properties"] def test_skips_var_args(self): def fn(x: int, *args, **kwargs): pass - schema = _build_json_schema_from_signature(fn) + schema, _ = _introspect_signature(fn) assert set(schema["properties"].keys()) == {"x"}