How can I use an Assistant Agent with a custom model

[AgustinGaliana]

My intention is to use an Assistant Agent with the following custom model,

class CustomClient:
    """
    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.
    """
    def __init__(self, config, **kwargs):
        self.api_url = os.environ.get("URL")
        self.api_key = os.environ.get("API_KEY")
        self.user_email = os.environ.get("USER_EMAIL")
        self.model = config.get("model", os.environ.get("MODEL", "modelo_personalizado"))

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

            message: Message

        choices: List[Choice]
        model: str

    def create(self, params)->ModelClientResponseProtocol:
        text = ''
        for i in range(len(params['messages'])-1):
            text += params['messages'][i+1]['content']
        headers = {
            "Content-Type": "application/json",
            "X-API-KEY": self.api_key,
        }

        data = {
            "model": self.model,
            "uuid": "research-" + os.urandom(8).hex(),
            "message": {
                "role": "user",
                "content": params['messages'][0]['content'],
            },
            "prompt": text,
            "tokens": 4000,
            "temperature": 0,
            "origin": "research",
            "origin_detail": "web",
            "language": "es",
            "user": self.user_email,
        }

        try:
            res = requests.post(url=self.api_url, json=data, headers=headers)
            res.raise_for_status()
            result = res.json()
            # Crear un objeto de respuesta que cumpla con ModelClientResponseProtocol
            response = self.ModelClientResponseProtocol()
            choice = self.ModelClientResponseProtocol()
            message = {}
            # message.content = result["content"]
            message['content'] = result["content"]
            choice.message = message
            response.choices = [choice]
            response.model = self.model
            uso = RequestUsage(prompt_tokens=0,completion_tokens=0)
            if "TERMINATE" in result["content"]:
                print(result["content"])
                return response
            else:
                return response
            

        except requests.exceptions.RequestException as e:
            raise RuntimeError(f"Error en la solicitud HTTP: {e}")

    def message_retrieval(self, response):
        """
        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.
        """
        return [choice.message for choice in response.choices]

    def cost(self, response) -> float:
        return 0.0  # Implementa el cálculo de costos si es necesario

    @staticmethod
    def get_usage(response):
        return {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "cost": 0.0,
            "model": response.model,
        }

Im using now the following call:

async def get_weather(city:str)-> str:
    """Get the weather of a specific city"""
    return f"The weather on the {city} is 72 degrees and sunny"

llm_config = {
    "config_list": [
        {
            "model": os.getenv("MODEL", "modelo_personalizado"),
            "model_client_cls": "CustomClient"
        }
    ]
}
agent = AssistantAgent(
    name="weather_agent",
    tools=[get_weather],
    llm_config=llm_config,
    system_message="You are a helpfull asistant",
    reflect_on_tool_use=True,
    model_client_stream=True,
    ## Usa hasta los ultimos 5 mensajes como memoria
    model_context=BufferedChatCompletionContext(buffer_size=5),
)


agent.register_model_client(model_client_cls=CustomClient)

async def main() -> None:
    await Console(agent.run_stream(task="What is the weather in New York?"))

asyncio.run(main())

This is the issue retrieved:

Traceback (most recent call last):
  File "c:\Users\6003345\Documents\Investigacion_Frameworks\AutoGen\AutoGen\Agustin\tools.py", line 21, in <module>
    agent = AssistantAgent(
            ^^^^^^^^^^^^^^^
  File "C:\Users\6003345\Documents\Investigacion_Frameworks\AutoGen\AutoGen\.venv_1\Lib\site-packages\autogen\agentchat\assistant_agent.py", line 69, in __init__   
    super().__init__(
TypeError: ConversableAgent.__init__() got an unexpected keyword argument 'tools'

¿Any ideas?

[AndreaTang123]

Hi! The issue could be that you’re passing tools to AssistantAgent, but it’s not recognized as a valid parameter for that class in the version you’re using. Maybe you can try to use tool_config instead of tool

agent = AssistantAgent(
    name="weather_agent",
    tool_config={"functions": [get_weather]},
    llm_config=llm_config,
    system_message="You are a helpful assistant",
    reflect_on_tool_use=True,
    model_client_stream=True,
    model_context=BufferedChatCompletionContext(buffer_size=5),
)