src/common/schemas.py

# Stage 간 JSON 인터페이스 (계약)
# 변경 시 3인 합의 필수

from pydantic import BaseModel, Field


# --- Stage 1 Models ---

class Models(BaseModel):
    diarization: str
    asr: str
    language_id: str
    alignment: str


class ProcessingInfo(BaseModel):
    processing_time_sec: float = Field(gt=0)
    models: Models
    device: str  # "cpu" | "cuda" | "cuda:N"
    language_chosen: str | None = None  # "ko" | "en", text-based decision used for alignment
    korean_ratio: float | None = Field(default=None, ge=0.0, le=1.0)


class Segment(BaseModel):
    segment_id: int = Field(ge=0)
    speaker_id: str
    start: float = Field(ge=0.0)
    end: float = Field(gt=0.0)
    text: str
    language: str  # "ko" | "en"
    audio_path: str
    confidence: float = Field(ge=0.0, le=1.0)


class Stage1Output(BaseModel):
    call_id: str
    duration: float = Field(gt=0.0)
    speakers: list[str]  # ["speaker_0", "speaker_1"]
    audio_path: str  # path to original audio file
    segments: list[Segment]
    processing_info: ProcessingInfo


class EmotionResult(BaseModel):
    speaker_id: str
    segment_id: int
    audio_emotion: str
    audio_confidence: float
    text_emotion: str
    text_confidence: float
    fused_emotion: str
    fused_confidence: float


class SpeakerSummary(BaseModel):
    dominant_emotion: str
    emotion_distribution: dict[str, float]
    avg_confidence: float


class Stage2Output(BaseModel):
    call_id: str
    emotions: list[EmotionResult]
    speaker_summaries: dict[str, SpeakerSummary]


class CharacterReaction(BaseModel):
    speaker_id: str
    solo_state: str
    pair_state: str
    role: str | None = None  # "giver" | "receiver" | None
    # Comforting strength (1..3). Default 2 preserves the P0 baseline (3 cycles
    # to heal). Higher = stronger comfort → receiver heals in fewer cycles.
    # Client mapping lives in app/src/hooks/healingThreshold.ts.
    intensity: int = Field(default=2, ge=1, le=3)


class GardenUpdate(BaseModel):
    growth_delta: int
    total_level: int
    mood: str


class RecapCard(BaseModel):
    title: str
    summary: str
    highlights: list[str]


class Stage3Output(BaseModel):
    call_id: str
    character_reactions: list[CharacterReaction]
    garden_update: GardenUpdate
    recap_card: RecapCard
    # Stage2 data surfaced at top level for frontend convenience
    emotions: list[EmotionResult] = Field(default_factory=list)
    speaker_summaries: dict[str, SpeakerSummary] = Field(default_factory=dict)