From 3e17e2b0f69698ed025e360ae33f0cd65c2028a5 Mon Sep 17 00:00:00 2001 From: Shine1i Date: Mon, 23 Feb 2026 18:46:02 +0100 Subject: [PATCH] refactor: enhance seed source handling with new source types and streamlined inspection flows --- studio/backend/models/data_recipe.py | 6 + studio/backend/routes/data_recipe.py | 89 ++ .../src/features/recipe-studio/api/index.ts | 14 + .../recipe-studio/blocks/definitions.ts | 36 +- .../features/recipe-studio/blocks/registry.ts | 2 +- .../recipe-studio/components/block-sheet.tsx | 12 +- .../components/recipe-graph-node.tsx | 20 +- .../dialogs/seed/seed-dialog.tsx | 802 +++++++++--------- .../recipe-studio/stores/recipe-studio.ts | 42 +- .../src/features/recipe-studio/types/index.ts | 4 + .../recipe-studio/utils/config-factories.ts | 5 + .../import/parsers/seed-config-parser.ts | 29 +- .../features/recipe-studio/utils/node-data.ts | 9 +- .../utils/payload/builders-seed.ts | 24 +- .../recipe-studio/utils/validation.ts | 9 +- 15 files changed, 664 insertions(+), 439 deletions(-) diff --git a/studio/backend/models/data_recipe.py b/studio/backend/models/data_recipe.py index 73eece0b29..9e501c15e2 100644 --- a/studio/backend/models/data_recipe.py +++ b/studio/backend/models/data_recipe.py @@ -45,6 +45,12 @@ class SeedInspectRequest(BaseModel): preview_size: int = Field(default=10, ge=1, le=50) +class SeedInspectUploadRequest(BaseModel): + filename: str = Field(min_length=1) + content_base64: str = Field(min_length=1) + preview_size: int = Field(default=10, ge=1, le=50) + + class SeedInspectResponse(BaseModel): dataset_name: str resolved_path: str diff --git a/studio/backend/routes/data_recipe.py b/studio/backend/routes/data_recipe.py index b70f3c8188..55a2e805b1 100644 --- a/studio/backend/routes/data_recipe.py +++ b/studio/backend/routes/data_recipe.py @@ -4,10 +4,13 @@ Data Recipe routes (DataDesigner runner). from __future__ import annotations +import base64 +import binascii import sys from itertools import islice from pathlib import Path from typing import Any +from uuid import uuid4 from fastapi import APIRouter, HTTPException, Query, Request from fastapi.responses import JSONResponse, StreamingResponse @@ -23,6 +26,7 @@ from models.data_recipe import ( JobCreateResponse, RecipePayload, SeedInspectRequest, + SeedInspectUploadRequest, SeedInspectResponse, ValidateError, ValidateResponse, @@ -31,6 +35,8 @@ from models.data_recipe import ( router = APIRouter() DATA_EXTS = (".parquet", ".jsonl", ".json", ".csv") DEFAULT_SPLIT = "train" +LOCAL_UPLOAD_EXTS = {".csv", ".json", ".jsonl"} +SEED_UPLOAD_DIR = Path.home() / ".cache" / "unsloth" / "data-recipe" / "seed-uploads" def _serialize_preview_value(value: Any) -> Any: @@ -147,6 +153,51 @@ def _extract_columns(rows: list[dict[str, Any]]) -> list[str]: return list(columns_seen.keys()) +def _sanitize_filename(filename: str) -> str: + name = Path(filename).name.strip().replace("\x00", "") + if not name: + return "seed_upload" + return name + + +def _decode_base64_payload(content_base64: str) -> bytes: + raw = content_base64.strip() + if "," in raw and raw.lower().startswith("data:"): + raw = raw.split(",", 1)[1] + try: + return base64.b64decode(raw, validate=True) + except binascii.Error as exc: + raise HTTPException(status_code=400, detail="invalid base64 payload") from exc + + +def _read_preview_rows_from_local_file(path: Path, preview_size: int) -> list[dict[str, Any]]: + try: + import pandas as pd + except Exception as exc: + raise HTTPException(status_code=500, detail=f"seed inspect dependencies unavailable: {exc}") from exc + + ext = path.suffix.lower() + try: + if ext == ".csv": + df = pd.read_csv(path, nrows=preview_size) + elif ext == ".jsonl": + df = pd.read_json(path, lines=True).head(preview_size) + elif ext == ".json": + try: + df = pd.read_json(path, lines=True).head(preview_size) + except Exception: + df = pd.read_json(path).head(preview_size) + else: + raise HTTPException(status_code=422, detail=f"unsupported file type: {ext}") + except HTTPException: + raise + except Exception as exc: + raise HTTPException(status_code=422, detail=f"seed inspect failed: {exc}") from exc + + rows = df.to_dict(orient="records") + return _serialize_preview_rows(rows) + + @router.post("/seed/inspect", response_model=SeedInspectResponse) def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse: dataset_name = payload.dataset_name.strip() @@ -223,6 +274,44 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse: ) +@router.post("/seed/inspect-upload", response_model=SeedInspectResponse) +def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectResponse: + filename = _sanitize_filename(payload.filename) + ext = Path(filename).suffix.lower() + if ext not in LOCAL_UPLOAD_EXTS: + allowed = ", ".join(sorted(LOCAL_UPLOAD_EXTS)) + raise HTTPException(status_code=400, detail=f"unsupported file type: {ext}. allowed: {allowed}") + + file_bytes = _decode_base64_payload(payload.content_base64) + if not file_bytes: + raise HTTPException(status_code=400, detail="empty upload payload") + max_size_bytes = 50 * 1024 * 1024 + if len(file_bytes) > max_size_bytes: + raise HTTPException(status_code=413, detail="file too large (max 50MB)") + + SEED_UPLOAD_DIR.mkdir(parents=True, exist_ok=True) + stored_name = f"{uuid4().hex}_{filename}" + stored_path = SEED_UPLOAD_DIR / stored_name + stored_path.write_bytes(file_bytes) + + preview_rows = _read_preview_rows_from_local_file( + stored_path, + int(payload.preview_size), + ) + if not preview_rows: + raise HTTPException(status_code=422, detail="dataset appears empty or unreadable") + columns = _extract_columns(preview_rows) + + return SeedInspectResponse( + dataset_name=filename, + resolved_path=str(stored_path), + columns=columns, + preview_rows=preview_rows, + split=None, + subset=None, + ) + + @router.post("/validate", response_model=ValidateResponse) def validate(payload: RecipePayload) -> ValidateResponse: recipe = payload.recipe diff --git a/studio/frontend/src/features/recipe-studio/api/index.ts b/studio/frontend/src/features/recipe-studio/api/index.ts index 1dac98bbe9..7d2837abd0 100644 --- a/studio/frontend/src/features/recipe-studio/api/index.ts +++ b/studio/frontend/src/features/recipe-studio/api/index.ts @@ -81,6 +81,14 @@ export type SeedInspectRequest = { preview_size?: number; }; +export type SeedInspectUploadRequest = { + filename: string; + // base64 payload without data URL prefix + content_base64: string; + // biome-ignore lint/style/useNamingConvention: api schema + preview_size?: number; +}; + export type SeedInspectResponse = { // biome-ignore lint/style/useNamingConvention: api schema dataset_name: string; @@ -232,6 +240,12 @@ export async function inspectSeedDataset( return postJson("/seed/inspect", payload); } +export async function inspectSeedUpload( + payload: SeedInspectUploadRequest, +): Promise { + return postJson("/seed/inspect-upload", payload); +} + export async function streamRecipeJobEvents(options: { jobId: string; signal: AbortSignal; diff --git a/studio/frontend/src/features/recipe-studio/blocks/definitions.ts b/studio/frontend/src/features/recipe-studio/blocks/definitions.ts index d91f5d6991..db78810734 100644 --- a/studio/frontend/src/features/recipe-studio/blocks/definitions.ts +++ b/studio/frontend/src/features/recipe-studio/blocks/definitions.ts @@ -15,7 +15,7 @@ import { TagsIcon, UserAccountIcon, } from "@hugeicons/core-free-icons"; -import type { LlmType, NodeConfig, SamplerType } from "../types"; +import type { LlmType, NodeConfig, SamplerType, SeedSourceType } from "../types"; import { makeExpressionConfig, makeLlmConfig, @@ -31,9 +31,14 @@ export type BlockType = | LlmType | "expression" | "seed" + | "seed_hf" + | "seed_local" + | "seed_unstructured" | "model_provider" | "model_config"; +export type SeedBlockType = "seed_hf" | "seed_local" | "seed_unstructured"; + type IconType = typeof CodeIcon; export type BlockGroup = { @@ -99,12 +104,30 @@ export const BLOCK_GROUPS: BlockGroup[] = [ const BLOCK_DEFINITIONS: BlockDefinition[] = [ { kind: "seed", - type: "seed", + type: "seed_hf", title: "Seed (Hugging Face)", description: "Load real rows from HF and use them as generation context.", icon: Plant01Icon, dialogKey: "seed", - createConfig: (id, existing) => makeSeedConfig(id, existing), + createConfig: (id, existing) => makeSeedConfig(id, existing, "hf"), + }, + { + kind: "seed", + type: "seed_local", + title: "Seed (Local File)", + description: "Upload CSV/JSON/JSONL and use rows as seed context.", + icon: Plant01Icon, + dialogKey: "seed", + createConfig: (id, existing) => makeSeedConfig(id, existing, "local"), + }, + { + kind: "seed", + type: "seed_unstructured", + title: "Seed (Unstructured)", + description: "Upload PDF/DOCX/TXT, chunk to text rows, then seed.", + icon: Plant01Icon, + dialogKey: "seed", + createConfig: (id, existing) => makeSeedConfig(id, existing, "unstructured"), }, { kind: "sampler", @@ -273,7 +296,12 @@ export function getBlockDefinitionForConfig( return null; } if (config.kind === "seed") { - return getBlockDefinition("seed", "seed"); + const seedType: Record = { + hf: "seed_hf", + local: "seed_local", + unstructured: "seed_unstructured", + }; + return getBlockDefinition("seed", seedType[config.seed_source_type ?? "hf"]); } if (config.kind === "sampler") { const samplerType = diff --git a/studio/frontend/src/features/recipe-studio/blocks/registry.ts b/studio/frontend/src/features/recipe-studio/blocks/registry.ts index c6798c20dc..3551981695 100644 --- a/studio/frontend/src/features/recipe-studio/blocks/registry.ts +++ b/studio/frontend/src/features/recipe-studio/blocks/registry.ts @@ -4,6 +4,7 @@ export type { BlockGroup, BlockKind, BlockType, + SeedBlockType, } from "./definitions"; export { BLOCK_GROUPS, @@ -12,4 +13,3 @@ export { getBlocksForKind, } from "./definitions"; export { renderBlockDialog } from "./render-dialog"; - diff --git a/studio/frontend/src/features/recipe-studio/components/block-sheet.tsx b/studio/frontend/src/features/recipe-studio/components/block-sheet.tsx index 66ecb678d8..f602fc55c7 100644 --- a/studio/frontend/src/features/recipe-studio/components/block-sheet.tsx +++ b/studio/frontend/src/features/recipe-studio/components/block-sheet.tsx @@ -20,7 +20,11 @@ import { HugeiconsIcon } from "@hugeicons/react"; import { type ReactElement, useMemo, useState } from "react"; import { RECIPE_FLOATING_ICON_BUTTON_CLASS } from "./recipe-floating-icon-button-class"; import type { LlmType, SamplerType } from "../types"; -import { BLOCK_GROUPS, getBlocksForKind } from "../blocks/registry"; +import { + BLOCK_GROUPS, + getBlocksForKind, + type SeedBlockType, +} from "../blocks/registry"; type SheetView = "root" | "sampler" | "seed" | "llm" | "expression" | "processor"; type SheetKind = "sampler" | "seed" | "llm" | "expression"; @@ -39,7 +43,7 @@ type BlockSheetProps = { open?: boolean; onOpenChange?: (open: boolean) => void; onAddSampler: (type: SamplerType) => void; - onAddSeed: () => void; + onAddSeed: (type: SeedBlockType) => void; onAddLlm: (type: LlmType) => void; onAddModelProvider: () => void; onAddModelConfig: () => void; @@ -221,7 +225,7 @@ export function BlockSheet({ } if (item.kind === "seed" && seedBlocks.length === 1) { setSheetOpen(false); - onAddSeed(); + onAddSeed(seedBlocks[0].type as SeedBlockType); return; } if (item.kind === "expression" && expressionBlocks.length === 1) { @@ -257,7 +261,7 @@ export function BlockSheet({ if (item.kind === "sampler") { onAddSampler(item.type as SamplerType); } else if (item.kind === "seed") { - onAddSeed(); + onAddSeed(item.type as SeedBlockType); } else if (item.kind === "llm") { if (item.type === "model_provider") { onAddModelProvider(); diff --git a/studio/frontend/src/features/recipe-studio/components/recipe-graph-node.tsx b/studio/frontend/src/features/recipe-studio/components/recipe-graph-node.tsx index 8967a37dcb..b571fec32d 100644 --- a/studio/frontend/src/features/recipe-studio/components/recipe-graph-node.tsx +++ b/studio/frontend/src/features/recipe-studio/components/recipe-graph-node.tsx @@ -172,13 +172,29 @@ function getConfigSummary(config: NodeConfig | undefined): string { } if (config.kind === "seed") { - if (config.hf_repo_id.trim()) { + const seedSourceType = config.seed_source_type ?? "hf"; + if (seedSourceType === "hf" && config.hf_repo_id.trim()) { return config.hf_repo_id.trim(); } + if (seedSourceType === "local" && config.local_file_name?.trim()) { + return config.local_file_name.trim(); + } + if ( + seedSourceType === "unstructured" && + config.unstructured_file_name?.trim() + ) { + return config.unstructured_file_name.trim(); + } if (config.hf_path.trim()) { return config.hf_path.trim(); } - return "Set HF dataset repo"; + if (seedSourceType === "hf") { + return "Set HF dataset repo"; + } + if (seedSourceType === "local") { + return "Upload CSV/JSON file"; + } + return "Upload PDF/DOCX/TXT file"; } return "Open details for config"; diff --git a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx index ab6e0137b0..f3acf51e73 100644 --- a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx +++ b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx @@ -1,12 +1,4 @@ import { Button } from "@/components/ui/button"; -import { - Combobox, - ComboboxContent, - ComboboxEmpty, - ComboboxInput, - ComboboxItem, - ComboboxList, -} from "@/components/ui/combobox"; import { Collapsible, CollapsibleContent, @@ -20,7 +12,6 @@ import { EmptyTitle, } from "@/components/ui/empty"; import { Input } from "@/components/ui/input"; -import { InputGroupAddon } from "@/components/ui/input-group"; import { Select, SelectContent, @@ -28,7 +19,6 @@ import { SelectTrigger, SelectValue, } from "@/components/ui/select"; -import { Spinner } from "@/components/ui/spinner"; import { Table, TableBody, @@ -43,18 +33,10 @@ import { TabsList, TabsTrigger, } from "@/components/ui/tabs"; -import { - useDebouncedValue, - useHfDatasetSearch, - useHfDatasetSplits, - useHfTokenValidation, - useInfiniteScroll, -} from "@/hooks"; -import { formatCompact } from "@/lib/utils"; -import { Search01Icon } from "@hugeicons/core-free-icons"; -import { HugeiconsIcon } from "@hugeicons/react"; -import { type ReactElement, useEffect, useMemo, useRef, useState } from "react"; -import { inspectSeedDataset } from "../../api"; +import mammoth from "mammoth"; +import { type ReactElement, useEffect, useMemo, useState } from "react"; +import { extractText, getDocumentProxy } from "unpdf"; +import { inspectSeedDataset, inspectSeedUpload } from "../../api"; import type { SeedConfig, SeedSamplingStrategy, @@ -73,6 +55,12 @@ const SELECTION_OPTIONS: Array<{ value: SeedSelectionType; label: string }> = [ { value: "partition_block", label: "Partition block" }, ]; +const LOCAL_ACCEPT = ".csv,.json,.jsonl"; +const UNSTRUCTURED_ACCEPT = ".txt,.pdf,.docx"; +const MAX_UPLOAD_BYTES = 50 * 1024 * 1024; +const CHUNK_SIZE = 1200; +const CHUNK_OVERLAP = 200; + type SeedDialogProps = { config: SeedConfig; onUpdate: (patch: Partial) => void; @@ -85,23 +73,6 @@ function getErrorMessage(error: unknown, fallback: string): string { return fallback; } -function getPreferredSplit( - availableSplits: string[], - currentSplit: string | undefined, -): string | null { - if (availableSplits.length === 0) return null; - if (availableSplits.length === 1) { - return availableSplits[0] === currentSplit ? null : availableSplits[0]; - } - if (!currentSplit && availableSplits.includes("train")) { - return "train"; - } - if (!currentSplit) { - return availableSplits[0]; - } - return null; -} - function stringifyCell(value: unknown): string { if (value === null || value === undefined) return ""; if (typeof value === "string") return value; @@ -113,123 +84,168 @@ function stringifyCell(value: unknown): string { } } +function chunkText(input: string): string[] { + const text = input.replace(/\r\n/g, "\n").replace(/\n{3,}/g, "\n\n").trim(); + if (!text) return []; + + const chunks: string[] = []; + let cursor = 0; + const overlap = Math.max(0, Math.min(CHUNK_OVERLAP, CHUNK_SIZE - 1)); + while (cursor < text.length) { + const end = Math.min(cursor + CHUNK_SIZE, text.length); + chunks.push(text.slice(cursor, end).trim()); + if (end >= text.length) break; + cursor = Math.max(0, end - overlap); + } + return chunks.filter(Boolean); +} + +async function fileToBase64Payload(file: File): Promise { + return new Promise((resolve, reject) => { + const reader = new FileReader(); + reader.onload = () => { + const value = String(reader.result ?? ""); + const parts = value.split(","); + resolve(parts.length > 1 ? parts[1] : value); + }; + reader.onerror = () => reject(new Error("Failed to read file")); + reader.readAsDataURL(file); + }); +} + +async function extractUnstructuredText(file: File): Promise { + const lower = file.name.toLowerCase(); + if (lower.endsWith(".txt")) { + return file.text(); + } + if (lower.endsWith(".pdf")) { + const buffer = new Uint8Array(await file.arrayBuffer()); + const pdf = await getDocumentProxy(buffer); + const { text } = await extractText(pdf, { mergePages: true }); + return text; + } + if (lower.endsWith(".docx")) { + const arrayBuffer = await file.arrayBuffer(); + const { value } = await mammoth.extractRawText({ arrayBuffer }); + return value; + } + throw new Error("Unsupported unstructured file type"); +} + export function SeedDialog({ config, onUpdate }: SeedDialogProps): ReactElement { - const [inputValue, setInputValue] = useState(""); const [inspectError, setInspectError] = useState(null); const [isInspecting, setIsInspecting] = useState(false); const [advancedOpen, setAdvancedOpen] = useState(false); const [previewRows, setPreviewRows] = useState[]>([]); + const [localFile, setLocalFile] = useState(null); + const [unstructuredFile, setUnstructuredFile] = useState(null); - const selectingRef = useRef(false); - const comboboxAnchorRef = useRef(null); - const debouncedQuery = useDebouncedValue(inputValue); - - const hfToken = config.hf_token?.trim() ?? ""; - const dataset = config.hf_repo_id.trim(); - - const { - results: hfResults, - isLoading, - isLoadingMore, - fetchMore, - error: hfSearchError, - } = useHfDatasetSearch(debouncedQuery, { - accessToken: hfToken || undefined, - }); - - const { error: tokenValidationError, isChecking: isCheckingToken } = - useHfTokenValidation(hfToken); - - const { scrollRef, sentinelRef } = useInfiniteScroll(fetchMore, hfResults.length); - - const { - subsets, - splits, - hasMultipleSubsets, - hasMultipleSplits, - isLoading: splitsLoading, - error: splitsError, - } = useHfDatasetSplits(dataset || null, config.hf_subset || null, { - accessToken: hfToken || undefined, - }); + const mode = config.seed_source_type ?? "hf"; useEffect(() => { - if (subsets.length === 1 && config.hf_subset !== subsets[0]) { - onUpdate({ hf_subset: subsets[0] }); - } - }, [subsets, config.hf_subset, onUpdate]); - - useEffect(() => { - if (splits.length === 0) return; - if (hasMultipleSubsets && !config.hf_subset) return; - const preferredSplit = getPreferredSplit(splits, config.hf_split); - if (preferredSplit) onUpdate({ hf_split: preferredSplit }); - }, [ - splits, - hasMultipleSubsets, - config.hf_subset, - config.hf_split, - onUpdate, - ]); - - const resultIds = useMemo(() => { - const ids = hfResults.map((result) => result.id); - if (dataset && !ids.includes(dataset)) { - ids.push(dataset); - } - return ids; - }, [hfResults, dataset]); + setInspectError(null); + setPreviewRows([]); + setLocalFile(null); + setUnstructuredFile(null); + }, [mode]); const samplingId = `${config.id}-sampling`; const selectionId = `${config.id}-selection`; const tokenId = `${config.id}-hf-token`; const subsetId = `${config.id}-hf-subset`; const splitId = `${config.id}-hf-split`; - - function handleDatasetSelect(id: string | null): void { - selectingRef.current = true; - const value = id ?? ""; - onUpdate({ - hf_repo_id: value, - hf_subset: "", - hf_split: "", - hf_path: "", - seed_columns: [], - }); - setInspectError(null); - setPreviewRows([]); - } - - function handleInputChange(value: string): void { - if (selectingRef.current) { - selectingRef.current = false; - return; - } - setInputValue(value); - } + const datasetId = `${config.id}-hf-dataset`; async function loadSeedMetadata(): Promise { - const datasetName = config.hf_repo_id.trim(); - if (!datasetName) { - setInspectError("Select a dataset first."); - return; - } - setInspectError(null); setIsInspecting(true); try { - const response = await inspectSeedDataset({ - dataset_name: datasetName, - hf_token: hfToken || undefined, - subset: config.hf_subset || undefined, - split: config.hf_split || "train", + if (mode === "hf") { + const datasetName = config.hf_repo_id.trim(); + if (!datasetName) { + throw new Error("Dataset repo is required."); + } + const response = await inspectSeedDataset({ + dataset_name: datasetName, + hf_token: config.hf_token?.trim() || undefined, + subset: config.hf_subset || undefined, + split: config.hf_split || "train", + preview_size: 10, + }); + onUpdate({ + hf_path: response.resolved_path, + seed_columns: response.columns, + hf_split: response.split ?? config.hf_split ?? "", + hf_subset: response.subset ?? config.hf_subset ?? "", + local_file_name: "", + unstructured_file_name: "", + }); + setPreviewRows(response.preview_rows ?? []); + return; + } + + if (mode === "local") { + if (!localFile) { + throw new Error("Select a local CSV/JSON/JSONL file first."); + } + if (localFile.size > MAX_UPLOAD_BYTES) { + throw new Error("File too large (max 50MB)."); + } + const payload = await fileToBase64Payload(localFile); + const response = await inspectSeedUpload({ + filename: localFile.name, + content_base64: payload, + preview_size: 10, + }); + onUpdate({ + hf_path: response.resolved_path, + seed_columns: response.columns, + hf_repo_id: "", + hf_subset: "", + hf_split: "", + local_file_name: localFile.name, + unstructured_file_name: "", + }); + setPreviewRows(response.preview_rows ?? []); + return; + } + + if (!unstructuredFile) { + throw new Error("Select a PDF/DOCX/TXT file first."); + } + if (unstructuredFile.size > MAX_UPLOAD_BYTES) { + throw new Error("File too large (max 50MB)."); + } + + const text = await extractUnstructuredText(unstructuredFile); + const chunks = chunkText(text); + if (chunks.length === 0) { + throw new Error("No text found in file."); + } + const jsonl = chunks + .map((chunk) => JSON.stringify({ chunk_text: chunk })) + .join("\n"); + const stem = + unstructuredFile.name.replace(/\.(pdf|docx|txt)$/i, "") || + "unstructured_seed"; + const jsonlFile = new File([jsonl], `${stem}.jsonl`, { + type: "application/json", + }); + + const payload = await fileToBase64Payload(jsonlFile); + const response = await inspectSeedUpload({ + filename: jsonlFile.name, + content_base64: payload, preview_size: 10, }); onUpdate({ hf_path: response.resolved_path, seed_columns: response.columns, - hf_split: response.split ?? config.hf_split ?? "", - hf_subset: response.subset ?? config.hf_subset ?? "", + hf_repo_id: "", + hf_subset: "", + hf_split: "", + local_file_name: "", + unstructured_file_name: unstructuredFile.name, }); setPreviewRows(response.preview_rows ?? []); } catch (error) { @@ -247,6 +263,13 @@ export function SeedDialog({ config, onUpdate }: SeedDialogProps): ReactElement return []; }, [config.seed_columns, previewRows]); + const canLoad = + mode === "hf" + ? Boolean(config.hf_repo_id.trim()) + : mode === "local" + ? Boolean(localFile) + : Boolean(unstructuredFile); + return ( @@ -256,168 +279,139 @@ export function SeedDialog({ config, onUpdate }: SeedDialogProps): ReactElement
-
- -
{ - if (event.key !== "Enter") return; - if (!(event.target instanceof HTMLInputElement)) return; - event.preventDefault(); - if (hfResults.length > 0) { - handleDatasetSelect(hfResults[0].id); - } else { - const text = event.target.value.trim(); - if (text) handleDatasetSelect(text); - } - }} - > - id} - autoHighlight={true} - > - - - - - - - {isLoading ? ( -
- Searching... -
- ) : ( - No datasets found - )} -
- - {(id: string) => { - const result = hfResults.find((entry) => entry.id === id); - let detail: string | null = null; - if (result?.totalExamples) { - detail = `${formatCompact(result.totalExamples)} rows`; - } else if (result?.sizeCategory) { - detail = result.sizeCategory; - } else if (result?.downloads != null) { - detail = `↓${formatCompact(result.downloads)}`; - } - return ( - - {id} - {detail && ( - {detail} - )} - - ); - }} - -
- {isLoadingMore && ( -
- -
- )} -
- - -
-
+ {mode === "hf" && ( + <> +
+ + + onUpdate({ + hf_repo_id: event.target.value, + hf_path: "", + seed_columns: [], + }) + } + /> +
-
- - onUpdate({ hf_token: event.target.value })} - /> - {(tokenValidationError ?? hfSearchError) && ( -

{tokenValidationError ?? hfSearchError}

- )} - {isCheckingToken && ( -

Checking token…

- )} -
+
+ + onUpdate({ hf_token: event.target.value })} + /> +
- {splitsLoading && dataset ? ( -
- Loading subsets and splits... -
- ) : null} +
+
+ + onUpdate({ hf_subset: event.target.value })} + /> +
+
+ + onUpdate({ hf_split: event.target.value })} + /> +
+
+ + )} - {splitsError ? ( -

- Could not fetch dataset splits: {splitsError} -

- ) : null} - - {hasMultipleSubsets && ( + {mode === "local" && (
- + { + const file = event.target.files?.[0] ?? null; + setLocalFile(file); + onUpdate({ hf_path: "", seed_columns: [] }); + }} + /> +

+ Upload-only. Max 50MB. +

+ {localFile && ( +

Selected: {localFile.name}

+ )}
)} - {hasMultipleSplits && ( + {mode === "unstructured" && (
- + { + const file = event.target.files?.[0] ?? null; + setUnstructuredFile(file); + onUpdate({ hf_path: "", seed_columns: [] }); + }} + /> +

+ Chunking uses chunk_text only. Max 50MB. +

+ {unstructuredFile && ( +

+ Selected: {unstructuredFile.name} +

+ )}
)} + + {inspectError &&

{inspectError}

} @@ -431,103 +425,103 @@ export function SeedDialog({ config, onUpdate }: SeedDialogProps): ReactElement -
- - -
+
+ + +
-
- - -
+
+ + +
- {config.selection_type === "index_range" && ( -
-
- - onUpdate({ selection_start: event.target.value })} - /> -
-
- - onUpdate({ selection_end: event.target.value })} - /> -
+ {config.selection_type === "index_range" && ( +
+
+ + onUpdate({ selection_start: event.target.value })} + />
- )} - - {config.selection_type === "partition_block" && ( -
-
- - onUpdate({ selection_index: event.target.value })} - /> -
-
- - - onUpdate({ selection_num_partitions: event.target.value }) - } - /> -
+
+ + onUpdate({ selection_end: event.target.value })} + />
- )} +
+ )} + + {config.selection_type === "partition_block" && ( +
+
+ + onUpdate({ selection_index: event.target.value })} + /> +
+
+ + + onUpdate({ selection_num_partitions: event.target.value }) + } + /> +
+
+ )}
@@ -541,7 +535,7 @@ export function SeedDialog({ config, onUpdate }: SeedDialogProps): ReactElement Seed preview - Click load to fetch 10 rows from the selected dataset. + Click load to fetch 10 rows from the selected seed source. @@ -550,7 +544,7 @@ export function SeedDialog({ config, onUpdate }: SeedDialogProps): ReactElement variant="outline" className="nodrag" onClick={() => void loadSeedMetadata()} - disabled={isInspecting || !dataset} + disabled={isInspecting || !canLoad} > {isInspecting ? "Loading..." : "Load 10 rows"} @@ -558,43 +552,39 @@ export function SeedDialog({ config, onUpdate }: SeedDialogProps): ReactElement
) : ( -
-
- +
+
+ Loaded columns: {previewColumns.join(", ") || "None"}
- - - - {previewColumns.map((column) => ( - - {column} - - ))} - - - - {previewRows.map((row, index) => ( - - {previewColumns.map((column) => ( - -
{stringifyCell(row[column])}
- +
+
+ + + {previewColumns.map((col) => ( + + {col} + ))} - ))} - -
+ + + {previewRows.map((row, rowIdx) => ( + + {previewColumns.map((col) => ( + + {stringifyCell(row[col])} + + ))} + + ))} + + +
)} - {inspectError &&

{inspectError}

}
diff --git a/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts b/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts index d583b7dd01..cfac089071 100644 --- a/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts +++ b/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts @@ -15,12 +15,14 @@ import type { LayoutDirection, LlmType, NodeConfig, + SeedSourceType, SamplerType, } from "../types"; import { getBlockDefinition, type BlockKind, type BlockType, + type SeedBlockType, } from "../blocks/registry"; import { deriveDisplayGraph } from "../utils/graph/derive-display-graph"; import { applyRecipeConnection, isValidRecipeConnection } from "../utils/graph"; @@ -64,7 +66,7 @@ type RecipeStudioState = { applyLayout: () => void; setLlmAuxVisibility: (id: string, visible: boolean) => void; addSamplerNode: (type: SamplerType) => void; - addSeedNode: () => void; + addSeedNode: (type: SeedBlockType) => void; addLlmNode: (type: LlmType) => void; addModelProviderNode: () => void; addModelConfigNode: () => void; @@ -290,21 +292,47 @@ export const useRecipeStudioStore = create((set, get) => ({ }), addSamplerNode: (type) => set((state) => buildAddedNodeState(state, "sampler", type)), - addSeedNode: () => + addSeedNode: (type) => set((state) => { const existing = Object.values(state.configs).find( (config) => config.kind === "seed", ); if (!existing) { - return buildAddedNodeState(state, "seed", "seed"); + return buildAddedNodeState(state, "seed", type); } + const nextSourceType: SeedSourceType = + type === "seed_local" + ? "local" + : type === "seed_unstructured" + ? "unstructured" + : "hf"; + + const nextConfig = { + ...existing, + seed_source_type: nextSourceType, + hf_repo_id: "", + hf_subset: "", + hf_split: "", + hf_path: "", + hf_token: "", + hf_endpoint: "https://huggingface.co", + local_file_name: "", + unstructured_file_name: "", + seed_columns: [], + }; return { + configs: { + ...state.configs, + [existing.id]: nextConfig, + }, + nodes: updateNodeData( + state.nodes.map((node) => ({ ...node, selected: node.id === existing.id })), + existing.id, + nextConfig, + state.layoutDirection, + ), activeConfigId: existing.id, dialogOpen: true, - nodes: state.nodes.map((node) => ({ - ...node, - selected: node.id === existing.id, - })), }; }), addLlmNode: (type) => set((state) => buildAddedNodeState(state, "llm", type)), diff --git a/studio/frontend/src/features/recipe-studio/types/index.ts b/studio/frontend/src/features/recipe-studio/types/index.ts index 45c23c4954..c019f393a8 100644 --- a/studio/frontend/src/features/recipe-studio/types/index.ts +++ b/studio/frontend/src/features/recipe-studio/types/index.ts @@ -20,6 +20,7 @@ export type LayoutDirection = "LR" | "TB"; export type SeedSamplingStrategy = "ordered" | "shuffle"; export type SeedSelectionType = "none" | "index_range" | "partition_block"; +export type SeedSourceType = "hf" | "local" | "unstructured"; export type RecipeNodeData = { title: string; @@ -219,6 +220,7 @@ export type SeedConfig = { kind: "seed"; name: string; drop?: boolean; + seed_source_type: SeedSourceType; // ui-only (serialized in seed_config) hf_repo_id: string; hf_subset?: string; @@ -226,6 +228,8 @@ export type SeedConfig = { hf_path: string; hf_token?: string; hf_endpoint?: string; + local_file_name?: string; + unstructured_file_name?: string; seed_splits?: string[]; // ui-only // biome-ignore lint/style/useNamingConvention: ui schema diff --git a/studio/frontend/src/features/recipe-studio/utils/config-factories.ts b/studio/frontend/src/features/recipe-studio/utils/config-factories.ts index ad2ed6a3e7..c987d39644 100644 --- a/studio/frontend/src/features/recipe-studio/utils/config-factories.ts +++ b/studio/frontend/src/features/recipe-studio/utils/config-factories.ts @@ -6,6 +6,7 @@ import type { ModelProviderConfig, NodeConfig, SeedConfig, + SeedSourceType, SamplerConfig, SamplerType, } from "../types"; @@ -279,18 +280,22 @@ export function makeExpressionConfig( export function makeSeedConfig( id: string, existing: NodeConfig[], + seedSourceType: SeedSourceType = "hf", ): SeedConfig { return { id, kind: "seed", name: nextName(existing, "seed"), drop: false, + seed_source_type: seedSourceType, hf_repo_id: "", hf_subset: "", hf_split: "", hf_path: "", hf_token: "", hf_endpoint: "https://huggingface.co", + local_file_name: "", + unstructured_file_name: "", seed_splits: [], seed_globs_by_split: {}, seed_columns: [], diff --git a/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts b/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts index 15df0bf36b..f13ad40b70 100644 --- a/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts +++ b/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts @@ -2,6 +2,7 @@ import type { SeedConfig, SeedSamplingStrategy, SeedSelectionType, + SeedSourceType, } from "../../../types"; import { isRecord, readString } from "../helpers"; @@ -17,12 +18,15 @@ function makeDefaultSeedConfig(id: string): SeedConfig { kind: "seed", name: "seed", drop: false, + seed_source_type: "hf", hf_repo_id: "", hf_subset: "", hf_split: "", hf_path: "", hf_token: "", hf_endpoint: "https://huggingface.co", + local_file_name: "", + unstructured_file_name: "", seed_splits: [], seed_globs_by_split: {}, seed_columns: [], @@ -55,16 +59,28 @@ function parseSeedSettings(seedConfigRaw: unknown): Partial { const sampling_strategy = normalizeSampling(seedConfigRaw.sampling_strategy); + let seed_source_type: SeedSourceType = "hf"; let hf_path = ""; let hf_token = ""; let hf_endpoint = "https://huggingface.co"; let hf_repo_id = ""; + let local_file_name = ""; + let unstructured_file_name = ""; const sourceRaw = seedConfigRaw.source; - if (isRecord(sourceRaw) && readString(sourceRaw.seed_type) === "hf") { - hf_path = readString(sourceRaw.path) ?? ""; - hf_token = readString(sourceRaw.token) ?? ""; - hf_endpoint = readString(sourceRaw.endpoint) ?? hf_endpoint; - hf_repo_id = inferRepoIdFromSeedPath(hf_path); + if (isRecord(sourceRaw)) { + const seedType = readString(sourceRaw.seed_type); + const sourcePath = readString(sourceRaw.path) ?? ""; + if (seedType === "hf") { + seed_source_type = "hf"; + hf_path = sourcePath; + hf_token = readString(sourceRaw.token) ?? ""; + hf_endpoint = readString(sourceRaw.endpoint) ?? hf_endpoint; + hf_repo_id = inferRepoIdFromSeedPath(hf_path); + } else if (seedType === "local") { + seed_source_type = "local"; + hf_path = sourcePath; + local_file_name = sourcePath.split("/").pop() ?? sourcePath; + } } let selection_type: SeedSelectionType = "none"; @@ -92,10 +108,13 @@ function parseSeedSettings(seedConfigRaw: unknown): Partial { } return { + seed_source_type, hf_repo_id, hf_path, hf_token, hf_endpoint, + local_file_name, + unstructured_file_name, sampling_strategy, selection_type, selection_start, diff --git a/studio/frontend/src/features/recipe-studio/utils/node-data.ts b/studio/frontend/src/features/recipe-studio/utils/node-data.ts index 79082b42df..dbf2cdca54 100644 --- a/studio/frontend/src/features/recipe-studio/utils/node-data.ts +++ b/studio/frontend/src/features/recipe-studio/utils/node-data.ts @@ -30,10 +30,17 @@ export function nodeDataFromConfig( }; } if (config.kind === "seed") { + const seedSourceType = config.seed_source_type ?? "hf"; + const subtype = + seedSourceType === "hf" + ? "Hugging Face" + : seedSourceType === "local" + ? "Local File" + : "Unstructured"; return { title: "Seed", kind: "seed", - subtype: "Hugging Face", + subtype, blockType: "seed", name: config.name, layoutDirection, diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/builders-seed.ts b/studio/frontend/src/features/recipe-studio/utils/payload/builders-seed.ts index 240932280b..2afd02af1b 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/builders-seed.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/builders-seed.ts @@ -12,6 +12,7 @@ export function buildSeedConfig( config: SeedConfig, errors: string[], ): Record | undefined { + const seedSourceType = config.seed_source_type ?? "hf"; const path = config.hf_path.trim(); if (!path) { return undefined; @@ -40,14 +41,23 @@ export function buildSeedConfig( selectionStrategy = { index, num_partitions: numPartitions }; } + const source = + seedSourceType === "hf" + ? { + // biome-ignore lint/style/useNamingConvention: api schema + seed_type: "hf", + path, + token, + endpoint, + } + : { + // biome-ignore lint/style/useNamingConvention: api schema + seed_type: "local", + path, + }; + return { - source: { - // biome-ignore lint/style/useNamingConvention: api schema - seed_type: "hf", - path, - token, - endpoint, - }, + source, // biome-ignore lint/style/useNamingConvention: api schema sampling_strategy: config.sampling_strategy, // biome-ignore lint/style/useNamingConvention: api schema diff --git a/studio/frontend/src/features/recipe-studio/utils/validation.ts b/studio/frontend/src/features/recipe-studio/utils/validation.ts index b8e8594f2c..c95a2eda1a 100644 --- a/studio/frontend/src/features/recipe-studio/utils/validation.ts +++ b/studio/frontend/src/features/recipe-studio/utils/validation.ts @@ -180,13 +180,18 @@ export function getConfigErrors(config: NodeConfig | null): string[] { } } if (config.kind === "seed") { - if (!config.hf_repo_id.trim()) { + const seedSourceType = config.seed_source_type ?? "hf"; + if (seedSourceType === "hf" && !config.hf_repo_id.trim()) { errors.push("Seed dataset repo is required."); } if (!config.hf_path.trim()) { errors.push("Seed metadata not loaded. Click 'Load columns + 10 rows'."); } - if (config.hf_endpoint?.trim() && !config.hf_endpoint.trim().startsWith("http")) { + if ( + seedSourceType === "hf" && + config.hf_endpoint?.trim() && + !config.hf_endpoint.trim().startsWith("http") + ) { errors.push("HF endpoint must start with http."); } if (config.drop && (config.seed_columns?.length ?? 0) === 0) {