Swarm: User-defined AfterWork string for agent selection using LLM

[marklysze]

Why are these changes needed?

In line with #26, this provides the developer with an option to specify a string prompt for an LLM to use to select the next agent in an agent's AfterWork hand off.

It extends the AfterWorkOption.SWARM_MANAGER (which, by default, uses the default auto speaker selection options) with an additional AfterWork parameter, next_agent_selection_msg, which becomes the selection prompt used by the LLM.

next_agent_selection_msg can be a string or a Callable that returns a string.

The signature for the Callable is:
def my_selection_message(agent: ConversableAgent, messages: List[Dict[str, Any]]) -> str

The string provided, or returned from the Callable, will have the following string substitutions:

• {agentlist} will be replaced by a comma-delimited list of agents
• context variable keys will be replaced by their values

The underlying groupchat's GroupChatManager will be used for the speaker selection and must, therefore, have an LLM Config. This can be set through the new swarm_manager_args parameter on a/initiate_swarm_chat.

Example:

context_variables = {
    "order_id": "ABC1234"
}

register_handoff(
    from=planner,
    hand_to=[
        AfterWork(agent=AfterWorkOption.SWARM_MANAGER, next_agent_selection_msg="Select the next agent from this list {agentlist}. From the list please select the order_mgmt_agent. Order Id is {order_id}"),
    ]
)

_, _, _ = initiate_swarm_chat(
    ...,
    swarm_manager_args={"llm_config": llm_config}
)

Related issue number

#26

Checks

• I've included any doc changes needed for https://docs.ag2.ai/. See https://docs.ag2.ai/docs/contributor-guide/documentation to build and test documentation locally.
• I've added tests (if relevant) corresponding to the changes introduced in this PR.
• I've made sure all auto checks have passed.

[marklysze]

Thanks for the initial review @sonichi, if you could have another look over now that would be appreciated.

[marklysze]

@yiranwu0, @linmou - if you're able to test this it would be appreciated!

[yiranwu0]

Look good to me overall. My only question is about the context_str class. I think this class might not be necessary at all, and also consider existing packages like dedent.

[yiranwu0]

Can we let the user to only pass in a str or a callable?
We can always convert the str to a ContextStr class, replace agent_list and context variables.

[marklysze]

Thanks for the review. You raise a good point and Chi and I had a chat about this as well prior. The thought behind making it an explicit choice to use ContextStr or a normal, non-substituted, string is so that the developer has complete control over the choice and they're not unsure of what will happen to the value.

@sonichi if you have anything further to add to this consideration?

[sonichi]

I guess the ambiguity happens when the str contains "{}", and whether it means a format string or not is ambiguous. That's my understanding of why ContextStr is introduced. Is that right? Shall we add the consideration into docstr?

[marklysze]

Yes, that's right, by using "{}" the user doesn't have to be concerned about unexpected replacements (except {agentlist], which is default Group Chat behaviour)

In terms of adding to docstr, do you mean the ContextStr docstring?

[marklysze]

@yiranwu0, did you have any further thoughts on this? Are you okay with making it an explicit choice for the developer to choose ContextStr when they want the substitutions and without when they don't?