kaxil commented on code in PR #70433:
URL: https://github.com/apache/airflow/pull/70433#discussion_r3706635804
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py:
##########
@@ -136,6 +136,12 @@ 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)
+ bytes_params = _bytes_param_names(method)
Review Comment:
+1 on the base64 contract agreed above. One constraint for the rework: the
decode has to stay here in `call_tool` rather than move into
`build_args_validator`, even though #70096 makes the validator look like the
natural home. Validated args flow through the whole toolset chain, and
`CachingToolset` fingerprints `tool_args` before delegating (`_digest` is
`json.dumps` with no `default=`), so `bytes` values there would turn every
bytes-param call into an unverifiable fingerprint under `durable=True`.
Building a new dict inside the innermost toolset, as this PR already does,
keeps bytes out of that path, so worth a short code comment to make the
placement survive the rework. Since `contentEncoding` is ignored by most
function-calling APIs, the base64 instruction should also be appended to each
parameter's schema description after the docstring enrichment loop in
`get_tools` (that loop overwrites descriptions, and e.g. `S3Hook.load_bytes`'s
":param bytes_data: bytes to set as content
for the key" would otherwise be all the model sees), plus a sentence in
`toolsets.rst` documenting the contract.
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py:
##########
@@ -173,6 +179,32 @@ 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 ``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
+
+
+def _bytes_param_names(method: Callable[..., Any]) -> frozenset[str]:
+ """Names of ``method``'s parameters whose resolved type is
``bytes``-like."""
+ sig = inspect.signature(method)
+ try:
+ hints = get_type_hints(method)
+ except Exception:
Review Comment:
Narrowing this to `(NameError, TypeError)` as agreed above documents the
failure mode but doesn't close it. `get_type_hints` is all-or-nothing: one
unresolvable annotation anywhere in the signature (co-params and return type
included) throws away the hints for every parameter, and the PEP 563 fallback
then checks the string `"bytes"`, so coercion silently turns off for the whole
method. `CloudKMSHook.encrypt` hits this today: `plaintext: bytes` and
`authenticated_data: bytes | None` sit next to `retry: Retry | _MethodDefault`
where `Retry` is imported under `TYPE_CHECKING`, so its bytes params would
never be coerced. Resolving annotations per parameter closes it.
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py:
##########
@@ -136,6 +136,12 @@ 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)
+ bytes_params = _bytes_param_names(method)
+ if bytes_params:
+ tool_args = {
+ key: value.encode("utf-8") if key in bytes_params and
isinstance(value, str) else value
Review Comment:
On the plan to re-raise decode failures as `ValueError` so the model can
retry: a plain `ValueError` raised inside `call_tool` never reaches the model.
pydantic-ai only turns `ModelRetry` (and `ValidationError` from the
args-validation stage) into retry prompts, anything else fails the run, and
`langchain_bridge.py` documents the same two-stage split for the LangChain
path. The `b64decode` failure needs to be raised as `ModelRetry`. Same
treatment for this encode: `validate_json` rejects lone surrogates upstream,
but the LangChain bridge validates an already-parsed dict via
`validate_python`, which lets them through, and `.encode("utf-8")` then raises
`UnicodeEncodeError`.
##########
providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py:
##########
@@ -203,6 +212,50 @@ def test_dispatches_with_prefix(self):
)
assert result == "contents of test.txt"
+ def test_coerces_str_to_bytes_for_bytes_param(self):
+ hook = _FakeHook()
+ ts = HookToolset(hook, allowed_methods=["upload_bytes"])
+ tools = asyncio.run(ts.get_tools(ctx=MagicMock()))
+
+ result = asyncio.run(
+ ts.call_tool(
+ "upload_bytes",
+ {"data": "hello world", "key": "greeting.txt"},
+ ctx=MagicMock(),
+ tool=tools["upload_bytes"],
+ )
+ )
+ assert result == "uploaded 11 bytes to greeting.txt (type=bytes)"
+
+ def test_coerces_str_to_bytes_for_optional_bytes_param(self):
+ hook = _FakeHook()
+ ts = HookToolset(hook, allowed_methods=["upload_optional_bytes"])
+ tools = asyncio.run(ts.get_tools(ctx=MagicMock()))
+
+ result = asyncio.run(
+ ts.call_tool(
+ "upload_optional_bytes",
+ {"data": "hi"},
+ ctx=MagicMock(),
+ tool=tools["upload_optional_bytes"],
+ )
+ )
+ assert result == "data type=bytes"
+
+
+class TestBytesParamNames:
Review Comment:
These three cases re-test what the two dispatch tests above already cover,
while the one subtle branch in `_resolves_to_bytes`, the `len(non_none) == 1`
guard that keeps `bytes | str` unions uncoerced, is untested. When the tests
get reworked for base64, parametrizing this class over `(bytes, bytes | None,
bytes | str, list[bytes], unannotated)` would cover that guard and drop the
duplication.
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py:
##########
@@ -173,6 +179,32 @@ 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 ``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
+
+
+def _bytes_param_names(method: Callable[..., Any]) -> frozenset[str]:
+ """Names of ``method``'s parameters whose resolved type is
``bytes``-like."""
+ sig = inspect.signature(method)
+ try:
+ hints = get_type_hints(method)
+ except Exception:
+ hints = {}
+ return frozenset(
+ name
+ for name, param in sig.parameters.items()
+ if name not in ("self", "cls") and _resolves_to_bytes(hints.get(name,
param.annotation))
Review Comment:
This pass doesn't skip `VAR_POSITIONAL`/`VAR_KEYWORD` the way
`_build_json_schema_from_signature` does, so `def m(self, *chunks: bytes,
**extra: bytes)` reports `{"chunks", "extra"}`, and a literal `"chunks"` key in
`tool_args` would get encoded while real `**kwargs` values stay `str`. Latent
(no in-tree hook exposes bytes varargs), but it shows the two introspection
passes can drift. When this moves to build time per the thread, consider having
`_build_json_schema_from_signature` return the bytes param names alongside the
schema so both come from one pass, and avoid an `lru_cache` keyed on bound
methods, which would pin hook instances for the life of the process.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
