airflow.providers.common.ai.operators.llm

Operator for general-purpose LLM calls.

Classes

LLMOperator

Call an LLM with a prompt and return the output.

Module Contents

class airflow.providers.common.ai.operators.llm.LLMOperator(*, prompt, llm_conn_id, model_id=None, fallback_conn_ids=None, system_prompt='', output_type=str, agent_params=None, usage_limits=None, require_approval=False, approval_timeout=None, on_approval_timeout='fail', allow_modifications=False, approval_notifiers=None, approval_assigned_users=None, serialize_output=False, **kwargs)[source]

Bases: airflow.providers.common.compat.sdk.BaseOperator, airflow.providers.common.ai.mixins.approval.LLMApprovalMixin

Call an LLM with a prompt and return the output.

Uses a PydanticAIHook for LLM access. Supports plain string output (default) and structured output via a Pydantic BaseModel. When output_type is a BaseModel subclass, the model instance is returned to XCom unchanged so downstream tasks can type-hint it directly (e.g. def downstream(result: MyModel) -> None). The class is auto-registered for deserialization in each process that parses the DAG, so no edit to [core] allowed_deserialization_classes is required. The Pydantic class must be defined at module scope: classes nested inside a function or @dag-decorated body cannot be deserialized from XCom.

Parameters:

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

  • 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.

  • fallback_conn_ids (list[str] | None) – Connection IDs to fail over to, in order, when the primary provider is unavailable. Overrides the fallback_conn_ids set in the connection’s extra field. None (default) reads the connection’s own extra field; an explicit [] disables a chain configured there. See PydanticAIHook for how blank entries in the list are dropped.

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

  • output_type (type) – Expected output type. Default str. Set to a Pydantic BaseModel subclass for structured output; the model instance is returned to XCom unchanged so downstream tasks can type-hint it directly. The class must be defined at module scope – nested classes cannot be deserialized from XCom.

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

  • usage_limits (pydantic_ai.usage.UsageLimits | dict[str, Any] | None) –

    Optional pydantic-ai UsageLimits enforced on the run, or a dict of the same fields (e.g. {"cost_limit": "{{ params.budget }}", "request_limit": 5}). The dict form is templated: each value is rendered by Jinja like any other template_fields entry, then coerced to that field’s type (Decimal, int, or bool). A value that cannot be coerced – a Variable that exists but is empty renders to "", a typo renders to a non-numeric string – fails the task with a ValueError naming the field and the rendered value, instead of silently disabling the limit. A UsageLimits instance passed directly is used as-is and is not templated or validated. None (default) means no enforcement.

    A dict that omits request_limit still gets pydantic-ai’s default of 50 requests – pass "request_limit": None explicitly for no request cap. This matches building a UsageLimits directly, but it is easy to miss when moving from usage_limits=None to a dict that only sets cost_limit. See LLMOperator for the full set of caveats.

  • require_approval (bool) – If True, the task defers after generating output and waits for a human reviewer to approve or reject via the HITL interface. Default False. Needs Airflow 3.1+.

  • approval_timeout (datetime.timedelta | None) – Maximum time to wait for a review. When exceeded, on_approval_timeout decides the outcome.

  • on_approval_timeout (Literal['fail', 'approve', 'reject']) – What to do when approval_timeout expires without a review. "fail" (default) fails the task with HITLTimeoutError; "approve" and "reject" answer the review with that option, so the task resumes as if a reviewer had chosen it. The chosen option is also pre-highlighted for the reviewer in the HITL form. Requires require_approval=True and a positive approval_timeout.

  • allow_modifications (bool) – If True, the reviewer can edit the output before approving. The modified value is returned as the task result. Default False.

  • approval_notifiers (airflow.providers.common.compat.notifier.BaseNotifier | collections.abc.Iterable[airflow.providers.common.compat.notifier.BaseNotifier] | None) – Notifiers called once the review is open, so a reviewer is told about it. Only takes effect with require_approval=True. A retry re-notifies with the regenerated output while the open review keeps the original subject and body. Default None.

  • approval_assigned_users (airflow.sdk.execution_time.hitl.HITLUser | collections.abc.Iterable[airflow.sdk.execution_time.hitl.HITLUser] | None) – Users allowed to answer the review, as {"id": ..., "name": ...} dicts where id is the auth manager’s user id. None (default) lets any user with the permission respond. The list is fixed when the review is first created. Needs Airflow 3.1+.

  • serialize_output (bool) – If True and output_type is a Pydantic BaseModel subclass, the model instance is dumped to a dict via model_dump() before being pushed to XCom. Default False – the Pydantic instance flows through XCom unchanged. Set to True when a downstream consumer needs the dict shape (e.g. sending to an external system that expects JSON-style payloads).

deserialization_allowed_class_fields: ClassVar[tuple[str, ...]] = ('output_type',)[source]
template_fields: collections.abc.Sequence[str] = ('prompt', 'llm_conn_id', 'model_id', 'fallback_conn_ids', 'system_prompt', 'agent_params',...[source]
prompt[source]
llm_conn_id[source]
model_id = None[source]
fallback_conn_ids = None[source]
system_prompt = ''[source]
output_type[source]
serialize_output = False[source]
agent_params[source]
usage_limits = None[source]
require_approval = False[source]
approval_timeout = None[source]
on_approval_timeout = 'fail'[source]
allow_modifications = False[source]
approval_notifiers[source]
approval_assigned_users: list[airflow.sdk.execution_time.hitl.HITLUser][source]
property llm_hook: airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook[source]

Return the correct PydanticAIHook subclass for the configured connection.

Delegates to get_hook() which looks up the connection’s conn_type and instantiates the matching subclass (e.g. PydanticAIAzureHook for pydanticai_azure connections).

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.

execute_complete(context, generated_output, event)[source]

Resume after human review and restore the Pydantic model for XCom consumers.

Was this entry helpful?