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, model_settings=None, output_retries=None, usage_limits=None, require_approval=False, approval_timeout=None, on_approval_timeout='fail', allow_modifications=False, approval_notifiers=None, approval_assigned_users=None, decision_policy=None, serialize_output=False, **kwargs)[source]

Bases: airflow.providers.common.ai.mixins.cancellable_run.CancellableAgentRunMixin, airflow.providers.common.compat.sdk.BaseOperator, airflow.providers.common.ai.mixins.approval.LLMApprovalMixin, airflow.providers.common.ai.mixins.usage_budget.UsageBudgetMixin

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.

No tool runs in the task: the operator accepts no function tools, toolsets or capabilities, makes one model call (repeated only when the output fails validation) and returns its output. For an agent that calls tools, use AgentOperator.

Once _clear_usage_budget has run, nothing may fail the attempt (raising DownstreamTasksSkipped to signal success is fine): a failure would be retried with a fresh budget. This applies to subclasses that override execute or execute_complete.

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 | types.UnionType) – Expected output type. Default str. A single type (str, int, list[str], a Pydantic BaseModel subclass) or a union (A | B). For a BaseModel subclass, 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.

  • model_settings (pydantic_ai.settings.ModelSettings | None) – Pydantic AI ModelSettings for the model call, such as temperature or max_tokens, including provider-specific settings. None (default) uses the model’s defaults.

  • output_retries (int | None) – How many times the model is asked again when its output fails validation against output_type. None (default) uses Pydantic AI’s default.

  • 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) sets no token, cost, or tool-call limits, but pydantic-ai still caps the run at its default request_limit of 50 requests.

    A dict that omits request_limit gets the same default of 50 requests – pass "request_limit": None explicitly for no request cap. This matches building a UsageLimits directly.

    On Airflow >= 3.3, a set usage_limits bounds the whole task instance: usage from earlier attempts counts against it. See Single prompts: LLMOperator and @task.llm for the full set of caveats.

  • require_approval (bool) – If True, the task pauses after generating output (awaiting_input on Airflow 3.3+, deferred on 3.1 and 3.2) 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 a review path (require_approval=True or a decision_policy that reviews) and a positive approval_timeout.

  • allow_modifications (bool) – If True, the reviewer can edit the output before approving. The edit is converted back into output_type and returned as the task result. For a type pydantic can validate, an edit that does not validate fails the task with ReviewedOutputValidationError; a retry runs the model again. 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 when a review is opened. A retry re-notifies with the regenerated output. On Airflow 3.4+ the reopened review shows it too; on Airflow 3.1 to 3.3 the review keeps the original subject, body and editable output, so approving it unedited returns the first attempt’s output. 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. On Airflow 3.1 to 3.3 the list is fixed when the review is first created; on 3.4+ a cleared or retried task replaces it. Needs Airflow 3.1+.

  • decision_policy (airflow.providers.common.ai.policies.decision.DecisionPolicy | None) – Experimental. A DecisionPolicy saying how confident the model has to be for the operator to return its answer by itself (min_confidence) and what happens otherwise (on_uncertain: "review" or "fail"). Confidence comes from models that report one per output field, such as a decision model, in provider_details. A structured output is judged by its least confident field among the fields that reported one; a field whose type reports none (a bounded float, where the probability is the answer) is not gated, and the record’s confidence map shows which fields were compared. When no field reports any confidence, as with a text model, the output counts as uncertain, so swapping the connection does not silently switch off a control the author set. Independent of require_approval, which always asks. on_uncertain="review" needs Airflow 3.1+, like require_approval. Default None: no gate, behaviour unchanged.

  • 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]
supports_decision_policy: ClassVar[bool] = True[source]
template_fields: collections.abc.Sequence[str] = ('prompt', 'llm_conn_id', 'model_id', 'fallback_conn_ids', 'system_prompt', 'model_settings',...[source]
decision_policy[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]
model_settings = None[source]
output_retries = None[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, decision=None)[source]

Resume after human review and convert the output back into output_type.

Raises:

ReviewedOutputValidationError – If a reviewer’s edit does not validate against output_type.

Was this entry helpful?