from __future__ import annotations

from amplitude_ai.testing import MockAmplitudeAI


def run_multi_agent_example() -> MockAmplitudeAI:
    mock = MockAmplitudeAI()
    orchestrator = mock.agent(
        "orchestrator",
        user_id="user-multi",
        session_id="session-multi-agent",
    )
    retriever = orchestrator.child("retriever")

    orchestrator.track_user_message("Find the top churn drivers this month.")
    retriever.track_ai_message(
        "Retrieved churn cohorts and top correlated behaviors.",
        model="gpt-4o-mini",
        provider="openai",
        latency_ms=180,
    )
    orchestrator.track_ai_message(
        "Top churn driver is onboarding drop-off in day 1.",
        model="gpt-4o",
        provider="openai",
        latency_ms=260,
    )
    return mock
