Custom Model Client support (#1345)

* add client interface, response protocol, and move code into openai client class

* add ability to register custom client

* tidy up code

* adding checks and errors, and more unit tests

* remove code

* fix error msg

* add use_docer False in notebook

* better error message

* add another example to custom model notebook

* rename and have register_client take model name too

* make Client protocol and remove inheritance

* renames

* update notebook

* add link

* rename and more error checking for registered agents

* adding message retrieval to client protocol for more flexible response

* fix failing openai test

* api_type req made model_client_cls requirement

* notebook cleanup and added blog

* remove raise error if client list is empty - client list will never be empty it will have placeholders

* rename Client -> ModelClient

* add forgotten file

* fix test by fetching internal client

* Update autogen/oai/client.py

Co-authored-by: Eric Zhu <[email protected]>

* don't add retrieval function to cache

* added placeholder cllient class during initial client init, and rewrote registration

* fix spelling

* Update autogen/agentchat/conversable_agent.py

Co-authored-by: Chi Wang <[email protected]>

* type hints, small fixes, docstr comment

* fix api type checking

---------

Co-authored-by: Eric Zhu <[email protected]>
Co-authored-by: Chi Wang <[email protected]>

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",