Source code for airflow.providers.common.ai.example_dags.example_decision_model
# 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.
"""
Example DAG using a decision model to route an incident.
A decision model answers typed questions and cannot write text, so it suits a branch
(pick one of these task ids) and a classification (pick one of these labels), and nothing
that produces prose. See the "Decision models" guide in this provider's documentation.
Prerequisites: ``pydantic-ai-slim`` 2.46.0 or later (2.53.0 for ``system-one:``), and a
connection ``decision_default`` with ``conn_type='pydanticai'`` whose model is a decision model,
either:
- TypeSafe Jev: ``password=<API key>``, ``extra='{"model": "typesafe:jev-1.13.0"}'``, and
``pip install 'apache-airflow-providers-common-ai[typesafe]'``, or
- a server answering the System One API: ``host=<server URL>``, ``password=<API key, if any>``,
``extra='{"model": "system-one:<model name>"}'``
"""
from __future__ import annotations
from enum import Enum
from pydantic_ai import UseEnumMemberDocstrings
from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
from airflow.providers.common.ai.operators.llm_branch import LLMBranchOperator
from airflow.providers.common.compat.sdk import dag, task
# Clear-cut: the error names the missing permission, so one remediation fits and the others
# do not. Branch options a decision model can separate are the ones to give it.
[docs]
PERMISSION_FAILURE = (
"Task export_report failed: botocore.exceptions.ClientError: An error occurred "
"(AccessDenied) when calling the PutObject operation: the task role "
"airflow-worker lacks s3:PutObject on arn:aws:s3:::reports-prod/daily/."
)
# Answerable but not decisively: a decision model lands on ``resource`` here with roughly
# two thirds of the probability, which is enough to be worth recording and not enough to act
# on unattended. That is the case the confidence gate below exists for.
[docs]
UNDERDETERMINED_INCIDENT = (
"Task load_orders failed: psycopg2.OperationalError: connection refused to "
"warehouse.internal:5432. Two other Dags writing to the same warehouse failed in the "
"last ten minutes. A Terraform apply touching the database security group merged "
"four hours ago."
)
# Acting automatically deserves a higher bar than flagging for review, so the two use
# different thresholds rather than one shared number.
# [START howto_decision_model_branch]
@dag(tags=["example", "decision"])
[docs]
def example_decision_model_branch():
"""Route the incident. The connection's model is the only thing that makes this a decision model."""
route = LLMBranchOperator(
task_id="route_failure",
prompt=PERMISSION_FAILURE,
llm_conn_id="decision_default",
system_prompt="Pick the remediation that addresses the cause, not the symptom.",
# Each description is what the model weighs that branch by. Some System One servers
# refuse an option without one, so describe every branch.
branches={
"grant_bucket_write": "The task's role lacks a permission it needs on the bucket.",
"restore_deleted_bucket": "The destination bucket or prefix no longer exists.",
"wait_and_retry": "A transient fault: throttling, a timeout, a brief outage.",
},
)
@task
def grant_bucket_write():
return "Adding s3:PutObject for the task role"
@task
def restore_deleted_bucket():
return "Recreating the destination bucket"
@task
def wait_and_retry():
return "Treating this as transient"
route >> [grant_bucket_write(), restore_deleted_bucket(), wait_and_retry()]
# [END howto_decision_model_branch]
example_decision_model_branch()
# [START howto_decision_model_confidence]
[docs]
class FailureCause(UseEnumMemberDocstrings, str, Enum):
"""Why an Airflow task failed."""
# Each docstring is what the model weighs that option by. Some System One servers
# refuse an option without one, so describe every option.
[docs]
transient = "transient"
"""A fault that clears by itself: a timeout, throttling, a dropped connection."""
"""A dependency is down or unreachable and needs fixing before a retry can work."""
[docs]
permanent = "permanent"
"""A bug or bad input that fails the same way however often it is retried."""
@dag(tags=["example", "decision"])
[docs]
def example_decision_model_confidence():
"""Classify a failure and escalate when the model says it does not know.
``LLMBranchOperator`` can gate its pick on one confidence bar through ``decision_policy``.
Reading the confidence in a task of your own, as here, lets you act differently in more
than one band.
"""
@task
def classify(log_line: str) -> dict:
agent = PydanticAIHook(llm_conn_id="decision_default").create_agent(
output_type=FailureCause,
instructions="Classify why this Airflow task failed.",
)
result = agent.run_sync(log_line)
# Confidence is reported per output field; a bare output type lands under
# "response". A bounded ``float`` output would report none at all -- there the
# probability is the answer -- so this ``or 0.0`` would read as no confidence
# rather than as a missing one. ``LLMOperator`` and ``LLMBranchOperator`` push the
# same confidence in their ``decision`` XCom.
details = result.response.provider_details or {}
return {
"category": result.output.value,
"confidence": (details.get("confidence") or {}).get("response"),
}
@task
def act(classification: dict) -> str:
confidence = classification["confidence"] or 0.0
if confidence >= ACT_ABOVE:
return f"Remediating {classification['category']} automatically"
if confidence >= REVIEW_ABOVE:
return f"Filing {classification['category']} for review"
return "Paging a human: the classification was a coin flip"
act(classify(UNDERDETERMINED_INCIDENT))
# [END howto_decision_model_confidence]
example_decision_model_confidence()