Lee-W commented on code in PR #71463:
URL: https://github.com/apache/airflow/pull/71463#discussion_r3771802449
##########
providers/anthropic/src/airflow/providers/anthropic/operators/agent.py:
##########
@@ -226,24 +237,83 @@ def execute_complete(self, context: Context, event: Any =
None) -> str:
self.session_id = event["session_id"]
status = event["status"]
if status == "timeout":
- self._archive_session(self.session_id)
+ self._tear_down(context, self.session_id)
raise AnthropicAgentSessionTimeout(event["message"])
if status == "error":
# The trigger yields "error" when polling gives up while the
session may still
# be running; archive it best-effort so its container does not
linger.
- self._archive_session(self.session_id)
+ self._tear_down(context, self.session_id)
raise _create_session_error(event["message"],
event.get("stop_reason"))
self.log.info("Session %s completed.", self.session_id)
+ self._push_usage(context, self.session_id)
return self.session_id
- def _archive_session(self, session_id: str | None) -> None:
- """Best-effort teardown of the server-side session (frees its
container)."""
+ def _tear_down(self, context: Context, session_id: str | None) -> None:
+ """
+ Archive the session and record what it spent, in that order.
+
+ Every failure path needs both. Teardown goes first because it is the
+ time-critical call, and ``sessions.archive`` returns the session
carrying its final
+ usage, so recording spend costs no extra request.
+ """
+ archived = self._archive_session(session_id)
+ self._push_usage(context, session_id, session=archived)
+
+ def _push_usage(self, context: Context, session_id: str | None, session:
Any = None) -> None:
+ """
+ Push the session's token/cost usage to XCom under ``usage``, best
effort.
+
+ Runs on failure as well as success: a session that stopped against its
budget is
+ exactly the one whose spend you want recorded. Never allowed to raise
-- a usage
+ read that fails must not mask the task's real outcome.
+
+ On the failure paths the session has just been archived, and
``sessions.archive``
+ returns the session with its final usage; pass it as ``session`` so
teardown is not
+ delayed by an extra retrieve against an API that may be why the task
is failing.
+ """
if not session_id:
return
+ # The whole body is guarded, not just the API call: this runs on the
failure path
+ # immediately before re-raising, so anything that throws here -- an
unavailable
+ # session, a serialization error, a context without a task instance --
would
+ # replace the exception the task should actually fail with.
+ try:
+ # XCom is cleared at the start of every attempt, so this key only
ever holds
+ # the last one. Stamping the attempt makes that visible rather
than silently
+ # under-reporting total spend across retries. Built as a new dict
rather than
+ # mutating what the hook returned.
+ summary = (
+ self.hook.summarize_usage(session)
+ if session is not None
+ else self.hook.get_session_usage(session_id)
+ )
+ usage = {**summary, "try_number": getattr(context["ti"],
"try_number", None)}
+ context["ti"].xcom_push(key="usage", value=usage)
+ cost = usage.get("list_cost")
+ self.log.info(
+ "Session %s used %s input / %s output tokens; list cost %s",
+ session_id,
+ usage.get("input_tokens"),
+ usage.get("output_tokens"),
+ f"{cost['currency']} {cost['amount']} (minor units)" if cost
else "unavailable",
+ )
+ except Exception as e:
+ self.log.warning("Could not record usage for session %s: %s",
session_id, e)
Review Comment:
```suggestion
self.log.exception("Could not record usage for session %s",
session_id)
```
##########
providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py:
##########
@@ -308,7 +315,100 @@ def test_drops_a_none_budget(self):
assert "budget" not in client.beta.sessions.create.call_args.kwargs
-class TestCreateSessionError:
+class TestGetSessionUsage:
+ """
+ Built on the real SDK usage models rather than mocks.
+
+ Constructing the real model is not on its own a check on field names: it
sets
+ ``extra="allow"``, so a misspelling is accepted as an extra field rather
than rejected.
+ ``test_usage_fields_exist_on_the_sdk_model`` is what pins the names to the
declared
+ schema. Only the session wrapper is a stand-in, since
``get_session_usage`` reads
+ nothing from it but ``.usage``.
+ """
+
+ def test_usage_fields_exist_on_the_sdk_model(self):
+ # The reported breakdown must be fields the SDK actually declares.
active_seconds
+ # lived on session.stats rather than session.usage as recently as
0.117.0, and the
+ # floor is an open-ended >= on a beta API -- if a field moves again,
summarize_usage
+ # would return it as None (or raise, swallowed into a warning) with no
other signal.
+ assert set(BetaManagedAgentsSessionUsage.model_fields) >= {
+ "input_tokens",
+ "output_tokens",
+ "cache_read_input_tokens",
+ "cache_creation",
+ "server_tool_use",
+ "active_seconds",
+ "list_cost",
+ }
+
+ @staticmethod
+ def _create_session(list_cost, cache_creation=None, server_tool_use=None):
+ usage = BetaManagedAgentsSessionUsage(
+ input_tokens=827,
+ output_tokens=17065,
+ cache_read_input_tokens=0,
+ active_seconds=91.2,
+ list_cost=list_cost,
+ cache_creation=cache_creation,
+ server_tool_use=server_tool_use,
+ )
+ return SimpleNamespace(usage=usage)
+
+ def test_flattens_to_json_safe_scalars(self):
+ hook, client = _make_hook()
+ client.beta.sessions.retrieve.return_value = self._create_session(
+ BetaMonetaryAmount(amount="44", currency="USD")
+ )
+ assert hook.get_session_usage("s") == {
Review Comment:
should we make "s" a `session_id` variable? I was confused first seeing "s"
and woundering where this "s" came from
##########
providers/anthropic/src/airflow/providers/anthropic/operators/agent.py:
##########
@@ -226,24 +237,83 @@ def execute_complete(self, context: Context, event: Any =
None) -> str:
self.session_id = event["session_id"]
status = event["status"]
if status == "timeout":
- self._archive_session(self.session_id)
+ self._tear_down(context, self.session_id)
raise AnthropicAgentSessionTimeout(event["message"])
if status == "error":
# The trigger yields "error" when polling gives up while the
session may still
# be running; archive it best-effort so its container does not
linger.
- self._archive_session(self.session_id)
+ self._tear_down(context, self.session_id)
raise _create_session_error(event["message"],
event.get("stop_reason"))
self.log.info("Session %s completed.", self.session_id)
+ self._push_usage(context, self.session_id)
return self.session_id
- def _archive_session(self, session_id: str | None) -> None:
- """Best-effort teardown of the server-side session (frees its
container)."""
+ def _tear_down(self, context: Context, session_id: str | None) -> None:
+ """
+ Archive the session and record what it spent, in that order.
+
+ Every failure path needs both. Teardown goes first because it is the
+ time-critical call, and ``sessions.archive`` returns the session
carrying its final
+ usage, so recording spend costs no extra request.
+ """
+ archived = self._archive_session(session_id)
+ self._push_usage(context, session_id, session=archived)
+
+ def _push_usage(self, context: Context, session_id: str | None, session:
Any = None) -> None:
+ """
+ Push the session's token/cost usage to XCom under ``usage``, best
effort.
+
+ Runs on failure as well as success: a session that stopped against its
budget is
+ exactly the one whose spend you want recorded. Never allowed to raise
-- a usage
+ read that fails must not mask the task's real outcome.
+
+ On the failure paths the session has just been archived, and
``sessions.archive``
+ returns the session with its final usage; pass it as ``session`` so
teardown is not
+ delayed by an extra retrieve against an API that may be why the task
is failing.
+ """
if not session_id:
return
+ # The whole body is guarded, not just the API call: this runs on the
failure path
+ # immediately before re-raising, so anything that throws here -- an
unavailable
+ # session, a serialization error, a context without a task instance --
would
+ # replace the exception the task should actually fail with.
+ try:
+ # XCom is cleared at the start of every attempt, so this key only
ever holds
+ # the last one. Stamping the attempt makes that visible rather
than silently
+ # under-reporting total spend across retries. Built as a new dict
rather than
+ # mutating what the hook returned.
+ summary = (
+ self.hook.summarize_usage(session)
+ if session is not None
+ else self.hook.get_session_usage(session_id)
+ )
+ usage = {**summary, "try_number": getattr(context["ti"],
"try_number", None)}
+ context["ti"].xcom_push(key="usage", value=usage)
+ cost = usage.get("list_cost")
+ self.log.info(
+ "Session %s used %s input / %s output tokens; list cost %s",
+ session_id,
+ usage.get("input_tokens"),
+ usage.get("output_tokens"),
+ f"{cost['currency']} {cost['amount']} (minor units)" if cost
else "unavailable",
+ )
+ except Exception as e:
+ self.log.warning("Could not record usage for session %s: %s",
session_id, e)
+
+ def _archive_session(self, session_id: str | None) -> Any:
Review Comment:
```suggestion
def _archive_session(self, session_id: str | None) ->
BetaManagedAgentsSession | None:
```
--
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]
