def test_supervisor_basic_workflow(
    include_agent_name: AgentNameMode | None,
    include_individual_agent_name: AgentNameMode | None,
) -> None:
    """Test basic supervisor workflow with two agents."""

    # output_mode = "last_message"
    @tool
    def add(a: float, b: float) -> float:
        """Add two numbers."""
        return a + b

    @tool
    def web_search(query: str) -> str:
        """Search the web for information."""
        return (
            "Here are the headcounts for each of the FAANG companies in 2024:\n"
            "1. **Facebook (Meta)**: 67,317 employees.\n"
            "2. **Apple**: 164,000 employees.\n"
            "3. **Amazon**: 1,551,000 employees.\n"
            "4. **Netflix**: 14,000 employees.\n"
            "5. **Google (Alphabet)**: 181,269 employees."
        )

    math_model: FakeChatModel = FakeChatModel(responses=math_agent_messages)
    if include_individual_agent_name:
        math_model = cast(
            FakeChatModel,
            with_agent_name(math_model.bind_tools([add]), include_individual_agent_name),
        )

    math_agent = create_react_agent(
        model=math_model,
        tools=[add],
        name="math_expert",
    )

    research_model = FakeChatModel(responses=research_agent_messages)
    if include_individual_agent_name:
        research_model = cast(
            FakeChatModel,
            with_agent_name(research_model.bind_tools([web_search]), include_individual_agent_name),
        )

    research_agent = create_react_agent(
        model=research_model,
        tools=[web_search],
        name="research_expert",
    )

    workflow = create_supervisor(
        [math_agent, research_agent],
        model=FakeChatModel(responses=supervisor_messages),
        include_agent_name=include_agent_name,
    )

    app = workflow.compile()
    assert app is not None

    result = app.invoke(
        {
            "messages": [
                HumanMessage(
                    content="what's the combined headcount of the FAANG companies in 2024?"
                )
            ]
        }
    )

    assert len(result["messages"]) == 12
    # first supervisor handoff
    assert result["messages"][1] == supervisor_messages[0]
    # last research agent message
    assert result["messages"][3] == research_agent_messages[-1]
    # next supervisor handoff
    assert result["messages"][6] == supervisor_messages[1]
    # last math agent message
    assert result["messages"][8] == math_agent_messages[-1]
    # final supervisor message
    assert result["messages"][11] == supervisor_messages[-1]

    # output_mode = "full_history"
    math_agent = create_react_agent(
        model=FakeChatModel(responses=math_agent_messages),
        tools=[add],
        name="math_expert",
    )

    research_agent = create_react_agent(
        model=FakeChatModel(responses=research_agent_messages),
        tools=[web_search],
        name="research_expert",
    )

    workflow_full_history = create_supervisor(
        [math_agent, research_agent],
        model=FakeChatModel(responses=supervisor_messages),
        output_mode="full_history",
    )
    app_full_history = workflow_full_history.compile()
    result_full_history = app_full_history.invoke(
        {
            "messages": [
                HumanMessage(
                    content="what's the combined headcount of the FAANG companies in 2024?"
                )
            ]
        }
    )

    assert len(result_full_history["messages"]) == 23
    # first supervisor handoff
    assert result_full_history["messages"][1] == supervisor_messages[0]
    # all research agent AI messages
    assert result_full_history["messages"][3] == research_agent_messages[0]
    assert result_full_history["messages"][5] == research_agent_messages[1]
    # next supervisor handoff
    assert result_full_history["messages"][8] == supervisor_messages[1]
    # all math agent AI messages
    assert result_full_history["messages"][10] == math_agent_messages[0]
    assert result_full_history["messages"][14] == math_agent_messages[1]
    assert result_full_history["messages"][17] == math_agent_messages[2]
    # final supervisor message
    assert result_full_history["messages"][-1] == supervisor_messages[-1]
Test basic supervisor workflow with two agents.
tests/test_supervisor.py