Fix #2643 - groupchat model registration

autogen/agentchat/groupchat.py

@@ -7,14 +7,14 @@
 from dataclasses import dataclass, field
 from typing import Callable, Dict, List, Literal, Optional, Tuple, Union
 
+from ..oai.client import ModelClient
 from ..code_utils import content_str
 from ..exception_utils import AgentNameConflict, NoEligibleSpeaker, UndefinedNextAgent
 from ..formatting_utils import colored
 from ..graph_utils import check_graph_validity, invert_disallowed_to_allowed
 from ..io.base import IOStream
 from ..runtime_logging import log_new_agent, logging_enabled
 from .agent import Agent
-from .chat import ChatResult
 from .conversable_agent import ConversableAgent
 
 logger = logging.getLogger(__name__)
@@ -101,6 +101,7 @@ def custom_speaker_selection_func(
 
     agents: List[Agent]
     messages: List[Dict]
+    model_client_cls: Optional[Union[ModelClient, List[ModelClient]]] = None
     max_round: Optional[int] = 10
     admin_name: Optional[str] = "Admin"
     func_call_filter: Optional[bool] = True
@@ -561,6 +562,39 @@ def _finalize_speaker(self, last_speaker: Agent, final: bool, name: str, agents:
         agent = self.agent_by_name(name)
         return agent if agent else self.next_agent(last_speaker, agents)
 
+    def _register_custom_model_clients(self, agent: ConversableAgent, selector: ConversableAgent):
+        for config in selector.llm_config['config_list']:
+            model_client_cls_to_match = config.get('model_client_cls')
+            if model_client_cls_to_match:
+                if not self.model_client_cls:
+                    raise ValueError(
+                        "A custom model was detected in the config but no 'model_client_cls' "
+                        "was supplied for registration in GroupChat."
+                    )
+
+                if isinstance(self.model_client_cls, list):
+                    # Register the first custom model client class matching the name specified in the config
+                    matching_model_cls = [
+                        client_cls
+                        for client_cls in self.model_client_cls
+                        if client_cls.__name__ == model_client_cls_to_match
+                    ]
+                    if len(set(matching_model_cls)) > 1:
+                        raise RuntimeError(
+                            f"More than one unique 'model_client_cls' with __name__ '{model_client_cls_to_match}'."
+                        )
+                    if not matching_model_cls:
+                        raise ValueError(
+                            "No model's __name__ matches the model client class "
+                            f"'{model_client_cls_to_match}' specified in llm_config."
+                        )
+                    model_client_cls = matching_model_cls[0]
+                else:
+                    # Register the only custom model client
+                    model_client_cls = self.model_client_cls
+
+                agent.register_model_client(model_client_cls)
+
     def _auto_select_speaker(
         self,
         last_speaker: Agent,
@@ -615,7 +649,14 @@ def validate_speaker_name(recipient, messages, sender, config) -> Tuple[bool, Un
         # Two-agent chat for speaker selection
 
         # Agent for checking the response from the speaker_select_agent
-        checking_agent = ConversableAgent("checking_agent", default_auto_reply=max_attempts)
+        checking_agent = ConversableAgent(
+            "checking_agent",
+            default_auto_reply=max_attempts,
+            llm_config=selector.llm_config
+        )
+
+        # Register any custom model passed in llm_config with the checking_agent
+        self._register_custom_model_clients(checking_agent, selector)
 
         # Register the speaker validation function with the checking agent
         checking_agent.register_reply(
@@ -633,6 +674,9 @@ def validate_speaker_name(recipient, messages, sender, config) -> Tuple[bool, Un
             human_input_mode="NEVER",  # Suppresses some extra terminal outputs, outputs will be handled by select_speaker_auto_verbose
         )
 
+        # Register any custom model passed in llm_config with the speaker_selection_agent
+        self._register_custom_model_clients(speaker_selection_agent, selector)
+
         # Run the speaker selection chat
         result = checking_agent.initiate_chat(
             speaker_selection_agent,
@@ -702,7 +746,14 @@ def validate_speaker_name(recipient, messages, sender, config) -> Tuple[bool, Un
         # Two-agent chat for speaker selection
 
         # Agent for checking the response from the speaker_select_agent
-        checking_agent = ConversableAgent("checking_agent", default_auto_reply=max_attempts)
+        checking_agent = ConversableAgent(
+            "checking_agent",
+            default_auto_reply=max_attempts,
+            llm_config=selector.llm_config
+        )
+
+        # Register any custom model passed in llm_config with the checking_agent
+        self._register_custom_model_clients(checking_agent, selector)
 
         # Register the speaker validation function with the checking agent
         checking_agent.register_reply(
@@ -720,6 +771,9 @@ def validate_speaker_name(recipient, messages, sender, config) -> Tuple[bool, Un
             human_input_mode="NEVER",  # Suppresses some extra terminal outputs, outputs will be handled by select_speaker_auto_verbose
         )
 
+        # Register any custom model passed in llm_config with the speaker_selection_agent
+        self._register_custom_model_clients(speaker_selection_agent, selector)
+
         # Run the speaker selection chat
         result = await checking_agent.a_initiate_chat(
             speaker_selection_agent,

website/docs/topics/groupchat/using_custom_models.md

@@ -0,0 +1,79 @@
+# Using Custom Models
+
+When using `GroupChatManager` we need to pass a `GroupChat` object in the constructor, a dataclass responsible for 
+gathering agents, preparing messages from prompt templates and selecting speakers 
+(eventually using `speaker_selection_method` as described [here](customized_speaker_selection.mdx)).
+
+To do so GroupChat internally initializes two instances of ConversableAgent.
+In order to control the model clients used by the agents instantiated within the GroupChat, which already receives the 
+`llm_config` passed to GroupChatManager, the optional `model_client_cls` attribute can be set.
+
+
+## Example
+First we need to define an `llm_config` and define some agents that will partake in the group chat:
+```python
+from autogen import GroupChat, ConversableAgent, GroupChatManager, UserProxyAgent
+from somewhere import MyModelClient
+
+
+# Define the custom model configuration
+llm_config = {
+    "config_list": [
+        {
+            "model": "gpt-3.5-turbo",
+            "model_client_cls": "MyModelClient"
+        }
+    ]
+}
+
+# Initialize the agents with the custom model
+agent1 = ConversableAgent(
+    name="Agent 1",
+    llm_config=llm_config
+)
+agent1.register_model_client(model_client_cls=MyModelClient)
+
+agent2 = ConversableAgent(
+    name="Agent 2",
+    llm_config=llm_config
+)
+agent2.register_model_client(model_client_cls=MyModelClient)
+
+agent3 = ConversableAgent(
+    name="Agent 2",
+    llm_config=llm_config
+)
+agent3.register_model_client(model_client_cls=MyModelClient)
+
+user_proxy = UserProxyAgent(name="user", llm_config=llm_config, code_execution_config={"use_docker": False})
+user_proxy.register_model_client(MyModelClient)
+```
+
+Note that the agents definition illustrated here is minimal and might not suit your needs. The only aim is to show a 
+basic setup for a group chat scenario.
+
+We then create a `GroupChat` and, if we want the underlying agents used by GroupChat to use our
+ custom client, we will pass it in the `model_client_cls` attribute. 
+
+Finally we create an instance of `GroupChatManager` and pass the config to it. This same config will be forwarded to 
+the GroupChat, that (if needed) will automatically handle registration of custom models only.
+
+```python
+# Create a GroupChat instance and add the agents
+group_chat = GroupChat(agents=[agent1, agent2, agent3], messages=[], model_client_cls=MyModelClient)
+
+# Create the GroupChatManager with the GroupChat, UserProxy, and model configuration
+chat_manager = GroupChatManager(groupchat=group_chat, llm_config=llm_config)
+chat_manager.register_model_client(model_client_cls=MyModelClient)
+
+# Initiate the chat using the UserProxy
+user_proxy.initiate_chat(chat_manager, initial_message="Suggest me the most trending papers in microbiology that you think might interest me")
+
+```
+
+This attribute can either be a class or a list of classes which adheres to the `ModelClient` protocol (see 
+[this link](../non-openai-models/about-using-nonopenai-models.md) for more info about defining a custom model client 
+class).
+
+Note that it is not necessary to define a `model_client_cls` when working with Azure OpenAI, OpenAI or other non-custom 
+models natively supported by the library.