ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

merge nightly

Browse source
This commit is contained in:
Shine1i 2026-02-27 10:31:37 +01:00
commit 3c728f5eb3
87 changed files with 4420 additions and 1394 deletions
Download patch file Download diff file Expand all files Collapse all files

6
.gitignore vendored
View file

@ -25,6 +25,12 @@ unsloth_training_checkpoints/
*.gguf
*.safetensors
# llama.cpp build (built by setup.sh, shared with unsloth-zoo export)
llama.cpp/
# Built binaries (llama-server etc.)
bin/
# IDE / Editors
.vscode/
.idea/

109
setup.sh
View file

@ -69,13 +69,14 @@ if [ "$NEED_NODE" = true ]; then
# Load nvm (source ~/.bashrc won't work inside a script)
export NVM_DIR="$HOME/.nvm"
set +u
[ -s "$NVM_DIR/nvm.sh" ] && \. "$NVM_DIR/nvm.sh"
# ── 3. Install Node LTS ──
echo "Installing Node LTS..."
run_quiet "nvm install" nvm install --lts
nvm use --lts > /dev/null 2>&1
set -u
# ── 4. Verify versions ──
NODE_MAJOR=$(node -v | sed 's/v//' | cut -d. -f1)
NPM_MAJOR=$(npm -v | cut -d. -f1)
@ -105,7 +106,9 @@ echo "✅ Frontend built to studio/frontend/dist"
echo ""
echo "Setting up Python environment..."
# ── 6a. Discover best Python <= 3.12.x ──
# ── 6a. Discover best Python >= 3.11 and < 3.14 (i.e. 3.11.x, 3.12.x, or 3.13.x) ──
MIN_PY_MINOR=11 # minimum minor version (>= 3.11)
MAX_PY_MINOR=13 # maximum minor version (< 3.14)
BEST_PY=""
BEST_MAJOR=0
BEST_MINOR=0
@ -115,7 +118,7 @@ for candidate in $(compgen -c python3 2>/dev/null | grep -E '^python3(\.[0-9]+)?
if ! command -v "$candidate" &>/dev/null; then
continue
fi
# Get version string, e.g. "Python 3.11.5"
# Get version string, e.g. "Python 3.12.5"
ver_str=$("$candidate" --version 2>&1 | awk '{print $2}')
py_major=$(echo "$ver_str" | cut -d. -f1)
py_minor=$(echo "$ver_str" | cut -d. -f2)
@ -125,8 +128,13 @@ for candidate in $(compgen -c python3 2>/dev/null | grep -E '^python3(\.[0-9]+)?
continue
fi
# Skip versions above 3.12
if [ "$py_minor" -gt 12 ] 2>/dev/null; then
# Skip versions below 3.12 (require > 3.11)
if [ "$py_minor" -lt "$MIN_PY_MINOR" ] 2>/dev/null; then
continue
fi
# Skip versions above 3.13 (require < 3.14)
if [ "$py_minor" -gt "$MAX_PY_MINOR" ] 2>/dev/null; then
continue
fi
@ -139,7 +147,7 @@ for candidate in $(compgen -c python3 2>/dev/null | grep -E '^python3(\.[0-9]+)?
done
if [ -z "$BEST_PY" ]; then
echo "❌ ERROR: No Python version <= 3.12.x found on this system."
echo "❌ ERROR: No Python version between 3.${MIN_PY_MINOR} and 3.${MAX_PY_MINOR} found on this system."
echo " Detected Python 3 installations:"
for candidate in $(compgen -c python3 2>/dev/null | grep -E '^python3(\.[0-9]+)?$' | sort -u); do
if command -v "$candidate" &>/dev/null; then
@ -147,13 +155,13 @@ if [ -z "$BEST_PY" ]; then
fi
done
echo ""
echo " Please install Python <= 3.12.x for maximum compatibility."
echo " Please install Python 3.${MIN_PY_MINOR} or 3.${MAX_PY_MINOR}."
echo " For example: sudo apt install python3.12 python3.12-venv"
exit 1
fi
BEST_VER=$("$BEST_PY" --version 2>&1 | awk '{print $2}')
echo "✅ Using $BEST_PY ($BEST_VER) — compatible (≤ 3.12.x)"
echo "✅ Using $BEST_PY ($BEST_VER) — compatible (3.${MIN_PY_MINOR}.x – 3.${MAX_PY_MINOR}.x)"
REQ_ROOT="$SCRIPT_DIR/studio/backend/requirements"
SINGLE_ENV_CONSTRAINTS="$REQ_ROOT/single-env/constraints.txt"
@ -211,7 +219,90 @@ else
fi
fi
# ── 8. Add shell alias (skip in Colab) ──
# ── 8. Build llama.cpp binaries for GGUF inference + export ──
# Builds in-tree at $REPO/llama.cpp/. This directory is shared with
# unsloth-zoo's GGUF export pipeline. We build:
# - llama-server: for GGUF model inference
# - llama-quantize: for GGUF export quantization (symlinked to root for check_llama_cpp())
LLAMA_CPP_DIR="$SCRIPT_DIR/llama.cpp"
LLAMA_SERVER_BIN="$LLAMA_CPP_DIR/build/bin/llama-server"
rm -rf "$LLAMA_CPP_DIR"
{
# Check prerequisites
if ! command -v cmake &>/dev/null; then
echo ""
echo "⚠️ cmake not found — skipping llama-server build (GGUF inference won't be available)"
echo " Install cmake and re-run setup.sh to enable GGUF inference."
elif ! command -v git &>/dev/null; then
echo ""
echo "⚠️ git not found — skipping llama-server build (GGUF inference won't be available)"
else
echo ""
echo "Building llama-server for GGUF inference..."
BUILD_OK=true
run_quiet "clone llama.cpp" git clone --depth 1 https://github.com/ggml-org/llama.cpp.git "$LLAMA_CPP_DIR" || BUILD_OK=false
if [ "$BUILD_OK" = true ]; then
CMAKE_ARGS=""
# Detect CUDA: check nvcc on PATH, then common install locations
NVCC_PATH=""
if command -v nvcc &>/dev/null; then
NVCC_PATH="$(command -v nvcc)"
elif [ -x /usr/local/cuda/bin/nvcc ]; then
NVCC_PATH="/usr/local/cuda/bin/nvcc"
export PATH="/usr/local/cuda/bin:$PATH"
elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then
# Pick the newest cuda-XX.X directory
NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)"
export PATH="$(dirname "$NVCC_PATH"):$PATH"
fi
if [ -n "$NVCC_PATH" ]; then
echo " Building with CUDA support (nvcc: $NVCC_PATH)..."
CMAKE_ARGS="-DGGML_CUDA=ON"
elif [ -d /usr/local/cuda ] || nvidia-smi &>/dev/null; then
echo " CUDA driver detected but nvcc not found — building CPU-only"
echo " To enable GPU: install cuda-toolkit or add nvcc to PATH"
else
echo " Building CPU-only (no CUDA detected)..."
fi
NCPU=$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4)
run_quiet "cmake llama.cpp" cmake -S "$LLAMA_CPP_DIR" -B "$LLAMA_CPP_DIR/build" $CMAKE_ARGS || BUILD_OK=false
fi
if [ "$BUILD_OK" = true ]; then
run_quiet "build llama-server" cmake --build "$LLAMA_CPP_DIR/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false
fi
# Also build llama-quantize (needed by unsloth-zoo's GGUF export pipeline)
if [ "$BUILD_OK" = true ]; then
run_quiet "build llama-quantize" cmake --build "$LLAMA_CPP_DIR/build" --config Release --target llama-quantize -j"$NCPU" || true
# Symlink to llama.cpp root — check_llama_cpp() looks for the binary there
QUANTIZE_BIN="$LLAMA_CPP_DIR/build/bin/llama-quantize"
if [ -f "$QUANTIZE_BIN" ]; then
ln -sf build/bin/llama-quantize "$LLAMA_CPP_DIR/llama-quantize"
fi
fi
if [ "$BUILD_OK" = true ]; then
if [ -f "$LLAMA_SERVER_BIN" ]; then
echo "✅ llama-server built at $LLAMA_SERVER_BIN"
else
echo "⚠️ llama-server binary not found after build — GGUF inference won't be available"
fi
if [ -f "$LLAMA_CPP_DIR/llama-quantize" ]; then
echo "✅ llama-quantize available for GGUF export"
fi
else
echo "⚠️ llama-server build failed — GGUF inference won't be available, but everything else works"
fi
fi
}
# ── 9. Add shell alias (skip in Colab) ──
# Note: venv activation does NOT persist across terminal sessions.
# This alias hardcodes the venv python path so users don't need to activate.
if [ "$IS_COLAB" = false ]; then

6
studio/backend/auth/authentication.py
View file

@ -1,5 +1,5 @@
import secrets
from datetime import UTC, datetime, timedelta
from datetime import datetime, timedelta, timezone
from typing import Optional
from fastapi import Depends, HTTPException, status
@ -34,7 +34,7 @@ def create_access_token(
Tokens are valid across restarts because SECRET_KEY is stored in SQLite.
"""
to_encode = {"sub": subject}
expire = datetime.now(UTC) + (
expire = datetime.now(timezone.utc) + (
expires_delta or timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
)
to_encode.update({"exp": expire})
@ -48,7 +48,7 @@ def create_refresh_token(subject: str) -> str:
Refresh tokens are opaque (not JWTs) and expire after REFRESH_TOKEN_EXPIRE_DAYS.
"""
token = secrets.token_urlsafe(48)
expires_at = datetime.now(UTC) + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS)
expires_at = datetime.now(timezone.utc) + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS)
save_refresh_token(token, subject, expires_at.isoformat())
return token

6
studio/backend/auth/storage.py
View file

