airflow.providers.common.ai.operators.llm_branch

LLM-driven branching operator.

Classes

LLMBranchOperator

Ask an LLM to choose which downstream task(s) to execute.

Module Contents

class airflow.providers.common.ai.operators.llm_branch.LLMBranchOperator(*, allow_multiple_branches=False, fail_on_reject=False, ignore_downstream_trigger_rules=False, **kwargs)[source]

Bases: airflow.providers.common.ai.operators.llm.LLMOperator, airflow.providers.standard.operators.branch.BranchMixIn

Ask an LLM to choose which downstream task(s) to execute.

Downstream task IDs are discovered automatically from the DAG topology and presented to the LLM as a constrained enum via pydantic-ai structured output. No text parsing or manual validation is needed.

Parameters:

  • prompt – The prompt to send to the LLM.

  • llm_conn_id – Connection ID for the LLM provider.

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

  • fallback_conn_ids – 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 – System-level instructions for the LLM agent.

  • allow_multiple_branches (bool) – When False (default) the LLM returns a single task ID. When True the LLM may return one or more task IDs.

  • fail_on_reject (bool) – If True, a rejected review fails the task instead of skipping the downstream tasks. Generally discouraged, as for ApprovalOperator. Only takes effect with require_approval=True. Default False.

  • ignore_downstream_trigger_rules (bool) – If True, a rejected review skips every downstream task rather than only the direct ones, so a task whose trigger rule would still run it is skipped too. Only takes effect with require_approval=True. Default False.

  • agent_params – Additional keyword arguments passed to the pydantic-ai Agent constructor (e.g. retries, model_settings, tools).

usage_limits is inherited from LLMOperator.

Human-in-the-Loop approval parameters are inherited from LLMOperator (require_approval, approval_timeout, on_approval_timeout, allow_modifications, approval_notifiers, approval_assigned_users). The task pauses after the LLM chooses the branch(es) and only skips the unselected downstream tasks once a reviewer approves. Rejecting the review skips the direct downstream tasks except teardowns, matching ApprovalOperator; set fail_on_reject=True to fail the task instead, or ignore_downstream_trigger_rules=True to skip every downstream task rather than only the direct ones. The review form lists the valid downstream task IDs; with allow_modifications=True the editable choice is rendered as a dropdown of those IDs (single-branch mode) or a multi-select of them (allow_multiple_branches=True), and the reviewed branch(es) are validated against the downstream task IDs before branching.

inherits_from_skipmixin = True[source]

Used to determine if an Operator is inherited from SkipMixin or its subclasses (e.g., BranchMixin).

template_fields: collections.abc.Sequence[str] = ('prompt', 'llm_conn_id', 'model_id', 'fallback_conn_ids', 'system_prompt', 'agent_params',...[source]
allow_multiple_branches = False[source]
fail_on_reject = False[source]
ignore_downstream_trigger_rules = False[source]
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, validating the reviewed choice before branching.

Was this entry helpful?