This is an automated email from the ASF dual-hosted git repository.
eladkal pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 77b21f34c05 Fix Common AI durable retries and tool-approval resumes on
pydantic-ai 2.50+ (#74313)
77b21f34c05 is described below
commit 77b21f34c05f7a8b17c43827327c912605d144cd
Author: Kaxil Naik <[email protected]>
AuthorDate: Tue Oct 6 08:08:56 2026 +0100
Fix Common AI durable retries and tool-approval resumes on pydantic-ai
2.50+ (#74313)
pydantic-ai 2.50 added a float field to RunUsage, audio_seconds. The task
state store loader required every count field to be an int, so the usage
it had saved on the first attempt (audio_seconds=0.0) failed to load on the
retry or resume with "is not an int (got 0.0)". Float fields now load as
floats.
---
.../providers/common/ai/utils/usage_budget.py | 12 ++++++-
.../unit/common/ai/utils/test_usage_budget.py | 42 ++++++++++++++++++++++
2 files changed, 53 insertions(+), 1 deletion(-)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/usage_budget.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/usage_budget.py
index cb8c43e412e..f49fde96385 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/usage_budget.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/usage_budget.py
@@ -72,7 +72,8 @@ def load_run_usage(raw: Any, *, key: str) -> RunUsage:
module still loads. Only the fields ``RunUsage`` currently declares are
read.
:raises ValueError: *raw* is not a dict, or a field has the wrong shape
(``cost``
- not a valid number, a count field not an int, ``details`` not a dict).
The
+ not a valid number, a count field not an int, a float field such as
+ ``audio_seconds`` not a number, ``details`` not a dict). The
message names *key* so the error points at which task state store key
to
delete to reset the budget.
"""
@@ -107,6 +108,15 @@ def load_run_usage(raw: Any, *, key: str) -> RunUsage:
# loaded RunUsage silently mutate the raw dict this was read from
(matters
# most for copy_run_usage's dump/load round trip of a live
RunUsage).
kwargs["details"] = dict(value)
+ elif field.type in (float, "float"):
+ # A measure rather than a count, such as ``audio_seconds``
(pydantic-ai 2.50+).
+ # usage.py uses postponed annotations, so ``field.type`` is the
string "float".
+ if not isinstance(value, (int, float)) or isinstance(value, bool):
+ raise ValueError(
+ f"{key!r}[{field.name!r}] in the task state store is not a
number (got {value!r}); "
+ "delete the key to reset."
+ )
+ kwargs[field.name] = float(value)
else:
if not isinstance(value, int) or isinstance(value, bool):
raise ValueError(
diff --git
a/providers/common/ai/tests/unit/common/ai/utils/test_usage_budget.py
b/providers/common/ai/tests/unit/common/ai/utils/test_usage_budget.py
index c95a2ddd4fa..9b6ee6e850a 100644
--- a/providers/common/ai/tests/unit/common/ai/utils/test_usage_budget.py
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_usage_budget.py
@@ -16,9 +16,11 @@
# under the License.
from __future__ import annotations
+import dataclasses
from decimal import Decimal
import pytest
+from pydantic import TypeAdapter
from pydantic_ai.usage import RunUsage
from airflow.providers.common.ai.utils.usage_budget import (
@@ -32,6 +34,12 @@ from airflow.providers.common.ai.utils.usage_budget import (
from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS
+# Float measures such as ``audio_seconds``, added to RunUsage in pydantic-ai
2.50.
+_FLOAT_FIELDS = [field.name for field in dataclasses.fields(RunUsage) if
field.type in (float, "float")]
+requires_float_field = pytest.mark.skipif(
+ not _FLOAT_FIELDS, reason="RunUsage has no float field before pydantic-ai
2.50"
+)
+
class TestDumpLoadRunUsage:
def test_round_trips_all_fields_including_decimal_cost(self):
@@ -82,6 +90,40 @@ class TestDumpLoadRunUsage:
load_run_usage(raw, key=USAGE_BUDGET_KEY)
assert USAGE_BUDGET_KEY in str(exc_info.value)
+ def test_usage_validated_by_pydantic_round_trips(self):
+ """
+ The tool-approval resume path rebuilds usage with
``TypeAdapter(RunUsage)``, which
+ normalizes every field to its declared type, e.g. a float field's
default ``0`` to
+ ``0.0``. Whatever shape that produces must survive the task state
store round trip.
+ """
+ usage = TypeAdapter(RunUsage).validate_python(
+ dump_run_usage(RunUsage(requests=2, input_tokens=10,
details={"reasoning": 3}))
+ )
+
+ assert copy_run_usage(usage) == usage
+ assert load_run_usage(dump_run_usage(usage), key=USAGE_BUDGET_KEY) ==
usage
+
+ @requires_float_field
+ @pytest.mark.parametrize("value", [0.0, 1.5, 2], ids=["zero", "fraction",
"int"])
+ def test_float_fields_load_as_float(self, value):
+ raw = dump_run_usage(RunUsage(requests=1))
+ raw.update(dict.fromkeys(_FLOAT_FIELDS, value))
+
+ loaded = load_run_usage(raw, key=USAGE_BUDGET_KEY)
+
+ for name in _FLOAT_FIELDS:
+ assert getattr(loaded, name) == value
+ assert isinstance(getattr(loaded, name), float)
+
+ @requires_float_field
+ @pytest.mark.parametrize("value", ["1.5", True], ids=["string", "bool"])
+ def test_float_field_not_a_number_raises_valueerror_naming_the_key(self,
value):
+ raw = {_FLOAT_FIELDS[0]: value}
+
+ with pytest.raises(ValueError, match="not a number") as exc_info:
+ load_run_usage(raw, key=USAGE_BUDGET_KEY)
+ assert USAGE_BUDGET_KEY in str(exc_info.value)
+
class TestCopyRunUsage:
def test_copy_is_independent_of_the_original(self):