@ -3,7 +3,7 @@ SQLite storage for authentication data (user credentials + JWT secret).
"""
import hashlib
import sqlite3
from datetime import UTC, datetime
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional, Tuple
@ -218,7 +218,7 @@ def verify_refresh_token(token: str) -> Optional[str]:
# Clean up any expired tokens while we're here
conn.execute(
"DELETE FROM refresh_tokens WHERE expires_at < ?",
(datetime.now(UTC).isoformat(),),
(datetime.now(timezone.utc).isoformat(),),
)
conn.commit()
@ -235,7 +235,7 @@ def verify_refresh_token(token: str) -> Optional[str]:
# Check expiry
expires_at = datetime.fromisoformat(row["expires_at"])
if datetime.now(UTC) > expires_at:
if datetime.now(timezone.utc) > expires_at:
conn.execute("DELETE FROM refresh_tokens WHERE id = ?", (row["id"],))
conn.commit()
return None

30
studio/backend/core/data_recipe/jobs/constants.py Normal file
View file

@ -0,0 +1,30 @@
from __future__ import annotations
# stages parsed from data-designer logs
STAGE_CREATE = "create"
STAGE_PREVIEW = "preview"
STAGE_DAG = "dag"
STAGE_HEALTHCHECK = "healthcheck"
STAGE_SAMPLING = "sampling"
STAGE_COLUMN_CONFIG = "column_config"
STAGE_GENERATING = "generating"
STAGE_BATCH = "batch"
STAGE_PROFILING = "profiling"
USAGE_RESET_STAGES = {
STAGE_CREATE,
STAGE_PREVIEW,
STAGE_DAG,
STAGE_HEALTHCHECK,
STAGE_SAMPLING,
STAGE_GENERATING,
STAGE_PROFILING,
}
# job event types emitted by worker/manager
EVENT_JOB_ENQUEUED = "job.enqueued"
EVENT_JOB_STARTED = "job.started"
EVENT_JOB_CANCELLING = "job.cancelling"
EVENT_JOB_CANCELLED = "job.cancelled"
EVENT_JOB_COMPLETED = "job.completed"
EVENT_JOB_ERROR = "job.error"

67
studio/backend/core/data_recipe/jobs/manager.py
View file

@ -13,6 +13,15 @@ from typing import Any
import multiprocessing as mp
from ..jsonable import to_jsonable
from .constants import (
EVENT_JOB_CANCELLING,
EVENT_JOB_CANCELLED,
EVENT_JOB_COMPLETED,
EVENT_JOB_ENQUEUED,
EVENT_JOB_ERROR,
EVENT_JOB_STARTED,
)
from .parse import apply_update, coerce_event, parse_log_message
from .types import Job
from .worker import run_job_process
@ -20,29 +29,6 @@ from .worker import run_job_process
_CTX = mp.get_context("spawn")
def _to_jsonable(value: Any) -> Any:
try:
import numpy as np # type: ignore
except Exception: # pragma: no cover
np = None # type: ignore
if np is not None:
if isinstance(value, np.ndarray):
return value.tolist()
if isinstance(value, np.generic):
return value.item()
if isinstance(value, dict):
return {str(k): _to_jsonable(v) for k, v in value.items()}
if isinstance(value, (list, tuple, set)):
return [_to_jsonable(v) for v in value]
if hasattr(value, "isoformat") and callable(value.isoformat):
try:
return value.isoformat()
except Exception:
pass
return value
@dataclass
class Subscription:
@ -123,7 +109,7 @@ class JobManager:
self._pump_thread = threading.Thread(target=self._pump_loop, daemon=True)
self._pump_thread.start()
self._emit({"type": "job.enqueued", "ts": time.time(), "job_id": job_id})
self._emit({"type": EVENT_JOB_ENQUEUED, "ts": time.time(), "job_id": job_id})
return job_id
def cancel(self, job_id: str) -> bool:
@ -134,15 +120,15 @@ class JobManager:
if self._proc is None or not self._proc.is_alive():
return True
self._job.status = "cancelling"
self._emit({"type": "job.cancelling", "ts": time.time(), "job_id": job_id})
self._emit({"type": EVENT_JOB_CANCELLING, "ts": time.time(), "job_id": job_id})
try:
self._proc.terminate()
except Exception:
except (AttributeError, OSError):
pass
return True
def get_status(self, job_id: str) -> dict | None:
"""UI-friendly snapshot. Poll this if you don't want SSE."""
"""UI friendly snapshot that we need. Alternative to sse kinda of and structured"""
with self._lock:
if self._job is None or self._job.job_id != job_id:
return None
@ -304,7 +290,7 @@ class JobManager:
).fetchdf()
finally:
conn.close()
except Exception:
except (RuntimeError, ValueError, duckdb.Error):
return None
for helper_col in ("filename", "__row_num__"):
@ -312,7 +298,7 @@ class JobManager:
dataframe = dataframe.drop(columns=[helper_col])
rows = dataframe.to_dict(orient="records")
return {"dataset": _to_jsonable(rows), "total": total}
return {"dataset": to_jsonable(rows), "total": total}
@staticmethod
def _load_dataset_page_with_data_designer(
@ -326,7 +312,7 @@ class JobManager:
dataframe = read_parquet_dataset(parquet_dir)
total = int(len(dataframe.index))
rows = dataframe.iloc[offset:offset + limit].to_dict(orient="records")
return {"dataset": _to_jsonable(rows), "total": total}
return {"dataset": to_jsonable(rows), "total": total}
def subscribe(self, job_id: str, *, after_seq: int | None = None) -> Subscription | None:
"""SSE subscribe: get replay buffer + live events stream."""
@ -355,7 +341,7 @@ class JobManager:
for q in self._subs:
try:
q.put_nowait(event)
except Exception:
except queue.Full:
stale.append(q)
if stale:
self._subs = [q for q in self._subs if q not in stale]
@ -374,7 +360,7 @@ class JobManager:
return coerce_event(q.get(timeout=timeout_sec))
except queue.Empty:
return None
except Exception:
except (EOFError, OSError, ValueError):
return None
@staticmethod
@ -386,7 +372,7 @@ class JobManager:
events.append(coerce_event(q.get_nowait()))
except queue.Empty:
return events
except Exception:
except (EOFError, OSError, ValueError):
return events
def _pump_loop(self) -> None:
@ -416,13 +402,10 @@ class JobManager:
self._job.status = "error"
self._job.error = self._job.error or "process exited"
self._job.finished_at = time.time()
self._emit(
{
"type": f"job.{self._job.status}",
"ts": time.time(),
"job_id": self._job.job_id,
}
event_type = (
EVENT_JOB_CANCELLED if self._job.status == "cancelled" else EVENT_JOB_ERROR
)
self._emit({"type": event_type, "ts": time.time(), "job_id": self._job.job_id})
return
def _handle_event(self, job: Job, event: dict) -> None:
@ -433,9 +416,9 @@ class JobManager:
with self._lock:
if self._job is None or self._job.job_id != job.job_id:
return
if et == "job.started":
if et == EVENT_JOB_STARTED:
self._job.status = "active"
if et == "job.completed":
if et == EVENT_JOB_COMPLETED:
self._job.status = "completed"
self._job.finished_at = time.time()
self._job.analysis = event.get("analysis")
@ -445,7 +428,7 @@ class JobManager:
if self._job.progress.total and self._job.progress.total > 0:
self._job.progress.done = self._job.progress.total
self._job.progress.percent = 100.0
if et == "job.error":
if et == EVENT_JOB_ERROR:
self._job.status = "error"
self._job.finished_at = time.time()
self._job.error = event.get("error") or "error"

49
studio/backend/core/data_recipe/jobs/parse.py
View file

@ -4,6 +4,18 @@ import re
from dataclasses import dataclass
from typing import Any
from .constants import (
STAGE_BATCH,
STAGE_COLUMN_CONFIG,
STAGE_CREATE,
STAGE_DAG,
STAGE_GENERATING,
STAGE_HEALTHCHECK,
STAGE_PREVIEW,
STAGE_PROFILING,
STAGE_SAMPLING,
USAGE_RESET_STAGES,
)
from .types import Job, ModelUsage, Progress
@ -27,8 +39,7 @@ class ParsedUpdate:
usage_rpm: float | None = None
usage_section_start: bool | None = None
# welp, best effort to parse the logs and convert them to structured information so we can access it and read it properly on the client
# i couldnt find what datadesigner do for progress tracking besides the logs so il probably raise a pr in their repo to add some sort of progress monitoring
# kinda of a bummber but currently only option, Best effort parser from data-designer logs -> structured status for UI.
_RE_SAMPLERS = re.compile(
r"Preparing samplers to generate (?P<rows>\d+) records across (?P<cols>\d+) columns"
)
@ -52,31 +63,31 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
m = _RE_SAMPLERS.search(msg)
if m:
return ParsedUpdate(
stage="sampling",
stage=STAGE_SAMPLING,
rows=int(m.group("rows")),
cols=int(m.group("cols")),
)
if "Sorting column configs into a Directed Acyclic Graph" in msg:
return ParsedUpdate(stage="dag")
return ParsedUpdate(stage=STAGE_DAG)
if "Running health checks for models" in msg:
return ParsedUpdate(stage="healthcheck")
return ParsedUpdate(stage=STAGE_HEALTHCHECK)
if "Preview generation in progress" in msg:
return ParsedUpdate(stage="preview")
return ParsedUpdate(stage=STAGE_PREVIEW)
if "Creating Data Designer dataset" in msg:
return ParsedUpdate(stage="create")
return ParsedUpdate(stage=STAGE_CREATE)
if "Measuring dataset column statistics" in msg:
return ParsedUpdate(stage="profiling")
return ParsedUpdate(stage=STAGE_PROFILING)
m = _RE_COLCFG.search(msg)
if m:
col = m.group("col")
return ParsedUpdate(stage="column_config", current_column=col)
return ParsedUpdate(stage=STAGE_COLUMN_CONFIG, current_column=col)
m = _RE_PROCESSING_COL.search(msg)
if m:
col = m.group("col")
return ParsedUpdate(stage="generating", current_column=col)
return ParsedUpdate(stage=STAGE_GENERATING, current_column=col)
m = _RE_PROGRESS.search(msg)
if m:
@ -89,12 +100,12 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
rate=float(m.group("rate")),
eta_sec=float(m.group("eta")),
)
return ParsedUpdate(stage="generating", progress=p)
return ParsedUpdate(stage=STAGE_GENERATING, progress=p)
m = _RE_BATCH.search(msg)
if m:
return ParsedUpdate(
stage="batch",
stage=STAGE_BATCH,
batch_idx=int(m.group("idx")),
batch_total=int(m.group("total")),
)
@ -132,7 +143,7 @@ def apply_update(job: Job, update: ParsedUpdate) -> None:
job.stage = update.stage
if update.current_column is not None:
job.current_column = update.current_column
if update.stage == "generating" and update.current_column not in job._seen_generation_columns:
if update.stage == STAGE_GENERATING and update.current_column not in job._seen_generation_columns:
job._seen_generation_columns.append(update.current_column)
if update.rows is not None:
job.rows = update.rows
@ -146,16 +157,8 @@ def apply_update(job: Job, update: ParsedUpdate) -> None:
if update.batch_total is not None:
job.batch.total = update.batch_total
if update.stage in {
"profiling",
"generating",
"sampling",
"healthcheck",
"dag",
"create",
"preview",
}:
# usage summary is a short block; reset once we move into the next stage.
if update.stage in USAGE_RESET_STAGES:
# usage summary is a short block so we reset once we move into the next stage.
job._in_usage_summary = False
if update.usage_section_start is not None:

48
studio/backend/core/data_recipe/jobs/worker.py
View file

@ -7,6 +7,8 @@ import traceback
from pathlib import Path
from typing import Any
from ..jsonable import to_jsonable
from .constants import EVENT_JOB_COMPLETED, EVENT_JOB_ERROR, EVENT_JOB_STARTED
from ..service import build_config_builder, create_data_designer
_PROJECT_ROOT = Path(__file__).resolve().parents[5]
@ -28,36 +30,10 @@ class _QueueLogHandler(logging.Handler):
"message": record.getMessage(),
}
self._q.put(event)
except Exception:
except (OSError, RuntimeError, ValueError):
pass
def _to_jsonable(value: Any) -> Any:
try:
import numpy as np # type: ignore
except Exception: # pragma: no cover
np = None # type: ignore
if np is not None:
if isinstance(value, np.ndarray):
return value.tolist()
if isinstance(value, np.generic):
return value.item()
if isinstance(value, dict):
return {str(k): _to_jsonable(v) for k, v in value.items()}
if isinstance(value, (list, tuple, set)):
return [_to_jsonable(v) for v in value]
if hasattr(value, "isoformat") and callable(value.isoformat):
try:
return value.isoformat()
except Exception:
pass
return value
def run_job_process(
*,
event_queue,
@ -68,7 +44,7 @@ def run_job_process(
Subprocess entrypoint.
Sends events to `event_queue`.
"""
event_queue.put({"type": "job.started", "ts": time.time()})
event_queue.put({"type": EVENT_JOB_STARTED, "ts": time.time()})
try:
from data_designer.config.run_config import RunConfig
@ -103,21 +79,21 @@ def run_job_process(
analysis = (
None
if results.analysis is None
else _to_jsonable(results.analysis.model_dump(mode="json"))
else to_jsonable(results.analysis.model_dump(mode="json"))
)
dataset = (
[]
if results.dataset is None
else _to_jsonable(results.dataset.to_dict(orient="records"))
else to_jsonable(results.dataset.to_dict(orient="records"))
)
processor_artifacts = (
None
if results.processor_artifacts is None
else _to_jsonable(results.processor_artifacts)
else to_jsonable(results.processor_artifacts)
)
event_queue.put(
{
"type": "job.completed",
"type": EVENT_JOB_COMPLETED,
"ts": time.time(),
"analysis": analysis,
"dataset": dataset,
@ -128,13 +104,13 @@ def run_job_process(
)
else:
results = designer.create(builder, num_records=rows, dataset_name=dataset_name)
analysis = _to_jsonable(results.load_analysis().model_dump(mode="json"))
analysis = to_jsonable(results.load_analysis().model_dump(mode="json"))
if merge_batches:
_merge_batches_to_single_parquet(results.artifact_storage.base_dataset_path)
artifact_path = str(results.artifact_storage.base_dataset_path)
event_queue.put(
{
"type": "job.completed",
"type": EVENT_JOB_COMPLETED,
"ts": time.time(),
"analysis": analysis,
"artifact_path": artifact_path,
@ -144,7 +120,7 @@ def run_job_process(
except Exception as exc:
event_queue.put(
{
"type": "job.error",
"type": EVENT_JOB_ERROR,
"ts": time.time(),
"error": str(exc),
"stack": traceback.format_exc(limit=20),
@ -160,7 +136,7 @@ def _merge_batches_to_single_parquet(base_dataset_path: Path) -> None:
try:
from data_designer.config.utils.io_helpers import read_parquet_dataset
except Exception:
except ImportError:
return
dataframe = read_parquet_dataset(parquet_dir)

30
studio/backend/core/data_recipe/jsonable.py Normal file
View file

@ -0,0 +1,30 @@
from __future__ import annotations
from typing import Any
def to_jsonable(value: Any) -> Any:
"""Convert numpy/pandas-ish values into plain JSON-safe values."""
try:
import numpy as np # type: ignore
except ImportError: # pragma: no cover
np = None # type: ignore
if np is not None:
if isinstance(value, np.ndarray):
return value.tolist()
if isinstance(value, np.generic):
return value.item()
if isinstance(value, dict):
return {str(k): to_jsonable(v) for k, v in value.items()}
if isinstance(value, (list, tuple, set)):
return [to_jsonable(v) for v in value]
if hasattr(value, "isoformat") and callable(value.isoformat):
try:
return value.isoformat()
except (TypeError, ValueError):
return value
return value

33
studio/backend/core/data_recipe/service.py
View file

@ -3,32 +3,7 @@ from __future__ import annotations
import os
from typing import Any
def _to_jsonable(value: Any) -> Any:
# pydantic/fastapi can't serialize numpy arrays/scalars.
try:
import numpy as np # type: ignore
except Exception: # pragma: no cover
np = None # type: ignore
if np is not None:
if isinstance(value, np.ndarray):
return value.tolist()
if isinstance(value, np.generic):
return value.item()
if isinstance(value, dict):
return {str(k): _to_jsonable(v) for k, v in value.items()}
if isinstance(value, (list, tuple, set)):
return [_to_jsonable(v) for v in value]
# pandas Timestamp/date-like
if hasattr(value, "isoformat") and callable(value.isoformat):
try:
return value.isoformat()
except Exception:
pass
return value
from .jsonable import to_jsonable
def build_model_providers(recipe: dict[str, Any]):
@ -158,17 +133,17 @@ def preview_recipe(
dataset: list[dict[str, Any]] = []
if results.dataset is not None:
raw_rows = results.dataset.to_dict(orient="records")
dataset = [_to_jsonable(row) for row in raw_rows]
dataset = [to_jsonable(row) for row in raw_rows]
artifacts = (
None
if results.processor_artifacts is None
else _to_jsonable(results.processor_artifacts)
else to_jsonable(results.processor_artifacts)
)
analysis = (
None
if results.analysis is None
else _to_jsonable(results.analysis.model_dump(mode="json"))
else to_jsonable(results.analysis.model_dump(mode="json"))
)
return dataset, artifacts, analysis

87
studio/backend/core/export/export.py
View file

@ -2,9 +2,11 @@
"""
Export backend - handles model exporting in various formats
"""
import glob
import json
import logging
import os
import shutil
from pathlib import Path
from typing import Optional, Tuple, List
from peft import PeftModel, PeftModelForCausalLM
@ -397,53 +399,62 @@ class ExportBackend:
# Save locally if requested
if save_directory:
logger.info(f"Saving GGUF model locally to: {save_directory}")
# Resolve to absolute path so unsloth's relative-path internals
# (check_llama_cpp, use_local_gguf, _download_convert_hf_to_gguf)
# all resolve against the repo root cwd, NOT the export directory.
abs_save_dir = os.path.abspath(save_directory)
logger.info(f"Saving GGUF model locally to: {abs_save_dir}")
# Create the directory if it doesn't exist
os.makedirs(save_directory, exist_ok=True)
os.makedirs(abs_save_dir, exist_ok=True)
# Get the base filename for the GGUF file
import shutil
original_dir = os.getcwd()
# On WSL, patch out sudo check before llama.cpp build
_apply_wsl_sudo_patch()
try:
# Change to target directory
os.chdir(save_directory)
logger.info(f"Changed directory to: {save_directory}")
# Snapshot existing .gguf files in cwd before conversion.
# unsloth's convert_to_gguf writes output files relative to
# cwd (repo root), so we diff afterwards and relocate them.
cwd = os.getcwd()
pre_existing_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf")))
# On WSL, patch out sudo check before llama.cpp build
_apply_wsl_sudo_patch()
# Pass absolute path — no os.chdir needed.
# unsloth saves intermediate HF model files into model_save_path,
# while check_llama_cpp("llama.cpp") resolves against cwd (repo root)
# where setup.sh already built llama.cpp with quantizer.
model_save_path = os.path.join(abs_save_dir, "model")
self.current_model.save_pretrained_gguf(
model_save_path,
self.current_tokenizer,
quantization_method=quant_method
)
# Now save (will save in current directory)
self.current_model.save_pretrained_gguf(
"model", # Base filename
self.current_tokenizer,
quantization_method=quant_method
)
# Relocate GGUF artifacts into the export directory.
# convert_to_gguf writes .gguf files to cwd (repo root)
# because --outfile is a relative path like "model.Q4_K_M.gguf".
new_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf"))) - pre_existing_ggufs
for src in sorted(new_ggufs):
dest = os.path.join(abs_save_dir, os.path.basename(src))
shutil.move(src, dest)
logger.info(f"Relocated GGUF: {os.path.basename(src)} → {abs_save_dir}/")
logger.info(f"GGUF model saved successfully in {save_directory}")
# Flatten any .gguf files from subdirectories into abs_save_dir.
# save_pretrained_gguf may create subdirs (e.g. model_gguf/)
# with a name different from model_save_path.
for sub in list(Path(abs_save_dir).iterdir()):
if not sub.is_dir():
continue
for src in sub.glob("*.gguf"):
dest = os.path.join(abs_save_dir, src.name)
shutil.move(str(src), dest)
logger.info(f"Relocated GGUF: {src.name} → {abs_save_dir}/")
# Clean up the subdirectory (intermediate HF files, etc.)
shutil.rmtree(str(sub), ignore_errors=True)
logger.info(f"Cleaned up subdirectory: {sub.name}")
# Check if llama.cpp directory was created here
llama_cpp_in_target = os.path.join(save_directory, "llama.cpp")
llama_cpp_in_original = os.path.join(original_dir, "llama.cpp")
# Write export metadata so the Chat page can identify the base model
self._write_export_metadata(abs_save_dir)
if os.path.exists(llama_cpp_in_target):
logger.info(f"Found llama.cpp directory in {save_directory}")
# Remove llama.cpp from original directory if it exists
if os.path.exists(llama_cpp_in_original):
logger.info(f"Removing existing llama.cpp in {original_dir}")
shutil.rmtree(llama_cpp_in_original)
# Move llama.cpp back to original directory
logger.info(f"Moving llama.cpp to {original_dir}")
shutil.move(llama_cpp_in_target, llama_cpp_in_original)
logger.info(f"Successfully moved llama.cpp back to original directory")
finally:
# Always change back to original directory
os.chdir(original_dir)
logger.info(f"Changed back to original directory: {original_dir}")
logger.info(f"GGUF model saved successfully in {abs_save_dir}")
# Push to hub if requested
if push_to_hub:

2
studio/backend/core/inference/__init__.py
View file

@ -2,8 +2,10 @@
Inference submodule - Inference backend for model loading and generation
"""
from .inference import InferenceBackend, get_inference_backend
from .llama_cpp import LlamaCppBackend
__all__ = [
'InferenceBackend',
'get_inference_backend',
'LlamaCppBackend',
]

420
studio/backend/core/inference/llama_cpp.py Normal file
View file

@ -0,0 +1,420 @@
"""
llama-server inference backend for GGUF models.
Manages a llama-server subprocess and proxies chat completions
through its OpenAI-compatible /v1/chat/completions endpoint.
"""
import atexit
import json
import logging
import shutil
import signal
import socket
import subprocess
import threading
import time
from pathlib import Path
from typing import Generator, Optional
import httpx
logger = logging.getLogger(__name__)
class LlamaCppBackend:
"""
Manages a llama-server subprocess for GGUF model inference.
Lifecycle:
1. load_model() — starts llama-server with the GGUF file
2. generate_chat_completion() — proxies to /v1/chat/completions, streams back
3. unload_model() — terminates llama-server subprocess
"""
def __init__(self):
self._process: Optional[subprocess.Popen] = None
self._port: Optional[int] = None
self._model_identifier: Optional[str] = None
self._gguf_path: Optional[str] = None
self._hf_repo: Optional[str] = None
self._hf_variant: Optional[str] = None
self._is_vision: bool = False
self._healthy = False
self._lock = threading.Lock()
atexit.register(self._cleanup)
# ── Properties ────────────────────────────────────────────────
@property
def is_loaded(self) -> bool:
return self._process is not None and self._healthy
@property
def base_url(self) -> str:
return f"http://127.0.0.1:{self._port}"
@property
def model_identifier(self) -> Optional[str]:
return self._model_identifier
@property
def is_vision(self) -> bool:
return self._is_vision
@property
def hf_variant(self) -> Optional[str]:
return self._hf_variant
# ── Binary discovery ──────────────────────────────────────────
@staticmethod
def _find_llama_server_binary() -> Optional[str]:
"""
Locate the llama-server binary.
Search order:
1. LLAMA_SERVER_PATH environment variable
2. ./llama.cpp/build/bin/llama-server (built by setup.sh in-tree)
3. llama-server on PATH (system install)
4. ./bin/llama-server (legacy: extracted binary)
"""
import os
# 1. Env var
env_path = os.environ.get("LLAMA_SERVER_PATH")
if env_path and Path(env_path).is_file():
return env_path
# Project root: llama_cpp.py → inference/ → core/ → backend/ → studio/ → root
project_root = Path(__file__).resolve().parents[4]
# 2. In-tree llama.cpp build (setup.sh builds here)
build_path = project_root / "llama.cpp" / "build" / "bin" / "llama-server"
if build_path.is_file():
return str(build_path)
# 3. System PATH
system_path = shutil.which("llama-server")
if system_path:
return system_path
# 4. Legacy: extracted to bin/
bin_path = project_root / "bin" / "llama-server"
if bin_path.is_file():
return str(bin_path)
return None
# ── Port allocation ───────────────────────────────────────────
@staticmethod
def _find_free_port() -> int:
"""Find an available TCP port."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("", 0))
return s.getsockname()[1]
# ── Lifecycle ─────────────────────────────────────────────────
def load_model(
self,
*,
# Local mode: pass a path to a .gguf file
gguf_path: Optional[str] = None,
# HF mode: let llama-server download via -hf "repo:quant"
hf_repo: Optional[str] = None,
hf_variant: Optional[str] = None,
hf_token: Optional[str] = None,
# Common
model_identifier: str,
is_vision: bool = False,
n_ctx: int = 4096,
n_gpu_layers: int = -1,
n_threads: Optional[int] = None,
) -> bool:
"""
Start llama-server with a GGUF model.
Two modes:
- Local: ``gguf_path="/path/to/model.gguf"`` → uses ``-m``
- HF: ``hf_repo="unsloth/gemma-3-4b-it-GGUF", hf_variant="Q4_K_M"`` → uses ``-hf``
In HF mode, llama-server handles downloading, caching, and
auto-loading mmproj files for vision models.
Returns True if server started and health check passed.
"""
with self._lock:
self._kill_process()
binary = self._find_llama_server_binary()
if not binary:
raise RuntimeError(
"llama-server binary not found. "
"Run setup.sh to build it, install llama.cpp, "
"or set LLAMA_SERVER_PATH environment variable."
)
self._port = self._find_free_port()
# Build command based on mode
if hf_repo:
hf_spec = f"{hf_repo}:{hf_variant}" if hf_variant else hf_repo
cmd = [
binary,
"-hf", hf_spec,
"--port", str(self._port),
"-c", str(n_ctx),
"-ngl", str(n_gpu_layers),
]
if hf_token:
cmd.extend(["--hf-token", hf_token])
elif gguf_path:
if not Path(gguf_path).is_file():
raise FileNotFoundError(f"GGUF file not found: {gguf_path}")
cmd = [
binary,
"-m", gguf_path,
"--port", str(self._port),
"-c", str(n_ctx),
"-ngl", str(n_gpu_layers),
]
else:
raise ValueError("Either gguf_path or hf_repo must be provided")
if n_threads is not None:
cmd.extend(["--threads", str(n_threads)])
logger.info(f"Starting llama-server: {' '.join(cmd)}")
# Set LD_LIBRARY_PATH so llama-server can find its shared libs
# (libmtmd.so, libllama.so, etc.) which live next to the binary
import os
env = os.environ.copy()
binary_dir = str(Path(binary).parent)
existing_ld = env.get("LD_LIBRARY_PATH", "")
env["LD_LIBRARY_PATH"] = f"{binary_dir}:{existing_ld}" if existing_ld else binary_dir
self._process = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
env=env,
)
self._gguf_path = gguf_path
self._hf_repo = hf_repo
self._hf_variant = hf_variant
self._is_vision = is_vision
self._model_identifier = model_identifier
# HF mode: llama-server downloads before becoming healthy — need longer timeout
timeout = 600.0 if hf_repo else 120.0
if not self._wait_for_health(timeout=timeout):
self._kill_process()
raise RuntimeError(
"llama-server failed to start. "
"Check that the GGUF file is valid and you have enough memory."
)
self._healthy = True
logger.info(
f"llama-server ready on port {self._port} "
f"for model '{model_identifier}'"
)
return True
def unload_model(self) -> bool:
"""Terminate the llama-server subprocess and clean up state."""
with self._lock:
self._kill_process()
logger.info(f"Unloaded GGUF model: {self._model_identifier}")
self._model_identifier = None
self._gguf_path = None
self._hf_repo = None
self._hf_variant = None
self._is_vision = False
self._port = None
self._healthy = False
return True
def _kill_process(self):
"""Terminate the subprocess if running."""
if self._process is None:
return
try:
self._process.terminate()
self._process.wait(timeout=5)
except subprocess.TimeoutExpired:
logger.warning("llama-server did not exit on SIGTERM, sending SIGKILL")
self._process.kill()
self._process.wait(timeout=5)
except Exception as e:
logger.warning(f"Error killing llama-server process: {e}")
finally:
self._process = None
def _cleanup(self):
"""atexit handler to ensure llama-server is terminated."""
self._kill_process()
def _wait_for_health(self, timeout: float = 120.0, interval: float = 0.5) -> bool:
"""
Poll llama-server's /health endpoint until it responds 200.
Also monitors subprocess for early exit/crash.
"""
deadline = time.monotonic() + timeout
url = f"http://127.0.0.1:{self._port}/health"
while time.monotonic() < deadline:
# Check if process crashed
if self._process.poll() is not None:
# Read remaining output for error info
output = self._process.stdout.read() if self._process.stdout else ""
logger.error(
f"llama-server exited with code {self._process.returncode}. "
f"Output: {output[:2000]}"
)
return False
try:
resp = httpx.get(url, timeout=2.0)
if resp.status_code == 200:
return True
except (httpx.ConnectError, httpx.TimeoutException):
pass
time.sleep(interval)
logger.error(f"llama-server health check timed out after {timeout}s")
return False
# ── Message building (OpenAI format) ──────────────────────────
@staticmethod
def _build_openai_messages(
messages: list[dict],
image_b64: Optional[str] = None,
) -> list[dict]:
"""
Build OpenAI-format messages, optionally injecting an image_url
content part into the last user message for vision models.
If no image is provided, returns messages as-is.
"""
if not image_b64:
return messages
# Find the last user message and convert to multimodal content parts
result = [msg.copy() for msg in messages]
last_user_idx = None
for i, msg in enumerate(result):
if msg["role"] == "user":
last_user_idx = i
if last_user_idx is not None:
text_content = result[last_user_idx].get("content", "")
result[last_user_idx]["content"] = [
{"type": "text", "text": text_content},
{
"type": "image_url",
"image_url": {
"url": f"data:image/png;base64,{image_b64}",
},
},
]
return result
# ── Generation (proxy to llama-server) ────────────────────────
def generate_chat_completion(
self,
messages: list[dict],
image_b64: Optional[str] = None,
temperature: float = 0.7,
top_p: float = 0.9,
top_k: int = 40,
min_p: float = 0.0,
max_tokens: int = 512,
repetition_penalty: float = 1.1,
stop: Optional[list[str]] = None,
cancel_event: Optional[threading.Event] = None,
) -> Generator[str, None, None]:
"""
Send a chat completion request to llama-server and stream tokens back.
Uses /v1/chat/completions — llama-server handles chat template
application and vision (multimodal image_url parts) natively.
Yields cumulative text (matching InferenceBackend's convention).
"""
if not self.is_loaded:
raise RuntimeError("llama-server is not loaded")
openai_messages = self._build_openai_messages(messages, image_b64)
payload = {
"messages": openai_messages,
"stream": True,
"temperature": temperature,
"top_p": top_p,
"top_k": top_k if top_k >= 0 else 0,
"min_p": min_p,
"max_tokens": max_tokens,
"repeat_penalty": repetition_penalty,
}
if stop:
payload["stop"] = stop
url = f"{self.base_url}/v1/chat/completions"
cumulative = ""
try:
with httpx.Client(timeout=None) as client:
with client.stream("POST", url, json=payload) as response:
if response.status_code != 200:
error_body = response.read().decode()
raise RuntimeError(
f"llama-server returned {response.status_code}: {error_body}"
)
buffer = ""
for raw_chunk in response.iter_text():
if cancel_event is not None and cancel_event.is_set():
break
buffer += raw_chunk
while "\n" in buffer:
line, buffer = buffer.split("\n", 1)
line = line.strip()
if not line:
continue
if line == "data: [DONE]":
return
if not line.startswith("data: "):
continue
try:
data = json.loads(line[6:])
choices = data.get("choices", [])
if choices:
delta = choices[0].get("delta", {})
token = delta.get("content", "")
if token:
cumulative += token
yield cumulative
except json.JSONDecodeError:
logger.debug(f"Skipping malformed SSE line: {line[:100]}")
except httpx.ConnectError:
raise RuntimeError("Lost connection to llama-server")
except Exception as e:
if cancel_event is not None and cancel_event.is_set():
return
raise

9
studio/backend/core/training/trainer.py
View file

@ -173,6 +173,15 @@ class UnslothTrainer:
token=hf_token,
)
logger.info("Loaded vision model")
# Diagnostic: check if FastVisionModel returned a real Processor or a raw tokenizer
from transformers import ProcessorMixin
tok = self.tokenizer
has_image_proc = isinstance(tok, ProcessorMixin) or hasattr(tok, "image_processor")
print(f"\n[VLM Diagnostic] FastVisionModel returned: {type(tok).__name__}")
print(f"[VLM Diagnostic] Is ProcessorMixin: {isinstance(tok, ProcessorMixin)}")
print(f"[VLM Diagnostic] Has image_processor: {hasattr(tok, 'image_processor')}")
print(f"[VLM Diagnostic] Usable as vision processor: {has_image_proc}\n")
else:
# Load text model - returns (model, tokenizer)
self.model, self.tokenizer = FastLanguageModel.from_pretrained(

4
studio/backend/models/inference.py
View file

@ -17,6 +17,7 @@ class LoadRequest(BaseModel):
max_seq_length: int = Field(2048, ge=128, le=32768, description="Maximum sequence length")
load_in_4bit: bool = Field(True, description="Load model in 4-bit quantization")
is_lora: bool = Field(False, description="Whether this is a LoRA adapter")
gguf_variant: Optional[str] = Field(None, description="GGUF quantization variant (e.g. 'Q4_K_M')")
class UnloadRequest(BaseModel):
@ -43,6 +44,7 @@ class LoadResponse(BaseModel):
display_name: str = Field(..., description="Display name of the model")
is_vision: bool = Field(False, description="Whether model is a vision model")
is_lora: bool = Field(False, description="Whether model is a LoRA adapter")
is_gguf: bool = Field(False, description="Whether model is a GGUF model (llama.cpp)")
inference: dict = Field(..., description="Inference parameters (temperature, top_p, top_k, min_p)")
@ -56,6 +58,8 @@ class InferenceStatusResponse(BaseModel):
"""Current inference backend status"""
active_model: Optional[str] = Field(None, description="Currently active model identifier")
is_vision: bool = Field(False, description="Whether the active model is a vision model")
is_gguf: bool = Field(False, description="Whether the active model is a GGUF model (llama.cpp)")
gguf_variant: Optional[str] = Field(None, description="GGUF quantization variant (e.g. Q4_K_M)")
loading: List[str] = Field(default_factory=list, description="Models currently being loaded")
loaded: List[str] = Field(default_factory=list, description="Models currently loaded")

18
studio/backend/models/models.py
View file

@ -53,6 +53,7 @@ class ModelDetails(BaseModel):
config: Optional[Dict[str, Any]] = Field(None, description="Model configuration dictionary")
is_vision: bool = Field(False, description="Whether model is a vision model")
is_lora: bool = Field(False, description="Whether model is a LoRA adapter")
is_gguf: bool = Field(False, description="Whether model is a GGUF model (llama.cpp format)")
base_model: Optional[str] = Field(None, description="Base model if this is a LoRA adapter")
@ -62,7 +63,7 @@ class LoRAInfo(BaseModel):
adapter_path: str = Field(..., description="Path to the LoRA adapter or exported model")
base_model: Optional[str] = Field(None, description="Base model identifier")
source: Optional[str] = Field(None, description="'training' or 'exported'")
export_type: Optional[str] = Field(None, description="'lora' or 'merged' (for exports)")
export_type: Optional[str] = Field(None, description="'lora', 'merged', or 'gguf' (for exports)")
class LoRAScanResponse(BaseModel):
@ -77,6 +78,21 @@ class ModelListResponse(BaseModel):
default_models: List[str] = Field(default_factory=list, description="List of default model IDs")
class GgufVariantDetail(BaseModel):
"""A single GGUF quantization variant in a HuggingFace repo."""
filename: str = Field(..., description="GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')")
quant: str = Field(..., description="Quantization label (e.g., 'Q4_K_M')")
size_bytes: int = Field(0, description="File size in bytes")
class GgufVariantsResponse(BaseModel):
"""Response for listing GGUF quantization variants in a HuggingFace repo."""
repo_id: str = Field(..., description="HuggingFace repo ID")
variants: List[GgufVariantDetail] = Field(default_factory=list, description="Available GGUF variants")
has_vision: bool = Field(False, description="Whether the model has vision support (mmproj files)")
default_variant: Optional[str] = Field(None, description="Recommended default quantization variant")
class LocalModelInfo(BaseModel):
"""Discovered local model candidate."""
id: str = Field(..., description="Identifier to use for loading/training")

23
studio/backend/routes/data_recipe/__init__.py Normal file
View file

@ -0,0 +1,23 @@
"""Data Recipe route package."""
from __future__ import annotations
import sys
from pathlib import Path
from fastapi import APIRouter
backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
sys.path.insert(0, str(backend_path))
from .jobs import router as jobs_router
from .seed import router as seed_router
from .validate import router as validate_router
router = APIRouter()
router.include_router(seed_router)
router.include_router(validate_router)
router.include_router(jobs_router)
__all__ = ["router"]

143
studio/backend/routes/data_recipe/jobs.py Normal file
View file

@ -0,0 +1,143 @@
"""Job lifecycle endpoints for data recipe."""
from __future__ import annotations
from typing import Any
from fastapi import APIRouter, HTTPException, Query, Request
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import ValidationError
from core.data_recipe.jobs import get_job_manager
from models.data_recipe import JobCreateResponse, RecipePayload
router = APIRouter()
@router.post("/jobs", response_class=JSONResponse, response_model=JobCreateResponse)
def create_job(payload: RecipePayload):
recipe = payload.recipe
if not recipe.get("columns"):
raise HTTPException(status_code=400, detail="Recipe must include columns.")
run: dict[str, Any] = payload.run or {}
run.pop("artifact_path", None)
run.pop("dataset_name", None)
execution_type = str(run.get("execution_type") or "full").strip().lower()
if execution_type not in {"preview", "full"}:
raise HTTPException(status_code=400, detail="invalid execution_type: must be 'preview' or 'full'")
run["execution_type"] = execution_type
run_config_raw = run.get("run_config")
if run_config_raw is not None:
try:
from data_designer.config.run_config import RunConfig
RunConfig.model_validate(run_config_raw)
except (ImportError, ValidationError, TypeError, ValueError) as exc:
raise HTTPException(status_code=400, detail=f"invalid run_config: {exc}") from exc
mgr = get_job_manager()
try:
job_id = mgr.start(recipe=recipe, run=run)
except RuntimeError as exc:
raise HTTPException(status_code=409, detail=str(exc)) from exc
except ValueError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
return {"job_id": job_id}
@router.get("/jobs/{job_id}/status")
def job_status(job_id: str):
mgr = get_job_manager()
state = mgr.get_status(job_id)
if state is None:
raise HTTPException(status_code=404, detail="job not found")
return state
@router.get("/jobs/current")
def current_job():
mgr = get_job_manager()
state = mgr.get_current_status()
if state is None:
raise HTTPException(status_code=404, detail="no job")
return state
@router.post("/jobs/{job_id}/cancel")
def cancel_job(job_id: str):
mgr = get_job_manager()
ok = mgr.cancel(job_id)
if not ok:
raise HTTPException(status_code=404, detail="job not found")
return mgr.get_status(job_id)
@router.get("/jobs/{job_id}/analysis")
def job_analysis(job_id: str):
mgr = get_job_manager()
analysis = mgr.get_analysis(job_id)
if analysis is None:
raise HTTPException(status_code=404, detail="analysis not ready")
return analysis
@router.get("/jobs/{job_id}/dataset")
def job_dataset(
job_id: str,
limit: int = Query(default=20, ge=1, le=500),
offset: int = Query(default=0, ge=0),
):
mgr = get_job_manager()
result = mgr.get_dataset(job_id, limit=limit, offset=offset)
if result is None:
raise HTTPException(status_code=404, detail="dataset not ready")
if "error" in result:
raise HTTPException(status_code=422, detail=result["error"])
return {
"dataset": result["dataset"],
"total": result["total"],
"limit": limit,
"offset": offset,
}
@router.get("/jobs/{job_id}/events")
async def job_events(request: Request, job_id: str):
mgr = get_job_manager()
last_id = request.headers.get("last-event-id")
after_seq: int | None = None
if last_id:
try:
after_seq = int(str(last_id).strip())
except (TypeError, ValueError):
after_seq = None
after_q = request.query_params.get("after")
if after_q:
try:
after_seq = int(str(after_q).strip())
except (TypeError, ValueError):
pass
sub = mgr.subscribe(job_id, after_seq=after_seq)
if sub is None:
raise HTTPException(status_code=404, detail="job not found")
async def gen():
try:
for event in sub.replay:
yield sub.format_sse(event)
while True:
if await request.is_disconnected():
break
event = await sub.next_event(timeout_sec=1.0)
if event is None:
continue
yield sub.format_sse(event)
finally:
mgr.unsubscribe(sub)
return StreamingResponse(gen(), media_type="text/event-stream")

269
studio/backend/routes/data_recipe.py → studio/backend/routes/data_recipe/seed.py
View file

@ -1,42 +1,24 @@
"""
Data Recipe routes (DataDesigner runner).
"""
"""Seed inspect endpoints for data recipe."""
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
from fastapi import APIRouter, HTTPException
# same thing as other files do
backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
sys.path.insert(0, str(backend_path))
from core.data_recipe.jobs import get_job_manager
from core.data_recipe.service import (
build_config_builder,
create_data_designer,
validate_recipe,
)
from models.data_recipe import (
JobCreateResponse,
RecipePayload,
SeedInspectRequest,
SeedInspectUploadRequest,
SeedInspectResponse,
ValidateError,
ValidateResponse,
SeedInspectUploadRequest,
)
router = APIRouter()
DATA_EXTS = (".parquet", ".jsonl", ".json", ".csv")
DEFAULT_SPLIT = "train"
LOCAL_UPLOAD_EXTS = {".csv", ".json", ".jsonl"}
@ -70,20 +52,21 @@ def _normalize_optional_text(value: str | None) -> str | None:
def _list_hf_data_files(*, dataset_name: str, token: str | None) -> list[str]:
try:
from huggingface_hub import HfApi
from huggingface_hub.utils import HfHubHTTPError
except ImportError:
return []
try:
api = HfApi()
repo_files = api.list_repo_files(dataset_name, repo_type="dataset", token=token)
return [file for file in repo_files if file.lower().endswith(DATA_EXTS)]
except Exception:
except (HfHubHTTPError, OSError, ValueError):
return []
def _select_best_file(data_files: list[str], split: str | None) -> str | None:
def _select_best_file(data_files: list[str]) -> str | None:
if not data_files:
return None
if not split:
return data_files[0]
split_lower = split.lower()
split_lower = DEFAULT_SPLIT
def score(path: str) -> tuple[int, int]:
name = path.lower()
@ -102,8 +85,8 @@ def _select_best_file(data_files: list[str], split: str | None) -> str | None:
return sorted(data_files, key=score)[0]
def _resolve_seed_hf_path(dataset_name: str, data_files: list[str], split: str | None) -> str | None:
selected = _select_best_file(data_files, split)
def _resolve_seed_hf_path(dataset_name: str, data_files: list[str]) -> str | None:
selected = _select_best_file(data_files)
if not selected:
return None
@ -177,7 +160,7 @@ def _decode_base64_payload(content_base64: str) -> bytes:
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:
except ImportError as exc:
raise HTTPException(status_code=500, detail=f"seed inspect dependencies unavailable: {exc}") from exc
ext = path.suffix.lower()
@ -189,68 +172,19 @@ def _read_preview_rows_from_local_file(path: Path, preview_size: int) -> list[di
elif ext == ".json":
try:
df = pd.read_json(path, lines=True).head(preview_size)
except Exception:
except ValueError:
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:
except (ValueError, OSError) 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)
def _collect_validation_errors(recipe: dict[str, Any]) -> list[ValidateError]:
try:
from data_designer.engine.compiler import (
_add_internal_row_id_column_if_needed,
_get_allowed_references,
_resolve_and_add_seed_columns,
)
from data_designer.engine.validation import (
ViolationLevel,
validate_data_designer_config,
)
except Exception:
return []
try:
builder = build_config_builder(recipe)
designer = create_data_designer(recipe)
resource_provider = designer._create_resource_provider( # type: ignore[attr-defined]
"validate-configuration",
builder,
)
config = builder.build()
_resolve_and_add_seed_columns(config, resource_provider.seed_reader)
_add_internal_row_id_column_if_needed(config)
violations = validate_data_designer_config(
columns=config.columns,
processor_configs=config.processors or [],
allowed_references=_get_allowed_references(config),
)
except Exception:
return []
errors: list[ValidateError] = []
for violation in violations:
if violation.level != ViolationLevel.ERROR:
continue
code = getattr(violation.type, "value", None)
path = violation.column if violation.column else None
message = str(violation.message).strip() or "Validation failed."
errors.append(
ValidateError(
message=message,
path=path,
code=code,
)
)
return errors
@router.post("/seed/inspect", response_model=SeedInspectResponse)
def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
dataset_name = payload.dataset_name.strip()
@ -259,10 +193,10 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
try:
from datasets import load_dataset
except Exception as exc:
except ImportError as exc:
raise HTTPException(status_code=500, detail=f"seed inspect dependencies unavailable: {exc}") from exc
split = (payload.split or DEFAULT_SPLIT).strip() or DEFAULT_SPLIT
split = DEFAULT_SPLIT
subset = _normalize_optional_text(payload.subset)
token = _normalize_optional_text(payload.hf_token)
preview_size = int(payload.preview_size)
@ -270,7 +204,7 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
preview_rows: list[dict[str, Any]] = []
data_files = _list_hf_data_files(dataset_name=dataset_name, token=token)
selected_file = _select_best_file(data_files, split)
selected_file = _select_best_file(data_files)
if selected_file:
try:
single_file_kwargs = _build_stream_load_kwargs(
@ -285,7 +219,7 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
load_kwargs=single_file_kwargs,
preview_size=preview_size,
)
except Exception:
except (ValueError, OSError, RuntimeError):
preview_rows = []
if not preview_rows:
@ -301,7 +235,7 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
load_kwargs=split_kwargs,
preview_size=preview_size,
)
except Exception as exc:
except (ValueError, OSError, RuntimeError) as exc:
raise HTTPException(status_code=422, detail=f"seed inspect failed: {exc}") from exc
if not preview_rows:
@ -310,10 +244,9 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
columns = _extract_columns(preview_rows)
if not data_files:
# Best effort path fallback when file list is unavailable.
resolved_path = f"datasets/{dataset_name}/**/*.parquet"
else:
resolved_path = _resolve_seed_hf_path(dataset_name, data_files, split)
resolved_path = _resolve_seed_hf_path(dataset_name, data_files)
if not resolved_path:
raise HTTPException(status_code=422, detail="unable to resolve seed dataset path")
@ -322,7 +255,7 @@ def inspect_seed_dataset(payload: SeedInspectRequest) -> SeedInspectResponse:
resolved_path=resolved_path,
columns=columns,
preview_rows=preview_rows,
split=split,
split=None,
subset=subset,
)
@ -363,159 +296,3 @@ def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectRespons
split=None,
subset=None,
)
@router.post("/validate", response_model=ValidateResponse)
def validate(payload: RecipePayload) -> ValidateResponse:
recipe = payload.recipe
if not recipe.get("columns"):
return ValidateResponse(
valid=False,
errors=[ValidateError(message="Recipe must include columns.")],
)
try:
validate_recipe(recipe)
except RuntimeError as exc:
raise HTTPException(status_code=503, detail=str(exc)) from exc
except Exception as exc:
detail = str(exc).strip() or "Validation failed."
parsed_errors = _collect_validation_errors(recipe)
return ValidateResponse(
valid=False,
errors=parsed_errors or [ValidateError(message=detail)],
raw_detail=detail,
)
return ValidateResponse(valid=True)
@router.post("/jobs", response_class=JSONResponse, response_model=JobCreateResponse)
def create_job(payload: RecipePayload):
recipe = payload.recipe
if not recipe.get("columns"):
raise HTTPException(status_code=400, detail="Recipe must include columns.")
run: dict[str, Any] = payload.run or {}
run.pop("artifact_path", None)
run.pop("dataset_name", None)
execution_type = str(run.get("execution_type") or "full").strip().lower()
if execution_type not in {"preview", "full"}:
raise HTTPException(status_code=400, detail="invalid execution_type: must be 'preview' or 'full'")
run["execution_type"] = execution_type
run_config_raw = run.get("run_config")
if run_config_raw is not None:
try:
from data_designer.config.run_config import RunConfig
RunConfig.model_validate(run_config_raw)
except Exception as exc:
raise HTTPException(status_code=400, detail=f"invalid run_config: {exc}") from exc
mgr = get_job_manager()
try:
job_id = mgr.start(recipe=recipe, run=run)
except RuntimeError as exc:
raise HTTPException(status_code=409, detail=str(exc)) from exc
except ValueError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
return {"job_id": job_id}
@router.get("/jobs/{job_id}/status")
def job_status(job_id: str):
mgr = get_job_manager()
state = mgr.get_status(job_id)
if state is None:
raise HTTPException(status_code=404, detail="job not found")
return state
@router.get("/jobs/current")
def current_job():
mgr = get_job_manager()
state = mgr.get_current_status()
if state is None:
raise HTTPException(status_code=404, detail="no job")
return state
@router.post("/jobs/{job_id}/cancel")
def cancel_job(job_id: str):
mgr = get_job_manager()
ok = mgr.cancel(job_id)
if not ok:
raise HTTPException(status_code=404, detail="job not found")
return mgr.get_status(job_id)
@router.get("/jobs/{job_id}/analysis")
def job_analysis(job_id: str):
mgr = get_job_manager()
analysis = mgr.get_analysis(job_id)
if analysis is None:
raise HTTPException(status_code=404, detail="analysis not ready")
return analysis
@router.get("/jobs/{job_id}/dataset")
def job_dataset(
job_id: str,
limit: int = Query(default=20, ge=1, le=500),
offset: int = Query(default=0, ge=0),
):
mgr = get_job_manager()
result = mgr.get_dataset(job_id, limit=limit, offset=offset)
if result is None:
raise HTTPException(status_code=404, detail="dataset not ready")
if "error" in result:
raise HTTPException(status_code=422, detail=result["error"])
return {
"dataset": result["dataset"],
"total": result["total"],
"limit": limit,
"offset": offset,
}
@router.get("/jobs/{job_id}/events")
async def job_events(request: Request, job_id: str):
mgr = get_job_manager()
last_id = request.headers.get("last-event-id")
after_seq: int | None = None
if last_id:
try:
after_seq = int(str(last_id).strip())
except Exception:
after_seq = None
# EventSource can't set custom headers on first connect after a full page refresh,
# so allow resume via query param too: /events?after=<seq>
after_q = request.query_params.get("after")
if after_q:
try:
after_seq = int(str(after_q).strip())
except Exception:
pass
sub = mgr.subscribe(job_id, after_seq=after_seq)
if sub is None:
raise HTTPException(status_code=404, detail="job not found")
async def gen():
try:
for event in sub.replay:
yield sub.format_sse(event)
while True:
if await request.is_disconnected():
break
event = await sub.next_event(timeout_sec=1.0)
if event is None:
continue
yield sub.format_sse(event)
finally:
mgr.unsubscribe(sub)
return StreamingResponse(gen(), media_type="text/event-stream")

90
studio/backend/routes/data_recipe/validate.py Normal file
View file

@ -0,0 +1,90 @@
"""Validation endpoints for data recipe."""
from __future__ import annotations
from typing import Any
from fastapi import APIRouter, HTTPException
from core.data_recipe.service import (
build_config_builder,
create_data_designer,
validate_recipe,
)
from models.data_recipe import RecipePayload, ValidateError, ValidateResponse
router = APIRouter()
def _collect_validation_errors(recipe: dict[str, Any]) -> list[ValidateError]:
try:
from data_designer.engine.compiler import (
_add_internal_row_id_column_if_needed,
_get_allowed_references,
_resolve_and_add_seed_columns,
)
from data_designer.engine.validation import (
ViolationLevel,
validate_data_designer_config,
)
except ImportError:
return []
try:
builder = build_config_builder(recipe)
designer = create_data_designer(recipe)
resource_provider = designer._create_resource_provider( # type: ignore[attr-defined]
"validate-configuration",
builder,
)
config = builder.build()
_resolve_and_add_seed_columns(config, resource_provider.seed_reader)
_add_internal_row_id_column_if_needed(config)
violations = validate_data_designer_config(
columns=config.columns,
processor_configs=config.processors or [],
allowed_references=_get_allowed_references(config),
)
except (TypeError, ValueError, AttributeError):
return []
errors: list[ValidateError] = []
for violation in violations:
if violation.level != ViolationLevel.ERROR:
continue
code = getattr(violation.type, "value", None)
path = violation.column if violation.column else None
message = str(violation.message).strip() or "Validation failed."
errors.append(
ValidateError(
message=message,
path=path,
code=code,
)
)
return errors
@router.post("/validate", response_model=ValidateResponse)
def validate(payload: RecipePayload) -> ValidateResponse:
recipe = payload.recipe
if not recipe.get("columns"):
return ValidateResponse(
valid=False,
errors=[ValidateError(message="Recipe must include columns.")],
)
try:
validate_recipe(recipe)
except RuntimeError as exc:
raise HTTPException(status_code=503, detail=str(exc)) from exc
except Exception as exc:
detail = str(exc).strip() or "Validation failed."
parsed_errors = _collect_validation_errors(recipe)
return ValidateResponse(
valid=False,
errors=parsed_errors or [ValidateError(message=detail)],
raw_detail=detail,
)
return ValidateResponse(valid=True)

318
studio/backend/routes/inference.py
View file

@ -5,7 +5,7 @@ import sys
import time
import uuid
from pathlib import Path
from fastapi import APIRouter, HTTPException, Request
from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi.responses import StreamingResponse, JSONResponse
from typing import Optional
import json
@ -23,6 +23,7 @@ if str(backend_path) not in sys.path:
# Import backend functions
try:
from core.inference import get_inference_backend
from core.inference.llama_cpp import LlamaCppBackend
from utils.models import ModelConfig
from utils.inference import load_inference_config
except ImportError:
@ -30,6 +31,7 @@ except ImportError:
if str(parent_backend) not in sys.path:
sys.path.insert(0, str(parent_backend))
from core.inference import get_inference_backend
from core.inference.llama_cpp import LlamaCppBackend
from utils.models import ModelConfig
from utils.inference import load_inference_config
@ -48,6 +50,7 @@ from models.inference import (
CompletionChoice,
CompletionMessage,
)
from auth.authentication import get_current_subject
router = APIRouter()
logger = logging.getLogger(__name__)
@ -61,60 +64,129 @@ if not logger.handlers:
logger.addHandler(handler)
logger.setLevel(logging.INFO)
# GGUF inference backend (llama-server)
_llama_cpp_backend = LlamaCppBackend()
def get_llama_cpp_backend() -> LlamaCppBackend:
return _llama_cpp_backend
@router.post("/load", response_model=LoadResponse)
async def load_model(request: LoadRequest):
async def load_model(
request: LoadRequest,
current_subject: str = Depends(get_current_subject),
):
"""
Load a model for inference.
The model_path should be a clean identifier from GET /models/list.
Returns inference configuration parameters (temperature, top_p, top_k, min_p)
from the model's YAML config, falling back to default.yaml for missing values.
GGUF models are loaded via llama-server (llama.cpp) instead of Unsloth.
"""
try:
backend = get_inference_backend()
# Create config using clean factory method
# is_lora is auto-detected from adapter_config.json on disk/HF
config = ModelConfig.from_identifier(
model_id=request.model_path,
hf_token=request.hf_token,
gguf_variant=request.gguf_variant,
)
if not config:
raise HTTPException(
status_code=400,
detail=f"Invalid model identifier: {request.model_path}"
)
# Load the model
# ── GGUF path: load via llama-server ──────────────────────
if config.is_gguf:
llama_backend = get_llama_cpp_backend()
unsloth_backend = get_inference_backend()
# Unload any active Unsloth model first to free VRAM
if unsloth_backend.active_model_name:
logger.info(f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF")
unsloth_backend.unload_model(unsloth_backend.active_model_name)
# Route to HF mode or local mode based on config
if config.gguf_hf_repo:
# HF mode: llama-server downloads via -hf "repo:quant"
success = llama_backend.load_model(
hf_repo=config.gguf_hf_repo,
hf_variant=config.gguf_variant,
hf_token=request.hf_token,
model_identifier=config.identifier,
is_vision=config.is_vision,
n_ctx=request.max_seq_length,
)
else:
# Local mode: llama-server loads via -m <path>
success = llama_backend.load_model(
gguf_path=config.gguf_file,
model_identifier=config.identifier,
is_vision=config.is_vision,
n_ctx=request.max_seq_length,
)
if not success:
raise HTTPException(
status_code=500,
detail=f"Failed to load GGUF model: {config.display_name}"
)
logger.info(f"Loaded GGUF model via llama-server: {config.identifier}")
inference_config = load_inference_config(config.identifier)
return LoadResponse(
status="loaded",
model=config.identifier,
display_name=config.display_name,
is_vision=config.is_vision,
is_lora=False,
is_gguf=True,
inference=inference_config,
)
# ── Standard path: load via Unsloth/transformers ──────────
backend = get_inference_backend()
# Unload any active GGUF model first
llama_backend = get_llama_cpp_backend()
if llama_backend.is_loaded:
logger.info("Unloading GGUF model before loading Unsloth model")
llama_backend.unload_model()
success = backend.load_model(
config=config,
max_seq_length=request.max_seq_length,
load_in_4bit=request.load_in_4bit,
hf_token=request.hf_token,
)
if not success:
raise HTTPException(
status_code=500,
detail=f"Failed to load model: {config.display_name}"
)
logger.info(f"Loaded model: {config.identifier}")
# Load inference configuration parameters
inference_config = load_inference_config(config.identifier)
return LoadResponse(
status="loaded",
model=config.identifier,
display_name=config.display_name,
is_vision=config.is_vision,
is_lora=config.is_lora,
is_gguf=False,
inference=inference_config,
)
except HTTPException:
raise
except Exception as e:
@ -126,16 +198,28 @@ async def load_model(request: LoadRequest):
@router.post("/unload", response_model=UnloadResponse)
async def unload_model(request: UnloadRequest):
async def unload_model(
request: UnloadRequest,
current_subject: str = Depends(get_current_subject),
):
"""
Unload a model from memory.
Routes to the correct backend (llama-server for GGUF, Unsloth otherwise).
"""
try:
# Check if the GGUF backend has this model loaded
llama_backend = get_llama_cpp_backend()
if llama_backend.is_loaded and llama_backend.model_identifier == request.model_path:
llama_backend.unload_model()
logger.info(f"Unloaded GGUF model: {request.model_path}")
return UnloadResponse(status="unloaded", model=request.model_path)
# Otherwise, unload from Unsloth backend
backend = get_inference_backend()
backend.unload_model(request.model_path)
logger.info(f"Unloaded model: {request.model_path}")
return UnloadResponse(status="unloaded", model=request.model_path)
except Exception as e:
logger.error(f"Error unloading model: {e}", exc_info=True)
raise HTTPException(
@ -145,7 +229,10 @@ async def unload_model(request: UnloadRequest):
@router.post("/generate/stream")
async def generate_stream(request: GenerateRequest):
async def generate_stream(
request: GenerateRequest,
current_subject: str = Depends(get_current_subject),
):
"""
Generate a chat response with Server-Sent Events (SSE) streaming.
@ -218,25 +305,43 @@ async def generate_stream(request: GenerateRequest):
@router.get("/status", response_model=InferenceStatusResponse)
async def get_status():
async def get_status(
current_subject: str = Depends(get_current_subject),
):
"""
Get current inference backend status.
Reports whichever backend (Unsloth or llama-server) is currently active.
"""
try:
llama_backend = get_llama_cpp_backend()
# If a GGUF model is loaded via llama-server, report that
if llama_backend.is_loaded:
return InferenceStatusResponse(
active_model=llama_backend.model_identifier,
is_vision=llama_backend.is_vision,
is_gguf=True,
gguf_variant=llama_backend.hf_variant,
loading=[],
loaded=[llama_backend.model_identifier],
)
# Otherwise, report Unsloth backend status
backend = get_inference_backend()
is_vision = False
if backend.active_model_name:
model_info = backend.models.get(backend.active_model_name, {})
is_vision = model_info.get("is_vision", False)
return InferenceStatusResponse(
active_model=backend.active_model_name,
is_vision=is_vision,
is_gguf=False,
loading=list(getattr(backend, 'loading_models', set())),
loaded=list(backend.models.keys()),
)
except Exception as e:
logger.error(f"Error getting status: {e}", exc_info=True)
raise HTTPException(
@ -306,7 +411,11 @@ def _extract_content_parts(
@router.post("/chat/completions")
async def openai_chat_completions(payload: ChatCompletionRequest, request: Request):
async def openai_chat_completions(
payload: ChatCompletionRequest,
request: Request,
current_subject: str = Depends(get_current_subject),
):
"""
OpenAI-compatible chat completions endpoint.
@ -315,29 +424,163 @@ async def openai_chat_completions(payload: ChatCompletionRequest, request: Reque
Streaming (default): returns SSE chunks matching OpenAI's format.
Non-streaming: returns a single ChatCompletion JSON object.
"""
backend = get_inference_backend()
if not backend.active_model_name:
raise HTTPException(
status_code=400,
detail="No model loaded. Call POST /inference/load first.",
)
Automatically routes to the correct backend:
- GGUF models → llama-server via LlamaCppBackend
- Other models → Unsloth/transformers via InferenceBackend
"""
llama_backend = get_llama_cpp_backend()
using_gguf = llama_backend.is_loaded
# ── Determine which backend is active ─────────────────────
if using_gguf:
model_name = llama_backend.model_identifier or payload.model
else:
backend = get_inference_backend()
if not backend.active_model_name:
raise HTTPException(
status_code=400,
detail="No model loaded. Call POST /inference/load first.",
)
model_name = backend.active_model_name or payload.model
# ── Parse messages (handles multimodal content parts) ─────
system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(
payload.messages
)
# If no non-system messages were provided, error out
if not chat_messages:
raise HTTPException(
status_code=400,
detail="At least one non-system message is required.",
)
# ── Decode image (from content parts OR legacy field) ─────
# Content-part images take priority; fall back to legacy field
# ── GGUF path: proxy to llama-server /v1/chat/completions ──
if using_gguf:
# Reject images if this GGUF model doesn't support vision
image_b64 = extracted_image_b64 or payload.image_base64
if image_b64 and not llama_backend.is_vision:
raise HTTPException(
status_code=400,
detail="Image provided but current GGUF model does not support vision.",
)
# Build message list with system prompt prepended
gguf_messages = []
if system_prompt:
gguf_messages.append({"role": "system", "content": system_prompt})
gguf_messages.extend(chat_messages)
cancel_event = threading.Event()
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
created = int(time.time())
def gguf_generate():
return llama_backend.generate_chat_completion(
messages=gguf_messages,
image_b64=image_b64,
temperature=payload.temperature,
top_p=payload.top_p,
top_k=payload.top_k,
min_p=payload.min_p,
max_tokens=payload.max_tokens or 512,
repetition_penalty=payload.repetition_penalty,
cancel_event=cancel_event,
)
if payload.stream:
async def gguf_stream_chunks():
try:
# First chunk: role
first_chunk = ChatCompletionChunk(
id=completion_id,
created=created,
model=model_name,
choices=[ChunkChoice(
delta=ChoiceDelta(role="assistant"),
finish_reason=None,
)],
)
yield f"data: {first_chunk.model_dump_json(exclude_none=True)}\n\n"
# Content chunks — llama backend yields cumulative text
prev_text = ""
for cumulative in gguf_generate():
if await request.is_disconnected():
cancel_event.set()
return
new_text = cumulative[len(prev_text):]
prev_text = cumulative
if not new_text:
continue
chunk = ChatCompletionChunk(
id=completion_id,
created=created,
model=model_name,
choices=[ChunkChoice(
delta=ChoiceDelta(content=new_text),
finish_reason=None,
)],
)
yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n"
# Final chunk
final_chunk = ChatCompletionChunk(
id=completion_id,
created=created,
model=model_name,
choices=[ChunkChoice(
delta=ChoiceDelta(),
finish_reason="stop",
)],
)
yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n"
yield "data: [DONE]\n\n"
except asyncio.CancelledError:
cancel_event.set()
raise
except Exception as e:
logger.error(f"Error during GGUF streaming: {e}", exc_info=True)
error_chunk = {
"error": {"message": str(e), "type": "server_error"},
}
yield f"data: {json.dumps(error_chunk)}\n\n"
return StreamingResponse(
gguf_stream_chunks(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)
else:
try:
full_text = ""
for token in gguf_generate():
full_text = token
response = ChatCompletion(
id=completion_id,
created=created,
model=model_name,
choices=[CompletionChoice(
message=CompletionMessage(content=full_text),
finish_reason="stop",
)],
)
return JSONResponse(content=response.model_dump())
except Exception as e:
logger.error(f"Error during GGUF completion: {e}", exc_info=True)
raise HTTPException(status_code=500, detail=str(e))
# ── Standard Unsloth path ─────────────────────────────────
# Decode image (from content parts OR legacy field)
image_b64 = extracted_image_b64 or payload.image_base64
image = None
@ -363,7 +606,7 @@ async def openai_chat_completions(payload: ChatCompletionRequest, request: Reque
except Exception as e:
raise HTTPException(status_code=400, detail=f"Failed to decode image: {e}")
# ── Shared generation kwargs ──────────────────────────────
# Shared generation kwargs
gen_kwargs = dict(
messages=chat_messages,
system_prompt=system_prompt,
@ -376,11 +619,10 @@ async def openai_chat_completions(payload: ChatCompletionRequest, request: Reque
repetition_penalty=payload.repetition_penalty,
)
# ── Choose generation path (adapter-controlled or standard) ──
# Choose generation path (adapter-controlled or standard)
cancel_event = threading.Event()
if payload.use_adapter is not None:
# Compare mode: toggle adapter state atomically with generation
def generate():
return backend.generate_with_adapter_control(
use_adapter=payload.use_adapter,
@ -388,11 +630,9 @@ async def openai_chat_completions(payload: ChatCompletionRequest, request: Reque
**gen_kwargs,
)
else:
# Standard path: no adapter toggling
def generate():
return backend.generate_chat_response(cancel_event=cancel_event, **gen_kwargs)
model_name = backend.active_model_name or payload.model
completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
created = int(time.time())
@ -400,7 +640,6 @@ async def openai_chat_completions(payload: ChatCompletionRequest, request: Reque
if payload.stream:
async def stream_chunks():
try:
# First chunk: send the role
first_chunk = ChatCompletionChunk(
id=completion_id,
created=created,
@ -412,8 +651,6 @@ async def openai_chat_completions(payload: ChatCompletionRequest, request: Reque
)
yield f"data: {first_chunk.model_dump_json(exclude_none=True)}\n\n"
# Content chunks — generate_chat_response yields cumulative
# text, so we diff to get incremental deltas.
prev_text = ""
for cumulative in generate():
if await request.is_disconnected():
@ -435,7 +672,6 @@ async def openai_chat_completions(payload: ChatCompletionRequest, request: Reque
)
yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n"
# Final chunk: finish_reason = stop
final_chunk = ChatCompletionChunk(
id=completion_id,
created=created,
@ -475,7 +711,7 @@ async def openai_chat_completions(payload: ChatCompletionRequest, request: Reque
try:
full_text = ""
for token in generate():
full_text = token # generate_stream yields cumulative text
full_text = token
response = ChatCompletion(
id=completion_id,

66
studio/backend/routes/models.py
View file

@ -23,8 +23,10 @@ try:
get_base_model_from_lora,
is_vision_model,
scan_checkpoints,
list_gguf_variants,
ModelConfig,
)
from utils.models.model_config import _pick_best_gguf, _extract_quant_label
from core.inference import get_inference_backend
except ImportError:
# Fallback: try to import from parent directory
@ -38,8 +40,10 @@ except ImportError:
get_base_model_from_lora,
is_vision_model,
scan_checkpoints,
list_gguf_variants,
ModelConfig,
)
from utils.models.model_config import _pick_best_gguf, _extract_quant_label
from core.inference import get_inference_backend
from models import (
@ -53,6 +57,7 @@ from models import (
LoRAInfo,
ModelListResponse,
)
from models.models import GgufVariantDetail, GgufVariantsResponse
from models.responses import LoRABaseModelResponse, VisionCheckResponse
router = APIRouter()
@ -90,6 +95,7 @@ def _scan_models_dir(models_dir: Path) -> List[LocalModelInfo]:
or (child / "adapter_config.json").exists()
or any(child.glob("*.safetensors"))
or any(child.glob("*.bin"))
or any(child.glob("*.gguf"))
)
if not has_model_files:
continue
@ -106,6 +112,23 @@ def _scan_models_dir(models_dir: Path) -> List[LocalModelInfo]:
updated_at=updated_at,
),
)
# Also scan for standalone .gguf files directly in the models directory
for gguf_file in models_dir.glob("*.gguf"):
if gguf_file.is_file():
try:
updated_at = gguf_file.stat().st_mtime
except OSError:
updated_at = None
found.append(
LocalModelInfo(
id=str(gguf_file),
display_name=gguf_file.stem,
path=str(gguf_file),
source="models_dir",
updated_at=updated_at,
),
)
return found
@ -399,6 +422,49 @@ async def check_vision_model(
detail=f"Failed to check vision model: {str(e)}"
)
@router.get("/gguf-variants", response_model=GgufVariantsResponse)
async def get_gguf_variants(
repo_id: str = Query(..., description="HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')"),
hf_token: Optional[str] = Query(None, description="HuggingFace token for private repos"),
current_subject: str = Depends(get_current_subject),
):
"""
List available GGUF quantization variants for a HuggingFace repo.
Returns all available quantization variants (Q4_K_M, Q8_0, BF16, etc.)
with file sizes, whether the model supports vision, and the recommended
default variant.
"""
try:
variants, has_vision = list_gguf_variants(repo_id, hf_token=hf_token)
# Determine default variant
filenames = [v.filename for v in variants]
best = _pick_best_gguf(filenames)
default_variant = _extract_quant_label(best) if best else None
return GgufVariantsResponse(
repo_id=repo_id,
variants=[
GgufVariantDetail(
filename=v.filename,
quant=v.quant,
size_bytes=v.size_bytes,
)
for v in variants
],
has_vision=has_vision,
default_variant=default_variant,
)
except Exception as e:
logger.error(f"Error listing GGUF variants for '{repo_id}': {e}", exc_info=True)
raise HTTPException(
status_code=500,
detail=f"Failed to list GGUF variants: {str(e)}",
)
@router.get("/checkpoints", response_model=CheckpointListResponse)
async def list_checkpoints(
outputs_dir: str = Query(

4
studio/backend/utils/models/__init__.py
View file

@ -3,12 +3,14 @@ Model and LoRA configuration handling
"""
from .model_config import (
ModelConfig,
GgufVariantInfo,
is_vision_model,
scan_trained_loras,
scan_exported_models,
load_model_defaults,
get_base_model_from_lora,
load_model_config,
list_gguf_variants,
MODEL_NAME_MAPPING,
UI_STATUS_INDICATORS,
)
@ -16,12 +18,14 @@ from .checkpoints import scan_checkpoints
__all__ = [
'ModelConfig',
'GgufVariantInfo',
'is_vision_model',
'scan_trained_loras',
'scan_exported_models',
'load_model_defaults',
'get_base_model_from_lora',
'load_model_config',
'list_gguf_variants',
'MODEL_NAME_MAPPING',
'UI_STATUS_INDICATORS',
'scan_checkpoints',

321
studio/backend/utils/models/model_config.py
View file

@ -422,6 +422,194 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
pass
def detect_gguf_model(path: str) -> Optional[str]:
"""
Check if the given local path is or contains a GGUF model file.
Handles two cases:
1. path is a direct .gguf file path
2. path is a directory containing .gguf files
Returns the full path to the .gguf file if found, None otherwise.
For HuggingFace repo detection, use detect_gguf_model_remote() instead.
"""
p = Path(path)
# Case 1: direct .gguf file
if p.suffix == ".gguf" and p.is_file():
return str(p.resolve())
# Case 2: directory containing .gguf files
if p.is_dir():
gguf_files = sorted(p.glob("*.gguf"), key=lambda f: f.stat().st_size, reverse=True)
if gguf_files:
return str(gguf_files[0].resolve())
return None
# Preferred GGUF quantization levels, in descending priority.
# Q4_K_M is a good default: small, fast, acceptable quality.
_GGUF_QUANT_PREFERENCE = [
"Q4_K_M", "Q4_K_S", "Q5_K_M", "Q5_K_S",
"Q6_K", "Q8_0", "Q3_K_M", "Q3_K_L", "Q2_K",
"F16", "BF16", "F32",
]
def _pick_best_gguf(filenames: list[str]) -> Optional[str]:
"""
Pick the best GGUF file from a list of filenames.
Prefers quantization levels in _GGUF_QUANT_PREFERENCE order.
Falls back to the first .gguf file found.
"""
gguf_files = [f for f in filenames if f.endswith(".gguf")]
if not gguf_files:
return None
# Try preferred quantization levels
for quant in _GGUF_QUANT_PREFERENCE:
for f in gguf_files:
if quant in f:
return f
# Fallback: first GGUF file
return gguf_files[0]
@dataclass
class GgufVariantInfo:
"""A single GGUF quantization variant from a HuggingFace repo."""
filename: str # e.g., "gemma-3-4b-it-Q4_K_M.gguf"
quant: str # e.g., "Q4_K_M" (extracted from filename)
size_bytes: int # file size
def _extract_quant_label(filename: str) -> str:
"""
Extract quantization label like Q4_K_M, IQ4_XS, BF16 from a GGUF filename.
Examples:
"gemma-3-4b-it-Q4_K_M.gguf" → "Q4_K_M"
"model-IQ4_NL.gguf" → "IQ4_NL"
"model-BF16.gguf" → "BF16"
"model-UD-IQ1_S.gguf" → "UD-IQ1_S"
"model-UD-TQ1_0.gguf" → "UD-TQ1_0"
"MXFP4_MOE/model-MXFP4_MOE-0001.gguf"→ "MXFP4_MOE"
"""
import re
# Use only the basename (rfilename may include directory)
basename = filename.rsplit("/", 1)[-1]
# Strip .gguf and any shard suffix (-00001-of-00010)
stem = re.sub(r'-\d{3,}-of-\d{3,}', '', basename.rsplit(".", 1)[0])
# Match known quantization patterns
match = re.search(
r'(UD-)?' # Optional UD- prefix (Ultra Discrete)
r'(MXFP[0-9]+(?:_[A-Z0-9]+)*' # MXFP variants: MXFP4, MXFP4_MOE
r'|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?' # IQ variants: IQ4_XS, IQ4_NL, IQ1_S
r'|TQ[0-9]+_[0-9]+' # Ternary quant: TQ1_0, TQ2_0
r'|Q[0-9]+_K_[A-Z]+' # K-quant: Q4_K_M, Q3_K_S
r'|Q[0-9]+_[0-9]+' # Standard: Q8_0, Q5_1
r'|Q[0-9]+_K' # Short K-quant: Q6_K
r'|BF16|F16|F32)', # Full precision
stem, re.IGNORECASE,
)
if match:
prefix = match.group(1) or ""
return f"{prefix}{match.group(2)}"
# Fallback: last segment after hyphen
return stem.split("-")[-1]
def list_gguf_variants(
repo_id: str,
hf_token: Optional[str] = None,
) -> tuple[list[GgufVariantInfo], bool]:
"""
List all GGUF quantization variants in a HuggingFace repo.
Separates main model files from mmproj (vision projection) files.
The presence of mmproj files indicates a vision-capable model.
Returns:
(variants, has_vision): list of non-mmproj GGUF variants + vision flag.
"""
from huggingface_hub import model_info as hf_model_info
info = hf_model_info(repo_id, token=hf_token, files_metadata=True)
variants: list[GgufVariantInfo] = []
has_vision = False
quant_totals: dict[str, int] = {} # quant -> total bytes
quant_first_file: dict[str, str] = {} # quant -> first filename (for display)
for sibling in info.siblings:
fname = sibling.rfilename
if not fname.endswith(".gguf"):
continue
size = sibling.size or 0
# mmproj files are vision projection models, not main model files
if "mmproj" in fname.lower():
has_vision = True
continue
quant = _extract_quant_label(fname)
quant_totals[quant] = quant_totals.get(quant, 0) + size
if quant not in quant_first_file:
quant_first_file[quant] = fname
for quant, total_size in quant_totals.items():
variants.append(GgufVariantInfo(
filename=quant_first_file[quant],
quant=quant,
size_bytes=total_size,
))
return variants, has_vision
def detect_gguf_model_remote(
repo_id: str,
hf_token: Optional[str] = None,
) -> Optional[str]:
"""
Check if a HuggingFace repo contains GGUF files.
Returns the filename of the best GGUF file in the repo, or None.
"""
try:
from huggingface_hub import model_info as hf_model_info
info = hf_model_info(repo_id, token=hf_token)
repo_files = [s.rfilename for s in info.siblings]
return _pick_best_gguf(repo_files)
except Exception as e:
logger.debug(f"Could not check GGUF files for '{repo_id}': {e}")
return None
def download_gguf_file(
repo_id: str,
filename: str,
hf_token: Optional[str] = None,
) -> str:
"""
Download a specific GGUF file from a HuggingFace repo.
Returns the local path to the downloaded file.
"""
from huggingface_hub import hf_hub_download
local_path = hf_hub_download(
repo_id=repo_id,
filename=filename,
token=hf_token,
)
return local_path
def scan_trained_loras(outputs_dir: str = "./outputs") -> List[Tuple[str, str]]:
"""
Scan outputs folder for trained LoRA adapters.
@ -467,14 +655,15 @@ def scan_trained_loras(outputs_dir: str = "./outputs") -> List[Tuple[str, str]]:
def scan_exported_models(exports_dir: str = "./exports") -> List[Tuple[str, str, str, Optional[str]]]:
"""
Scan exports folder for exported models (merged, LoRA, base).
Skips GGUF-only exports (not loadable by Unsloth inference backend).
Scan exports folder for exported models (merged, LoRA, GGUF).
The exports directory is two levels deep: {run}/{checkpoint}/
Supports two directory layouts:
- Two-level: {run}/{checkpoint}/ (merged & LoRA exports)
- Flat: {name}-finetune-gguf/ (GGUF exports)
Returns:
List of tuples: [(display_name, model_path, export_type, base_model), ...]
export_type: "lora" | "merged"
export_type: "lora" | "merged" | "gguf"
"""
results = []
exports_path = Path(exports_dir)
@ -486,6 +675,26 @@ def scan_exported_models(exports_dir: str = "./exports") -> List[Tuple[str, str,
for run_dir in exports_path.iterdir():
if not run_dir.is_dir():
continue
# Check for flat GGUF export (e.g. exports/gemma-3-4b-it-finetune-gguf/)
gguf_files = list(run_dir.glob("*.gguf"))
if gguf_files:
base_model = None
export_meta = run_dir / "export_metadata.json"
try:
if export_meta.exists():
meta = json.loads(export_meta.read_text())
base_model = meta.get("base_model")
except Exception:
pass
display_name = run_dir.name
model_path = str(gguf_files[0]) # path to the .gguf file
results.append((display_name, model_path, "gguf", base_model))
logger.debug(f"Found GGUF export: {display_name}")
continue
# Two-level: {run}/{checkpoint}/
for checkpoint_dir in run_dir.iterdir():
if not checkpoint_dir.is_dir():
continue
@ -510,7 +719,6 @@ def scan_exported_models(exports_dir: str = "./exports") -> List[Tuple[str, str,
pass
elif config_file.exists() and has_weights:
export_type = "merged"
# Read base model from export_metadata.json (written at export time)
export_meta = checkpoint_dir / "export_metadata.json"
try:
if export_meta.exists():
@ -519,7 +727,25 @@ def scan_exported_models(exports_dir: str = "./exports") -> List[Tuple[str, str,
except Exception:
pass
elif has_gguf:
# GGUF-only — not loadable by current inference backend
export_type = "gguf"
gguf_list = list(checkpoint_dir.glob("*.gguf"))
# Check checkpoint_dir first, then fall back to parent run_dir
# (export.py writes metadata to the top-level export directory)
for meta_dir in (checkpoint_dir, run_dir):
export_meta = meta_dir / "export_metadata.json"
try:
if export_meta.exists():
meta = json.loads(export_meta.read_text())
base_model = meta.get("base_model")
if base_model:
break
except Exception:
pass
display_name = f"{run_dir.name} / {checkpoint_dir.name}"
model_path = str(gguf_list[0]) if gguf_list else str(checkpoint_dir)
results.append((display_name, model_path, export_type, base_model))
logger.debug(f"Found GGUF export: {display_name}")
continue
else:
continue
@ -679,6 +905,10 @@ class ModelConfig:
is_cached: bool # Is this already in HF cache?
is_vision: bool # Is this a vision model?
is_lora: bool # Is this a lora adapter?
is_gguf: bool = False # Is this a GGUF model?
gguf_file: Optional[str] = None # Full path to the .gguf file (local mode)
gguf_hf_repo: Optional[str] = None # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF")
gguf_variant: Optional[str] = None # Quantization variant (e.g. "Q4_K_M")
base_model: Optional[str] = None # Base model (for LoRAs)
@classmethod
@ -734,37 +964,102 @@ class ModelConfig:
cls,
model_id: str,
hf_token: Optional[str] = None,
is_lora: bool = False
is_lora: bool = False,
gguf_variant: Optional[str] = None,
) -> Optional['ModelConfig']:
"""
Create ModelConfig from a clean model identifier.
For FastAPI routes where the frontend sends sanitized model paths.
No Gradio dropdown parsing - expects clean identifiers like:
- "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"
- "./outputs/my_lora_adapter"
- "/absolute/path/to/model"
Args:
model_id: Clean model identifier (HF repo name or local path)
hf_token: Optional HF token for vision detection on gated models
is_lora: Whether this is a LoRA adapter
gguf_variant: Optional GGUF quantization variant (e.g. "Q4_K_M").
For remote GGUF repos, specifies which quant to load via -hf.
If None, auto-selects using _pick_best_gguf().
Returns:
ModelConfig or None if configuration cannot be created
"""
if not model_id or not model_id.strip():
return None
identifier = model_id.strip()
is_local = is_local_path(identifier)
path = normalize_path(identifier) if is_local else identifier
# Add unsloth/ prefix for shorthand HF models
if not is_local and "/" not in identifier:
identifier = f"unsloth/{identifier}"
path = identifier
# Auto-detect GGUF models (check before LoRA/vision detection)
if is_local:
gguf_file = detect_gguf_model(path)
if gguf_file:
display_name = Path(gguf_file).stem
logger.info(f"Detected local GGUF model: {gguf_file}")
return cls(
identifier=identifier,
display_name=display_name,
path=path,
is_local=True,
is_cached=True,
is_vision=False,
is_lora=False,
is_gguf=True,
gguf_file=gguf_file,
)
else:
# Check if the HF repo contains GGUF files
gguf_filename = detect_gguf_model_remote(identifier, hf_token=hf_token)
if gguf_filename:
# Preflight: verify llama-server binary exists BEFORE user waits
# for a multi-GB download that llama-server handles natively
from core.inference.llama_cpp import LlamaCppBackend
if not LlamaCppBackend._find_llama_server_binary():
raise RuntimeError(
"llama-server binary not found — cannot load GGUF models. "
"Run setup.sh to build it, or set LLAMA_SERVER_PATH."
)
# Use list_gguf_variants() to detect vision & resolve variant
variants, has_vision = list_gguf_variants(identifier, hf_token=hf_token)
variant = gguf_variant
if not variant:
# Auto-select best quantization
variant_filenames = [v.filename for v in variants]
best = _pick_best_gguf(variant_filenames)
if best:
variant = _extract_quant_label(best)
else:
variant = "Q4_K_M" # Fallback — llama-server's own default
display_name = f"{identifier.split('/')[-1]} ({variant})"
logger.info(
f"Detected remote GGUF repo '{identifier}', "
f"variant={variant}, vision={has_vision}"
)
return cls(
identifier=identifier,
display_name=display_name,
path=identifier,
is_local=False,
is_cached=False,
is_vision=has_vision,
is_lora=False,
is_gguf=True,
gguf_file=None,
gguf_hf_repo=identifier,
gguf_variant=variant,
)
# Auto-detect LoRA for local paths (check adapter_config.json on disk)
if not is_lora and is_local:
detected_base = get_base_model_from_lora(path)

6
studio/frontend/src/components/assistant-ui/attachment.tsx
View file

@ -149,10 +149,8 @@ const AttachmentUI: FC = () => {
return "Document";
case "file":
return "File";
default: {
const _exhaustiveCheck: never = type;
throw new Error(`Unknown attachment type: ${_exhaustiveCheck}`);
}
default:
throw new Error(`Unknown attachment type: ${type as string}`);
}
});

14
studio/frontend/src/components/assistant-ui/model-selector.tsx
View file

@ -27,6 +27,7 @@ interface ModelSelectorProps {
loraModels?: LoraModelOption[];
value?: string;
defaultValue?: string;
activeGgufVariant?: string | null;
onValueChange?: (value: string, meta: ModelSelectorChangeMeta) => void;
onEject?: () => void;
variant?: "outline" | "ghost" | "muted";
@ -158,6 +159,7 @@ export function ModelSelector({
loraModels = [],
value,
defaultValue,
activeGgufVariant,
onValueChange,
onEject,
variant = "outline",
@ -202,9 +204,15 @@ export function ModelSelector({
return all;
}, [loraModels, models]);
const currentModel = selected
? optionById.get(selected) ?? { id: selected, name: selected }
: undefined;
const currentModel = useMemo(() => {
if (!selected) return undefined;
const found = optionById.get(selected);
if (activeGgufVariant) {
const desc = `GGUF · ${activeGgufVariant}`;
return found ? { ...found, description: desc } : { id: selected, name: selected, description: desc };
}
return found ?? { id: selected, name: selected };
}, [selected, optionById, activeGgufVariant]);
function handleSelect(id: string, meta: ModelSelectorChangeMeta) {
if (onValueChange) {

238
studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
View file

@ -5,6 +5,8 @@ import {
TooltipContent,
TooltipTrigger,
} from "@/components/ui/tooltip";
import { listGgufVariants } from "@/features/chat/api/chat-api";
import type { GgufVariantDetail } from "@/features/chat/types/api";
import {
useDebouncedValue,
useGpuInfo,
@ -17,7 +19,7 @@ import type { VramFitStatus } from "@/lib/vram";
import { checkVramFit, estimateLoadingVram } from "@/lib/vram";
import { Search01Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useMemo, useState, type ReactNode } from "react";
import { useCallback, useEffect, useMemo, useState, type ReactNode } from "react";
import type {
LoraModelOption,
ModelOption,
@ -36,6 +38,15 @@ function ListLabel({ children }: { children: ReactNode }) {
);
}
/** Format bytes to a human-readable size string. */
function formatBytes(bytes: number): string {
if (bytes === 0) return "0 B";
const units = ["B", "KB", "MB", "GB", "TB"];
const i = Math.floor(Math.log(bytes) / Math.log(1024));
const value = bytes / 1024 ** i;
return `${value.toFixed(value < 10 ? 1 : 0)} ${units[i]}`;
}
function ModelRow({
label,
meta,
@ -114,6 +125,143 @@ function ModelRow({
return content;
}
// ── GGUF Variant Expander ────────────────────────────────────
function GgufVariantExpander({
repoId,
onSelect,
gpuGb,
}: {
repoId: string;
onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
gpuGb?: number;
}) {
const [variants, setVariants] = useState<GgufVariantDetail[] | null>(null);
const [defaultVariant, setDefaultVariant] = useState<string | null>(null);
const [hasVision, setHasVision] = useState(false);
const [loading, setLoading] = useState(true);
const [error, setError] = useState<string | null>(null);
useEffect(() => {
let canceled = false;
setLoading(true);
setError(null);
listGgufVariants(repoId)
.then((res) => {
if (canceled) return;
setVariants(res.variants);
setDefaultVariant(res.default_variant);
setHasVision(res.has_vision);
})
.catch((err) => {
if (canceled) return;
setError(err instanceof Error ? err.message : "Failed to load variants");
})
.finally(() => {
if (!canceled) setLoading(false);
});
return () => {
canceled = true;
};
}, [repoId]);
const handleVariantClick = useCallback(
(quant: string) => {
onSelect(repoId, {
source: "hub",
isLora: false,
ggufVariant: quant,
});
},
[repoId, onSelect],
);
if (loading) {
return (
<div className="flex items-center gap-2 px-5 py-2">
<Spinner className="size-3 text-muted-foreground" />
<span className="text-xs text-muted-foreground">Loading variants…</span>
</div>
);
}
if (error) {
return (
<div className="px-5 py-2 text-xs text-destructive">{error}</div>
);
}
if (!variants || variants.length === 0) {
return (
<div className="px-5 py-2 text-xs text-muted-foreground">
No GGUF variants found.
</div>
);
}
return (
<div className="pl-4 border-l-2 border-accent/50 ml-3 my-1">
<div className="px-2 py-1 flex items-center gap-1.5">
<span className="text-[10px] font-semibold uppercase tracking-wider text-muted-foreground">
Quantizations
</span>
{hasVision && (
<span className="text-[9px] font-medium text-blue-400">Vision</span>
)}
</div>
{variants.map((v) => {
const sizeGb = v.size_bytes / (1024 ** 3);
const fitStatus = gpuGb != null && gpuGb > 0 && sizeGb > 0
? checkVramFit(sizeGb, gpuGb)
: null;
return (
<button
key={v.filename}
type="button"
onClick={() => handleVariantClick(v.quant)}
className={cn(
"flex w-full items-center justify-between gap-2 rounded-md px-2.5 py-1 text-left text-sm transition-colors hover:bg-accent",
)}
>
<span className="min-w-0 flex-1 truncate font-mono text-xs">
{v.quant}
{v.quant === defaultVariant && (
<span className="ml-1.5 text-[9px] font-sans font-medium text-primary/70">
recommended
</span>
)}
</span>
<span className="flex items-center gap-1.5 shrink-0">
{fitStatus === "exceeds" && (
<span className="text-[9px] font-medium text-red-400">OOM</span>
)}
{fitStatus === "tight" && (
<span className="text-[9px] font-medium text-amber-400">TIGHT</span>
)}
{fitStatus === "fits" && (
<span className="text-[9px] font-medium text-emerald-500/90">FIT</span>
)}
<span className="text-[10px] text-muted-foreground">
{formatBytes(v.size_bytes)}
</span>
</span>
</button>
);
})}
</div>
);
}
// ── Detect GGUF repos by naming convention ────────────────────
function isGgufRepo(id: string): boolean {
return id.toUpperCase().includes("-GGUF");
}
// ── Hub Model Picker ──────────────────────────────────────────
export function HubModelPicker({
models,
value,
@ -130,6 +278,9 @@ export function HubModelPicker({
debouncedQuery,
);
// Track which GGUF repo is expanded for variant selection
const [expandedGguf, setExpandedGguf] = useState<string | null>(null);
const recommendedIds = useMemo(
() => dedupe([...models.map((model) => model.id), value ?? ""]),
[models, value],
@ -199,6 +350,19 @@ export function HubModelPicker({
const { scrollRef, sentinelRef } = useInfiniteScroll(fetchMore, results.length);
/** Handle clicking a model row — GGUF repos expand, others load directly. */
const handleModelClick = useCallback(
(id: string) => {
if (isGgufRepo(id)) {
// Toggle GGUF variant expander
setExpandedGguf((prev) => (prev === id ? null : id));
} else {
onSelect(id, { source: "hub", isLora: false });
}
},
[onSelect],
);
return (
<div className="space-y-2">
<div className="relative">
@ -230,18 +394,24 @@ export function HubModelPicker({
recommendedIds.map((id) => {
const vram = recommendedVramMap.get(id);
return (
<ModelRow
key={id}
label={id}
meta={vram?.detail ?? undefined}
selected={value === id}
onClick={() =>
onSelect(id, { source: "hub", isLora: false })
}
vramStatus={vram?.status ?? null}
vramEst={vram?.est}
gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
/>
<div key={id}>
<ModelRow
label={id}
meta={
isGgufRepo(id)
? "GGUF"
: vram?.detail ?? undefined
}
selected={value === id}
onClick={() => handleModelClick(id)}
vramStatus={isGgufRepo(id) ? null : vram?.status ?? null}
vramEst={isGgufRepo(id) ? undefined : vram?.est}
gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
/>
{expandedGguf === id && (
<GgufVariantExpander repoId={id} onSelect={onSelect} gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} />
)}
</div>
);
})
)}
@ -259,18 +429,24 @@ export function HubModelPicker({
hfIds.map((id) => {
const vram = vramMap.get(id);
return (
<ModelRow
key={id}
label={id}
meta={metricsById.get(id)}
selected={value === id}
onClick={() =>
onSelect(id, { source: "hub", isLora: false })
}
vramStatus={vram?.status ?? null}
vramEst={vram?.est}
gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
/>
<div key={id}>
<ModelRow
label={id}
meta={
isGgufRepo(id)
? "GGUF"
: metricsById.get(id)
}
selected={value === id}
onClick={() => handleModelClick(id)}
vramStatus={isGgufRepo(id) ? null : vram?.status ?? null}
vramEst={isGgufRepo(id) ? undefined : vram?.est}
gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
/>
{expandedGguf === id && (
<GgufVariantExpander repoId={id} onSelect={onSelect} gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} />
)}
</div>
);
})
)}
@ -368,9 +544,12 @@ export function LoraModelPicker({
{adapters.map((adapter) => {
const isExported = adapter.source === "exported";
const isMerged = adapter.exportType === "merged";
const tag = isExported
? isMerged ? "Merged" : "LoRA"
: "LoRA";
const isGguf = adapter.exportType === "gguf";
const tag = isGguf
? "GGUF"
: isExported
? isMerged ? "Merged" : "LoRA"
: "LoRA";
const meta = isExported ? `${tag} · Exported` : tag;
return (
<ModelRow
@ -380,7 +559,7 @@ export function LoraModelPicker({
selected={value === adapter.id}
onClick={() => onSelect(adapter.id, {
source: isExported ? "exported" : "lora",
isLora: !isMerged,
isLora: !isMerged && !isGguf,
})}
/>
);
@ -393,4 +572,3 @@ export function LoraModelPicker({
</div>
);
}

3
studio/frontend/src/components/assistant-ui/model-selector/types.ts
View file

@ -11,11 +11,12 @@ export interface LoraModelOption extends ModelOption {
baseModel?: string;
updatedAt?: number;
source?: "training" | "exported";
exportType?: "lora" | "merged";
exportType?: "lora" | "merged" | "gguf";
}
export interface ModelSelectorChangeMeta {
source: "hub" | "lora" | "exported";
isLora: boolean;
ggufVariant?: string;
}

42
studio/frontend/src/components/ui/combobox.tsx
View file

@ -140,10 +140,10 @@ function ComboboxInput({
);
}
function ComboboxContent({
className,
side = "bottom",
sideOffset = 6,
function ComboboxContent({
className,
side = "bottom",
sideOffset = 6,
align = "start",
alignOffset = 0,
anchor,
@ -159,23 +159,23 @@ function ComboboxContent({
const dialogContainer = useDialogPortalContainer();
return (
<ComboboxPrimitive.Portal container={container ?? dialogContainer ?? undefined}>
<ComboboxPrimitive.Positioner
side={side}
sideOffset={sideOffset}
align={align}
alignOffset={alignOffset}
anchor={anchor}
className="isolate z-50"
>
<ComboboxPrimitive.Popup
data-slot="combobox-content"
data-chips={!!anchor}
className={cn(
"bg-popover text-popover-foreground data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 shadow-border ring-1 ring-border *:data-[slot=input-group]:bg-input/30 max-h-72 min-w-36 overflow-hidden rounded-xl corner-squircle duration-100 *:data-[slot=input-group]:m-1 *:data-[slot=input-group]:mb-0 *:data-[slot=input-group]:h-9 *:data-[slot=input-group]:border-none *:data-[slot=input-group]:shadow-none group/combobox-content relative max-h-(--available-height) w-(--anchor-width) max-w-(--available-width) min-w-[calc(var(--anchor-width)+--spacing(7))] origin-(--transform-origin) data-[chips=true]:min-w-(--anchor-width)",
className,
)}
{...props}
/>
<ComboboxPrimitive.Positioner
side={side}
sideOffset={sideOffset}
align={align}
alignOffset={alignOffset}
anchor={anchor}
className="isolate z-[120] pointer-events-auto"
>
<ComboboxPrimitive.Popup
data-slot="combobox-content"
data-chips={!!anchor}
className={cn(
"bg-popover text-popover-foreground data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 shadow-border ring-1 ring-border *:data-[slot=input-group]:bg-input/30 max-h-72 min-w-36 overflow-hidden rounded-xl corner-squircle duration-100 *:data-[slot=input-group]:m-1 *:data-[slot=input-group]:mb-0 *:data-[slot=input-group]:h-9 *:data-[slot=input-group]:border-none *:data-[slot=input-group]:shadow-none group/combobox-content relative pointer-events-auto max-h-(--available-height) w-(--anchor-width) max-w-(--available-width) min-w-[calc(var(--anchor-width)+--spacing(7))] origin-(--transform-origin) data-[chips=true]:min-w-(--anchor-width)",
className,
)}
{...props}
/>
</ComboboxPrimitive.Positioner>
</ComboboxPrimitive.Portal>
);

10
studio/frontend/src/config/training.ts
View file

@ -39,6 +39,11 @@ export const MODEL_TYPES: ReadonlyArray<{
label: string;
description: string;
}> = [
{
value: "text",
label: "Text",
description: "Language models",
},
{
value: "vision",
label: "Vision",
@ -54,11 +59,6 @@ export const MODEL_TYPES: ReadonlyArray<{
label: "Embeddings",
description: "Text embedding models",
},
{
value: "text",
label: "Text",
description: "Language models",
},
];
export const CONTEXT_LENGTHS = [512, 1024, 2048, 4096, 8192, 16384, 32768];

11
studio/frontend/src/features/chat/api/chat-api.ts
View file

@ -1,5 +1,6 @@
import { authFetch } from "@/features/auth";
import type {
GgufVariantsResponse,
InferenceStatusResponse,
ListLorasResponse,
ListModelsResponse,
@ -74,6 +75,16 @@ export async function unloadModel(payload: UnloadModelRequest): Promise<void> {
await parseJsonOrThrow<unknown>(response);
}
export async function listGgufVariants(
repoId: string,
hfToken?: string,
): Promise<GgufVariantsResponse> {
const params = new URLSearchParams({ repo_id: repoId });
if (hfToken) params.set("hf_token", hfToken);
const response = await authFetch(`/api/models/gguf-variants?${params}`);
return parseJsonOrThrow<GgufVariantsResponse>(response);
}
function parseSseEvent(rawEvent: string): string[] {
const dataLines: string[] = [];
for (const line of rawEvent.split(/\r?\n/)) {

13
studio/frontend/src/features/chat/chat-page.tsx
View file

@ -314,6 +314,7 @@ export function ChatPage(): ReactElement {
);
const inferenceParams = useChatRuntimeStore((state) => state.params);
const setInferenceParams = useChatRuntimeStore((state) => state.setParams);
const activeGgufVariant = useChatRuntimeStore((state) => state.activeGgufVariant);
const autoTitle = useChatRuntimeStore((state) => state.autoTitle);
const setAutoTitle = useChatRuntimeStore((state) => state.setAutoTitle);
const modelsFromStore = useChatRuntimeStore((state) => state.models);
@ -335,10 +336,11 @@ export function ChatPage(): ReactElement {
}, [inferenceParams.checkpoint, lorasFromStore]);
const handleCheckpointChange = useCallback(
(value: string, meta?: { isLora: boolean }) => {
const currentCheckpoint =
useChatRuntimeStore.getState().params.checkpoint;
if (!value || value === currentCheckpoint) return;
(value: string, meta?: { isLora: boolean; ggufVariant?: string }) => {
const store = useChatRuntimeStore.getState();
const currentCheckpoint = store.params.checkpoint;
const currentVariant = store.activeGgufVariant;
if (!value || (value === currentCheckpoint && (meta?.ggufVariant ?? null) === (currentVariant ?? null))) return;
void (async () => {
let switchNote: string | undefined;
const activeThreadId = await resolveActiveSingleThreadId(view);
@ -367,10 +369,10 @@ export function ChatPage(): ReactElement {
duration: 6000,
});
}
await selectModel({
id: value,
isLora: meta?.isLora,
ggufVariant: meta?.ggufVariant,
});
})();
},
@ -591,6 +593,7 @@ export function ChatPage(): ReactElement {
models={models}
loraModels={loraModels}
value={inferenceParams.checkpoint}
activeGgufVariant={activeGgufVariant}
onValueChange={handleCheckpointChange}
onEject={handleEject}
variant="ghost"

6
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -165,7 +165,11 @@ export function ChatSettingsPanel({
function applyPreset(name: string) {
const p = presets.find((pr) => pr.name === name);
if (p) {
onParamsChange({ ...p.params, systemPrompt: params.systemPrompt });
onParamsChange({
...p.params,
systemPrompt: params.systemPrompt,
checkpoint: params.checkpoint,
});
setActivePreset(name);
}
}

17
studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
View file

@ -20,6 +20,7 @@ const DEFAULT_MODEL_MAX_SEQ_LENGTH = 2048;
type SelectedModelInput = {
id: string;
isLora?: boolean;
ggufVariant?: string;
loadingDescription?: string;
};
@ -42,11 +43,13 @@ function stripTrailingEpoch(input: string): string {
function describeModel(model: {
is_lora?: boolean;
is_vision?: boolean;
is_gguf?: boolean;
}): string | undefined {
const tags: string[] = [];
if (model.is_gguf) tags.push("GGUF");
if (model.is_lora) tags.push("LoRA");
if (model.is_vision) tags.push("Vision");
if (!model.is_lora && !model.is_vision) tags.push("Base");
if (!model.is_lora && !model.is_vision && !model.is_gguf) tags.push("Base");
return tags.join(" · ");
}
@ -55,6 +58,7 @@ function toChatModelSummary(model: {
name?: string | null;
is_lora?: boolean;
is_vision?: boolean;
is_gguf?: boolean;
}): ChatModelSummary {
return {
id: model.id,
@ -62,6 +66,7 @@ function toChatModelSummary(model: {
description: describeModel(model),
isLora: Boolean(model.is_lora),
isVision: Boolean(model.is_vision),
isGguf: Boolean(model.is_gguf),
};
}
@ -70,7 +75,7 @@ function toLoraSummary(lora: {
adapter_path: string;
base_model?: string | null;
source?: "training" | "exported" | null;
export_type?: "lora" | "merged" | null;
export_type?: "lora" | "merged" | "gguf" | null;
}): ChatLoraSummary {
const idTail = lora.adapter_path.split("/").filter(Boolean).at(-1) ?? "";
const updatedAt =
@ -139,7 +144,7 @@ export function useChatModelRuntime() {
setLoras(lorasRes.loras.map(toLoraSummary));
if (statusRes.active_model) {
setCheckpoint(statusRes.active_model);
setCheckpoint(statusRes.active_model, statusRes.gguf_variant);
}
} catch (error) {
const message =
@ -154,7 +159,10 @@ export function useChatModelRuntime() {
const selectModel = useCallback(
async (selection: string | SelectedModelInput) => {
const modelId = typeof selection === "string" ? selection : selection.id;
if (!modelId || params.checkpoint === modelId) {
const ggufVariant =
typeof selection === "string" ? undefined : selection.ggufVariant;
const currentVariant = useChatRuntimeStore.getState().activeGgufVariant;
if (!modelId || (params.checkpoint === modelId && (ggufVariant ?? null) === (currentVariant ?? null))) {
return;
}
@ -193,6 +201,7 @@ export function useChatModelRuntime() {
max_seq_length: DEFAULT_MODEL_MAX_SEQ_LENGTH,
load_in_4bit: true,
is_lora: isLora,
gguf_variant: ggufVariant ?? null,
});
const currentParams = useChatRuntimeStore.getState().params;

8
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -39,13 +39,14 @@ type ChatRuntimeStore = {
runningByThreadId: Record<string, boolean>;
autoTitle: boolean;
modelsError: string | null;
activeGgufVariant: string | null;
setParams: (params: InferenceParams) => void;
setModels: (models: ChatModelSummary[]) => void;
setLoras: (loras: ChatLoraSummary[]) => void;
setThreadRunning: (threadId: string, running: boolean) => void;
setAutoTitle: (enabled: boolean) => void;
setModelsError: (error: string | null) => void;
setCheckpoint: (modelId: string) => void;
setCheckpoint: (modelId: string, ggufVariant?: string | null) => void;
clearCheckpoint: () => void;
};
@ -56,6 +57,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
runningByThreadId: {},
autoTitle: loadBool(AUTO_TITLE_KEY, false),
modelsError: null,
activeGgufVariant: null,
setParams: (params) => set({ params }),
setModels: (models) => set({ models }),
setLoras: (loras) => set({ loras }),
@ -75,12 +77,13 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
return { autoTitle };
}),
setModelsError: (modelsError) => set({ modelsError }),
setCheckpoint: (modelId) =>
setCheckpoint: (modelId, ggufVariant) =>
set((state) => ({
params: {
...state.params,
checkpoint: modelId,
},
activeGgufVariant: ggufVariant ?? null,
})),
clearCheckpoint: () =>
set((state) => ({
@ -88,5 +91,6 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
...state.params,
checkpoint: "",
},
activeGgufVariant: null,
})),
}));

20
studio/frontend/src/features/chat/types/api.ts
View file

@ -3,6 +3,7 @@ export interface BackendModelDetails {
name?: string | null;
is_vision?: boolean;
is_lora?: boolean;
is_gguf?: boolean;
}
export interface ListModelsResponse {
@ -15,7 +16,7 @@ export interface BackendLoraInfo {
adapter_path: string;
base_model?: string | null;
source?: "training" | "exported" | null;
export_type?: "lora" | "merged" | null;
export_type?: "lora" | "merged" | "gguf" | null;
}
export interface ListLorasResponse {
@ -29,6 +30,20 @@ export interface LoadModelRequest {
max_seq_length: number;
load_in_4bit: boolean;
is_lora: boolean;
gguf_variant?: string | null;
}
export interface GgufVariantDetail {
filename: string;
quant: string;
size_bytes: number;
}
export interface GgufVariantsResponse {
repo_id: string;
variants: GgufVariantDetail[];
has_vision: boolean;
default_variant: string | null;
}
export interface LoadModelResponse {
@ -37,6 +52,7 @@ export interface LoadModelResponse {
display_name: string;
is_vision: boolean;
is_lora: boolean;
is_gguf?: boolean;
inference?: {
temperature?: number;
top_p?: number;
@ -52,6 +68,8 @@ export interface UnloadModelRequest {
export interface InferenceStatusResponse {
active_model: string | null;
is_vision: boolean;
is_gguf?: boolean;
gguf_variant?: string | null;
loading: string[];
loaded: string[];
}

3
studio/frontend/src/features/chat/types/runtime.ts
View file

@ -26,6 +26,7 @@ export interface ChatModelSummary {
description?: string;
isVision: boolean;
isLora: boolean;
isGguf?: boolean;
}
export interface ChatLoraSummary {
@ -34,5 +35,5 @@ export interface ChatLoraSummary {
baseModel: string;
updatedAt?: number;
source?: "training" | "exported";
exportType?: "lora" | "merged";
exportType?: "lora" | "merged" | "gguf";
}

7
studio/frontend/src/features/export/export-page.tsx
View file

@ -155,7 +155,12 @@ export function ExportPage() {
setExportError(null);
setExportSuccess(false);
const saveDir = `./exports/${selectedModelIdx ?? "model"}/${checkpoint}`;
// For GGUF, use a flat folder like "exports/gemma-3-4b-it-finetune-gguf"
// For other formats, nest under training-run/checkpoint
const saveDir =
exportMethod === "gguf"
? `./exports/${(baseModelName.split("/").pop() ?? selectedModelIdx ?? "model")}-finetune-gguf`
: `./exports/${selectedModelIdx ?? "model"}/${checkpoint}`;
const pushToHub = destination === "hub";
const repoId = pushToHub && hfUsername && modelName
? `${hfUsername}/${modelName}`

24
studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx
View file

@ -38,7 +38,7 @@ import {
useHfTokenValidation,
useInfiniteScroll,
} from "@/hooks";
import { cn, formatCompact } from "@/lib/utils";
import { cn } from "@/lib/utils";
import {
HfDatasetSubsetSplitSelectors,
useTrainingConfigStore,
@ -79,6 +79,7 @@ export function DatasetStep() {
setDatasetEvalSplit,
uploadedFile,
setUploadedFile,
modelType,
} = useTrainingConfigStore(
useShallow((s) => ({
hfToken: s.hfToken,
@ -97,6 +98,7 @@ export function DatasetStep() {
setDatasetEvalSplit: s.setDatasetEvalSplit,
uploadedFile: s.uploadedFile,
setUploadedFile: s.setUploadedFile,
modelType: s.modelType,
})),
);
@ -110,6 +112,7 @@ export function DatasetStep() {
fetchMore,
error: hfSearchError,
} = useHfDatasetSearch(debouncedQuery, {
modelType,
accessToken: hfToken || undefined,
});
@ -251,16 +254,8 @@ export function DatasetStep() {
>
<ComboboxList className="p-1 !max-h-none !overflow-visible">
{(id: string) => {
const r = hfResults.find((r) => r.id === id);
const detail = r?.totalExamples
? `${formatCompact(r.totalExamples)} rows`
: (r?.sizeCategory ?? null);
return (
<ComboboxItem
key={id}
value={id}
className="gap-2"
>
<ComboboxItem key={id} value={id} className="gap-2">
<Tooltip>
<TooltipTrigger asChild={true}>
<span className="block min-w-0 flex-1 truncate">
@ -274,15 +269,6 @@ export function DatasetStep() {
{id}
</TooltipContent>
</Tooltip>
{detail ? (
<span className="ml-auto text-[10px] text-muted-foreground shrink-0">
{detail}
</span>
) : r?.downloads != null ? (
<span className="ml-auto text-[10px] text-muted-foreground shrink-0">
↓{formatCompact(r.downloads)}
</span>
) : null}
</ComboboxItem>
);
}}

1
studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx
View file

@ -85,6 +85,7 @@ export function ModelSelectionStep() {
} = useHfModelSearch(debouncedQuery, {
task,
accessToken: hfToken || undefined,
excludeGguf: true,
});
const { error: tokenValidationError, isChecking: isCheckingToken } =

2
studio/frontend/src/features/recipe-studio/blocks/definitions.ts
View file

@ -125,7 +125,7 @@ const BLOCK_DEFINITIONS: BlockDefinition[] = [
{
kind: "seed",
type: "seed_local",
title: "Local file",
title: "Structured file",
description: "Upload CSV/JSON/JSONL and use rows as seed context.",
icon: DocumentCodeIcon,
dialogKey: "seed",

241
studio/frontend/src/features/recipe-studio/components/block-sheet.tsx
View file

@ -1,4 +1,5 @@
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import {
Sheet,
SheetContent,
@ -12,17 +13,24 @@ import {
CodeIcon,
Copy02Icon,
type Database02Icon,
DragDropVerticalIcon,
PlusSignIcon,
Tick02Icon,
Upload01Icon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { type ReactElement, useMemo, useState } from "react";
import {
type DragEvent as ReactDragEvent,
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,
type BlockType,
type SeedBlockType,
} from "../blocks/registry";
@ -62,6 +70,12 @@ type BlockSheetProps = {
onImport: () => void;
};
export const RECIPE_BLOCK_DND_MIME = "application/x-recipe-studio-block";
export type RecipeBlockDragPayload = {
kind: SheetKind;
type: BlockType;
};
function getSheetTitle(sheetView: SheetView): string {
if (sheetView === "root") {
return "Add a block";
@ -103,6 +117,15 @@ const ROOT_GROUPS: RootGroup[] = [
icon: CodeIcon,
},
];
const SEARCHABLE_KINDS: SheetKind[] = [
"sampler",
"seed",
"llm",
"expression",
"note",
];
const PROCESSOR_TITLE = "Schema Transform";
const PROCESSOR_DESCRIPTION = "Transform final dataset schema.";
function BlockSheetButton({
icon,
@ -110,22 +133,30 @@ function BlockSheetButton({
description,
onClick,
isActive = false,
draggable = false,
onDragStart,
trailing = "chevron",
}: {
icon: typeof Database02Icon;
title: string;
description: string;
onClick: () => void;
isActive?: boolean;
draggable?: boolean;
onDragStart?: (event: ReactDragEvent<HTMLButtonElement>) => void;
trailing?: "chevron" | "drag" | "none";
}): ReactElement {
return (
<button
type="button"
onClick={onClick}
draggable={draggable}
onDragStart={onDragStart}
className={`flex w-full items-center gap-3 border-l-2 bg-background px-3 py-3 text-left transition hover:bg-muted/35 ${
isActive
? "border-emerald-500"
: "border-transparent hover:border-border/60"
}`}
} ${draggable ? "cursor-grab active:cursor-grabbing" : ""}`}
>
<div className="flex size-9 items-center justify-center rounded-xl text-foreground/70">
<HugeiconsIcon icon={icon} className="size-5" />
@ -134,10 +165,18 @@ function BlockSheetButton({
<p className="text-sm font-semibold text-foreground">{title}</p>
<p className="text-[11px] text-muted-foreground">{description}</p>
</div>
<HugeiconsIcon
icon={ArrowRight01Icon}
className="size-3.5 text-muted-foreground"
/>
{trailing === "chevron" ? (
<HugeiconsIcon
icon={ArrowRight01Icon}
className="size-3.5 text-muted-foreground"
/>
) : trailing === "drag" ? (
<HugeiconsIcon
icon={DragDropVerticalIcon}
strokeWidth={3.5}
className="size-5 text-foreground"
/>
) : null}
</button>
);
}
@ -162,11 +201,17 @@ export function BlockSheet({
}: BlockSheetProps): ReactElement {
const sheetTitle = getSheetTitle(sheetView);
const [uncontrolledOpen, setUncontrolledOpen] = useState(false);
const [search, setSearch] = useState("");
const expressionBlocks = useMemo(() => getBlocksForKind("expression"), []);
const noteBlocks = useMemo(() => getBlocksForKind("note"), []);
const seedBlocks = useMemo(() => getBlocksForKind("seed"), []);
const isControlled = typeof open === "boolean";
const sheetOpen = isControlled ? (open as boolean) : uncontrolledOpen;
const normalizedSearch = search.trim().toLowerCase();
const hasSearch = normalizedSearch.length > 0;
const isProcessorView = sheetView === "processor";
const isRootView = sheetView === "root";
const isScopedBlockView = !isRootView && !isProcessorView;
const setSheetOpen = (nextOpen: boolean) => {
if (!isControlled) {
@ -174,6 +219,95 @@ export function BlockSheet({
}
onOpenChange?.(nextOpen);
};
const matchesSearch = (title: string, description: string) =>
title.toLowerCase().includes(normalizedSearch) ||
description.toLowerCase().includes(normalizedSearch);
const searchableBlocks = useMemo(
() => SEARCHABLE_KINDS.flatMap((kind) => getBlocksForKind(kind)),
[],
);
const rootSearchBlocks = useMemo(() => {
if (!hasSearch) {
return [];
}
return searchableBlocks.filter((item) =>
matchesSearch(item.title, item.description),
);
}, [hasSearch, searchableBlocks, normalizedSearch]);
const scopedBlocks = useMemo(() => {
if (!isScopedBlockView) {
return [];
}
const blocks = getBlocksForKind(VIEW_KIND[sheetView] ?? "sampler");
if (!hasSearch) {
return blocks;
}
return blocks.filter((item) => matchesSearch(item.title, item.description));
}, [hasSearch, isScopedBlockView, normalizedSearch, sheetView]);
const rootGroups = useMemo(() => {
if (!hasSearch) {
return ROOT_GROUPS;
}
return ROOT_GROUPS.filter((group) => {
if (matchesSearch(group.title, group.description)) {
return true;
}
if (group.kind === "processor") {
return matchesSearch(PROCESSOR_TITLE, PROCESSOR_DESCRIPTION);
}
return getBlocksForKind(group.kind).some((item) =>
matchesSearch(item.title, item.description),
);
});
}, [hasSearch, normalizedSearch]);
const showNoMatches =
(isRootView && hasSearch && rootSearchBlocks.length === 0) ||
(isScopedBlockView && scopedBlocks.length === 0) ||
(isProcessorView &&
hasSearch &&
!matchesSearch(PROCESSOR_TITLE, PROCESSOR_DESCRIPTION));
const buildDragStart =
(kind: SheetKind, type: BlockType) =>
(event: ReactDragEvent<HTMLButtonElement>) => {
const payload: RecipeBlockDragPayload = { kind, type };
const serialized = JSON.stringify(payload);
event.dataTransfer.setData(RECIPE_BLOCK_DND_MIME, serialized);
event.dataTransfer.setData("text/plain", serialized);
event.dataTransfer.effectAllowed = "copy";
};
const getTrailing = (_kind: SheetKind): "drag" => "drag";
const onBlockClick = (kind: SheetKind, type: BlockType) => {
setSheetOpen(false);
if (kind === "sampler") {
onAddSampler(type as SamplerType);
return;
}
if (kind === "seed") {
onAddSeed(type as SeedBlockType);
return;
}
if (kind === "llm") {
if (type === "model_provider") {
onAddModelProvider();
return;
}
if (type === "model_config") {
onAddModelConfig();
return;
}
onAddLlm(type as LlmType);
return;
}
if (kind === "expression") {
onAddExpression();
return;
}
onAddMarkdownNote();
};
return (
<div className="flex flex-col items-end gap-2">
@ -183,6 +317,7 @@ export function BlockSheet({
setSheetOpen(nextOpen);
if (nextOpen) {
onViewChange("root");
setSearch("");
}
}}
>
@ -206,7 +341,7 @@ export function BlockSheet({
className="absolute gap-0 p-0 shadow-none"
overlayClassName="bg-transparent pointer-events-none backdrop-blur-none supports-backdrop-filter:backdrop-blur-none"
>
<SheetHeader className="border-b border-border/60 px-6 py-5">
<SheetHeader className="px-6 py-5">
<div className="flex items-center gap-2">
{sheetView !== "root" && (
<Button
@ -220,17 +355,52 @@ export function BlockSheet({
)}
<SheetTitle>{sheetTitle}</SheetTitle>
</div>
<Input
value={search}
onChange={(event) => setSearch(event.target.value)}
placeholder="Search blocks..."
className="corner-squircle mt-3 h-9"
/>
</SheetHeader>
<div className=" py-4">
<div className="mt-4 flex flex-col gap-2">
{sheetView === "root" &&
ROOT_GROUPS.map((item, index) => (
{isRootView &&
hasSearch &&
rootSearchBlocks.map((item, index) => (
<BlockSheetButton
key={`${item.kind}:${item.type}`}
icon={item.icon}
title={item.title}
description={item.description}
isActive={index === 0}
draggable={true}
onDragStart={buildDragStart(item.kind, item.type)}
trailing={getTrailing(item.kind)}
onClick={() => onBlockClick(item.kind, item.type)}
/>
))}
{isRootView &&
!hasSearch &&
rootGroups.map((item, index) => (
<BlockSheetButton
key={item.kind}
icon={item.icon}
title={item.title}
description={item.description}
isActive={index === 0}
draggable={item.kind === "expression" || item.kind === "note"}
onDragStart={
item.kind === "expression" && expressionBlocks[0]
? buildDragStart("expression", expressionBlocks[0].type)
: item.kind === "note" && noteBlocks[0]
? buildDragStart("note", noteBlocks[0].type)
: undefined
}
trailing={
item.kind === "expression" || item.kind === "note"
? "drag"
: "chevron"
}
onClick={() => {
if (item.kind === "processor") {
setSheetOpen(false);
@ -256,18 +426,20 @@ export function BlockSheet({
}}
/>
))}
{sheetView === "processor" && (
<BlockSheetButton
icon={CodeIcon}
title="Schema Transform"
description="Transform final dataset schema."
isActive={true}
onClick={onOpenProcessors}
/>
{isProcessorView && (
(!hasSearch ||
matchesSearch(PROCESSOR_TITLE, PROCESSOR_DESCRIPTION)) && (
<BlockSheetButton
icon={CodeIcon}
title={PROCESSOR_TITLE}
description={PROCESSOR_DESCRIPTION}
isActive={true}
onClick={onOpenProcessors}
/>
)
)}
{sheetView !== "root" &&
sheetView !== "processor" &&
getBlocksForKind(VIEW_KIND[sheetView] ?? "sampler").map(
{isScopedBlockView &&
scopedBlocks.map(
(item, index) => (
<BlockSheetButton
key={item.type}
@ -275,29 +447,18 @@ export function BlockSheet({
title={item.title}
description={item.description}
isActive={index === 0}
onClick={() => {
setSheetOpen(false);
if (item.kind === "sampler") {
onAddSampler(item.type as SamplerType);
} else if (item.kind === "seed") {
onAddSeed(item.type as SeedBlockType);
} else if (item.kind === "llm") {
if (item.type === "model_provider") {
onAddModelProvider();
} else if (item.type === "model_config") {
onAddModelConfig();
} else {
onAddLlm(item.type as LlmType);
}
} else if (item.kind === "expression") {
onAddExpression();
} else {
onAddMarkdownNote();
}
}}
draggable={true}
onDragStart={buildDragStart(item.kind, item.type)}
trailing={getTrailing(item.kind)}
onClick={() => onBlockClick(item.kind, item.type)}
/>
),
)}
{showNoMatches && (
<p className="px-3 py-2 text-xs text-muted-foreground">
No blocks match.
</p>
)}
</div>
</div>
</SheetContent>

31
studio/frontend/src/features/recipe-studio/components/chip-input.tsx
View file

@ -1,7 +1,15 @@
import { Button } from "@/components/ui/button";
import { Cancel01Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { type KeyboardEvent, type ReactElement, useId, useMemo, useState } from "react";
import {
type KeyboardEvent,
type ReactElement,
useEffect,
useId,
useMemo,
useRef,
useState,
} from "react";
type ChipInputProps = {
values: string[];
@ -19,12 +27,28 @@ export function ChipInput({
suggestions,
}: ChipInputProps): ReactElement {
const [draft, setDraft] = useState("");
const [isWrapped, setIsWrapped] = useState(false);
const containerRef = useRef<HTMLDivElement | null>(null);
const listId = useId();
const suggestionSet = useMemo(
() => new Set((suggestions ?? []).map((value) => value.trim())),
[suggestions],
);
useEffect(() => {
const element = containerRef.current;
if (!element) {
return;
}
const syncWrapped = () => {
setIsWrapped(element.clientHeight > 44);
};
syncWrapped();
const observer = new ResizeObserver(syncWrapped);
observer.observe(element);
return () => observer.disconnect();
}, [values.length, draft]);
function addValue(rawValue: string, allowAny: boolean): void {
const trimmed = rawValue.trim();
if (!trimmed) {
@ -55,7 +79,10 @@ export function ChipInput({
}
return (
<div className="bg-input/30 border-input focus-within:border-ring focus-within:ring-ring/50 flex min-h-9 flex-wrap items-center gap-1.5 rounded-4xl border bg-clip-padding px-1.5 py-1.5 text-sm transition-colors focus-within:ring-[3px]">
<div
ref={containerRef}
className={`bg-input/30 border-input focus-within:border-ring focus-within:ring-ring/50 flex min-h-9 flex-wrap items-center gap-1.5 border bg-clip-padding px-1.5 py-1.5 text-sm transition-colors focus-within:ring-[3px] ${isWrapped ? "corner-squircle rounded-xl" : "rounded-4xl"}`}
>
{values.map((value, index) => (
<span
key={`${value}-${index}`}

18
studio/frontend/src/features/recipe-studio/components/controls/layout-controls.tsx
View file

@ -5,6 +5,7 @@ import {
useUpdateNodeInternals,
} from "@xyflow/react";
import { Button } from "@/components/ui/button";
import { getFitNodeIdsIgnoringNotes } from "../../utils/graph/fit-view";
type LayoutControlsProps = {
direction: "LR" | "TB";
@ -32,20 +33,29 @@ export function LayoutControls({
requestAnimationFrame(() => {
refreshNodeInternals();
requestAnimationFrame(() => {
fitView({ duration: 250 });
fitView({
duration: 250,
nodes: getFitNodeIdsIgnoringNotes(getNodes()),
});
});
});
}, [fitView, onLayout, refreshNodeInternals]);
}, [fitView, getNodes, onLayout, refreshNodeInternals]);
const handleToggleDirection = useCallback(() => {
onToggleDirection();
requestAnimationFrame(() => {
refreshNodeInternals();
onLayout();
requestAnimationFrame(() => {
refreshNodeInternals();
requestAnimationFrame(() => {
fitView({
duration: 250,
nodes: getFitNodeIdsIgnoringNotes(getNodes()),
});
});
});
});
}, [onToggleDirection, refreshNodeInternals]);
}, [fitView, getNodes, onLayout, onToggleDirection, refreshNodeInternals]);
return (
<Panel position="top-left" className="m-3 flex items-center gap-2">

10
studio/frontend/src/features/recipe-studio/components/controls/viewport-controls.tsx
View file

@ -2,6 +2,7 @@ import { type ReactElement, useCallback } from "react";
import { Lock, LockOpen, Maximize2, Minus, Plus } from "lucide-react";
import { Panel, useReactFlow } from "@xyflow/react";
import { Button } from "@/components/ui/button";
import { getFitNodeIdsIgnoringNotes } from "../../utils/graph/fit-view";
import { RECIPE_FLOATING_ICON_BUTTON_CLASS } from "../recipe-floating-icon-button-class";
type ViewportControlsProps = {
@ -13,7 +14,7 @@ export function ViewportControls({
interactive,
onToggleInteractive,
}: ViewportControlsProps): ReactElement {
const { zoomIn, zoomOut, fitView } = useReactFlow();
const { zoomIn, zoomOut, fitView, getNodes } = useReactFlow();
const handleZoomIn = useCallback(() => {
zoomIn({ duration: 150 });
@ -24,8 +25,11 @@ export function ViewportControls({
}, [zoomOut]);
const handleFitView = useCallback(() => {
fitView({ duration: 250 });
}, [fitView]);
fitView({
duration: 250,
nodes: getFitNodeIdsIgnoringNotes(getNodes()),
});
}, [fitView, getNodes]);
return (
<Panel position="bottom-left" className="m-3 flex items-center gap-2">

29
studio/frontend/src/features/recipe-studio/components/inline/inline-expression.tsx
View file

@ -1,4 +1,3 @@
import { Badge } from "@/components/ui/badge";
import { Input } from "@/components/ui/input";
import {
Select,
@ -10,7 +9,9 @@ import {
import type { ReactElement } from "react";
import { useRecipeStudioStore } from "../../stores/recipe-studio";
import type { ExpressionConfig, ExpressionDtype } from "../../types";
import { findInvalidJinjaReferences } from "../../utils/refs";
import { getAvailableVariableEntries } from "../../utils/variables";
import { AvailableReferencesInline } from "../shared/available-references-inline";
import { InlineField } from "./inline-field";
type InlineExpressionProps = {
@ -26,6 +27,10 @@ export function InlineExpression({
}: InlineExpressionProps): ReactElement {
const configs = useRecipeStudioStore((state) => state.configs);
const vars = getAvailableVariableEntries(configs, config.id);
const invalidRefs = findInvalidJinjaReferences(
config.expr,
vars.map((entry) => entry.name),
);
return (
<div className="space-y-3">
@ -52,32 +57,14 @@ export function InlineExpression({
<InlineField label="Expression">
<Input
className="nodrag h-8 w-full text-xs"
aria-invalid={invalidRefs.length > 0}
placeholder="{{ column_name }}"
value={config.expr}
onChange={(event) => onUpdate({ expr: event.target.value })}
/>
</InlineField>
</div>
{vars.length > 0 && (
<div className="space-y-1">
<p className="text-[10px] font-medium text-muted-foreground">Available references</p>
<div className="flex flex-wrap gap-1">
{vars.map((v) => (
<Badge
key={`${v.source}:${v.name}`}
variant="secondary"
className={
v.source === "seed"
? "corner-squircle h-4 border-blue-500/25 bg-blue-500/10 px-1.5 font-mono text-[10px] text-blue-700 dark:text-blue-300"
: "corner-squircle h-4 px-1.5 font-mono text-[10px]"
}
>
{v.name}
</Badge>
))}
</div>
</div>
)}
<AvailableReferencesInline entries={vars} />
</div>
);
}

26
studio/frontend/src/features/recipe-studio/components/inline/inline-model.tsx
View file

@ -14,19 +14,6 @@ export function InlineModel(props: InlineModelProps): ReactElement {
if (props.config.kind === "model_provider") {
return (
<div className="grid gap-3 sm:grid-cols-2">
<InlineField label="Provider type">
<Input
className="nodrag h-8 w-full text-xs"
placeholder="openai-compatible"
value={props.config.provider_type}
onChange={(event) =>
props.onUpdate({
// biome-ignore lint/style/useNamingConvention: api schema
provider_type: event.target.value,
})
}
/>
</InlineField>
<InlineField label="Endpoint">
<Input
className="nodrag h-8 w-full text-xs"
@ -35,6 +22,19 @@ export function InlineModel(props: InlineModelProps): ReactElement {
onChange={(event) => props.onUpdate({ endpoint: event.target.value })}
/>
</InlineField>
<InlineField label="API key">
<Input
className="nodrag h-8 w-full text-xs"
placeholder="Optional"
value={props.config.api_key ?? ""}
onChange={(event) =>
props.onUpdate({
// biome-ignore lint/style/useNamingConvention: api schema
api_key: event.target.value,
})
}
/>
</InlineField>
</div>
);
}

3
studio/frontend/src/features/recipe-studio/components/inline/inline-policy.ts
View file

@ -27,6 +27,9 @@ export function getConfigUiMode(
}
return "dialog";
}
if (config.kind === "seed") {
return "inline";
}
if (config.kind === "expression") {
return "inline";
}

66
studio/frontend/src/features/recipe-studio/components/inline/inline-seed.tsx Normal file
View file

@ -0,0 +1,66 @@
import { DocumentAttachmentIcon, DocumentCodeIcon, Plant01Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import type { ReactElement } from "react";
import type { SeedConfig } from "../../types";
import { HfDatasetCombobox } from "../shared/hf-dataset-combobox";
import { InlineField } from "./inline-field";
type InlineSeedProps = {
config: SeedConfig;
onUpdate: (patch: Partial<SeedConfig>) => void;
};
export function InlineSeed({ config, onUpdate }: InlineSeedProps): ReactElement {
const mode = config.seed_source_type ?? "hf";
if (mode === "hf") {
return (
<div className="space-y-2">
<InlineField label="Dataset">
<HfDatasetCombobox
value={config.hf_repo_id}
accessToken={config.hf_token?.trim() || undefined}
onValueChange={(next) =>
onUpdate({
hf_repo_id: next,
hf_path: "",
seed_columns: [],
seed_drop_columns: [],
seed_preview_rows: [],
})
}
placeholder="org/repo"
/>
</InlineField>
<p className="text-[11px] text-muted-foreground">
Load columns in dialog.
</p>
</div>
);
}
const isLocal = mode === "local";
const fileName = isLocal
? config.local_file_name?.trim()
: config.unstructured_file_name?.trim();
return (
<div className="corner-squircle flex items-center gap-2 rounded-md border border-border/60 bg-muted/30 px-2 py-2">
<div className="corner-squircle rounded-md bg-primary/10 p-1.5 text-primary">
<HugeiconsIcon
icon={isLocal ? DocumentCodeIcon : DocumentAttachmentIcon}
className="size-3.5"
/>
</div>
<div className="min-w-0">
<p className="truncate text-xs font-medium">
{fileName || "No file selected"}
</p>
<p className="text-[11px] text-muted-foreground">
{isLocal ? "Structured file" : "Unstructured document"} · configure in dialog
</p>
</div>
<HugeiconsIcon icon={Plant01Icon} className="ml-auto size-3.5 text-muted-foreground/60" />
</div>
);
}

136
studio/frontend/src/features/recipe-studio/components/recipe-graph-aux-node.tsx
View file

@ -1,24 +1,21 @@
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Textarea } from "@/components/ui/textarea";
import {
Handle,
NodeResizer,
Position,
type Node,
type NodeProps,
useUpdateNodeInternals,
} from "@xyflow/react";
import { memo, type ReactElement, useEffect } from "react";
import { MAX_NODE_WIDTH, MIN_NODE_WIDTH } from "../constants";
import { useRecipeStudioStore } from "../stores/recipe-studio";
import type { LayoutDirection, LlmConfig, Score, ScoreOption } from "../types";
import {
AUX_HANDLE_CLASS,
getAuxSourceHandlePosition,
} from "../utils/handle-layout";
import type { LlmConfig, Score, ScoreOption } from "../types";
import { AUX_HANDLE_CLASS } from "../utils/handle-layout";
import { HANDLE_IDS } from "../utils/handles";
import { findInvalidJinjaReferences } from "../utils/refs";
import { getAvailableVariableEntries } from "../utils/variables";
import { AvailableReferencesInline } from "./shared/available-references-inline";
import { BaseNode, BaseNodeContent, BaseNodeHeader, BaseNodeHeaderTitle } from "./rf-ui/base-node";
type PromptField = "prompt" | "system_prompt";
@ -28,14 +25,12 @@ type PromptInputNodeData = {
llmId: string;
field: PromptField;
title: string;
layoutDirection: LayoutDirection;
};
type JudgeScoreNodeData = {
kind: "llm-judge-score";
llmId: string;
scoreIndex: number;
layoutDirection: LayoutDirection;
};
export type RecipeGraphAuxNodeData = PromptInputNodeData | JudgeScoreNodeData;
@ -62,30 +57,12 @@ function updateOptionAt(
);
}
function AuxVariableBadges({ llmId }: { llmId: string }): ReactElement | null {
const configs = useRecipeStudioStore((state) => state.configs);
const vars = getAvailableVariableEntries(configs, llmId);
if (vars.length === 0) return null;
return (
<div className="space-y-1">
<p className="text-[10px] font-medium text-muted-foreground">Available references</p>
<div className="flex flex-wrap gap-1">
{vars.map((v) => (
<Badge
key={`${v.source}:${v.name}`}
variant="secondary"
className={
v.source === "seed"
? "corner-squircle h-4 border-blue-500/25 bg-blue-500/10 px-1.5 font-mono text-[10px] text-blue-700 dark:text-blue-300"
: "corner-squircle h-4 px-1.5 font-mono text-[10px]"
}
>
{v.name}
</Badge>
))}
</div>
</div>
);
function AuxVariableBadges({
entries,
}: {
entries: ReturnType<typeof getAvailableVariableEntries>;
}): ReactElement | null {
return <AvailableReferencesInline entries={entries} />;
}
function AuxNodeBase({
@ -93,6 +70,7 @@ function AuxNodeBase({
data,
}: NodeProps<RecipeGraphAuxNodeType>): ReactElement | null {
const config = useRecipeStudioStore((state) => state.configs[data.llmId]);
const configs = useRecipeStudioStore((state) => state.configs);
const updateConfig = useRecipeStudioStore((state) => state.updateConfig);
const updateNodeInternals = useUpdateNodeInternals();
@ -104,30 +82,58 @@ function AuxNodeBase({
return null;
}
const sourcePosition = getAuxSourceHandlePosition(data.layoutDirection);
const sourceHandles = (
<>
<Handle
id={HANDLE_IDS.llmInputOutLeft}
type="source"
position={Position.Left}
isConnectable={false}
isConnectableStart={false}
className={AUX_HANDLE_CLASS}
/>
<Handle
id={HANDLE_IDS.llmInputOutRight}
type="source"
position={Position.Right}
isConnectable={false}
isConnectableStart={false}
className={AUX_HANDLE_CLASS}
/>
<Handle
id={HANDLE_IDS.llmInputOutTop}
type="source"
position={Position.Top}
isConnectable={false}
isConnectableStart={false}
className={AUX_HANDLE_CLASS}
/>
<Handle
id={HANDLE_IDS.llmInputOutBottom}
type="source"
position={Position.Bottom}
isConnectable={false}
isConnectableStart={false}
className={AUX_HANDLE_CLASS}
/>
</>
);
if (data.kind === "llm-prompt-input") {
const value = data.field === "prompt" ? config.prompt : config.system_prompt;
const variableEntries = getAvailableVariableEntries(configs, data.llmId);
const availableRefs = variableEntries.map((entry) => entry.name);
const hasInvalidRefs =
findInvalidJinjaReferences(value, availableRefs).length > 0;
return (
<BaseNode className="corner-squircle w-full min-w-0 rounded-lg border-border/60 bg-card shadow-sm">
<NodeResizer
isVisible={true}
minWidth={MIN_NODE_WIDTH}
minHeight={120}
maxWidth={MAX_NODE_WIDTH}
maxHeight={520}
color="var(--primary)"
lineClassName="!border-transparent !shadow-none"
lineStyle={{ opacity: 0 }}
handleClassName="!h-3 !w-3 !border-transparent !bg-transparent"
handleStyle={{ opacity: 0 }}
/>
<BaseNodeHeader className="border-b border-border/50 px-3 py-2">
<BaseNodeHeaderTitle className="text-xs">{data.title}</BaseNodeHeaderTitle>
</BaseNodeHeader>
<BaseNodeContent className="gap-2 px-3 py-2">
<Textarea
className="corner-squircle nodrag max-h-40 min-h-[88px] w-full resize-none overflow-y-auto text-xs"
className="corner-squircle nodrag nowheel max-h-40 min-h-[88px] w-full resize-none overflow-y-auto text-xs"
aria-invalid={hasInvalidRefs}
value={value}
onChange={(event) =>
updateConfig(data.llmId, {
@ -135,16 +141,9 @@ function AuxNodeBase({
} as Partial<LlmConfig>)
}
/>
<AuxVariableBadges llmId={data.llmId} />
<AuxVariableBadges entries={variableEntries} />
</BaseNodeContent>
<Handle
id={HANDLE_IDS.llmInputOut}
type="source"
position={sourcePosition}
isConnectable={false}
isConnectableStart={false}
className={AUX_HANDLE_CLASS}
/>
{sourceHandles}
</BaseNode>
);
}
@ -190,18 +189,6 @@ function AuxNodeBase({
return (
<BaseNode className="corner-squircle w-full min-w-0 rounded-lg border-border/60 bg-card shadow-sm">
<NodeResizer
isVisible={true}
minWidth={MIN_NODE_WIDTH}
minHeight={120}
maxWidth={MAX_NODE_WIDTH}
maxHeight={640}
color="var(--primary)"
lineClassName="!border-transparent !shadow-none"
lineStyle={{ opacity: 0 }}
handleClassName="!h-3 !w-3 !border-transparent !bg-transparent"
handleStyle={{ opacity: 0 }}
/>
<BaseNodeHeader className="border-b border-border/50 px-3 py-2">
<BaseNodeHeaderTitle className="text-xs">
{score.name.trim() || `Scorer ${data.scoreIndex + 1}`}
@ -218,7 +205,7 @@ function AuxNodeBase({
onChange={(event) => updateScore({ name: event.target.value })}
/>
<Textarea
className="corner-squircle nodrag max-h-32 min-h-[56px] w-full resize-none overflow-y-auto text-xs"
className="corner-squircle nodrag nowheel max-h-32 min-h-[56px] w-full resize-none overflow-y-auto text-xs"
placeholder="Score description"
value={score.description}
onChange={(event) => updateScore({ description: event.target.value })}
@ -260,14 +247,7 @@ function AuxNodeBase({
</Button>
</div>
</BaseNodeContent>
<Handle
id={HANDLE_IDS.llmInputOut}
type="source"
position={sourcePosition}
isConnectable={false}
isConnectableStart={false}
className={AUX_HANDLE_CLASS}
/>
{sourceHandles}
</BaseNode>
);
}

131
studio/frontend/src/features/recipe-studio/components/recipe-graph-node.tsx
View file

@ -20,7 +20,6 @@ import {
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import {
Handle,
NodeResizer,
Position,
useUpdateNodeInternals,
@ -32,18 +31,18 @@ import { useRecipeStudioStore } from "../stores/recipe-studio";
import type {
RecipeNode as RecipeGraphNodeType,
LlmType,
LayoutDirection,
NodeConfig,
SamplerType,
} from "../types";
import { NODE_HANDLE_CLASS } from "../utils/handle-layout";
import { getLlmJudgeScoreHandleId, HANDLE_IDS } from "../utils/handles";
import { HANDLE_IDS } from "../utils/handles";
import { InlineCategoryBadges } from "./inline/inline-category-badges";
import { InlineExpression } from "./inline/inline-expression";
import { InlineLlm } from "./inline/inline-llm";
import { InlineModel } from "./inline/inline-model";
import { isInlineConfig } from "./inline/inline-policy";
import { InlineSampler } from "./inline/inline-sampler";
import { InlineSeed } from "./inline/inline-seed";
import {
BaseNode,
BaseNodeContent,
@ -220,7 +219,7 @@ function getConfigSummary(config: NodeConfig | undefined): string {
return "Set HF dataset repo";
}
if (seedSourceType === "local") {
return "Upload CSV/JSON file";
return "Upload structured file";
}
return "Upload PDF/DOCX/TXT file";
}
@ -259,6 +258,9 @@ function renderNodeBody(
if (config.kind === "expression") {
return <InlineExpression config={config} onUpdate={onUpdate} />;
}
if (config.kind === "seed") {
return <InlineSeed config={config} onUpdate={onUpdate} />;
}
}
if (config?.kind === "sampler" && config.sampler_type === "category") {
@ -268,89 +270,6 @@ function renderNodeBody(
return <p className="text-xs text-muted-foreground">{summary}</p>;
}
type LlmInputHandleItem = {
id: string;
label: string;
};
function getLlmInputHandleItems(config: NodeConfig | undefined): LlmInputHandleItem[] {
if (!(config && config.kind === "llm")) {
return [];
}
const items: LlmInputHandleItem[] = [];
if (config.system_prompt.trim()) {
items.push({ id: HANDLE_IDS.llmSystemIn, label: "System" });
}
if (config.prompt.trim()) {
items.push({ id: HANDLE_IDS.llmPromptIn, label: "Prompt" });
}
if (config.llm_type === "judge") {
(config.scores ?? []).forEach((score, index) => {
items.push({
id: getLlmJudgeScoreHandleId(index),
label: score.name.trim() || `Score ${index + 1}`,
});
});
}
return items;
}
type LlmInputHandlesProps = {
items: LlmInputHandleItem[];
layoutDirection: LayoutDirection;
};
function LlmInputHandles({
items,
layoutDirection,
}: LlmInputHandlesProps): ReactElement | null {
if (items.length === 0) {
return null;
}
const isTopBottom = layoutDirection === "TB";
if (isTopBottom) {
return (
<div className="flex flex-wrap gap-2 pb-1">
{items.map((item) => (
<div
key={item.id}
className="pointer-events-none relative flex min-w-[80px] flex-1 justify-center pt-2"
>
<Handle
id={item.id}
type="target"
position={Position.Top}
className={NODE_HANDLE_CLASS}
style={{ left: "50%", top: 0, transform: "translate(-50%, -50%)" }}
/>
<span className="text-[10px] text-muted-foreground">{item.label}</span>
</div>
))}
</div>
);
}
return (
<div className="space-y-1 pb-1">
{items.map((item) => (
<div key={item.id} className="pointer-events-none relative min-w-0 pl-3">
<Handle
id={item.id}
type="target"
position={Position.Left}
className={NODE_HANDLE_CLASS}
style={{ left: -3, top: "50%", transform: "translate(-50%, -50%)" }}
/>
<span className="block truncate text-[10px] text-muted-foreground">
{item.label}
</span>
</div>
))}
</div>
);
}
function RecipeGraphNodeBase({
id,
data,
@ -418,7 +337,6 @@ function RecipeGraphNodeBase({
data.kind === "model_config" || data.kind === "model_provider";
const summary = getConfigSummary(config);
const nodeBody = renderNodeBody(config, summary, updateConfig);
const llmInputHandles = llmAuxVisible ? getLlmInputHandleItems(config) : [];
const canShowLlmAux =
config?.kind === "llm" &&
(Boolean(config.prompt.trim()) ||
@ -491,7 +409,6 @@ function RecipeGraphNodeBase({
</BaseNodeHeader>
<BaseNodeContent className="gap-2 px-3 py-2">
<LlmInputHandles items={llmInputHandles} layoutDirection={layoutDirection} />
{nodeBody}
</BaseNodeContent>
@ -506,6 +423,15 @@ function RecipeGraphNodeBase({
labelClassName="sr-only"
handleClassName={NODE_HANDLE_CLASS}
/>
<LabeledHandle
id={HANDLE_IDS.dataOutLeft}
title="Data output"
type="source"
position={Position.Left}
className="absolute inset-0 pointer-events-none opacity-0"
labelClassName="sr-only"
handleClassName={NODE_HANDLE_CLASS}
/>
<LabeledHandle
id={HANDLE_IDS.dataInTop}
title="Data input"
@ -515,6 +441,15 @@ function RecipeGraphNodeBase({
labelClassName="sr-only"
handleClassName={NODE_HANDLE_CLASS}
/>
<LabeledHandle
id={HANDLE_IDS.dataOutTop}
title="Data output"
type="source"
position={Position.Top}
className="absolute inset-0 pointer-events-none opacity-0"
labelClassName="sr-only"
handleClassName={NODE_HANDLE_CLASS}
/>
<LabeledHandle
id={HANDLE_IDS.dataOut}
title="Data output"
@ -524,6 +459,15 @@ function RecipeGraphNodeBase({
labelClassName="sr-only"
handleClassName={NODE_HANDLE_CLASS}
/>
<LabeledHandle
id={HANDLE_IDS.dataInRight}
title="Data input"
type="target"
position={Position.Right}
className="absolute inset-0 pointer-events-none opacity-0"
labelClassName="sr-only"
handleClassName={NODE_HANDLE_CLASS}
/>
<LabeledHandle
id={HANDLE_IDS.dataOutBottom}
title="Data output"
@ -533,6 +477,15 @@ function RecipeGraphNodeBase({
labelClassName="sr-only"
handleClassName={NODE_HANDLE_CLASS}
/>
<LabeledHandle
id={HANDLE_IDS.dataInBottom}
title="Data input"
type="target"
position={Position.Bottom}
className="absolute inset-0 pointer-events-none opacity-0"
labelClassName="sr-only"
handleClassName={NODE_HANDLE_CLASS}
/>
</>
)}

129
studio/frontend/src/features/recipe-studio/components/shared/available-references-inline.tsx Normal file
View file

@ -0,0 +1,129 @@
import { Badge } from "@/components/ui/badge";
import { type ReactElement, useLayoutEffect, useRef, useState } from "react";
import type { AvailableVariableEntry } from "../../utils/variables";
type AvailableReferencesInlineProps = {
entries: AvailableVariableEntry[];
};
const MAX_ROWS = 2;
export function AvailableReferencesInline({
entries,
}: AvailableReferencesInlineProps): ReactElement | null {
const [expanded, setExpanded] = useState(false);
const [collapsedCount, setCollapsedCount] = useState(entries.length);
const wrapperRef = useRef<HTMLDivElement | null>(null);
const measureRefs = useRef<Array<HTMLSpanElement | null>>([]);
useLayoutEffect(() => {
if (expanded) {
return;
}
const wrapper = wrapperRef.current;
const items = measureRefs.current.filter(
(node): node is HTMLSpanElement => Boolean(node),
);
if (!(wrapper && items.length > 0)) {
setCollapsedCount(entries.length);
return;
}
const compute = () => {
const rowTops: number[] = [];
let cutoff = items.length;
for (let i = 0; i < items.length; i += 1) {
const top = items[i].offsetTop;
if (!rowTops.some((value) => Math.abs(value - top) <= 1)) {
rowTops.push(top);
}
if (rowTops.length > MAX_ROWS) {
cutoff = i;
break;
}
}
if (cutoff < items.length) {
cutoff = Math.max(0, cutoff - 1);
}
setCollapsedCount(cutoff);
};
compute();
const observer = new ResizeObserver(compute);
observer.observe(wrapper);
return () => observer.disconnect();
}, [entries.length, expanded]);
if (entries.length === 0) {
return null;
}
const shown = expanded ? entries : entries.slice(0, collapsedCount);
const hiddenCount = Math.max(0, entries.length - shown.length);
return (
<div className="space-y-1">
<p className="text-[10px] font-medium text-muted-foreground">
Available references
</p>
<div ref={wrapperRef} className="relative">
{!expanded && (
<div className="invisible pointer-events-none absolute inset-0 -z-10">
<div className="flex flex-wrap gap-1">
{entries.map((entry, index) => (
<Badge
// biome-ignore lint/suspicious/noArrayIndexKey: static measurement mirror
key={`${entry.source}:${entry.name}:${index}`}
ref={(node) => {
measureRefs.current[index] = node;
}}
variant="secondary"
className={
entry.source === "seed"
? "corner-squircle h-4 border-blue-500/25 bg-blue-500/10 px-1.5 font-mono text-[10px] text-blue-700 dark:text-blue-300"
: "corner-squircle h-4 px-1.5 font-mono text-[10px]"
}
>
{entry.name}
</Badge>
))}
</div>
</div>
)}
<div className="flex flex-wrap gap-1">
{shown.map((entry) => (
<Badge
key={`${entry.source}:${entry.name}`}
variant="secondary"
className={
entry.source === "seed"
? "corner-squircle h-4 border-blue-500/25 bg-blue-500/10 px-1.5 font-mono text-[10px] text-blue-700 dark:text-blue-300"
: "corner-squircle h-4 px-1.5 font-mono text-[10px]"
}
>
{entry.name}
</Badge>
))}
{!expanded && hiddenCount > 0 && (
<button
type="button"
className="corner-squircle h-4 px-1.5 text-[10px] text-muted-foreground hover:text-foreground"
onClick={() => setExpanded(true)}
>
+{hiddenCount} more
</button>
)}
{expanded && collapsedCount < entries.length && (
<button
type="button"
className="corner-squircle h-4 px-1.5 text-[10px] text-muted-foreground hover:text-foreground"
onClick={() => setExpanded(false)}
>
Show less
</button>
)}
</div>
</div>
</div>
);
}

122
studio/frontend/src/features/recipe-studio/components/shared/hf-dataset-combobox.tsx Normal file
View file

@ -0,0 +1,122 @@
import {
Combobox,
ComboboxContent,
ComboboxEmpty,
ComboboxInput,
ComboboxItem,
ComboboxList,
} from "@/components/ui/combobox";
import { Spinner } from "@/components/ui/spinner";
import { useDebouncedValue, useHfDatasetSearch } from "@/hooks";
import { type ReactElement, useEffect, useMemo, useRef, useState } from "react";
type HfDatasetComboboxProps = {
value: string;
onValueChange: (value: string) => void;
accessToken?: string;
inputId?: string;
placeholder?: string;
className?: string;
};
export function HfDatasetCombobox({
value,
onValueChange,
accessToken,
inputId,
placeholder = "Search datasets...",
className,
}: HfDatasetComboboxProps): ReactElement {
const [inputValue, setInputValue] = useState(value);
const selectingRef = useRef(false);
const anchorRef = useRef<HTMLDivElement>(null);
const debouncedQuery = useDebouncedValue(inputValue);
useEffect(() => {
setInputValue(value);
}, [value]);
const { results, isLoading, error } = useHfDatasetSearch(debouncedQuery, {
accessToken,
});
const items = useMemo(() => {
const ids = results.map((item) => item.id);
const selected = value.trim();
if (selected && !ids.includes(selected)) {
ids.push(selected);
}
return ids;
}, [results, value]);
return (
<div
ref={anchorRef}
className={className}
onKeyDown={(event) => {
if (event.key !== "Enter") return;
if (!(event.target instanceof HTMLInputElement)) return;
event.preventDefault();
if (items.length > 0) {
onValueChange(items[0]);
return;
}
const typed = event.target.value.trim();
if (typed) {
onValueChange(typed);
}
}}
>
<Combobox
items={items}
filteredItems={items}
filter={null}
value={value.trim() ? value : null}
onValueChange={(next) => onValueChange(next ?? "")}
onInputValueChange={(next) => {
if (selectingRef.current) {
selectingRef.current = false;
return;
}
setInputValue(next);
}}
itemToStringValue={(item) => item}
autoHighlight={true}
>
<ComboboxInput
id={inputId}
className="nodrag w-full"
placeholder={placeholder}
/>
<ComboboxContent anchor={anchorRef}>
{isLoading ? (
<div className="flex items-center gap-2 px-2 py-3 text-xs text-muted-foreground">
<Spinner className="size-3.5" />
Searching...
</div>
) : (
<ComboboxEmpty>No datasets found</ComboboxEmpty>
)}
<ComboboxList>
{(id: string) => (
<ComboboxItem
key={id}
value={id}
onPointerDown={() => {
selectingRef.current = true;
}}
>
{id}
</ComboboxItem>
)}
</ComboboxList>
</ComboboxContent>
</Combobox>
{error && (
<p className="mt-1 text-xs text-destructive">
{error}
</p>
)}
</div>
);
}

26
studio/frontend/src/features/recipe-studio/dialogs/expression/expression-dialog.tsx
View file

@ -7,7 +7,11 @@ import {
} from "@/components/ui/select";
import { Textarea } from "@/components/ui/textarea";
import type { ReactElement } from "react";
import { useMemo } from "react";
import { useRecipeStudioStore } from "../../stores/recipe-studio";
import type { ExpressionConfig, ExpressionDtype } from "../../types";
import { findInvalidJinjaReferences } from "../../utils/refs";
import { getAvailableVariables } from "../../utils/variables";
import { AvailableVariables } from "../shared/available-variables";
import { FieldLabel } from "../shared/field-label";
import { NameField } from "../shared/name-field";
@ -23,8 +27,21 @@ export function ExpressionDialog({
config,
onUpdate,
}: ExpressionDialogProps): ReactElement {
const configs = useRecipeStudioStore((state) => state.configs);
const dtypeId = `${config.id}-dtype`;
const exprId = `${config.id}-expr`;
const validReferences = useMemo(
() => getAvailableVariables(configs, config.id),
[configs, config.id],
);
const invalidExprRefs = useMemo(
() => findInvalidJinjaReferences(config.expr, validReferences),
[config.expr, validReferences],
);
const invalidExprText = invalidExprRefs
.slice(0, 3)
.map((ref) => `{{ ${ref} }}`)
.join(", ");
const updateField = <K extends keyof ExpressionConfig>(
key: K,
value: ExpressionConfig[K],
@ -71,10 +88,19 @@ export function ExpressionDialog({
<Textarea
id={exprId}
className="corner-squircle nodrag"
aria-invalid={invalidExprRefs.length > 0}
placeholder="{{ category_1 }} - {{ subcategory_1 }}"
value={config.expr}
onChange={(event) => updateField("expr", event.target.value)}
/>
{invalidExprRefs.length > 0 && (
<p className="text-xs text-destructive">
Unknown reference: {invalidExprText}
{invalidExprRefs.length > 3
? ` +${invalidExprRefs.length - 3} more`
: ""}
</p>
)}
<p className="text-xs text-muted-foreground">
Use Jinja2. Reference columns like {"{{ column_name }}"}.
</p>

4
studio/frontend/src/features/recipe-studio/dialogs/import-dialog.tsx
View file

@ -50,7 +50,7 @@ export function ImportDialog({
position="absolute"
overlayPosition="absolute"
overlayClassName="bg-transparent"
className="corner-squircle max-h-[650px] overflow-auto sm:max-w-2xl"
className="corner-squircle max-h-[650px] overflow-auto sm:max-w-2xl shadow-border"
>
<DialogHeader>
<DialogTitle>Import recipe</DialogTitle>
@ -63,7 +63,7 @@ export function ImportDialog({
/>
<Textarea
id={payloadId}
className="corner-squircle nodrag min-h-[220px]"
className="corner-squircle nodrag min-h-[220px] max-h-[450px]"
placeholder='{"recipe": { "columns": [] }}'
value={value}
onChange={(event) => setValue(event.target.value)}

48
studio/frontend/src/features/recipe-studio/dialogs/llm/general-tab.tsx
View file

@ -14,8 +14,11 @@ import {
SelectValue,
} from "@/components/ui/select";
import { Textarea } from "@/components/ui/textarea";
import { type ReactElement, type RefObject } from "react";
import { type ReactElement, type RefObject, useMemo } from "react";
import { useRecipeStudioStore } from "../../stores/recipe-studio";
import type { LlmConfig } from "../../types";
import { findInvalidJinjaReferences } from "../../utils/refs";
import { getAvailableVariables } from "../../utils/variables";
import { AvailableVariables } from "../shared/available-variables";
import { FieldLabel } from "../shared/field-label";
import { NameField } from "../shared/name-field";
@ -54,6 +57,7 @@ export function LlmGeneralTab({
modelAliasAnchorRef,
onUpdate,
}: LlmGeneralTabProps): ReactElement {
const configs = useRecipeStudioStore((state) => state.configs);
const modelAliasId = `${config.id}-model-alias`;
const codeLangId = `${config.id}-code-lang`;
const promptId = `${config.id}-prompt`;
@ -61,6 +65,26 @@ export function LlmGeneralTab({
const systemPromptId = `${config.id}-system-prompt`;
const hasModelConfigs = modelConfigAliases.length > 0;
const hasModelProviders = modelProviderOptions.length > 0;
const validReferences = useMemo(
() => getAvailableVariables(configs, config.id),
[configs, config.id],
);
const invalidPromptRefs = useMemo(
() => findInvalidJinjaReferences(config.prompt, validReferences),
[config.prompt, validReferences],
);
const invalidSystemRefs = useMemo(
() => findInvalidJinjaReferences(config.system_prompt, validReferences),
[config.system_prompt, validReferences],
);
const invalidPromptText = invalidPromptRefs
.slice(0, 3)
.map((ref) => `{{ ${ref} }}`)
.join(", ");
const invalidSystemText = invalidSystemRefs
.slice(0, 3)
.map((ref) => `{{ ${ref} }}`)
.join(", ");
return (
<div className="space-y-4">
@ -147,10 +171,19 @@ export function LlmGeneralTab({
/>
<Textarea
id={promptId}
className="corner-squircle nodrag"
className="corner-squircle nodrag max-h-[450px] overflow-auto"
aria-invalid={invalidPromptRefs.length > 0}
value={config.prompt}
onChange={(event) => onUpdate({ prompt: event.target.value })}
/>
{invalidPromptRefs.length > 0 && (
<p className="text-xs text-destructive">
Unknown reference: {invalidPromptText}
{invalidPromptRefs.length > 3
? ` +${invalidPromptRefs.length - 3} more`
: ""}
</p>
)}
</div>
{config.llm_type === "structured" && (
<div className="grid gap-2">
@ -177,10 +210,19 @@ export function LlmGeneralTab({
/>
<Textarea
id={systemPromptId}
className="corner-squircle nodrag"
className="corner-squircle nodrag max-h-[450px] overflow-auto"
aria-invalid={invalidSystemRefs.length > 0}
value={config.system_prompt}
onChange={(event) => onUpdate({ system_prompt: event.target.value })}
/>
{invalidSystemRefs.length > 0 && (
<p className="text-xs text-destructive">
Unknown reference: {invalidSystemText}
{invalidSystemRefs.length > 3
? ` +${invalidSystemRefs.length - 3} more`
: ""}
</p>
)}
</div>
</div>
);

124
studio/frontend/src/features/recipe-studio/dialogs/models/model-provider-dialog.tsx
View file

@ -1,6 +1,11 @@
import {
Collapsible,
CollapsibleContent,
CollapsibleTrigger,
} from "@/components/ui/collapsible";
import { Input } from "@/components/ui/input";
import { Textarea } from "@/components/ui/textarea";
import type { ReactElement } from "react";
import { type ReactElement, useState } from "react";
import type { ModelProviderConfig } from "../../types";
import { FieldLabel } from "../shared/field-label";
import { NameField } from "../shared/name-field";
@ -14,8 +19,8 @@ export function ModelProviderDialog({
config,
onUpdate,
}: ModelProviderDialogProps): ReactElement {
const [optionalOpen, setOptionalOpen] = useState(false);
const endpointId = `${config.id}-endpoint`;
const providerTypeId = `${config.id}-provider-type`;
const apiKeyEnvId = `${config.id}-api-key-env`;
const apiKeyId = `${config.id}-api-key`;
const extraHeadersId = `${config.id}-extra-headers`;
@ -33,22 +38,6 @@ export function ModelProviderDialog({
value={config.name}
onChange={(value) => onUpdate({ name: value })}
/>
<div className="grid gap-2">
<FieldLabel
label="Provider type"
htmlFor={providerTypeId}
hint="Provider adapter type, e.g. openai or openrouter."
/>
<Input
id={providerTypeId}
className="nodrag"
placeholder="openai"
value={config.provider_type}
onChange={(event) =>
updateField("provider_type", event.target.value)
}
/>
</div>
<div className="grid gap-2">
<FieldLabel
label="Endpoint"
@ -63,20 +52,6 @@ export function ModelProviderDialog({
onChange={(event) => updateField("endpoint", event.target.value)}
/>
</div>
<div className="grid gap-2">
<FieldLabel
label="API key env (optional)"
htmlFor={apiKeyEnvId}
hint="Env var name to read secret key from runtime."
/>
<Input
id={apiKeyEnvId}
className="nodrag"
placeholder="OPENAI_API_KEY"
value={config.api_key_env ?? ""}
onChange={(event) => updateField("api_key_env", event.target.value)}
/>
</div>
<div className="grid gap-2">
<FieldLabel
label="API key (optional)"
@ -90,36 +65,61 @@ export function ModelProviderDialog({
onChange={(event) => updateField("api_key", event.target.value)}
/>
</div>
<div className="grid gap-2">
<FieldLabel
label="Extra headers (JSON)"
htmlFor={extraHeadersId}
hint="Optional request headers merged into every call."
/>
<Textarea
id={extraHeadersId}
className="corner-squircle nodrag"
placeholder='{"X-Header": "value"}'
value={config.extra_headers ?? ""}
onChange={(event) =>
updateField("extra_headers", event.target.value)
}
/>
</div>
<div className="grid gap-2">
<FieldLabel
label="Extra body (JSON)"
htmlFor={extraBodyId}
hint="Optional payload fields merged into requests."
/>
<Textarea
id={extraBodyId}
className="corner-squircle nodrag"
placeholder='{"key": "value"}'
value={config.extra_body ?? ""}
onChange={(event) => updateField("extra_body", event.target.value)}
/>
</div>
<Collapsible open={optionalOpen} onOpenChange={setOptionalOpen}>
<CollapsibleTrigger asChild={true}>
<button
type="button"
className="flex w-full items-center justify-between text-left text-xs text-muted-foreground"
>
<span className="font-semibold uppercase">Optional</span>
<span>{optionalOpen ? "Hide" : "Show"}</span>
</button>
</CollapsibleTrigger>
<CollapsibleContent className="mt-3 space-y-4">
<div className="grid gap-2">
<FieldLabel
label="API key env (optional)"
htmlFor={apiKeyEnvId}
hint="Env var name to read secret key from runtime."
/>
<Input
id={apiKeyEnvId}
className="nodrag"
placeholder="OPENAI_API_KEY"
value={config.api_key_env ?? ""}
onChange={(event) => updateField("api_key_env", event.target.value)}
/>
</div>
<div className="grid gap-2">
<FieldLabel
label="Extra headers (JSON)"
htmlFor={extraHeadersId}
hint="Optional request headers merged into every call."
/>
<Textarea
id={extraHeadersId}
className="corner-squircle nodrag"
placeholder='{"X-Header": "value"}'
value={config.extra_headers ?? ""}
onChange={(event) => updateField("extra_headers", event.target.value)}
/>
</div>
<div className="grid gap-2">
<FieldLabel
label="Extra body (JSON)"
htmlFor={extraBodyId}
hint="Optional payload fields merged into requests."
/>
<Textarea
id={extraBodyId}
className="corner-squircle nodrag"
placeholder='{"key": "value"}'
value={config.extra_body ?? ""}
onChange={(event) => updateField("extra_body", event.target.value)}
/>
</div>
</CollapsibleContent>
</Collapsible>
</div>
);
}

4
studio/frontend/src/features/recipe-studio/dialogs/preview-dialog.tsx
View file

@ -285,7 +285,7 @@ export function RunDialog({
position="absolute"
overlayPosition="absolute"
overlayClassName="bg-transparent"
className="corner-squircle sm:max-w-2xl"
className="corner-squircle sm:max-w-2xl shadow-border"
>
<DialogHeader>
<DialogTitle>{kindLabel} settings</DialogTitle>
@ -294,7 +294,7 @@ export function RunDialog({
</p>
</DialogHeader>
<div className="flex items-center justify-between rounded-xl border bg-muted/20 px-3 py-2 text-sm">
<div className="flex items-center justify-between text-sm">
<span className="font-medium text-foreground">Preview mode</span>
<Switch
checked={kind === "preview"}

2
studio/frontend/src/features/recipe-studio/dialogs/processors-dialog.tsx
View file

@ -68,7 +68,7 @@ export function ProcessorsDialog({
position="absolute"
overlayPosition="absolute"
overlayClassName="bg-transparent"
className="corner-squircle max-h-[650px] overflow-auto sm:max-w-2xl"
className="corner-squircle max-h-[650px] overflow-auto sm:max-w-2xl shadow-border"
>
<VisuallyHidden.Root>
<DialogTitle>Processors</DialogTitle>

22
studio/frontend/src/features/recipe-studio/dialogs/samplers/category-dialog.tsx
View file

@ -133,12 +133,15 @@ export function CategoryDialog({
Add values first, then set optional weights.
</p>
) : (
<div className="grid gap-2">
<div className="grid gap-2 sm:grid-cols-2 lg:grid-cols-3">
{(config.values ?? []).map((value, index) => (
<div key={`${value}-weight`} className="flex items-center gap-3">
<span className="max-w-20 truncate text-xs text-muted-foreground">
<div key={`${value}-weight`} className="space-y-1">
<p
className="truncate text-xs text-muted-foreground"
title={value}
>
{value}
</span>
</p>
<Input
type="number"
className="nodrag w-full"
@ -236,15 +239,18 @@ export function CategoryDialog({
<p className="text-xs font-semibold uppercase text-muted-foreground">
Rule weights (optional)
</p>
<div className="grid gap-2">
<div className="grid gap-2 sm:grid-cols-2 lg:grid-cols-3">
{(params.values ?? []).map((value, index) => (
<div
key={`${condition}-${value}-${index}-weight`}
className="flex items-center gap-3"
className="space-y-1"
>
<span className="w-28 truncate text-xs text-muted-foreground">
<p
className="truncate text-xs text-muted-foreground"
title={value}
>
{value}
</span>
</p>
<Input
type="number"
className="nodrag"

61
studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx
View file

@ -45,6 +45,7 @@ import type {
SeedSamplingStrategy,
SeedSelectionType,
} from "../../types";
import { HfDatasetCombobox } from "../../components/shared/hf-dataset-combobox";
import { FieldLabel } from "../shared/field-label";
const SAMPLING_OPTIONS: Array<{ value: SeedSamplingStrategy; label: string }> = [
@ -209,8 +210,6 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl
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`;
const datasetId = `${config.id}-hf-dataset`;
const chunkSizeId = `${config.id}-chunk-size`;
const chunkOverlapId = `${config.id}-chunk-overlap`;
@ -221,10 +220,8 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl
if (mode === "hf") {
const dataset = config.hf_repo_id.trim();
if (!dataset) return null;
const subset = config.hf_subset?.trim() ?? "";
const split = config.hf_split?.trim() || "train";
const token = config.hf_token?.trim() ?? "";
return `hf:${dataset}|${subset}|${split}|${token}`;
return `hf:${dataset}|${token}`;
}
if (mode === "local") {
if (!localFile) return null;
@ -255,8 +252,7 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl
const response = await inspectSeedDataset({
dataset_name: datasetName,
hf_token: config.hf_token?.trim() || undefined,
subset: config.hf_subset || undefined,
split: config.hf_split || "train",
subset: undefined,
preview_size: 10,
});
onUpdate({
@ -266,8 +262,8 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl
response.columns.includes(name),
),
seed_preview_rows: response.preview_rows ?? [],
hf_split: response.split ?? config.hf_split ?? "",
hf_subset: response.subset ?? config.hf_subset ?? "",
hf_split: "",
hf_subset: "",
local_file_name: "",
unstructured_file_name: "",
});
@ -416,14 +412,17 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl
hint="Hugging Face dataset repo id (org/repo)."
/>
<div className="flex items-center gap-2">
<Input
id={datasetId}
className="nodrag flex-1"
placeholder="org/repo"
<HfDatasetCombobox
inputId={datasetId}
className="flex-1"
value={config.hf_repo_id}
onChange={(event) =>
accessToken={config.hf_token?.trim() || undefined}
placeholder="org/repo"
onValueChange={(nextValue) =>
onUpdate({
hf_repo_id: event.target.value,
hf_repo_id: nextValue,
hf_subset: "",
hf_split: "",
hf_path: "",
seed_columns: [],
seed_drop_columns: [],
@ -458,43 +457,13 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl
/>
</div>
<div className="grid grid-cols-2 gap-3">
<div className="grid gap-2">
<FieldLabel
label="Subset (optional)"
htmlFor={subsetId}
hint="Dataset config/subset name."
/>
<Input
id={subsetId}
className="nodrag"
placeholder="default"
value={config.hf_subset ?? ""}
onChange={(event) => onUpdate({ hf_subset: event.target.value })}
/>
</div>
<div className="grid gap-2">
<FieldLabel
label="Split"
htmlFor={splitId}
hint="Split to inspect (default train)."
/>
<Input
id={splitId}
className="nodrag"
placeholder="train"
value={config.hf_split ?? ""}
onChange={(event) => onUpdate({ hf_split: event.target.value })}
/>
</div>
</div>
</>
)}
{mode === "local" && (
<div className="grid gap-2">
<FieldLabel
label="Local file"
label="Structured file"
hint="Upload CSV, JSON, or JSONL seed file."
/>
<div className="flex items-center gap-2">

56
studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts
View file

@ -64,6 +64,58 @@ function stripApiKeys(value: unknown): unknown {
return output;
}
function sanitizeSeedForShare(payload: unknown): unknown {
if (!payload || typeof payload !== "object") {
return payload;
}
const root = payload as Record<string, unknown>;
const recipe =
root.recipe && typeof root.recipe === "object"
? (root.recipe as Record<string, unknown>)
: null;
const ui =
root.ui && typeof root.ui === "object"
? (root.ui as Record<string, unknown>)
: null;
const seedConfig =
recipe?.seed_config && typeof recipe.seed_config === "object"
? (recipe.seed_config as Record<string, unknown>)
: null;
const source =
seedConfig?.source && typeof seedConfig.source === "object"
? (seedConfig.source as Record<string, unknown>)
: null;
if (source && "token" in source) {
delete source.token;
}
const uiSourceType =
typeof ui?.seed_source_type === "string" ? ui.seed_source_type : null;
const sourceType =
typeof source?.seed_type === "string" ? source.seed_type : null;
const shouldResetLocalState =
sourceType === "local" ||
uiSourceType === "local" ||
uiSourceType === "unstructured";
if (shouldResetLocalState) {
if (source && "path" in source) {
source.path = "";
}
if (ui) {
ui.seed_columns = [];
ui.seed_drop_columns = [];
ui.seed_preview_rows = [];
ui.local_file_name = "";
ui.unstructured_file_name = "";
}
}
return root;
}
export function useRecipePersistence({
recipeId,
initialRecipeName,
@ -160,14 +212,14 @@ export function useRecipePersistence({
const copyRecipe = useCallback(async (): Promise<void> => {
setCopied(false);
try {
const safePayload = stripApiKeys(payloadResult.payload);
const safePayload = sanitizeSeedForShare(stripApiKeys(payloadResult.payload));
const ok = await copyTextToClipboard(JSON.stringify(safePayload, null, 2));
if (!ok) {
throw new Error("Clipboard not available.");
}
setCopied(true);
window.setTimeout(() => setCopied(false), 1500);
toastSuccess("Payload copied");
toastSuccess("👨‍🍳 Recipe copied");
} catch (error) {
console.error("Copy failed:", error);
toastError("Copy failed", "Could not copy payload.");

223
studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx
View file

@ -18,16 +18,22 @@ import {
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import {
type DragEvent as ReactDragEvent,
type ReactElement,
useCallback,
useEffect,
useMemo,
useRef,
useState,
} from "react";
import { useShallow } from "zustand/react/shallow";
import "@xyflow/react/dist/style.css";
import { RecipeGraphAuxNode, type RecipeGraphAuxNodeData } from "./components/recipe-graph-aux-node";
import { BlockSheet } from "./components/block-sheet";
import {
BlockSheet,
RECIPE_BLOCK_DND_MIME,
type RecipeBlockDragPayload,
} from "./components/block-sheet";
import { LayoutControls } from "./components/controls/layout-controls";
import { ViewportControls } from "./components/controls/viewport-controls";
import { ExecutionsView } from "./components/executions/executions-view";
@ -44,10 +50,14 @@ import { ProcessorsDialog } from "./dialogs/processors-dialog";
import { useRecipeStudioActions } from "./hooks/use-recipe-studio-actions";
import { useRecipeStudioStore } from "./stores/recipe-studio";
import type {
LlmType,
RecipeNode as RecipeBuilderNode,
RecipeNodeData,
SamplerType,
} from "./types";
import type { SeedBlockType } from "./blocks/registry";
import { deriveDisplayGraph } from "./utils/graph/derive-display-graph";
import { getFitNodeIdsIgnoringNotes } from "./utils/graph/fit-view";
import { buildRecipePayload } from "./utils/payload";
import type { RecipePayload } from "./utils/payload/types";
import { buildDefaultSchemaTransform } from "./utils/processors";
@ -63,6 +73,35 @@ import type { RecipeStudioView } from "./execution-types";
const NODE_TYPES: NodeTypes = { builder: RecipeNode, aux: RecipeGraphAuxNode };
const EDGE_TYPES: EdgeTypes = { canvas: DataEdge, semantic: RecipeGraphSemanticEdge };
const SUPPORTED_DRAG_KINDS: RecipeBlockDragPayload["kind"][] = [
"sampler",
"seed",
"llm",
"expression",
"note",
];
function parseRecipeBlockDragPayload(raw: string): RecipeBlockDragPayload | null {
try {
const parsed = JSON.parse(raw) as {
kind?: RecipeBlockDragPayload["kind"];
type?: RecipeBlockDragPayload["type"];
};
if (
!parsed.kind ||
!parsed.type ||
!SUPPORTED_DRAG_KINDS.includes(parsed.kind)
) {
return null;
}
return {
kind: parsed.kind,
type: parsed.type,
};
} catch {
return null;
}
}
export type PersistRecipeInput = {
id: string | null;
@ -94,7 +133,6 @@ export function RecipeStudioPage({
nodes,
edges,
auxNodePositions,
auxNodeSizes,
llmAuxVisibility,
configs,
processors,
@ -125,15 +163,11 @@ export function RecipeStudioPage({
setLayoutDirection,
applyLayout,
setAuxNodePosition,
setAuxNodeSize,
syncAuxNodePositions,
syncAuxNodeSizes,
} = useRecipeStudioStore(
useShallow((state) => ({
nodes: state.nodes,
edges: state.edges,
auxNodePositions: state.auxNodePositions,
auxNodeSizes: state.auxNodeSizes,
llmAuxVisibility: state.llmAuxVisibility,
configs: state.configs,
processors: state.processors,
@ -164,14 +198,12 @@ export function RecipeStudioPage({
setLayoutDirection: state.setLayoutDirection,
applyLayout: state.applyLayout,
setAuxNodePosition: state.setAuxNodePosition,
setAuxNodeSize: state.setAuxNodeSize,
syncAuxNodePositions: state.syncAuxNodePositions,
syncAuxNodeSizes: state.syncAuxNodeSizes,
})),
);
const [sheetContainer, setSheetContainer] = useState<HTMLDivElement | null>(
null,
);
const flowContainerRef = useRef<HTMLDivElement | null>(null);
const [blockSheetOpen, setBlockSheetOpen] = useState(false);
const [activeView, setActiveView] = useState<RecipeStudioView>("editor");
const [processorsOpen, setProcessorsOpen] = useState(false);
@ -179,6 +211,7 @@ export function RecipeStudioPage({
const [reactFlowInstance, setReactFlowInstance] = useState<
ReactFlowInstance<Node<RecipeNodeData | RecipeGraphAuxNodeData>, Edge> | null
>(null);
const lastProcessedFitTickRef = useRef(0);
const handleExecutionStart = useCallback(() => {
setActiveView("executions");
}, []);
@ -202,12 +235,10 @@ export function RecipeStudioPage({
configs,
layoutDirection,
auxNodePositions,
auxNodeSizes,
llmAuxVisibility,
});
}, [
auxNodePositions,
auxNodeSizes,
configs,
edges,
layoutDirection,
@ -218,12 +249,6 @@ export function RecipeStudioPage({
() => displayGraph.nodes.map((node) => node.id),
[displayGraph.nodes],
);
useEffect(() => {
syncAuxNodePositions(displayGraph.auxNodeIds, displayGraph.auxDefaults);
}, [displayGraph.auxDefaults, displayGraph.auxNodeIds, syncAuxNodePositions]);
useEffect(() => {
syncAuxNodeSizes(displayGraph.auxNodeIds);
}, [displayGraph.auxNodeIds, syncAuxNodeSizes]);
const handleNodeClick = useCallback(
(_: unknown, node: Node<RecipeNodeData | RecipeGraphAuxNodeData>) => {
@ -250,7 +275,7 @@ export function RecipeStudioPage({
const handleNodesChange = useCallback(
(changes: NodeChange<Node<RecipeNodeData | RecipeGraphAuxNodeData>>[]) => {
applyAuxNodeChanges(changes, { setAuxNodePosition, setAuxNodeSize });
applyAuxNodeChanges(changes, { setAuxNodePosition });
const next = filterNodeChangesByIds(
changes as NodeChange<RecipeBuilderNode>[],
baseNodeIds,
@ -259,7 +284,7 @@ export function RecipeStudioPage({
onNodesChange(next);
}
},
[baseNodeIds, onNodesChange, setAuxNodePosition, setAuxNodeSize],
[baseNodeIds, onNodesChange, setAuxNodePosition],
);
const handleEdgesChange = useCallback(
@ -272,6 +297,116 @@ export function RecipeStudioPage({
[baseEdgeIds, onEdgesChange],
);
const handleDragOver = useCallback((event: ReactDragEvent<HTMLDivElement>) => {
if (
!event.dataTransfer.types.includes(RECIPE_BLOCK_DND_MIME) &&
!event.dataTransfer.types.includes("text/plain")
) {
return;
}
event.preventDefault();
event.dataTransfer.dropEffect = "copy";
}, []);
const handleDrop = useCallback(
(event: ReactDragEvent<HTMLDivElement>) => {
if (!reactFlowInstance) {
return;
}
const raw =
event.dataTransfer.getData(RECIPE_BLOCK_DND_MIME) ||
event.dataTransfer.getData("text/plain");
if (!raw) {
return;
}
const payload = parseRecipeBlockDragPayload(raw);
if (!payload) {
return;
}
event.preventDefault();
const position = reactFlowInstance.screenToFlowPosition({
x: event.clientX,
y: event.clientY,
});
if (payload.kind === "sampler") {
addSamplerNode(payload.type as SamplerType, position, false);
return;
}
if (payload.kind === "seed") {
addSeedNode(payload.type as SeedBlockType, position, false);
return;
}
if (payload.kind === "expression") {
addExpressionNode(position, false);
return;
}
if (payload.kind === "note") {
addMarkdownNoteNode(position, false);
return;
}
if (payload.type === "model_provider") {
addModelProviderNode(position, false);
return;
}
if (payload.type === "model_config") {
addModelConfigNode(position, false);
return;
}
addLlmNode(payload.type as LlmType, position, false);
},
[
addExpressionNode,
addLlmNode,
addMarkdownNoteNode,
addModelConfigNode,
addModelProviderNode,
addSamplerNode,
addSeedNode,
reactFlowInstance,
],
);
const getViewportCenterPosition = useCallback(() => {
if (!reactFlowInstance || !flowContainerRef.current) {
return undefined;
}
const rect = flowContainerRef.current.getBoundingClientRect();
return reactFlowInstance.screenToFlowPosition({
x: rect.left + rect.width / 2,
y: rect.top + rect.height / 2,
});
}, [reactFlowInstance]);
const handleAddSamplerFromSheet = useCallback(
(type: SamplerType) => {
addSamplerNode(type, getViewportCenterPosition());
},
[addSamplerNode, getViewportCenterPosition],
);
const handleAddSeedFromSheet = useCallback(
(type: SeedBlockType) => {
addSeedNode(type, getViewportCenterPosition());
},
[addSeedNode, getViewportCenterPosition],
);
const handleAddLlmFromSheet = useCallback(
(type: LlmType) => {
addLlmNode(type, getViewportCenterPosition());
},
[addLlmNode, getViewportCenterPosition],
);
const handleAddModelProviderFromSheet = useCallback(() => {
addModelProviderNode(getViewportCenterPosition());
}, [addModelProviderNode, getViewportCenterPosition]);
const handleAddModelConfigFromSheet = useCallback(() => {
addModelConfigNode(getViewportCenterPosition());
}, [addModelConfigNode, getViewportCenterPosition]);
const handleAddExpressionFromSheet = useCallback(() => {
addExpressionNode(getViewportCenterPosition());
}, [addExpressionNode, getViewportCenterPosition]);
const handleAddMarkdownNoteFromSheet = useCallback(() => {
addMarkdownNoteNode(getViewportCenterPosition());
}, [addMarkdownNoteNode, getViewportCenterPosition]);
const configList = useMemo(() => Object.values(configs), [configs]);
const config = activeConfigId ? configs[activeConfigId] : null;
const dialogOptions = useMemo(
@ -288,8 +423,16 @@ export function RecipeStudioPage({
}, []);
const payloadResult = useMemo(
() => buildRecipePayload(configs, nodes, edges, processors, layoutDirection),
[configs, edges, layoutDirection, nodes, processors],
() =>
buildRecipePayload(
configs,
nodes,
edges,
processors,
layoutDirection,
auxNodePositions,
),
[auxNodePositions, configs, edges, layoutDirection, nodes, processors],
);
const getCurrentPayloadFromStore = useCallback((): RecipePayload => {
const state = useRecipeStudioStore.getState();
@ -299,6 +442,7 @@ export function RecipeStudioPage({
state.edges,
state.processors,
state.layoutDirection,
state.auxNodePositions,
).payload;
}, []);
const {
@ -371,13 +515,23 @@ export function RecipeStudioPage({
runDialogKind === "preview" ? previewLoading : fullLoading;
useEffect(() => {
if (!reactFlowInstance || activeView !== "editor" || fitViewTick === 0) {
if (!reactFlowInstance || fitViewTick === 0 || activeView !== "editor") {
return;
}
if (lastProcessedFitTickRef.current === fitViewTick) {
return;
}
lastProcessedFitTickRef.current = fitViewTick;
let frame2 = 0;
let frame3 = 0;
const frame1 = window.requestAnimationFrame(() => {
frame2 = window.requestAnimationFrame(() => {
reactFlowInstance.fitView({ duration: 250 });
frame3 = window.requestAnimationFrame(() => {
reactFlowInstance.fitView({
duration: 320,
nodes: getFitNodeIdsIgnoringNotes(reactFlowInstance.getNodes()),
});
});
});
});
return () => {
@ -385,6 +539,9 @@ export function RecipeStudioPage({
if (frame2) {
window.cancelAnimationFrame(frame2);
}
if (frame3) {
window.cancelAnimationFrame(frame3);
}
};
}, [activeView, fitViewTick, reactFlowInstance]);
@ -407,10 +564,12 @@ export function RecipeStudioPage({
void persistRecipe();
}}
/>
<div className="h-[75vh] w-full rounded-t-none">
<div className="h-[75vh] w-full rounded-t-none" ref={flowContainerRef}>
{activeView === "editor" ? (
<ReactFlow<Node<RecipeNodeData | RecipeGraphAuxNodeData>, Edge>
onInit={setReactFlowInstance}
onDragOver={handleDragOver}
onDrop={handleDrop}
nodes={displayGraph.nodes}
edges={displayGraph.edges}
nodeTypes={NODE_TYPES}
@ -428,7 +587,7 @@ export function RecipeStudioPage({
nodesDraggable={interactive}
nodesConnectable={interactive}
elementsSelectable={interactive}
fitView={true}
fitView={false}
className="h-full w-full rounded-t-none"
>
<LayoutControls
@ -474,13 +633,13 @@ export function RecipeStudioPage({
onViewChange={setSheetView}
open={blockSheetOpen}
onOpenChange={setBlockSheetOpen}
onAddSampler={addSamplerNode}
onAddSeed={addSeedNode}
onAddLlm={addLlmNode}
onAddModelProvider={addModelProviderNode}
onAddModelConfig={addModelConfigNode}
onAddExpression={addExpressionNode}
onAddMarkdownNote={addMarkdownNoteNode}
onAddSampler={handleAddSamplerFromSheet}
onAddSeed={handleAddSeedFromSheet}
onAddLlm={handleAddLlmFromSheet}
onAddModelProvider={handleAddModelProviderFromSheet}
onAddModelConfig={handleAddModelConfigFromSheet}
onAddExpression={handleAddExpressionFromSheet}
onAddMarkdownNote={handleAddMarkdownNoteFromSheet}
onOpenProcessors={openProcessorsFromSheet}
copied={copied}
onCopy={copyRecipe}

61
studio/frontend/src/features/recipe-studio/stores/helpers/aux-sync.ts
View file

@ -1,61 +0,0 @@
import type { XYPosition } from "@xyflow/react";
export function syncPositionsRecord(
prev: Record<string, XYPosition>,
activeIds: string[],
defaults: Record<string, XYPosition>,
): Record<string, XYPosition> {
const next: Record<string, XYPosition> = {};
for (const id of activeIds) {
const existing = prev[id];
if (existing) {
next[id] = existing;
continue;
}
const fallback = defaults[id];
if (fallback) {
next[id] = fallback;
}
}
const prevIds = Object.keys(prev);
const nextIds = Object.keys(next);
if (prevIds.length !== nextIds.length) {
return next;
}
for (const id of nextIds) {
const a = prev[id];
const b = next[id];
if (!(a && b && a.x === b.x && a.y === b.y)) {
return next;
}
}
return prev;
}
export function syncSizesRecord(
prev: Record<string, { width: number; height: number }>,
activeIds: string[],
): Record<string, { width: number; height: number }> {
const active = new Set(activeIds);
const next: Record<string, { width: number; height: number }> = {};
for (const [id, size] of Object.entries(prev)) {
if (active.has(id)) {
next[id] = size;
}
}
const prevIds = Object.keys(prev);
const nextIds = Object.keys(next);
if (prevIds.length !== nextIds.length) {
return next;
}
for (const id of nextIds) {
const a = prev[id];
const b = next[id];
if (!(a && b && a.width === b.width && a.height === b.height)) {
return next;
}
}
return prev;
}

397
studio/frontend/src/features/recipe-studio/stores/helpers/model-infra-layout.ts Normal file
View file

@ -0,0 +1,397 @@
import type { Edge, XYPosition } from "@xyflow/react";
import { DEFAULT_NODE_HEIGHT, DEFAULT_NODE_WIDTH } from "../../constants";
import type { LayoutDirection, NodeConfig, RecipeNode } from "../../types";
import { HANDLE_IDS, normalizeRecipeHandleId } from "../../utils/handles";
import { readNodeHeight, readNodeWidth } from "../../utils/rf-node-dimensions";
type Rect = {
x: number;
y: number;
width: number;
height: number;
};
type Bounds = {
minX: number;
maxX: number;
minY: number;
maxY: number;
};
function toRect(node: RecipeNode): Rect {
return {
x: node.position.x,
y: node.position.y,
width: readNodeWidth(node) ?? DEFAULT_NODE_WIDTH,
height: readNodeHeight(node) ?? DEFAULT_NODE_HEIGHT,
};
}
function intersects(a: Rect, b: Rect, pad = 18): boolean {
return !(
a.x + a.width + pad <= b.x ||
b.x + b.width + pad <= a.x ||
a.y + a.height + pad <= b.y ||
b.y + b.height + pad <= a.y
);
}
function findNonOverlappingPosition(
preferred: XYPosition,
width: number,
height: number,
occupied: Rect[],
): XYPosition {
const step = 24;
for (let ring = 0; ring <= 16; ring += 1) {
for (let dx = -ring; dx <= ring; dx += 1) {
for (let dy = -ring; dy <= ring; dy += 1) {
if (ring > 0 && Math.max(Math.abs(dx), Math.abs(dy)) !== ring) {
continue;
}
const candidate = {
x: preferred.x + dx * step,
y: preferred.y + dy * step,
};
const rect = {
x: candidate.x,
y: candidate.y,
width,
height,
};
if (!occupied.some((item) => intersects(rect, item))) {
return candidate;
}
}
}
}
return preferred;
}
function isProviderToConfigEdge(edge: Edge, configs: Record<string, NodeConfig>): boolean {
const source = configs[edge.source];
const target = configs[edge.target];
return source?.kind === "model_provider" && target?.kind === "model_config";
}
function isConfigToLlmEdge(edge: Edge, configs: Record<string, NodeConfig>): boolean {
const source = configs[edge.source];
const target = configs[edge.target];
return source?.kind === "model_config" && target?.kind === "llm";
}
function usageKey(nodeId: string, handleId: string): string {
return `${nodeId}::${handleId}`;
}
function incrementUsage(map: Map<string, number>, nodeId: string, handleId: string): void {
const key = usageKey(nodeId, handleId);
map.set(key, (map.get(key) ?? 0) + 1);
}
function decrementUsage(map: Map<string, number>, nodeId: string, handleId: string): void {
const key = usageKey(nodeId, handleId);
map.set(key, Math.max(0, (map.get(key) ?? 0) - 1));
}
function getUsage(map: Map<string, number>, nodeId: string, handleId: string): number {
return map.get(usageKey(nodeId, handleId)) ?? 0;
}
function pickHandleByUsage(
candidates: string[],
nodeId: string,
usageMap: Map<string, number>,
): string {
const free = candidates.filter((handleId) => getUsage(usageMap, nodeId, handleId) === 0);
if (free.length > 0) {
return free[0];
}
let bestHandle = candidates[0];
let bestCount = Number.POSITIVE_INFINITY;
for (const handleId of candidates) {
const count = getUsage(usageMap, nodeId, handleId);
if (count < bestCount) {
bestHandle = handleId;
bestCount = count;
}
}
return bestHandle;
}
function applyEdgeWithHandles(
edge: Edge,
sourceHandle: string,
targetHandle: string,
sourceUsage: Map<string, number>,
targetUsage: Map<string, number>,
): Edge {
incrementUsage(sourceUsage, edge.source, sourceHandle);
incrementUsage(targetUsage, edge.target, targetHandle);
return { ...edge, sourceHandle, targetHandle, type: "semantic" };
}
function getNodeCenter(node: RecipeNode): { x: number; y: number } {
const width = readNodeWidth(node) ?? DEFAULT_NODE_WIDTH;
const height = readNodeHeight(node) ?? DEFAULT_NODE_HEIGHT;
return {
x: node.position.x + width / 2,
y: node.position.y + height / 2,
};
}
function collectBounds(ids: string[], nodesById: Map<string, RecipeNode>): Bounds | null {
const rects = ids
.map((id) => nodesById.get(id))
.flatMap((node) => (node ? [toRect(node)] : []));
if (rects.length === 0) {
return null;
}
return rects.reduce<Bounds>(
(acc, rect) => ({
minX: Math.min(acc.minX, rect.x),
maxX: Math.max(acc.maxX, rect.x + rect.width),
minY: Math.min(acc.minY, rect.y),
maxY: Math.max(acc.maxY, rect.y + rect.height),
}),
{
minX: rects[0].x,
maxX: rects[0].x + rects[0].width,
minY: rects[0].y,
maxY: rects[0].y + rects[0].height,
},
);
}
function sortPreferredLlmTargetHandles(
direction: LayoutDirection,
sourceNode: RecipeNode | undefined,
targetNode: RecipeNode | undefined,
): string[] {
const sourceCenter = sourceNode ? getNodeCenter(sourceNode) : { x: 0, y: 0 };
const targetCenter = targetNode ? getNodeCenter(targetNode) : { x: 0, y: 0 };
if (direction === "TB") {
const horizontalFirst =
sourceCenter.x <= targetCenter.x
? [HANDLE_IDS.dataIn, HANDLE_IDS.dataInRight]
: [HANDLE_IDS.dataInRight, HANDLE_IDS.dataIn];
return [...horizontalFirst, HANDLE_IDS.dataInTop, HANDLE_IDS.dataInBottom];
}
const verticalFirst =
sourceCenter.y <= targetCenter.y
? [HANDLE_IDS.dataInTop, HANDLE_IDS.dataInBottom]
: [HANDLE_IDS.dataInBottom, HANDLE_IDS.dataInTop];
return [...verticalFirst, HANDLE_IDS.dataIn, HANDLE_IDS.dataInRight];
}
function getProviderSourceHandleCandidates(direction: LayoutDirection): string[] {
return direction === "TB"
? [HANDLE_IDS.semanticOut, HANDLE_IDS.semanticOutBottom]
: [HANDLE_IDS.semanticOutBottom, HANDLE_IDS.semanticOut];
}
function getProviderTargetHandleCandidates(direction: LayoutDirection): string[] {
return direction === "TB"
? [HANDLE_IDS.semanticIn, HANDLE_IDS.semanticInTop]
: [HANDLE_IDS.semanticInTop, HANDLE_IDS.semanticIn];
}
function getConfigSourceHandleCandidates(direction: LayoutDirection): string[] {
return direction === "TB" ? [HANDLE_IDS.semanticOut] : [HANDLE_IDS.semanticOutBottom];
}
export function optimizeModelInfraEdgeHandles(
edges: Edge[],
nodes: RecipeNode[],
configs: Record<string, NodeConfig>,
direction: LayoutDirection,
): Edge[] {
const nodesById = new Map(nodes.map((node) => [node.id, node] as const));
const sourceUsage = new Map<string, number>();
const targetUsage = new Map<string, number>();
for (const edge of edges) {
const sourceHandle = normalizeRecipeHandleId(edge.sourceHandle);
const targetHandle = normalizeRecipeHandleId(edge.targetHandle);
if (sourceHandle) {
incrementUsage(sourceUsage, edge.source, sourceHandle);
}
if (targetHandle) {
incrementUsage(targetUsage, edge.target, targetHandle);
}
}
const nextEdges: Edge[] = [];
for (const edge of edges) {
const source = configs[edge.source];
const target = configs[edge.target];
if (!(source && target)) {
nextEdges.push(edge);
continue;
}
const sourceHandleBefore = normalizeRecipeHandleId(edge.sourceHandle);
const targetHandleBefore = normalizeRecipeHandleId(edge.targetHandle);
const isModelSemantic =
isProviderToConfigEdge(edge, configs) || isConfigToLlmEdge(edge, configs);
if (!isModelSemantic) {
nextEdges.push(edge);
continue;
}
if (sourceHandleBefore) {
decrementUsage(sourceUsage, edge.source, sourceHandleBefore);
}
if (targetHandleBefore) {
decrementUsage(targetUsage, edge.target, targetHandleBefore);
}
if (isProviderToConfigEdge(edge, configs)) {
const sourceCandidates = getProviderSourceHandleCandidates(direction);
const targetCandidates = getProviderTargetHandleCandidates(direction);
const sourceHandle = pickHandleByUsage(sourceCandidates, edge.source, sourceUsage);
const targetHandle = pickHandleByUsage(targetCandidates, edge.target, targetUsage);
nextEdges.push(
applyEdgeWithHandles(
edge,
sourceHandle,
targetHandle,
sourceUsage,
targetUsage,
),
);
continue;
}
const sourceCandidates = getConfigSourceHandleCandidates(direction);
const targetCandidates = sortPreferredLlmTargetHandles(
direction,
nodesById.get(edge.source),
nodesById.get(edge.target),
);
const sourceHandle = pickHandleByUsage(sourceCandidates, edge.source, sourceUsage);
const targetHandle = pickHandleByUsage(targetCandidates, edge.target, targetUsage);
nextEdges.push(
applyEdgeWithHandles(
edge,
sourceHandle,
targetHandle,
sourceUsage,
targetUsage,
),
);
}
return nextEdges;
}
export function centerModelInfraNodes(
nodes: RecipeNode[],
edges: Edge[],
configs: Record<string, NodeConfig>,
direction: LayoutDirection,
): RecipeNode[] {
const nodesById = new Map(nodes.map((node) => [node.id, node] as const));
const configToLlmIds = new Map<string, string[]>();
const providerToConfigIds = new Map<string, string[]>();
for (const edge of edges) {
if (isProviderToConfigEdge(edge, configs)) {
const entries = providerToConfigIds.get(edge.source) ?? [];
if (!entries.includes(edge.target)) {
entries.push(edge.target);
}
providerToConfigIds.set(edge.source, entries);
continue;
}
if (isConfigToLlmEdge(edge, configs)) {
const entries = configToLlmIds.get(edge.source) ?? [];
if (!entries.includes(edge.target)) {
entries.push(edge.target);
}
configToLlmIds.set(edge.source, entries);
}
}
const modelConfigIds = Object.values(configs)
.filter((config) => config.kind === "model_config" && nodesById.has(config.id))
.map((config) => config.id);
const modelProviderIds = Object.values(configs)
.filter((config) => config.kind === "model_provider" && nodesById.has(config.id))
.map((config) => config.id);
const occupiedById = new Map(nodes.map((node) => [node.id, toRect(node)] as const));
const clusterGap = 72;
const placeNode = (nodeId: string, preferred: XYPosition): void => {
const currentNode = nodesById.get(nodeId);
if (!currentNode) {
return;
}
const width = readNodeWidth(currentNode) ?? DEFAULT_NODE_WIDTH;
const height = readNodeHeight(currentNode) ?? DEFAULT_NODE_HEIGHT;
occupiedById.delete(nodeId);
const position = findNonOverlappingPosition(
preferred,
width,
height,
Array.from(occupiedById.values()),
);
const nextNode = { ...currentNode, position };
nodesById.set(nodeId, nextNode);
occupiedById.set(nodeId, {
x: position.x,
y: position.y,
width,
height,
});
};
for (const modelConfigId of modelConfigIds) {
const llmIds = configToLlmIds.get(modelConfigId) ?? [];
const targetBounds = collectBounds(llmIds, nodesById);
const modelConfigNode = nodesById.get(modelConfigId);
if (!(targetBounds && modelConfigNode)) {
continue;
}
const width = readNodeWidth(modelConfigNode) ?? DEFAULT_NODE_WIDTH;
const height = readNodeHeight(modelConfigNode) ?? DEFAULT_NODE_HEIGHT;
const preferred =
direction === "LR"
? {
x: (targetBounds.minX + targetBounds.maxX) / 2 - width / 2,
y: targetBounds.minY - height - clusterGap,
}
: {
x: targetBounds.minX - width - clusterGap,
y: (targetBounds.minY + targetBounds.maxY) / 2 - height / 2,
};
placeNode(modelConfigId, preferred);
}
for (const modelProviderId of modelProviderIds) {
const configIds = providerToConfigIds.get(modelProviderId) ?? [];
const targetBounds = collectBounds(configIds, nodesById);
const modelProviderNode = nodesById.get(modelProviderId);
if (!(targetBounds && modelProviderNode)) {
continue;
}
const width = readNodeWidth(modelProviderNode) ?? DEFAULT_NODE_WIDTH;
const height = readNodeHeight(modelProviderNode) ?? DEFAULT_NODE_HEIGHT;
const preferred =
direction === "LR"
? {
x: (targetBounds.minX + targetBounds.maxX) / 2 - width / 2,
y: targetBounds.minY - height - clusterGap,
}
: {
x: targetBounds.minX - width - clusterGap,
y: (targetBounds.minY + targetBounds.maxY) / 2 - height / 2,
};
placeNode(modelProviderId, preferred);
}
return nodes.map((node) => nodesById.get(node.id) ?? node);
}

9
studio/frontend/src/features/recipe-studio/stores/helpers/node-updates.ts
View file

@ -1,3 +1,4 @@
import type { XYPosition } from "@xyflow/react";
import { DEFAULT_NODE_WIDTH } from "../../constants";
import type {
RecipeNode,
@ -40,11 +41,13 @@ export function buildNodeUpdate(
state: NodeUpdateState,
config: NodeConfig,
layoutDirection: LayoutDirection,
position?: XYPosition,
openDialog = true,
): NodeUpdateResult {
const node: RecipeNode = {
id: config.id,
type: "builder",
position: { x: 0, y: state.nextY },
position: position ?? { x: 0, y: state.nextY },
data: nodeDataFromConfig(config, layoutDirection),
style: { width: DEFAULT_NODE_WIDTH },
selected: true,
@ -54,9 +57,9 @@ export function buildNodeUpdate(
configs: { ...state.configs, [config.id]: config },
nodes: [...state.nodes.map((item) => ({ ...item, selected: false })), node],
nextId: state.nextId + 1,
nextY: state.nextY + 140,
nextY: position ? state.nextY : state.nextY + 140,
activeConfigId: config.id,
dialogOpen: mode === "dialog",
dialogOpen: openDialog && mode === "dialog",
};
}

227
studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts
View file

@ -26,10 +26,17 @@ import {
} from "../blocks/registry";
import { deriveDisplayGraph } from "../utils/graph/derive-display-graph";
import { applyRecipeConnection, isValidRecipeConnection } from "../utils/graph";
import { HANDLE_IDS, remapRecipeEdgeHandlesForLayout } from "../utils/handles";
import {
HANDLE_IDS,
normalizeRecipeHandleId,
remapRecipeEdgeHandlesForLayout,
} from "../utils/handles";
import type { RecipeSnapshot } from "../utils/import";
import { getLayoutedElements } from "../utils/layout";
import { syncPositionsRecord, syncSizesRecord } from "./helpers/aux-sync";
import {
centerModelInfraNodes,
optimizeModelInfraEdgeHandles,
} from "./helpers/model-infra-layout";
import { applyEdgeRemovals, applyNodeRemovals } from "./helpers/removals";
import {
applyRenameToConfigs,
@ -53,7 +60,6 @@ type RecipeStudioState = {
nodes: RecipeNode[];
edges: Edge[];
auxNodePositions: Record<string, XYPosition>;
auxNodeSizes: Record<string, { width: number; height: number }>;
llmAuxVisibility: Record<string, boolean>;
configs: Record<string, NodeConfig>;
processors: RecipeProcessorConfig[];
@ -73,25 +79,24 @@ type RecipeStudioState = {
setLayoutDirection: (direction: LayoutDirection) => void;
applyLayout: () => void;
setLlmAuxVisibility: (id: string, visible: boolean) => void;
addSamplerNode: (type: SamplerType) => void;
addSeedNode: (type: SeedBlockType) => void;
addLlmNode: (type: LlmType) => void;
addModelProviderNode: () => void;
addModelConfigNode: () => void;
addExpressionNode: () => void;
addMarkdownNoteNode: () => void;
addSamplerNode: (
type: SamplerType,
position?: XYPosition,
openDialog?: boolean,
) => void;
addSeedNode: (
type: SeedBlockType,
position?: XYPosition,
openDialog?: boolean,
) => void;
addLlmNode: (type: LlmType, position?: XYPosition, openDialog?: boolean) => void;
addModelProviderNode: (position?: XYPosition, openDialog?: boolean) => void;
addModelConfigNode: (position?: XYPosition, openDialog?: boolean) => void;
addExpressionNode: (position?: XYPosition, openDialog?: boolean) => void;
addMarkdownNoteNode: (position?: XYPosition, openDialog?: boolean) => void;
updateConfig: (id: string, patch: Partial<NodeConfig>) => void;
loadRecipe: (snapshot: RecipeSnapshot) => void;
setAuxNodePosition: (id: string, position: XYPosition) => void;
setAuxNodeSize: (
id: string,
size: { width: number; height: number },
) => void;
syncAuxNodePositions: (
activeIds: string[],
defaults: Record<string, XYPosition>,
) => void;
syncAuxNodeSizes: (activeIds: string[]) => void;
onNodesChange: (changes: NodeChange<RecipeNode>[]) => void;
onEdgesChange: (changes: EdgeChange<Edge>[]) => void;
onConnect: (connection: Connection) => void;
@ -102,7 +107,6 @@ const INITIAL_STATE = {
nodes: [],
edges: [],
auxNodePositions: {},
auxNodeSizes: {},
llmAuxVisibility: {},
configs: {},
processors: [],
@ -118,7 +122,6 @@ const INITIAL_STATE = {
| "nodes"
| "edges"
| "auxNodePositions"
| "auxNodeSizes"
| "llmAuxVisibility"
| "configs"
| "processors"
@ -135,6 +138,8 @@ function buildAddedNodeState(
state: RecipeStudioState,
kind: BlockKind,
type: BlockType,
position?: XYPosition,
openDialog = true,
): Partial<RecipeStudioState> | RecipeStudioState {
const id = `n${state.nextId}`;
const existing = Object.values(state.configs);
@ -143,7 +148,13 @@ function buildAddedNodeState(
return state;
}
const config = definition.createConfig(id, existing);
return buildNodeUpdate(state, config, state.layoutDirection);
return buildNodeUpdate(
state,
config,
state.layoutDirection,
position,
openDialog,
);
}
function getAddedNodeContext(
@ -219,6 +230,17 @@ function connectSemantic(
};
}
function isModelSemanticEdge(edge: Edge, configs: Record<string, NodeConfig>): boolean {
const source = configs[edge.source];
const target = configs[edge.target];
return Boolean(
source &&
target &&
((source.kind === "model_provider" && target.kind === "model_config") ||
(source.kind === "model_config" && target.kind === "llm")),
);
}
export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
...INITIAL_STATE,
setSheetView: (view) => set({ sheetView: view }),
@ -230,10 +252,19 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
setLayoutDirection: (direction) =>
set((state) => ({
layoutDirection: direction,
edges: state.edges.map((edge) => ({
...edge,
...remapRecipeEdgeHandlesForLayout(edge, direction),
})),
edges: state.edges.map((edge) => {
if (isModelSemanticEdge(edge, state.configs)) {
return {
...edge,
sourceHandle: normalizeRecipeHandleId(edge.sourceHandle),
targetHandle: normalizeRecipeHandleId(edge.targetHandle),
};
}
return {
...edge,
...remapRecipeEdgeHandlesForLayout(edge, direction),
};
}),
nodes: applyLayoutDirectionToNodes(
state.nodes,
state.configs,
@ -243,13 +274,13 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
applyLayout: () =>
set((state) => {
const isTopBottom = state.layoutDirection === "TB";
const displayGraph = deriveDisplayGraph({
nodes: state.nodes,
edges: state.edges,
configs: state.configs,
layoutDirection: state.layoutDirection,
auxNodePositions: state.auxNodePositions,
auxNodeSizes: state.auxNodeSizes,
auxNodePositions: {},
llmAuxVisibility: state.llmAuxVisibility,
});
const { nodes } = getLayoutedElements(displayGraph.nodes, displayGraph.edges, {
@ -267,27 +298,23 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
}
return { ...node, position };
});
const nextAuxNodePositions: Record<string, XYPosition> = {};
for (const auxId of displayGraph.auxNodeIds) {
const existing = state.auxNodePositions[auxId];
const layouted = layoutedPositions.get(auxId);
if (layouted) {
nextAuxNodePositions[auxId] = layouted;
continue;
}
if (existing) {
nextAuxNodePositions[auxId] = existing;
continue;
}
const fallback = displayGraph.auxDefaults[auxId];
if (fallback) {
nextAuxNodePositions[auxId] = fallback;
}
}
const centeredNodes = centerModelInfraNodes(
nextNodes,
state.edges,
state.configs,
state.layoutDirection,
);
const optimizedEdges = optimizeModelInfraEdgeHandles(
state.edges,
centeredNodes,
state.configs,
state.layoutDirection,
);
return {
auxNodePositions: nextAuxNodePositions,
auxNodePositions: {},
edges: optimizedEdges,
nodes: applyLayoutDirectionToNodes(
nextNodes,
centeredNodes,
state.configs,
state.layoutDirection,
),
@ -305,15 +332,23 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
},
};
}),
addSamplerNode: (type) =>
set((state) => buildAddedNodeState(state, "sampler", type)),
addSeedNode: (type) =>
addSamplerNode: (type, position, openDialog = true) =>
set((state) =>
buildAddedNodeState(state, "sampler", type, position, openDialog),
),
addSeedNode: (type, position, openDialog = true) =>
set((state) => {
const existing = Object.values(state.configs).find(
(config) => config.kind === "seed",
);
if (!existing) {
return buildAddedNodeState(state, "seed", type);
return buildAddedNodeState(
state,
"seed",
type,
position,
openDialog,
);
}
let nextSourceType: SeedSourceType = "hf";
if (type === "seed_local") {
@ -351,13 +386,22 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
state.layoutDirection,
),
activeConfigId: existing.id,
dialogOpen: true,
dialogOpen: openDialog,
};
}),
addLlmNode: (type) => set((state) => buildAddedNodeState(state, "llm", type)),
addModelProviderNode: () =>
addLlmNode: (type, position, openDialog = true) =>
set((state) =>
buildAddedNodeState(state, "llm", type, position, openDialog),
),
addModelProviderNode: (position, openDialog = true) =>
set((state) => {
const added = buildAddedNodeState(state, "llm", "model_provider");
const added = buildAddedNodeState(
state,
"llm",
"model_provider",
position,
openDialog,
);
const context = getAddedNodeContext(added);
if (!context) {
return added;
@ -369,7 +413,7 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
config.kind === "model_config" &&
!config.provider.trim(),
);
if (unboundModelConfigs.length > 0) {
if (!position && unboundModelConfigs.length > 0) {
nodes = placeNodeNear(
nodes,
context.newNodeId,
@ -390,9 +434,15 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
}
return { ...added, nodes, edges, configs };
}),
addModelConfigNode: () =>
addModelConfigNode: (position, openDialog = true) =>
set((state) => {
const added = buildAddedNodeState(state, "llm", "model_config");
const added = buildAddedNodeState(
state,
"llm",
"model_config",
position,
openDialog,
);
const context = getAddedNodeContext(added);
if (!context) {
return added;
@ -405,7 +455,7 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
const unboundLlms = Object.values(configs).filter(
(config) => config.kind === "llm" && !config.model_alias.trim(),
);
if (providers.length === 1) {
if (!position && providers.length === 1) {
nodes = placeNodeNear(
nodes,
context.newNodeId,
@ -413,7 +463,7 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
state.layoutDirection,
"after",
);
} else if (unboundLlms.length > 0) {
} else if (!position && unboundLlms.length > 0) {
nodes = placeNodeNear(
nodes,
context.newNodeId,
@ -444,10 +494,26 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
}
return { ...added, nodes, edges, configs };
}),
addExpressionNode: () =>
set((state) => buildAddedNodeState(state, "expression", "expression")),
addMarkdownNoteNode: () =>
set((state) => buildAddedNodeState(state, "note", "markdown_note")),
addExpressionNode: (position, openDialog = true) =>
set((state) =>
buildAddedNodeState(
state,
"expression",
"expression",
position,
openDialog,
),
),
addMarkdownNoteNode: (position, openDialog = true) =>
set((state) =>
buildAddedNodeState(
state,
"note",
"markdown_note",
position,
openDialog,
),
),
loadRecipe: (snapshot) =>
set((state) => ({
configs: snapshot.configs,
@ -461,8 +527,7 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
layoutDirection: snapshot.layoutDirection,
nextId: snapshot.nextId,
nextY: snapshot.nextY,
auxNodePositions: {},
auxNodeSizes: {},
auxNodePositions: snapshot.auxNodePositions ?? {},
llmAuxVisibility: {},
activeConfigId: null,
dialogOpen: false,
@ -482,36 +547,6 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
},
};
}),
setAuxNodeSize: (id, size) =>
set((state) => {
const width = Math.max(1, size.width);
const height = Math.max(1, size.height);
const current = state.auxNodeSizes[id];
if (current && current.width === width && current.height === height) {
return state;
}
return {
auxNodeSizes: {
...state.auxNodeSizes,
[id]: { width, height },
},
};
}),
syncAuxNodePositions: (activeIds, defaults) =>
set((state) => {
const next = syncPositionsRecord(state.auxNodePositions, activeIds, defaults);
if (next === state.auxNodePositions) {
return state;
}
return {
auxNodePositions: next,
};
}),
syncAuxNodeSizes: (activeIds) =>
set((state) => {
const next = syncSizesRecord(state.auxNodeSizes, activeIds);
return next === state.auxNodeSizes ? state : { auxNodeSizes: next };
}),
updateConfig: (id, patch) => {
const applyUpdate = (state: RecipeStudioState) => {
const current = state.configs[id];

353
studio/frontend/src/features/recipe-studio/utils/graph/derive-display-graph.ts
View file

@ -7,7 +7,6 @@ import {
getDefaultDataTargetHandle,
getDefaultSemanticSourceHandle,
getDefaultSemanticTargetHandle,
getLlmJudgeScoreHandleId,
HANDLE_IDS,
isDataSourceHandle,
isDataTargetHandle,
@ -24,15 +23,12 @@ type DisplayGraphInput = {
configs: Record<string, NodeConfig>;
layoutDirection: LayoutDirection;
auxNodePositions: Record<string, XYPosition>;
auxNodeSizes: Record<string, { width: number; height: number }>;
llmAuxVisibility: Record<string, boolean>;
};
export type DisplayGraph = {
nodes: Array<Node<RecipeNode["data"] | RecipeGraphAuxNodeData>>;
edges: Edge[];
auxNodeIds: string[];
auxDefaults: Record<string, XYPosition>;
};
function normalizeEdge(
@ -99,7 +95,6 @@ function normalizeEdge(
type AuxNodeItem = {
key: string;
targetHandle: string;
data: RecipeGraphAuxNodeData;
};
@ -136,42 +131,235 @@ function findNonOverlappingPosition(
preferred: XYPosition,
width: number,
height: number,
direction: LayoutDirection,
occupied: Rect[],
): XYPosition {
const primaryStep =
direction === "TB"
? { x: 0, y: -(height + 24) }
: { x: -(width + 24), y: 0 };
const lateralUnit =
direction === "TB"
? { x: Math.max(48, Math.round(width * 0.3)), y: 0 }
: { x: 0, y: Math.max(40, Math.round(height * 0.35)) };
const lateralPattern = [0, 1, -1, 2, -2];
for (let ring = 0; ring <= 8; ring += 1) {
for (const lateral of lateralPattern) {
const candidate = {
x: preferred.x + primaryStep.x * ring + lateralUnit.x * lateral,
y: preferred.y + primaryStep.y * ring + lateralUnit.y * lateral,
};
const rect = toRect(candidate, width, height);
if (!occupied.some((other) => intersects(rect, other))) {
return candidate;
const step = 24;
for (let ring = 0; ring <= 10; ring += 1) {
for (let dx = -ring; dx <= ring; dx += 1) {
for (let dy = -ring; dy <= ring; dy += 1) {
if (ring > 0 && Math.max(Math.abs(dx), Math.abs(dy)) !== ring) {
continue;
}
const candidate = {
x: preferred.x + dx * step,
y: preferred.y + dy * step,
};
const rect = toRect(candidate, width, height);
if (!occupied.some((other) => intersects(rect, other))) {
return candidate;
}
}
}
}
return preferred;
}
type HandleSide = "left" | "right" | "top" | "bottom";
const SIDE_TO_TARGET_HANDLE: Record<HandleSide, string> = {
left: HANDLE_IDS.dataIn,
right: HANDLE_IDS.dataInRight,
top: HANDLE_IDS.dataInTop,
bottom: HANDLE_IDS.dataInBottom,
};
function getTargetSide(
handleId: string | null | undefined,
direction: LayoutDirection,
): HandleSide {
const normalized = normalizeRecipeHandleId(handleId);
if (!normalized) {
return direction === "TB" ? "top" : "left";
}
if (
normalized === HANDLE_IDS.dataInRight ||
normalized === HANDLE_IDS.semanticInRight
) {
return "right";
}
if (
normalized === HANDLE_IDS.dataInBottom ||
normalized === HANDLE_IDS.semanticInBottom
) {
return "bottom";
}
if (
normalized === HANDLE_IDS.dataInTop ||
normalized === HANDLE_IDS.semanticInTop
) {
return "top";
}
return "left";
}
function getSourceSide(
handleId: string | null | undefined,
direction: LayoutDirection,
): HandleSide {
const normalized = normalizeRecipeHandleId(handleId);
if (!normalized) {
return direction === "TB" ? "bottom" : "right";
}
if (
normalized === HANDLE_IDS.dataOutLeft ||
normalized === HANDLE_IDS.semanticOutLeft
) {
return "left";
}
if (
normalized === HANDLE_IDS.dataOutTop ||
normalized === HANDLE_IDS.semanticOutTop
) {
return "top";
}
if (
normalized === HANDLE_IDS.dataOutBottom ||
normalized === HANDLE_IDS.semanticOutBottom
) {
return "bottom";
}
return "right";
}
function pickAuxTargetHandle(
llmId: string,
direction: LayoutDirection,
edges: Edge[],
): string {
const occupied = new Set<HandleSide>();
for (const edge of edges) {
if (edge.source.startsWith("aux-") || edge.target.startsWith("aux-")) {
continue;
}
if (edge.target === llmId) {
occupied.add(getTargetSide(edge.targetHandle, direction));
}
if (edge.source === llmId) {
occupied.add(getSourceSide(edge.sourceHandle, direction));
}
}
const priority: HandleSide[] =
direction === "LR"
? ["left", "right", "bottom", "top"]
: ["top", "bottom", "right", "left"];
for (const side of priority) {
if (!occupied.has(side)) {
return SIDE_TO_TARGET_HANDLE[side];
}
}
const fallback: HandleSide = direction === "LR" ? "bottom" : "right";
return SIDE_TO_TARGET_HANDLE[fallback];
}
function getHandleSideFromTargetHandle(targetHandle: string): HandleSide {
if (targetHandle === HANDLE_IDS.dataInRight) {
return "right";
}
if (targetHandle === HANDLE_IDS.dataInTop) {
return "top";
}
if (targetHandle === HANDLE_IDS.dataInBottom) {
return "bottom";
}
return "left";
}
function pickAuxSourceHandle(
auxPosition: XYPosition,
auxWidth: number,
auxHeight: number,
llmPosition: XYPosition,
llmWidth: number,
llmHeight: number,
): string {
const auxCenter = {
x: auxPosition.x + auxWidth / 2,
y: auxPosition.y + auxHeight / 2,
};
const llmCenter = {
x: llmPosition.x + llmWidth / 2,
y: llmPosition.y + llmHeight / 2,
};
const dx = llmCenter.x - auxCenter.x;
const dy = llmCenter.y - auxCenter.y;
if (Math.abs(dx) >= Math.abs(dy)) {
return dx >= 0 ? HANDLE_IDS.llmInputOutRight : HANDLE_IDS.llmInputOutLeft;
}
return dy >= 0 ? HANDLE_IDS.llmInputOutBottom : HANDLE_IDS.llmInputOutTop;
}
type AppendAuxNodeAndEdgeInput = {
auxNodes: Node<RecipeGraphAuxNodeData>[];
auxEdges: Edge[];
entry: {
item: AuxNodeItem;
auxId: string;
width: number;
height: number;
};
position: XYPosition;
parentNode: Node<RecipeNode["data"] | RecipeGraphAuxNodeData>;
parentWidth: number;
parentHeight: number;
auxTargetHandle: string;
};
function appendAuxNodeAndEdge({
auxNodes,
auxEdges,
entry,
position,
parentNode,
parentWidth,
parentHeight,
auxTargetHandle,
}: AppendAuxNodeAndEdgeInput): void {
auxNodes.push({
id: entry.auxId,
type: "aux",
data: entry.item.data,
position,
width: entry.width,
height: entry.height,
style: {
width: entry.width,
height: entry.height,
},
draggable: true,
selectable: true,
focusable: true,
connectable: false,
});
auxEdges.push({
id: `e-${entry.auxId}-${parentNode.id}`,
source: entry.auxId,
sourceHandle: pickAuxSourceHandle(
position,
entry.width,
entry.height,
parentNode.position,
parentWidth,
parentHeight,
),
target: parentNode.id,
targetHandle: auxTargetHandle,
type: "canvas",
data: { path: "auto" },
selectable: false,
focusable: false,
});
}
export function deriveDisplayGraph({
nodes,
edges,
configs,
layoutDirection,
auxNodePositions,
auxNodeSizes,
llmAuxVisibility,
}: DisplayGraphInput): DisplayGraph {
const displayNodes = nodes.map((node) => {
@ -190,8 +378,6 @@ export function deriveDisplayGraph({
});
const auxNodes: Node<RecipeGraphAuxNodeData>[] = [];
const auxEdges: Edge[] = [];
const auxDefaults: Record<string, XYPosition> = {};
const auxNodeIds: string[] = [];
const occupiedRects: Rect[] = displayNodes.map((node) =>
toRect(
node.position,
@ -209,18 +395,18 @@ export function deriveDisplayGraph({
continue;
}
const llmDirection = node.data.layoutDirection ?? layoutDirection;
const auxTargetHandle = pickAuxTargetHandle(node.id, llmDirection, edges);
const auxTargetSide = getHandleSideFromTargetHandle(auxTargetHandle);
const items: AuxNodeItem[] = [];
if (config.system_prompt.trim()) {
items.push({
key: "system",
targetHandle: HANDLE_IDS.llmSystemIn,
data: {
kind: "llm-prompt-input",
llmId: config.id,
field: "system_prompt",
title: "System Prompt",
layoutDirection: llmDirection,
},
});
}
@ -228,13 +414,11 @@ export function deriveDisplayGraph({
if (config.prompt.trim()) {
items.push({
key: "prompt",
targetHandle: HANDLE_IDS.llmPromptIn,
data: {
kind: "llm-prompt-input",
llmId: config.id,
field: "prompt",
title: "Prompt",
layoutDirection: llmDirection,
},
});
}
@ -243,12 +427,10 @@ export function deriveDisplayGraph({
(config.scores ?? []).forEach((_score, scoreIndex) => {
items.push({
key: `score-${scoreIndex}`,
targetHandle: getLlmJudgeScoreHandleId(scoreIndex),
data: {
kind: "llm-judge-score",
llmId: config.id,
scoreIndex,
layoutDirection: llmDirection,
},
});
});
@ -262,19 +444,20 @@ export function deriveDisplayGraph({
const parentHeight = readNodeHeight(node) ?? DEFAULT_NODE_HEIGHT;
const itemsWithLayout = items.map((item) => {
const auxId = `aux-${node.id}-${item.key}`;
const savedSize = auxNodeSizes[auxId];
return {
item,
auxId,
width: savedSize?.width ?? DEFAULT_NODE_WIDTH,
height: savedSize?.height ?? DEFAULT_NODE_HEIGHT,
width: DEFAULT_NODE_WIDTH,
height: DEFAULT_NODE_HEIGHT,
};
});
const gap = 24;
const sideOffset = 48;
const stackHorizontal =
auxTargetSide === "top" || auxTargetSide === "bottom";
if (llmDirection === "TB") {
if (stackHorizontal) {
const totalWidth =
itemsWithLayout.reduce((sum, entry) => sum + entry.width, 0) +
(itemsWithLayout.length - 1) * gap;
@ -284,51 +467,30 @@ export function deriveDisplayGraph({
for (const entry of itemsWithLayout) {
const preferredPosition = {
x: xCursor,
y: node.position.y - entry.height - sideOffset,
y:
auxTargetSide === "top"
? node.position.y - entry.height - sideOffset
: node.position.y + parentHeight + sideOffset,
};
const defaultPosition = findNonOverlappingPosition(
preferredPosition,
entry.width,
entry.height,
llmDirection,
occupiedRects,
);
const position = auxNodePositions[entry.auxId] ?? defaultPosition;
xCursor += entry.width + gap;
auxNodeIds.push(entry.auxId);
if (!auxNodePositions[entry.auxId]) {
auxDefaults[entry.auxId] = defaultPosition;
}
occupiedRects.push(toRect(position, entry.width, entry.height));
auxNodes.push({
id: entry.auxId,
type: "aux",
data: entry.item.data,
appendAuxNodeAndEdge({
auxNodes,
auxEdges,
entry,
position,
width: entry.width,
height: entry.height,
style: {
width: entry.width,
height: entry.height,
},
draggable: true,
selectable: true,
focusable: true,
connectable: true,
});
auxEdges.push({
id: `e-${entry.auxId}-${node.id}`,
source: entry.auxId,
sourceHandle: HANDLE_IDS.llmInputOut,
target: node.id,
targetHandle: entry.item.targetHandle,
type: "canvas",
data: { path: "auto" },
selectable: false,
focusable: false,
parentNode: node,
parentWidth,
parentHeight,
auxTargetHandle,
});
}
continue;
@ -338,7 +500,10 @@ export function deriveDisplayGraph({
itemsWithLayout.reduce((sum, entry) => sum + entry.height, 0) +
(itemsWithLayout.length - 1) * gap;
const maxWidth = Math.max(...itemsWithLayout.map((entry) => entry.width));
const baseX = node.position.x - maxWidth - sideOffset;
const baseX =
auxTargetSide === "right"
? node.position.x + parentWidth + sideOffset
: node.position.x - maxWidth - sideOffset;
let yCursor = node.position.y + (parentHeight - totalHeight) / 2;
for (const entry of itemsWithLayout) {
@ -350,45 +515,21 @@ export function deriveDisplayGraph({
preferredPosition,
entry.width,
entry.height,
llmDirection,
occupiedRects,
);
const position = auxNodePositions[entry.auxId] ?? defaultPosition;
yCursor += entry.height + gap;
auxNodeIds.push(entry.auxId);
if (!auxNodePositions[entry.auxId]) {
auxDefaults[entry.auxId] = defaultPosition;
}
occupiedRects.push(toRect(position, entry.width, entry.height));
auxNodes.push({
id: entry.auxId,
type: "aux",
data: entry.item.data,
appendAuxNodeAndEdge({
auxNodes,
auxEdges,
entry,
position,
width: entry.width,
height: entry.height,
style: {
width: entry.width,
height: entry.height,
},
draggable: true,
selectable: true,
focusable: true,
connectable: true,
});
auxEdges.push({
id: `e-${entry.auxId}-${node.id}`,
source: entry.auxId,
sourceHandle: HANDLE_IDS.llmInputOut,
target: node.id,
targetHandle: entry.item.targetHandle,
type: "canvas",
data: { path: "auto" },
selectable: false,
focusable: false,
parentNode: node,
parentWidth,
parentHeight,
auxTargetHandle,
});
}
}
@ -398,7 +539,5 @@ export function deriveDisplayGraph({
edges: [...edges, ...auxEdges].map((edge) =>
normalizeEdge(edge, configs, layoutDirection),
),
auxNodeIds,
auxDefaults,
};
}

17
studio/frontend/src/features/recipe-studio/utils/graph/fit-view.ts Normal file
View file

@ -0,0 +1,17 @@
import type { Node } from "@xyflow/react";
function isMarkdownNoteNode(node: Node): boolean {
if (node.type !== "builder") {
return false;
}
if (!node.data || typeof node.data !== "object") {
return false;
}
return (node.data as { kind?: string }).kind === "note";
}
export function getFitNodeIdsIgnoringNotes(nodes: Node[]): Array<{ id: string }> {
const nodesWithoutNotes = nodes.filter((node) => !isMarkdownNoteNode(node));
const targetNodes = nodesWithoutNotes.length > 0 ? nodesWithoutNotes : nodes;
return targetNodes.map((node) => ({ id: node.id }));
}

101
studio/frontend/src/features/recipe-studio/utils/graph/recipe-graph-connection.ts
View file

@ -1,10 +1,12 @@
import { type Connection, type Edge, addEdge } from "@xyflow/react";
import type { NodeConfig, SamplerConfig } from "../../types";
import {
HANDLE_IDS,
isDataSourceHandle,
isDataTargetHandle,
isSemanticSourceHandle,
isSemanticTargetHandle,
normalizeRecipeHandleId,
} from "../handles";
import { isSemanticRelation } from "./relations";
import {
@ -127,6 +129,97 @@ function isCompetingIncomingEdge(
return source.kind === "sampler" && source.sampler_type === "datetime";
}
function isModelSemanticRelation(source: NodeConfig, target: NodeConfig): boolean {
return (
(source.kind === "model_provider" && target.kind === "model_config") ||
(source.kind === "model_config" && target.kind === "llm")
);
}
function countHandleUsage(
edges: Edge[],
nodeId: string,
handleId: string,
lane: "source" | "target",
): number {
return edges.reduce((count, edge) => {
const edgeNodeId = lane === "source" ? edge.source : edge.target;
if (edgeNodeId !== nodeId) {
return count;
}
const edgeHandleId =
lane === "source"
? normalizeRecipeHandleId(edge.sourceHandle)
: normalizeRecipeHandleId(edge.targetHandle);
return edgeHandleId === handleId ? count + 1 : count;
}, 0);
}
function pickLeastUsedHandle(
candidates: string[],
requested: string | null,
usageFor: (handleId: string) => number,
): string {
let bestHandle = candidates[0];
let bestCount = Number.POSITIVE_INFINITY;
const requestedNormalized = requested
? normalizeRecipeHandleId(requested)
: null;
for (const candidate of candidates) {
const usage = usageFor(candidate);
if (usage < bestCount) {
bestHandle = candidate;
bestCount = usage;
continue;
}
if (usage === bestCount && requestedNormalized === candidate) {
bestHandle = candidate;
}
}
return bestHandle;
}
function chooseModelSemanticHandles(
connection: Connection,
source: NodeConfig,
target: NodeConfig,
edges: Edge[],
): Connection {
if (!isModelSemanticRelation(source, target)) {
return connection;
}
const sourceCandidates = [HANDLE_IDS.semanticOut, HANDLE_IDS.semanticOutBottom];
const targetCandidates =
target.kind === "model_config"
? [HANDLE_IDS.semanticIn, HANDLE_IDS.semanticInTop]
: [
HANDLE_IDS.dataIn,
HANDLE_IDS.dataInTop,
HANDLE_IDS.dataInRight,
HANDLE_IDS.dataInBottom,
];
const sourceHandle = pickLeastUsedHandle(
sourceCandidates,
connection.sourceHandle ?? null,
(handleId) => countHandleUsage(edges, source.id, handleId, "source"),
);
const targetHandle = pickLeastUsedHandle(
targetCandidates,
connection.targetHandle ?? null,
(handleId) => countHandleUsage(edges, target.id, handleId, "target"),
);
return {
...connection,
sourceHandle,
targetHandle,
};
}
export function isValidRecipeConnection(
connection: Connection,
configs: Record<string, NodeConfig>,
@ -177,8 +270,14 @@ export function applyRecipeConnection(
!isCompetingIncomingEdge(edge, target.id, singleRefRelation, configs),
)
: edges;
const resolvedConnection = chooseModelSemanticHandles(
connection,
source,
target,
nextBaseEdges,
);
const nextEdges = addEdge(
{ ...connection, type: semanticRelation ? "semantic" : "canvas" },
{ ...resolvedConnection, type: semanticRelation ? "semantic" : "canvas" },
nextBaseEdges,
);
if (source.kind === "model_provider" && target.kind === "model_config") {

31
studio/frontend/src/features/recipe-studio/utils/handle-layout.ts
View file

@ -1,36 +1,5 @@
import { Position } from "@xyflow/react";
import type { LayoutDirection } from "../types";
export const NODE_HANDLE_CLASS =
"pointer-events-auto !size-2.5 !border-border/80 !bg-muted shadow-sm hover:!border-primary/70 hover:!bg-primary/20";
export const AUX_HANDLE_CLASS =
"!size-2 !border-border/80 !bg-muted/80 shadow-sm";
export type NodeHandleLayout = {
isTopBottom: boolean;
dataInPosition: Position;
dataOutPosition: Position;
semanticInPosition: Position;
semanticOutPosition: Position;
};
export function getNodeHandleLayout(
direction: LayoutDirection,
): NodeHandleLayout {
const isTopBottom = direction === "TB";
return {
isTopBottom,
dataInPosition: isTopBottom ? Position.Top : Position.Left,
dataOutPosition: isTopBottom ? Position.Bottom : Position.Right,
semanticInPosition: isTopBottom ? Position.Left : Position.Top,
semanticOutPosition: isTopBottom ? Position.Right : Position.Bottom,
};
}
export function getAuxSourceHandlePosition(
direction: LayoutDirection,
): Position {
return direction === "TB" ? Position.Bottom : Position.Right;
}

11
studio/frontend/src/features/recipe-studio/utils/handles.ts
View file

@ -23,17 +23,14 @@ export const HANDLE_IDS = {
semanticOutBottom: "semantic-out-bottom",
semanticOutRight: "semantic-out-right",
// llm prompt/scorer lanes
llmPromptIn: "llm-prompt-in",
llmSystemIn: "llm-system-in",
llmInputOut: "llm-input-out",
llmInputOutLeft: "llm-input-out-left",
llmInputOutRight: "llm-input-out-right",
llmInputOutTop: "llm-input-out-top",
llmInputOutBottom: "llm-input-out-bottom",
} as const;
export type RecipeHandleId = (typeof HANDLE_IDS)[keyof typeof HANDLE_IDS];
export function getLlmJudgeScoreHandleId(index: number): string {
return `llm-judge-score-in-${index}`;
}
const LEGACY_HANDLE_ALIAS_MAP: Record<string, string> = {
[HANDLE_IDS.semanticInLeft]: HANDLE_IDS.semanticIn,
[HANDLE_IDS.semanticOutRight]: HANDLE_IDS.semanticOut,

12
studio/frontend/src/features/recipe-studio/utils/import/importer.ts
View file

@ -453,7 +453,7 @@ export function importRecipePayload(input: string): ImportResult {
return { errors, snapshot: null };
}
const { layouts, edges: uiEdges, layoutDirection } = parseUi(ui);
const { layouts, auxNodes, edges: uiEdges, layoutDirection } = parseUi(ui);
const resolvedLayoutDirection = layoutDirection ?? "LR";
const nodes = buildNodes(configs, layouts);
const edges = buildEdges(
@ -462,6 +462,15 @@ export function importRecipePayload(input: string): ImportResult {
uiEdges,
resolvedLayoutDirection,
);
const auxNodePositions = Object.fromEntries(
auxNodes.flatMap((item) => {
const llmId = nameToId.get(item.llm);
if (!llmId) {
return [];
}
return [[`aux-${llmId}-${item.key}`, { x: item.x, y: item.y }]];
}),
);
const maxY = nodes.reduce(
(acc, node) => Math.max(acc, node.position.y),
@ -474,6 +483,7 @@ export function importRecipePayload(input: string): ImportResult {
configs: Object.fromEntries(configs.map((config) => [config.id, config])),
nodes,
edges,
auxNodePositions,
processors,
layoutDirection: resolvedLayoutDirection,
nextId,

3
studio/frontend/src/features/recipe-studio/utils/import/types.ts
View file

@ -1,4 +1,4 @@
import type { Edge } from "@xyflow/react";
import type { Edge, XYPosition } from "@xyflow/react";
import type {
LayoutDirection,
RecipeNode,
@ -10,6 +10,7 @@ export type RecipeSnapshot = {
configs: Record<string, NodeConfig>;
nodes: RecipeNode[];
edges: Edge[];
auxNodePositions: Record<string, XYPosition>;
processors: RecipeProcessorConfig[];
layoutDirection: LayoutDirection;
nextId: number;

32
studio/frontend/src/features/recipe-studio/utils/import/ui.ts
View file

@ -7,14 +7,23 @@ import { isRecord, readString } from "./helpers";
type UiInput = {
nodes?: unknown;
edges?: unknown;
aux_nodes?: unknown;
layout_direction?: unknown;
layoutDirection?: unknown;
};
type ParsedAuxNode = {
llm: string;
key: string;
x: number;
y: number;
};
export function parseUi(
ui: UiInput | null,
): {
layouts: Map<string, { x: number; y: number; width?: number }>;
auxNodes: ParsedAuxNode[];
edges: Array<{
from: string;
to: string;
@ -25,6 +34,7 @@ export function parseUi(
layoutDirection: "LR" | "TB" | null;
} {
const layouts = new Map<string, { x: number; y: number; width?: number }>();
const auxNodes: ParsedAuxNode[] = [];
const edges: Array<{
from: string;
to: string;
@ -72,6 +82,21 @@ export function parseUi(
}
}
}
if (ui && Array.isArray(ui.aux_nodes)) {
for (const node of ui.aux_nodes) {
if (!isRecord(node)) {
continue;
}
const llm = readString(node.llm);
const key = readString(node.key);
const x = typeof node.x === "number" ? node.x : null;
const y = typeof node.y === "number" ? node.y : null;
if (!(llm && key && x !== null && y !== null)) {
continue;
}
auxNodes.push({ llm, key, x, y });
}
}
const layoutDirectionRaw =
readString(ui?.layout_direction) ?? readString(ui?.layoutDirection);
const layoutDirection =
@ -81,7 +106,12 @@ export function parseUi(
? "LR"
: null;
return { layouts, edges: edges.length > 0 ? edges : null, layoutDirection };
return {
layouts,
auxNodes,
edges: edges.length > 0 ? edges : null,
layoutDirection,
};
}
export function buildNodes(

14
studio/frontend/src/features/recipe-studio/utils/layout.ts
View file

@ -1,6 +1,8 @@
import dagre from "@dagrejs/dagre";
import type { Edge, Node } from "@xyflow/react";
import { DEFAULT_NODE_HEIGHT, DEFAULT_NODE_WIDTH } from "../constants";
import type { LayoutDirection } from "../types";
import { readNodeHeight, readNodeWidth } from "./rf-node-dimensions";
type LayoutOptions = {
direction?: LayoutDirection;
@ -21,8 +23,8 @@ export function getLayoutedElements<TNode extends Node>(
nodesep = 80,
ranksep = 80,
edgesep = 28,
nodeWidth = 220,
nodeHeight = 64,
nodeWidth = DEFAULT_NODE_WIDTH,
nodeHeight = DEFAULT_NODE_HEIGHT,
} = options;
const graph = new dagre.graphlib.Graph();
@ -36,8 +38,8 @@ export function getLayoutedElements<TNode extends Node>(
});
nodes.forEach((node) => {
const width = node.measured?.width ?? nodeWidth;
const height = node.measured?.height ?? nodeHeight;
const width = readNodeWidth(node) ?? nodeWidth;
const height = readNodeHeight(node) ?? nodeHeight;
graph.setNode(node.id, { width, height });
});
@ -54,8 +56,8 @@ export function getLayoutedElements<TNode extends Node>(
const layoutedNodes = nodes.map((node) => {
const pos = graph.node(node.id);
const width = node.measured?.width ?? nodeWidth;
const height = node.measured?.height ?? nodeHeight;
const width = readNodeWidth(node) ?? nodeWidth;
const height = readNodeHeight(node) ?? nodeHeight;
return {
...node,
position: {

12
studio/frontend/src/features/recipe-studio/utils/node-data.ts
View file

@ -41,18 +41,18 @@ export function nodeDataFromConfig(
}
if (config.kind === "seed") {
const seedSourceType = config.seed_source_type ?? "hf";
const subtype =
const sourceLabel =
seedSourceType === "hf"
? "Hugging Face"
? "Hugging Face dataset"
: seedSourceType === "local"
? "Local File"
: "Unstructured";
? "Structured file"
: "Unstructured document";
return {
title: "Seed",
kind: "seed",
subtype,
subtype: sourceLabel,
blockType: "seed",
name: config.name,
name: sourceLabel,
layoutDirection,
};
}

25
studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts
View file

@ -1,4 +1,4 @@
import type { Edge } from "@xyflow/react";
import type { Edge, XYPosition } from "@xyflow/react";
import type {
LayoutDirection,
ModelConfig,
@ -70,6 +70,7 @@ export function buildRecipePayload(
edges: Edge[],
processors: RecipeProcessorConfig[] = [],
layoutDirection: LayoutDirection = "LR",
auxNodePositions: Record<string, XYPosition> = {},
): RecipePayloadResult {
const errors: string[] = [];
const columns: Record<string, unknown>[] = [];
@ -270,6 +271,27 @@ export function buildRecipePayload(
},
];
});
const uiAuxNodes = Object.entries(auxNodePositions).flatMap(
([auxId, position]) => {
const match = /^aux-([^-]+)-(.+)$/.exec(auxId);
if (!match) {
return [];
}
const [, llmId, key] = match;
const llmConfig = configs[llmId];
if (!(llmConfig && llmConfig.kind === "llm")) {
return [];
}
return [
{
llm: llmConfig.name,
key,
x: position.x,
y: position.y,
},
];
},
);
const recipeProcessors = buildProcessors(processors, errors);
const seedConfig = firstSeed ? buildSeedConfig(firstSeed, errors) : undefined;
const seedDropProcessor = firstSeed
@ -306,6 +328,7 @@ export function buildRecipePayload(
nodes: uiNodes,
edges: uiEdges,
layout_direction: layoutDirection,
...(uiAuxNodes.length > 0 && { aux_nodes: uiAuxNodes }),
...(firstSeed && { seed_source_type: firstSeed.seed_source_type }),
...(firstSeed && { seed_columns: firstSeed.seed_columns ?? [] }),
...(firstSeed && {

2
studio/frontend/src/features/recipe-studio/utils/payload/builders-model.ts
View file

@ -19,7 +19,7 @@ export function buildModelProvider(
name: config.name,
endpoint: config.endpoint,
// biome-ignore lint/style/useNamingConvention: api schema
provider_type: config.provider_type,
provider_type: "openai",
// biome-ignore lint/style/useNamingConvention: api schema
api_key_env: config.api_key_env?.trim() || undefined,
// biome-ignore lint/style/useNamingConvention: api schema

3
studio/frontend/src/features/recipe-studio/utils/payload/builders-seed.ts
View file

@ -14,9 +14,6 @@ export function buildSeedConfig(
): Record<string, unknown> | undefined {
const seedSourceType = config.seed_source_type ?? "hf";
const path = config.hf_path.trim();
if (!path) {
return undefined;
}
const endpoint = config.hf_endpoint?.trim() || "https://huggingface.co";
const token = config.hf_token?.trim() || null;

7
studio/frontend/src/features/recipe-studio/utils/payload/types.ts
View file

@ -52,6 +52,13 @@ export type RecipePayload = {
layout_direction?: "LR" | "TB";
// ui-only, used to preserve seed block mode across imports/refresh
seed_source_type?: "hf" | "local" | "unstructured";
// ui-only, persisted aux node positions by llm name + aux key
aux_nodes?: Array<{
llm: string;
key: string;
x: number;
y: number;
}>;
// ui-only, seed metadata cached for refresh/import UX
seed_columns?: string[];
seed_drop_columns?: string[];

33
studio/frontend/src/features/recipe-studio/utils/reactflow-changes.ts
View file

@ -5,43 +5,25 @@ import type {
NodeChange,
XYPosition,
} from "@xyflow/react";
import type { RecipeGraphAuxNodeData } from "../components/recipe-graph-aux-node";
import type { RecipeNodeData } from "../types";
type AnyNode = Node<RecipeNodeData | RecipeGraphAuxNodeData>;
export function applyAuxNodeChanges(
changes: NodeChange<AnyNode>[],
export function applyAuxNodeChanges<T extends Node>(
changes: NodeChange<T>[],
actions: {
setAuxNodePosition: (id: string, position: XYPosition) => void;
setAuxNodeSize: (
id: string,
size: { width: number; height: number },
) => void;
},
): void {
for (const change of changes) {
if (!("id" in change) || !change.id.startsWith("aux-")) {
continue;
}
if (change.type === "position") {
const nextPosition = change.position ?? change.positionAbsolute;
if (nextPosition) {
actions.setAuxNodePosition(change.id, nextPosition);
}
if (change.type !== "position") {
continue;
}
if (
change.type === "dimensions" &&
change.dimensions &&
change.dimensions.width > 0 &&
change.dimensions.height > 0
) {
actions.setAuxNodeSize(change.id, {
width: change.dimensions.width,
height: change.dimensions.height,
});
const nextPosition = change.position ?? change.positionAbsolute;
if (!nextPosition) {
continue;
}
actions.setAuxNodePosition(change.id, nextPosition);
}
}
@ -62,4 +44,3 @@ export function filterEdgeChangesByIds(
(change): change is EdgeChange<Edge> => "id" in change && ids.has(change.id),
);
}

34
studio/frontend/src/features/recipe-studio/utils/refs.ts
View file

@ -1,4 +1,7 @@
const JINJA_REF_RE = /{{\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*}}/g;
const JINJA_EXPR_RE = /{{\s*([^{}]+?)\s*}}/g;
const SIMPLE_JINJA_EXPR_RE = /^[a-zA-Z_][a-zA-Z0-9_.]*$/;
const PLAIN_JINJA_EXPR_RE = /^[a-zA-Z0-9_.\s-]+$/;
function escapeRegExp(value: string): string {
return value.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
@ -17,6 +20,37 @@ export function extractRefs(template: string): string[] {
return Array.from(refs);
}
export function findInvalidJinjaReferences(
template: string,
validReferences: string[],
): string[] {
if (!template) {
return [];
}
const validSet = new Set(
validReferences.map((name) => name.trim()).filter(Boolean),
);
const invalid = new Set<string>();
for (const match of template.matchAll(JINJA_EXPR_RE)) {
const expr = (match[1] ?? "").trim();
if (!expr) {
continue;
}
if (SIMPLE_JINJA_EXPR_RE.test(expr)) {
if (!validSet.has(expr)) {
invalid.add(expr);
}
continue;
}
if (PLAIN_JINJA_EXPR_RE.test(expr)) {
invalid.add(expr);
}
}
return Array.from(invalid);
}
export function replaceRef(
template: string,
from: string,

24
studio/frontend/src/features/studio/sections/dataset-section.tsx
View file

@ -33,7 +33,6 @@ import {
useHfTokenValidation,
useInfiniteScroll,
} from "@/hooks";
import { formatCompact } from "@/lib/utils";
import {
HfDatasetSubsetSplitSelectors,
useDatasetPreviewDialogStore,
@ -75,6 +74,7 @@ export function DatasetSection() {
datasetEvalSplit,
setDatasetEvalSplit,
hfToken,
modelType,
} = useTrainingConfigStore(
useShallow((s) => ({
dataset: s.dataset,
@ -88,6 +88,7 @@ export function DatasetSection() {
datasetEvalSplit: s.datasetEvalSplit,
setDatasetEvalSplit: s.setDatasetEvalSplit,
hfToken: s.hfToken,
modelType: s.modelType,
})),
);
@ -116,6 +117,7 @@ export function DatasetSection() {
fetchMore,
error: hfSearchError,
} = useHfDatasetSearch(debouncedQuery, {
modelType,
accessToken: hfToken || undefined,
});
@ -221,21 +223,8 @@ export function DatasetSection() {
>
<ComboboxList className="p-1 !max-h-none !overflow-visible">
{(id: string) => {
const r = hfResults.find((ds) => ds.id === id);
let detail: string | null = null;
if (r?.totalExamples) {
detail = `${formatCompact(r.totalExamples)} rows`;
} else if (r?.sizeCategory) {
detail = r.sizeCategory;
} else if (r?.downloads != null) {
detail = `↓${formatCompact(r.downloads)}`;
}
return (
<ComboboxItem
key={id}
value={id}
className="gap-2"
>
<ComboboxItem key={id} value={id} className="gap-2">
<Tooltip>
<TooltipTrigger asChild={true}>
<span className="block min-w-0 flex-1 truncate">
@ -249,11 +238,6 @@ export function DatasetSection() {
{id}
</TooltipContent>
</Tooltip>
{detail && (
<span className="ml-auto text-[10px] text-muted-foreground shrink-0">
{detail}
</span>
)}
</ComboboxItem>
);
}}

22
studio/frontend/src/features/studio/sections/model-section.tsx
View file

@ -159,6 +159,7 @@ export function ModelSection() {
} = useHfModelSearch(debouncedQuery, {
task,
accessToken: hfToken || undefined,
excludeGguf: true,
});
const { error: tokenValidationError, isChecking: isCheckingToken } =
@ -172,14 +173,25 @@ export function ModelSection() {
return ids;
}, [hfResults, selectedModel]);
// Filter out GGUF models — they can't be used for training
const trainableLocalModels = useMemo(
() =>
localModels.filter((m) => {
if (m.path.endsWith(".gguf")) return false;
if (m.id.toLowerCase().includes("-gguf")) return false;
return true;
}),
[localModels],
);
const localMetaById = useMemo(() => {
const map = new Map<string, LocalModelInfo>();
for (const model of localModels) map.set(model.id, model);
for (const model of trainableLocalModels) map.set(model.id, model);
return map;
}, [localModels]);
}, [trainableLocalModels]);
const localResultIds = useMemo(() => {
const ids = localModels.map((model) => model.id);
const ids = trainableLocalModels.map((model) => model.id);
const manual = localModelInput.trim();
if (manual && !ids.includes(manual)) {
ids.unshift(manual);
@ -346,8 +358,8 @@ export function ModelSection() {
<p className="text-[10px] text-red-500">{localModelsError}</p>
) : (
<p className="text-[10px] text-muted-foreground">
{localModels.length > 0
? `${localModels.length} local/cached models found`
{trainableLocalModels.length > 0
? `${trainableLocalModels.length} local/cached models found`
: "No local models found. Enter path manually."}
</p>
)}

242
studio/frontend/src/hooks/use-hf-dataset-search.ts
View file

@ -1,5 +1,6 @@
import { listDatasets } from "@huggingface/hub";
import { useCallback } from "react";
import { useCallback, useMemo } from "react";
import type { ModelType } from "@/types/training";
import { useHfPaginatedSearch } from "./use-hf-paginated-search";
interface DatasetInfoSplit {
@ -47,6 +48,8 @@ export interface HfDatasetResult {
likes: number;
totalExamples?: number;
sizeCategory?: string;
taskCategories: string[];
plainTags: string[];
}
function mapDataset(raw: unknown): HfDatasetResult {
@ -54,32 +57,261 @@ function mapDataset(raw: unknown): HfDatasetResult {
name: string;
downloads: number;
likes: number;
tags?: string[];
cardData?: unknown;
};
const card = ds.cardData as CardDataWithInfo | undefined;
const tags = ds.tags ?? [];
const taskCategories = tags
.filter((t) => t.startsWith("task_categories:"))
.map((t) => t.slice("task_categories:".length));
const plainTags = tags.filter((t) => !t.includes(":"));
return {
id: ds.name,
downloads: ds.downloads,
likes: ds.likes,
totalExamples: extractTotalExamples(card),
sizeCategory: card?.size_categories?.[0],
taskCategories,
plainTags,
};
}
function withTrendingSort(
input: Parameters<typeof fetch>[0],
init?: Parameters<typeof fetch>[1],
): ReturnType<typeof fetch> {
const rawUrl =
typeof input === "string"
? input
: input instanceof URL
? input.toString()
: input.url;
const url = new URL(rawUrl);
if (!url.searchParams.has("sort")) {
url.searchParams.set("sort", "trendingScore");
}
if (!url.searchParams.has("direction")) {
url.searchParams.set("direction", "-1");
}
return fetch(url, init);
}
type DatasetRelevance = "incompatible" | "neutral" | "boosted";
const BOOSTED_TASK_CATEGORIES: Record<ModelType, Set<string>> = {
text: new Set([
"text-generation",
"text2text-generation",
"question-answering",
"summarization",
"conversational",
]),
vision: new Set([
"image-text-to-text",
"visual-question-answering",
"image-to-text",
"image-captioning",
]),
tts: new Set([
"text-to-speech",
"text-to-audio",
"automatic-speech-recognition",
]),
embeddings: new Set([
"feature-extraction",
"sentence-similarity",
"text-retrieval",
]),
};
const INCOMPATIBLE_TASKS_ALL_MODELS = new Set([
"text-to-3d",
"image-to-3d",
"robotics",
"reinforcement-learning",
"tabular-classification",
"tabular-regression",
"time-series-forecasting",
]);
const PRETRAINING_PLAIN_TAGS = new Set(["pretraining", "pre-training"]);
const OCR_PLAIN_TAGS = new Set(["ocr", "document-ocr"]);
const PRETRAINING_SIZE_CATEGORIES = new Set([
"5M<n<10M",
"10M<n<100M",
"100M<n<1B",
"1B<n<10B",
"10B<n<100B",
"100B<n<1T",
"n>1T",
]);
const OCR_OR_VISION_TEXT_TASKS = new Set([
"image-to-text",
"image-captioning",
"visual-question-answering",
"document-question-answering",
]);
const INCOMPATIBLE_TASKS_BY_MODEL: Record<ModelType, Set<string>> = {
text: new Set([
"text-to-image",
"image-to-image",
"image-to-video",
"text-to-video",
"image-classification",
"image-feature-extraction",
"image-text-to-image",
"zero-shot-image-classification",
"keypoint-detection",
"object-detection",
"image-segmentation",
"depth-estimation",
"text-to-speech",
"text-to-audio",
"audio-classification",
"audio-to-audio",
"automatic-speech-recognition",
"video-classification",
"visual-document-retrieval",
]),
vision: new Set([
"text-to-speech",
"text-to-audio",
"audio-classification",
"audio-to-audio",
"automatic-speech-recognition",
]),
tts: new Set([
"text-to-image",
"image-to-image",
"image-to-video",
"text-to-video",
"image-classification",
"image-feature-extraction",
"image-text-to-image",
"zero-shot-image-classification",
"keypoint-detection",
"object-detection",
"image-segmentation",
"depth-estimation",
"video-classification",
"visual-document-retrieval",
]),
embeddings: new Set([
"text-to-image",
"image-to-image",
"image-to-video",
"text-to-video",
"image-classification",
"image-feature-extraction",
"image-text-to-image",
"zero-shot-image-classification",
"keypoint-detection",
"object-detection",
"image-segmentation",
"depth-estimation",
"text-to-speech",
"text-to-audio",
"audio-classification",
"audio-to-audio",
"automatic-speech-recognition",
"video-classification",
"visual-document-retrieval",
]),
};
function isPretrainingDataset(dataset: HfDatasetResult): boolean {
if (dataset.plainTags.some((t) => PRETRAINING_PLAIN_TAGS.has(t.toLowerCase())))
return true;
if (
dataset.sizeCategory &&
PRETRAINING_SIZE_CATEGORIES.has(dataset.sizeCategory)
)
return true;
return false;
}
function rankDatasetRelevance(
dataset: HfDatasetResult,
modelType: ModelType,
): DatasetRelevance {
if (isPretrainingDataset(dataset)) return "incompatible";
// Keep OCR / vision-text corpora out of non-vision defaults.
if (modelType !== "vision") {
if (
dataset.plainTags.some((t) => OCR_PLAIN_TAGS.has(t.toLowerCase())) ||
dataset.taskCategories.some((t) => OCR_OR_VISION_TEXT_TASKS.has(t))
) {
return "incompatible";
}
}
const { taskCategories } = dataset;
if (taskCategories.length === 0) return "neutral";
const boosted = BOOSTED_TASK_CATEGORIES[modelType];
const modelIncompat = INCOMPATIBLE_TASKS_BY_MODEL[modelType];
if (taskCategories.some((t) => boosted.has(t))) return "boosted";
if (
taskCategories.every(
(t) => INCOMPATIBLE_TASKS_ALL_MODELS.has(t) || modelIncompat.has(t),
)
)
return "incompatible";
return "neutral";
}
function isOcrOrVisionTextDataset(dataset: HfDatasetResult): boolean {
return (
dataset.plainTags.some((t) => OCR_PLAIN_TAGS.has(t.toLowerCase())) ||
dataset.taskCategories.some((t) => OCR_OR_VISION_TEXT_TASKS.has(t))
);
}
export function useHfDatasetSearch(
query: string,
options?: { accessToken?: string },
options?: { modelType?: ModelType | null; accessToken?: string },
) {
const { accessToken } = options ?? {};
const { modelType, accessToken } = options ?? {};
const createIter = useCallback(
() =>
listDatasets({
search: query.trim() ? { query } : {},
additionalFields: ["cardData"],
additionalFields: ["cardData", "tags"],
fetch: withTrendingSort,
...(accessToken ? { credentials: { accessToken } } : {}),
}) as AsyncGenerator<unknown>,
[query, accessToken],
);
return useHfPaginatedSearch(createIter, mapDataset);
const search = useHfPaginatedSearch(createIter, mapDataset);
const results = useMemo(() => {
const hideOcr = modelType !== "vision";
const baseResults = hideOcr
? search.results.filter((ds) => !isOcrOrVisionTextDataset(ds))
: search.results;
if (!modelType) return baseResults;
const boosted: HfDatasetResult[] = [];
const neutral: HfDatasetResult[] = [];
for (const ds of baseResults) {
const relevance = rankDatasetRelevance(ds, modelType);
if (relevance === "boosted") boosted.push(ds);
else if (relevance !== "incompatible") neutral.push(ds);
}
return [...boosted, ...neutral];
}, [search.results, modelType]);
return { ...search, results };
}

43
studio/frontend/src/hooks/use-hf-model-search.ts
View file

@ -11,7 +11,6 @@ export interface HfModelResult {
}
const EXCLUDED_TAGS = new Set([
"gguf",
"gptq",
"awq",
"exl2",
@ -45,22 +44,27 @@ function withPopularitySort(
return fetch(url, init);
}
function mapModel(raw: unknown): HfModelResult | null {
const m = raw as {
name: string;
downloads: number;
likes: number;
safetensors?: { total: number };
tags?: string[];
};
if (m.tags?.some((t) => EXCLUDED_TAGS.has(t))) {
return null;
}
return {
id: m.name,
downloads: m.downloads,
likes: m.likes,
totalParams: m.safetensors?.total,
function makeMapModel(excludeGguf: boolean) {
return (raw: unknown): HfModelResult | null => {
const m = raw as {
name: string;
downloads: number;
likes: number;
safetensors?: { total: number };
tags?: string[];
};
if (m.tags?.some((t) => EXCLUDED_TAGS.has(t))) {
return null;
}
if (excludeGguf && m.tags?.includes("gguf")) {
return null;
}
return {
id: m.name,
downloads: m.downloads,
likes: m.likes,
totalParams: m.safetensors?.total,
};
};
}
@ -113,9 +117,9 @@ async function* mergedModelIterator(
export function useHfModelSearch(
query: string,
options?: { task?: PipelineType; accessToken?: string },
options?: { task?: PipelineType; accessToken?: string; excludeGguf?: boolean },
) {
const { task, accessToken } = options ?? {};
const { task, accessToken, excludeGguf = false } = options ?? {};
const createIter = useCallback(
() => {
@ -135,6 +139,7 @@ export function useHfModelSearch(
[query, task, accessToken],
);
const mapModel = useMemo(() => makeMapModel(excludeGguf), [excludeGguf]);
const search = useHfPaginatedSearch(createIter, mapModel);
// Secondary sort guarantee: unsloth models always float to the top