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

import datetime
import enum
from typing import Any, Optional

from pydantic import Field
from scale_egp.utils.model_utils import BaseModel, Entity, hiddenstr

ModelState = enum.StrEnum(
    "ModelState", [("ENABLED", "ENABLED"), ("PENDING", "PENDING"), ("DISABLED", "DISABLED")]
)
ModelVendor = enum.StrEnum(
    "ModelVendor",
    [
        ("OPENAI", "OPENAI"),
        ("COHERE", "COHERE"),
        ("ANTHROPIC", "ANTHROPIC"),
        ("LLMENGINE", "LLMENGINE"),
        ("OTHER", "OTHER"),
    ],
)
ModelEndpointType = enum.StrEnum(
    "ModelEndpointType",
    [("SYNC", "SYNC"), ("ASYNC", "ASYNC"), ("STREAMING", "STREAMING"), ("BATCH", "BATCH")],
)
ModelType = enum.StrEnum(
    "ModelType",
    [
        ("COMPLETION", "COMPLETION"),
        ("CHAT_COMPLETION", "CHAT_COMPLETION"),
        ("AGENT", "AGENT"),
        ("EMBEDDING", "EMBEDDING"),
        ("RERANKING", "RERANKING"),
    ],
)


class ModelRequest(BaseModel):
    name: str = Field(...)
    model_state: ModelState = Field(...)
    vendor: ModelVendor = Field(...)
    vendor_name: str = Field(...)
    base_model_id: Optional[str] = Field(None)


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

    name: str = Field(...)
    model_state: ModelState = Field(...)
    vendor: ModelVendor = Field(...)
    vendor_name: str = Field(...)
    base_model_id: Optional[str] = Field(None)
    id: str = Field(..., description="The unique identifier of the entity.")
    created_at: datetime.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 ScaleModelTemplateRequest(BaseModel):
    name: str = Field(...)
    endpoint_type: ModelEndpointType = Field(...)
    model_type: ModelType = Field(...)
    weights_uri: str = Field(...)
    docker_image: str = Field(...)
    docker_env_vars: Optional[dict[str, Any]] = Field(None)
    launch_parameters: dict[str, Any] = Field(...)


class ScaleModelTemplate(Entity):
    name: str = Field(...)
    endpoint_type: ModelEndpointType = Field(...)
    model_type: ModelType = Field(...)
    weights_uri: str = Field(...)
    docker_image: str = Field(...)
    docker_env_vars: Optional[dict[str, Any]] = Field(None)
    launch_parameters: dict[str, Any] = Field(...)
    id: str = Field(..., description="The unique identifier of the entity.")
    created_at: datetime.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 ScaleModelInstanceRequest(BaseModel):
    model_id: str = Field(...)
    model_template_id: str = Field(...)
    endpoint_url: Optional[hiddenstr] = Field(None)


class ScaleModelInstance(Entity):
    model_id: str = Field(...)
    model_template_id: str = Field(...)
    id: str = Field(..., description="The unique identifier of the entity.")
    created_at: datetime.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 ModelUsageRequest(BaseModel):
    model_id: str = Field(...)
    user_id: str = Field(...)
    started_at: datetime.datetime = Field(...)
    finished_at: Optional[datetime.datetime] = Field(None)


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