如何在AutoGen中使用自定义的大模型

背景

AutoGen原生只支持国外的大模型，如OpenAI, Claude, Mistral等，不支持国内的大模型。但是国内有一些大模型做的还是不错的，尤其是考虑的价格因素之后，国内的大模型性价比很好，我这两天就在想办法集成国内的大模型。

虽然AutoGen不直接支持国内的大模型，但是它支持自定义大模型（custom model）。可以参考这个博客：AutoGen with Custom Models: Empowering Users to Use Their Own Inference Mechanism

但是博客中的案例代码不是很直观，我在这篇博客中记录一下具体怎么接入国内的大模型，并给出案例代码。

自定义模型类

AutoGen允许自定义模型类，只要符合它的协议就行。

具体的协议要求在 autogen.oai.client.ModelClient 中，代码如下：

class ModelClient(Protocol):
    """
    A client class must implement the following methods:
    - create must return a response object that implements the ModelClientResponseProtocol
    - cost must return the cost of the response
    - get_usage must return a dict with the following keys:
        - prompt_tokens
        - completion_tokens
        - total_tokens
        - cost
        - model

    This class is used to create a client that can be used by OpenAIWrapper.
    The response returned from create must adhere to the ModelClientResponseProtocol but can be extended however needed.
    The message_retrieval method must be implemented to return a list of str or a list of messages from the response.
    """

    RESPONSE_USAGE_KEYS = ["prompt_tokens", "completion_tokens", "total_tokens", "cost", "model"]

    class ModelClientResponseProtocol(Protocol):
        class Choice(Protocol):
            class Message(Protocol):
                content: Optional[str]

            message: Message

        choices: List[Choice]
        model: str

    def create(self, params: Dict[str, Any]) -> ModelClientResponseProtocol: ...  # pragma: no cover

    def message_retrieval(
        self, response: ModelClientResponseProtocol
    ) -> Union[List[str], List[ModelClient.ModelClientResponseProtocol.Choice.Message]]:
        """
        Retrieve and return a list of strings or a list of Choice.Message from the response.

        NOTE: if a list of Choice.Message is returned, it currently needs to contain the fields of OpenAI's ChatCompletion Message object,
        since that is expected for function or tool calling in the rest of the codebase at the moment, unless a custom agent is being used.
        """
        ...  # pragma: no cover

    def cost(self, response: ModelClientResponseProtocol) -> float: ...  # pragma: no cover

    @staticmethod
    def get_usage(response: ModelClientResponseProtocol) -> Dict:
        """Return usage summary of the response using RESPONSE_USAGE_KEYS."""
        ...  # pragma: no cover

直白点说，这个协议有四个要求：

1. 自定义的类中有create()函数，并且这个函数的返回应当是ModelClientResponseProtocol的一种实现
2. 要有message_retrieval()函数，用于处理响应，并且返回一个列表，聊表中包含字符串或者message对象
3. 要有cost()函数，返回消耗的费用
4. 要有get_usage()函数，返回一些字典，key应该来自于[“prompt_tokens”, “completion_tokens”, “total_tokens”, “cost”, “model”]。这个主要用于分析，如果不需要分析使用情况，可以反馈空。

实际案例

我在这里使用的UNIAPI(一个大模型代理)托管的claude模型，但是国内的大模型可以完全套用下面的代码。

代码如下：

"""
本代码用于展示如何自定义一个模型，本模型基于UniAPI，
但是任何支持HTTPS调用的大模型都可以套用以下代码
"""

from autogen.agentchat import AssistantAgent, UserProxyAgent
from autogen.oai.openai_utils import config_list_from_json
from types import SimpleNamespace
import requests
import os


class UniAPIModelClient:
    def __init__(self, config, **kwargs):
        print(f"CustomModelClient config: {config}")
        self.api_key = config.get("api_key")
        self.api_url = "https://api.uniapi.me/v1/chat/completions"
        self.model = config.get("model", "gpt-3.5-turbo")
        self.max_tokens = config.get("max_tokens", 1200)
        self.temperature = config.get("temperature", 0.8)
        self.top_p = config.get("top_p", 1)
        self.presence_penalty = config.get("presence_penalty", 1)

        print(f"Initialized CustomModelClient with model {self.model}")

    def create(self, params):
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

        data = {
            "max_tokens": self.max_tokens,
            "model": self.model,
            "temperature": self.temperature,
            "top_p": self.top_p,
            "presence_penalty": self.presence_penalty,
            "messages": params.get("messages", []),
        }

        response = requests.post(self.api_url, headers=headers, json=data)
        response.raise_for_status()  # Raise an exception for HTTP errors

        api_response = response.json()

        # Convert API response to SimpleNamespace for compatibility
        client_response = SimpleNamespace()
        client_response.choices = []
        client_response.model = self.model

        for choice in api_response.get("choices", []):
            client_choice = SimpleNamespace()
            client_choice.message = SimpleNamespace()
            client_choice.message.content = choice.get("message", {}).get("content")
            client_choice.message.function_call = None
            client_response.choices.append(client_choice)

        return client_response

    def message_retrieval(self, response):
        """Retrieve the messages from the response."""
        choices = response.choices
        return [choice.message.content for choice in choices]

    def cost(self, response) -> float:
        """Calculate the cost of the response."""
        # Implement cost calculation if available from your API
        response.cost = 0
        return 0

    @staticmethod
    def get_usage(response):
        # Implement usage tracking if available from your API
        return {}


config_list_custom = config_list_from_json(
    "UNIAPI_CONFIG_LIST.json",
    filter_dict={"model_client_cls": ["UniAPIModelClient"]},
)

assistant = AssistantAgent("assistant", llm_config={"config_list": config_list_custom})
user_proxy = UserProxyAgent(
    "user_proxy",
    code_execution_config={
        "work_dir": "coding",
        "use_docker": False,
    },
)

assistant.register_model_client(model_client_cls=UniAPIModelClient)
user_proxy.initiate_chat(
    assistant,
    message="Write python code to print hello world",
)

如果想要修改为其他模型，唯一的要求是，这个模型支持HTTP调用，然后把 self.api_url = "https://api.uniapi.me/v1/chat/completions" 替换成你自己的值。

在运行上面的案例代码之前，需要创建 UNIAPI_CONFIG_LIST.json 文件，并且可以被程序读取到。其格式如下：

[
    {
        "model": "claude-3-5-sonnet-20240620",
        "api_key": "xxxxxxxxxxxxxxxxxxxxxxxxxxx",
        "temperature": 0.8,
        "max_tokens": 4000,
        "model_client_cls": "UniAPIModelClient"
    }
]

其实这个json本质上就是一个大模型的配置，指定一些必要的参数，其中 model_client_cls 的值要是自定义的模型类的名字，这里不能写错。

以上就是如何在AutoGen使用自定义大模型的全部内容了。

我在这篇博客中只给了具体的案例代码，没有关于更深层次的解读，感兴趣可以阅读官网的文档。

这里想吐槽一下，AutoGen的文档不咋地，不少案例代码都是旧的，没有跟着代码及时更新，有不少坑。