From 11ebd3c42aa4c8a5e2b8a3fdaf36be15510f29a6 Mon Sep 17 00:00:00 2001 From: Davor Runje Date: Fri, 12 Jan 2024 10:37:31 +0000 Subject: [PATCH] added decorator for reply functions --- .gitignore | 1 + autogen/agentchat/conversable_agent.py | 110 +++++++++++++++++++++++ test/agentchat/test_conversable_agent.py | 39 +++++++- 3 files changed, 148 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 145904ed8..26175af00 100644 --- a/.gitignore +++ b/.gitignore @@ -177,3 +177,4 @@ test/test_files/agenteval-in-out/out/ test/agentchat/test_agent_scripts/* test/agentchat/tsp_add_new_point.py autogen/extensions/* +test/executed_openai_notebook_output.txt diff --git a/autogen/agentchat/conversable_agent.py b/autogen/agentchat/conversable_agent.py index 7422d05a3..09adc1855 100644 --- a/autogen/agentchat/conversable_agent.py +++ b/autogen/agentchat/conversable_agent.py @@ -103,6 +103,45 @@ class _Register: """ return self._agent.register_for_llm(name=name, description=description) + def for_reply( + self, + trigger: Union[Type[Agent], str, Agent, Callable[[Agent], bool], List], + *, + position: int = 0, + config: Optional[Any] = None, + reset_config: Optional[Callable] = None, + ): + """Register a reply function. + + The reply function will be called when the trigger matches the sender. + The function registered later will be checked earlier by default. + To change the order, set the position to a positive integer. + + Args: + trigger (Agent class, str, Agent instance, callable, or list): the trigger. + - If a class is provided, the reply function will be called when the sender is an instance of the class. + - If a string is provided, the reply function will be called when the sender's name matches the string. + - If an agent instance is provided, the reply function will be called when the sender is the agent instance. + - If a callable is provided, the reply function will be called when the callable returns True. + - If a list is provided, the reply function will be called when any of the triggers in the list is activated. + - If None is provided, the reply function will be called only when the sender is None. + Note: Be sure to register `None` as a trigger if you would like to trigger an auto-reply function with non-empty messages and `sender=None`. + ``` + position (int): the position of the reply function in the reply function list. + The function registered later will be checked earlier by default. + To change the order, set the position to a positive integer. + config (Any): the config to be passed to the reply function. + When an agent is reset, the config will be reset to the original value. + reset_config (Callable): the function to reset the config. + The function returns None. Signature: ```def reset_config(config: Any)``` + """ + return self._agent._register_for_reply( + trigger=trigger, + position=position, + config=config, + reset_config=reset_config, + ) + class ConversableAgent(Agent): """(In preview) A class for generic conversable agents which can be configured as assistant or user proxy. @@ -289,6 +328,77 @@ class ConversableAgent(Agent): }, ) + def _register_for_reply( + self, + *, + trigger: Union[Type[Agent], str, Agent, Callable[[Agent], bool], List], + position: int = 0, + config: Optional[Any] = None, + reset_config: Optional[Callable] = None, + ) -> Callable[[F], F]: + """Decorator factory for registering a reply function to be used by an agent. + + The reply function will be called when the trigger matches the sender. + The function registered later will be checked earlier by default. + To change the order, set the position to a positive integer. + + Args: + trigger (Agent class, str, Agent instance, callable, or list): the trigger. + - If a class is provided, the reply function will be called when the sender is an instance of the class. + - If a string is provided, the reply function will be called when the sender's name matches the string. + - If an agent instance is provided, the reply function will be called when the sender is the agent instance. + - If a callable is provided, the reply function will be called when the callable returns True. + - If a list is provided, the reply function will be called when any of the triggers in the list is activated. + - If None is provided, the reply function will be called only when the sender is None. + Note: Be sure to register `None` as a trigger if you would like to trigger an auto-reply function with non-empty messages and `sender=None`. + position (int): the position of the reply function in the reply function list. + The function registered later will be checked earlier by default. + To change the order, set the position to a positive integer. + config (Any): the config to be passed to the reply function. + When an agent is reset, the config will be reset to the original value. + reset_config (Callable): the function to reset the config. + The function returns None. Signature: ```def reset_config(config: Any)``` + + Returns: + The decorator for registering a function to be used by an agent. + + Examples: + ``` + agent0 = ConversableAgent("a0", max_consecutive_auto_reply=0, llm_config=False, human_input_mode="NEVER") + agent1 = ConversableAgent("a1", max_consecutive_auto_reply=0, llm_config=False, human_input_mode="NEVER") + @agent0.register.for_reply() + def reply_with_hi(recipient, messages, sender, config): + return (True, "hello") + agent1.initiate_chat(agent, message="hi") + assert agent1.last_message(agent)["content"] == "hello" + ``` + + """ + + def _decorator(reply_func: F) -> F: + """Decorator for registering a function to be used by an agent. + + Args: + reply_func: the function to be registered with the following signature + ```python + def reply_func( + recipient: ConversableAgent, + messages: Optional[List[Dict]] = None, + sender: Optional[Agent] = None, + config: Optional[Any] = None, + ) -> Tuple[bool, Union[str, Dict, None]]: + ``` + + Returns: + The function to be registered as a reply function + + """ + self.register_reply(trigger, reply_func, position, config, reset_config) + + return reply_func + + return _decorator + @property def system_message(self) -> Union[str, List]: """Return the system message.""" diff --git a/test/agentchat/test_conversable_agent.py b/test/agentchat/test_conversable_agent.py index 749f065b9..27a843c2b 100644 --- a/test/agentchat/test_conversable_agent.py +++ b/test/agentchat/test_conversable_agent.py @@ -2,7 +2,7 @@ import asyncio import copy import sys import time -from typing import Any, Callable, Dict, Literal +from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union import unittest import pytest @@ -11,7 +11,8 @@ from pydantic import BaseModel, Field from typing_extensions import Annotated import autogen -from autogen.agentchat import ConversableAgent, UserProxyAgent +from autogen.agentchat import Agent, ConversableAgent, UserProxyAgent + from test_assistant_agent import KEY_LOC, OAI_CONFIG_LIST from conftest import skip_openai @@ -72,6 +73,40 @@ def test_trigger(): pytest.raises(ValueError, agent._match_trigger, 1, agent1) +def test__register_for_reply() -> None: + agent = ConversableAgent("a0", max_consecutive_auto_reply=0, llm_config=False, human_input_mode="NEVER") + agent1 = ConversableAgent("a1", max_consecutive_auto_reply=0, llm_config=False, human_input_mode="NEVER") + + @agent._register_for_reply(trigger=agent1) + def reply_function( + recipient: ConversableAgent, + messages: Optional[List[Dict[str, Any]]], + sender: Optional[Agent], + config: Optional[Any], + ) -> Tuple[bool, Optional[Union[str, Dict[str, Any]]]]: + return (True, "hello") + + agent1.initiate_chat(agent, message="hi") + assert agent1.last_message(agent)["content"] == "hello" + + +def test_register_dot_for_reply() -> None: + agent = ConversableAgent("a0", max_consecutive_auto_reply=0, llm_config=False, human_input_mode="NEVER") + agent1 = ConversableAgent("a1", max_consecutive_auto_reply=0, llm_config=False, human_input_mode="NEVER") + + @agent.register.for_reply(trigger=agent1) + def reply_function( + recipient: ConversableAgent, + messages: Optional[List[Dict[str, Any]]], + sender: Optional[Agent], + config: Optional[Any], + ) -> Tuple[bool, Optional[Union[str, Dict[str, Any]]]]: + return (True, "hello") + + agent1.initiate_chat(agent, message="hi") + assert agent1.last_message(agent)["content"] == "hello" + + def test_context(): agent = ConversableAgent("a0", max_consecutive_auto_reply=0, llm_config=False, human_input_mode="NEVER") agent1 = ConversableAgent("a1", max_consecutive_auto_reply=0, llm_config=False, human_input_mode="NEVER")