AutoGen with Custom Models: Empowering Personalized Inference

AutoGen now supports custom models, enabling users to define and integrate their own models for a more flexible and personalized inference mechanism. By following a specific protocol, you can incorporate your custom model with AutoGen and respond to prompts using any model, API call, or hardcoded response.

NOTE: Depending on the model used, adjusting the Agent’s default prompts may be necessary.

Quickstart Guide

To get started, follow the interactive notebook that demonstrates loading a local model from HuggingFace into AutoGen and making necessary class modifications.

Step 1: Create the Custom Model Client Class

Begin by creating a model client class adhering to the ModelClient protocol defined in client.py. The new model client class should implement the following methods:

create(): Returns a response object that follows the ModelClientResponseProtocol.
message_retrieval(): Processes the response object and returns a list of strings or message objects.
cost(): Returns the cost of the response.
get_usage(): Returns a dictionary with keys from RESPONSE_USAGE_KEYS (e.g., “prompt_tokens”, “completion_tokens”, “total_tokens”, “cost”, “model”).

Step 2: Add Configuration to the `OAI_CONFIG_LIST`

Ensure that the model_client_cls field is set to the new class name (e.g., "model_client_cls":"CustomModelClient"). Other fields will be forwarded to the class constructor, giving you control over specifying parameters.

Step 3: Register the New Custom Model to the Agent

After adding the configuration with "model_client_cls":"<class name>" to an Agent’s config list, register the corresponding model before initializing the conversation:

python

model_client_cls=CustomModelClient

This must match the entry in the OAI_CONFIG_LIST. If the model client is not registered before chat initialization, an error will occur.

Protocol Details

Your custom model class must adhere to the ModelClient protocol and response structure defined in client.py. This protocol aligns with the OpenAI response structure to ensure resilience to future changes, starting with minimal requirements for easier adoption.

Troubleshooting Steps

If issues arise, follow this checklist:

Ensure your custom model class follows the ModelClient protocol and response protocol.
The create() method must return an inference response that adheres to the ModelClientResponseProtocol.
The message_retrieval() method should return a list of strings or message objects compatible with OpenAI’s ChatCompletion Message object format.
The cost() method should return an integer; return 0 if cost tracking isn’t needed.
The get_usage() method should return a dictionary; return an empty dictionary if usage tracking isn’t needed.
Confirm an entry in the OAI_CONFIG_LIST with the "model_client_cls":"<custom-model-class-name>" field.
Ensure the client is registered with the corresponding config entry and new class:

python

agent.register_model_client(model_client_cls=<class-of-custom-model>, [other optional args])

Verify that all custom models defined in the OAI_CONFIG_LIST are registered.

Additional troubleshooting may be required in your custom code.

Conclusion

With custom model support, AutoGen offers greater flexibility and power for AI applications. Whether you’re using a self-trained model or a specific pre-trained model, AutoGen can meet your needs. Happy coding!

AutoGen with Custom Models: Empowering Personalized Inference