src/models.py - master - Anna_brv/ASR-Service

src/models.py
210 строк8 KB
Anna
first_commit
30 апр 2026, 21:52
30 апр 2026, 21:52283885c
"""Pydantic models for API requests and responses."""

from datetime import datetime
from enum import StrEnum
from typing import Any, List, Optional

from pydantic import BaseModel, Field


class TaskStatusEnum(StrEnum):
    """Enum for task status."""

    pending = "pending"
    queued = "queued"
    processing = "processing"
    completed = "completed"
    failed = "failed"
    cancelled = "cancelled"


class TranscribeRequest(BaseModel):
    """Request model for transcription job."""

    model_name: Optional[str] = Field(
        default=None,
        description="name of the whisperx model to use (e.g., 'large-v2'). If not specified, default model will be used.",
    )
    diarize: bool = Field(
        default=False,
        description="enable speaker diarization (requires pyannote model).",
    )
    language: Optional[str] = Field(
        default="ru",
        description="audio language (ISO 639-1 code).",
    )


class TaskResponse(BaseModel):
    """Response model for task creation."""

    task_id: str = Field(..., description="unique identifier of the transcription task")
    status: TaskStatusEnum = Field(
        ..., description="current status of the task"
    )
    message: str = Field(..., description="status message")
    created_at: datetime = Field(..., description="task creation timestamp")


class TaskStatusResponse(BaseModel):
    """Response model for task status."""

    status: TaskStatusEnum = Field(..., description="current status of the task")
    progress: Optional[int] = Field(
        None, description="progress percentage (0-100), available during processing"
    )
    message: Optional[str] = Field(None, description="additional status message")


class TranscriptionSpeaker(BaseModel):
    """Model for a single speaker."""

    id: int = Field(..., description="unique identifier of the speaker")
    speaker: str = Field(..., description="speaker label (e.g., 'SPEAKER_00')")


class TranscriptionSegment(BaseModel):
    """Model for a single transcript segment."""

    start: float = Field(..., description="start time in seconds")
    end: float = Field(..., description="end time in seconds")
    text: str = Field(..., description="transcribed text for this segment")
    speaker: Optional[TranscriptionSpeaker] = Field(None, description="speaker label if diarization was used")


class TranscriptionResult(BaseModel):
    """Response model for transcription result."""

    task_id: str = Field(..., description="unique identifier of the task")
    title: Optional[str] = Field(None, description="title of the transcription")
    duration: Optional[float] = Field(None, description="duration of the transcription in seconds")
    status: TaskStatusEnum = Field(..., description="current status of the task")
    lang: str = Field(..., description="detected language of the transcription")
    speakers: List[TranscriptionSpeaker] = Field(default_factory=list, description="list of speaker objects")
    segments: List[TranscriptionSegment] = Field(..., description="list of transcript segments")
    full_text: str = Field(..., description="complete transcribed text")
    created_at: datetime = Field(..., description="task creation timestamp")
    updated_at: Optional[datetime] = Field(None, description="last update timestamp")


class TaskResult(BaseModel):
    """Result model for ASR service transcription."""

    text: str = Field(..., description="full transcription text")
    segments: List[dict] = Field(default_factory=list, description="segments with timestamps")
    speakers: Optional[List[dict]] = Field(None, description="speaker information if diarization enabled")
    duration: Optional[float] = Field(None, description="audio duration in seconds")
    model_name: str = Field(..., description="model name used for transcription")


class ModelsResponse(BaseModel):
    """Response model for available models."""

    models: List[str] = Field(..., description="list of available model names")


class ModelLoadResponse(BaseModel):
    """Response model for model loading."""

    model_name: str = Field(..., description="name of the loaded model")
    status: str = Field(..., description="loading status ('loaded', 'already_loaded', 'failed')")
    message: Optional[str] = Field(None, description="additional information about the operation")


class TaskInfo(BaseModel):
    """Model for task information."""

    task_id: str = Field(..., description="unique identifier of the task")
    status: TaskStatusEnum = Field(..., description="current status")
    created_at: datetime = Field(..., description="task creation timestamp")
    started_at: Optional[datetime] = Field(None, description="task start timestamp")
    ended_at: Optional[datetime] = Field(None, description="task end timestamp")
    result: Optional[TranscriptionResult] = Field(None, description="transcription result")
    error: Optional[str] = Field(None, description="error message if failed")
    file_name: Optional[str] = Field(None, description="original file name")
    updated_at: Optional[datetime] = Field(None, description="last update timestamp")


class TasksListResponse(BaseModel):
    """Response model for tasks list."""

    tasks: List[TaskInfo] = Field(..., description="list of tasks")
    total: int = Field(..., description="total number of tasks")
    page: int = Field(1, description="current page number")
    page_size: int = Field(20, description="page size")


class CancelTaskResponse(BaseModel):
    """Response model for task actions."""

    task_id: str = Field(..., description="unique identifier of the task")
    status: str = Field(..., description="result of the action")
    message: str = Field(..., description="additional information")


class RestartTaskResponse(BaseModel):
    """Response model for task restart."""

    task_id: str = Field(..., description="unique identifier of the task")
    new_task_id: str = Field(..., description="new task identifier")
    status: TaskStatusEnum = Field(..., description="status of the new task")


class ClearQueueResponse(BaseModel):
    """Response model for queue clearing."""

    cleared_count: int = Field(..., description="number of tasks cleared")
    status: str = Field(..., description="result of the operation")


class ModelInfo(BaseModel):
    """Model for model information."""

    name: str = Field(..., description="model name")
    loaded: bool = Field(..., description="whether model is loaded in cache")
    is_default: bool = Field(..., description="whether this is the default model")
    supports_diarization: bool = Field(..., description="whether model supports diarization")


class ModelListResponse(BaseModel):
    """Response model for model list."""

    models: List[ModelInfo] = Field(..., description="list of models")
    default_model: str = Field(..., description="default model name")


class LoadModelRequest(BaseModel):
    """Request model for loading a model."""

    model_name: str = Field(..., description="name of the model to load")


class LoadModelResponse(BaseModel):
    """Response model for model loading."""

    model_name: str = Field(..., description="name of the loaded model")
    status: str = Field(..., description="loading status")
    message: str = Field(..., description="status message")


class UnloadModelResponse(BaseModel):
    """Response model for model unloading."""

    model_name: str = Field(..., description="name of the unloaded model")
    status: str = Field(..., description="unloading status")
    message: str = Field(..., description="status message")


class HealthResponse(BaseModel):
    """Response model for health check."""

    status: str = Field(..., description="service status")
    version: str = Field(..., description="service version")
    redis_connected: bool = Field(..., description="whether Redis is connected")
    queue_size: int = Field(..., description="number of tasks in queue")


class ErrorResponse(BaseModel):
    """Error response model."""

    detail: str = Field(..., description="error description")