kaxil commented on code in PR #70096:
URL: https://github.com/apache/airflow/pull/70096#discussion_r3682576543
##########
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:
The widened `except (ModelRetry, ValidationError)` spans both
`_validate(kwargs)` and `toolset.call_tool(...)`, so a `ValidationError` raised
inside the tool body (a `HookToolset` method, an MCP client validating a server
response, a custom toolset) now gets fed back as a retry too. pydantic-ai 2.13
two-stages this: arg-validation retries a `ValidationError`, but `_raw_execute`
(`tool_manager.py:758`) catches only `ModelRetry` around `call_tool`, so a
tool-body `ValidationError` propagates rather than retrying. For a
non-idempotent tool that already ran a side effect before raising, the retry
re-invokes it. SQL/DataFusion aren't affected (SQL wraps everything in
`ModelRetry`, DataFusion doesn't raise `ValidationError`); this is the
Hook/MCP/custom path the docstring covers. Catching `ValidationError` only
around `_validate` and `ModelRetry` around `call_tool` keeps the mirror-native
behaviour. Non-blocking.
##########
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:
`extra_behavior="ignore"` silently drops undeclared args. Native pydantic-ai
uses `forbid` for fixed signatures (`_function_schema.py:358`, `allow` only
when there's a `**kwargs`), so a model that sends `region` when the field is
`region_name` gets a bounded retry to fix itself instead of a silent partial
call. Since this PR is about turning malformed calls into retries, was `ignore`
deliberate to avoid spurious retries when a model over-supplies fields, or
worth aligning with native's `forbid` for fixed signatures?
--
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]
