zozo123 commented on code in PR #70096:
URL: https://github.com/apache/airflow/pull/70096#discussion_r3682806932
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py:
##########
@@ -257,7 +254,7 @@ async def get_tools(self, ctx: RunContext[Any]) ->
dict[str, ToolsetTool[Any]]:
toolset=self,
tool_def=tool_def,
max_retries=1,
- args_validator=_PASSTHROUGH_VALIDATOR,
+ args_validator=build_args_validator(schema),
Review Comment:
Good catch — addressed in e53d7e9 (feed arg-validation `ValidationError`
through the bridge retry path) and refined in e3ebc028f6 to two-stage catching
so we don't over-retry tool-body errors (see sibling thread).
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/langchain_bridge.py:
##########
@@ -174,15 +182,15 @@ def _handle_retry(error: ModelRetry) -> str:
def _sync_call(**kwargs: Any) -> Any:
try:
result = _run_coro_sync(toolset.call_tool(name, _validate(kwargs),
ctx, toolset_tool))
- except ModelRetry as e:
+ except (ModelRetry, ValidationError) as e:
Review Comment:
Agreed — that was an over-broad catch. Fixed in e3ebc028f6:
`ValidationError` is now caught only around `_validate`, and `ModelRetry` only
around `call_tool`, matching pydantic-ai's two-stage `ToolManager` behaviour. A
tool-body `ValidationError` (Hook/MCP/custom) propagates instead of being fed
back, so a non-idempotent tool that already ran a side effect is not
re-invoked. Added a regression test for that path and tightened the bridge
docstring/comments accordingly.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/tool_definition.py:
##########
@@ -42,3 +43,70 @@ def return_schema_kwargs(schema: dict[str, Any]) ->
dict[str, Any]:
if _SUPPORTS_RETURN_SCHEMA:
return {"return_schema": schema}
return {}
+
+
+def _fragment_to_core_schema(fragment: dict[str, Any]) ->
core_schema.CoreSchema:
+ any_of = fragment.get("anyOf")
+ if isinstance(any_of, list):
+ choices: list[core_schema.CoreSchema | tuple[core_schema.CoreSchema,
str]] = [
+ _fragment_to_core_schema(choice) for choice in any_of if
isinstance(choice, dict)
+ ]
+ return core_schema.union_schema(choices) if choices else
core_schema.any_schema()
+
+ schema_type = fragment.get("type")
+ if isinstance(schema_type, list):
+ choices = [
+ _fragment_to_core_schema({**fragment, "type": item})
+ for item in schema_type
+ if isinstance(item, str)
+ ]
+ return core_schema.union_schema(choices) if choices else
core_schema.any_schema()
+
+ match schema_type:
+ case "string":
+ return core_schema.str_schema()
+ case "integer":
+ return core_schema.int_schema()
+ case "number":
+ return core_schema.float_schema()
+ case "boolean":
+ return core_schema.bool_schema()
+ case "null":
+ return core_schema.none_schema()
+ case "array":
+ items = fragment.get("items")
+ return core_schema.list_schema(
+ _fragment_to_core_schema(items) if isinstance(items, dict)
else None
+ )
+ case "object":
+ return _object_fragment_to_core_schema(fragment)
+ case _:
+ return core_schema.any_schema()
+
+
+def _object_fragment_to_core_schema(fragment: dict[str, Any]) ->
core_schema.CoreSchema:
+ """
+ Convert a JSON Schema ``object`` fragment to a core schema.
+
+ A fragment with no ``properties`` key is an untyped object (e.g. from a
+ ``dict[K, V]`` annotation): accept any dict rather than stripping its
+ contents. When ``properties`` is present, build a typed-dict that validates
+ each declared field recursively — nested objects are handled the same way
+ arrays already recurse into ``items``.
+ """
+ if "properties" not in fragment:
+ return core_schema.dict_schema()
+ required = set(fragment.get("required", []))
+ fields = {
+ name: core_schema.typed_dict_field(_fragment_to_core_schema(prop),
required=name in required)
+ for name, prop in fragment["properties"].items()
+ }
+ extra_behavior: Literal["allow", "ignore"] = (
Review Comment:
Worth aligning with native — `ignore` was the wrong default for this PR's
goal. Switched to `forbid` for fixed signatures in e3ebc028f6 (`allow` only
when the schema has `additionalProperties: true`, i.e. methods with
`**kwargs`). A mistyped field like `region` vs `region_name` now becomes a
bounded retry instead of a silent partial call. Updated the unit tests and PR
description to match.
--
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]
