Function handle_model_switching

api_utils/model_switching.py:44–86

(
    req_id: str, context: RequestContext
)

async def handle_model_switching(
    req_id: str, context: RequestContext
) -> RequestContext:
    set_request_id(req_id)
    if not context["needs_model_switching"]:
        return context

    logger = context["logger"]
    page = context["page"]
    model_switching_lock = context["model_switching_lock"]
    model_id_to_use = context["model_id_to_use"]

    # Assert non-None values required for model switching
    assert page is not None, "Page must be ready for model switching"
    assert model_id_to_use is not None, "Target model ID must be set"

    async with model_switching_lock:
        if state.current_ai_studio_model_id != model_id_to_use:
            logger.info(
                f"[{req_id}] Preparing to switch model: {state.current_ai_studio_model_id} -> {model_id_to_use}"
            )
            from browser_utils import switch_ai_studio_model

            switch_success = await switch_ai_studio_model(page, model_id_to_use, req_id)
            if switch_success:
                state.current_ai_studio_model_id = model_id_to_use
                context["model_actually_switched"] = True
                context["current_ai_studio_model_id"] = model_id_to_use
                logger.info(
                    f"[{req_id}] ✅ Model switched successfully: {state.current_ai_studio_model_id}"
                )
            else:
                # Current model ID should exist when switching fails
                current_model = state.current_ai_studio_model_id or "unknown"
                await _handle_model_switch_failure(
                    req_id,
                    page,
                    model_id_to_use,
                    current_model,
                    logger,
                )

    return context


async def _handle_model_switch_failure(