refactor(recipe-studio): add image preview support for dataset and LLM configurations p2

studio/backend/core/data_recipe/service.py

@@ -1,10 +1,123 @@
 from __future__ import annotations
 
+import base64
+import io
 import os
+from pathlib import Path
 from typing import Any
 
 from .jsonable import to_jsonable
 
+_IMAGE_CONTEXT_PATCHED = False
+
+
+def _encode_bytes_to_base64(value: bytes | bytearray) -> str:
+    return base64.b64encode(bytes(value)).decode("utf-8")
+
+
+def _load_image_file_to_base64(path_value: str, *, base_path: str | None = None) -> str | None:
+    try:
+        path = Path(path_value)
+        candidates: list[Path] = []
+        if path.is_absolute():
+            candidates.append(path)
+        else:
+            if base_path:
+                candidates.append(Path(base_path) / path)
+            candidates.append(Path.cwd() / path)
+
+        for candidate in candidates:
+            if not candidate.exists() or not candidate.is_file():
+                continue
+            with candidate.open("rb") as f:
+                return _encode_bytes_to_base64(f.read())
+    except (OSError, TypeError, ValueError):
+        return None
+    return None
+
+
+def _pil_image_to_base64(value: Any) -> str | None:
+    try:
+        from PIL.Image import Image as PILImage  # type: ignore
+    except ImportError:
+        return None
+    if not isinstance(value, PILImage):
+        return None
+    buffer = io.BytesIO()
+    image_format = str(getattr(value, "format", "") or "").upper()
+    if image_format not in {"PNG", "JPEG", "JPG", "WEBP", "GIF"}:
+        image_format = "PNG"
+    value.save(buffer, format=image_format)
+    return _encode_bytes_to_base64(buffer.getvalue())
+
+
+def _normalize_image_context_value(value: Any, *, base_path: str | None = None) -> Any:
+    if isinstance(value, str):
+        return value
+
+    if isinstance(value, (bytes, bytearray)):
+        return _encode_bytes_to_base64(value)
+
+    pil_base64 = _pil_image_to_base64(value)
+    if pil_base64 is not None:
+        return pil_base64
+
+    if isinstance(value, dict):
+        url = value.get("url")
+        if isinstance(url, str):
+            return url
+
+        image_url = value.get("image_url")
+        if isinstance(image_url, str):
+            return image_url
+        if isinstance(image_url, dict):
+            nested_url = image_url.get("url")
+            if isinstance(nested_url, str):
+                return nested_url
+
+        inline_data = value.get("data")
+        if isinstance(inline_data, str):
+            return inline_data
+
+        raw_bytes = value.get("bytes")
+        if isinstance(raw_bytes, (bytes, bytearray)):
+            return _encode_bytes_to_base64(raw_bytes)
+        if isinstance(raw_bytes, str) and raw_bytes.strip():
+            return raw_bytes
+
+        path_value = value.get("path")
+        if isinstance(path_value, str) and path_value.strip():
+            if as_base64 := _load_image_file_to_base64(path_value, base_path=base_path):
+                return as_base64
+            return path_value
+
+    return value
+
+
+def _apply_data_designer_image_context_patch() -> None:
+    global _IMAGE_CONTEXT_PATCHED
+    if _IMAGE_CONTEXT_PATCHED:
+        return
+
+    try:
+        from data_designer.config.models import ImageContext
+    except ImportError:
+        return
+
+    if getattr(ImageContext, "_unsloth_image_context_patch_applied", False):
+        _IMAGE_CONTEXT_PATCHED = True
+        return
+
+    original_auto_resolve = ImageContext._auto_resolve_context_value
+
+    def _patched_auto_resolve(self: Any, context_value: Any, base_path: str | None) -> Any:
+        normalized = _normalize_image_context_value(context_value, base_path=base_path)
+        return original_auto_resolve(self, normalized, base_path)
+
+    ImageContext._auto_resolve_context_value = _patched_auto_resolve
+    setattr(ImageContext, "_unsloth_image_context_patch_applied", True)
+    _IMAGE_CONTEXT_PATCHED = True
+
 
 def build_model_providers(recipe: dict[str, Any]):
     from data_designer.config.default_model_settings import get_default_providers
