Enhance LLM model toggle functionality with migration support

Updated the toggle LLM model API to include an optional migration feature, allowing workflows to be migrated to a specified replacement model when disabling a model. Refactored related request and response models to accommodate this change. Improved error handling and logging for better debugging. Updated frontend actions and components to support the new migration parameter.

autogpt_platform/backend/backend/server/v2/admin/llm_routes.py

@@ -132,18 +132,37 @@ async def update_llm_model(
 @router.patch(
     "/models/{model_id}/toggle",
     summary="Toggle LLM model availability",
-    response_model=llm_model.LlmModel,
+    response_model=llm_model.ToggleLlmModelResponse,
 )
 async def toggle_llm_model(
     model_id: str,
     request: llm_model.ToggleLlmModelRequest,
 ):
+    """
+    Toggle a model's enabled status, optionally migrating workflows when disabling.
+
+    If disabling a model and `migrate_to_slug` is provided, all workflows using
+    this model will be migrated to the specified replacement model before disabling.
+    """
     try:
-        model = await llm_db.toggle_model(
-            model_id=model_id, is_enabled=request.is_enabled
+        result = await llm_db.toggle_model(
+            model_id=model_id,
+            is_enabled=request.is_enabled,
+            migrate_to_slug=request.migrate_to_slug,
         )
         await _refresh_runtime_state()
-        return model
+        if result.nodes_migrated > 0:
+            logger.info(
+                "Toggled model '%s' to %s and migrated %d nodes to '%s'",
+                result.model.slug,
+                "enabled" if request.is_enabled else "disabled",
+                result.nodes_migrated,
+                result.migrated_to_slug,
+            )
+        return result
+    except ValueError as exc:
+        logger.warning("Model toggle validation failed: %s", exc)
+        raise fastapi.HTTPException(status_code=400, detail=str(exc)) from exc
     except Exception as exc:
         logger.exception("Failed to toggle LLM model %s: %s", model_id, exc)
         raise fastapi.HTTPException(

autogpt_platform/backend/backend/server/v2/llm/db.py

@@ -219,13 +219,91 @@ async def update_model(
     return _map_model(record)
 
 
-async def toggle_model(model_id: str, is_enabled: bool) -> llm_model.LlmModel:
-    record = await prisma.models.LlmModel.prisma().update(
-        where={"id": model_id},
-        data={"isEnabled": is_enabled},
-        include={"Costs": True},
+async def toggle_model(
+    model_id: str, is_enabled: bool, migrate_to_slug: str | None = None
+) -> llm_model.ToggleLlmModelResponse:
+    """
+    Toggle a model's enabled status, optionally migrating workflows when disabling.
+
+    Args:
+        model_id: UUID of the model to toggle
+        is_enabled: New enabled status
+        migrate_to_slug: If disabling and this is provided, migrate all workflows
+                         using this model to the specified replacement model
+
+    Returns:
+        ToggleLlmModelResponse with the updated model and optional migration stats
+    """
+    import prisma as prisma_module
+
+    # Get the model being toggled
+    model = await prisma.models.LlmModel.prisma().find_unique(
+        where={"id": model_id}, include={"Costs": True}
+    )
+    if not model:
+        raise ValueError(f"Model with id '{model_id}' not found")
+
+    nodes_migrated = 0
+
+    # If disabling with migration, perform migration first
+    if not is_enabled and migrate_to_slug:
+        # Validate replacement model exists and is enabled
+        replacement = await prisma.models.LlmModel.prisma().find_unique(
+            where={"slug": migrate_to_slug}
+        )
+        if not replacement:
+            raise ValueError(f"Replacement model '{migrate_to_slug}' not found")
+        if not replacement.isEnabled:
+            raise ValueError(
+                f"Replacement model '{migrate_to_slug}' is disabled. "
+                f"Please enable it before using it as a replacement."
+            )
+
+        # Count affected nodes
+        count_result = await prisma_module.get_client().query_raw(
+            """
+            SELECT COUNT(*) as count
+            FROM "AgentNode"
+            WHERE "constantInput"::jsonb->>'model' = $1
+            """,
+            model.slug,
+        )
+        nodes_migrated = int(count_result[0]["count"]) if count_result else 0
+
+        # Perform migration and toggle atomically
+        async with transaction() as tx:
+            if nodes_migrated > 0:
+                await tx.execute_raw(
+                    """
+                    UPDATE "AgentNode"
+                    SET "constantInput" = JSONB_SET(
+                        "constantInput"::jsonb,
+                        '{model}',
+                        to_jsonb($1::text)
+                    )
+                    WHERE "constantInput"::jsonb->>'model' = $2
+                    """,
+                    migrate_to_slug,
+                    model.slug,
+                )
+            record = await tx.llmmodel.update(
+                where={"id": model_id},
+                data={"isEnabled": is_enabled},
+                include={"Costs": True},
+            )
+    else:
+        # Simple toggle without migration
+        record = await prisma.models.LlmModel.prisma().update(
+            where={"id": model_id},
+            data={"isEnabled": is_enabled},
+            include={"Costs": True},
+        )
+
+    return llm_model.ToggleLlmModelResponse(
+        model=_map_model(record),
+        nodes_migrated=nodes_migrated,
+        migrated_to_slug=migrate_to_slug if nodes_migrated > 0 else None,
     )
-    return _map_model(record)
 
 
 async def get_model_usage(model_id: str) -> llm_model.LlmModelUsageResponse:

autogpt_platform/backend/backend/server/v2/llm/model.py

@@ -106,6 +106,13 @@ class UpdateLlmModelRequest(pydantic.BaseModel):
 
 class ToggleLlmModelRequest(pydantic.BaseModel):
     is_enabled: bool
+    migrate_to_slug: Optional[str] = None
+
+
+class ToggleLlmModelResponse(pydantic.BaseModel):
+    model: LlmModel
+    nodes_migrated: int = 0
+    migrated_to_slug: Optional[str] = None
 
 
 class DeleteLlmModelResponse(pydantic.BaseModel):