m2-l3.py

from langchain.agents import create_agent
from langchain.messages import HumanMessage
from langchain.tools import tool

QUESTION = HumanMessage(content="What's the square root of 456?")

@tool
def square_root(x: float) -> float:
    """
    Calculate the square root of a number.
    """
    return x ** 0.5

@tool
def square(x: float) -> float:
    """
    Calculate the square of a number.
    """
    return x ** 2

subagent_1 = create_agent(
    model="gpt-5-nano",
    tools=[square_root],
)

subagent_2 = create_agent(
    model="gpt-5-nano",
    tools=[square],
)

@tool
def call_subagent_1(x: float) -> float:
    """
    Call the subagent 1 to calculate the square root of a number.
    """
    response = subagent_1.invoke({"messages": [HumanMessage(content=f"Calculate the square root of {x}?")]})
    return response["messages"][-1].content

@tool
def call_subagent_2(x: float) -> float:
    """
    Call the subagent 2 to calculate the square of a number.
    """
    response = subagent_2.invoke({"messages": [HumanMessage(content=f"Calculate the square of {x}?")]})
    return response["messages"][-1].content

main_agent = create_agent(
    model="gpt-5-nano",
    tools=[call_subagent_1, call_subagent_2],
    system_prompt="You are a helpful assistant who can call subagents to calculate the square root and square of a number.",
)
response = main_agent.invoke({"messages": [QUESTION]})
print(response["messages"][-1].content)