From bdc825298d2e99cd37ec4673eaf7ea34f83c4c37 Mon Sep 17 00:00:00 2001 From: Shine1i Date: Tue, 3 Mar 2026 11:11:26 +0100 Subject: [PATCH] feat(seed): backend unstructured seed reader + server-side chunking, remove client chunk splitter --- studio/backend/core/data_recipe/service.py | 3 + .../core/data_recipe/unstructured_seed.py | 186 ++++++++++++++++++ .../data_recipe/unstructured_seed_plugin.py | 90 +++++++++ studio/backend/models/data_recipe.py | 3 + studio/backend/routes/data_recipe/seed.py | 53 ++++- studio/frontend/bun.lock | 3 - studio/frontend/package.json | 1 - .../src/features/recipe-studio/api/index.ts | 6 + .../dialogs/seed/seed-dialog.tsx | 60 +++--- .../hooks/use-recipe-persistence.ts | 1 + .../import/parsers/seed-config-parser.ts | 14 +- .../utils/payload/builders-seed.ts | 38 +++- .../recipe-studio/utils/validation.ts | 15 ++ 13 files changed, 421 insertions(+), 52 deletions(-) create mode 100644 studio/backend/core/data_recipe/unstructured_seed.py create mode 100644 studio/backend/core/data_recipe/unstructured_seed_plugin.py diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index 5367c3c0fc..54e3305e3f 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -7,6 +7,7 @@ from pathlib import Path from typing import Any from .jsonable import to_jsonable +from .unstructured_seed_plugin import ensure_unstructured_seed_plugin_registered _IMAGE_CONTEXT_PATCHED = False @@ -99,6 +100,8 @@ def _apply_data_designer_image_context_patch() -> None: if _IMAGE_CONTEXT_PATCHED: return + ensure_unstructured_seed_plugin_registered() + try: from data_designer.config.models import ImageContext except ImportError: diff --git a/studio/backend/core/data_recipe/unstructured_seed.py b/studio/backend/core/data_recipe/unstructured_seed.py new file mode 100644 index 0000000000..ee486fc496 --- /dev/null +++ b/studio/backend/core/data_recipe/unstructured_seed.py @@ -0,0 +1,186 @@ +from __future__ import annotations + +import hashlib +import re +from pathlib import Path +from typing import Any + + +DEFAULT_CHUNK_SIZE = 1200 +DEFAULT_CHUNK_OVERLAP = 200 +MAX_CHUNK_SIZE = 20000 +_MIN_BREAK_RATIO = 0.6 +_CACHE_DIR = Path.home() / ".cache" / "unsloth" / "data-recipe" / "unstructured-seed-cache" + + +def resolve_chunking( + chunk_size: Any, + chunk_overlap: Any, +) -> tuple[int, int]: + size = _to_int(chunk_size, DEFAULT_CHUNK_SIZE) + size = max(1, min(size, MAX_CHUNK_SIZE)) + overlap = _to_int(chunk_overlap, DEFAULT_CHUNK_OVERLAP) + overlap = max(0, min(overlap, max(0, size - 1))) + return size, overlap + + +def build_unstructured_preview_rows( + *, + source_path: Path, + preview_size: int, + chunk_size: Any, + chunk_overlap: Any, +) -> list[dict[str, str]]: + parquet_path, rows = materialize_unstructured_seed_dataset( + source_path=source_path, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + ) + count = max(0, int(preview_size)) + if rows: + return rows[:count] + + try: + import pandas as pd + except ImportError as exc: # pragma: no cover + raise RuntimeError(f"pandas is required for unstructured seed processing: {exc}") from exc + dataframe = pd.read_parquet(parquet_path).head(count) + return [ + {"chunk_text": str(value.get("chunk_text", "")).strip()} + for value in dataframe.to_dict(orient="records") + if str(value.get("chunk_text", "")).strip() + ] + + +def materialize_unstructured_seed_dataset( + *, + source_path: Path, + chunk_size: Any, + chunk_overlap: Any, +) -> tuple[Path, list[dict[str, str]]]: + resolved = source_path.expanduser().resolve() + if not resolved.is_file(): + raise FileNotFoundError(f"unstructured seed file not found: {resolved}") + + size, overlap = resolve_chunking(chunk_size, chunk_overlap) + key = _compute_cache_key( + source_path=resolved, + chunk_size=size, + chunk_overlap=overlap, + ) + parquet_path = _CACHE_DIR / f"{key}.parquet" + if parquet_path.exists(): + return parquet_path, [] + + text = load_unstructured_text_file(resolved) + chunks = split_text_into_chunks( + text=text, + chunk_size=size, + chunk_overlap=overlap, + ) + if not chunks: + raise ValueError("No text found in unstructured seed source.") + + rows = [{"chunk_text": chunk} for chunk in chunks] + _CACHE_DIR.mkdir(parents=True, exist_ok=True) + try: + import pandas as pd + except ImportError as exc: # pragma: no cover + raise RuntimeError(f"pandas is required for unstructured seed processing: {exc}") from exc + + tmp_path = _CACHE_DIR / f"{key}.tmp.parquet" + pd.DataFrame(rows).to_parquet(tmp_path, index=False) + tmp_path.replace(parquet_path) + return parquet_path, rows + + +def load_unstructured_text_file(path: Path) -> str: + ext = path.suffix.lower() + if ext not in {".txt", ".md"}: + raise ValueError(f"Unsupported unstructured seed file type: {ext}") + + raw = path.read_text(encoding="utf-8", errors="ignore") + return normalize_unstructured_text(raw) + + +def normalize_unstructured_text(text: str) -> str: + normalized = text.replace("\r\n", "\n").replace("\r", "\n") + return re.sub(r"\n{3,}", "\n\n", normalized).strip() + + +def split_text_into_chunks( + *, + text: str, + chunk_size: int, + chunk_overlap: int, +) -> list[str]: + if not text: + return [] + if chunk_size <= 0: + return [text] + + chunks: list[str] = [] + start = 0 + min_break_index = int(chunk_size * _MIN_BREAK_RATIO) + text_len = len(text) + while start < text_len: + end = min(text_len, start + chunk_size) + if end < text_len: + window = text[start:end] + cut = _find_break_index(window, min_break_index) + if cut is not None and cut > 0: + end = start + cut + + if end <= start: + end = min(text_len, start + chunk_size) + + chunk = text[start:end].strip() + if chunk: + chunks.append(chunk) + if end >= text_len: + break + + next_start = end - chunk_overlap + if next_start <= start: + next_start = end + start = max(0, next_start) + + return chunks + + +def _find_break_index(window: str, min_index: int) -> int | None: + breakpoints = ["\n\n", "\n", " "] + for token in breakpoints: + idx = window.rfind(token) + if idx >= min_index: + return idx + len(token) + return None + + +def _to_int(value: Any, fallback: int) -> int: + if isinstance(value, bool): + return fallback + try: + parsed = int(str(value).strip()) + except (TypeError, ValueError): + return fallback + return parsed + + +def _compute_cache_key( + *, + source_path: Path, + chunk_size: int, + chunk_overlap: int, +) -> str: + stat = source_path.stat() + payload = "|".join( + [ + str(source_path), + str(stat.st_size), + str(stat.st_mtime_ns), + str(chunk_size), + str(chunk_overlap), + ] + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() diff --git a/studio/backend/core/data_recipe/unstructured_seed_plugin.py b/studio/backend/core/data_recipe/unstructured_seed_plugin.py new file mode 100644 index 0000000000..a956ef4236 --- /dev/null +++ b/studio/backend/core/data_recipe/unstructured_seed_plugin.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Literal + +from pydantic import Field, field_validator + +from .unstructured_seed import ( + DEFAULT_CHUNK_OVERLAP, + DEFAULT_CHUNK_SIZE, + materialize_unstructured_seed_dataset, + resolve_chunking, +) + +try: + import data_designer.lazy_heavy_imports as lazy + from data_designer.config.seed_source import SeedSource + from data_designer.engine.resources.seed_reader import SeedReader +except ImportError: # pragma: no cover + lazy = None + + class SeedSource: # type: ignore[no-redef] + pass + + class SeedReader: # type: ignore[no-redef] + @classmethod + def __class_getitem__(cls, _item): + return cls + + +class UnstructuredSeedSource(SeedSource): + seed_type: Literal["unstructured"] = "unstructured" + path: str = Field(..., min_length=1) + chunk_size: int = DEFAULT_CHUNK_SIZE + chunk_overlap: int = DEFAULT_CHUNK_OVERLAP + + @field_validator("path", mode="after") + @classmethod + def _validate_path(cls, value: str) -> str: + path = Path(value).expanduser() + if not path.is_file(): + raise ValueError(f"Unstructured seed path is not a file: {path}") + return value + + @field_validator("chunk_size", mode="after") + @classmethod + def _validate_chunk_size(cls, value: int) -> int: + size, _ = resolve_chunking(value, 0) + return size + + @field_validator("chunk_overlap", mode="after") + @classmethod + def _validate_chunk_overlap(cls, value: int, info) -> int: + size = info.data.get("chunk_size", cls.model_fields["chunk_size"].default) + _, overlap = resolve_chunking(size, value) + return overlap + + +class UnstructuredSeedReader(SeedReader[UnstructuredSeedSource]): + def create_duckdb_connection(self): + if lazy is None: + raise RuntimeError("data_designer is not available") + return lazy.duckdb.connect() + + def get_dataset_uri(self) -> str: + path, _ = materialize_unstructured_seed_dataset( + source_path=Path(self.source.path), + chunk_size=self.source.chunk_size, + chunk_overlap=self.source.chunk_overlap, + ) + return str(path) + + +def ensure_unstructured_seed_plugin_registered() -> None: + try: + from data_designer.plugins.plugin import Plugin, PluginType + from data_designer.plugins.registry import PluginRegistry + except ImportError: + return + + registry = PluginRegistry() + if registry.plugin_exists("unstructured"): + return + + plugin = Plugin( + impl_qualified_name="core.data_recipe.unstructured_seed_plugin.UnstructuredSeedReader", + config_qualified_name="core.data_recipe.unstructured_seed_plugin.UnstructuredSeedSource", + plugin_type=PluginType.SEED_READER, + ) + registry._plugins[plugin.name] = plugin # type: ignore[attr-defined] diff --git a/studio/backend/models/data_recipe.py b/studio/backend/models/data_recipe.py index 9e501c15e2..01e7c8cd1f 100644 --- a/studio/backend/models/data_recipe.py +++ b/studio/backend/models/data_recipe.py @@ -49,6 +49,9 @@ 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) + seed_source_type: str | None = None + unstructured_chunk_size: int | None = Field(default=None, ge=1, le=20000) + unstructured_chunk_overlap: int | None = Field(default=None, ge=0, le=20000) class SeedInspectResponse(BaseModel): diff --git a/studio/backend/routes/data_recipe/seed.py b/studio/backend/routes/data_recipe/seed.py index 1b8df3bc9b..3a824827cb 100644 --- a/studio/backend/routes/data_recipe/seed.py +++ b/studio/backend/routes/data_recipe/seed.py @@ -10,6 +10,10 @@ from typing import Any from uuid import uuid4 from fastapi import APIRouter, HTTPException +from core.data_recipe.unstructured_seed import ( + build_unstructured_preview_rows, + resolve_chunking, +) from core.data_recipe.jsonable import to_preview_jsonable from models.data_recipe import ( @@ -23,6 +27,7 @@ router = APIRouter() DATA_EXTS = (".parquet", ".jsonl", ".json", ".csv") DEFAULT_SPLIT = "train" LOCAL_UPLOAD_EXTS = {".csv", ".json", ".jsonl"} +UNSTRUCTURED_UPLOAD_EXTS = {".txt", ".md"} SEED_UPLOAD_DIR = Path.home() / ".cache" / "unsloth" / "data-recipe" / "seed-uploads" @@ -180,6 +185,26 @@ def _read_preview_rows_from_local_file(path: Path, preview_size: int) -> list[di return _serialize_preview_rows(rows) +def _read_preview_rows_from_unstructured_file( + *, + path: Path, + preview_size: int, + chunk_size: int | None, + chunk_overlap: int | None, +) -> list[dict[str, Any]]: + size, overlap = resolve_chunking(chunk_size, chunk_overlap) + try: + rows = build_unstructured_preview_rows( + source_path=path, + preview_size=preview_size, + chunk_size=size, + chunk_overlap=overlap, + ) + except (FileNotFoundError, RuntimeError, ValueError, OSError) as exc: + raise HTTPException(status_code=422, detail=f"seed inspect failed: {exc}") from exc + 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() @@ -257,11 +282,17 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse: @router.post("/seed/inspect-upload", response_model=SeedInspectResponse) def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectResponse: + seed_source_type = _normalize_optional_text(payload.seed_source_type) or "local" 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}") + if seed_source_type == "unstructured": + if ext not in UNSTRUCTURED_UPLOAD_EXTS: + allowed = ", ".join(sorted(UNSTRUCTURED_UPLOAD_EXTS)) + raise HTTPException(status_code=400, detail=f"unsupported file type: {ext}. allowed: {allowed}") + else: + 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: @@ -275,10 +306,18 @@ def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectRespons 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 seed_source_type == "unstructured": + preview_rows = _read_preview_rows_from_unstructured_file( + path=stored_path, + preview_size=int(payload.preview_size), + chunk_size=payload.unstructured_chunk_size, + chunk_overlap=payload.unstructured_chunk_overlap, + ) + else: + 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) diff --git a/studio/frontend/bun.lock b/studio/frontend/bun.lock index f200156928..e096f34cd7 100644 --- a/studio/frontend/bun.lock +++ b/studio/frontend/bun.lock @@ -18,7 +18,6 @@ "@hugeicons/react": "^1.1.5", "@huggingface/hub": "^2.9.0", "@langchain/core": "^1.1.27", - "@langchain/textsplitters": "^1.0.1", "@radix-ui/react-checkbox": "^1.3.3", "@radix-ui/react-label": "^2.1.8", "@radix-ui/react-select": "^2.2.6", @@ -389,8 +388,6 @@ "@langchain/core": ["@langchain/core@1.1.28", "", { "dependencies": { "@cfworker/json-schema": "^4.0.2", "ansi-styles": "^5.0.0", "camelcase": "6", "decamelize": "1.2.0", "js-tiktoken": "^1.0.12", "langsmith": ">=0.5.0 <1.0.0", "mustache": "^4.2.0", "p-queue": "^6.6.2", "uuid": "^10.0.0", "zod": "^3.25.76 || ^4" } }, "sha512-6FAGdezEp8zHY92LtnsAiv54KaG41nBdsuukk+R+1484edV20cVOyIc36ANuGKPx0pmYFCBWhCUdO0jxB/zn2Q=="], - "@langchain/textsplitters": ["@langchain/textsplitters@1.0.1", "", { "dependencies": { "js-tiktoken": "^1.0.12" }, "peerDependencies": { "@langchain/core": "^1.0.0" } }, "sha512-rheJlB01iVtrOUzttscutRgLybPH9qR79EyzBEbf1u97ljWyuxQfCwIWK+SjoQTM9O8M7GGLLRBSYE26Jmcoww=="], - "@mermaid-js/parser": ["@mermaid-js/parser@1.0.0", "", { "dependencies": { "langium": "^4.0.0" } }, "sha512-vvK0Hi/VWndxoh03Mmz6wa1KDriSPjS2XMZL/1l19HFwygiObEEoEwSDxOqyLzzAI6J2PU3261JjTMTO7x+BPw=="], "@modelcontextprotocol/sdk": ["@modelcontextprotocol/sdk@1.27.1", "", { "dependencies": { "@hono/node-server": "^1.19.9", "ajv": "^8.17.1", "ajv-formats": "^3.0.1", "content-type": "^1.0.5", "cors": "^2.8.5", "cross-spawn": "^7.0.5", "eventsource": "^3.0.2", "eventsource-parser": "^3.0.0", "express": "^5.2.1", "express-rate-limit": "^8.2.1", "hono": "^4.11.4", "jose": "^6.1.3", "json-schema-typed": "^8.0.2", "pkce-challenge": "^5.0.0", "raw-body": "^3.0.0", "zod": "^3.25 || ^4.0", "zod-to-json-schema": "^3.25.1" }, "peerDependencies": { "@cfworker/json-schema": "^4.1.1" }, "optionalPeers": ["@cfworker/json-schema"] }, "sha512-sr6GbP+4edBwFndLbM60gf07z0FQ79gaExpnsjMGePXqFcSSb7t6iscpjk9DhFhwd+mTEQrzNafGP8/iGGFYaA=="], diff --git a/studio/frontend/package.json b/studio/frontend/package.json index f4b636c2c6..432f25b1e6 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -26,7 +26,6 @@ "@hugeicons/react": "^1.1.5", "@huggingface/hub": "^2.9.0", "@langchain/core": "^1.1.27", - "@langchain/textsplitters": "^1.0.1", "@radix-ui/react-checkbox": "^1.3.3", "@radix-ui/react-label": "^2.1.8", "@radix-ui/react-select": "^2.2.6", diff --git a/studio/frontend/src/features/recipe-studio/api/index.ts b/studio/frontend/src/features/recipe-studio/api/index.ts index 88a6c91d21..1b3c186659 100644 --- a/studio/frontend/src/features/recipe-studio/api/index.ts +++ b/studio/frontend/src/features/recipe-studio/api/index.ts @@ -89,6 +89,12 @@ export type SeedInspectUploadRequest = { content_base64: string; // biome-ignore lint/style/useNamingConvention: api schema preview_size?: number; + // biome-ignore lint/style/useNamingConvention: api schema + seed_source_type?: "local" | "unstructured"; + // biome-ignore lint/style/useNamingConvention: api schema + unstructured_chunk_size?: number; + // biome-ignore lint/style/useNamingConvention: api schema + unstructured_chunk_overlap?: number; }; export type SeedInspectResponse = { 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 b08e390b0b..5f73e3543e 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 @@ -34,7 +34,6 @@ import { TabsList, TabsTrigger, } from "@/components/ui/tabs"; -import { RecursiveCharacterTextSplitter } from "@langchain/textsplitters"; import mammoth from "mammoth"; import { type ReactElement, useCallback, useEffect, useMemo, useRef, useState } from "react"; import { extractText, getDocumentProxy } from "unpdf"; @@ -138,22 +137,6 @@ function resolveChunking(config: SeedConfig): { return { chunkSize, chunkOverlap }; } -async function chunkText( - input: string, - chunkSize: number, - chunkOverlap: number, -): Promise { - const text = input.replace(/\r\n/g, "\n").replace(/\n{3,}/g, "\n\n").trim(); - if (!text) return []; - - const splitter = new RecursiveCharacterTextSplitter({ - chunkSize, - chunkOverlap, - }); - const chunks = await splitter.splitText(text); - return chunks.map((chunk) => chunk.trim()).filter(Boolean); -} - async function fileToBase64Payload(file: File): Promise { return new Promise((resolve, reject) => { const reader = new FileReader(); @@ -186,6 +169,23 @@ async function extractUnstructuredText(file: File): Promise { throw new Error("Unsupported unstructured file type"); } +async function toUnstructuredUploadFile(file: File): Promise { + const lower = file.name.toLowerCase(); + if (lower.endsWith(".txt") || lower.endsWith(".md")) { + return file; + } + + const text = (await extractUnstructuredText(file)).trim(); + if (!text) { + throw new Error("No text found in file."); + } + const normalized = text.replace(/\r\n/g, "\n").replace(/\r/g, "\n"); + const stem = file.name.replace(/\.(pdf|docx)$/i, "") || "unstructured_seed"; + return new File([normalized], `${stem}.txt`, { + type: "text/plain", + }); +} + export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactElement { const [inspectError, setInspectError] = useState(null); const [isInspecting, setIsInspecting] = useState(false); @@ -312,27 +312,19 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl throw new Error("File too large (max 50MB)."); } - const text = await extractUnstructuredText(unstructuredFile); const { chunkSize, chunkOverlap } = resolveChunking(config); - const chunks = await chunkText(text, chunkSize, chunkOverlap); - if (chunks.length === 0) { - throw new Error("No text found in file."); + const uploadFile = await toUnstructuredUploadFile(unstructuredFile); + if (uploadFile.size > MAX_UPLOAD_BYTES) { + throw new Error("Processed text is too large (max 50MB)."); } - 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 payload = await fileToBase64Payload(uploadFile); const response = await inspectSeedUpload({ - filename: jsonlFile.name, + filename: uploadFile.name, content_base64: payload, preview_size: 10, + seed_source_type: "unstructured", + unstructured_chunk_size: chunkSize, + unstructured_chunk_overlap: chunkOverlap, }); onUpdate({ hf_path: response.resolved_path, @@ -550,7 +542,7 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl

