# !!!!! DO NOT EVER CHANGE THIS FILE MANUALLY -- AUTOGENERATED by orm2pydantic.py !!!!!
# Generated on 2023-12-23 from module egp_api_backend.server.internal.orm.orm_models in the egp-api-backend package for use in the EGP python SDK
# To regenerate this file, run:
# scaleapi/packages/egp-api-backend/scripts/orm2pydantic.sh

from datetime import datetime
from typing import Any, Dict, Optional

from pydantic import Field
from scale_egp.sdk.models.model_schemas.model_enums import ModelEndpointType, ModelType, ModelVendor
from scale_egp.sdk.models.model_schemas.model_vendor_configuration import LaunchVendorConfiguration
from scale_egp.sdk.models.model_schemas.parameter_schema import ParameterSchema
from scale_egp.utils.model_utils import BaseModel, Entity


class ModelTemplateRequest(BaseModel):
    name: str
    endpoint_type: ModelEndpointType
    model_type: ModelType
    vendor_configuration: LaunchVendorConfiguration
    model_creation_parameters_schema: Optional[ParameterSchema] = Field(
        None,
        description="The field names and types of available parameter fields which may be specified during model creation",
    )
    model_request_parameters_schema: Optional[ParameterSchema] = Field(
        None,
        description="The field names and types of available parameter fields which may be specified in a model execution API's `model_request_parameters` field.",
    )
    account_id: str = Field(..., description="The ID of the account that owns the given entity.")


class ModelTemplate(Entity):
    """
    This is a template for types of models that can be quickly customized by end users.
    It allows users to upload static docker images that can run specific types of models.
    These docker images will expose parameters that can be injected at ModelAlias creation
    time to customize the functionality. A common example of this is to use a
    HuggingFace LLM template, but swap out model weights for a finetuned model.
    """

    name: str
    endpoint_type: ModelEndpointType
    model_type: ModelType
    vendor_configuration: LaunchVendorConfiguration
    model_creation_parameters_schema: Optional[ParameterSchema] = Field(
        None,
        description="The field names and types of available parameter fields which may be specified during model creation",
    )
    model_request_parameters_schema: Optional[ParameterSchema] = Field(
        None,
        description="The field names and types of available parameter fields which may be specified in a model execution API's `model_request_parameters` field.",
    )
    id: str = Field(..., description="The unique identifier of the entity.")
    created_at: datetime = Field(
        ..., description="The date and time when the entity was created in ISO format."
    )
    account_id: str = Field(..., description="The ID of the account that owns the given entity.")
    created_by_user_id: str = Field(..., description="The user who originally created the entity.")


class ModelEndpointRequest(BaseModel):
    name: str


class ModelEndpoint(Entity):
    name: str
    status: str
    id: str = Field(..., description="The unique identifier of the entity.")
    created_at: datetime = Field(
        ..., description="The date and time when the entity was created in ISO format."
    )
    created_by_user_id: str = Field(..., description="The user who originally created the entity.")


class ModelAliasRequest(BaseModel):
    name: str
    model_vendor: Optional[ModelVendor] = Field(None)
    base_model_id: Optional[str] = Field(None)
    description: Optional[str] = Field(None)
    model_template_id: Optional[str] = Field(None)
    model_creation_parameters: Optional[dict[str, Any]] = Field(None)
    account_id: str = Field(..., description="The ID of the account that owns the given entity.")


class ModelAlias(Entity):
    """
    Entity for all models, including both self-hosted and 3rd party, base- and fine-tuned models.
    """

    name: str
    model_vendor: Optional[ModelVendor] = Field(None)
    base_model_id: Optional[str] = Field(None)
    description: Optional[str] = Field(None)
    model_template_id: Optional[str] = Field(None)
    model_creation_parameters: Optional[dict[str, Any]] = Field(None)
    id: str = Field(..., description="The unique identifier of the entity.")
    created_at: datetime = Field(
        ..., description="The date and time when the entity was created in ISO format."
    )
    account_id: str = Field(..., description="The ID of the account that owns the given entity.")
    created_by_user_id: str = Field(..., description="The user who originally created the entity.")
    status: str = Field(..., description="Status of the model's deployment.")
    request_schema: Dict[str, Any] = Field(
        ..., description="JSON schema for the requests to the model"
    )
    response_schema: Dict[str, Any] = Field(
        ..., description="JSON schema for the response to the model"
    )


class ModelUsageRequest(BaseModel):
    model_id: str
    user_id: str
    started_at: datetime
    finished_at: Optional[datetime] = Field(None)


class ModelUsage(Entity):
    model_id: str
    user_id: str
    started_at: datetime
    finished_at: Optional[datetime] = Field(None)
    id: str = Field(..., description="The unique identifier of the entity.")
    created_at: datetime = Field(
        ..., description="The date and time when the entity was created in ISO format."
    )
