diff --git a/autogen/agentchat/conversable_agent.py b/autogen/agentchat/conversable_agent.py index a3726678a45..bf10f465358 100644 --- a/autogen/agentchat/conversable_agent.py +++ b/autogen/agentchat/conversable_agent.py @@ -2039,15 +2039,18 @@ async def a_generate_init_message(self, **context) -> Union[str, Dict]: self._process_carryover(context) return context["message"] - def register_function(self, function_map: Dict[str, Callable]): + def register_function(self, function_map: Dict[str, Union[Callable, None]]): """Register functions to the agent. Args: - function_map: a dictionary mapping function names to functions. + function_map: a dictionary mapping function names to functions. if function_map[name] is None, the function will be removed from the function_map. """ - for name in function_map.keys(): + for name, func in function_map.items(): self._assert_valid_name(name) - self._function_map.update(function_map) + if func is None and name not in self._function_map.keys(): + warnings.warn(f"The function {name} to remove doesn't exist", name) + + self._function_map = {name: func for name, func in function_map.items() if func is not None} def update_function_signature(self, func_sig: Union[str, Dict], is_remove: None): """update a function_signature in the LLM configuration for function_call. diff --git a/test/agentchat/test_conversable_agent.py b/test/agentchat/test_conversable_agent.py index b71a6341c87..bc7546b2ef3 100644 --- a/test/agentchat/test_conversable_agent.py +++ b/test/agentchat/test_conversable_agent.py @@ -559,6 +559,16 @@ def exec_sh(script: str) -> None: assert agent.function_map["python"] == exec_python assert agent.function_map["sh"] == exec_sh + # remove the functions + agent.register_function( + function_map={ + "python": None, + } + ) + + assert set(agent.function_map.keys()) == {"sh"} + assert agent.function_map["sh"] == exec_sh + def test__wrap_function_sync(): CurrencySymbol = Literal["USD", "EUR"]