unsloth/studio/backend/routes/datasets.py

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""Datasets API routes."""

import base64
import io
import json
import sys
from contextlib import suppress
from pathlib import Path
from uuid import uuid4
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException, Query, UploadFile
import re as _re
import structlog
from loggers import get_logger

_VALID_REPO_ID = _re.compile(r"^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$")


def _is_valid_repo_id(repo_id: str) -> bool:
    return bool(_VALID_REPO_ID.fullmatch(repo_id))


_dataset_size_cache: dict[str, int] = {}


def _get_dataset_size_cached(repo_id: str) -> int:
    if repo_id in _dataset_size_cache:
        return _dataset_size_cache[repo_id]
    try:
        from huggingface_hub import dataset_info as hf_dataset_info

        info = hf_dataset_info(repo_id, token = None, files_metadata = True)
        total = sum(s.size for s in info.siblings if getattr(s, "size", None))
        _dataset_size_cache[repo_id] = total
        return total
    except Exception:
        return 0


def _resolve_hf_cache_realpath(repo_dir: Path) -> Optional[str]:
    """Resolved realpath for a HF cache repo dir: most-recent snapshot, else cache root.

    Mirrors routes/models.py; duplicated here to keep this module self-contained.
    """
    try:
        snapshots_dir = repo_dir / "snapshots"
        if snapshots_dir.is_dir():
            snaps = [s for s in snapshots_dir.iterdir() if s.is_dir()]
            if snaps:
                latest = max(snaps, key = lambda s: s.stat().st_mtime)
                return str(latest.resolve())
        return str(repo_dir.resolve())
    except Exception:
        return None


backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
    sys.path.insert(0, str(backend_path))

from utils.datasets import check_dataset_format
from utils.upload_limits import get_upload_limit_bytes, get_upload_limit_label
from auth.authentication import get_current_subject

router = APIRouter()
logger = get_logger(__name__)


from models.datasets import (
    AiAssistMappingRequest,
    AiAssistMappingResponse,
    CheckFormatRequest,
    CheckFormatResponse,
    LocalDatasetItem,
    LocalDatasetsResponse,
    UploadDatasetResponse,
)
from utils.paths import (
    dataset_uploads_root,
    ensure_dir,
    recipe_datasets_root,
    resolve_dataset_path,
)


def _serialize_preview_value(value):
    """Make a value JSON-safe for the client preview."""
    if value is None or isinstance(value, (str, int, float, bool)):
        return value

    try:
        from PIL.Image import Image as PILImage
        if isinstance(value, PILImage):
            buffer = io.BytesIO()
            value.convert("RGB").save(buffer, format = "JPEG", quality = 85)
            return {
                "type": "image",
                "mime": "image/jpeg",
                "width": value.width,
                "height": value.height,
                "data": base64.b64encode(buffer.getvalue()).decode("ascii"),
            }
    except Exception:
        pass

    if isinstance(value, dict):
        return {str(key): _serialize_preview_value(item) for key, item in value.items()}

    if isinstance(value, (list, tuple)):
        return [_serialize_preview_value(item) for item in value]

    return str(value)


def _serialize_preview_rows(rows):
    return [
        {str(key): _serialize_preview_value(value) for key, value in dict(row).items()}
        for row in rows
    ]


# Data-file extensions for single-file preview. Tier 1 only uses tabular
# files; archives/text/config fall through to full load_dataset.
_COLUMNAR_EXTS = (".parquet", ".arrow")
_RECORD_EXTS = (".jsonl", ".csv", ".tsv")
_JSON_EXTS = (".json",)
_TABULAR_EXTS = _COLUMNAR_EXTS + _RECORD_EXTS + _JSON_EXTS
LOCAL_FILE_EXTS = (".json", ".jsonl", ".csv", ".parquet")
LOCAL_UPLOAD_EXTS = {".csv", ".json", ".jsonl", ".parquet"}
# sync: training dataset upload limits are exposed by /api/settings/upload-limit
LOCAL_DATASETS_ROOT = recipe_datasets_root()
DATASET_UPLOAD_DIR = dataset_uploads_root()


def _safe_read_metadata(path: Path) -> dict | None:
    try:
        payload = json.loads(path.read_text(encoding = "utf-8"))
    except (OSError, ValueError, TypeError):
        return None
    if not isinstance(payload, dict):
        return None
    return payload


_HF_PREVIEW_EXT_PRIORITY = {
    ".parquet": 0,
    ".arrow": 0,
    ".jsonl": 1,
    ".csv": 2,
    ".tsv": 3,
    ".json": 4,
}
_HF_NON_DATA_EXACT_FILENAMES = {
    ".gitattributes",
    "builder_config.json",
    "config.json",
    "dataset_info.json",
    "dataset_infos.json",
    "metadata.json",
}
_HF_NON_DATA_CARD_FILENAMES = {"card.json", "dataset_card.json"}


def _normalize_hf_repo_path(path: str) -> str:
    return path.strip().replace("\\", "/").lstrip("./")


def _hf_preview_extension(path: str) -> str | None:
    lower = path.lower()
    for ext in _HF_PREVIEW_EXT_PRIORITY:
        if lower.endswith(ext):
            return ext
    return None


def _is_known_hf_non_data_file(path: str) -> bool:
    name = Path(path).name.lower()
    if name in _HF_NON_DATA_EXACT_FILENAMES:
        return True
    if name in _HF_NON_DATA_CARD_FILENAMES:
        return True
    if name == "readme" or name.startswith("readme."):
        return True
    if name.endswith("_config.json") or name.endswith("-config.json"):
        return True
    if name.endswith("_card.json") or name.endswith("-card.json"):
        return True
    return False


def _is_hf_preview_data_file(path: str) -> bool:
    normalized = _normalize_hf_repo_path(path)
    if not normalized or _is_known_hf_non_data_file(normalized):
        return False
    return _hf_preview_extension(normalized) is not None


def _extract_hf_metadata_data_paths(metadata: dict | None) -> list[str]:
    if not metadata:
        return []
    file_paths = metadata.get("file_paths")
    if not isinstance(file_paths, dict):
        return []
    raw_data_paths = file_paths.get("data")
    if isinstance(raw_data_paths, str):
        values = [raw_data_paths]
    elif isinstance(raw_data_paths, list):
        values = raw_data_paths
    else:
        return []

    paths: list[str] = []
    for value in values:
        if not isinstance(value, str):
            continue
        normalized = _normalize_hf_repo_path(value)
        if normalized:
            paths.append(normalized)
    return paths


def _select_best_hf_preview_candidate(
    candidates: list[str], *, subset: str | None, split: str | None
) -> str | None:
    if not candidates:
        return None
    subset_lower = subset.lower() if subset else None
    split_lower = split.lower() if split else None

    def score(path: str) -> tuple[int, int, int, int, str]:
        ext = _hf_preview_extension(path)
        ext_priority = _HF_PREVIEW_EXT_PRIORITY[ext] if ext else 99
        stem = Path(path).stem.lower()
        path_lower = path.lower()

        subset_miss = 0
        if subset_lower:
            subset_miss = 0 if subset_lower in stem or subset_lower in path_lower else 1

        split_miss = 0
        if split_lower:
            split_hit = (
                stem == split_lower
                or stem.startswith(f"{split_lower}_")
                or stem.startswith(f"{split_lower}-")
                or f"/{split_lower}/" in path_lower
                or f"_{split_lower}." in path_lower
                or f"-{split_lower}." in path_lower
                or f"/{split_lower}." in path_lower
                or f"/{split_lower}_" in path_lower
                or f"/{split_lower}-" in path_lower
            )
            split_miss = 0 if split_hit else 1

        return (subset_miss, split_miss, ext_priority, len(path), path)

    return sorted(candidates, key = score)[0]


def _select_hf_preview_file(
    repo_files: list[str], *, metadata: dict | None, subset: str | None, split: str | None
) -> str | None:
    normalized_repo_files = [_normalize_hf_repo_path(path) for path in repo_files]
    repo_file_set = set(normalized_repo_files)

    metadata_candidates = [
        path
        for path in _extract_hf_metadata_data_paths(metadata)
        if path in repo_file_set and _is_hf_preview_data_file(path)
    ]
    if metadata_candidates:
        return _select_best_hf_preview_candidate(metadata_candidates, subset = subset, split = split)

    data_candidates = [path for path in normalized_repo_files if _is_hf_preview_data_file(path)]
    return _select_best_hf_preview_candidate(data_candidates, subset = subset, split = split)


def _download_hf_metadata(*, repo_id: str, repo_files: list[str], token: str | None) -> dict | None:
    metadata_file = next(
        (
            path
            for path in repo_files
            if Path(_normalize_hf_repo_path(path)).name.lower() == "metadata.json"
        ),
        None,
    )
    if not metadata_file:
        return None

    try:
        from huggingface_hub import hf_hub_download
        local_path = hf_hub_download(
            repo_id = repo_id,
            filename = metadata_file,
            repo_type = "dataset",
            token = token,
        )
    except Exception as exc:
        logger.warning(f"Could not read HF dataset metadata for {repo_id}: {exc}")
        return None

    return _safe_read_metadata(Path(local_path))


def _safe_read_rows_from_metadata(payload: dict | None) -> int | None:
    if not payload:
        return None
    for key in ("actual_num_records", "target_num_records"):
        value = payload.get(key)
        if isinstance(value, int):
            return value
    return None


def _safe_read_metadata_summary(payload: dict | None) -> dict | None:
    if not payload:
        return None

    actual_num_records = (
        payload.get("actual_num_records")
        if isinstance(payload.get("actual_num_records"), int)
        else None
    )
    target_num_records = (
        payload.get("target_num_records")
        if isinstance(payload.get("target_num_records"), int)
        else actual_num_records
    )

    columns: list[str] | None = None
    schema = payload.get("schema")
    if isinstance(schema, dict):
        columns = [str(key) for key in schema.keys()]
    if not columns:
        stats = payload.get("column_statistics")
        if isinstance(stats, list):
            derived = [
                str(item.get("column_name"))
                for item in stats
                if isinstance(item, dict) and item.get("column_name")
            ]
            columns = derived or None

    parquet_files_count = None
    file_paths = payload.get("file_paths")
    if isinstance(file_paths, dict):
        parquet_files = file_paths.get("parquet-files")
        if isinstance(parquet_files, list):
            parquet_files_count = len(parquet_files)

    total_num_batches = (
        payload.get("total_num_batches")
        if isinstance(payload.get("total_num_batches"), int)
        else parquet_files_count
    )
    num_completed_batches = (
        payload.get("num_completed_batches")
        if isinstance(payload.get("num_completed_batches"), int)
        else total_num_batches
    )

    return {
        "actual_num_records": actual_num_records,
        "target_num_records": target_num_records,
        "total_num_batches": total_num_batches,
        "num_completed_batches": num_completed_batches,
        "columns": columns,
    }


def _build_local_dataset_items() -> list[LocalDatasetItem]:
    if not LOCAL_DATASETS_ROOT.exists():
        return []

    items: list[LocalDatasetItem] = []
    for entry in LOCAL_DATASETS_ROOT.iterdir():
        if not entry.is_dir() or not entry.name.startswith("recipe_"):
            continue
        parquet_dir = entry / "parquet-files"
        if not parquet_dir.exists() or not any(parquet_dir.glob("*.parquet")):
            continue

        rows = None
        metadata_summary = None
        metadata_path = entry / "metadata.json"
        if metadata_path.exists():
            metadata_payload = _safe_read_metadata(metadata_path)
            rows = _safe_read_rows_from_metadata(metadata_payload)
            metadata_summary = _safe_read_metadata_summary(metadata_payload)

        try:
            updated_at = entry.stat().st_mtime
        except OSError:
            updated_at = None

        items.append(
            LocalDatasetItem(
                id = entry.name,
                label = entry.name,
                path = str(parquet_dir.resolve()),
                rows = rows,
                updated_at = updated_at,
                metadata = metadata_summary,
            )
        )

    items.sort(key = lambda item: item.updated_at or 0, reverse = True)
    return items


def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_size: int):
    # Non-streaming loads take the cached builder lock; use the EACCES-safe wrapper.
    from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset

    if dataset_path.is_dir():
        parquet_dir = (
            dataset_path / "parquet-files"
            if (dataset_path / "parquet-files").exists()
            else dataset_path
        )
        parquet_files = sorted(parquet_dir.glob("*.parquet"))
        if parquet_files:
            dataset = load_dataset(
                "parquet",
                data_files = [str(path) for path in parquet_files],
                split = train_split,
            )
            total_rows = len(dataset)
            preview_slice = dataset.select(range(min(preview_size, total_rows)))
            return preview_slice, total_rows
        else:
            candidate_files: list[Path] = []
            for ext in LOCAL_FILE_EXTS:
                candidate_files.extend(sorted(dataset_path.glob(f"*{ext}")))
            if not candidate_files:
                raise HTTPException(
                    status_code = 400,
                    detail = "Unsupported local dataset directory (expected parquet/json/jsonl/csv files)",
                )
            dataset_path = candidate_files[0]

    if dataset_path.suffix in [".json", ".jsonl"]:
        dataset = load_dataset("json", data_files = str(dataset_path), split = train_split)
    elif dataset_path.suffix == ".csv":
        dataset = load_dataset("csv", data_files = str(dataset_path), split = train_split)
    elif dataset_path.suffix == ".parquet":
        dataset = load_dataset("parquet", data_files = str(dataset_path), split = train_split)
    else:
        raise HTTPException(
            status_code = 400, detail = f"Unsupported file format: {dataset_path.suffix}"
        )

    total_rows = len(dataset)
    preview_slice = dataset.select(range(min(preview_size, total_rows)))
    return preview_slice, total_rows


def _sanitize_filename(filename: str) -> str:
    name = Path(filename).name.strip().replace("\x00", "")
    if not name:
        return "dataset_upload"
    return name


@router.post("/upload", response_model = UploadDatasetResponse)
async def upload_dataset(
    file: UploadFile, current_subject: str = Depends(get_current_subject)
) -> UploadDatasetResponse:
    filename = _sanitize_filename(file.filename or "dataset_upload")
    ext = Path(filename).suffix.lower()
    if ext not in LOCAL_UPLOAD_EXTS:
        allowed = ", ".join(sorted(LOCAL_UPLOAD_EXTS))
        raise HTTPException(
            status_code = 400,
            detail = f"Unsupported file type: {ext}. Allowed: {allowed}",
        )

    ensure_dir(DATASET_UPLOAD_DIR)
    stem = Path(filename).stem
    stored_name = f"{uuid4().hex}_{stem}{ext}"
    stored_path = DATASET_UPLOAD_DIR / stored_name

    # Stream to disk in chunks to avoid holding the whole file in memory. The
    # route-level cap gives a clear training-dataset error and avoids leaving
    # oversized partial files in the Unsloth uploads directory.
    upload_limit_bytes = get_upload_limit_bytes()
    total_bytes = 0
    upload_complete = False
    try:
        with open(stored_path, "wb") as f:
            while chunk := await file.read(1024 * 1024):
                total_bytes += len(chunk)
                if total_bytes > upload_limit_bytes:
                    raise HTTPException(
                        status_code = 413,
                        detail = (
                            "Training dataset upload too large. "
                            f"Maximum is {get_upload_limit_label()}."
                        ),
                    )
                f.write(chunk)
        upload_complete = True
    finally:
        if not upload_complete:
            with suppress(OSError):
                stored_path.unlink(missing_ok = True)

    if stored_path.stat().st_size == 0:
        stored_path.unlink(missing_ok = True)
        raise HTTPException(status_code = 400, detail = "Empty upload payload")

    return UploadDatasetResponse(filename = filename, stored_path = str(stored_path))


@router.get("/local", response_model = LocalDatasetsResponse)
def list_local_datasets(
    current_subject: str = Depends(get_current_subject),
) -> LocalDatasetsResponse:
    return LocalDatasetsResponse(datasets = _build_local_dataset_items())


@router.get("/download-progress")
async def get_dataset_download_progress(
    repo_id: str = Query(..., description = "HuggingFace dataset repo ID, e.g. 'unsloth/LaTeX_OCR'"),
    current_subject: str = Depends(get_current_subject),
):
    """Return download progress for a HuggingFace dataset repo.

    Mirrors ``GET /api/models/download-progress`` but scans the
    ``datasets--owner--name`` cache dir under HF_HUB_CACHE, where in-progress
    download bytes are visible. Returns ``cache_path`` so the UI can show it.
    """
    _empty = {
        "downloaded_bytes": 0,
        "expected_bytes": 0,
        "progress": 0,
        "cache_path": None,
    }
    try:
        if not _is_valid_repo_id(repo_id):
            return _empty

        from huggingface_hub import constants as hf_constants

        cache_dir = Path(hf_constants.HF_HUB_CACHE)
        target = f"datasets--{repo_id.replace('/', '--')}".lower()
        completed_bytes = 0
        in_progress_bytes = 0
        cache_path: Optional[str] = None

        if cache_dir.is_dir():
            for entry in cache_dir.iterdir():
                if entry.name.lower() != target:
                    continue
                cache_path = _resolve_hf_cache_realpath(entry)
                blobs_dir = entry / "blobs"
                if not blobs_dir.is_dir():
                    break
                for f in blobs_dir.iterdir():
                    if not f.is_file():
                        continue
                    if f.name.endswith(".incomplete"):
                        in_progress_bytes += f.stat().st_size
                    else:
                        completed_bytes += f.stat().st_size
                break

        downloaded_bytes = completed_bytes + in_progress_bytes
        if downloaded_bytes == 0:
            return {**_empty, "cache_path": cache_path}

        expected_bytes = _get_dataset_size_cached(repo_id)
        if expected_bytes <= 0:
            return {
                "downloaded_bytes": downloaded_bytes,
                "expected_bytes": 0,
                "progress": 0,
                "cache_path": cache_path,
            }

        # 95% threshold (as in the model endpoint): HF blob dedup makes
        # completed_bytes drift under expected_bytes; inter-file gaps look "done".
        if completed_bytes >= expected_bytes * 0.95:
            progress = 1.0
        else:
            progress = min(downloaded_bytes / expected_bytes, 0.99)
        return {
            "downloaded_bytes": downloaded_bytes,
            "expected_bytes": expected_bytes,
            "progress": round(progress, 3),
            "cache_path": cache_path,
        }
    except Exception as e:
        logger.warning(f"Error checking dataset download progress for {repo_id}: {e}")
        return _empty


@router.post("/check-format", response_model = CheckFormatResponse)
def check_format(request: CheckFormatRequest, current_subject: str = Depends(get_current_subject)):
    """Check if a dataset requires manual column mapping.

    HuggingFace strategy:
      1. list_repo_files -> select one tabular data file -> load_dataset
         (avoids resolving thousands of files; ~2-4 s).
      2. Full streaming load_dataset as a last-resort fallback.

    Local files load directly. Plain `def` (not async) so FastAPI runs it in a
    thread-pool, keeping blocking IO off the event loop.
    """
    try:
        from itertools import islice
        from datasets import Dataset, load_dataset
        from utils.datasets import format_dataset

        PREVIEW_SIZE = 10

        logger.info(f"Checking format for dataset: {request.dataset_name}")

        dataset_path = resolve_dataset_path(request.dataset_name)
        total_rows = None

        if dataset_path.exists():
            # ── Local file ──────────────────────────────────────────
            train_split = request.train_split or "train"
            preview_slice, total_rows = _load_local_preview_slice(
                dataset_path = dataset_path,
                train_split = train_split,
                preview_size = PREVIEW_SIZE,
            )
        else:
            # ── HuggingFace dataset ─────────────────────────────────
            # Tier 1: list_repo_files -> load one selected tabular data file
            preview_slice = None

            try:
                from huggingface_hub import HfApi

                api = HfApi()
                repo_files = api.list_repo_files(
                    request.dataset_name,
                    repo_type = "dataset",
                    token = request.hf_token or None,
                )
                metadata = _download_hf_metadata(
                    repo_id = request.dataset_name,
                    repo_files = repo_files,
                    token = request.hf_token or None,
                )
                selected_file = _select_hf_preview_file(
                    repo_files,
                    metadata = metadata,
                    subset = request.subset,
                    split = request.train_split or "train",
                )

                if selected_file:
                    logger.info(f"Tier 1: loading single file {selected_file}")
                    load_kwargs = {
                        "path": request.dataset_name,
                        "data_files": [selected_file],
                        "split": "train",
                        "streaming": True,
                    }
                    if request.hf_token:
                        load_kwargs["token"] = request.hf_token

                    streamed_ds = load_dataset(**load_kwargs)
                    rows = list(islice(streamed_ds, PREVIEW_SIZE))
                    if rows:
                        preview_slice = Dataset.from_list(rows)
            except Exception as e:
                logger.warning(f"Tier 1 (single-file) failed: {e}")

            if preview_slice is None:
                # Tier 2: full streaming (resolves all files; slow for large repos)
                logger.info("Tier 2: falling back to full streaming load_dataset")
                load_kwargs = {
                    "path": request.dataset_name,
                    "split": request.train_split,
                    "streaming": True,
                }
                if request.subset:
                    load_kwargs["name"] = request.subset
                if request.hf_token:
                    load_kwargs["token"] = request.hf_token

                streamed_ds = load_dataset(**load_kwargs)

                rows = list(islice(streamed_ds, PREVIEW_SIZE))
                if not rows:
                    raise HTTPException(
                        status_code = 400,
                        detail = "Dataset appears to be empty or could not be streamed",
                    )

                preview_slice = Dataset.from_list(rows)
            total_rows = None

        result = check_dataset_format(preview_slice, is_vlm = request.is_vlm)

        logger.info(
            f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}, is_image={result.get('is_image', False)}"
        )

        preview_samples = None
        if not result["requires_manual_mapping"]:
            if result.get("suggested_mapping"):
                # Heuristic-detected: show raw data so columns match the API response.
                # Column stripping happens at training time, not preview.
                preview_samples = _serialize_preview_rows(preview_slice)
            else:
                try:
                    format_result = format_dataset(
                        preview_slice,
                        format_type = "auto",
                        num_proc = None,  # Only 10 preview rows
                    )
                    processed = format_result["dataset"]
                    preview_samples = _serialize_preview_rows(processed)
                except Exception as e:
                    logger.warning(f"Processed preview generation failed (non-fatal): {e}")
                    preview_samples = _serialize_preview_rows(preview_slice)
        else:
            preview_samples = _serialize_preview_rows(preview_slice)

        # Warnings from check_dataset_format plus URL-based image detection.
        warning = result.get("warning")
        image_col = result.get("detected_image_column")
        if image_col and image_col in (result.get("columns") or []):
            try:
                sample_val = preview_slice[0][image_col]
                if isinstance(sample_val, str) and sample_val.startswith(("http://", "https://")):
                    url_warning = (
                        "This dataset contains image URLs instead of embedded images. "
                        "Images will be downloaded during training, which may be slow for large datasets."
                    )
                    logger.info(f"URL-based image column detected: {image_col}")
                    warning = f"{warning} {url_warning}" if warning else url_warning
            except Exception:
                pass

        return CheckFormatResponse(
            requires_manual_mapping = result["requires_manual_mapping"],
            detected_format = result["detected_format"],
            columns = result["columns"],
            is_image = result.get("is_image", False),
            is_audio = result.get("is_audio", False),
            multimodal_columns = result.get("multimodal_columns"),
            suggested_mapping = result.get("suggested_mapping"),
            detected_image_column = result.get("detected_image_column"),
            detected_audio_column = result.get("detected_audio_column"),
            detected_text_column = result.get("detected_text_column"),
            detected_speaker_column = result.get("detected_speaker_column"),
            chat_column = result.get("chat_column"),
            preview_samples = preview_samples,
            total_rows = total_rows,
            warning = warning,
        )

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error checking dataset format: {e}", exc_info = True)
        raise HTTPException(status_code = 500, detail = "Failed to check dataset format")


@router.post("/ai-assist-mapping", response_model = AiAssistMappingResponse)
def ai_assist_mapping(
    request: AiAssistMappingRequest, current_subject: str = Depends(get_current_subject)
):
    """Run LLM-assisted dataset conversion advisor (user-triggered).

    Multi-pass analysis with a 7B helper model: classify dataset type, generate
    conversion strategy, validate quality. Falls back to simple column
    classification if the advisor fails.
    """
    try:
        from utils.datasets.llm_assist import llm_conversion_advisor

        # Truncate sample values for the LLM prompt.
        truncated = [
            {col: str(s.get(col, ""))[:200] for col in request.columns} for s in request.samples[:5]
        ]

        result = llm_conversion_advisor(
            column_names = request.columns,
            samples = truncated,
            dataset_name = request.dataset_name,
            hf_token = request.hf_token,
            model_name = request.model_name,
            model_type = request.model_type,
        )

        if result and result.get("success"):
            return AiAssistMappingResponse(
                success = True,
                suggested_mapping = result.get("suggested_mapping"),
                system_prompt = result.get("system_prompt"),
                user_template = result.get("user_template"),
                assistant_template = result.get("assistant_template"),
                label_mapping = result.get("label_mapping"),
                dataset_type = result.get("dataset_type"),
                is_conversational = result.get("is_conversational"),
                user_notification = result.get("user_notification"),
            )

        return AiAssistMappingResponse(
            success = False,
            warning = "AI could not determine column roles. Please assign them manually.",
        )

    except Exception as e:
        logger.error(f"AI assist mapping failed: {e}", exc_info = True)
        raise HTTPException(status_code = 500, detail = "AI assist failed")