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(*, branches=None, 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.

  • branches (collections.abc.Mapping[str, airflow.providers.common.ai.policies.decision.BranchOption | str] | None) – Optional mapping of downstream task ID to a BranchOption, or to a string as shorthand for its description. The description travels in the output schema next to the option, so the model reads “here is an option, here is what it means” rather than guessing from the task ID; min_confidence on an option is a bar for that branch alone. A downstream task without an entry is presented by its ID alone and takes the policy’s bar. A key that is not a downstream task ID fails the task before the model is called. Descriptions support Jinja templating.

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

  • decision_policy – A DecisionPolicy: the confidence a pick needs for the operator to branch on it without a person (min_confidence) and what happens under it (on_uncertain: "review" or "fail"). Confidence comes from models that report one, such as a classifier model (TypeSafe’s); a text model reports none, which counts as uncertain, so swapping the connection does not silently switch off a control the author set. A branch’s own min_confidence overrides the policy’s for that pick; with allow_multiple_branches the strictest bar among the picked branches applies. require_approval=True still sends every pick to a person regardless. Default None: no gate.

  • 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 when a review is opened. 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 when a review is opened. 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]
branches = None[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, decision=None)[source]

Resume after human review, validating the reviewed choice before branching.

Was this entry helpful?