Custom Model Client support

autogen/agentchat/contrib/gpt_assistant_agent.py

@@ -56,7 +56,7 @@ def __init__(
         oai_wrapper = OpenAIWrapper(**llm_config)
         if len(oai_wrapper._clients) > 1:
             logger.warning("GPT Assistant only supports one OpenAI client. Using the first client in the list.")
-        self._openai_client = oai_wrapper._clients[0]
+        self._openai_client = oai_wrapper._clients[0]._oai_client
         openai_assistant_id = llm_config.get("assistant_id", None)
         if openai_assistant_id is None:
             # try to find assistant by name first

autogen/agentchat/conversable_agent.py

@@ -9,7 +9,7 @@
 from typing import Any, Awaitable, Callable, Dict, List, Literal, Optional, Tuple, Type, TypeVar, Union
 import warnings
 
-from .. import OpenAIWrapper
+from .. import OpenAIWrapper, ModelClient
 from ..cache.cache import Cache
 from ..code_utils import (
     DEFAULT_MODEL,
@@ -1946,6 +1946,15 @@ def _decorator(func: F) -> F:
 
         return _decorator
 
+    def register_model_client(self, model_client_cls: ModelClient, **kwargs):
+        """Register a model client.
+
+        Args:
+            model_client_cls: A custom client class that follows the Client interface
+            **kwargs: The kwargs for the custom client class to be initialized with
+        """
+        self.client.register_model_client(model_client_cls, **kwargs)
+
     def register_hook(self, hookable_method: Callable, hook: Callable):
         """
         Registers a hook to be called by a hookable method, in order to add a capability to the agent.

autogen/oai/__init__.py

@@ -1,4 +1,4 @@
-from autogen.oai.client import OpenAIWrapper
+from autogen.oai.client import OpenAIWrapper, ModelClient
 from autogen.oai.completion import Completion, ChatCompletion
 from autogen.oai.openai_utils import (
     get_config_list,
@@ -12,6 +12,7 @@
 
 __all__ = [
     "OpenAIWrapper",
+    "ModelClient",
     "Completion",
     "ChatCompletion",
     "get_config_list",