Agents as tools

ORCHESTRATOR_PROMPT = """You are CEO Alice's executive assistant at DataStream Corp.

Route queries to specialists:
- Database queries, employee counts, department stats -> use data_agent
- Weather forecasts, temperature, weather conditions (any location) -> use weather_agent. It covers US
  locations only; when it says it can only look up US weather, pass that sentence on to Alice unchanged.
- Simple company questions -> answer directly

After the specialists reply, ALWAYS finish with a final answer for Alice that combines every specialist result
in your own words. Never end your turn with an empty message. If a specialist reports that it failed, tell Alice
which part could not be answered; never invent numbers or use placeholders such as [count].

DataStream Corp is a technology company with Sales, HR, and Engineering departments."""


SPECIALIST_FAILED = "The {specialist} specialist failed: {reason}. No data is available from it for this question."
"""Returned to the orchestrator instead of "" so a failed specialist is reported, never filled in with guesses."""


def consult(specialist: Agent, query: str, results: list[str], name: str) -> str:
    """Ask a specialist via `ask_specialist`; an exception or empty answer becomes an explicit failure message.

    Args:
        specialist: The specialist agent behind an orchestrator tool.
        query: The question the orchestrator routed to it.
        results: Shared list of specialist answers for this turn.
        name: Short specialist name for the failure message, e.g. "data".

    Returns:
        The specialist's answer, or "The <name> specialist failed: ..." so the orchestrator never guesses.
    """
    try:
        answer = ask_specialist(specialist, query, results).strip()
        reason = "it returned an empty answer"
    except Exception as error:
        answer, reason = "", f"{type(error).__name__}: {error}"
    if not answer:
        answer = SPECIALIST_FAILED.format(specialist=name, reason=reason)
        results.append(answer)
    return answer


def build_orchestrator(db_tools: list, specialist_results: list[str]) -> Agent:
    """Orchestrator whose two tools are specialist agents; their answers are appended to `specialist_results`."""

    @tool
    def data_agent(query: str) -> str:
        """Query and analyze the DataStream Corp database: employee counts, department statistics, any SQL question.

        Args:
            query: A database or data analysis question.
        """
        specialist = Agent(
            model=make_model(),
            system_prompt="You are a data specialist. Query the database to answer questions.",
            tools=db_tools,
            hooks=[ApprovalHook()],
            callback_handler=None,
        )
        return consult(specialist, query, specialist_results, "data")

    @tool
    def weather_agent(query: str) -> str:
        """Get the weather forecast for a US location (National Weather Service); says so for non-US places.

        Args:
            query: A weather-related question, including the place.
        """
        specialist = Agent(
            model=make_model(),
            system_prompt=WEATHER_PROMPT,
            tools=[get_us_forecast],
            callback_handler=None,
        )
        return consult(specialist, query, specialist_results, "weather")

    return Agent(
        name="Executive Assistant",
        model=make_model(),
        system_prompt=ORCHESTRATOR_PROMPT,
        tools=[data_agent, weather_agent],
    )

US_ONLY_REPLY = "I can only look up US weather (National Weather Service)."


WEATHER_PROMPT = f"""You are a weather specialist for US locations only, using the National Weather Service.
- For a place in the United States, call get_us_forecast with its latitude and longitude in decimal degrees,
  then answer concisely from the forecast periods it returns.
- For any place outside the United States, do not call the tool; reply exactly: {US_ONLY_REPLY}
- If the tool returns an error, say that the forecast is unavailable; never guess the weather."""


@tool
def get_us_forecast(latitude: float, longitude: float) -> str:
    """Get the National Weather Service forecast for a location in the United States (next few periods).

    Args:
        latitude: Latitude of the US location in decimal degrees, e.g. 47.6062 for Seattle.
        longitude: Longitude of the US location in decimal degrees, e.g. -122.3321 for Seattle.
    """
    try:
        with make_client() as client:
            place, periods = point_forecast(client, latitude, longitude)
    except NotUsLocationError:
        return f"Error: no NWS forecast for {latitude},{longitude}. {US_ONLY_REPLY}"
    except httpx.TimeoutException:
        return f"Error: the National Weather Service did not answer within {TIMEOUT_SECONDS:g}s; try again."
    except (httpx.HTTPError, KeyError, ValueError) as error:
        return f"Error: the National Weather Service request failed ({type(error).__name__}: {error})."
    lines = [format_period(period) for period in periods[:FORECAST_PERIODS]]
    return "\n".join([f"NWS forecast for {place}:", *lines])