src/customer_intelligence/ml/schemas.py

"""Pydantic request/response schemas for ML prediction endpoints."""

from pydantic import BaseModel, Field


class ChurnPredictRequest(BaseModel):
    customer_unique_id: str


class TopFactor(BaseModel):
    feature: str
    impact: float  # SHAP value: positive = increases churn risk, negative = reduces it


class ChurnPredictResponse(BaseModel):
    customer_unique_id: str
    churn_probability: float = Field(ge=0.0, le=1.0)
    churn_label: bool
    model_version: str
    top_factors: list[TopFactor] | None = None


class SegmentPredictRequest(BaseModel):
    customer_unique_id: str


class SegmentPredictResponse(BaseModel):
    customer_unique_id: str
    segment_id: int
    segment_name: str
    model_version: str


class BatchChurnRequest(BaseModel):
    customer_ids: list[str] = Field(max_length=1000)


class BatchChurnResponse(BaseModel):
    predictions: list[ChurnPredictResponse]