diff --git a/python/packages/autogen-agentchat/tests/test_group_chat.py b/python/packages/autogen-agentchat/tests/test_group_chat.py index 6a3c91b805c3..42162af7c67b 100644 --- a/python/packages/autogen-agentchat/tests/test_group_chat.py +++ b/python/packages/autogen-agentchat/tests/test_group_chat.py @@ -238,8 +238,13 @@ async def test_round_robin_group_chat_state() -> None: await team2.load_state(state) state2 = await team2.save_state() assert state == state2 - assert agent3._model_context == agent1._model_context # pyright: ignore - assert agent4._model_context == agent2._model_context # pyright: ignore + + agent1_model_ctx_messages = await agent1._model_context.get_messages() # pyright: ignore + agent2_model_ctx_messages = await agent2._model_context.get_messages() # pyright: ignore + agent3_model_ctx_messages = await agent3._model_context.get_messages() # pyright: ignore + agent4_model_ctx_messages = await agent4._model_context.get_messages() # pyright: ignore + assert agent3_model_ctx_messages == agent1_model_ctx_messages + assert agent4_model_ctx_messages == agent2_model_ctx_messages manager_1 = await team1._runtime.try_get_underlying_agent_instance( # pyright: ignore AgentId("group_chat_manager", team1._team_id), # pyright: ignore RoundRobinGroupChatManager, # pyright: ignore @@ -335,7 +340,7 @@ async def test_round_robin_group_chat_with_tools(monkeypatch: pytest.MonkeyPatch assert result.stop_reason is not None and result.stop_reason == "Text 'TERMINATE' mentioned" # Test streaming. - tool_use_agent._model_context.clear() # pyright: ignore + await tool_use_agent._model_context.clear() # pyright: ignore mock.reset() index = 0 await team.reset() @@ -349,7 +354,7 @@ async def test_round_robin_group_chat_with_tools(monkeypatch: pytest.MonkeyPatch index += 1 # Test Console. - tool_use_agent._model_context.clear() # pyright: ignore + await tool_use_agent._model_context.clear() # pyright: ignore mock.reset() index = 0 await team.reset() @@ -577,8 +582,13 @@ async def test_selector_group_chat_state() -> None: await team2.load_state(state) state2 = await team2.save_state() assert state == state2 - assert agent3._model_context == agent1._model_context # pyright: ignore - assert agent4._model_context == agent2._model_context # pyright: ignore + + agent1_model_ctx_messages = await agent1._model_context.get_messages() # pyright: ignore + agent2_model_ctx_messages = await agent2._model_context.get_messages() # pyright: ignore + agent3_model_ctx_messages = await agent3._model_context.get_messages() # pyright: ignore + agent4_model_ctx_messages = await agent4._model_context.get_messages() # pyright: ignore + assert agent3_model_ctx_messages == agent1_model_ctx_messages + assert agent4_model_ctx_messages == agent2_model_ctx_messages manager_1 = await team1._runtime.try_get_underlying_agent_instance( # pyright: ignore AgentId("group_chat_manager", team1._team_id), # pyright: ignore SelectorGroupChatManager, # pyright: ignore @@ -929,7 +939,7 @@ async def test_swarm_handoff_using_tool_calls(monkeypatch: pytest.MonkeyPatch) - assert result.stop_reason is not None and result.stop_reason == "Text 'TERMINATE' mentioned" # Test streaming. - agent1._model_context.clear() # pyright: ignore + await agent1._model_context.clear() # pyright: ignore mock.reset() index = 0 await team.reset() @@ -942,7 +952,7 @@ async def test_swarm_handoff_using_tool_calls(monkeypatch: pytest.MonkeyPatch) - index += 1 # Test Console - agent1._model_context.clear() # pyright: ignore + await agent1._model_context.clear() # pyright: ignore mock.reset() index = 0 await team.reset()