当前位置：首页 > news >正文

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

news 来源：原创 2024/9/20 9:02:30

也可在我的个人博客上查看：https://panzhixiang.cn/2024/autogen-custom-model/

背景

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- modelThis 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: Messagechoices: List[Choice]model: strdef create(self, params: Dict[str, Any]) -> ModelClientResponseProtocol: ...  # pragma: no coverdef 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 coverdef cost(self, response: ModelClientResponseProtocol) -> float: ...  # pragma: no cover@staticmethoddef get_usage(response: ModelClientResponseProtocol) -> Dict:"""Return usage summary of the response using RESPONSE_USAGE_KEYS."""...  # pragma: no cover

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

自定义的类中有create()函数，并且这个函数的返回应当是ModelClientResponseProtocol的一种实现
要有message_retrieval()函数，用于处理响应，并且返回一个列表，聊表中包含字符串或者message对象
要有cost()函数，返回消耗的费用
要有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 osclass 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 errorsapi_response = response.json()# Convert API response to SimpleNamespace for compatibilityclient_response = SimpleNamespace()client_response.choices = []client_response.model = self.modelfor 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 = Noneclient_response.choices.append(client_choice)return client_responsedef message_retrieval(self, response):"""Retrieve the messages from the response."""choices = response.choicesreturn [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 APIresponse.cost = 0return 0@staticmethoddef get_usage(response):# Implement usage tracking if available from your APIreturn {}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"}
]