airflow.providers.common.ai.operators.agent

Operator for running pydantic-ai agents with tools and multi-turn reasoning.

Classes

AgentOperator

Run a pydantic-ai Agent with tools and multi-turn reasoning.

Module Contents

class airflow.providers.common.ai.operators.agent.AgentOperator(*, prompt, llm_conn_id, model_id=None, system_prompt='', output_type=str, toolsets=None, enable_tool_logging=True, agent_params=None, **kwargs)[source]

Bases: airflow.providers.common.compat.sdk.BaseOperator

Run a pydantic-ai Agent with tools and multi-turn reasoning.

Provide llm_conn_id and optional toolsets to let the operator build and run the agent. The agent reasons about the prompt, calls tools in a multi-turn loop, and returns a final answer.

Parameters:

  • prompt (str) – The prompt to send to the agent.

  • llm_conn_id (str) – Connection ID for the LLM provider.

  • model_id (str | None) – Model identifier (e.g. "openai:gpt-5"). Overrides the model stored in the connection’s extra field.

  • system_prompt (str) – System-level instructions for the agent.

  • output_type (type) – Expected output type. Default str. Set to a Pydantic BaseModel subclass for structured output.

  • toolsets (list[pydantic_ai.toolsets.abstract.AbstractToolset] | None) – List of pydantic-ai toolsets the agent can use (e.g. SQLToolset, HookToolset).

  • enable_tool_logging (bool) – When True (default), wraps each toolset in a LoggingToolset that logs tool calls with timing at INFO level and arguments at DEBUG level. Set to False to disable.

  • agent_params (dict[str, Any] | None) – Additional keyword arguments passed to the pydantic-ai Agent constructor (e.g. retries, model_settings).

template_fields: collections.abc.Sequence[str] = ('prompt', 'llm_conn_id', 'model_id', 'system_prompt', 'agent_params')[source]
prompt[source]
llm_conn_id[source]
model_id = None[source]
system_prompt = ''[source]
output_type[source]
toolsets = None[source]
enable_tool_logging = True[source]
agent_params[source]
property llm_hook: airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook[source]

Return PydanticAIHook for the configured LLM connection.

execute(context)[source]

Derive when creating an operator.

The main method to execute the task. Context is the same dictionary used as when rendering jinja templates.

Refer to get_template_context for more context.

Was this entry helpful?