@@ -2061,6 +2061,7 @@ def test_manager_resume_messages():
2061
2061
with pytest .raises (Exception ):
2062
2062
return_agent , return_message = manager .resume (messages = "Let's get this conversation started." )
2063
2063
2064
+
2064
2065
def test_custom_model_client ():
2065
2066
class CustomModelClient :
2066
2067
def __init__ (self , config , ** kwargs ):
@@ -2097,18 +2098,24 @@ def get_usage(response):
2097
2098
llm_config = {"config_list" : [{"model" : "test_model_name" , "model_client_cls" : "CustomModelClient" }]}
2098
2099
2099
2100
group_chat = autogen .GroupChat (
2100
- agents = [], messages = [], max_round = 3 , select_speaker_auto_llm_config = llm_config ,
2101
- select_speaker_auto_model_client_cls = CustomModelClient
2101
+ agents = [],
2102
+ messages = [],
2103
+ max_round = 3 ,
2104
+ select_speaker_auto_llm_config = llm_config ,
2105
+ select_speaker_auto_model_client_cls = CustomModelClient ,
2102
2106
)
2103
2107
2104
- checking_agent , speaker_selection_agent = group_chat ._create_internal_agents (agents = [], messages = [], max_attempts = 3 , validate_speaker_name = (True , "test" ))
2108
+ checking_agent , speaker_selection_agent = group_chat ._create_internal_agents (
2109
+ agents = [], messages = [], max_attempts = 3 , validate_speaker_name = (True , "test" )
2110
+ )
2105
2111
2106
2112
# Check that the custom model client is assigned to the speaker selection agent
2107
2113
assert isinstance (speaker_selection_agent .client ._clients [0 ], CustomModelClient )
2108
2114
2109
2115
# Check that the LLM Config is assigned
2110
2116
assert speaker_selection_agent .client ._config_list == llm_config ["config_list" ]
2111
2117
2118
+
2112
2119
def test_select_speaker_transform_messages ():
2113
2120
"""Tests adding transform messages to a GroupChat for speaker selection when in 'auto' mode"""
2114
2121
@@ -2148,6 +2155,7 @@ def test_select_speaker_transform_messages():
2148
2155
2149
2156
assert groupchat_none ._speaker_selection_transforms is None
2150
2157
2158
+
2151
2159
if __name__ == "__main__" :
2152
2160
# test_func_call_groupchat()
2153
2161
# test_broadcast()
0 commit comments