mediaproc/core/schema/models/detect_jobs.py

"""
Detection Job and Checkpoint Schema Definitions

Source of truth for detection pipeline job tracking and stage checkpoints.
Follows the TranscodeJob/ChunkJob pattern.
"""

from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional
from uuid import UUID


class DetectJobStatus(str, Enum):
    PENDING = "pending"
    RUNNING = "running"
    PAUSED = "paused"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class RunType(str, Enum):
    INITIAL = "initial"
    REPLAY = "replay"
    RETRY = "retry"


@dataclass
class DetectJob:
    """
    A detection pipeline job.

    Each invocation of the pipeline (initial run, replay, retry) creates a DetectJob.
    Jobs for the same source video are linked via parent_job_id.
    """

    id: UUID

    # Input
    source_asset_id: UUID
    video_path: str
    profile_name: str = "soccer_broadcast"

    # Run lineage
    parent_job_id: Optional[UUID] = None  # links all runs for the same source
    run_type: RunType = RunType.INITIAL
    replay_from_stage: Optional[str] = None  # null for initial runs
    config_overrides: Dict[str, Any] = field(default_factory=dict)

    # Status
    status: DetectJobStatus = DetectJobStatus.PENDING
    current_stage: Optional[str] = None
    progress: float = 0.0
    error_message: Optional[str] = None

    # Results summary
    total_detections: int = 0
    brands_found: int = 0
    cloud_llm_calls: int = 0
    estimated_cost_usd: float = 0.0

    # Worker tracking
    celery_task_id: Optional[str] = None
    priority: int = 0

    # Timestamps
    created_at: Optional[datetime] = None
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None


@dataclass
class Timeline:
    """
    The frame sequence from a source video.

    Independent of stages — exists before any stage runs.
    Stages annotate the timeline, they don't own it.
    Frames are stored in MinIO as JPEGs.
    """

    id: UUID
    source_asset_id: Optional[UUID] = None
    source_video: str = ""
    profile_name: str = ""
    fps: float = 2.0

    # Frame metadata (images in MinIO, metadata here)
    frames_prefix: str = ""  # s3: timelines/{id}/frames/
    frames_manifest: Dict[int, str] = field(default_factory=dict)  # seq → s3 key
    frames_meta: List[Dict[str, Any]] = field(default_factory=list)

    created_at: Optional[datetime] = None


@dataclass
class Checkpoint:
    """
    A snapshot of pipeline state on a timeline.

    Stage outputs stored as JSONB — each stage serializes to JSON,
    the checkpoint stores it without knowing the shape.

    parent_id forms a tree: multiple children from the same parent
    = different config tries from the same starting point.
    """

    id: UUID
    timeline_id: UUID
    parent_id: Optional[UUID] = None  # null = root checkpoint

    # Stage outputs — JSONB per stage, opaque to the checkpoint layer
    stage_outputs: Dict[str, Any] = field(default_factory=dict)

    # Config that produced this checkpoint
    config_overrides: Dict[str, Any] = field(default_factory=dict)

    # Pipeline state
    stats: Dict[str, Any] = field(default_factory=dict)

    # Scenario bookmark
    is_scenario: bool = False
    scenario_label: str = ""

    created_at: Optional[datetime] = None


class BrandSource(str, Enum):
    """How a brand was first identified."""
    OCR = "ocr"
    VLM = "local_vlm"
    CLOUD = "cloud_llm"
    MANUAL = "manual"  # user-added via UI


@dataclass
class KnownBrand:
    """
    A brand discovered or registered in the system.

    Global — not per-source. Accumulates across all pipeline runs.
    Aliases enable fuzzy matching without re-escalating to VLM.
    """

    id: UUID
    canonical_name: str              # normalized display name
    aliases: List[str] = field(default_factory=list)  # known spellings/variants
    first_source: BrandSource = BrandSource.OCR
    total_occurrences: int = 0
    confirmed: bool = False          # manually confirmed by user

    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None


@dataclass
class SourceBrandSighting:
    """
    A brand seen in a specific source (video/asset).

    Per-source session cache — avoids re-escalating the same brand
    on subsequent frames or re-runs of the same source.
    """

    id: UUID
    source_asset_id: UUID            # the video this sighting belongs to
    brand_id: UUID                   # FK to KnownBrand
    brand_name: str                  # denormalized for fast lookup
    first_seen_timestamp: float = 0.0
    last_seen_timestamp: float = 0.0
    occurrences: int = 0
    detection_source: BrandSource = BrandSource.OCR
    avg_confidence: float = 0.0

    created_at: Optional[datetime] = None