from autogen import initiate_swarm_chat, ConversableAgent, register_hand_off, OnCondition, AfterWork, AfterWorkOption
from autogen.agents.experimental import WebSurferAgent

# LLM configuration for our agent to select the tools and craft the message
llm_config = {"model": "gpt-4o-mini"}

triage_agent = ConversableAgent(
    name="triage_agent",
    llm_config=llm_config,
)

websurfer_agent = WebSurferAgent(
    name="websurfer_agent",
    llm_config=llm_config,
)

register_hand_off(
    agent=triage_agent,
    hand_to=[
        OnCondition(
            target=websurfer_agent,
            condition="Find information on the web."
        ),
        AfterWork(AfterWorkOption.TERMINATE)
    ]
)

register_hand_off(
    agent=websurfer_agent,
    hand_to=[
        AfterWork(
            agent=triage_agent
        )
    ]
)

result, cv, last_agent = initiate_swarm_chat(
    initial_agent=triage_agent,
    agents=[triage_agent, websurfer_agent],
    messages="Find the latest weather for Sydney, NSW, Australia.",
    max_rounds=10,
)

_User (to chat_manager):

Find the latest weather for Sydney, NSW, Australia.

--------------------------------------------------------------------------------

Next speaker: triage_agent


>>>>>>>> USING AUTO REPLY...
triage_agent (to chat_manager):

***** Suggested tool call (call_mNKxUjLz9hdouBN2nCIniWiW): transfer_triage_agent_to_websurfer_agent *****
Arguments: 
{}
*********************************************************************************************************

--------------------------------------------------------------------------------

Next speaker: _Swarm_Tool_Executor


>>>>>>>> EXECUTING FUNCTION transfer_triage_agent_to_websurfer_agent...
Call ID: call_mNKxUjLz9hdouBN2nCIniWiW
Input arguments: {}
_Swarm_Tool_Executor (to chat_manager):

***** Response from calling tool (call_mNKxUjLz9hdouBN2nCIniWiW) *****
Swarm agent --> websurfer_agent
**********************************************************************

--------------------------------------------------------------------------------

Next speaker: websurfer_agent


>>>>>>>> USING AUTO REPLY...
websurfer_agent (to chat_manager):

***** Suggested tool call (call_gGNHUdnycZscUIMar58NWmdT): browser_use *****
Arguments: 
{"task":"Find the latest weather report for Sydney, NSW, Australia."}
****************************************************************************

--------------------------------------------------------------------------------

Next speaker: _Swarm_Tool_Executor

_Swarm_Tool_Executor (to chat_manager):

***** Response from calling tool (call_gGNHUdnycZscUIMar58NWmdT) *****
Error: Function browser_use not found.
**********************************************************************

--------------------------------------------------------------------------------