kaxil commented on code in PR #71463:
URL: https://github.com/apache/airflow/pull/71463#discussion_r3766976194
##########
providers/anthropic/tests/unit/anthropic/operators/test_agent.py:
##########
@@ -260,10 +262,143 @@ def test_budget_is_templated(self):
assert "budget" in AnthropicAgentSessionOperator.template_fields
+class TestUsageXCom:
+ USAGE = {
+ "input_tokens": 827,
+ "output_tokens": 17065,
+ "cache_read_input_tokens": 0,
+ "active_seconds": 91.2,
+ "list_cost": {"amount": "44", "currency": "USD"},
+ }
+
+ @staticmethod
+ def _op(**kwargs):
Review Comment:
Taken, but one step further, because the duplication turned out to be the
real problem.
`TestBudgetParam` and `TestUsageXCom` each had a byte-identical builder --
one named `_make_op` (from your suggestion on #71462) and this one named `_op`.
Renaming this one would have left two copies of the same function in the same
file under two different prefixes.
So it is now a single module-level `_create_op(**kwargs) ->
AnthropicAgentSessionOperator` with the signature you suggested, and both
classes use it. That also settles the prefix: `_create_` matches
`_create_session`, `_create_context` and `_create_session_error` from your
other comments, so the file has one convention rather than two.
Flagging it because it means #71462 changed too -- the helper is defined
there now.
##########
providers/anthropic/src/airflow/providers/anthropic/operators/agent.py:
##########
@@ -178,7 +183,11 @@ def execute(self, context: Context) -> str | None:
)
except Exception:
# send_event failed after create_session allocated the container;
tear it down.
- self._archive_session(session.id)
+ # The event may still have been accepted server-side, so the agent
can already be
+ # spending -- record usage here too rather than leaving this the
one failure path
+ # with no cost trail.
+ archived = self._archive_session(session.id)
+ self._push_usage(context, session.id, session=archived)
Review Comment:
Good catch -- it was four sites, not a few. Extracted `_tear_down(context,
session_id)`, which archives and then records usage; all four failure paths
call it.
Keeping the order inside the helper is the point: teardown runs first
because it is the time-critical call on a failing task, and `sessions.archive`
returns the session carrying its final usage, so recording spend costs no extra
request. Having that in one place means the ordering cannot drift between paths.
--
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]
