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):

autogpt_platform/frontend/src/app/(platform)/admin/llms/actions.ts

@@ -134,11 +134,14 @@ export async function updateLlmModelAction(formData: FormData) {
   revalidatePath(ADMIN_LLM_PATH);
 }
 
-export async function toggleLlmModelAction(formData: FormData) {
+export async function toggleLlmModelAction(formData: FormData): Promise<void> {
   const modelId = String(formData.get("model_id"));
   const shouldEnable = formData.get("is_enabled") === "true";
+  const migrateToSlug = formData.get("migrate_to_slug");
+
   const payload: ToggleLlmModelRequest = {
     is_enabled: shouldEnable,
+    migrate_to_slug: migrateToSlug ? String(migrateToSlug) : undefined,
   };
   const api = new BackendApi();
   await api.toggleAdminLlmModel(modelId, payload);

autogpt_platform/frontend/src/app/(platform)/admin/llms/components/DisableModelModal.tsx

@@ -0,0 +1,223 @@
+"use client";
+
+import { useState } from "react";
+import { Dialog } from "@/components/molecules/Dialog/Dialog";
+import { Button } from "@/components/atoms/Button/Button";
+import type { LlmModel } from "@/lib/autogpt-server-api/types";
+import { toggleLlmModelAction } from "../actions";
+
+export function DisableModelModal({
+  model,
+  availableModels,
+}: {
+  model: LlmModel;
+  availableModels: LlmModel[];
+}) {
+  const [open, setOpen] = useState(false);
+  const [isDisabling, setIsDisabling] = useState(false);
+  const [error, setError] = useState<string | null>(null);
+  const [usageCount, setUsageCount] = useState<number | null>(null);
+  const [selectedMigration, setSelectedMigration] = useState<string>("");
+  const [wantsMigration, setWantsMigration] = useState(false);
+
+  // Filter out the current model and disabled models from replacement options
+  const migrationOptions = availableModels.filter(
+    (m) => m.id !== model.id && m.is_enabled
+  );
+
+  async function fetchUsage() {
+    try {
+      const BackendApi = (await import("@/lib/autogpt-server-api")).default;
+      const api = new BackendApi();
+      const usage = await api.getAdminLlmModelUsage(model.id);
+      setUsageCount(usage.node_count);
+    } catch {
+      setUsageCount(null);
+    }
+  }
+
+  async function handleDisable(formData: FormData) {
+    setIsDisabling(true);
+    setError(null);
+    try {
+      await toggleLlmModelAction(formData);
+      setOpen(false);
+    } catch (err) {
+      setError(err instanceof Error ? err.message : "Failed to disable model");
+    } finally {
+      setIsDisabling(false);
+    }
+  }
+
+  function resetState() {
+    setError(null);
+    setSelectedMigration("");
+    setWantsMigration(false);
+  }
+
+  const hasUsage = usageCount !== null && usageCount > 0;
+
+  return (
+    <Dialog
+      title="Disable Model"
+      controlled={{
+        isOpen: open,
+        set: async (isOpen) => {
+          setOpen(isOpen);
+          if (isOpen) {
+            setUsageCount(null);
+            resetState();
+            await fetchUsage();
+          }
+        },
+      }}
+      styling={{ maxWidth: "550px" }}
+    >
+      <Dialog.Trigger>
+        <Button type="button" variant="outline" size="small" className="min-w-0">
+          Disable
+        </Button>
+      </Dialog.Trigger>
+      <Dialog.Content>
+        <div className="mb-4 text-sm text-muted-foreground">
+          Disabling a model will hide it from users when creating new workflows.
+        </div>
+
+        <div className="space-y-4">
+          <div className="rounded-lg border border-yellow-200 bg-yellow-50 p-4 dark:border-yellow-900 dark:bg-yellow-950">
+            <div className="flex items-start gap-3">
+              <div className="flex-shrink-0 text-yellow-600">⚠️</div>
+              <div className="text-sm text-yellow-800 dark:text-yellow-200">
+                <p className="font-semibold">You are about to disable:</p>
+                <p className="mt-1">
+                  <span className="font-medium">{model.display_name}</span>{" "}
+                  <span className="text-yellow-600 dark:text-yellow-400">
+                    ({model.slug})
+                  </span>
+                </p>
+                {usageCount === null ? (
+                  <p className="mt-2 text-yellow-600 dark:text-yellow-400">
+                    Loading usage data...
+                  </p>
+                ) : usageCount > 0 ? (
+                  <p className="mt-2 font-semibold">
+                    📊 Impact: {usageCount} block{usageCount !== 1 ? "s" : ""}{" "}
+                    currently use this model
+                  </p>
+                ) : (
+                  <p className="mt-2">
+                    ✓ No workflows are currently using this model.
+                  </p>
+                )}
+              </div>
+            </div>
+          </div>
+
+          {hasUsage && (
+            <div className="rounded-lg border border-border bg-muted/50 p-4">
+              <label className="flex items-start gap-3">
+                <input
+                  type="checkbox"
+                  checked={wantsMigration}
+                  onChange={(e) => {
+                    setWantsMigration(e.target.checked);
+                    if (!e.target.checked) {
+                      setSelectedMigration("");
+                    }
+                  }}
+                  className="mt-1"
+                />
+                <div className="text-sm">
+                  <span className="font-medium">
+                    Migrate existing workflows to another model
+                  </span>
+                  <p className="mt-1 text-muted-foreground">
+                    If unchecked, existing workflows will use automatic fallback
+                    to an enabled model from the same provider.
+                  </p>
+                </div>
+              </label>
+
+              {wantsMigration && (
+                <div className="mt-4">
+                  <label className="text-sm font-medium">
+                    <span className="mb-2 block">
+                      Select Replacement Model{" "}
+                      <span className="text-red-500">*</span>
+                    </span>
+                    <select
+                      required
+                      value={selectedMigration}
+                      onChange={(e) => setSelectedMigration(e.target.value)}
+                      className="w-full rounded border border-input bg-background p-2 text-sm"
+                    >
+                      <option value="">-- Choose a replacement model --</option>
+                      {migrationOptions.map((m) => (
+                        <option key={m.id} value={m.slug}>
+                          {m.display_name} ({m.slug})
+                        </option>
+                      ))}
+                    </select>
+                    {migrationOptions.length === 0 && (
+                      <p className="mt-2 text-xs text-red-600">
+                        No other enabled models available for migration.
+                      </p>
+                    )}
+                  </label>
+                </div>
+              )}
+            </div>
+          )}
+
+          <form action={handleDisable} className="space-y-4">
+            <input type="hidden" name="model_id" value={model.id} />
+            <input type="hidden" name="is_enabled" value="false" />
+            {wantsMigration && selectedMigration && (
+              <input
+                type="hidden"
+                name="migrate_to_slug"
+                value={selectedMigration}
+              />
+            )}
+
+            {error && (
+              <div className="rounded-lg border border-red-200 bg-red-50 p-3 text-sm text-red-800 dark:border-red-900 dark:bg-red-950 dark:text-red-200">
+                {error}
+              </div>
+            )}
+
+            <Dialog.Footer>
+              <Button
+                variant="ghost"
+                size="small"
+                onClick={() => {
+                  setOpen(false);
+                  resetState();
+                }}
+                disabled={isDisabling}
+              >
+                Cancel
+              </Button>
+              <Button
+                type="submit"
+                variant="primary"
+                size="small"
+                disabled={
+                  isDisabling ||
+                  (wantsMigration && !selectedMigration) ||
+                  usageCount === null
+                }
+              >
+                {isDisabling
+                  ? "Disabling..."
+                  : wantsMigration && selectedMigration
+                    ? "Disable & Migrate"
+                    : "Disable Model"}
+              </Button>
+            </Dialog.Footer>
+          </form>
+        </div>
+      </Dialog.Content>
+    </Dialog>
+  );
+}

autogpt_platform/frontend/src/app/(platform)/admin/llms/components/ModelsTable.tsx

@@ -10,6 +10,7 @@ import {
 import { Button } from "@/components/atoms/Button/Button";
 import { toggleLlmModelAction } from "../actions";
 import { DeleteModelModal } from "./DeleteModelModal";
+import { DisableModelModal } from "./DisableModelModal";
 import { EditModelModal } from "./EditModelModal";
 
 export function ModelsTable({
@@ -105,10 +106,14 @@ export function ModelsTable({
                 </TableCell>
                 <TableCell>
                   <div className="flex items-center justify-end gap-2">
-                    <ToggleModelButton
-                      modelId={model.id}
-                      isEnabled={model.is_enabled}
-                    />
+                    {model.is_enabled ? (
+                      <DisableModelModal
+                        model={model}
+                        availableModels={models}
+                      />
+                    ) : (
+                      <EnableModelButton modelId={model.id} />
+                    )}
                     <EditModelModal model={model} providers={providers} />
                     <DeleteModelModal
                       model={model}
@@ -125,24 +130,13 @@ export function ModelsTable({
   );
 }
 
-function ToggleModelButton({
-  modelId,
-  isEnabled,
-}: {
-  modelId: string;
-  isEnabled: boolean;
-}) {
+function EnableModelButton({ modelId }: { modelId: string }) {
   return (
     <form action={toggleLlmModelAction} className="inline">
       <input type="hidden" name="model_id" value={modelId} />
-      <input type="hidden" name="is_enabled" value={(!isEnabled).toString()} />
-      <Button
-        type="submit"
-        variant="outline"
-        size="small"
-        className="min-w-0"
-      >
-        {isEnabled ? "Disable" : "Enable"}
+      <input type="hidden" name="is_enabled" value="true" />
+      <Button type="submit" variant="outline" size="small" className="min-w-0">
+        Enable
       </Button>
     </form>
   );

autogpt_platform/frontend/src/lib/autogpt-server-api/client.ts

@@ -56,6 +56,7 @@ import type {
   CreateLlmModelRequest,
   UpdateLlmModelRequest,
   ToggleLlmModelRequest,
+  ToggleLlmModelResponse,
   UpsertLlmProviderRequest,
   LlmModelsResponse,
   LlmProvider,
@@ -463,7 +464,7 @@ export default class BackendAPI {
   toggleAdminLlmModel(
     modelId: string,
     payload: ToggleLlmModelRequest,
-  ): Promise<LlmModel> {
+  ): Promise<ToggleLlmModelResponse> {
     return this._request(
       "PATCH",
       `/llm/admin/llm/models/${modelId}/toggle`,

autogpt_platform/frontend/src/lib/autogpt-server-api/types.ts

@@ -349,6 +349,13 @@ export type UpdateLlmModelRequest = {
 
 export type ToggleLlmModelRequest = {
   is_enabled: boolean;
+  migrate_to_slug?: string;
+};
+
+export type ToggleLlmModelResponse = {
+  model: LlmModel;
+  nodes_migrated: number;
+  migrated_to_slug?: string | null;
 };
 
 export type BlockIOOneOfSubSchema = {