From b7ee065ffd199ddf18d2df6127336d90fafa2122 Mon Sep 17 00:00:00 2001 From: Shine1i Date: Sun, 1 Mar 2026 11:21:10 +0100 Subject: [PATCH] refactor(recipe-studio): add image preview support for dataset and LLM configurations p2 --- studio/backend/core/data_recipe/service.py | 115 +++++++++++++++ .../executions/execution-data-tab.tsx | 13 +- .../executions/executions-view-helpers.ts | 13 ++ .../components/executions/executions-view.tsx | 8 +- .../recipe-studio/utils/image-preview.ts | 135 ++++++++++++++---- 5 files changed, 242 insertions(+), 42 deletions(-) diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index 2d11cb2845..5367c3c0fc 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/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( diff --git a/studio/frontend/src/features/recipe-studio/components/executions/execution-data-tab.tsx b/studio/frontend/src/features/recipe-studio/components/executions/execution-data-tab.tsx index e68a047c85..3d4cc11ab6 100644 --- a/studio/frontend/src/features/recipe-studio/components/executions/execution-data-tab.tsx +++ b/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; } diff --git a/studio/frontend/src/features/recipe-studio/components/executions/executions-view-helpers.ts b/studio/frontend/src/features/recipe-studio/components/executions/executions-view-helpers.ts index b5b12f0075..d75ac4f7ce 100644 --- a/studio/frontend/src/features/recipe-studio/components/executions/executions-view-helpers.ts +++ b/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, + 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; } diff --git a/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx b/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx index 6ebfa6af5b..255d66b72b 100644 --- a/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx +++ b/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; diff --git a/studio/frontend/src/features/recipe-studio/utils/image-preview.ts b/studio/frontend/src/features/recipe-studio/utils/image-preview.ts index 61099b8f4a..0f13bc3411 100644 --- a/studio/frontend/src/features/recipe-studio/utils/image-preview.ts +++ b/studio/frontend/src/features/recipe-studio/utils/image-preview.ts @@ -6,6 +6,8 @@ type PreviewImagePayload = { data?: unknown; }; +type UnknownRecord = Record; + 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; }