From 891739a56ababeb68e8777d2ce927f50ceceeb61 Mon Sep 17 00:00:00 2001 From: Shine1i Date: Sun, 1 Mar 2026 12:01:48 +0100 Subject: [PATCH] feat(recipe-studio): add LLM trace modes and reasoning content extraction support --- studio/backend/core/data_recipe/jsonable.py | 68 +++++++++++++--- .../recipe-studio/dialogs/llm/general-tab.tsx | 78 ++++++++++++++++++- .../src/features/recipe-studio/types/index.ts | 6 ++ .../recipe-studio/utils/config-factories.ts | 4 + .../utils/import/parsers/llm-parser.ts | 15 ++++ .../utils/payload/builders-llm.ts | 6 +- .../recipe-studio/utils/validation.ts | 8 ++ 7 files changed, 172 insertions(+), 13 deletions(-) diff --git a/studio/backend/core/data_recipe/jsonable.py b/studio/backend/core/data_recipe/jsonable.py index a6828b2274..b45a378028 100644 --- a/studio/backend/core/data_recipe/jsonable.py +++ b/studio/backend/core/data_recipe/jsonable.py @@ -2,9 +2,62 @@ from __future__ import annotations import base64 import io +from pathlib import Path from typing import Any +def _pil_to_preview_payload(image: Any) -> dict[str, Any]: + buffer = io.BytesIO() + image.convert("RGB").save(buffer, format="JPEG", quality=85) + return { + "type": "image", + "mime": "image/jpeg", + "width": image.width, + "height": image.height, + "data": base64.b64encode(buffer.getvalue()).decode("ascii"), + } + + +def _open_pil_image_from_bytes(raw_bytes: bytes): + from PIL import Image # type: ignore + + with Image.open(io.BytesIO(raw_bytes)) as image: + return image.copy() + + +def _to_pil_from_hf_image_dict(value: Any) -> Any | None: + if not isinstance(value, dict): + return None + + raw_bytes = value.get("bytes") + if isinstance(raw_bytes, (bytes, bytearray)) and len(raw_bytes) > 0: + try: + return _open_pil_image_from_bytes(bytes(raw_bytes)) + except (OSError, ValueError): + pass + if ( + isinstance(raw_bytes, list) + and len(raw_bytes) > 0 + and all(isinstance(item, int) and 0 <= item <= 255 for item in raw_bytes) + ): + try: + return _open_pil_image_from_bytes(bytes(raw_bytes)) + except (OSError, ValueError): + pass + + path_value = value.get("path") + if isinstance(path_value, str) and path_value.strip(): + try: + from PIL import Image # type: ignore + + with Image.open(Path(path_value)) as image: + return image.copy() + except (OSError, ValueError, TypeError): + return None + + return None + + def to_jsonable(value: Any) -> Any: """Convert numpy/pandas-ish values into plain JSON-safe values.""" try: @@ -39,17 +92,12 @@ def _to_preview_image_payload(value: Any) -> dict[str, Any] | None: return None if not isinstance(value, PILImage): - return None + hf_image = _to_pil_from_hf_image_dict(value) + if hf_image is None: + return None + value = hf_image - buffer = io.BytesIO() - value.convert("RGB").save(buffer, format="JPEG", quality=85) - return { - "type": "image", - "mime": "image/jpeg", - "width": value.width, - "height": value.height, - "data": base64.b64encode(buffer.getvalue()).decode("ascii"), - } + return _pil_to_preview_payload(value) def to_preview_jsonable(value: Any) -> Any: diff --git a/studio/frontend/src/features/recipe-studio/dialogs/llm/general-tab.tsx b/studio/frontend/src/features/recipe-studio/dialogs/llm/general-tab.tsx index 547e8745ce..c974944491 100644 --- a/studio/frontend/src/features/recipe-studio/dialogs/llm/general-tab.tsx +++ b/studio/frontend/src/features/recipe-studio/dialogs/llm/general-tab.tsx @@ -6,6 +6,11 @@ import { ComboboxItem, ComboboxList, } from "@/components/ui/combobox"; +import { + Collapsible, + CollapsibleContent, + CollapsibleTrigger, +} from "@/components/ui/collapsible"; import { Switch } from "@/components/ui/switch"; import { Select, @@ -15,7 +20,7 @@ import { SelectValue, } from "@/components/ui/select"; import { Textarea } from "@/components/ui/textarea"; -import { type ReactElement, type RefObject, useMemo } from "react"; +import { type ReactElement, type RefObject, useMemo, useState } from "react"; import { useRecipeStudioStore } from "../../stores/recipe-studio"; import { isLikelyImageValue } from "../../utils/image-preview"; import type { LlmConfig } from "../../types"; @@ -44,6 +49,15 @@ const CODE_LANG_OPTIONS = [ "sql:ansi", ]; +const TRACE_MODE_OPTIONS = ["none", "last_message", "all_messages"] as const; + +function normalizeTraceMode(value: string): LlmConfig["with_trace"] { + if (value === "last_message" || value === "all_messages") { + return value; + } + return "none"; +} + type LlmGeneralTabProps = { config: LlmConfig; modelConfigAliases: string[]; @@ -120,6 +134,9 @@ export function LlmGeneralTab({ }; const imageContextToggleId = `${config.id}-image-context-enabled`; const imageContextColumnId = `${config.id}-image-context-column`; + const traceModeId = `${config.id}-trace-mode`; + const reasoningToggleId = `${config.id}-reasoning-content`; + const [advancedOpen, setAdvancedOpen] = useState(false); return (
@@ -285,6 +302,65 @@ export function LlmGeneralTab({
)} + + + + + +
+ + +
+
+
+ +
+ + onUpdate({ + // biome-ignore lint/style/useNamingConvention: api schema + extract_reasoning_content: checked, + }) + } + /> +
+ + {config.llm_type === "structured" && (
, name: string, @@ -64,6 +72,9 @@ export function parseLlm( } } + const withTrace = parseTraceMode(column.with_trace); + const extractReasoningContent = column.extract_reasoning_content === true; + return { id, kind: "llm", @@ -82,6 +93,10 @@ export function parseLlm( output_format: normalizeOutputFormat(column.output_format), // biome-ignore lint/style/useNamingConvention: api schema tool_alias: readString(column.tool_alias) ?? "", + // biome-ignore lint/style/useNamingConvention: api schema + with_trace: withTrace, + // biome-ignore lint/style/useNamingConvention: api schema + extract_reasoning_content: extractReasoningContent, scores: llmType === "judge" ? scores : undefined, // biome-ignore lint/style/useNamingConvention: ui schema image_context: imageContext, diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/builders-llm.ts b/studio/frontend/src/features/recipe-studio/utils/payload/builders-llm.ts index 8556d83157..f142a63701 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/builders-llm.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/builders-llm.ts @@ -39,6 +39,10 @@ export function buildLlmColumn( multi_modal_context: buildImageContext(config, errors), // biome-ignore lint/style/useNamingConvention: api schema tool_alias: toolAlias || undefined, + // biome-ignore lint/style/useNamingConvention: api schema + with_trace: config.with_trace ?? "none", + // biome-ignore lint/style/useNamingConvention: api schema + extract_reasoning_content: config.extract_reasoning_content === true, }; if (config.llm_type === "code") { @@ -103,8 +107,6 @@ export function buildLlmColumn( // biome-ignore lint/style/useNamingConvention: api schema column_type: "llm-text", ...base, - // biome-ignore lint/style/useNamingConvention: api schema - with_trace: "none", }; } diff --git a/studio/frontend/src/features/recipe-studio/utils/validation.ts b/studio/frontend/src/features/recipe-studio/utils/validation.ts index d0a683aa42..9950f19937 100644 --- a/studio/frontend/src/features/recipe-studio/utils/validation.ts +++ b/studio/frontend/src/features/recipe-studio/utils/validation.ts @@ -1,6 +1,8 @@ import type { NodeConfig } from "../types"; import { isValidSex, parseAgeRange, parseIntNumber, parseNumber } from "./parse"; +const TRACE_MODES = new Set(["none", "last_message", "all_messages"]); + // biome-ignore lint/complexity/noExcessiveCognitiveComplexity: validation rules export function getConfigErrors(config: NodeConfig | null): string[] { if (!config) { @@ -178,6 +180,12 @@ export function getConfigErrors(config: NodeConfig | null): string[] { errors.push("Image context column is required."); } } + if ( + config.with_trace && + !TRACE_MODES.has(config.with_trace) + ) { + errors.push("Trace mode must be none, last_message, or all_messages."); + } } if (config.kind === "expression") { if (!config.expr.trim()) {