from typing import Literal, Optional

from pydantic import BaseModel, Field


class RobotPlan(BaseModel):
    steps: list[str] = Field(
        description="Последовательность высокоуровневых действий робота",
        min_length=1
    )


class ChatMessage(BaseModel):
    role: Literal["system",
                  "user",
                  "assistant"]

    content: str


class ChatRecord(BaseModel):
    messages: list[ChatMessage]

class Prompts(BaseModel):
    system_prompt: str
    instructions: str
    user_prompt: str


class ModelResponse(BaseModel):
    model_response: str
    assistant_content: Optional[str] = None

class ModelInfo(BaseModel):
    provider: str
    model: str

class GenerationSettings(BaseModel):
    temperature: float
    max_tokens: int

class BaseDatasetRecord(BaseModel):
    prompts: Prompts
    response: ModelResponse
    model_info: ModelInfo
    generation: GenerationSettings

class RobotDatasetRecord(BaseDatasetRecord):
    task: str