Lee-W commented on code in PR #71463:
URL: https://github.com/apache/airflow/pull/71463#discussion_r3765412048
##########
providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py:
##########
@@ -646,6 +646,40 @@ def get_session(self, session_id: str) ->
BetaManagedAgentsSession:
self._require_first_party("Managed Agents")
return self._first_party_conn.beta.sessions.retrieve(session_id)
+ def get_session_usage(self, session_id: str) -> dict[str, Any]:
Review Comment:
should we make it a property instead?
##########
providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py:
##########
@@ -308,6 +315,94 @@ def test_drops_a_none_budget(self):
assert "budget" not in client.beta.sessions.create.call_args.kwargs
+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 _session(list_cost, cache_creation=None, server_tool_use=None):
Review Comment:
```suggestion
def _create_session(list_cost, cache_creation=None,
server_tool_use=None):
```
##########
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:
should we make this a method? it appears a few times already.
##########
providers/anthropic/tests/unit/anthropic/operators/test_agent.py:
##########
@@ -98,7 +100,7 @@ def test_message_sends_user_message_and_waits(self,
mock_hook_prop):
"sess_1", {"type": "user.message", "content": [{"type": "text",
"text": "summarize"}]}
)
hook.wait_for_session.assert_called_once()
- context["ti"].xcom_push.assert_called_once_with(key="session_id",
value="sess_1")
+ context["ti"].xcom_push.assert_any_call(key="session_id",
value="sess_1")
Review Comment:
why do we make this change
##########
providers/anthropic/tests/unit/anthropic/operators/test_agent.py:
##########
@@ -34,8 +34,10 @@
pytest.importorskip("anthropic")
-def _context():
- return {"ti": mock.MagicMock()}
+def _context(try_number=1):
Review Comment:
```suggestion
def _create_context(try_number=1):
```
##########
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:
```suggestion
def _create_op(**kwargs) -> AnthropicAgentSessionOperator:
```
##########
providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py:
##########
@@ -308,6 +315,94 @@ def test_drops_a_none_budget(self):
assert "budget" not in client.beta.sessions.create.call_args.kwargs
+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 _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._session(
+ BetaMonetaryAmount(amount="44", currency="USD")
+ )
+ assert hook.get_session_usage("s") == {
+ "input_tokens": 827,
+ "output_tokens": 17065,
+ "cache_read_input_tokens": 0,
+ "cache_creation": None,
+ "server_tool_use": None,
+ "active_seconds": 91.2,
+ "list_cost": {"amount": "44", "currency": "USD"},
+ }
+
+ def test_reports_every_billable_dimension(self):
+ # list_cost is None exactly when a caller must price the run from
usage, so the
+ # expensive side (cache writes, server tool calls) cannot be missing.
+ hook, client = _make_hook()
+ client.beta.sessions.retrieve.return_value = self._session(
+ None,
+ cache_creation=BetaManagedAgentsCacheCreationUsage(
+ ephemeral_5m_input_tokens=120, ephemeral_1h_input_tokens=340
+ ),
+
server_tool_use=BetaManagedAgentsServerToolUsage(web_search_requests=3,
web_fetch_requests=5),
+ )
+ usage = hook.get_session_usage("s")
Review Comment:
should we just assert the whole usage dict insetad?
--
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]
