|
21 | 21 | from strands.agent.conversation_manager.sliding_window_conversation_manager import SlidingWindowConversationManager |
22 | 22 | from strands.agent.state import AgentState |
23 | 23 | from strands.handlers.callback_handler import PrintingCallbackHandler, null_callback_handler |
24 | | -from strands.hooks import BeforeInvocationEvent, BeforeModelCallEvent, BeforeToolCallEvent |
| 24 | +from strands.hooks import ( |
| 25 | + AfterReduceContextEvent, |
| 26 | + BeforeInvocationEvent, |
| 27 | + BeforeModelCallEvent, |
| 28 | + BeforeReduceContextEvent, |
| 29 | + BeforeToolCallEvent, |
| 30 | +) |
25 | 31 | from strands.interrupt import Interrupt |
26 | 32 | from strands.models.bedrock import DEFAULT_BEDROCK_MODEL_ID, BedrockModel |
27 | 33 | from strands.session.repository_session_manager import RepositorySessionManager |
@@ -529,6 +535,78 @@ def test_agent__call__retry_with_reduced_context(mock_model, agent, tool, agener |
529 | 535 | assert conversation_manager_spy.apply_management.call_count == 1 |
530 | 536 |
|
531 | 537 |
|
| 538 | +def test_agent__call__emits_reduce_context_events(mock_model, agent, agenerator): |
| 539 | + """Verify Before/AfterReduceContextEvent fire with correct metadata when overflow triggers reduction.""" |
| 540 | + messages: Messages = [ |
| 541 | + {"role": "user", "content": [{"text": "Hello!"}]}, |
| 542 | + {"role": "assistant", "content": [{"text": "Hi!"}]}, |
| 543 | + {"role": "user", "content": [{"text": "Whats your favorite color?"}]}, |
| 544 | + {"role": "assistant", "content": [{"text": "Blue!"}]}, |
| 545 | + ] |
| 546 | + agent.messages = messages |
| 547 | + |
| 548 | + before_events: list[BeforeReduceContextEvent] = [] |
| 549 | + after_events: list[AfterReduceContextEvent] = [] |
| 550 | + |
| 551 | + agent.hooks.add_callback(BeforeReduceContextEvent, before_events.append) |
| 552 | + agent.hooks.add_callback(AfterReduceContextEvent, after_events.append) |
| 553 | + |
| 554 | + trigger_exception = ContextWindowOverflowException(RuntimeError("Input is too long for requested model")) |
| 555 | + mock_model.mock_stream.side_effect = [ |
| 556 | + trigger_exception, |
| 557 | + agenerator( |
| 558 | + [ |
| 559 | + {"contentBlockStart": {"start": {}}}, |
| 560 | + {"contentBlockDelta": {"delta": {"text": "Green!"}}}, |
| 561 | + {"contentBlockStop": {}}, |
| 562 | + {"messageStop": {"stopReason": "end_turn"}}, |
| 563 | + ] |
| 564 | + ), |
| 565 | + ] |
| 566 | + |
| 567 | + agent("And now?") |
| 568 | + |
| 569 | + assert len(before_events) == 1 |
| 570 | + assert len(after_events) == 1 |
| 571 | + |
| 572 | + before_event = before_events[0] |
| 573 | + assert before_event.agent is agent |
| 574 | + assert before_event.exception is trigger_exception |
| 575 | + # Before reduction runs, the prompt "And now?" has already been appended to messages (5 total). |
| 576 | + assert before_event.message_count == 5 |
| 577 | + |
| 578 | + after_event = after_events[0] |
| 579 | + assert after_event.agent is agent |
| 580 | + assert after_event.exception is trigger_exception |
| 581 | + assert after_event.message_count_before == 5 |
| 582 | + assert after_event.message_count_after < after_event.message_count_before |
| 583 | + assert after_event.messages_removed == after_event.message_count_before - after_event.message_count_after |
| 584 | + assert after_event.messages_removed > 0 |
| 585 | + |
| 586 | + |
| 587 | +def test_agent__call__no_reduce_context_events_on_success(mock_model, agent, agenerator): |
| 588 | + """Verify reduce-context events are NOT fired on a normal successful invocation.""" |
| 589 | + before_events: list[BeforeReduceContextEvent] = [] |
| 590 | + after_events: list[AfterReduceContextEvent] = [] |
| 591 | + |
| 592 | + agent.hooks.add_callback(BeforeReduceContextEvent, before_events.append) |
| 593 | + agent.hooks.add_callback(AfterReduceContextEvent, after_events.append) |
| 594 | + |
| 595 | + mock_model.mock_stream.return_value = agenerator( |
| 596 | + [ |
| 597 | + {"contentBlockStart": {"start": {}}}, |
| 598 | + {"contentBlockDelta": {"delta": {"text": "ok"}}}, |
| 599 | + {"contentBlockStop": {}}, |
| 600 | + {"messageStop": {"stopReason": "end_turn"}}, |
| 601 | + ] |
| 602 | + ) |
| 603 | + |
| 604 | + agent("Hello?") |
| 605 | + |
| 606 | + assert before_events == [] |
| 607 | + assert after_events == [] |
| 608 | + |
| 609 | + |
532 | 610 | def test_agent__call__always_sliding_window_conversation_manager_doesnt_infinite_loop(mock_model, agent, tool): |
533 | 611 | conversation_manager = SlidingWindowConversationManager(window_size=500, should_truncate_results=False) |
534 | 612 | conversation_manager_spy = unittest.mock.Mock(wraps=conversation_manager) |
|
0 commit comments