Skip to content

Commit

Permalink
fix
Browse files Browse the repository at this point in the history
  • Loading branch information
“skzhang1” committed Feb 26, 2024
1 parent 07520aa commit 42cb1fb
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 4 deletions.
11 changes: 7 additions & 4 deletions autogen/agentchat/conversable_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
10 changes: 10 additions & 0 deletions test/agentchat/test_conversable_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down

0 comments on commit 42cb1fb

Please sign in to comment.