Function test_task_tool_mock_main_agent

tests/main/test_task_tool.py:24–71

Test that when MockAgent uses TaskTool, it properly spawns a sub-agent that can use tools and complete tasks.

()

def test_task_tool_mock_main_agent():
    """
    Test that when MockAgent uses TaskTool, it  properly spawns a sub-agent
    that can use tools and complete tasks.
    """

    # Configure the main agent to use TaskTool:
    # The MockLM has a fixed response, which is the TaskTool request
    main_config = ChatAgentConfig(
        llm=MockLMConfig(
            default_response=TaskTool(
                system_message=f"""
                    You are a calculator assistant. When asked to 
                    calculate, use the TOOL `{MultiplierTool.name()}` to multiply the 
                    numbers, then use the TOOL `{DoneTool.name()}` to return the result
                    """,
                prompt="Multiply 5 and 7",
                model="gpt-4.1-mini",
                tools=["multiplier_tool"],
                max_iterations=5,
            ).model_dump_json()
        ),
        name="MainAgent",
    )
    main_agent = ChatAgent(main_config)

    # Enable TaskTool and MultiplierTool for the main agent.
    # The MultiplierTool must be enabled for the main agent,
    # since the TaskTool handler will create a sub-agent that uses it,
    # and the handler only has access to tools enabled for the main agent.
    main_agent.enable_message([TaskTool, MultiplierTool], use=True, handle=True)

    # Create main task that stops after handling the task_tool
    task = Task(
        main_agent,
        name="MainTask",
        interactive=False,
        config=TaskConfig(
            done_sequences=["T,A"],  # LLM (Tool (TaskTool)), Agent(Handled) -> done
        ),
    )

    # Run the task
    result = task.run(msg="Please calculate something")

    # The result should be from the sub-agent's execution
    assert result is not None, "Task should return a result"
    assert "35" in result.content, "Result should contain the multiplication result"


class NebrowskiTool(ToolMessage):