from pydantic import Field

from .base import BaseModel
from .harm import HarmCategory, HarmProbability
from .part import TextPart


class Content(BaseModel):
    role: str = Field(..., description="Role of the author of the content.")

    parts: list[TextPart] = Field(..., description="Parts of the content.")


class SafetyRating(BaseModel):
    category: HarmCategory = Field(..., description="The category of the content.")

    probability: HarmProbability = Field(
        ..., description="The probability of the content falling into the category."
    )


class GenerationCandidate(BaseModel):
    index: int = Field(..., description="Index of the candidate in the list.")

    content: Content | None = Field(
        None, description="Content generated by the model."
    )

    finish_reason: str | None = Field(
        None, alias="finishReason", description="The reason the generation stopped."
    )

    safety_ratings: list[SafetyRating] | None = Field(
        None, alias="safetyRatings", description="Safety ratings of the content."
    )


class GenerationResponse(BaseModel):
    candidates: list[GenerationCandidate] = Field(
        [], description="The generated responses."
    )
