# 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 DAGs demonstrating LLMOperator and @task.llm usage."""
from __future__ import annotations
from datetime import timedelta
from decimal import Decimal
from pydantic import BaseModel
from pydantic_ai.usage import UsageLimits
from airflow.providers.common.ai.operators.llm import LLMOperator
from airflow.providers.common.compat.notifier import BaseNotifier
from airflow.providers.common.compat.sdk import dag, task
# [START howto_operator_llm_structured_output_class]
# Pydantic output classes must be defined at module scope so they survive
# XCom serialization (their qualname is used to re-import them downstream).
[docs]
class Entities(BaseModel):
"""Named entities extracted from a text."""
# [END howto_operator_llm_structured_output_class]
# [START howto_operator_llm_basic]
@dag(tags=["example"])
[docs]
def example_llm_operator():
LLMOperator(
task_id="summarize",
prompt="Summarize the key findings from the Q4 earnings report.",
llm_conn_id="pydanticai_default",
system_prompt="You are a financial analyst. Be concise.",
)
# [END howto_operator_llm_basic]
example_llm_operator()
# [START howto_operator_llm_structured]
@dag(tags=["example"])
[docs]
def example_llm_operator_structured():
LLMOperator(
task_id="extract_entities",
prompt="Extract all named entities from the article.",
llm_conn_id="pydanticai_default",
system_prompt="Extract named entities.",
output_type=Entities,
)
# [END howto_operator_llm_structured]
example_llm_operator_structured()
# [START howto_operator_llm_agent_params]
@dag(tags=["example"])
[docs]
def example_llm_operator_agent_params():
LLMOperator(
task_id="creative_writing",
prompt="Write a haiku about data pipelines.",
llm_conn_id="pydanticai_default",
system_prompt="You are a creative writer.",
agent_params={"model_settings": {"temperature": 0.9}, "retries": 3},
)
# [END howto_operator_llm_agent_params]
example_llm_operator_agent_params()
# [START howto_decorator_llm]
@dag(tags=["example"])
[docs]
def example_llm_decorator():
@task.llm(llm_conn_id="pydanticai_default", system_prompt="Summarize concisely.")
def summarize(text: str):
return f"Summarize this article: {text}"
summarize("Apache Airflow is a platform for programmatically authoring...")
# [END howto_decorator_llm]
example_llm_decorator()
# [START howto_decorator_llm_structured]
@dag(tags=["example"])
[docs]
def example_llm_decorator_structured():
@task.llm(
llm_conn_id="pydanticai_default",
system_prompt="Extract named entities.",
output_type=Entities,
)
def extract(text: str):
return f"Extract entities from: {text}"
extract("Alice visited Paris and met Bob in London.")
# [END howto_decorator_llm_structured]
example_llm_decorator_structured()
# [START howto_operator_llm_usage_limits]
@dag(tags=["example"])
[docs]
def example_llm_operator_usage_limits():
LLMOperator(
task_id="capped_summary",
prompt="Summarize the attached design doc in three bullet points.",
llm_conn_id="pydanticai_default",
system_prompt="You are a concise technical reviewer.",
# Fail the task if the run exceeds 5 model requests, 4_000 input
# tokens, or 1_000 output tokens. Useful for guardrails on shared
# connections or untrusted prompts.
usage_limits=UsageLimits(
request_limit=5,
input_tokens_limit=4_000,
output_tokens_limit=1_000,
# Fail the task if the run's estimated USD cost exceeds $0.50.
# See docs/operators/llm.rst for caveats (not a hard guarantee;
# not enforced for models pydantic-ai can't price, which log a
# warning instead of failing the run).
cost_limit=Decimal("0.50"),
),
)
# [END howto_operator_llm_usage_limits]
example_llm_operator_usage_limits()
# [START howto_operator_llm_templated_usage_limits]
@dag(tags=["example"])
[docs]
def example_llm_operator_templated_usage_limits():
LLMOperator(
task_id="capped_summary",
prompt="Summarize the trade-offs of a message queue vs. direct HTTP calls in three bullet points.",
llm_conn_id="pydanticai_default",
system_prompt="You are a concise technical reviewer.",
# A plain dict lets every UsageLimits field be templated -- e.g. driven by
# an Airflow Variable so the budget can change per environment without
# editing the Dag. This caps a single task run, not a day's total spend --
# each run gets the full budget again. Use var.value.get() with a default
# so the example doesn't fail outright if the Variable isn't set.
usage_limits={
"cost_limit": "{{ var.value.get('llm_cost_cap_per_task', '0.50') }}",
"request_limit": 5,
},
)
# [END howto_operator_llm_templated_usage_limits]
example_llm_operator_templated_usage_limits()
# [START howto_operator_llm_approval]
[docs]
class LogNotifier(BaseNotifier):
[docs]
template_fields = ("message",)
def __init__(self, message: str) -> None:
super().__init__()
[docs]
def notify(self, context) -> None:
self.log.info(self.message)
@dag(tags=["example"])
[docs]
def example_llm_operator_approval():
LLMOperator(
task_id="summarize_with_approval",
prompt="Summarize the quarterly financial report for stakeholders.",
llm_conn_id="pydanticai_default",
system_prompt="You are a financial analyst. Be concise and accurate.",
require_approval=True,
approval_timeout=timedelta(hours=24),
on_approval_timeout="approve",
allow_modifications=True,
approval_notifiers=LogNotifier(message="{{ task.subject }}\n{{ task.body }}"),
)
# [END howto_operator_llm_approval]
example_llm_operator_approval()