- Chunking uses chunk_text only. Max 50MB. + File is converted to text, then chunked server-side into chunk_text rows. Max 50MB.

{(unstructuredFile?.name || config.unstructured_file_name?.trim()) && ( diff --git a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts index 740ab54c74..4546deb130 100644 --- a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts +++ b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts @@ -97,6 +97,7 @@ function sanitizeSeedForShare(payload: unknown): unknown { typeof source?.seed_type === "string" ? source.seed_type : null; const shouldResetLocalState = sourceType === "local" || + sourceType === "unstructured" || uiSourceType === "local" || uiSourceType === "unstructured"; 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 5cdf5b61bd..e8035c30b4 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 @@ -4,7 +4,7 @@ import type { SeedSelectionType, SeedSourceType, } from "../../../types"; -import { isRecord, readString } from "../helpers"; +import { isRecord, readNumberString, readString } from "../helpers"; function normalizeSampling(value: unknown): SeedSamplingStrategy { const raw = readString(value); @@ -69,7 +69,9 @@ function parseSeedSettings(seedConfigRaw: unknown): Partial { let hf_endpoint = "https://huggingface.co"; let hf_repo_id = ""; let local_file_name = ""; - const unstructured_file_name = ""; + let unstructured_file_name = ""; + let unstructured_chunk_size = "1200"; + let unstructured_chunk_overlap = "200"; const sourceRaw = seedConfigRaw.source; if (isRecord(sourceRaw)) { const seedType = readString(sourceRaw.seed_type); @@ -84,6 +86,12 @@ function parseSeedSettings(seedConfigRaw: unknown): Partial { seed_source_type = "local"; hf_path = sourcePath; local_file_name = sourcePath.split("/").pop() ?? sourcePath; + } else if (seedType === "unstructured") { + seed_source_type = "unstructured"; + hf_path = sourcePath; + unstructured_file_name = sourcePath.split("/").pop() ?? sourcePath; + unstructured_chunk_size = readNumberString(sourceRaw.chunk_size) || "1200"; + unstructured_chunk_overlap = readNumberString(sourceRaw.chunk_overlap) || "200"; } } @@ -119,6 +127,8 @@ function parseSeedSettings(seedConfigRaw: unknown): Partial { hf_endpoint, local_file_name, unstructured_file_name, + unstructured_chunk_size, + unstructured_chunk_overlap, sampling_strategy, selection_type, selection_start, 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 a07e6c2ee8..16c47e5b76 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 @@ -1,5 +1,9 @@ import type { NodeConfig, SeedConfig } from "../../types"; +const DEFAULT_CHUNK_SIZE = 1200; +const DEFAULT_CHUNK_OVERLAP = 200; +const MAX_CHUNK_SIZE = 20000; + function parseIntStrict(value: string | undefined): number | null { const trimmed = value?.trim(); if (!trimmed) return null; @@ -8,6 +12,17 @@ function parseIntStrict(value: string | undefined): number | null { return num; } +function resolveChunking(config: SeedConfig): { chunkSize: number; chunkOverlap: number } { + const rawSize = parseIntStrict(config.unstructured_chunk_size); + const rawOverlap = parseIntStrict(config.unstructured_chunk_overlap); + const chunkSize = Math.min(MAX_CHUNK_SIZE, Math.max(1, rawSize ?? DEFAULT_CHUNK_SIZE)); + const chunkOverlap = Math.min( + Math.max(0, chunkSize - 1), + Math.max(0, rawOverlap ?? DEFAULT_CHUNK_OVERLAP), + ); + return { chunkSize, chunkOverlap }; +} + export function buildSeedConfig( config: SeedConfig, errors: string[], @@ -47,11 +62,24 @@ export function buildSeedConfig( token, endpoint, } - : { - // biome-ignore lint/style/useNamingConvention: api schema - seed_type: "local", - path, - }; + : seedSourceType === "unstructured" + ? (() => { + const { chunkSize, chunkOverlap } = resolveChunking(config); + return { + // biome-ignore lint/style/useNamingConvention: api schema + seed_type: "unstructured", + path, + // biome-ignore lint/style/useNamingConvention: api schema + chunk_size: chunkSize, + // biome-ignore lint/style/useNamingConvention: api schema + chunk_overlap: chunkOverlap, + }; + })() + : { + // biome-ignore lint/style/useNamingConvention: api schema + seed_type: "local", + path, + }; return { source, diff --git a/studio/frontend/src/features/recipe-studio/utils/validation.ts b/studio/frontend/src/features/recipe-studio/utils/validation.ts index 4030d922be..8d62c45eb7 100644 --- a/studio/frontend/src/features/recipe-studio/utils/validation.ts +++ b/studio/frontend/src/features/recipe-studio/utils/validation.ts @@ -226,6 +226,21 @@ export function getConfigErrors(config: NodeConfig | null): string[] { if (config.drop && (config.seed_columns?.length ?? 0) === 0) { errors.push("Seed drop needs loaded columns."); } + const chunkSizeRaw = Number(config.unstructured_chunk_size); + const chunkOverlapRaw = Number(config.unstructured_chunk_overlap); + if (!Number.isFinite(chunkSizeRaw) || Math.floor(chunkSizeRaw) < 1) { + errors.push("Chunk size must be an integer >= 1."); + } + if (!Number.isFinite(chunkOverlapRaw) || Math.floor(chunkOverlapRaw) < 0) { + errors.push("Chunk overlap must be an integer >= 0."); + } + if ( + Number.isFinite(chunkSizeRaw) && + Number.isFinite(chunkOverlapRaw) && + Math.floor(chunkOverlapRaw) >= Math.floor(chunkSizeRaw) + ) { + errors.push("Chunk overlap must be less than chunk size."); + } } else { const selectedDropColumns = (config.seed_drop_columns ?? []) .map((value) => value.trim())