• Home
• Common AI Operators
• LLMBranchOperator

↑↓ Navigate ⏎ Select Esc Close

LLMBranchOperator

Use LLMBranchOperator for LLM-driven branching — where the LLM decides which downstream task(s) to execute.

The operator discovers downstream tasks automatically from the Dag topology and presents them to the LLM as a constrained enum via pydantic-ai structured output. No text parsing or manual validation is needed.

See also

Connection configuration

Basic Usage

Connect the operator to downstream tasks. The LLM chooses which branch to execute based on the prompt:

airflow/providers/common/ai/example_dags/example_llm_branch.py

@dag(tags=["example"])
def example_llm_branch_operator():
    route = LLMBranchOperator(
        task_id="route_ticket",
        prompt="User says: 'My password reset email never arrived.'",
        llm_conn_id="pydanticai_default",
        system_prompt="Route support tickets to the right team.",
    )

    @task
    def handle_billing():
        return "Handling billing issue"

    @task
    def handle_auth():
        return "Handling auth issue"

    @task
    def handle_general():
        return "Handling general issue"

    route >> [handle_billing(), handle_auth(), handle_general()]

Multiple Branches

Set allow_multiple_branches=True to let the LLM select more than one downstream task. All selected branches run; unselected branches are skipped:

airflow/providers/common/ai/example_dags/example_llm_branch.py

@dag(tags=["example"])
def example_llm_branch_multi():
    route = LLMBranchOperator(
        task_id="classify",
        prompt="This product is great but shipping was slow and the box was damaged.",
        llm_conn_id="pydanticai_default",
        system_prompt="Select all applicable categories for this customer review.",
        allow_multiple_branches=True,
    )

    @task
    def handle_positive():
        return "Processing positive feedback"

    @task
    def handle_shipping():
        return "Escalating shipping issue"

    @task
    def handle_packaging():
        return "Escalating packaging issue"

    route >> [handle_positive(), handle_shipping(), handle_packaging()]

TaskFlow Decorator

The @task.llm_branch decorator wraps LLMBranchOperator. The function returns the prompt string; all other parameters are passed to the operator:

airflow/providers/common/ai/example_dags/example_llm_branch.py

@dag(tags=["example"])
def example_llm_branch_decorator():
    @task.llm_branch(
        llm_conn_id="pydanticai_default",
        system_prompt="Route support tickets to the right team.",
    )
    def route_ticket(message: str):
        return f"Route this support ticket: {message}"

    @task
    def handle_billing():
        return "Handling billing issue"

    @task
    def handle_auth():
        return "Handling auth issue"

    @task
    def handle_general():
        return "Handling general issue"

    route_ticket("I was charged twice for my subscription.") >> [
        handle_billing(),
        handle_auth(),
        handle_general(),
    ]

The callable may also return a non-empty Sequence[UserContent] for multimodal inputs – see @task.agent multimodal prompts.

With multiple branches:

airflow/providers/common/ai/example_dags/example_llm_branch.py

@dag(tags=["example"])
def example_llm_branch_decorator_multi():
    @task.llm_branch(
        llm_conn_id="pydanticai_default",
        system_prompt="Select all applicable categories for this customer review.",
        allow_multiple_branches=True,
    )
    def classify_review(review: str):
        return f"Classify this review: {review}"

    @task
    def handle_positive():
        return "Processing positive feedback"

    @task
    def handle_shipping():
        return "Escalating shipping issue"

    @task
    def handle_packaging():
        return "Escalating packaging issue"

    classify_review("Great product but shipping was slow.") >> [
        handle_positive(),
        handle_shipping(),
        handle_packaging(),
    ]

Human-in-the-Loop Approval

Set require_approval=True to pause the task after the LLM chooses the branch(es) and wait for a human reviewer to approve the choice before any downstream task is skipped. The review form shows the LLM’s choice and the valid downstream task IDs. When allow_modifications=True, the reviewer can also change the choice — rendered as a dropdown of the downstream task IDs, or a multi-select of them with allow_multiple_branches=True. The reviewed branch(es) are validated against the downstream task IDs before branching:

airflow/providers/common/ai/example_dags/example_llm_branch.py

@dag(tags=["example"])
def example_llm_branch_approval():
    route = LLMBranchOperator(
        task_id="route_with_approval",
        prompt="User says: 'I was charged twice for my subscription.'",
        llm_conn_id="pydanticai_default",
        system_prompt="Route support tickets to the right team.",
        require_approval=True,
        approval_timeout=timedelta(hours=24),
        allow_modifications=True,
    )

    @task
    def handle_billing():
        return "Handling billing issue"

    @task
    def handle_auth():
        return "Handling auth issue"

    @task
    def handle_general():
        return "Handling general issue"

    route >> [handle_billing(), handle_auth(), handle_general()]

Rejecting the review, or letting approval_timeout expire, fails the task (HITLRejectException / HITLTimeoutError), so downstream tasks end up upstream_failed rather than skipped.

require_approval=True requires a string prompt: a decorated callable returning a Sequence[UserContent] raises TypeError before the LLM call.

approval_timeout and the rest of the approval behaviour are inherited from LLMOperator.

How It Works

At execution time, the operator:

1. Reads self.downstream_task_ids from the Dag topology.

2. Creates a dynamic Enum with one member per downstream task ID.

3. Passes that enum as output_type to pydantic-ai, constraining the LLM to valid task IDs only.

4. Converts the LLM’s structured output to task ID string(s) and calls do_branch() to skip non-selected downstream tasks.

Parameters

• prompt: The prompt to send to the LLM (operator) or the return value of the decorated function (decorator).

• llm_conn_id: Airflow connection ID for the LLM provider.

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

• system_prompt: System-level instructions for the agent. Supports Jinja templating.

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

• agent_params: Additional keyword arguments passed to the pydantic-ai Agent constructor (e.g. retries, model_settings). Supports Jinja templating.

• require_approval: If True, the task pauses after the LLM chooses the branch(es) and waits for human review before branching. Default False.

• approval_timeout: Maximum time to wait for a review (timedelta). None means wait indefinitely. Default None.

• allow_modifications: If True, the reviewer can change the chosen branch(es) before approving. Default False.

Logging

After each LLM call, the operator logs a summary with model name, token usage, and request count at INFO level. See AgentOperator — Logging for details on the log format.

Previous Next

Was this entry helpful?