diff --git a/providers/common/ai/README.rst b/providers/common/ai/README.rst index 9d1488e3afeb3..313d4fe2adbb4 100644 --- a/providers/common/ai/README.rst +++ b/providers/common/ai/README.rst @@ -56,7 +56,7 @@ PIP package Version required ``apache-airflow`` ``>=3.0.0`` ``apache-airflow-providers-common-compat`` ``>=1.15.0`` ``apache-airflow-providers-standard`` ``>=1.12.1`` -``pydantic-ai-slim`` ``>=2.0.0`` +``pydantic-ai-slim`` ``>=2.23.0`` ========================================== ================== Optional cross provider package dependencies diff --git a/providers/common/ai/docs/index.rst b/providers/common/ai/docs/index.rst index 4b0f82608b36d..30fd3c70a88ce 100644 --- a/providers/common/ai/docs/index.rst +++ b/providers/common/ai/docs/index.rst @@ -204,7 +204,7 @@ PIP package Version required ``apache-airflow`` ``>=3.0.0`` ``apache-airflow-providers-common-compat`` ``>=1.15.0`` ``apache-airflow-providers-standard`` ``>=1.12.1`` -``pydantic-ai-slim`` ``>=2.0.0`` +``pydantic-ai-slim`` ``>=2.23.0`` ========================================== ================== Optional cross provider package dependencies @@ -241,11 +241,11 @@ Install them when installing from PyPI. For example: ============== ======================================================================================================================================= Extra Dependencies ============== ======================================================================================================================================= -``anthropic`` ``pydantic-ai-slim[anthropic]>=2.0.0`` -``bedrock`` ``pydantic-ai-slim[bedrock]>=2.0.0`` -``google`` ``pydantic-ai-slim[google]>=2.0.0`` -``openai`` ``pydantic-ai-slim[openai]>=2.0.0`` -``mcp`` ``pydantic-ai-slim[mcp]>=2.0.0`` +``anthropic`` ``pydantic-ai-slim[anthropic]>=2.23.0`` +``bedrock`` ``pydantic-ai-slim[bedrock]>=2.23.0`` +``google`` ``pydantic-ai-slim[google]>=2.23.0`` +``openai`` ``pydantic-ai-slim[openai]>=2.23.0`` +``mcp`` ``pydantic-ai-slim[mcp]>=2.23.0`` ``code-mode`` ``pydantic-ai-harness[codemode]>=0.3.0`` ``shields`` ``pydantic-ai-shields>=0.3.4`` ``skills`` ``apache-airflow-providers-git>=0.4.0``, ``pydantic-ai-skills>=1.2.0`` diff --git a/providers/common/ai/docs/observability.rst b/providers/common/ai/docs/observability.rst index 54d5f5f7e9a36..5715b58be5b44 100644 --- a/providers/common/ai/docs/observability.rst +++ b/providers/common/ai/docs/observability.rst @@ -51,6 +51,9 @@ How it works * **Content is off by default.** Only token counts, model id, latency, tool names, and finish reason are recorded. Prompt and completion text is never emitted unless you opt in (see below). +* **Cost is already on the span.** pydantic-ai's own instrumentation sets a + best-effort ``operation.cost`` attribute on the model-call span whenever it + can price the response -- no provider configuration is needed for this. .. note:: diff --git a/providers/common/ai/docs/operators/agent.rst b/providers/common/ai/docs/operators/agent.rst index b99f196efdc98..3b50064a9d3ce 100644 --- a/providers/common/ai/docs/operators/agent.rst +++ b/providers/common/ai/docs/operators/agent.rst @@ -487,7 +487,25 @@ Parameters agent run (initial run, durable replay, and HITL regeneration). Use it to cap requests, tokens, or tool calls per task -- agents are particularly prone to runaway tool loops, so ``tool_calls_limit`` is a useful guardrail. - See :ref:`howto/operator:llm` for an example. Default ``None``. + It also supports a per-run USD ``cost_limit``; see :ref:`howto/operator:llm` + for the caveats (not a hard guarantee, silently inert for unpriced models) + and an example. Default ``None``. + + .. warning:: + With ``durable=True``, a task retry replays cached model steps instead of + re-calling the model -- but pydantic-ai still adds each replayed step's + cost to the retry's own usage total, since it cannot distinguish a replay + from a live call. A ``cost_limit`` therefore counts already-paid-for + replayed cost against every retry's fresh budget, leaving less headroom + for the new calls the retry actually makes. And if the limit is lowered + between attempts -- easy to do by accident, since ``max_cost`` is + templated -- a retry can exceed it with zero new model calls. The + ``LLM run cost`` line in the task log reports the run's cumulative cost + for the same reason, not what this attempt actually spent. +- ``max_cost``: Convenience per-run USD cost cap, as a templated alternative to + ``usage_limits.cost_limit`` (``usage_limits`` itself cannot be templated). Overrides + ``cost_limit`` on ``usage_limits`` if both are set; every other field on + ``usage_limits`` is preserved. Default ``None`` (``usage_limits`` unchanged). - ``durable``: When ``True``, enables step-level caching of model responses and tool results. On retry, cached steps are replayed instead of re-executing expensive LLM calls. On Airflow >= 3.3 the cache uses the task state store (no diff --git a/providers/common/ai/docs/operators/llm.rst b/providers/common/ai/docs/operators/llm.rst index 426a2573206fc..85a7e13d98e50 100644 --- a/providers/common/ai/docs/operators/llm.rst +++ b/providers/common/ai/docs/operators/llm.rst @@ -119,6 +119,17 @@ calls within a single task. :start-after: [START howto_operator_llm_usage_limits] :end-before: [END howto_operator_llm_usage_limits] +``usage_limits`` cannot be templated -- it's an object, not a scalar. For a per-run +cost cap, use ``max_cost`` -- a plain number (``max_cost=0.5``), and templatable: + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm.py + :language: python + :start-after: [START howto_operator_llm_max_cost] + :end-before: [END howto_operator_llm_max_cost] + +``max_cost`` overrides ``cost_limit`` on ``usage_limits`` (building one if +``usage_limits`` is ``None``); every other field on ``usage_limits`` is left as-is. + Common knobs on ``UsageLimits``: - ``request_limit`` — max model requests per run (caps retry/tool-loop blow-ups). @@ -129,6 +140,19 @@ Common knobs on ``UsageLimits``: - ``input_tokens_limit`` / ``output_tokens_limit`` — per-run token caps. - ``total_tokens_limit`` — combined input + output cap. - ``tool_calls_limit`` — max tool invocations (``AgentOperator`` only). +- ``cost_limit`` — a ``Decimal`` cap on the run's estimated USD cost. This is **not** a + hard guarantee against overspend: the response that crosses the limit has already been + produced and billed — pydantic-ai checks the accumulated cost *after* each response and + then fails the run with ``UsageLimitExceeded``. It protects you from further spend, not + from the request that broke the budget; even a single-request run fails as soon as that + request's cost pushes the total over the limit. For self-hosted or unknown + models (e.g. Ollama, custom endpoints) pydantic-ai cannot price the response, so cost + is ``None`` and ``cost_limit`` silently has no effect (a ``CostNotFoundWarning`` is + emitted instead of a failure). And like the other knobs above, setting ``cost_limit`` + alone still inherits the ``request_limit=50`` default — see the ``request_limit`` note + above. Note that ``cost_limit`` (and ``max_cost``) only cap the operator's own LLM calls -- + the meta-agent that ``LLMRetryPolicy`` runs to classify a failed task is a separate, + uncapped LLM call; see :doc:`../retry_policies`. When the limit is hit pydantic-ai raises ``UsageLimitExceeded``, which propagates to Airflow as a task failure — Airflow's standard retry policy @@ -219,6 +243,10 @@ Parameters constructor (e.g. ``retries``, ``model_settings``, ``tools``). Supports Jinja templating. - ``usage_limits``: Optional pydantic-ai ``UsageLimits`` enforced on the run. Fails the task when token / request / tool-call budgets are exceeded. Default ``None``. +- ``max_cost``: Convenience per-run USD cost cap, as a templated alternative to + ``usage_limits.cost_limit`` (``usage_limits`` itself cannot be templated). Overrides + ``cost_limit`` on ``usage_limits`` if both are set; every other field on + ``usage_limits`` is preserved. Default ``None`` (``usage_limits`` unchanged). - ``require_approval``: If ``True``, the task defers after generating output and waits for human review. Default ``False``. - ``approval_timeout``: Maximum time to wait for a review (``timedelta``). ``None`` diff --git a/providers/common/ai/docs/retry_policies.rst b/providers/common/ai/docs/retry_policies.rst index c1962e41ce869..01269f6f1b492 100644 --- a/providers/common/ai/docs/retry_policies.rst +++ b/providers/common/ai/docs/retry_policies.rst @@ -79,6 +79,12 @@ When a task fails, ``LLMRetryPolicy``: 3. Based on the classification, returns RETRY (with a suggested delay) or FAIL 4. The classification reason is logged in the task logs +This classification call is a separate LLM request, made by ``LLMRetryPolicy`` +itself rather than by an operator -- it is not subject to an operator's +``usage_limits`` or ``max_cost``, and it runs on every task failure regardless +of any cost cap configured on the failing task. It is bounded by ``timeout`` +and ``max_exception_length``, but not by a cost limit. + If the LLM call fails (provider down, timeout, bad credentials), the policy falls back to ``fallback_rules`` if configured, or to the task's standard retry behaviour. diff --git a/providers/common/ai/pyproject.toml b/providers/common/ai/pyproject.toml index aa806092833df..cc294539ade92 100644 --- a/providers/common/ai/pyproject.toml +++ b/providers/common/ai/pyproject.toml @@ -69,18 +69,19 @@ dependencies = [ "apache-airflow>=3.0.0", "apache-airflow-providers-common-compat>=1.15.0", "apache-airflow-providers-standard>=1.12.1", - # Requires the pydantic-ai 2.x agent/instrumentation API (see #69122). - "pydantic-ai-slim>=2.0.0", + # Requires the pydantic-ai cost API (RunUsage.cost, UsageLimits.cost_limit), + # landed in 2.23.0 via https://github.com/pydantic/pydantic-ai/pull/2684. + "pydantic-ai-slim>=2.23.0", ] # The optional dependencies should be modified in place in the generated file # Any change in the dependencies is preserved when the file is regenerated [project.optional-dependencies] -"anthropic" = ["pydantic-ai-slim[anthropic]>=2.0.0"] -"bedrock" = ["pydantic-ai-slim[bedrock]>=2.0.0"] -"google" = ["pydantic-ai-slim[google]>=2.0.0"] -"openai" = ["pydantic-ai-slim[openai]>=2.0.0"] -"mcp" = ["pydantic-ai-slim[mcp]>=2.0.0"] +"anthropic" = ["pydantic-ai-slim[anthropic]>=2.23.0"] +"bedrock" = ["pydantic-ai-slim[bedrock]>=2.23.0"] +"google" = ["pydantic-ai-slim[google]>=2.23.0"] +"openai" = ["pydantic-ai-slim[openai]>=2.23.0"] +"mcp" = ["pydantic-ai-slim[mcp]>=2.23.0"] # Code mode: collapse tool calls into a single `run_code` tool that the model # drives by writing Python, executed in the Monty sandbox (pydantic-monty). # Enables AgentOperator(code_mode=True). Monty is pre-1.0; pinned here as an @@ -137,7 +138,7 @@ dev = [ "apache-airflow-providers-standard", # Additional devel dependencies (do not remove this line and add extra development dependencies) "sqlglot>=30.0.0", - "pydantic-ai-slim[mcp]>=2.0.0", + "pydantic-ai-slim[mcp]>=2.23.0", "pydantic-ai-skills>=1.2.0", "apache-airflow-providers-common-sql[datafusion]", "langchain>=1.0.0", diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py index 545a138e9df9f..766864b86136b 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py +++ b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py @@ -19,6 +19,7 @@ from __future__ import annotations from datetime import timedelta +from decimal import Decimal from pydantic import BaseModel from pydantic_ai.usage import UsageLimits @@ -139,6 +140,10 @@ def example_llm_operator_usage_limits(): 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, + # silently inert for models pydantic-ai can't price). + cost_limit=Decimal("0.50"), ), ) @@ -148,6 +153,27 @@ def example_llm_operator_usage_limits(): example_llm_operator_usage_limits() +# [START howto_operator_llm_max_cost] +@dag(tags=["example"]) +def example_llm_operator_max_cost(): + 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.", + # Unlike usage_limits (a UsageLimits object), max_cost is a scalar and + # can 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. + max_cost="{{ var.value.llm_max_cost_per_task }}", + ) + + +# [END howto_operator_llm_max_cost] + +example_llm_operator_max_cost() + + # [START howto_operator_llm_approval] @dag(tags=["example"]) def example_llm_operator_approval(): diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py index 83a03bfda02f7..c69f2f22f744b 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py @@ -31,6 +31,7 @@ from airflow.providers.common.ai.mixins.hitl_review import HITLReviewMixin from airflow.providers.common.ai.utils.logging import log_run_summary, wrap_toolsets_for_logging from airflow.providers.common.ai.utils.output_type import rehydrate_pydantic_output +from airflow.providers.common.ai.utils.usage import resolve_usage_limits from airflow.providers.common.compat.sdk import ( AirflowOptionalProviderFeatureException, BaseOperator, @@ -48,6 +49,8 @@ _CORE_WALKER = False if TYPE_CHECKING: + from decimal import Decimal + from pydantic_ai import Agent from pydantic_ai.messages import ModelMessage from pydantic_ai.toolsets.abstract import AbstractToolset @@ -148,6 +151,14 @@ class AgentOperator(BaseOperator, HITLReviewMixin): ``UsageLimits(request_limit=..., total_tokens_limit=..., tool_calls_limit=..., ...)`` to fail the task when the agent exceeds the configured token, request, or tool budget. ``None`` (default) means no enforcement. + :param max_cost: Convenience per-run USD cost cap, as a templated alternative + to ``usage_limits.cost_limit`` (``usage_limits`` itself cannot be + templated). When set, overrides ``cost_limit`` on ``usage_limits`` + (building one if ``usage_limits`` is ``None``); every other field on + ``usage_limits`` is left untouched. ``None`` (default) leaves + ``usage_limits`` unchanged. See :ref:`howto/operator:llm` for general + ``cost_limit`` caveats, and :ref:`howto/operator:agent` for the + ``durable=True`` replay double-counting warning. :param durable: When ``True``, enables step-level caching of model responses and tool results for durable execution. On retry, cached steps are replayed instead of re-executing. Each cached step is @@ -232,6 +243,7 @@ class AgentOperator(BaseOperator, HITLReviewMixin): "system_prompt", "agent_params", "message_history", + "max_cost", ) operator_extra_links = (HITLReviewLink(),) @@ -248,6 +260,7 @@ def __init__( enable_tool_logging: bool = True, agent_params: dict[str, Any] | None = None, usage_limits: UsageLimits | None = None, + max_cost: Decimal | float | str | None = None, durable: bool = False, code_mode: bool = False, message_history: list[ModelMessage] | str | bytes | None = None, @@ -274,6 +287,7 @@ def __init__( self.enable_tool_logging = enable_tool_logging self.agent_params = agent_params or {} self.usage_limits = usage_limits + self.max_cost = max_cost self.message_history = message_history self.durable = durable @@ -446,7 +460,7 @@ def execute(self, context: Context) -> Any: agent = self._build_agent() - run_kwargs: dict[str, Any] = {"usage_limits": self.usage_limits} + run_kwargs: dict[str, Any] = {"usage_limits": resolve_usage_limits(self.usage_limits, self.max_cost)} history = self._resolve_message_history() if history is not None: run_kwargs["message_history"] = history @@ -553,7 +567,11 @@ def regenerate_with_feedback(self, *, feedback: str, message_history: Any) -> tu """Re-run the agent with *feedback* appended to the conversation history.""" agent = self._build_agent() messages = message_history or [] - result = agent.run_sync(feedback, message_history=messages, usage_limits=self.usage_limits) + result = agent.run_sync( + feedback, + message_history=messages, + usage_limits=resolve_usage_limits(self.usage_limits, self.max_cost), + ) log_run_summary(self.log, result) output = result.output diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py index 36cd616c1b386..f425d523c8d8d 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py @@ -29,6 +29,7 @@ from airflow.providers.common.ai.mixins.approval import LLMApprovalMixin from airflow.providers.common.ai.utils.logging import log_run_summary from airflow.providers.common.ai.utils.output_type import rehydrate_pydantic_output +from airflow.providers.common.ai.utils.usage import resolve_usage_limits from airflow.providers.common.compat.sdk import BaseOperator try: @@ -41,6 +42,8 @@ _CORE_WALKER = False if TYPE_CHECKING: + from decimal import Decimal + from pydantic_ai import Agent from pydantic_ai.usage import UsageLimits @@ -80,6 +83,12 @@ class LLMOperator(BaseOperator, LLMApprovalMixin): ``UsageLimits(request_limit=..., total_tokens_limit=..., ...)`` to fail the task when the agent exceeds the configured token, request, or tool budget. ``None`` (default) means no enforcement. + :param max_cost: Convenience per-run USD cost cap, as a templated alternative + to ``usage_limits.cost_limit`` (``usage_limits`` itself cannot be + templated). When set, overrides ``cost_limit`` on ``usage_limits`` + (building one if ``usage_limits`` is ``None``); every other field on + ``usage_limits`` is left untouched. ``None`` (default) leaves + ``usage_limits`` unchanged. See :ref:`howto/operator:llm` for caveats. :param require_approval: If ``True``, the task defers after generating output and waits for a human reviewer to approve or reject via the HITL interface. Default ``False``. @@ -104,6 +113,7 @@ class LLMOperator(BaseOperator, LLMApprovalMixin): "model_id", "system_prompt", "agent_params", + "max_cost", ) def __init__( @@ -116,6 +126,7 @@ def __init__( output_type: type = str, agent_params: dict[str, Any] | None = None, usage_limits: UsageLimits | None = None, + max_cost: Decimal | float | str | None = None, require_approval: bool = False, approval_timeout: timedelta | None = None, allow_modifications: bool = False, @@ -135,6 +146,7 @@ def __init__( self._serialize_model_output = serialize_output or not _CORE_WALKER self.agent_params = agent_params or {} self.usage_limits = usage_limits + self.max_cost = max_cost self.require_approval = require_approval self.approval_timeout = approval_timeout self.allow_modifications = allow_modifications @@ -161,7 +173,9 @@ def execute(self, context: Context) -> Any: agent: Agent[object, Any] = self.llm_hook.create_agent( output_type=self.output_type, instructions=self.system_prompt, **self.agent_params ) - result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) + result = agent.run_sync( + self.prompt, usage_limits=resolve_usage_limits(self.usage_limits, self.max_cost) + ) log_run_summary(self.log, result) output = result.output diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py index 88c299432594b..7bdf1c7f8b5fd 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py @@ -25,6 +25,7 @@ from airflow.providers.common.ai.operators.llm import LLMOperator from airflow.providers.common.ai.utils.logging import log_run_summary +from airflow.providers.common.ai.utils.usage import resolve_usage_limits from airflow.providers.standard.operators.branch import BranchMixIn if TYPE_CHECKING: @@ -96,7 +97,9 @@ def execute(self, context: Context) -> str | Iterable[str] | None: instructions=self.system_prompt, **self.agent_params, ) - result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) + result = agent.run_sync( + self.prompt, usage_limits=resolve_usage_limits(self.usage_limits, self.max_cost) + ) log_run_summary(self.log, result) output = result.output diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py index 8e919acf13b27..ad3a017885d94 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py @@ -26,6 +26,7 @@ from airflow.providers.common.ai.operators.llm import LLMOperator from airflow.providers.common.ai.utils.file_analysis import build_file_analysis_request from airflow.providers.common.ai.utils.logging import log_run_summary +from airflow.providers.common.ai.utils.usage import resolve_usage_limits if TYPE_CHECKING: from pydantic_ai import Agent @@ -134,7 +135,9 @@ def execute(self, context: Context) -> Any: instructions=self._build_system_prompt(), **self.agent_params, ) - result = agent.run_sync(request.user_content, usage_limits=self.usage_limits) + result = agent.run_sync( + request.user_content, usage_limits=resolve_usage_limits(self.usage_limits, self.max_cost) + ) log_run_summary(self.log, result) output = result.output diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py index 7992419296fa2..e73ccce86f658 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py @@ -28,6 +28,7 @@ from airflow.providers.common.ai.operators.llm import LLMOperator from airflow.providers.common.ai.utils.logging import log_run_summary +from airflow.providers.common.ai.utils.usage import resolve_usage_limits from airflow.providers.common.compat.sdk import AirflowException, BaseHook if TYPE_CHECKING: @@ -322,7 +323,9 @@ def execute(self, context: Context) -> dict[str, Any]: **self.agent_params, ) self.log.info("Running LLM schema comparison...") - result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) + result = agent.run_sync( + self.prompt, usage_limits=resolve_usage_limits(self.usage_limits, self.max_cost) + ) log_run_summary(self.log, result) output = result.output diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py index 262d52ff8d0f4..960cb243ec975 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py @@ -35,6 +35,7 @@ from airflow.providers.common.ai.operators.llm import LLMOperator from airflow.providers.common.ai.utils.logging import log_run_summary +from airflow.providers.common.ai.utils.usage import resolve_usage_limits from airflow.providers.common.compat.sdk import BaseHook if TYPE_CHECKING: @@ -144,7 +145,9 @@ def execute(self, context: Context) -> str: agent = self.llm_hook.create_agent( output_type=str, instructions=full_system_prompt, **self.agent_params ) - result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) + result = agent.run_sync( + self.prompt, usage_limits=resolve_usage_limits(self.usage_limits, self.max_cost) + ) log_run_summary(self.log, result) sql = self._strip_llm_output(result.output, dialect=self._resolved_dialect) diff --git a/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py b/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py index 52dcf1d9ea068..2b7fa49634a15 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py +++ b/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py @@ -48,6 +48,8 @@ def log_run_summary(logger: Logger | logging.Logger, result: AgentRunResult[Any] usage.output_tokens, usage.total_tokens, ) + if usage.cost is not None: + logger.info("LLM run cost: $%s (USD, best-effort)", usage.cost) tool_names = _extract_tool_sequence(result) if tool_names: diff --git a/providers/common/ai/src/airflow/providers/common/ai/utils/usage.py b/providers/common/ai/src/airflow/providers/common/ai/utils/usage.py new file mode 100644 index 0000000000000..23cc9b98e8c21 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/utils/usage.py @@ -0,0 +1,69 @@ +# 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. +"""Helpers for merging the ``max_cost`` convenience parameter into ``UsageLimits``.""" + +from __future__ import annotations + +import dataclasses +from decimal import Decimal, InvalidOperation + +from pydantic_ai.usage import UsageLimits + + +def resolve_usage_limits( + usage_limits: UsageLimits | None, max_cost: Decimal | float | str | None +) -> UsageLimits | None: + """ + Merge ``max_cost`` into ``usage_limits.cost_limit``, leaving other fields untouched. + + ``max_cost`` exists only because ``usage_limits`` (a ``UsageLimits`` object) cannot be + templated -- it isn't a scalar and isn't in any operator's ``template_fields``. ``max_cost`` + is the templatable escape hatch for the single most common knob, which also means its value + is not always in the Dag author's control (an unset Airflow Variable renders to ``""``, a + typo renders to a non-numeric string). + + - ``max_cost is None``: ``usage_limits`` is returned unchanged (same object, by identity) -- + this keeps every existing ``usage_limits=None`` assertion in the test suite a true no-op. + - ``max_cost`` has a value: builds (or copies) a ``UsageLimits`` with ``cost_limit`` set to + ``Decimal(str(max_cost))`` -- never ``Decimal(max_cost)``, which would bake in binary-float + noise for a value like ``0.1``. ``max_cost`` overrides any ``cost_limit`` already present + on ``usage_limits``; every other field is preserved as-is. + + :raises ValueError: if ``max_cost`` cannot be parsed as a number, is not finite + (``inf``, ``-inf``, or ``nan``), or is negative. + """ + if max_cost is None: + return usage_limits + + try: + cost_limit = Decimal(str(max_cost)) + except InvalidOperation: + raise ValueError( + f"max_cost must be a number or a numeric string (got {max_cost!r}); " + "if it is templated, check the rendered value." + ) from None + if not cost_limit.is_finite(): + raise ValueError( + f"max_cost must be a finite number (got {max_cost!r}); " + "a non-finite value would silently disable the cost cap." + ) + if cost_limit < 0: + raise ValueError(f"max_cost must not be negative (got {max_cost!r})") + + if usage_limits is None: + return UsageLimits(cost_limit=cost_limit) + return dataclasses.replace(usage_limits, cost_limit=cost_limit) diff --git a/providers/common/ai/tests/unit/common/ai/conftest.py b/providers/common/ai/tests/unit/common/ai/conftest.py index e01654f907548..eca4a26bd2c43 100644 --- a/providers/common/ai/tests/unit/common/ai/conftest.py +++ b/providers/common/ai/tests/unit/common/ai/conftest.py @@ -31,15 +31,22 @@ def isolate_hook_lineage_collector(hook_lineage_collector): return None -def make_mock_run_result(output): +def make_mock_run_result(output, *, cost=None): """Create a mock AgentRunResult compatible with log_run_summary. Returns a MagicMock with .output, .usage, .response, and .all_messages() configured so that log_run_summary can read them without error. + + ``cost`` defaults to ``None`` and must be set explicitly -- a MagicMock + attribute left unconfigured returns a new (truthy, ``is not None``) + MagicMock, which would silently push every caller through + ``log_run_summary``'s cost-logging branch. """ mock_result = MagicMock() mock_result.output = output - mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0) + mock_result.usage = MagicMock( + requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0, cost=cost + ) mock_result.response = MagicMock(model_name="test-model") mock_result.all_messages.return_value = [] return mock_result diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py b/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py index 1a9d82bd065bb..84fed220a6268 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py @@ -44,7 +44,9 @@ def _make_mock_run_result(output): """Create a mock AgentRunResult compatible with log_run_summary.""" mock_result = MagicMock() mock_result.output = output - mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0) + mock_result.usage = MagicMock( + requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0, cost=None + ) mock_result.response = MagicMock(model_name="test-model") mock_result.all_messages.return_value = [] return mock_result diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py b/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py index 9a1c389fef513..49897bccbe433 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py @@ -28,7 +28,9 @@ def _make_mock_run_result(output): """Create a mock AgentRunResult compatible with log_run_summary.""" mock_result = MagicMock() mock_result.output = output - mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0) + mock_result.usage = MagicMock( + requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0, cost=None + ) mock_result.response = MagicMock(model_name="test-model") mock_result.all_messages.return_value = [] return mock_result diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py index 2243315dcf344..6514cc1547c7c 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py @@ -30,7 +30,9 @@ def _make_mock_run_result(output): """Create a mock AgentRunResult compatible with log_run_summary.""" mock_result = MagicMock() mock_result.output = output - mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0) + mock_result.usage = MagicMock( + requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0, cost=None + ) mock_result.response = MagicMock(model_name="test-model") mock_result.all_messages.return_value = [] return mock_result diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py index cac472b765763..2346ed6a0cf67 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py @@ -28,12 +28,13 @@ def _make_mock_run_result(output): mock_result = MagicMock(spec=["output", "usage", "response", "all_messages"]) mock_result.output = output mock_result.usage = MagicMock( - spec=["requests", "tool_calls", "input_tokens", "output_tokens", "total_tokens"], + spec=["requests", "tool_calls", "input_tokens", "output_tokens", "total_tokens", "cost"], requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0, + cost=None, ) mock_result.response = MagicMock(spec=["model_name"], model_name="test-model") mock_result.all_messages.return_value = [] diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py index 0664b6c52b18c..03c6f250dc0bf 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py @@ -32,7 +32,9 @@ def _make_mock_run_result(output): """Create a mock AgentRunResult compatible with log_run_summary.""" mock_result = MagicMock() mock_result.output = output - mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0) + mock_result.usage = MagicMock( + requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0, cost=None + ) mock_result.response = MagicMock(model_name="test-model") mock_result.all_messages.return_value = [] return mock_result diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py index d44b5d016f56f..2c8b8a53525ee 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py @@ -28,7 +28,9 @@ def _make_mock_run_result(output): """Create a mock AgentRunResult compatible with log_run_summary.""" mock_result = MagicMock() mock_result.output = output - mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0) + mock_result.usage = MagicMock( + requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0, cost=None + ) mock_result.response = MagicMock(model_name="test-model") mock_result.all_messages.return_value = [] return mock_result diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_replay_cost.py b/providers/common/ai/tests/unit/common/ai/durable/test_replay_cost.py new file mode 100644 index 0000000000000..20dc21130004c --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/durable/test_replay_cost.py @@ -0,0 +1,132 @@ +# 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. +""" +Empirical check of whether durable replay double-counts cost against ``cost_limit``. + +pydantic-ai's graph appends *every* model response's usage to the run's +``RunUsage`` in ``_agent_graph.py::_append_response`` -- it cannot distinguish a +response that came from a live model call from one ``CachingModel`` replayed +from the durable cache. Each Airflow task attempt starts a fresh ``RunUsage`` +(a new ``agent.run`` call), so a step that was already paid for in a prior, +crashed attempt gets its cost added again to the retry's own usage total -- +even though the retry made zero new model calls for that step. These tests +exercise the real ``CachingModel`` + ``DurableStorage`` + pydantic-ai ``Agent`` +stack (no mocked cost arithmetic) to confirm this, rather than relying on +reading ``_agent_graph.py`` / ``_cost.py`` and assuming. +""" + +from __future__ import annotations + +from decimal import Decimal +from unittest.mock import patch + +import pytest +from pydantic_ai import Agent +from pydantic_ai.exceptions import UsageLimitExceeded +from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart +from pydantic_ai.models.function import AgentInfo, FunctionModel +from pydantic_ai.usage import RequestUsage, UsageLimits + +from airflow.providers.common.ai.durable.caching_model import CachingModel +from airflow.providers.common.ai.durable.step_counter import DurableStepCounter +from airflow.providers.common.ai.durable.storage import DurableStorage +from airflow.sdk import ObjectStoragePath + +PRICED_COST = Decimal("0.10") + + +def _priced_model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + return ModelResponse( + parts=[TextPart(content="the answer")], + usage=RequestUsage(input_tokens=100, output_tokens=50, cost=PRICED_COST), + ) + + +@pytest.fixture +def durable_storage(tmp_path): + """A real, file-backed DurableStorage -- exercises the actual JSON round-trip.""" + with patch("airflow.providers.common.ai.durable.storage._get_base_path") as mock_base: + mock_base.return_value = ObjectStoragePath(f"file://{tmp_path.as_posix()}") + yield DurableStorage(dag_id="dag", task_id="task", run_id="run_1", map_index=-1) + + +async def _run_one_attempt(storage: DurableStorage, *, cost_limit: Decimal | None = None): + """Simulate one Airflow task attempt: fresh Agent + fresh DurableStepCounter, shared cache.""" + counter = DurableStepCounter() + caching = CachingModel(FunctionModel(_priced_model_fn), storage=storage, counter=counter) + agent = Agent(model=caching) + result = await agent.run("What is the answer?", usage_limits=UsageLimits(cost_limit=cost_limit)) + return result, counter + + +def _reopen_storage() -> DurableStorage: + """Build a fresh ``DurableStorage`` for the same dag/task/run -- simulates a new Airflow + task attempt (new process) reloading the durable cache from disk via the public + constructor, rather than reaching into the private ``_cache`` attribute.""" + return DurableStorage(dag_id="dag", task_id="task", run_id="run_1", map_index=-1) + + +class TestDurableReplayCostDuplication: + @pytest.mark.asyncio + async def test_replayed_step_cost_is_recounted_on_retry(self, durable_storage): + """A second attempt that only replays cached steps still reports the replayed cost + as its own usage -- pydantic-ai cannot tell a replay from a live call.""" + result1, counter1 = await _run_one_attempt(durable_storage) + assert counter1.cached_model == 1 + assert counter1.replayed_model == 0 + assert result1.usage.cost == PRICED_COST + + # New attempt: fresh process, so the cache is reloaded from disk via a new + # DurableStorage -- this is what actually happens on an Airflow task retry. + result2, counter2 = await _run_one_attempt(_reopen_storage()) + + # Zero new model calls this attempt ... + assert counter2.cached_model == 0 + assert counter2.replayed_model == 1 + # ... yet the replayed step's cost is counted again, identically to attempt 1. + assert result2.usage.cost == PRICED_COST + + @pytest.mark.asyncio + async def test_retry_with_zero_new_spend_still_raises_usage_limit_exceeded(self, durable_storage): + """A retry that makes no new model calls can still raise UsageLimitExceeded, + purely from replayed cost -- because check_cost() sees the run's cumulative + usage, not "money spent in this attempt".""" + # Attempt 1 stays comfortably under budget so it completes normally. + await _run_one_attempt(durable_storage, cost_limit=PRICED_COST * 2) + + # Attempt 2 sets a limit below the already-paid-for replayed cost: zero new + # spend, yet the replayed step alone pushes the cumulative usage over it. + cost_limit = PRICED_COST / 2 + with pytest.raises(UsageLimitExceeded): + await _run_one_attempt(_reopen_storage(), cost_limit=cost_limit) + + +class TestDurableStorageCostRoundTrip: + def test_decimal_cost_survives_json_round_trip(self, durable_storage): + """DurableStorage serializes the whole cache blob as JSON; confirm a Decimal + ``usage.cost`` is not silently lost or coerced to float/None by that round-trip.""" + response = ModelResponse( + parts=[TextPart(content="hi")], + usage=RequestUsage(input_tokens=1, output_tokens=1, cost=Decimal("0.0123456789")), + ) + durable_storage.save_model_response("model_step_0", response, fingerprint="fp") + + loaded, _fingerprint = _reopen_storage().load_model_response("model_step_0") + + assert loaded is not None + assert loaded.usage.cost == Decimal("0.0123456789") + assert isinstance(loaded.usage.cost, Decimal) diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py index 760bacb8cc7f6..44f5d159fad40 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py @@ -18,6 +18,7 @@ import sys from datetime import timedelta +from decimal import Decimal from unittest.mock import MagicMock, patch import pytest @@ -67,7 +68,9 @@ def _make_mock_run_result(output): """Create a mock AgentRunResult compatible with log_run_summary.""" mock_result = MagicMock() mock_result.output = output - mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0) + mock_result.usage = MagicMock( + requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0, cost=None + ) mock_result.response = MagicMock(model_name="test-model") mock_result.all_messages.return_value = [] return mock_result @@ -141,6 +144,7 @@ def test_template_fields(self): "system_prompt", "agent_params", "message_history", + "max_cost", } assert set(AgentOperator.template_fields) == expected @@ -163,6 +167,23 @@ def test_execute_forwards_usage_limits_to_run_sync(self, mock_hook_cls): mock_agent.run_sync.assert_called_once_with("run", usage_limits=limits) + @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + def test_execute_forwards_max_cost_as_cost_limit(self, mock_hook_cls): + """``max_cost`` is resolved into ``usage_limits.cost_limit`` before ``run_sync``.""" + mock_agent = _make_mock_agent("ok") + mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + + op = AgentOperator( + task_id="test", + prompt="run", + llm_conn_id="my_llm", + max_cost=0.5, + ) + op.execute(context=MagicMock()) + + _, kwargs = mock_agent.run_sync.call_args + assert kwargs["usage_limits"].cost_limit == Decimal("0.5") + @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) def test_regenerate_with_feedback_forwards_usage_limits(self, mock_hook_cls): """``usage_limits`` is also forwarded by ``regenerate_with_feedback``.""" @@ -184,6 +205,23 @@ def test_regenerate_with_feedback_forwards_usage_limits(self, mock_hook_cls): usage_limits=limits, ) + @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + def test_regenerate_with_feedback_forwards_max_cost(self, mock_hook_cls): + """``max_cost`` is also resolved into ``usage_limits.cost_limit`` by ``regenerate_with_feedback``.""" + mock_agent = _make_mock_agent("revised") + mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + + op = AgentOperator( + task_id="test", + prompt="run", + llm_conn_id="my_llm", + max_cost=0.5, + ) + op.regenerate_with_feedback(feedback="Add detail", message_history=[]) + + _, kwargs = mock_agent.run_sync.call_args + assert kwargs["usage_limits"].cost_limit == Decimal("0.5") + @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) def test_execute_creates_agent_from_hook(self, mock_hook_cls): mock_agent = _make_mock_agent("The answer is 42.") diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py index e2004b4031c32..27eb499ff9e43 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py @@ -17,6 +17,7 @@ from __future__ import annotations from datetime import timedelta +from decimal import Decimal from unittest.mock import MagicMock, patch from uuid import uuid4 @@ -69,7 +70,9 @@ def _make_mock_run_result(output): """Create a mock AgentRunResult compatible with log_run_summary.""" mock_result = MagicMock() mock_result.output = output - mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0) + mock_result.usage = MagicMock( + requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0, cost=None + ) mock_result.response = MagicMock(model_name="test-model") mock_result.all_messages.return_value = [] return mock_result @@ -77,7 +80,7 @@ def _make_mock_run_result(output): class TestLLMOperator: def test_template_fields(self): - expected = {"prompt", "llm_conn_id", "model_id", "system_prompt", "agent_params"} + expected = {"prompt", "llm_conn_id", "model_id", "system_prompt", "agent_params", "max_cost"} assert set(LLMOperator.template_fields) == expected @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) @@ -115,6 +118,24 @@ def test_execute_forwards_usage_limits_to_run_sync(self, mock_hook_cls): mock_agent.run_sync.assert_called_once_with("Summarize", usage_limits=limits) + @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + def test_execute_forwards_max_cost_as_cost_limit(self, mock_hook_cls): + """``max_cost`` is resolved into ``usage_limits.cost_limit`` before ``run_sync``.""" + mock_agent = MagicMock(spec=["run_sync"]) + mock_agent.run_sync.return_value = _make_mock_run_result("ok") + mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + + op = LLMOperator( + task_id="test", + prompt="Summarize", + llm_conn_id="my_llm", + max_cost=0.5, + ) + op.execute(context=MagicMock()) + + _, kwargs = mock_agent.run_sync.call_args + assert kwargs["usage_limits"].cost_limit == Decimal("0.5") + @requires_typed_xcom @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) def test_execute_structured_output_with_all_params(self, mock_hook_cls): diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py index 7f2d11bf26336..8deccb5b4b748 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +from decimal import Decimal from enum import Enum from unittest.mock import MagicMock, patch from uuid import uuid4 @@ -42,7 +43,9 @@ def _make_mock_run_result(output): """Create a mock AgentRunResult compatible with log_run_summary.""" mock_result = MagicMock() mock_result.output = output - mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0) + mock_result.usage = MagicMock( + requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0, cost=None + ) mock_result.response = MagicMock(model_name="test-model") mock_result.all_messages.return_value = [] return mock_result @@ -92,6 +95,30 @@ def test_execute_single_branch(self, mock_hook_cls, mock_do_branch): mock_do_branch.assert_called_once_with(ctx, "task_a") mock_agent.run_sync.assert_called_once_with("Pick a branch", usage_limits=None) + @patch.object(LLMBranchOperator, "do_branch") + @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + def test_execute_forwards_max_cost_as_cost_limit(self, mock_hook_cls, mock_do_branch): + """``max_cost`` is resolved into ``usage_limits.cost_limit`` before ``run_sync``.""" + downstream_enum = Enum("DownstreamTasks", {"task_a": "task_a", "task_b": "task_b"}) + + mock_agent = MagicMock(spec=["run_sync"]) + mock_agent.run_sync.return_value = _make_mock_run_result(downstream_enum.task_a) + mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_do_branch.return_value = "task_a" + + op = LLMBranchOperator( + task_id="test", + prompt="Pick a branch", + llm_conn_id="my_llm", + max_cost=0.5, + ) + op.downstream_task_ids = {"task_a", "task_b"} + + op.execute(MagicMock()) + + _, kwargs = mock_agent.run_sync.call_args + assert kwargs["usage_limits"].cost_limit == Decimal("0.5") + @patch.object(LLMBranchOperator, "do_branch") @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) def test_execute_multi_branch(self, mock_hook_cls, mock_do_branch): diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py index 9e692b420f9ea..d650218f77b15 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py @@ -17,6 +17,7 @@ from __future__ import annotations from datetime import timedelta +from decimal import Decimal from unittest.mock import MagicMock, patch from uuid import uuid4 @@ -56,12 +57,13 @@ def _make_mock_run_result(output): mock_result = MagicMock(spec=["output", "usage", "response", "all_messages"]) mock_result.output = output mock_result.usage = MagicMock( - spec=["requests", "tool_calls", "input_tokens", "output_tokens", "total_tokens"], + spec=["requests", "tool_calls", "input_tokens", "output_tokens", "total_tokens", "cost"], requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0, + cost=None, ) mock_result.response = MagicMock(spec=["model_name"], model_name="test-model") mock_result.all_messages.return_value = [] @@ -85,6 +87,7 @@ def test_template_fields(self): "model_id", "system_prompt", "agent_params", + "max_cost", "file_path", "file_conn_id", } @@ -126,6 +129,33 @@ def test_execute_returns_string_output(self, mock_build_request, mock_hook_cls): ) mock_agent.run_sync.assert_called_once_with("prepared prompt", usage_limits=None) + @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch( + "airflow.providers.common.ai.operators.llm_file_analysis.build_file_analysis_request", autospec=True + ) + def test_execute_forwards_max_cost_as_cost_limit(self, mock_build_request, mock_hook_cls): + """``max_cost`` is resolved into ``usage_limits.cost_limit`` before ``run_sync``.""" + mock_build_request.return_value = FileAnalysisRequest( + user_content="prepared prompt", + resolved_paths=["/tmp/app.log"], + total_size_bytes=10, + ) + mock_agent = MagicMock(spec=["run_sync"]) + mock_agent.run_sync.return_value = _make_mock_run_result("Analysis complete") + mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + + op = LLMFileAnalysisOperator( + task_id="test", + prompt="Summarize the file", + llm_conn_id="my_llm", + file_path="/tmp/app.log", + max_cost=0.5, + ) + op.execute(context={}) + + _, kwargs = mock_agent.run_sync.call_args + assert kwargs["usage_limits"].cost_limit == Decimal("0.5") + @requires_typed_xcom @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) @patch( diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py index 656457de2958a..d6f81764c04b1 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py @@ -17,6 +17,7 @@ from __future__ import annotations import json +from decimal import Decimal from unittest import mock from unittest.mock import MagicMock from uuid import uuid4 @@ -51,7 +52,9 @@ def _make_mock_run_result(output): """Create a mock AgentRunResult compatible with log_run_summary.""" mock_result = MagicMock() mock_result.output = output - mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0) + mock_result.usage = MagicMock( + requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0, cost=None + ) mock_result.response = MagicMock(model_name="test-model") mock_result.all_messages.return_value = [] return mock_result @@ -292,6 +295,41 @@ def test_execute(self, mock_build_system_prompt, mock_build_schema_context): mock_agent.run_sync.assert_called_once_with("user_prompt", usage_limits=None) assert result == {"compatible": True, "mismatches": [], "summary": "All good"} + @mock.patch( + "airflow.providers.common.ai.operators.llm_schema_compare.LLMSchemaCompareOperator._build_schema_context" + ) + @mock.patch( + "airflow.providers.common.ai.operators.llm_schema_compare.LLMSchemaCompareOperator._build_system_prompt" + ) + def test_execute_forwards_max_cost_as_cost_limit( + self, mock_build_system_prompt, mock_build_schema_context + ): + """``max_cost`` is resolved into ``usage_limits.cost_limit`` before ``run_sync``.""" + mock_build_schema_context.return_value = "schema_context" + mock_build_system_prompt.return_value = "system_prompt" + + op = LLMSchemaCompareOperator( + task_id="test", + prompt="user_prompt", + llm_conn_id="llm_conn", + db_conn_ids=["postgres_default", "snowflake_default"], + table_names=["orders"], + max_cost=0.5, + ) + + mock_llm_hook = mock.Mock() + mock_agent = mock.Mock() + mock_agent.run_sync.return_value = _make_mock_run_result( + SchemaCompareResult(compatible=True, mismatches=[], summary="All good") + ) + mock_llm_hook.create_agent.return_value = mock_agent + op.llm_hook = mock_llm_hook + + op.execute(context={}) + + _, kwargs = mock_agent.run_sync.call_args + assert kwargs["usage_limits"].cost_limit == Decimal("0.5") + @mock.patch( "airflow.providers.common.ai.operators.llm_schema_compare.LLMSchemaCompareOperator._get_db_hook" ) diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py index a6fe35277adfd..84a903b5b40b4 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py @@ -19,6 +19,7 @@ import subprocess import sys from datetime import timedelta +from decimal import Decimal from unittest.mock import MagicMock, PropertyMock, patch from uuid import uuid4 @@ -47,7 +48,9 @@ def _make_mock_run_result(output): """Create a mock AgentRunResult compatible with log_run_summary.""" mock_result = MagicMock() mock_result.output = output - mock_result.usage = MagicMock(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0) + mock_result.usage = MagicMock( + requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0, cost=None + ) mock_result.response = MagicMock(model_name="test-model") mock_result.all_messages.return_value = [] return mock_result @@ -207,6 +210,7 @@ def test_template_fields_include_parent_and_sql_specific(self): "model_id", "system_prompt", "agent_params", + "max_cost", "db_conn_id", "table_names", "schema_context", @@ -230,6 +234,24 @@ def test_execute_with_schema_context(self, mock_hook_cls): assert result == "SELECT id, name FROM users WHERE active = true" mock_agent.run_sync.assert_called_once_with("Get active users", usage_limits=None) + @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + def test_execute_forwards_max_cost_as_cost_limit(self, mock_hook_cls): + """``max_cost`` is resolved into ``usage_limits.cost_limit`` before ``run_sync``.""" + mock_agent = _make_mock_agent("SELECT id, name FROM users WHERE active = true") + mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + + op = LLMSQLQueryOperator( + task_id="test", + prompt="Get active users", + llm_conn_id="my_llm", + schema_context="Table: users\nColumns: id INT, name TEXT, active BOOLEAN", + max_cost=0.5, + ) + op.execute(context=MagicMock()) + + _, kwargs = mock_agent.run_sync.call_args + assert kwargs["usage_limits"].cost_limit == Decimal("0.5") + @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) def test_execute_validation_blocks_unsafe_sql(self, mock_hook_cls): """Validation catches unsafe SQL generated by the LLM.""" diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_logging.py b/providers/common/ai/tests/unit/common/ai/utils/test_logging.py index bbc91df0850d8..3a24121ac02a8 100644 --- a/providers/common/ai/tests/unit/common/ai/utils/test_logging.py +++ b/providers/common/ai/tests/unit/common/ai/utils/test_logging.py @@ -17,6 +17,7 @@ from __future__ import annotations import logging +from decimal import Decimal from unittest.mock import MagicMock from pydantic import BaseModel @@ -34,8 +35,14 @@ ) -def _make_mock_result(model_name="gpt-5", tool_names=None, usage_kwargs=None): - """Build a mock AgentRunResult with usage, response, and messages.""" +def _make_mock_result(model_name="gpt-5", tool_names=None, usage_kwargs=None, cost=None): + """Build a mock AgentRunResult with usage, response, and messages. + + ``cost`` defaults to ``None`` and must be set explicitly -- a MagicMock + attribute left unconfigured returns a new (truthy, ``is not None``) + MagicMock, which would silently push every caller through + ``log_run_summary``'s cost-logging branch. + """ usage_kwargs = usage_kwargs or { "requests": 4, "tool_calls": 3, @@ -44,7 +51,7 @@ def _make_mock_result(model_name="gpt-5", tool_names=None, usage_kwargs=None): "total_tokens": 3359, } result = MagicMock() - result.usage = MagicMock(**usage_kwargs) + result.usage = MagicMock(cost=cost, **usage_kwargs) result.response = MagicMock(model_name=model_name) messages: list = [] @@ -97,6 +104,27 @@ def test_no_tools_skips_sequence_line(self, caplog): assert len(records) == 2 # summary line + endgroup (no tool sequence) assert records[-1].message == "::endgroup::" + def test_cost_none_does_not_log_cost_line(self, caplog): + """cost is None means "unpriceable", not "free" -- no fragment, not a $0 line.""" + logger = logging.getLogger("test.log_run_summary") + result = _make_mock_result(cost=None) + + with caplog.at_level(logging.INFO, logger="test.log_run_summary"): + log_run_summary(logger, result) + + records = [r for r in caplog.records if r.name == "test.log_run_summary"] + assert not any("LLM run cost" in r.message for r in records) + + def test_cost_set_logs_cost_line_with_value(self, caplog): + logger = logging.getLogger("test.log_run_summary") + result = _make_mock_result(cost=Decimal("0.0123")) + + with caplog.at_level(logging.INFO, logger="test.log_run_summary"): + log_run_summary(logger, result) + + records = [r for r in caplog.records if r.name == "test.log_run_summary"] + assert records[1].message == "LLM run cost: $0.0123 (USD, best-effort)" + class TestLogOutputDebug: def test_logs_string_output(self, caplog): diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_usage.py b/providers/common/ai/tests/unit/common/ai/utils/test_usage.py new file mode 100644 index 0000000000000..26c92cd377735 --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/utils/test_usage.py @@ -0,0 +1,108 @@ +# 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. +from __future__ import annotations + +from decimal import Decimal + +import pytest +from pydantic_ai.usage import UsageLimits + +from airflow.providers.common.ai.utils.usage import resolve_usage_limits + + +class TestResolveUsageLimitsNoMaxCost: + def test_max_cost_none_returns_usage_limits_unchanged(self): + """max_cost=None must be a true no-op: the exact same object comes back.""" + limits = UsageLimits(request_limit=3) + assert resolve_usage_limits(limits, None) is limits + + def test_max_cost_none_with_usage_limits_none_returns_none(self): + assert resolve_usage_limits(None, None) is None + + +class TestResolveUsageLimitsWithMaxCost: + def test_max_cost_alone_builds_usage_limits(self): + result = resolve_usage_limits(None, 0.5) + assert isinstance(result, UsageLimits) + assert result.cost_limit == Decimal("0.5") + + def test_max_cost_overrides_existing_cost_limit(self): + limits = UsageLimits(cost_limit=Decimal("999"), request_limit=7) + result = resolve_usage_limits(limits, 0.5) + assert result is not limits + assert result.cost_limit == Decimal("0.5") + assert result.request_limit == 7 + + def test_max_cost_preserves_other_fields(self): + limits = UsageLimits(request_limit=3, input_tokens_limit=4_000, tool_calls_limit=2) + result = resolve_usage_limits(limits, 1.25) + assert result.request_limit == 3 + assert result.input_tokens_limit == 4_000 + assert result.tool_calls_limit == 2 + assert result.cost_limit == Decimal("1.25") + + @pytest.mark.parametrize( + ("max_cost", "expected"), + [ + (0.1, Decimal("0.1")), + ("0.1", Decimal("0.1")), + (1, Decimal("1")), + ("2.50", Decimal("2.50")), + (Decimal("2.50"), Decimal("2.50")), + ], + ) + def test_max_cost_accepts_float_str_and_decimal(self, max_cost, expected): + result = resolve_usage_limits(None, max_cost) + assert result.cost_limit == expected + + def test_max_cost_uses_decimal_str_not_decimal_float(self): + """Decimal(str(x)) must be used, not Decimal(x) -- the latter bakes in + binary-float noise for a value like 0.1.""" + result = resolve_usage_limits(None, 0.1) + assert result.cost_limit == Decimal("0.1") + assert result.cost_limit != Decimal(0.1) + + def test_max_cost_zero_is_accepted(self): + """0 is a valid (if unusual) cap and must not be rejected as falsy or negative.""" + result = resolve_usage_limits(None, 0) + assert result.cost_limit == Decimal("0") + + +class TestResolveUsageLimitsInvalidMaxCost: + @pytest.mark.parametrize("max_cost", ["", "n/a", "$0.50"]) + def test_non_numeric_max_cost_raises_value_error_naming_the_value(self, max_cost): + """The error must name ``max_cost`` and the offending value -- a bare + ``decimal.InvalidOperation`` traceback gives a Dag author no clue which + parameter (often a mistyped or unset Airflow Variable) broke.""" + with pytest.raises(ValueError, match="max_cost") as exc_info: + resolve_usage_limits(None, max_cost) + assert repr(max_cost) in str(exc_info.value) + + def test_negative_max_cost_raises_value_error(self): + with pytest.raises(ValueError, match="max_cost must not be negative"): + resolve_usage_limits(None, -1) + + @pytest.mark.parametrize( + "max_cost", + ["inf", "-inf", "nan", "Infinity", float("inf"), float("nan")], + ) + def test_non_finite_max_cost_raises_value_error_naming_the_value(self, max_cost): + """A non-finite cost_limit would compare as never-exceeded, silently + disabling the cap the Dag author thinks they configured.""" + with pytest.raises(ValueError, match="max_cost") as exc_info: + resolve_usage_limits(None, max_cost) + assert repr(max_cost) in str(exc_info.value) diff --git a/uv.lock b/uv.lock index 9f17a033e820a..e6dd38ecdddba 100644 --- a/uv.lock +++ b/uv.lock @@ -4483,12 +4483,12 @@ requires-dist = [ { name = "pydantic-ai-harness", extras = ["codemode"], marker = "extra == 'code-mode'", specifier = ">=0.3.0" }, { name = "pydantic-ai-shields", marker = "extra == 'shields'", specifier = ">=0.3.4" }, { name = "pydantic-ai-skills", marker = "extra == 'skills'", specifier = ">=1.2.0" }, - { name = "pydantic-ai-slim", specifier = ">=2.0.0" }, - { name = "pydantic-ai-slim", extras = ["anthropic"], marker = "extra == 'anthropic'", specifier = ">=2.0.0" }, - { name = "pydantic-ai-slim", extras = ["bedrock"], marker = "extra == 'bedrock'", specifier = ">=2.0.0" }, - { name = "pydantic-ai-slim", extras = ["google"], marker = "extra == 'google'", specifier = ">=2.0.0" }, - { name = "pydantic-ai-slim", extras = ["mcp"], marker = "extra == 'mcp'", specifier = ">=2.0.0" }, - { name = "pydantic-ai-slim", extras = ["openai"], marker = "extra == 'openai'", specifier = ">=2.0.0" }, + { name = "pydantic-ai-slim", specifier = ">=2.23.0" }, + { name = "pydantic-ai-slim", extras = ["anthropic"], marker = "extra == 'anthropic'", specifier = ">=2.23.0" }, + { name = "pydantic-ai-slim", extras = ["bedrock"], marker = "extra == 'bedrock'", specifier = ">=2.23.0" }, + { name = "pydantic-ai-slim", extras = ["google"], marker = "extra == 'google'", specifier = ">=2.23.0" }, + { name = "pydantic-ai-slim", extras = ["mcp"], marker = "extra == 'mcp'", specifier = ">=2.23.0" }, + { name = "pydantic-ai-slim", extras = ["openai"], marker = "extra == 'openai'", specifier = ">=2.23.0" }, { name = "pypdf", marker = "extra == 'pdf'", specifier = ">=4.0.0" }, { name = "python-docx", marker = "extra == 'docx'", specifier = ">=1.0.0" }, { name = "sqlglot", marker = "extra == 'sql'", specifier = ">=30.0.0" }, @@ -4510,7 +4510,7 @@ dev = [ { name = "llama-index-embeddings-openai", specifier = ">=0.6.0" }, { name = "llama-index-llms-openai", specifier = ">=0.6.0" }, { name = "pydantic-ai-skills", specifier = ">=1.2.0" }, - { name = "pydantic-ai-slim", extras = ["mcp"], specifier = ">=2.0.0" }, + { name = "pydantic-ai-slim", extras = ["mcp"], specifier = ">=2.23.0" }, { name = "sqlglot", specifier = ">=30.0.0" }, ] docs = [{ name = "apache-airflow-devel-common", extras = ["docs"], editable = "devel-common" }] @@ -12470,15 +12470,15 @@ wheels = [ [[package]] name = "genai-prices" -version = "0.0.71" +version = "0.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "httpx2" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/1d/e4/5072862613fba039da2b7c981a8649c6c6bbcb2863bd8bc81617c09ce5ee/genai_prices-0.0.71.tar.gz", hash = "sha256:de4db34ec38404f9ef383cb1ab29e204d16ccf27071af0b16d5747ee7affe36b", size = 82105, upload-time = "2026-07-10T00:38:30.491Z" } +sdist = { url = "https://files.pythonhosted.org/packages/83/9b/85e646305a90a2da18f1edf055498668391e71f9849d3e1754d66559a311/genai_prices-0.1.1.tar.gz", hash = "sha256:54a2237691e0aaefb057d10a0c3c20160accc9fc09521c64c03fcdb7a4a69f68", size = 91182, upload-time = "2026-08-01T09:02:49.552Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/0d/98/c06c1318f6834a26268a2d4280e4183f60c5ea92152aea841feff29826ff/genai_prices-0.0.71-py3-none-any.whl", hash = "sha256:1d13111563af2b1ce43ccfacf77b7ac3216ad704c644408a56e11b181fe0d128", size = 84586, upload-time = "2026-07-10T00:38:29.252Z" }, + { url = "https://files.pythonhosted.org/packages/6e/5e/cfe36dff790ffad6aeff8a069b6f36743987ac17053579035ee0a67635dd/genai_prices-0.1.1-py3-none-any.whl", hash = "sha256:de2e3d8ea3ca1d0d292025995c598da447a74e94f22cd3342df46941aeb5416b", size = 95300, upload-time = "2026-08-01T09:02:48.308Z" }, ] [[package]] @@ -19288,9 +19288,10 @@ wheels = [ [[package]] name = "pydantic-ai-slim" -version = "2.13.0" +version = "2.25.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "anyio" }, { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, { name = "genai-prices" }, { name = "griffelib" }, @@ -19300,9 +19301,9 @@ dependencies = [ { name = "pydantic-graph" }, { name = "typing-inspection" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/cc/cc/aeaef08962f86e4ad76f8f741138590fe9415b3e5085193dbcc621f76518/pydantic_ai_slim-2.13.0.tar.gz", hash = "sha256:3d5f4e5012dc0a4b0e9f76268ce8afd734191b48dbccbf221ac7d030d9e8fcae", size = 839266, upload-time = "2026-07-18T02:55:47.076Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e1/ec/6e186e59a9beede41a9e347f69bd61a833de724aad8a175da8da088d220a/pydantic_ai_slim-2.25.0.tar.gz", hash = "sha256:4f5a36f29e2b346d4b793bf3b983aba17ec19f24015bb811ee40815e98155417", size = 974671, upload-time = "2026-08-06T03:20:33.01Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/77/4e/24f932dfe7c29c23fc875c774acfae4330218362705faa9f590f43ef40da/pydantic_ai_slim-2.13.0-py3-none-any.whl", hash = "sha256:b56b4bf2d5bfcadc0f83a5da16c750b013a2d79cebc2170501ed82435acf39f1", size = 1019269, upload-time = "2026-07-18T02:55:40.017Z" }, + { url = "https://files.pythonhosted.org/packages/04/46/e168f03ec04a933b6b0ff3e02c7e608c24da8efb03d0e1f4d0b486bfef09/pydantic_ai_slim-2.25.0-py3-none-any.whl", hash = "sha256:9b69d1af463a63a88ea3c3567b38a09e8208efe73a36c8f5d5d5515939a88acd", size = 1169288, upload-time = "2026-08-06T03:20:24.91Z" }, ] [package.optional-dependencies] @@ -19454,17 +19455,18 @@ wheels = [ [[package]] name = "pydantic-graph" -version = "2.13.0" +version = "2.25.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "anyio" }, { name = "httpx" }, { name = "logfire-api" }, { name = "pydantic" }, { name = "typing-inspection" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/76/94/ff70ba2c05dddb4c503c6adae5e3f019b01b896e44805bda3b8f84feadfb/pydantic_graph-2.13.0.tar.gz", hash = "sha256:0b77975fef41c993744d06f6e56c49afdbfedb23e8984fe46b2f962dcf90f556", size = 43939, upload-time = "2026-07-18T02:55:49.153Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e9/a0/78b22670f9c9608939a27c9de647dba28b6fca864f8883db30f34832ab9b/pydantic_graph-2.25.0.tar.gz", hash = "sha256:1e1d61556ec0d5fdc02d307380f6ad4ac96d0bba9e5eac0881bae42466d3db8a", size = 45179, upload-time = "2026-08-06T03:20:35.284Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b4/87/c70957d6c519d1cb7375ffdcbe7bf28cf755a2efbf53863c9c4595c91e8e/pydantic_graph-2.13.0-py3-none-any.whl", hash = "sha256:65c51707c36fcc146d2d4130ed26c3186803419b126f72709b77181d0bb0e82e", size = 51656, upload-time = "2026-07-18T02:55:43.041Z" }, + { url = "https://files.pythonhosted.org/packages/d1/45/e3f93b1a33fea3e3989de9491a57a69bd93aae0c5b9137f05113e752a8ef/pydantic_graph-2.25.0-py3-none-any.whl", hash = "sha256:87017851610746f76463b0b1fd257286425f3b4feac1fd17e50b0370ae76c2cf", size = 52661, upload-time = "2026-08-06T03:20:28.332Z" }, ] [[package]]