# SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """Pydantic schemas for Model Management API""" from pydantic import BaseModel, Field from typing import Optional, List, Dict, Any, Literal ModelType = Literal["text", "vision", "audio", "embeddings"] class CheckpointInfo(BaseModel): """Information about a discovered checkpoint directory.""" display_name: str = Field(..., description = "User-friendly checkpoint name (folder name)") path: str = Field(..., description = "Full path to the checkpoint directory") loss: Optional[float] = Field(None, description = "Training loss at this checkpoint") class ModelCheckpoints(BaseModel): """A training run and its associated checkpoints.""" name: str = Field(..., description = "Training run folder name") checkpoints: List[CheckpointInfo] = Field( default_factory = list, description = "List of checkpoints for this training run (final + intermediate)", ) base_model: Optional[str] = Field( None, description = "Base model name from adapter_config.json or config.json", ) peft_type: Optional[str] = Field( None, description = "PEFT type (e.g. LORA) if adapter training, None for full fine-tune", ) lora_rank: Optional[int] = Field( None, description = "LoRA rank (r) if applicable", ) is_quantized: bool = Field( False, description = "Whether the model uses BNB quantization (e.g. bnb-4bit)", ) class CheckpointListResponse(BaseModel): """Response for listing available checkpoints in an outputs directory.""" outputs_dir: str = Field(..., description = "Directory that was scanned") models: List[ModelCheckpoints] = Field( default_factory = list, description = "List of training runs with their checkpoints", ) class ExportSizeResponse(BaseModel): """Model fp16/bf16-equivalent size; size fields are null when unknown.""" model: str = Field(..., description = "Model id or path the estimate was computed for") fp16_bytes: Optional[int] = Field( None, description = "Estimated FP16/BF16-equivalent on-disk size in bytes, or null if unknown", ) total_params: Optional[int] = Field( None, description = "Estimated total parameter count (fp16_bytes // 2), or null if unknown", ) source: str = Field( "unavailable", description = "How the estimate was derived (e.g. safetensors, config, local, vllm, unavailable)", ) class ModelDetails(BaseModel): """Model configuration and metadata; used for both list and detail views""" id: str = Field(..., description = "Model identifier") model_name: Optional[str] = Field( None, description = "Model identifier (alias for id, for backward compatibility)" ) name: Optional[str] = Field(None, description = "Display name for the model") config: Optional[Dict[str, Any]] = Field(None, description = "Model configuration dictionary") is_vision: bool = Field(False, description = "Whether model is a vision model") is_embedding: bool = Field( False, description = "Whether model is an embedding/sentence-transformer 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)") is_mlx: bool = Field( False, description = "Whether model is served via the MLX backend (Apple Silicon)" ) is_audio: bool = Field(False, description = "Whether model is a TTS audio model") audio_type: Optional[str] = Field(None, description = "Audio codec type: snac, csm, bicodec, dac") has_audio_input: bool = Field(False, description = "Whether model accepts audio input (ASR)") model_type: Optional[ModelType] = Field( None, description = "Collapsed model modality: text, vision, audio, or embeddings" ) base_model: Optional[str] = Field(None, description = "Base model if this is a LoRA adapter") max_position_embeddings: Optional[int] = Field( None, description = "Maximum context length supported by the model" ) model_size_bytes: Optional[int] = Field( None, description = "Total size of model weight files in bytes" ) class LoRAInfo(BaseModel): """LoRA adapter or exported model information""" display_name: str = Field(..., description = "Display name for the LoRA") 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', 'merged', or 'gguf' (for exports)" ) class LoRAScanResponse(BaseModel): """Response schema for scanning trained LoRA adapters""" loras: List[LoRAInfo] = Field(default_factory = list, description = "List of found LoRA adapters") outputs_dir: str = Field(..., description = "Directory that was scanned") class ModelListResponse(BaseModel): """Response schema for listing models""" models: List[ModelDetails] = Field(default_factory = list, description = "List of models") 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") download_size_bytes: int = Field(0, description = "Total bytes needed to download this variant") downloaded: bool = Field( False, description = "Whether this variant is already in the local HF cache" ) update_available: bool = Field( False, description = "Whether a newer version of this variant is available on HF" ) 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" ) context_length: Optional[int] = Field( None, description = "Native max context from GGUF metadata; set once a variant is downloaded", ) class LocalModelInfo(BaseModel): """Discovered local model candidate.""" id: str = Field(..., description = "Identifier to use for loading/training") display_name: str = Field(..., description = "Display label") path: str = Field(..., description = "Local path where model data was discovered") source: Literal["models_dir", "hf_cache", "lmstudio", "custom"] = Field( ..., description = "Discovery source", ) model_id: Optional[str] = Field( None, description = "HF repo id for cached models, e.g. org/model", ) model_format: Optional[str] = Field( None, description = "Detected weights format ('gguf' when known). Lets the UI " "classify scanned folders whose name lacks a -GGUF suffix.", ) updated_at: Optional[float] = Field( None, description = "Unix timestamp of latest observed update", ) class LocalModelListResponse(BaseModel): """Response schema for listing local/cached models.""" models_dir: str = Field(..., description = "Directory scanned for custom local models") hf_cache_dir: Optional[str] = Field( None, description = "HF cache root that was scanned", ) lmstudio_dirs: List[str] = Field( default_factory = list, description = "LM Studio model directories that were scanned", ) models: List[LocalModelInfo] = Field( default_factory = list, description = "Discovered local/cached models", ) class AddScanFolderRequest(BaseModel): """Request body for adding a custom scan folder.""" path: str = Field(..., description = "Absolute or relative directory path to scan for models") class ScanFolderInfo(BaseModel): """A registered custom model scan folder.""" id: int = Field(..., description = "Database row ID") path: str = Field(..., description = "Normalized absolute path") created_at: str = Field(..., description = "ISO 8601 creation timestamp") class BrowseEntry(BaseModel): """A directory entry surfaced by the folder browser.""" name: str = Field(..., description = "Entry name (basename, not full path)") has_models: bool = Field( False, description = ( "Hint that the directory likely contains models " "(*.gguf, *.safetensors, config.json, or HF-style " "`models--*` subfolders). Used by the UI to highlight " "promising candidates; the scanner itself is authoritative." ), ) hidden: bool = Field( False, description = "Name starts with a dot (e.g. `.cache`)", ) class BrowseFoldersResponse(BaseModel): """Response schema for the folder browser endpoint.""" current: str = Field(..., description = "Absolute path of the directory just listed") parent: Optional[str] = Field( None, description = ( "Parent directory of `current`, or null if `current` is the " "filesystem root. The frontend uses this to render an `Up` row." ), ) entries: List[BrowseEntry] = Field( default_factory = list, description = ( "Subdirectories of `current`. Sorted with model-bearing " "directories first, then alphabetically case-insensitive; " "hidden entries come last within each group." ), ) suggestions: List[str] = Field( default_factory = list, description = ( "Handy starting points (home, HF cache, already-registered " "scan folders). Rendered as quick-pick chips above the list." ), ) truncated: bool = Field( False, description = ( "True when the listing was capped because the directory had " "more subfolders than the server is willing to enumerate in " "one request. The UI should show a hint telling the user to " "narrow their path." ), ) model_files_here: int = Field( 0, description = ( "Count of GGUF/safetensors files immediately inside " "``current``. Used by the UI to surface a hint on leaf " "model directories (which otherwise look `empty` because " "they contain only files, no subdirectories)." ), )