使用AutoGen AgentChat与基于LangChain的自定义客户端和Hugging Face模型

import json
import os
from types import SimpleNamespace

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_huggingface import ChatHuggingFace, HuggingFacePipeline

from autogen import AssistantAgent, UserProxyAgent, config_list_from_json

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) -> ModelClientResponseProtocol:
        ...

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

    def cost(self, response: ModelClientResponseProtocol) -> float:
        ...

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

# custom client with custom model loader


class CustomModelClient:
    """Custom model client implementation for LangChain integration with AutoGen."""

    def __init__(self, config, **kwargs):
        """Initialize the CustomModelClient."""
        print(f"CustomModelClient config: {config}")
        self.device = config.get("device", "cpu")

        gen_config_params = config.get("params", {})
        self.model_name = config["model"]
        pipeline = HuggingFacePipeline.from_model_id(
            model_id=self.model_name,
            task="text-generation",
            pipeline_kwargs=gen_config_params,
            device=self.device,
        )
        self.model = ChatHuggingFace(llm=pipeline)
        print(f"Loaded model {config['model']} to {self.device}")

    def _to_chatml_format(self, message):
        """Convert message to ChatML format."""
        if message["role"] == "system":
            return SystemMessage(content=message["content"])
        if message["role"] == "assistant":
            return AIMessage(content=message["content"])
        if message["role"] == "user":
            return HumanMessage(content=message["content"])
        raise ValueError(f"Unknown message type: {type(message)}")

    def create(self, params):
        """Create a response using the model."""
        if params.get("stream", False) and "messages" in params:
            raise NotImplementedError("Local models do not support streaming.")

        num_of_responses = params.get("n", 1)
        response = SimpleNamespace()
        inputs = [self._to_chatml_format(m) for m in params["messages"]]
        response.choices = []
        response.model = self.model_name

        for _ in range(num_of_responses):
            outputs = self.model.invoke(inputs)
            text = outputs.content
            choice = SimpleNamespace()
            choice.message = SimpleNamespace()
            choice.message.content = text
            choice.message.function_call = None
            response.choices.append(choice)

        return response

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

    def cost(self, response) -> float:
        """Calculate the cost of the response."""
        response.cost = 0
        return 0

    @staticmethod
    def get_usage(response):
        """Get usage statistics."""
        return {}

os.environ["OAI_CONFIG_LIST"] = json.dumps(
    [
        {
            "model": "mistralai/Mistral-7B-Instruct-v0.2",
            "model_client_cls": "CustomModelClient",
            "device": 0,
            "n": 1,
            "params": {
                "max_new_tokens": 500,
                "top_k": 50,
                "temperature": 0.1,
                "do_sample": True,
                "return_full_text": False,
            },
        }
    ]
)

config_list_custom = config_list_from_json(
    "OAI_CONFIG_LIST",
    filter_dict={"model_client_cls": ["CustomModelClient"]},
)

import getpass

from huggingface_hub import login

# The Mistral-7B-Instruct-v0.2 is a gated model which requires API token to access
login(token=getpass.getpass("Enter your HuggingFace API Token"))

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

assistant.register_model_client(model_client_cls=CustomModelClient)

user_proxy.initiate_chat(assistant, message="Write python code to print Hello World!")