diff --git a/.gitignore b/.gitignore index 08f9d8ee6b..044775e846 100755 --- a/.gitignore +++ b/.gitignore @@ -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/ diff --git a/setup.sh b/setup.sh index 3e6b41db9f..8314ef6d75 100755 --- a/setup.sh +++ b/setup.sh @@ -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 diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index 6ea668db3e..e725834630 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -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 diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index faea6266e3..e5a486bca2 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -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 diff --git a/studio/backend/core/data_recipe/jobs/constants.py b/studio/backend/core/data_recipe/jobs/constants.py new file mode 100644 index 0000000000..5d4081abbc --- /dev/null +++ b/studio/backend/core/data_recipe/jobs/constants.py @@ -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" diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index eb8c10bb81..dbaa620004 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -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" diff --git a/studio/backend/core/data_recipe/jobs/parse.py b/studio/backend/core/data_recipe/jobs/parse.py index 99d1a85a79..6e2142adf2 100644 --- a/studio/backend/core/data_recipe/jobs/parse.py +++ b/studio/backend/core/data_recipe/jobs/parse.py @@ -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\d+) records across (?P\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: diff --git a/studio/backend/core/data_recipe/jobs/worker.py b/studio/backend/core/data_recipe/jobs/worker.py index ac27cd0d0b..8c0996b140 100644 --- a/studio/backend/core/data_recipe/jobs/worker.py +++ b/studio/backend/core/data_recipe/jobs/worker.py @@ -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) diff --git a/studio/backend/core/data_recipe/jsonable.py b/studio/backend/core/data_recipe/jsonable.py new file mode 100644 index 0000000000..aa6e1d6b2e --- /dev/null +++ b/studio/backend/core/data_recipe/jsonable.py @@ -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 diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index b04c6bbf00..2d11cb2845 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -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 diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 231566cc22..bc4e267f75 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -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: diff --git a/studio/backend/core/inference/__init__.py b/studio/backend/core/inference/__init__.py index 494229a087..ff8b75d36a 100644 --- a/studio/backend/core/inference/__init__.py +++ b/studio/backend/core/inference/__init__.py @@ -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', ] diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py new file mode 100644 index 0000000000..68bf871590 --- /dev/null +++ b/studio/backend/core/inference/llama_cpp.py @@ -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 diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 43cb5ff1a1..e4e7a474be 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -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( diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index d2d98d7944..3a908dcfc3 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -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") diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index 10d87d7825..8c7d0c037d 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -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") diff --git a/studio/backend/routes/data_recipe/__init__.py b/studio/backend/routes/data_recipe/__init__.py new file mode 100644 index 0000000000..dc0301d1c9 --- /dev/null +++ b/studio/backend/routes/data_recipe/__init__.py @@ -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"] diff --git a/studio/backend/routes/data_recipe/jobs.py b/studio/backend/routes/data_recipe/jobs.py new file mode 100644 index 0000000000..ffbded9474 --- /dev/null +++ b/studio/backend/routes/data_recipe/jobs.py @@ -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") diff --git a/studio/backend/routes/data_recipe.py b/studio/backend/routes/data_recipe/seed.py similarity index 53% rename from studio/backend/routes/data_recipe.py rename to studio/backend/routes/data_recipe/seed.py index 8a713f2cec..eb02ab2bbf 100644 --- a/studio/backend/routes/data_recipe.py +++ b/studio/backend/routes/data_recipe/seed.py @@ -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= - 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") diff --git a/studio/backend/routes/data_recipe/validate.py b/studio/backend/routes/data_recipe/validate.py new file mode 100644 index 0000000000..a8755f9410 --- /dev/null +++ b/studio/backend/routes/data_recipe/validate.py @@ -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) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index c30a1638d4..8d1c6667c0 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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 + 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, diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 8c96ee5f28..2c7f0e846f 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -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( diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index 006deb99c0..92e65cf67c 100644 --- a/studio/backend/utils/models/__init__.py +++ b/studio/backend/utils/models/__init__.py @@ -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', diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index f7b95fad11..b84a14d226 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -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) diff --git a/studio/frontend/src/components/assistant-ui/attachment.tsx b/studio/frontend/src/components/assistant-ui/attachment.tsx index c53c134ea2..94e30fb4ac 100644 --- a/studio/frontend/src/components/assistant-ui/attachment.tsx +++ b/studio/frontend/src/components/assistant-ui/attachment.tsx @@ -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}`); } }); diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index d470034832..c6685f0214 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -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) { diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 319f4fd484..7bf23dfbeb 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -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(null); + const [defaultVariant, setDefaultVariant] = useState(null); + const [hasVision, setHasVision] = useState(false); + const [loading, setLoading] = useState(true); + const [error, setError] = useState(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 ( +
+ + Loading variants… +
+ ); + } + + if (error) { + return ( +
{error}
+ ); + } + + if (!variants || variants.length === 0) { + return ( +
+ No GGUF variants found. +
+ ); + } + + return ( +
+
+ + Quantizations + + {hasVision && ( + Vision + )} +
+ {variants.map((v) => { + const sizeGb = v.size_bytes / (1024 ** 3); + const fitStatus = gpuGb != null && gpuGb > 0 && sizeGb > 0 + ? checkVramFit(sizeGb, gpuGb) + : null; + return ( + + ); + })} +
+ ); +} + +// ── 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(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 (
@@ -230,18 +394,24 @@ export function HubModelPicker({ recommendedIds.map((id) => { const vram = recommendedVramMap.get(id); return ( - - onSelect(id, { source: "hub", isLora: false }) - } - vramStatus={vram?.status ?? null} - vramEst={vram?.est} - gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} - /> +
+ handleModelClick(id)} + vramStatus={isGgufRepo(id) ? null : vram?.status ?? null} + vramEst={isGgufRepo(id) ? undefined : vram?.est} + gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} + /> + {expandedGguf === id && ( + + )} +
); }) )} @@ -259,18 +429,24 @@ export function HubModelPicker({ hfIds.map((id) => { const vram = vramMap.get(id); return ( - - onSelect(id, { source: "hub", isLora: false }) - } - vramStatus={vram?.status ?? null} - vramEst={vram?.est} - gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} - /> +
+ handleModelClick(id)} + vramStatus={isGgufRepo(id) ? null : vram?.status ?? null} + vramEst={isGgufRepo(id) ? undefined : vram?.est} + gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} + /> + {expandedGguf === id && ( + + )} +
); }) )} @@ -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 ( onSelect(adapter.id, { source: isExported ? "exported" : "lora", - isLora: !isMerged, + isLora: !isMerged && !isGguf, })} /> ); @@ -393,4 +572,3 @@ export function LoraModelPicker({
); } - diff --git a/studio/frontend/src/components/assistant-ui/model-selector/types.ts b/studio/frontend/src/components/assistant-ui/model-selector/types.ts index b8df0f6c6c..0e8cf5fb4d 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/types.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/types.ts @@ -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; } diff --git a/studio/frontend/src/components/ui/combobox.tsx b/studio/frontend/src/components/ui/combobox.tsx index 157fe3ab68..34b66ea64e 100644 --- a/studio/frontend/src/components/ui/combobox.tsx +++ b/studio/frontend/src/components/ui/combobox.tsx @@ -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 ( - - + + ); diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index 33249a044a..4df5107022 100644 --- a/studio/frontend/src/config/training.ts +++ b/studio/frontend/src/config/training.ts @@ -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]; diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 72baf9a6f6..5d5a9551ef 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -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 { await parseJsonOrThrow(response); } +export async function listGgufVariants( + repoId: string, + hfToken?: string, +): Promise { + 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(response); +} + function parseSseEvent(rawEvent: string): string[] { const dataLines: string[] = []; for (const line of rawEvent.split(/\r?\n/)) { diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 1ef3c6bec5..14a457ad76 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -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" diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 2d7ae781fa..1bc61cd25d 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -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); } } diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index aafbf00310..fece047cd8 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -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; diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 03825b5861..2e3b43d606 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -39,13 +39,14 @@ type ChatRuntimeStore = { runningByThreadId: Record; 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((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((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((set) => ({ ...state.params, checkpoint: "", }, + activeGgufVariant: null, })), })); diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 003c3b3629..edf8fecae3 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -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[]; } diff --git a/studio/frontend/src/features/chat/types/runtime.ts b/studio/frontend/src/features/chat/types/runtime.ts index 48f147d9bc..710898713b 100644 --- a/studio/frontend/src/features/chat/types/runtime.ts +++ b/studio/frontend/src/features/chat/types/runtime.ts @@ -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"; } diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 9d3e5053de..fa43c5eb98 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -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}` diff --git a/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx b/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx index 7fc42fd264..f5ba2eb898 100644 --- a/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx @@ -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() { > {(id: string) => { - const r = hfResults.find((r) => r.id === id); - const detail = r?.totalExamples - ? `${formatCompact(r.totalExamples)} rows` - : (r?.sizeCategory ?? null); return ( - + @@ -274,15 +269,6 @@ export function DatasetStep() { {id} - {detail ? ( - - {detail} - - ) : r?.downloads != null ? ( - - ↓{formatCompact(r.downloads)} - - ) : null} ); }} diff --git a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx index 747311c9a5..5cf4ebc0b5 100644 --- a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx @@ -85,6 +85,7 @@ export function ModelSelectionStep() { } = useHfModelSearch(debouncedQuery, { task, accessToken: hfToken || undefined, + excludeGguf: true, }); const { error: tokenValidationError, isChecking: isCheckingToken } = diff --git a/studio/frontend/src/features/recipe-studio/blocks/definitions.ts b/studio/frontend/src/features/recipe-studio/blocks/definitions.ts index 75d860290e..073cce928e 100644 --- a/studio/frontend/src/features/recipe-studio/blocks/definitions.ts +++ b/studio/frontend/src/features/recipe-studio/blocks/definitions.ts @@ -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", diff --git a/studio/frontend/src/features/recipe-studio/components/block-sheet.tsx b/studio/frontend/src/features/recipe-studio/components/block-sheet.tsx index d9e278c05d..be301a36fd 100644 --- a/studio/frontend/src/features/recipe-studio/components/block-sheet.tsx +++ b/studio/frontend/src/features/recipe-studio/components/block-sheet.tsx @@ -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) => void; + trailing?: "chevron" | "drag" | "none"; }): ReactElement { return ( ); } @@ -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) => { + 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 (
@@ -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" > - +
{sheetView !== "root" && (
+ setSearch(event.target.value)} + placeholder="Search blocks..." + className="corner-squircle mt-3 h-9" + />
- {sheetView === "root" && - ROOT_GROUPS.map((item, index) => ( + {isRootView && + hasSearch && + rootSearchBlocks.map((item, index) => ( + onBlockClick(item.kind, item.type)} + /> + ))} + {isRootView && + !hasSearch && + rootGroups.map((item, index) => ( { if (item.kind === "processor") { setSheetOpen(false); @@ -256,18 +426,20 @@ export function BlockSheet({ }} /> ))} - {sheetView === "processor" && ( - + {isProcessorView && ( + (!hasSearch || + matchesSearch(PROCESSOR_TITLE, PROCESSOR_DESCRIPTION)) && ( + + ) )} - {sheetView !== "root" && - sheetView !== "processor" && - getBlocksForKind(VIEW_KIND[sheetView] ?? "sampler").map( + {isScopedBlockView && + scopedBlocks.map( (item, index) => ( { - 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 && ( +

+ No blocks match. +

+ )}
diff --git a/studio/frontend/src/features/recipe-studio/components/chip-input.tsx b/studio/frontend/src/features/recipe-studio/components/chip-input.tsx index 5162b73e03..a7dd5ccc16 100644 --- a/studio/frontend/src/features/recipe-studio/components/chip-input.tsx +++ b/studio/frontend/src/features/recipe-studio/components/chip-input.tsx @@ -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(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 ( -
+
{values.map((value, index) => ( { 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 ( diff --git a/studio/frontend/src/features/recipe-studio/components/controls/viewport-controls.tsx b/studio/frontend/src/features/recipe-studio/components/controls/viewport-controls.tsx index 27e225f39e..e104841100 100644 --- a/studio/frontend/src/features/recipe-studio/components/controls/viewport-controls.tsx +++ b/studio/frontend/src/features/recipe-studio/components/controls/viewport-controls.tsx @@ -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 ( diff --git a/studio/frontend/src/features/recipe-studio/components/inline/inline-expression.tsx b/studio/frontend/src/features/recipe-studio/components/inline/inline-expression.tsx index 94c5470f49..206a068a78 100644 --- a/studio/frontend/src/features/recipe-studio/components/inline/inline-expression.tsx +++ b/studio/frontend/src/features/recipe-studio/components/inline/inline-expression.tsx @@ -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 (
@@ -52,32 +57,14 @@ export function InlineExpression({ 0} placeholder="{{ column_name }}" value={config.expr} onChange={(event) => onUpdate({ expr: event.target.value })} />
- {vars.length > 0 && ( -
-

Available references

-
- {vars.map((v) => ( - - {v.name} - - ))} -
-
- )} +
); } diff --git a/studio/frontend/src/features/recipe-studio/components/inline/inline-model.tsx b/studio/frontend/src/features/recipe-studio/components/inline/inline-model.tsx index 5316b0f297..2c2c2ffd41 100644 --- a/studio/frontend/src/features/recipe-studio/components/inline/inline-model.tsx +++ b/studio/frontend/src/features/recipe-studio/components/inline/inline-model.tsx @@ -14,19 +14,6 @@ export function InlineModel(props: InlineModelProps): ReactElement { if (props.config.kind === "model_provider") { return (
- - - props.onUpdate({ - // biome-ignore lint/style/useNamingConvention: api schema - provider_type: event.target.value, - }) - } - /> - props.onUpdate({ endpoint: event.target.value })} /> + + + props.onUpdate({ + // biome-ignore lint/style/useNamingConvention: api schema + api_key: event.target.value, + }) + } + /> +
); } diff --git a/studio/frontend/src/features/recipe-studio/components/inline/inline-policy.ts b/studio/frontend/src/features/recipe-studio/components/inline/inline-policy.ts index 97f441edce..17f22b6557 100644 --- a/studio/frontend/src/features/recipe-studio/components/inline/inline-policy.ts +++ b/studio/frontend/src/features/recipe-studio/components/inline/inline-policy.ts @@ -27,6 +27,9 @@ export function getConfigUiMode( } return "dialog"; } + if (config.kind === "seed") { + return "inline"; + } if (config.kind === "expression") { return "inline"; } diff --git a/studio/frontend/src/features/recipe-studio/components/inline/inline-seed.tsx b/studio/frontend/src/features/recipe-studio/components/inline/inline-seed.tsx new file mode 100644 index 0000000000..619faa7704 --- /dev/null +++ b/studio/frontend/src/features/recipe-studio/components/inline/inline-seed.tsx @@ -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) => void; +}; + +export function InlineSeed({ config, onUpdate }: InlineSeedProps): ReactElement { + const mode = config.seed_source_type ?? "hf"; + + if (mode === "hf") { + return ( +
+ + + onUpdate({ + hf_repo_id: next, + hf_path: "", + seed_columns: [], + seed_drop_columns: [], + seed_preview_rows: [], + }) + } + placeholder="org/repo" + /> + +

+ Load columns in dialog. +

+
+ ); + } + + const isLocal = mode === "local"; + const fileName = isLocal + ? config.local_file_name?.trim() + : config.unstructured_file_name?.trim(); + + return ( +
+
+ +
+
+

+ {fileName || "No file selected"} +

+

+ {isLocal ? "Structured file" : "Unstructured document"} · configure in dialog +

+
+ +
+ ); +} diff --git a/studio/frontend/src/features/recipe-studio/components/recipe-graph-aux-node.tsx b/studio/frontend/src/features/recipe-studio/components/recipe-graph-aux-node.tsx index c7f14ebbc4..bbb5dead41 100644 --- a/studio/frontend/src/features/recipe-studio/components/recipe-graph-aux-node.tsx +++ b/studio/frontend/src/features/recipe-studio/components/recipe-graph-aux-node.tsx @@ -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 ( -
-

Available references

-
- {vars.map((v) => ( - - {v.name} - - ))} -
-
- ); +function AuxVariableBadges({ + entries, +}: { + entries: ReturnType; +}): ReactElement | null { + return ; } function AuxNodeBase({ @@ -93,6 +70,7 @@ function AuxNodeBase({ data, }: NodeProps): 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 = ( + <> + + + + + + ); 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 ( - {data.title}