Lee-W commented on code in PR #71462:
URL: https://github.com/apache/airflow/pull/71462#discussion_r3764752222
##########
providers/anthropic/docs/operators/anthropic.rst:
##########
@@ -175,15 +177,34 @@ model requests once its tracked list cost reaches the
ceiling:
agent_id="agt_...",
environment_id="env_...",
message="Summarise yesterday's incidents.",
- session_kwargs={
- "budget": {
- "type": "limit",
- # Minor units as an integer string: "2500" is $25.00.
- "max_list_cost": {"amount": "2500", "currency": "USD"},
- }
- },
+ budget=25.00, # $25.00
Review Comment:
While $25,00 means USD, I feel something like `USD$ 25.00` might make this
comment actually helpful.
##########
providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py:
##########
@@ -123,6 +126,49 @@ def is_terminal(cls, status: str) -> bool:
#: ``session.status_idle`` stop reason emitted when a session stops against
its budget.
BUDGET_REACHED = "budget_reached"
+#: What a caller may pass as a session budget: an amount in USD, or the raw
API payload.
Review Comment:
Just wnat ot make sure we want to use `#:` syntax (I know there's already
one above). But this is not that frequently used in airflow yet.
##########
providers/anthropic/tests/unit/anthropic/operators/test_agent.py:
##########
@@ -199,6 +199,67 @@ def test_deferrable_defers_with_trigger(self,
mock_hook_prop):
hook.wait_for_session.assert_not_called()
+class TestBudgetParam:
+ @staticmethod
+ def _op(**kwargs):
Review Comment:
```suggestion
def _make_op(**kwargs):
```
##########
providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py:
##########
@@ -123,6 +126,49 @@ def is_terminal(cls, status: str) -> bool:
#: ``session.status_idle`` stop reason emitted when a session stops against
its budget.
BUDGET_REACHED = "budget_reached"
+#: What a caller may pass as a session budget: an amount in USD, or the raw
API payload.
+BudgetSpec = str | int | float | Decimal | Mapping[str, Any]
+
+
+def build_budget(budget: BudgetSpec) -> dict[str, Any]:
Review Comment:
Should we make it return a TypedDict instead?
##########
providers/anthropic/docs/operators/anthropic.rst:
##########
@@ -147,9 +147,12 @@ Parameters
* ``poll_interval`` — seconds between session status checks.
* ``timeout`` — seconds to wait for a terminal status; defaults to 24 hours.
* ``vault_ids`` — vault IDs providing MCP/credential access to the session.
+* ``budget`` -- spend ceiling for the session, in US dollars (``25.00``) or as
the raw API
+ payload (a mapping). Templated. See `Session budgets`_ below.
Review Comment:
**Templated** here is correct in Airflow context, but one might be surprised
seeing only this word in the middle of other descriptions. might be better to
extend the description of templated a bit
##########
providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py:
##########
@@ -123,6 +126,49 @@ def is_terminal(cls, status: str) -> bool:
#: ``session.status_idle`` stop reason emitted when a session stops against
its budget.
BUDGET_REACHED = "budget_reached"
+#: What a caller may pass as a session budget: an amount in USD, or the raw
API payload.
+BudgetSpec = str | int | float | Decimal | Mapping[str, Any]
+
+
+def build_budget(budget: BudgetSpec) -> dict[str, Any]:
+ """
+ Normalize a session budget into the API's ``max_list_cost`` payload.
+
+ A scalar is read as **US dollars** (``25``, ``25.0`` and ``"25.00"`` all
mean $25.00).
+ The API wants minor units as an integer decimal string, so the conversion
runs through
+ :class:`~decimal.Decimal` (never binary float) and rejects an amount finer
than a cent
+ rather than silently rounding money. A mapping is deep-copied and
otherwise returned
+ unchanged, so a raw payload the provider has not caught up with stays
usable without a
+ provider release.
+
+ .. warning::
+ The ceiling is a stop trigger, not a cap: it is checked between model
requests, so
+ a request already in flight can carry the session well past it.
+ """
+ if isinstance(budget, Mapping):
+ # Deep, not ``dict()``: a shallow copy leaves the nested
``max_list_cost`` aliased
+ # to the caller's object, so a later edit of the returned payload
would reach back
+ # into a templated operator field.
+ return deepcopy(dict(budget))
+ if isinstance(budget, bool):
+ raise ValueError(f"Invalid budget {budget!r}: expected an amount in
USD or a mapping.")
+ try:
+ dollars = Decimal(str(budget))
+ except (InvalidOperation, ValueError) as e:
+ raise ValueError(f"Invalid budget {budget!r}: not a decimal amount in
USD.") from e
Review Comment:
```suggestion
try:
dollars = Decimal(str(budget))
except (InvalidOperation, ValueError) as e:
raise ValueError(f"Invalid budget {budget!r}: not a decimal amount
in USD.") from e
```
Should we just go with this? "True" is not a decimal amount either
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