@@ -75,6 +188,7 @@ def build_mcp_providers(
 
 
 def build_config_builder(recipe: dict[str, Any]):
+    _apply_data_designer_image_context_patch()
     from data_designer.config import DataDesignerConfigBuilder
     from data_designer.config.processors import ProcessorType
 
@@ -107,6 +221,7 @@ def create_data_designer(
     *,
     artifact_path: str | None = None,
 ):
+    _apply_data_designer_image_context_patch()
     from data_designer.interface.data_designer import DataDesigner
 
     return DataDesigner(

studio/frontend/src/features/recipe-studio/components/executions/execution-data-tab.tsx

@@ -10,10 +10,9 @@ import {
   DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu";
 import { cn } from "@/lib/utils";
-import { resolveImagePreview } from "../../utils/image-preview";
 import { isExecutionInProgress } from "../../executions/execution-helpers";
 import type { RecipeExecutionRecord } from "../../execution-types";
-import { formatCellValue, isExpandableCellValue } from "./executions-view-helpers";
+import { hasExpandableTextCell } from "./executions-view-helpers";
 
 type ExecutionDataTabProps = {
   execution: RecipeExecutionRecord;
@@ -130,10 +129,7 @@ export function ExecutionDataTab({
             columns={tableColumns}
             data={datasetRowsForTable}
             getRowClassName={(row, _rowIndex, rowId) => {
-              const canExpand = visibleDatasetColumnNames.some((columnName) =>
-                !resolveImagePreview(row[columnName]) &&
-                isExpandableCellValue(formatCellValue(row[columnName])),
-              );
+              const canExpand = hasExpandableTextCell(row, visibleDatasetColumnNames);
               if (!canExpand) {
                 return undefined;
               }
@@ -143,10 +139,7 @@ export function ExecutionDataTab({
               );
             }}
             onRowClick={(row, _rowIndex, rowId) => {
-              const canExpand = visibleDatasetColumnNames.some((columnName) =>
-                !resolveImagePreview(row[columnName]) &&
-                isExpandableCellValue(formatCellValue(row[columnName])),
-              );
+              const canExpand = hasExpandableTextCell(row, visibleDatasetColumnNames);
               if (!canExpand || !selectedExecutionIdSafe) {
                 return;
               }

studio/frontend/src/features/recipe-studio/components/executions/executions-view-helpers.ts

@@ -3,6 +3,7 @@ import type {
   RecipeExecutionStatus,
 } from "../../execution-types";
 import { isExecutionInProgress } from "../../executions/execution-helpers";
+import { resolveImagePreview } from "../../utils/image-preview";
 
 export type AnalysisColumnStat = {
   column_name: string;
@@ -55,6 +56,18 @@ export function truncateCellValue(value: string): string {
   return `${value.slice(0, 180).trimEnd()}...`;
 }
 
+export function hasExpandableTextCell(
+  row: Record<string, unknown>,
+  visibleColumnNames: string[],
+): boolean {
+  return visibleColumnNames.some((columnName) => {
+    if (resolveImagePreview(row[columnName])) {
+      return false;
+    }
+    return isExpandableCellValue(formatCellValue(row[columnName]));
+  });
+}
+
 function parseNumber(value: unknown): number | null {
   return typeof value === "number" && Number.isFinite(value) ? value : null;
 }

studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx

@@ -28,7 +28,7 @@ import {
   formatPercent,
   formatStatus,
   formatTimestamp,
-  isExpandableCellValue,
+  hasExpandableTextCell,
   parseAnalysisColumns,
   parseModelUsageRows,
   statusTone,
@@ -144,9 +144,9 @@ export function ExecutionsView({
         }
         const value = formatCellValue(rawValue);
         const rowExpanded = Boolean(expandedDatasetRows[row.id]);
-        const rowHasExpandableCell = visibleDatasetColumnNames.some((columnName) =>
-          !resolveImagePreview(row.original[columnName]) &&
-          isExpandableCellValue(formatCellValue(row.original[columnName])),
+        const rowHasExpandableCell = hasExpandableTextCell(
+          row.original,
+          visibleDatasetColumnNames,
         );
         const showTruncated = rowHasExpandableCell && !rowExpanded;
 

studio/frontend/src/features/recipe-studio/utils/image-preview.ts

@@ -6,6 +6,8 @@ type PreviewImagePayload = {
   data?: unknown;
 };
 
+type UnknownRecord = Record<string, unknown>;
+
 export type ImagePreviewResult =
   | { kind: "ready"; src: string }
   | { kind: "too_large"; estimatedBytes: number };
@@ -56,40 +58,43 @@ function toDataUrlFromBase64(base64: string, mime: string): string {
   return `data:${mime};base64,${normalizeBase64(base64)}`;
 }
 
-function resolveImagePayloadObject(
+function isRecord(value: unknown): value is UnknownRecord {
+  return Boolean(value) && typeof value === "object" && !Array.isArray(value);
+}
+
+function isByteArray(value: unknown): value is number[] {
+  if (!Array.isArray(value) || value.length === 0) {
+    return false;
+  }
+  return value.every(
+    (item) => typeof item === "number" && Number.isInteger(item) && item >= 0 && item <= 255,
+  );
+}
+
+function byteArrayToBase64(bytes: number[]): string {
+  let binary = "";
+  const chunkSize = 0x8000;
+  for (let idx = 0; idx < bytes.length; idx += chunkSize) {
+    const chunk = bytes.slice(idx, idx + chunkSize);
+    binary += String.fromCharCode(...chunk);
+  }
+  return btoa(binary);
+}
+
+function resolveStringCandidate(
   value: unknown,
   maxBytes: number,
 ): ImagePreviewResult | null {
-  if (!value || typeof value !== "object") {
-    return null;
-  }
-  const payload = value as PreviewImagePayload;
-  if (payload.type !== "image" || typeof payload.data !== "string") {
-    return null;
-  }
-  const mime = typeof payload.mime === "string" ? payload.mime : "image/jpeg";
-  const estimatedBytes = estimateBase64Bytes(payload.data);
-  if (estimatedBytes > maxBytes) {
-    return { kind: "too_large", estimatedBytes };
-  }
-  return {
-    kind: "ready",
-    src: toDataUrlFromBase64(payload.data, mime),
-  };
-}
-
-export function resolveImagePreview(
-  value: unknown,
-  maxBytes = MAX_IMAGE_PREVIEW_BYTES,
-): ImagePreviewResult | null {
-  const payloadPreview = resolveImagePayloadObject(value, maxBytes);
-  if (payloadPreview) {
-    return payloadPreview;
-  }
-
   if (typeof value !== "string") {
     return null;
   }
+  return resolveImagePreviewFromString(value, maxBytes);
+}
+
+function resolveImagePreviewFromString(
+  value: string,
+  maxBytes: number,
+): ImagePreviewResult | null {
   const trimmed = value.trim();
   if (!trimmed) {
     return null;
@@ -121,6 +126,80 @@ export function resolveImagePreview(
   return null;
 }
 
+function resolveImagePayloadObject(value: unknown, maxBytes: number): ImagePreviewResult | null {
+  if (!isRecord(value)) {
+    return null;
+  }
+  const payload = value as PreviewImagePayload;
+  if (payload.type === "image" && typeof payload.data === "string") {
+    const mime = typeof payload.mime === "string" ? payload.mime : "image/jpeg";
+    const estimatedBytes = estimateBase64Bytes(payload.data);
+    if (estimatedBytes > maxBytes) {
+      return { kind: "too_large", estimatedBytes };
+    }
+    return {
+      kind: "ready",
+      src: toDataUrlFromBase64(payload.data, mime),
+    };
+  }
+
+  const imageUrl = value.image_url;
+  const directImageUrl = resolveStringCandidate(imageUrl, maxBytes);
+  if (directImageUrl !== null) {
+    return directImageUrl;
+  }
+  if (isRecord(imageUrl)) {
+    const nestedImageUrl = resolveStringCandidate(imageUrl.url, maxBytes);
+    if (nestedImageUrl !== null) {
+      return nestedImageUrl;
+    }
+  }
+
+  const scalarCandidates = [
+    value.url,
+    value.data,
+    value.bytes,
+    value.base64,
+    value.base64_image,
+    value.image,
+    value.path,
+  ];
+  for (const candidate of scalarCandidates) {
+    const resolved = resolveStringCandidate(candidate, maxBytes);
+    if (resolved !== null) {
+      return resolved;
+    }
+  }
+
+  if (isByteArray(value.bytes)) {
+    const resolved = resolveStringCandidate(byteArrayToBase64(value.bytes), maxBytes);
+    if (resolved !== null) {
+      return resolved;
+    }
+  }
+
+  if (isRecord(value.image)) {
+    return resolveImagePayloadObject(value.image, maxBytes);
+  }
+
+  return null;
+}
+
+export function resolveImagePreview(
+  value: unknown,
+  maxBytes = MAX_IMAGE_PREVIEW_BYTES,
+): ImagePreviewResult | null {
+  const payloadPreview = resolveImagePayloadObject(value, maxBytes);
+  if (payloadPreview) {
+    return payloadPreview;
+  }
+
+  if (typeof value !== "string") {
+    return null;
+  }
+  return resolveImagePreviewFromString(value, maxBytes);
+}
+
 export function isLikelyImageValue(value: unknown): boolean {
   return resolveImagePreview(value, Number.POSITIVE_INFINITY) !== null;
 }