mirror of https://github.com/microsoft/autogen.git
147 lines
6.7 KiB
Python
147 lines
6.7 KiB
Python
import pytest
|
|
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
|
|
from autogen_core.model_context import (
|
|
BufferedChatCompletionContext,
|
|
HeadAndTailChatCompletionContext,
|
|
TokenLimitedChatCompletionContext,
|
|
UnboundedChatCompletionContext,
|
|
)
|
|
from autogen_ext.models.openai import OpenAIChatCompletionClient
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_termination_declarative() -> None:
|
|
"""Test that termination conditions can be declared and serialized properly."""
|
|
# 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()
|
|
assert isinstance(max_config, ComponentModel)
|
|
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)
|
|
|
|
# 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)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_completion_context_declarative() -> None:
|
|
unbounded_context = UnboundedChatCompletionContext()
|
|
buffered_context = BufferedChatCompletionContext(buffer_size=5)
|
|
head_tail_context = HeadAndTailChatCompletionContext(head_size=3, tail_size=2)
|
|
model_client = OpenAIChatCompletionClient(model="gpt-4o", api_key="test_key")
|
|
token_limited_context = TokenLimitedChatCompletionContext(model_client=model_client, token_limit=5)
|
|
|
|
# Test serialization
|
|
unbounded_config = unbounded_context.dump_component()
|
|
assert unbounded_config.provider == "autogen_core.model_context.UnboundedChatCompletionContext"
|
|
|
|
buffered_config = buffered_context.dump_component()
|
|
assert buffered_config.provider == "autogen_core.model_context.BufferedChatCompletionContext"
|
|
assert buffered_config.config["buffer_size"] == 5
|
|
|
|
head_tail_config = head_tail_context.dump_component()
|
|
assert head_tail_config.provider == "autogen_core.model_context.HeadAndTailChatCompletionContext"
|
|
assert head_tail_config.config["head_size"] == 3
|
|
assert head_tail_config.config["tail_size"] == 2
|
|
|
|
token_limited_config = token_limited_context.dump_component()
|
|
assert token_limited_config.provider == "autogen_core.model_context.TokenLimitedChatCompletionContext"
|
|
assert token_limited_config.config["token_limit"] == 5
|
|
assert (
|
|
token_limited_config.config["model_client"]["provider"]
|
|
== "autogen_ext.models.openai.OpenAIChatCompletionClient"
|
|
)
|
|
|
|
# Test deserialization
|
|
loaded_unbounded = ComponentLoader.load_component(unbounded_config, UnboundedChatCompletionContext)
|
|
assert isinstance(loaded_unbounded, UnboundedChatCompletionContext)
|
|
|
|
loaded_buffered = ComponentLoader.load_component(buffered_config, BufferedChatCompletionContext)
|
|
|
|
assert isinstance(loaded_buffered, BufferedChatCompletionContext)
|
|
|
|
loaded_head_tail = ComponentLoader.load_component(head_tail_config, HeadAndTailChatCompletionContext)
|
|
|
|
assert isinstance(loaded_head_tail, HeadAndTailChatCompletionContext)
|
|
|
|
loaded_token_limited = ComponentLoader.load_component(token_limited_config, TokenLimitedChatCompletionContext)
|
|
assert isinstance(loaded_token_limited, TokenLimitedChatCompletionContext)
|