This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch hooktoolset-datetime-args in repository https://gitbox.apache.org/repos/asf/airflow.git
commit d9766b8f793f83a2ef79a0a84e7fa59d4acd0de0 Author: Kaxil Naik <[email protected]> AuthorDate: Wed Oct 14 19:28:54 2026 +0100 Fix HookToolset calls to hook methods that take a datetime A datetime parameter, such as from_datetime of S3Hook.list_keys, was passed the string the model sent, and the hook raised TypeError. It is now offered as an ISO 8601 date-time string and parsed before the call; a value with no UTC offset or one that does not parse goes back to the model as a retry. The bytes decoding and this share one table. --- providers/common/ai/docs/toolsets/hook.rst | 16 ++ .../airflow/providers/common/ai/toolsets/hook.py | 133 +++++++++---- .../ai/tests/unit/common/ai/toolsets/test_hook.py | 216 +++++++++++++++------ 3 files changed, 275 insertions(+), 90 deletions(-) diff --git a/providers/common/ai/docs/toolsets/hook.rst b/providers/common/ai/docs/toolsets/hook.rst index 089d8949630..955a1072ce8 100644 --- a/providers/common/ai/docs/toolsets/hook.rst +++ b/providers/common/ai/docs/toolsets/hook.rst @@ -194,6 +194,22 @@ The value of ``base64`` can be passed to a ``bytes`` parameter as it is; text ha to be encoded by the model first. ``bytes`` nested in a returned list or dict are rendered as UTF-8 text, with any byte that does not decode replaced. +.. _hook-toolset-datetimes: + +Dates and times +--------------- + +JSON has no date type either. A parameter typed ``datetime`` or ``datetime | None``, +such as ``from_datetime`` of ``S3Hook.list_keys``, is offered to the model as a string +in ``date-time`` format whose description asks for ISO 8601 with a UTC offset, and the +value is parsed into a timezone-aware ``datetime`` before the method is called. A value +with no offset, such as ``2026-09-01T00:00:00``, or one that does not parse, is refused +while the call's arguments are validated, and the model is asked to correct the call. +The toolset does not assume a timezone for it, because hooks disagree on one: +``S3Hook.list_keys`` and ``GCSHook.list_by_timespan`` raise ``TypeError`` for a naive +value, while ``GCSHook.is_updated_after`` treats it as UTC. A ``datetime`` nested in +another type, such as ``list[datetime]``, is not converted. + When to choose it ----------------- 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 5a6fa60c0c3..aa35dcc80b6 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 @@ -25,6 +25,8 @@ import inspect import re import types import warnings +from dataclasses import dataclass +from datetime import datetime from typing import TYPE_CHECKING, Any, Union, get_args, get_origin, get_type_hints from pydantic_ai.exceptions import ModelRetry @@ -58,6 +60,50 @@ _TYPE_MAP: dict[type, dict[str, Any]] = { } _BASE64_PARAM_NOTE = "Provide this value base64-encoded." +_DATETIME_PARAM_NOTE = "An ISO 8601 date-time with a UTC offset, such as 2026-09-01T00:00:00Z." + + +def _decode_base64(value: str) -> bytes: + # Whitespace never carries data, so line-wrapped base64 is accepted. validate=True + # because the default discards any other character outside the alphabet instead of + # failing, which would store the wrong bytes without a word. + return base64.b64decode("".join(value.split()), validate=True) + + +def _parse_aware_datetime(value: str) -> datetime: + parsed = datetime.fromisoformat(value) + if parsed.utcoffset() is None: + # A naive value is ambiguous: S3Hook.list_keys and GCSHook.list_by_timespan raise + # TypeError comparing it with an aware time, other hooks assume UTC for it. + raise ValueError(f"{value!r} has no UTC offset") + return parsed + + +@dataclass(frozen=True) +class _StringEncoding: + """How a parameter type JSON cannot carry is sent as a string, and turned back.""" + + schema: dict[str, Any] + """The JSON Schema of the string.""" + note: str + """Added to the parameter's description.""" + expected: str + """What the value must be, for the retry prompt.""" + decode: Callable[[str], Any] + """Raises ``ValueError`` for a value it cannot decode.""" + + +_STRING_ENCODINGS: dict[type, _StringEncoding] = { + bytes: _StringEncoding( + {"type": "string"}, _BASE64_PARAM_NOTE, "base64-encoded binary data", _decode_base64 + ), + datetime: _StringEncoding( + {"type": "string", "format": "date-time"}, + _DATETIME_PARAM_NOTE, + "an ISO 8601 date-time with a UTC offset", + _parse_aware_datetime, + ), +} class HookToolset(AirflowToolset): @@ -175,7 +221,7 @@ class HookToolset(AirflowToolset): if conn_attr is None or not hasattr(hook, conn_attr): conn_attr = "conn_id" if hasattr(hook, "conn_id") else None self._conn_attr = conn_attr - self._bytes_params = { + self._encoded_params = { method_name: _introspect_signature(getattr(hook, method_name))[1] for method_name in allowed_methods } @@ -208,7 +254,7 @@ class HookToolset(AirflowToolset): method = getattr(self._hook, method_name) tool_name = f"{self._name_prefix}{method_name}" - json_schema, bytes_params = _introspect_signature(method) + json_schema, encoded_params = _introspect_signature(method) description = _extract_description(method) param_docs = _parse_param_docs(method.__doc__ or "") properties = json_schema["properties"] @@ -220,10 +266,10 @@ class HookToolset(AirflowToolset): # After the loop above, which would otherwise overwrite the note. The encoding is # stated in the description rather than as ``contentEncoding``: Gemini drops that # keyword, and pydantic-ai passes it through to OpenAI strict mode. - for param_name in bytes_params: + for param_name, encoding in encoded_params.items(): existing = properties[param_name].get("description") properties[param_name]["description"] = ( - f"{existing} {_BASE64_PARAM_NOTE}" if existing else _BASE64_PARAM_NOTE + f"{existing} {encoding.note}" if existing else encoding.note ) _drop_properties(json_schema, self._pinned) # The validator accepts the pinned names, with any value, so a model that sends one @@ -254,18 +300,20 @@ class HookToolset(AirflowToolset): args_validator=build_args_validator(args_schema), # Refused during validation, so an approval gate never asks about such a call. args_validator_func=( - functools.partial(self._validate_args, bytes_params) - if self._pinned or bytes_params + functools.partial(self._validate_args, encoded_params) + if self._pinned or encoded_params else None ), ) return tools - def _validate_args(self, bytes_params: frozenset[str], ctx: RunContext[Any], /, **tool_args: Any) -> None: + def _validate_args( + self, encoded_params: dict[str, _StringEncoding], ctx: RunContext[Any], /, **tool_args: Any + ) -> None: self._refuse_pinned(ctx, **tool_args) - # Only checked here: the decoded bytes are not kept, so what the toolset chain passes + # Only checked here: the decoded values are not kept, so what the toolset chain passes # on, and the durable journal fingerprints, stays JSON. - _decode_bytes_args(tool_args, bytes_params) + _decode_string_args(tool_args, encoded_params) def _refuse_pinned(self, ctx: RunContext[Any], /, **tool_args: Any) -> None: if supplied := sorted(self._pinned.keys() & tool_args.keys()): @@ -289,8 +337,8 @@ class HookToolset(AirflowToolset): self._refuse_pinned(ctx, **tool_args) # Decoded here rather than in the args validator: validated arguments travel the whole # toolset chain, and the durable journal fingerprints them with a plain json.dumps, so bytes - # upstream of this point would make every binary call unverifiable on durable replay. - tool_args = _decode_bytes_args(tool_args, self._bytes_params[method_name]) + # or datetimes upstream of this point would make such calls unverifiable on durable replay. + tool_args = _decode_string_args(tool_args, self._encoded_params[method_name]) # A copy per call, so a method that modifies an argument it is given cannot change the pin. result = await self.run_blocking(method, **tool_args, **copy.deepcopy(self._pinned)) return serialize_for_llm(result) @@ -329,31 +377,41 @@ def _python_type_to_json_schema(annotation: Any) -> dict[str, Any]: return dict(schema) if schema else {} -def _resolves_to_bytes(annotation: Any) -> bool: - """Whether ``annotation`` is ``bytes`` or ``bytes | None``.""" - if annotation is bytes: - return True +def _string_encoding(annotation: Any) -> _StringEncoding | None: + """Return the encoding of ``annotation`` when it is ``T`` or ``T | None`` for a ``T`` in ``_STRING_ENCODINGS``.""" if get_origin(annotation) in (types.UnionType, Union): non_none = [arg for arg in get_args(annotation) if arg is not type(None)] - return len(non_none) == 1 and _resolves_to_bytes(non_none[0]) - return False + return _string_encoding(non_none[0]) if len(non_none) == 1 else None + return _STRING_ENCODINGS.get(annotation) -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.""" +def _string_encoded_schema(annotation: Any, encoding: _StringEncoding) -> dict[str, Any]: + """Return ``encoding.schema`` for ``T``, and it or null for ``T | None``.""" + if get_origin(annotation) in (types.UnionType, Union): + return { + "anyOf": [ + {"type": "null"} if arg is type(None) else _string_encoded_schema(arg, encoding) + for arg in get_args(annotation) + ] + } + return dict(encoding.schema) + + +def _decode_string_args( + tool_args: dict[str, Any], encoded_params: dict[str, _StringEncoding] +) -> dict[str, Any]: + """Decode the strings the model supplied for parameters of a type JSON cannot carry.""" decoded = dict(tool_args) - for name in bytes_params & tool_args.keys(): + for name in encoded_params.keys() & tool_args.keys(): value = tool_args[name] if not isinstance(value, str): continue + encoding = encoded_params[name] try: - # Whitespace never carries data, so line-wrapped base64 is accepted. validate=True - # because the default discards any other character outside the alphabet instead of - # failing, which would store the wrong bytes without a word. - decoded[name] = base64.b64decode("".join(value.split()), validate=True) + decoded[name] = encoding.decode(value) except ValueError as e: # ModelRetry asks the model to correct the call; a ValueError would fail the run. - raise ModelRetry(f"Parameter {name!r} must be base64-encoded binary data: {e}.") from e + raise ModelRetry(f"Parameter {name!r} must be {encoding.expected}: {e}.") from e return decoded @@ -380,19 +438,21 @@ def _resolve_annotations(method: Callable[..., Any]) -> dict[str, Any]: return hints -def _introspect_signature(method: Callable[..., Any]) -> tuple[dict[str, Any], frozenset[str]]: +def _introspect_signature( + method: Callable[..., Any], +) -> tuple[dict[str, Any], dict[str, _StringEncoding]]: """ - Build ``method``'s JSON Schema and the names of its ``bytes``-typed parameters. + Build ``method``'s JSON Schema and the string encoding of each parameter that needs one. - Both come from one pass, so the parameters described to the model as base64 are - exactly the ones decoded on the way back. + Both come from one pass, so the parameters described to the model as an encoded + string are exactly the ones decoded on the way back. """ sig = inspect.signature(method) hints = _resolve_annotations(method) properties: dict[str, Any] = {} required: list[str] = [] - bytes_params: set[str] = set() + encoded_params: dict[str, _StringEncoding] = {} allows_additional_properties = False for name, param in sig.parameters.items(): @@ -405,10 +465,13 @@ def _introspect_signature(method: Callable[..., Any]) -> tuple[dict[str, Any], f continue annotation = hints.get(name, param.annotation) - prop = _python_type_to_json_schema(annotation) - properties[name] = prop - if _resolves_to_bytes(annotation): - bytes_params.add(name) + if encoding := _string_encoding(annotation): + # Only here, not in _TYPE_MAP: a datetime nested in a list is not decoded, so it is + # not described as an encoded string either. + properties[name] = _string_encoded_schema(annotation, encoding) + encoded_params[name] = encoding + else: + properties[name] = _python_type_to_json_schema(annotation) if param.default is inspect.Parameter.empty: required.append(name) @@ -418,7 +481,7 @@ def _introspect_signature(method: Callable[..., Any]) -> tuple[dict[str, Any], f schema["required"] = required if allows_additional_properties: schema["additionalProperties"] = True - return schema, frozenset(bytes_params) + return schema, encoded_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 0c5b3314430..cdd31796a1e 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 @@ -22,6 +22,7 @@ import contextlib import json import re import threading +from datetime import UTC, datetime, timedelta, timezone from typing import TYPE_CHECKING from unittest.mock import MagicMock @@ -36,6 +37,7 @@ from pydantic_core import ValidationError from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.providers.common.ai.toolsets.hook import ( _BASE64_PARAM_NOTE, + _DATETIME_PARAM_NOTE, HookToolset, _extract_description, _introspect_signature, @@ -89,6 +91,13 @@ class _FakeHook: """Return the bytes it is given.""" return data + def list_modified(self, since: datetime, until: datetime | None = None) -> str: + """List objects modified in a period. + + :param since: Start of the period. + """ + return f"since={since!r} until={until!r}" + class TestHookToolsetInit: def test_requires_non_empty_allowed_methods(self): @@ -295,6 +304,21 @@ class TestHookToolsetGetTools: "description": _BASE64_PARAM_NOTE, } + def test_datetime_params_ask_for_iso_8601_with_an_offset(self): + ts = HookToolset(_FakeHook(), allowed_methods=["list_modified"]) + tools = asyncio.run(ts.get_tools(ctx=MagicMock(spec=RunContext))) + + props = tools["list_modified"].tool_def.parameters_json_schema["properties"] + assert props["since"] == { + "type": "string", + "format": "date-time", + "description": f"Start of the period. {_DATETIME_PARAM_NOTE}", + } + assert props["until"] == { + "anyOf": [{"type": "string", "format": "date-time"}, {"type": "null"}], + "description": _DATETIME_PARAM_NOTE, + } + class TestHookToolsetArgsValidator: @pytest.fixture @@ -400,14 +424,14 @@ class TestHookToolsetCallTool: assert calls[0] != loop_thread -class TestHookToolsetBytes: - @staticmethod - def _call(method_name, tool_args): - ts = HookToolset(_FakeHook(), allowed_methods=[method_name]) - ctx = MagicMock(spec=RunContext) - tools = asyncio.run(ts.get_tools(ctx=ctx)) - return asyncio.run(ts.call_tool(method_name, tool_args, ctx=ctx, tool=tools[method_name])) +def _call_tool(method_name, tool_args): + ts = HookToolset(_FakeHook(), allowed_methods=[method_name]) + ctx = MagicMock(spec=RunContext) + tools = asyncio.run(ts.get_tools(ctx=ctx)) + return asyncio.run(ts.call_tool(method_name, tool_args, ctx=ctx, tool=tools[method_name])) + +class TestHookToolsetBytes: @pytest.mark.parametrize( ("method_name", "tool_args", "expected"), [ @@ -429,7 +453,7 @@ class TestHookToolsetBytes: ], ) def test_decodes_base64_for_bytes_params(self, method_name, tool_args, expected): - assert self._call(method_name, tool_args) == expected + assert _call_tool(method_name, tool_args) == expected @pytest.mark.parametrize( "value", @@ -442,64 +466,128 @@ class TestHookToolsetBytes: ) def test_an_undecodable_value_asks_the_model_to_retry(self, value): with pytest.raises(ModelRetry, match="'data' must be base64-encoded"): - self._call("upload_bytes", {"data": value, "key": "greeting.txt"}) + _call_tool("upload_bytes", {"data": value, "key": "greeting.txt"}) def test_binary_round_trips_byte_exact(self): payload = b"\x89PNG\r\n\x1a\n\xff\x00" - result = self._call("echo_bytes", {"data": base64.b64encode(payload).decode("ascii")}) + result = _call_tool("echo_bytes", {"data": base64.b64encode(payload).decode("ascii")}) assert base64.b64decode(json.loads(result)["base64"]) == payload def test_a_utf8_result_reaches_the_model_as_text(self): - result = self._call("echo_bytes", {"data": base64.b64encode("naïve,1\n".encode()).decode("ascii")}) + result = _call_tool("echo_bytes", {"data": base64.b64encode("naïve,1\n".encode()).decode("ascii")}) assert result == "naïve,1\n" - def test_a_bad_value_is_corrected_by_the_model_in_a_real_run(self): - hook = _FakeHook() - ts = HookToolset(hook, allowed_methods=["upload_bytes"]) - attempts = iter( - [{"data": "hello world", "key": "k"}, {"data": base64.b64encode(b"hello").decode(), "key": "k"}] - ) + def test_a_pinned_bytes_argument_is_passed_as_given(self): + ts = HookToolset(_FakeHook(), allowed_methods=["upload_bytes"], pinned_arguments={"data": b"raw"}) + ctx = MagicMock(spec=RunContext) + tools = asyncio.run(ts.get_tools(ctx=ctx)) + + result = asyncio.run(ts.call_tool("upload_bytes", {"key": "k"}, ctx=ctx, tool=tools["upload_bytes"])) + + assert result == "uploaded 3 bytes to k (type=bytes)" + + +class TestHookToolsetDatetime: + @pytest.mark.parametrize( + ("tool_args", "expected_since", "expected_until"), + [ + pytest.param( + {"since": "2026-09-01T00:00:00Z"}, + datetime(2026, 9, 1, tzinfo=UTC), + None, + id="utc-z", + ), + pytest.param( + {"since": "2026-09-01T05:30:00+05:30", "until": "2026-09-30T23:59:59+00:00"}, + datetime(2026, 9, 1, 5, 30, tzinfo=timezone(timedelta(hours=5, minutes=30))), + datetime(2026, 9, 30, 23, 59, 59, tzinfo=UTC), + id="offsets", + ), + pytest.param( + {"since": "2026-09-01T00:00:00Z", "until": None}, + datetime(2026, 9, 1, tzinfo=UTC), + None, + id="explicit-null", + ), + ], + ) + def test_parses_iso_8601_for_datetime_params(self, tool_args, expected_since, expected_until): + assert _call_tool("list_modified", tool_args) == f"since={expected_since!r} until={expected_until!r}" + + @pytest.mark.parametrize( + ("value", "reason"), + [ + pytest.param("2026-09-01T00:00:00", "has no UTC offset", id="naive"), + pytest.param("2026-09-01", "has no UTC offset", id="date-only"), + pytest.param("last tuesday", "Invalid isoformat string", id="not-iso-8601"), + ], + ) + def test_a_value_without_an_offset_or_unparsable_asks_the_model_to_retry(self, value, reason): + with pytest.raises( + ModelRetry, match=f"'since' must be an ISO 8601 date-time with a UTC offset: .*{reason}" + ): + _call_tool("list_modified", {"since": value}) + + +_ENCODED_ARG_CASES = [ + pytest.param( + "upload_bytes", + {"data": "hello world", "key": "k"}, + {"data": base64.b64encode(b"hello").decode(), "key": "k"}, + "'data' must be base64-encoded binary data", + "uploaded 5 bytes to k (type=bytes)", + id="bytes", + ), + pytest.param( + "list_modified", + {"since": "2026-09-01T00:00:00"}, + {"since": "2026-09-01T00:00:00Z"}, + "'since' must be an ISO 8601 date-time with a UTC offset", + f"since={datetime(2026, 9, 1, tzinfo=UTC)!r} until=None", + id="datetime", + ), +] + + +class TestHookToolsetEncodedArgs: + @pytest.mark.parametrize(("method_name", "bad", "good", "retry", "result"), _ENCODED_ARG_CASES) + def test_a_bad_value_is_corrected_by_the_model_in_a_real_run(self, method_name, bad, good, retry, result): + ts = HookToolset(_FakeHook(), allowed_methods=[method_name]) + attempts = iter([bad, good]) def model(messages, info): retried = [p for m in messages for p in m.parts if isinstance(p, RetryPromptPart)] returned = [p for m in messages for p in m.parts if isinstance(p, ToolReturnPart)] if returned: return ModelResponse(parts=[TextPart(f"{retried[0].content} | {returned[0].content}")]) - return ModelResponse(parts=[ToolCallPart("upload_bytes", next(attempts), tool_call_id="c")]) + return ModelResponse(parts=[ToolCallPart(method_name, next(attempts), tool_call_id="c")]) - answer = Agent(FunctionModel(model), toolsets=[ts]).run_sync("upload").output + answer = Agent(FunctionModel(model), toolsets=[ts]).run_sync("go").output - assert "'data' must be base64-encoded binary data" in answer - assert answer.endswith("uploaded 5 bytes to k (type=bytes)") + assert retry in answer + assert answer.endswith(result) - def test_an_approval_gate_is_never_asked_about_an_undecodable_value(self): + @pytest.mark.parametrize(("method_name", "bad", "good", "retry", "result"), _ENCODED_ARG_CASES) + def test_an_approval_gate_is_never_asked_about_an_undecodable_value( + self, method_name, bad, good, retry, result + ): """A task instance can pause for approval only once, so the request must be one that can run.""" - ts = HookToolset(_FakeHook(), allowed_methods=["upload_bytes"]) - valid = {"data": base64.b64encode(b"hello").decode(), "key": "k"} - attempts = iter([{"data": "hello world", "key": "k"}, valid]) + ts = HookToolset(_FakeHook(), allowed_methods=[method_name]) + attempts = iter([bad, good]) def model(messages, info): - return ModelResponse(parts=[ToolCallPart("upload_bytes", next(attempts), tool_call_id="c")]) + return ModelResponse(parts=[ToolCallPart(method_name, next(attempts), tool_call_id="c")]) agent = Agent( FunctionModel(model), toolsets=[ts.approval_required()], output_type=[str, DeferredToolRequests] ) - output = agent.run_sync("upload").output + output = agent.run_sync("go").output assert isinstance(output, DeferredToolRequests) - assert [call.args for call in output.approvals] == [valid] - - def test_a_pinned_bytes_argument_is_passed_as_given(self): - ts = HookToolset(_FakeHook(), allowed_methods=["upload_bytes"], pinned_arguments={"data": b"raw"}) - ctx = MagicMock(spec=RunContext) - tools = asyncio.run(ts.get_tools(ctx=ctx)) - - result = asyncio.run(ts.call_tool("upload_bytes", {"key": "k"}, ctx=ctx, tool=tools["upload_bytes"])) - - assert result == "uploaded 3 bytes to k (type=bytes)" + assert [call.args for call in output.approvals] == [good] def _takes_bytes(value: bytes): ... @@ -520,34 +608,52 @@ def _takes_str(value: str): ... def _takes_unannotated(value): ... -class TestBytesParamDetection: +def _takes_datetime(value: datetime): ... + + +def _takes_optional_datetime(value: datetime | None = None): ... + + +def _takes_datetime_or_str(value: datetime | str): ... + + +def _takes_list_of_datetimes(value: list[datetime]): ... + + +class TestStringEncodedParamDetection: @pytest.mark.parametrize( - ("func", "is_bytes_param"), + ("func", "note"), [ - 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"), + pytest.param(_takes_bytes, _BASE64_PARAM_NOTE, id="bytes"), + pytest.param(_takes_optional_bytes, _BASE64_PARAM_NOTE, id="optional-bytes"), + pytest.param(_takes_bytes_or_str, None, id="bytes-or-str"), + pytest.param(_takes_list_of_bytes, None, id="list-of-bytes"), + pytest.param(_takes_datetime, _DATETIME_PARAM_NOTE, id="datetime"), + pytest.param(_takes_optional_datetime, _DATETIME_PARAM_NOTE, id="optional-datetime"), + pytest.param(_takes_datetime_or_str, None, id="datetime-or-str"), + pytest.param(_takes_list_of_datetimes, None, id="list-of-datetimes"), + pytest.param(_takes_str, None, id="str"), + pytest.param(_takes_unannotated, None, id="unannotated"), ], ) - def test_only_bytes_and_optional_bytes_are_decoded(self, func, is_bytes_param): + def test_only_a_type_or_optional_type_is_decoded(self, func, note): hook = type("_Hook", (), {"method": staticmethod(func)})() ts = HookToolset(hook, allowed_methods=["method"]) tools = asyncio.run(ts.get_tools(ctx=MagicMock(spec=RunContext))) prop = tools["method"].tool_def.parameters_json_schema["properties"]["value"] + encoding = _introspect_signature(func)[1].get("value") - assert ("value" in _introspect_signature(func)[1]) is is_bytes_param - # A parameter described as base64 but never decoded would hand the hook the encoded text. - assert (_BASE64_PARAM_NOTE in prop.get("description", "")) is is_bytes_param + # A parameter described as encoded but never decoded would hand the hook the encoded text. + assert prop.get("description") == note + assert (encoding.note if encoding else None) == note + assert ('"date-time"' in json.dumps(prop)) is (note == _DATETIME_PARAM_NOTE) def test_var_args_are_never_decoded(self): def fn(*chunks: bytes, **extra: bytes): ... - schema, bytes_params = _introspect_signature(fn) + schema, encoded_params = _introspect_signature(fn) - assert bytes_params == frozenset() + assert encoded_params == {} assert schema["properties"] == {} def test_an_unresolvable_annotation_does_not_hide_the_other_params(self): @@ -555,9 +661,9 @@ class TestBytesParamDetection: # TYPE_CHECKING-only import. def fn(data: bytes, precision: Decimal | None = None): ... - schema, bytes_params = _introspect_signature(fn) + schema, encoded_params = _introspect_signature(fn) - assert bytes_params == {"data"} + assert encoded_params.keys() == {"data"} assert schema["properties"]["data"]["type"] == "string" def test_a_name_missing_from_an_imported_module_does_not_hide_the_other_params(self): @@ -566,9 +672,9 @@ class TestBytesParamDetection: # Raises AttributeError, not NameError, when evaluated: ``json`` is imported. fn.__annotations__["other"] = "json.NoSuchType | None" - schema, bytes_params = _introspect_signature(fn) + schema, encoded_params = _introspect_signature(fn) - assert bytes_params == {"data"} + assert encoded_params.keys() == {"data"} assert schema["properties"]["other"] == {}
