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"}