schema.py

from pydantic import BaseModel


class TrainRequest(BaseModel):

    agent_id: str
    file_url: str
    file_type: str | None = None
    file_name: str | None = None
    agent_type: str | None = None


class AskRequest(BaseModel):

    agent_id: str
    message: str