 .. Licensed to the Apache Software Foundation (ASF) under one
    or more contributor license agreements.  See the NOTICE file
    distributed with this work for additional information
    regarding copyright ownership.  The ASF licenses this file
    to you under the Apache License, Version 2.0 (the
    "License"); you may not use this file except in compliance
    with the License.  You may obtain a copy of the License at

 ..   http://www.apache.org/licenses/LICENSE-2.0

 .. Unless required by applicable law or agreed to in writing,
    software distributed under the License is distributed on an
    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    KIND, either express or implied.  See the License for the
    specific language governing permissions and limitations
    under the License.

.. _howto/operator:llm_branch:

``LLMBranchOperator``
=====================

Use :class:`~airflow.providers.common.ai.operators.llm_branch.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.

.. seealso::
    :ref:`Connection configuration <howto/connection:pydanticai>`

Basic Usage
-----------

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

.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_branch.py
    :language: python
    :start-after: [START howto_operator_llm_branch_basic]
    :end-before: [END howto_operator_llm_branch_basic]

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:

.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_branch.py
    :language: python
    :start-after: [START howto_operator_llm_branch_multi]
    :end-before: [END howto_operator_llm_branch_multi]

TaskFlow Decorator
------------------

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

.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_branch.py
    :language: python
    :start-after: [START howto_decorator_llm_branch]
    :end-before: [END howto_decorator_llm_branch]

With multiple branches:

.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_branch.py
    :language: python
    :start-after: [START howto_decorator_llm_branch_multi]
    :end-before: [END howto_decorator_llm_branch_multi]

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.

Logging
-------

After each LLM call, the operator logs a summary with model name, token usage,
and request count at INFO level. See :ref:`AgentOperator — Logging <howto/operator:agent>`
for details on the log format.
