from __future__ import annotations

from amplitude_ai.testing import MockAmplitudeAI


def run_framework_integration_example() -> MockAmplitudeAI:
    mock = MockAmplitudeAI()
    req_context = {
        "user_id": "user-framework",
        "session_id": "session-framework",
        "route": "/api/chat",
    }

    agent = mock.agent(
        "framework-handler",
        user_id=req_context["user_id"],
        session_id=req_context["session_id"],
        context={"route": req_context["route"]},
    )
    agent.track_user_message("What is my weekly active user trend?")
    agent.track_ai_message(
        "Weekly active users are stable with a slight upward trend.",
        model="gpt-4o-mini",
        provider="openai",
        latency_ms=210,
    )
    return mock
