Lee-W commented on code in PR #71461:
URL: https://github.com/apache/airflow/pull/71461#discussion_r3773797821
##########
providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py:
##########
@@ -636,12 +696,13 @@ def wait_for_session(
idle event on a ``message`` run (defeats the start race).
:param poll_interval: Seconds to sleep between polls.
:param timeout: Maximum seconds to wait before raising
:class:`AnthropicAgentSessionTimeout`.
+ :raises AnthropicSessionBudgetExceeded: If the session stopped against
its budget.
"""
start = time.monotonic()
consecutive_failures = 0
while True:
try:
- done, error_message = self.poll_session_completion(
+ poll = self.poll_session_completion(
Review Comment:
```suggestion
poll_result = self.poll_session_completion(
```
##########
providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py:
##########
@@ -156,39 +159,77 @@ class TestPollSessionCompletion:
def test_terminated_is_error(self):
hook, client = _make_hook()
client.beta.sessions.retrieve.return_value = _session("terminated")
- done, err = hook.poll_session_completion("s")
- assert done is True
- assert err is not None
+ poll = hook.poll_session_completion("s")
Review Comment:
```suggestion
poll_result = hook.poll_session_completion(session_id="s")
```
or let's try something like
```python
session_id = "s"
...
poll_result = hook.poll_session_completion(self.session_id)
```
it would be easier to understand
##########
providers/anthropic/src/airflow/providers/anthropic/triggers/agent.py:
##########
@@ -94,7 +94,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
while True:
try:
# poll_session_completion does blocking SDK HTTP calls; run
off the event loop.
- done, error_message = await asyncio.to_thread(
+ poll = await asyncio.to_thread(
Review Comment:
```suggestion
poll_result = await asyncio.to_thread(
```
##########
providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py:
##########
@@ -156,39 +159,77 @@ class TestPollSessionCompletion:
def test_terminated_is_error(self):
hook, client = _make_hook()
client.beta.sessions.retrieve.return_value = _session("terminated")
- done, err = hook.poll_session_completion("s")
- assert done is True
- assert err is not None
+ poll = hook.poll_session_completion("s")
+ assert poll.done is True
+ assert poll.error_message is not None
+ assert poll.stop_reason is None
def test_message_end_turn_success(self):
hook, client = _make_hook()
client.beta.sessions.retrieve.return_value = _session("idle")
client.beta.sessions.events.list.return_value =
[_idle_event("end_turn")]
- assert hook.poll_session_completion("s", kickoff_event_id="evt_kick")
== (True, None)
+ assert hook.poll_session_completion("s", kickoff_event_id="evt_kick")
== SessionPollResult(
+ done=True, error_message=None, stop_reason="end_turn"
+ )
@pytest.mark.parametrize("reason", ["requires_action",
"retries_exhausted"])
def test_message_blocked_is_error(self, reason):
hook, client = _make_hook()
client.beta.sessions.retrieve.return_value = _session("idle")
client.beta.sessions.events.list.return_value = [_idle_event(reason)]
- done, err = hook.poll_session_completion("s",
kickoff_event_id="evt_kick")
- assert done is True
- assert err is not None
- assert reason in err
+ poll = hook.poll_session_completion("s", kickoff_event_id="evt_kick")
+ assert poll.done is True
+ assert poll.error_message is not None
+ assert reason in poll.error_message
+ assert poll.stop_reason == reason
def test_message_no_response_yet_not_done(self):
# newest event is our kickoff (agent hasn't responded) -> keep waiting
(start race)
hook, client = _make_hook()
client.beta.sessions.retrieve.return_value = _session("idle")
client.beta.sessions.events.list.return_value =
[mock.MagicMock(type="user.message", id="evt_kick")]
- assert hook.poll_session_completion("s", kickoff_event_id="evt_kick")
== (False, None)
+ assert hook.poll_session_completion("s", kickoff_event_id="evt_kick")
== SessionPollResult(
+ done=False, error_message=None, stop_reason=None
+ )
def test_outcome_satisfied_skips_event_check(self):
hook, client = _make_hook()
client.beta.sessions.retrieve.return_value = _session("idle",
["satisfied"])
- assert hook.poll_session_completion("s", expect_outcome=True) ==
(True, None)
+ assert hook.poll_session_completion("s", expect_outcome=True) ==
SessionPollResult(
+ done=True, error_message=None, stop_reason=None
+ )
client.beta.sessions.events.list.assert_not_called()
+ def test_budget_reached_names_both_causes(self):
+ # A budget stop must not be reported as "configure an autonomous
agent": that advice
+ # is wrong, and the no-list-price cause is invisible without being
named.
+ hook, client = _make_hook()
+ client.beta.sessions.retrieve.return_value = _session("idle")
+ client.beta.sessions.events.list.return_value =
[_idle_event("budget_reached")]
+ poll = hook.poll_session_completion("s", kickoff_event_id="evt_kick")
+ assert poll.done is True
+ assert poll.stop_reason == "budget_reached"
+ assert "budget" in poll.error_message
+ assert "no list price" in poll.error_message
+ assert "autonomous agent" not in poll.error_message
Review Comment:
```suggestion
assert poll == SessionPollResult(
done=True, error_message=ANY, stop_reason="budget_reached"
)
for msg_keyword in ["budget", "no list price", "autonomous agent"]:
assert msg_keyword in poll.error_message
```
--
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]
