This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 4e8ab38306f Fix HookToolset calls to hook methods that take a datetime
(#74759)
4e8ab38306f is described below
commit 4e8ab38306f2bec5262e2da465224d4635deb287
Author: Kaxil Naik <[email protected]>
AuthorDate: Thu Oct 15 09:24:15 2026 +0100
Fix HookToolset calls to hook methods that take a datetime (#74759)
---
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"] == {}
