From f431b73e95c80dddf2737ef29f079fcb25bb0b0a Mon Sep 17 00:00:00 2001 From: Victor Dibia Date: Fri, 10 Jan 2025 18:46:51 -0800 Subject: [PATCH] update tests --- .../tests/test_declarative_components.py | 105 ++++++++++++------ 1 file changed, 74 insertions(+), 31 deletions(-) diff --git a/python/packages/autogen-agentchat/tests/test_declarative_components.py b/python/packages/autogen-agentchat/tests/test_declarative_components.py index 0bf0aa1bfa10..35cf54f86416 100644 --- a/python/packages/autogen-agentchat/tests/test_declarative_components.py +++ b/python/packages/autogen-agentchat/tests/test_declarative_components.py @@ -1,7 +1,15 @@ import pytest -from autogen_agentchat.base import OrTerminationCondition -from autogen_agentchat.conditions import MaxMessageTermination, StopMessageTermination -from autogen_agentchat.messages import StopMessage +from autogen_agentchat.base import AndTerminationCondition +from autogen_agentchat.conditions import ( + ExternalTermination, + HandoffTermination, + MaxMessageTermination, + SourceMatchTermination, + StopMessageTermination, + TextMentionTermination, + TimeoutTermination, + TokenUsageTermination, +) from autogen_core import ComponentLoader, ComponentModel @@ -11,6 +19,12 @@ async def test_termination_declarative() -> None: # Create basic termination conditions max_term = MaxMessageTermination(5) stop_term = StopMessageTermination() + text_term = TextMentionTermination("stop") + token_term = TokenUsageTermination(max_total_token=100, max_prompt_token=50, max_completion_token=100) + handoff_term = HandoffTermination(target="human") + timeout_term = TimeoutTermination(timeout_seconds=30) + external_term = ExternalTermination() + source_term = SourceMatchTermination(sources=["human"]) # Test basic serialization max_config = max_term.dump_component() @@ -18,34 +32,63 @@ async def test_termination_declarative() -> None: assert max_config.provider == "autogen_agentchat.conditions.MaxMessageTermination" assert max_config.config.get("max_messages") == 5 + # Test serialization of new conditions + text_config = text_term.dump_component() + assert text_config.provider == "autogen_agentchat.conditions.TextMentionTermination" + assert text_config.config.get("text") == "stop" + + token_config = token_term.dump_component() + assert token_config.provider == "autogen_agentchat.conditions.TokenUsageTermination" + assert token_config.config.get("max_total_token") == 100 + + handoff_config = handoff_term.dump_component() + assert handoff_config.provider == "autogen_agentchat.conditions.HandoffTermination" + assert handoff_config.config.get("target") == "human" + + timeout_config = timeout_term.dump_component() + assert timeout_config.provider == "autogen_agentchat.conditions.TimeoutTermination" + assert timeout_config.config.get("timeout_seconds") == 30 + + external_config = external_term.dump_component() + assert external_config.provider == "autogen_agentchat.conditions.ExternalTermination" + + source_config = source_term.dump_component() + assert source_config.provider == "autogen_agentchat.conditions.SourceMatchTermination" + assert source_config.config.get("sources") == ["human"] + # Test basic deserialization loaded_max = ComponentLoader.load_component(max_config, MaxMessageTermination) assert isinstance(loaded_max, MaxMessageTermination) - # Use public interface to verify state - messages = [StopMessage(content="msg", source="test") for _ in range(5)] - result = await loaded_max(messages) - assert isinstance(result, StopMessage) - - # Test composition and complex serialization - or_term = max_term | stop_term - or_config = or_term.dump_component() - assert or_config.provider == "autogen_agentchat.base.OrTerminationCondition" - assert len(or_config.config["conditions"]) == 2 - - # Verify nested conditions are correctly serialized - conditions = or_config.config["conditions"] - assert conditions[0]["provider"] == "autogen_agentchat.conditions.MaxMessageTermination" - assert conditions[1]["provider"] == "autogen_agentchat.conditions.StopMessageTermination" - - # Test behavior of loaded composite condition - loaded_or = OrTerminationCondition.load_component(or_config) - - # Test with stop message - stop_messages = [StopMessage(content="stop", source="test")] - result = await loaded_or(stop_messages) - assert isinstance(result, StopMessage) - assert loaded_or.terminated - - # Test reset functionality - await loaded_or.reset() - assert not loaded_or.terminated + + # Test deserialization of new conditions + loaded_text = ComponentLoader.load_component(text_config, TextMentionTermination) + assert isinstance(loaded_text, TextMentionTermination) + + loaded_token = ComponentLoader.load_component(token_config, TokenUsageTermination) + assert isinstance(loaded_token, TokenUsageTermination) + + loaded_handoff = ComponentLoader.load_component(handoff_config, HandoffTermination) + assert isinstance(loaded_handoff, HandoffTermination) + + loaded_timeout = ComponentLoader.load_component(timeout_config, TimeoutTermination) + assert isinstance(loaded_timeout, TimeoutTermination) + + loaded_external = ComponentLoader.load_component(external_config, ExternalTermination) + assert isinstance(loaded_external, ExternalTermination) + + loaded_source = ComponentLoader.load_component(source_config, SourceMatchTermination) + assert isinstance(loaded_source, SourceMatchTermination) + + # Test composition with new conditions + composite_term = (max_term | stop_term) & (token_term | handoff_term) + composite_config = composite_term.dump_component() + + assert composite_config.provider == "autogen_agentchat.base.AndTerminationCondition" + conditions = composite_config.config["conditions"] + assert len(conditions) == 2 + assert conditions[0]["provider"] == "autogen_agentchat.base.OrTerminationCondition" + assert conditions[1]["provider"] == "autogen_agentchat.base.OrTerminationCondition" + + # Test loading complex composition + loaded_composite = ComponentLoader.load_component(composite_config) + assert isinstance(loaded_composite, AndTerminationCondition)