From 7ee4381936c7ed9e6beb160b3c736833b7b0c2fb Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 11 Feb 2026 12:39:58 +0000 Subject: [PATCH] move inline pydantic models - fix existing models routes integration --- studio/backend/models/__init__.py | 38 +++++++++++++++- studio/backend/models/inference.py | 54 +++++++++++++++++++++++ studio/backend/models/responses.py | 38 ++++++++++++++++ studio/backend/models/users.py | 29 +++--------- studio/backend/routes/inference.py | 71 ++++++------------------------ studio/backend/routes/models.py | 24 +++++----- studio/backend/routes/training.py | 38 ++++++++-------- 7 files changed, 179 insertions(+), 113 deletions(-) create mode 100644 studio/backend/models/inference.py create mode 100644 studio/backend/models/responses.py diff --git a/studio/backend/models/__init__.py b/studio/backend/models/__init__.py index 54696c9f99..584fb9f7c2 100644 --- a/studio/backend/models/__init__.py +++ b/studio/backend/models/__init__.py @@ -16,8 +16,28 @@ from .models import ( from .auth import ( AuthSetupRequest, AuthLoginRequest, + RefreshTokenRequest, AuthStatusResponse, ) +from .users import Token +from .datasets import ( + CheckFormatRequest, + CheckFormatResponse, +) +from .inference import ( + LoadRequest, + UnloadRequest, + GenerateRequest, + LoadResponse, + UnloadResponse, + InferenceStatusResponse, +) +from .responses import ( + TrainingStopResponse, + TrainingMetricsResponse, + LoRABaseModelResponse, + VisionCheckResponse, +) __all__ = [ # Training schemas @@ -33,6 +53,22 @@ __all__ = [ # Auth schemas "AuthSetupRequest", "AuthLoginRequest", + "RefreshTokenRequest", "AuthStatusResponse", + "Token", + # Dataset schemas + "CheckFormatRequest", + "CheckFormatResponse", + # Inference schemas + "LoadRequest", + "UnloadRequest", + "GenerateRequest", + "LoadResponse", + "UnloadResponse", + "InferenceStatusResponse", + # Response schemas + "TrainingStopResponse", + "TrainingMetricsResponse", + "LoRABaseModelResponse", + "VisionCheckResponse", ] - diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py new file mode 100644 index 0000000000..3917e38e36 --- /dev/null +++ b/studio/backend/models/inference.py @@ -0,0 +1,54 @@ +""" +Pydantic schemas for Inference API +""" +from pydantic import BaseModel, Field +from typing import Optional, List + + +class LoadRequest(BaseModel): + """Request to load a model for inference""" + model_path: str = Field(..., description="Model identifier or local path") + hf_token: Optional[str] = Field(None, description="HuggingFace token for gated models") + 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") + + +class UnloadRequest(BaseModel): + """Request to unload a model""" + model_path: str = Field(..., description="Model identifier to unload") + + +class GenerateRequest(BaseModel): + """Request for text generation""" + messages: List[dict] = Field(..., description="Chat messages in OpenAI format") + system_prompt: str = Field("You are a helpful AI assistant.", description="System prompt") + temperature: float = Field(0.7, ge=0.0, le=2.0, description="Sampling temperature") + top_p: float = Field(0.9, ge=0.0, le=1.0, description="Top-p sampling") + top_k: int = Field(40, ge=1, le=100, description="Top-k sampling") + max_new_tokens: int = Field(512, ge=1, le=4096, description="Maximum tokens to generate") + repetition_penalty: float = Field(1.1, ge=1.0, le=2.0, description="Repetition penalty") + image_base64: Optional[str] = Field(None, description="Base64 encoded image for vision models") + + +class LoadResponse(BaseModel): + """Response after loading a model""" + status: str = Field(..., description="Load status") + model: str = Field(..., description="Model identifier") + 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") + + +class UnloadResponse(BaseModel): + """Response after unloading a model""" + status: str = Field(..., description="Unload status") + model: str = Field(..., description="Model identifier that was unloaded") + + +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") + 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/responses.py b/studio/backend/models/responses.py new file mode 100644 index 0000000000..2aa798c5c9 --- /dev/null +++ b/studio/backend/models/responses.py @@ -0,0 +1,38 @@ +""" +Pydantic response schemas for endpoints that previously returned raw dicts. +These are small response models for training and model management routes. +""" +from pydantic import BaseModel, Field +from typing import Optional, List + + +# --- Training route response models --- + +class TrainingStopResponse(BaseModel): + """Response for stopping a training job""" + status: str = Field(..., description="Current status: 'stopped' or 'idle'") + message: str = Field(..., description="Human-readable status message") + + +class TrainingMetricsResponse(BaseModel): + """Response for training metrics history""" + loss_history: List[float] = Field(default_factory=list, description="Loss values per step") + lr_history: List[float] = Field(default_factory=list, description="Learning rate per step") + step_history: List[int] = Field(default_factory=list, description="Step numbers") + current_loss: Optional[float] = Field(None, description="Most recent loss value") + current_lr: Optional[float] = Field(None, description="Most recent learning rate") + current_step: Optional[int] = Field(None, description="Most recent step number") + + +# --- Model management route response models --- + +class LoRABaseModelResponse(BaseModel): + """Response for getting a LoRA's base model""" + lora_path: str = Field(..., description="Path to the LoRA adapter") + base_model: str = Field(..., description="Base model identifier") + + +class VisionCheckResponse(BaseModel): + """Response for checking if a model is a vision model""" + model_name: str = Field(..., description="Model identifier") + is_vision: bool = Field(..., description="Whether the model is a vision model") diff --git a/studio/backend/models/users.py b/studio/backend/models/users.py index 2f404211a7..c90e27942f 100644 --- a/studio/backend/models/users.py +++ b/studio/backend/models/users.py @@ -1,34 +1,15 @@ -"""Pydantic models for user-related API endpoints. +"""Pydantic models for authentication tokens. -This module defines the data models used for user authentication and management -in the FastAPI application. +This module defines the Token response model used by auth routes. """ from pydantic import BaseModel, Field -class User(BaseModel): - """Basic user model containing username.""" - - username: str - - -class UserInDB(BaseModel): - """User model with password for database storage.""" - - password: str - - class Token(BaseModel): """Authentication token model with access and refresh tokens.""" - access_token: str - refresh_token: str - token_type: str - - -class TokenData(BaseModel): - """Token payload model containing username.""" - - username: str | None = None + access_token: str = Field(..., description="JWT access token (60 min expiry)") + refresh_token: str = Field(..., description="Opaque refresh token (7 day expiry)") + token_type: str = Field(..., description="Token type, always 'bearer'") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index b80a163585..17d2011205 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5,8 +5,7 @@ import sys from pathlib import Path from fastapi import APIRouter, HTTPException from fastapi.responses import StreamingResponse -from pydantic import BaseModel, Field -from typing import Optional, List +from typing import Optional import json import logging @@ -26,6 +25,15 @@ except ImportError: from core.inference import get_inference_backend from utils.models import ModelConfig +from models.inference import ( + LoadRequest, + UnloadRequest, + GenerateRequest, + LoadResponse, + UnloadResponse, + InferenceStatusResponse, +) + router = APIRouter() logger = logging.getLogger(__name__) @@ -39,57 +47,6 @@ if not logger.handlers: logger.setLevel(logging.INFO) -# ============================================ -# Request/Response Models -# ============================================ - -class LoadRequest(BaseModel): - """Request to load a model for inference""" - model_path: str = Field(..., description="Model identifier or local path") - hf_token: Optional[str] = Field(None, description="HuggingFace token for gated models") - 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") - - -class UnloadRequest(BaseModel): - """Request to unload a model""" - model_path: str = Field(..., description="Model identifier to unload") - - -class GenerateRequest(BaseModel): - """Request for text generation""" - messages: List[dict] = Field(..., description="Chat messages in OpenAI format") - system_prompt: str = Field("You are a helpful AI assistant.", description="System prompt") - temperature: float = Field(0.7, ge=0.0, le=2.0, description="Sampling temperature") - top_p: float = Field(0.9, ge=0.0, le=1.0, description="Top-p sampling") - top_k: int = Field(40, ge=1, le=100, description="Top-k sampling") - max_new_tokens: int = Field(512, ge=1, le=4096, description="Maximum tokens to generate") - repetition_penalty: float = Field(1.1, ge=1.0, le=2.0, description="Repetition penalty") - image_base64: Optional[str] = Field(None, description="Base64 encoded image for vision models") - - -class LoadResponse(BaseModel): - """Response after loading a model""" - status: str - model: str - display_name: str - is_vision: bool - is_lora: bool - - -class StatusResponse(BaseModel): - """Current inference backend status""" - active_model: Optional[str] - is_vision: bool - loading: List[str] - loaded: List[str] - - -# ============================================ -# Routes -# ============================================ - @router.post("/load", response_model=LoadResponse) async def load_model(request: LoadRequest): """ @@ -147,7 +104,7 @@ async def load_model(request: LoadRequest): ) -@router.post("/unload") +@router.post("/unload", response_model=UnloadResponse) async def unload_model(request: UnloadRequest): """ Unload a model from memory. @@ -156,7 +113,7 @@ async def unload_model(request: UnloadRequest): backend = get_inference_backend() backend.unload_model(request.model_path) logger.info(f"Unloaded model: {request.model_path}") - return {"status": "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) @@ -239,7 +196,7 @@ async def generate_stream(request: GenerateRequest): ) -@router.get("/status", response_model=StatusResponse) +@router.get("/status", response_model=InferenceStatusResponse) async def get_status(): """ Get current inference backend status. @@ -252,7 +209,7 @@ async def get_status(): model_info = backend.models.get(backend.active_model_name, {}) is_vision = model_info.get("is_vision", False) - return StatusResponse( + return InferenceStatusResponse( active_model=backend.active_model_name, is_vision=is_vision, loading=list(getattr(backend, 'loading_models', set())), diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 713cbdc033..1bfd08f38c 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -7,8 +7,6 @@ from fastapi import APIRouter, Depends, HTTPException, Query from typing import List, Optional import logging -from pydantic import BaseModel - # Add backend directory to path backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: @@ -46,6 +44,8 @@ from models import ( LoRAInfo, ModelListResponse, ) +from models.responses import LoRABaseModelResponse, VisionCheckResponse + router = APIRouter() logger = logging.getLogger(__name__) @@ -207,7 +207,7 @@ async def scan_loras( ) -@router.get("/loras/{lora_path:path}/base-model") +@router.get("/loras/{lora_path:path}/base-model", response_model=LoRABaseModelResponse) async def get_lora_base_model( lora_path: str, current_subject: str = Depends(get_current_subject), @@ -226,10 +226,10 @@ async def get_lora_base_model( detail=f"Could not determine base model for LoRA: {lora_path}" ) - return { - "lora_path": lora_path, - "base_model": base_model - } + return LoRABaseModelResponse( + lora_path=lora_path, + base_model=base_model, + ) except HTTPException: raise @@ -241,7 +241,7 @@ async def get_lora_base_model( ) -@router.get("/check-vision/{model_name:path}") +@router.get("/check-vision/{model_name:path}", response_model=VisionCheckResponse) async def check_vision_model( model_name: str, current_subject: str = Depends(get_current_subject), @@ -254,10 +254,10 @@ async def check_vision_model( try: is_vision = is_vision_model(model_name) - return { - "model_name": model_name, - "is_vision": is_vision - } + return VisionCheckResponse( + model_name=model_name, + is_vision=is_vision, + ) except Exception as e: logger.error(f"Error checking vision model: {e}", exc_info=True) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 6e2fb3c355..9f6c298a4d 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -36,6 +36,7 @@ from models import ( TrainingStatus, TrainingProgress, ) +from models.responses import TrainingStopResponse, TrainingMetricsResponse router = APIRouter() logger = logging.getLogger(__name__) @@ -242,7 +243,7 @@ async def start_training( ) -@router.post("/stop") +@router.post("/stop", response_model=TrainingStopResponse) async def stop_training( current_subject: str = Depends(get_current_subject), ): @@ -253,18 +254,18 @@ async def stop_training( backend = get_training_backend() if not backend.is_training_active(): - return { - "status": "idle", - "message": "No training job is currently running" - } + return TrainingStopResponse( + status="idle", + message="No training job is currently running" + ) # Call backend stop method backend.stop_training() - return { - "status": "stopped", - "message": "Training job stopped successfully" - } + return TrainingStopResponse( + status="stopped", + message="Training job stopped successfully" + ) except Exception as e: logger.error(f"Error stopping training: {e}", exc_info=True) @@ -353,7 +354,7 @@ async def get_training_status( ) -@router.get("/metrics") +@router.get("/metrics", response_model=TrainingMetricsResponse) async def get_training_metrics( current_subject: str = Depends(get_current_subject), ): @@ -373,15 +374,14 @@ async def get_training_metrics( current_lr = lr_history[-1] if lr_history else None current_step = step_history[-1] if step_history else None - # Keep metrics as a simple JSON payload instead of a Pydantic model - return { - "loss_history": loss_history, - "lr_history": lr_history, - "step_history": step_history, - "current_loss": current_loss, - "current_lr": current_lr, - "current_step": current_step, - } + return TrainingMetricsResponse( + loss_history=loss_history, + lr_history=lr_history, + step_history=step_history, + current_loss=current_loss, + current_lr=current_lr, + current_step=current_step, + ) except Exception as e: logger.error(f"Error getting training metrics: {e}", exc_info=True)