airflow.providers.common.ai.utils.tool_definition¶

Version-tolerant helpers for building pydantic-ai ToolDefinition objects.

Functions¶

`return_schema_kwargs`(schema)	Return `{"return_schema": schema}` when pydantic-ai supports the field, else `{}`.
`code_arg_kwargs`(arg_name, language)	Mark a tool as a code-execution surface whose `arg_name` argument holds code.
`build_args_validator`(parameters_json_schema)	Build an argument validator from the schema advertised to the model.

Module Contents¶

airflow.providers.common.ai.utils.tool_definition.return_schema_kwargs(schema)[source]¶

Return {"return_schema": schema} when pydantic-ai supports the field, else {}.

return_schema lets CodeMode (the Monty sandbox) render a typed function signature for a tool (-> str) instead of -> Any, which helps the model write correct code. It has no effect outside code mode.

Parameters:: schema (dict[str, Any]) – A JSON Schema fragment describing the tool’s return value.

airflow.providers.common.ai.utils.tool_definition.code_arg_kwargs(arg_name, language)[source]¶

Mark a tool as a code-execution surface whose arg_name argument holds code.

Two consumers read this metadata. Instrumentation renders the argument as code in traces, tagged with language. More importantly, CodeMode uses it to decide which tools stay native: a tool that itself executes code must not be folded into run_code, because the model would then have to write a script whose argument is a second script quoted as a string. Without the marker, a sandbox tool is folded in and the generated orchestration runs in-process on the worker – the opposite of what a sandbox tool is for.

Returns {} when the installed pydantic-ai predates ToolDefinition.metadata, matching return_schema_kwargs().

Parameters:

arg_name (str) – The tool argument that carries the code or command.
language (str) – Language tag for the argument, e.g. "shell".

airflow.providers.common.ai.utils.tool_definition.build_args_validator(parameters_json_schema)[source]¶

Build an argument validator from the schema advertised to the model.