airflow.providers.common.ai.mixins.usage_budget¶

Classes¶

`UsageBudgetHostProtocol`	Protocol for the operator attributes and methods `UsageBudgetMixin` relies on.
`UsageBudgetMixin`	Mixin that bounds `usage_limits` across every attempt of a task instance.

Module Contents¶

class airflow.providers.common.ai.mixins.usage_budget.UsageBudgetHostProtocol[source]¶

Bases: Protocol

Protocol for the operator attributes and methods UsageBudgetMixin relies on.

log: logging.Logger[source]¶

usage_limits: pydantic_ai.usage.UsageLimits | dict[str, Any] | None[source]¶

run_agent_sync(agent, user_prompt, **run_kwargs)[source]¶

class airflow.providers.common.ai.mixins.usage_budget.UsageBudgetMixin[source]¶

Mixin that bounds usage_limits across every attempt of a task instance.

On Airflow >= 3.3 with usage_limits set, cumulative usage is persisted in the task state store (see TaskStateStoreUsageBudget) so a task retry resumes from the usage earlier attempts already spent instead of starting a fresh count.

Operators that use this mixin must provide usage_limits, log and run_agent_sync (see CancellableAgentRunMixin).