Function switch_ai_studio_model

browser_utils/models/switcher.py:21–210

Switch AI Studio model

(page: AsyncPage, model_id: str, req_id: str)

async def switch_ai_studio_model(page: AsyncPage, model_id: str, req_id: str) -> bool:
    """Switch AI Studio model"""
    logger.info(f"[Model] Switching to -> {model_id}")
    original_prefs_str: Optional[str] = None
    _original_prompt_model: Optional[str] = None
    new_chat_url = f"https://{AI_STUDIO_URL_PATTERN}prompts/new_chat"

    try:
        original_prefs_str = await page.evaluate(
            "() => localStorage.getItem('aiStudioUserPreference')"
        )
        if original_prefs_str:
            try:
                original_prefs_obj = json.loads(original_prefs_str)
                _original_prompt_model = original_prefs_obj.get("promptModel")
            except json.JSONDecodeError:
                logger.warning(
                    "Failed to parse original aiStudioUserPreference JSON string."
                )
                original_prefs_str = None

        current_prefs_for_modification = (
            json.loads(original_prefs_str) if original_prefs_str else {}
        )
        full_model_path = f"models/{model_id}"

        if current_prefs_for_modification.get("promptModel") == full_model_path:
            logger.debug(f"[Model] Already at target model {model_id}")
            if page.url != new_chat_url:
                logger.debug(
                    f"[Model] URL is not new_chat, navigating to {new_chat_url}"
                )
                await page.goto(
                    new_chat_url, wait_until="domcontentloaded", timeout=30000
                )
                await expect_async(page.locator(INPUT_SELECTOR)).to_be_visible(
                    timeout=30000
                )
            return True

        logger.debug(
            f"[Model] Updating localStorage.promptModel: {current_prefs_for_modification.get('promptModel', 'unknown')} -> {full_model_path}"
        )
        current_prefs_for_modification["promptModel"] = full_model_path
        await page.evaluate(
            "(prefsStr) => localStorage.setItem('aiStudioUserPreference', prefsStr)",
            json.dumps(current_prefs_for_modification),
        )

        # Use new forced setting feature
        logger.debug("[State] Applying forced UI state settings...")
        ui_state_success = await _verify_and_apply_ui_state(page, req_id)
        if not ui_state_success:
            logger.warning(
                "UI state setting failed, but continuing model switching flow"
            )

        # To maintain compatibility, also update current prefs object