from __future__ import annotations

from amplitude_ai.testing import MockAmplitudeAI


def run_wrap_openai_example() -> MockAmplitudeAI:
    mock = MockAmplitudeAI()
    agent = mock.agent("wrapped-openai-agent", user_id="user-wrap")
    agent.track_user_message("Summarize today's funnel drop-off.", session_id="session-wrap-openai")
    agent.track_ai_message(
        "The largest drop-off is step 2 -> 3, concentrated on mobile web.",
        session_id="session-wrap-openai",
        model="gpt-4o",
        provider="openai",
        latency_ms=340,
    )
    return mock
