# generated by datamodel-codegen:
#   filename:  artifacts.json
#   timestamp: 2025-01-03T09:11:54+00:00

from __future__ import annotations

from enum import Enum
from typing import Any, Dict, List, Literal, Optional, Union

from truefoundry.pydantic_v1 import BaseModel, Field, conint, constr


class Agent(BaseModel):
    type: Literal["agent"] = Field(..., description="+value=agent")
    name: constr(regex=r"^[a-zA-Z][a-zA-Z0-9\-]{1,30}[a-zA-Z0-9]$") = Field(
        ..., description="Names cannot repeat accross agent\n+uiType=Hidden"
    )
    available_tools: List[constr(min_length=1)] = Field(
        ..., description="+sort=30\n+uiType=Hidden"
    )
    goal: constr(min_length=1, max_length=128) = Field(
        ..., description="+sort=10\n+uiType=TextArea"
    )
    instruction: constr(min_length=1, max_length=2620) = Field(
        ...,
        description='`instruction` is the system prompt for now. (2.5 * 1024)\n+sort=20\n+uiType=AgentInstructions\n+uiProps={"helpText":"Use the syntax ${Tool FQN} to reference a tool, and ${AGENT FQN} to reference another agent"}',
    )
    model_id: constr(min_length=1) = Field(
        ...,
        description='+sort=40\n+uiType=EnabledModelSelector\n+uiProps={"searchable":true,"modelType":"chat","providerType":"openai"}',
    )


class Method(str, Enum):
    """
    +sort=50
    +uiType=Hidden
    """

    get = "get"
    post = "post"
    put = "put"
    delete = "delete"
    patch = "patch"


class AgentWithFQN(Agent):
    id: str
    fqn: str


class BaseArtifactVersion(BaseModel):
    description: Optional[constr(max_length=512)] = Field(
        None,
        description="+label=Description\n+usage=Description of the artifact or model version\n+docs=Description of the artifact or model version",
    )
    metadata: Dict[str, Any] = Field(
        ...,
        description="+label=Metadata\n+docs=Metadata for the artifact or model version\n+usage=Metadata for the artifact or model version\n+uiType=JsonInput",
    )
    version_alias: Optional[
        constr(regex=r"^v[a-zA-Z0-9.-]*([a-zA-Z0-9]+)$", max_length=128)
    ] = Field(
        None,
        description="+label=Version Alias\n+usage=The version alias for artifact or model version which should start with 'v' followed by alphanumeric and it can include '.' and '-' in between (e.g. v1.0.0, v-prod, v-dev, etc)\n+docs=The version alias for artifact or model version which should start with 'v' followed by alphanumeric and it can include '.' and '-' in between (e.g. v1.0.0, v-prod, v-dev, etc)\n+message=The version alias should start with 'v' followed by alphanumeric and it can include '.' and '-' in between (e.g. v1.0.0, v-prod, v-dev, etc)",
    )


class MimeType(str, Enum):
    """
    +label=MIME Type
    +usage=MIME type of the content
    """

    text_plain = "text/plain"
    application_json = "application/json"
    image_png = "image/png"
    image_jpeg = "image/jpeg"
    application_x_directory = "application/x-directory"


class BlobStorageReference(BaseModel):
    """
    +usage=Blob Storage Location
    +label=Blob Storage Location
    +docs=Defines the structure for blob storage content references, including type, path, and mime_type
    """

    type: Literal["blob-storage"]
    path: constr(regex=r"^.{1,}$") = Field(
        ..., description="+label=Path\n+usage=Path to the content in blob storage"
    )
    mime_type: MimeType = Field(
        ..., description="+label=MIME Type\n+usage=MIME type of the content"
    )


class ExternalArtifactSource(BaseModel):
    """
    +label=External
    """

    type: Literal["external"] = Field(
        ...,
        description="+label=Type\n+usage=Type of the source\n+value=external\n+type=External",
    )
    uri: str = Field(..., description="+label=URI\n+usage=URI of the Artifact source")


class FastAIFramework(BaseModel):
    """
    +docs=FastAI framework for the model version
    +label=FastAI
    """

    type: Literal["fastai"] = Field(
        ..., description="+label=Type\n+usage=Type of the framework\n+value=fastai"
    )


class GluonFramework(BaseModel):
    """
    +docs=Gluon framework for the model version
    +label=Gluon
    """

    type: Literal["gluon"] = Field(
        ..., description="+label=Type\n+usage=Type of the framework\n+value=gluon"
    )


class H2OFramework(BaseModel):
    """
    +docs=H2O framework for the model version
    +label=H2O
    """

    type: Literal["h2o"] = Field(
        ..., description="+label=Type\n+usage=Type of the framework\n+value=h2o"
    )


class ImageUrl(BaseModel):
    """
    +label=URL for the image
    +usage=The URL for the image, must be a non-empty valid HTTPS URL or a data URL
    """

    url: Union[
        constr(
            regex=r"^\b((https?://)?(?:www\.)?[-a-zA-Z0-9@:%._\+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}(?:[-a-zA-Z0-9()@:%_\+.~#?&/=]*)|data:image/[a-zA-Z]+;base64,[a-zA-Z0-9+/=\s]+)$"
        ),
        BlobStorageReference,
    ] = Field(
        ...,
        description="+label=URL for the image\n+usage=The URL for the image, must be a non-empty valid HTTPS URL or a data URL",
    )


class ImageContentPart(BaseModel):
    """
    +usage=Image URL
    +label=Image URL
    """

    type: Literal["image_url"]
    image_url: ImageUrl = Field(
        ...,
        description="+label=URL for the image\n+usage=The URL for the image, must be a non-empty valid HTTPS URL or a data URL",
    )


class KerasFramework(BaseModel):
    """
    +docs=Keras framework for the model version
    +label=Keras
    """

    type: Literal["keras"] = Field(
        ..., description="+label=Type\n+usage=Type of the framework\n+value=keras"
    )


class LightGBMFramework(BaseModel):
    """
    +docs=LightGBM framework for the model version
    +label=LightGBM
    """

    type: Literal["lightgbm"] = Field(
        ..., description="+label=Type\n+usage=Type of the framework\n+value=lightgbm"
    )


class Parameters(BaseModel):
    """
    +usage=Parameters for the provider
    +label=Parameters
    +docs=Key-value pairs to store additional parameters for the provider
    """

    max_tokens: Optional[int] = None
    temperature: Optional[float] = None
    top_k: Optional[float] = None
    top_p: Optional[float] = None
    stop: Optional[Union[List[str], str]] = None


class ModelConfiguration(BaseModel):
    """
    +label=Model Configuration
    +icon=fa-cogs:#326ce5
    +message=Configuration details for the provider and corresponding config
    +usage=Define the provider and its configuration
    +docs=Configuration settings specific to the provider, including model settings and other parameters.
    """

    provider: str = Field(
        ...,
        description='+label=Provider Name\n+usage=Name of the provider, must be non-empty, e.g., "openai", "google_gemini"',
    )
    model: constr(regex=r"^[a-zA-Z][a-zA-Z0-9\-]{1,30}[a-zA-Z0-9]$") = Field(
        ..., description="+label=Model name\n+usage=Model name, must be non-empty"
    )
    parameters: Optional[Parameters] = Field(
        None,
        description="+usage=Parameters for the provider\n+label=Parameters\n+docs=Key-value pairs to store additional parameters for the provider",
    )
    extra_parameters: Optional[Dict[str, Any]] = Field(
        None,
        description="+usage=Extra parameters for the provider\n+label=Extra Parameters\n+docs=Additional parameters for the provider",
    )


class ModelVersionEnvironment(BaseModel):
    """
    +label=Environment
    """

    python_version: Optional[constr(regex=r"^\d+(\.\d+){1,2}([\-\.a-z0-9]+)?$")] = (
        Field(
            None,
            description="+label=Python Version\n+usage=Python version for the model version",
        )
    )
    pip_packages: Optional[List[str]] = Field(
        None,
        description="+label=PIP Packages\n+usage=PIP packages for the model version",
    )


class ONNXFramework(BaseModel):
    """
    +docs=ONNX framework for the model version
    +label=ONNX
    """

    type: Literal["onnx"] = Field(
        ..., description="+label=Type\n+usage=Type of the framework\n+value=onnx"
    )


class PaddleFramework(BaseModel):
    """
    +docs=PaddlePaddle framework for the model version
    +label=Paddle
    """

    type: Literal["paddle"] = Field(
        ..., description="+label=Type\n+usage=Type of the framework\n+value=paddle"
    )


class PyTorchFramework(BaseModel):
    """
    +docs=PyTorch framework for the model version
    +label=PyTorch
    """

    type: Literal["pytorch"] = Field(
        ..., description="+label=Type\n+usage=Type of the framework\n+value=pytorch"
    )


class InferMethodName(str, Enum):
    """
    +label=Inference Method Name
    +usage=Name of the method used for inference
    """

    predict = "predict"
    predict_proba = "predict_proba"


class SklearnModelSchema(BaseModel):
    """
    +label=Sklearn Model Schema
    """

    infer_method_name: InferMethodName = Field(
        ...,
        description="+label=Inference Method Name\n+usage=Name of the method used for inference",
    )
    inputs: List[Dict[str, Any]] = Field(
        ..., description="+label= Input Schema\n+usage=Schema of the input"
    )
    outputs: List[Dict[str, Any]] = Field(
        ..., description="+label= Output Schema\n+usage=Schema of the output"
    )


class SklearnSerializationFormat(str, Enum):
    """
    +label=Serialization format
    +usage=Serialization format used for sklearn models
    """

    cloudpickle = "cloudpickle"
    joblib = "joblib"
    pickle = "pickle"


class SpaCyFramework(BaseModel):
    """
    +docs=spaCy framework for the model version
    +label=SpaCy
    """

    type: Literal["spacy"] = Field(
        ..., description="+label=Type\n+usage=Type of the framework\n+value=spacy"
    )


class StatsModelsFramework(BaseModel):
    """
    +docs=StatsModels framework for the model version
    +label=StatsModels
    """

    type: Literal["statsmodels"] = Field(
        ..., description="+label=Type\n+usage=Type of the framework\n+value=statsmodels"
    )


class SystemMessage(BaseModel):
    """
    +usage=System message
    +docs=Defines the structure of a system message, including role and content
    +label=System Message
    """

    role: Literal["system"]
    content: Union[constr(regex=r"^.[\s\S]*$"), BlobStorageReference] = Field(
        ..., description="+label=Content\n+usage=Text content for the system message"
    )
    name: Optional[str] = Field(
        None, description="+label=Name\n+usage=Name of the system"
    )


class TensorFlowFramework(BaseModel):
    """
    +docs=TensorFlow framework for the model version
    +label=TensorFlow
    """

    type: Literal["tensorflow"] = Field(
        ..., description="+label=Type\n+usage=Type of the framework\n+value=tensorflow"
    )


class TextContentPart(BaseModel):
    """
    +usage=Text content
    +label=Text content
    """

    type: Literal["text"]
    text: Union[constr(regex=r"^.[\s\S]*$"), BlobStorageReference]


class LibraryName(str, Enum):
    """
    +label=Library Name
    +usage=Name of the library for the framework
    """

    transformers = "transformers"
    sentence_transformers = "sentence-transformers"
    diffusers = "diffusers"


class TransformersFramework(BaseModel):
    """
    +docs=Transformers framework for the model version
    +label=Transformers
    +value=transformers
    """

    type: Literal["transformers"] = Field(
        ...,
        description="+label=Type\n+usage=Type of the framework\n+value=transformers",
    )
    library_name: Optional[LibraryName] = Field(
        "transformers",
        description="+label=Library Name\n+usage=Name of the library for the framework",
    )
    pipeline_tag: Optional[str] = Field(
        None,
        description="+label=Pipeline Tag\n+usage=The `pipeline()` task this model can be used with e.g. `text-generation`. See [huggingface docs](https://huggingface.co/docs/transformers/main/en/main_classes/pipelines#transformers.pipeline.task) for all possible values\n+docs=Pipeline tag for the framework",
    )
    base_model: Optional[str] = Field(
        None,
        description="+label=Base Model\n+usage=Base model Id. If this is a finetuned model, this points to the base model used for finetuning\n+docs=Base model Id. If this is a finetuned model, this points to the base model used for finetuning",
    )


class TrueFoundryArtifactSource(BaseModel):
    """
    +usage=Source for the Artifact or Model
    +label=Upload
    """

    type: Literal["truefoundry"] = Field(
        ...,
        description="+usage=Type of the source\n+type=Upload\n+value=truefoundry\n+label=Upload",
    )
    uri: str = Field(
        ...,
        description='+label=URI\n+uiType=ModelUploadInput\n+uiProps={"hideClear":true}',
    )


class UserMessage(BaseModel):
    """
    +usage=User message
    +docs=Defines the structure of a user message, including role and content
    +label=User Message
    """

    role: Literal["user"]
    content: Union[
        constr(regex=r"^.[\s\S]*$"),
        BlobStorageReference,
        List[Union[TextContentPart, ImageContentPart]],
    ] = Field(
        ..., description="+label=Content\n+usage=Text content for the user message"
    )
    name: Optional[str] = Field(
        None, description="+label=Name\n+usage=Name of the user message"
    )


class XGBoostModelSchema(BaseModel):
    """
    +label=XGBoost Model Schema
    """

    infer_method_name: Literal["predict"] = Field(
        ...,
        description="+label=Inference Method Name\n+usage=Name of the method used for inference",
    )
    inputs: List[Dict[str, Any]] = Field(
        ..., description="+label= Input Schema\n+usage=Schema of the input"
    )
    outputs: List[Dict[str, Any]] = Field(
        ..., description="+label= Output Schema\n+usage=Schema of the output"
    )


class XGBoostSerializationFormat(str, Enum):
    """
    +label=Serialization format
    +usage=Serialization format used for XGBoost models
    """

    cloudpickle = "cloudpickle"
    joblib = "joblib"
    pickle = "pickle"
    json = "json"


class AgentOpenAPITool(BaseModel):
    type: Literal["openapi-tool"] = Field(..., description="+value=openapi-tool")
    name: constr(regex=r"^[a-zA-Z][a-zA-Z0-9\-]{1,30}[a-zA-Z0-9]$") = Field(
        ..., description="Names cannot repeat accross tools\n+uiType=Hidden"
    )
    description: constr(min_length=1, max_length=128) = Field(
        ..., description="+sort=10\n+uiType=TextArea"
    )
    openapi_spec: Union[BlobStorageReference, Dict[str, Any]] = Field(
        ..., description="+sort=20\n+uiType=OpenapiSchema"
    )
    base_url: constr(
        regex=r"^(https?://)?(?:www\.)?[-a-zA-Z0-9@:%._\+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}(?:[-a-zA-Z0-9()@:%_\+.~#?&/=]*)$"
    ) = Field(..., description="+sort=30")
    path: constr(min_length=1) = Field(
        ...,
        description="+sort=40\n+uiType=MethodPathSelector\n+label=Method + Path\n+usage=Select one Method and Path for this tool",
    )
    method: Method = Field(..., description="+sort=50\n+uiType=Hidden")
    headers: Optional[Dict[str, str]] = Field(
        None,
        description='+sort=60\n+uiType=KV\n+uiProps={"allowSecrets":true,"secretConfig":{"enableNew":true,"hideOptions":true}}',
    )


class AgentOpenAPIToolWithFQN(AgentOpenAPITool):
    id: str
    fqn: str


class ArtifactVersion(BaseArtifactVersion):
    type: Literal["artifact-version"] = Field(
        ...,
        description='+label=Type\n+usage=Artifact Version\n+value="artifact-version"',
    )
    source: Union[TrueFoundryArtifactSource, ExternalArtifactSource] = Field(
        ...,
        description="+label=Artifact Source\n+usage=Source for the Artifact version\n+uiType=Group",
    )
    step: conint(ge=0) = Field(0, description="+label=Step")


class AssistantMessage(BaseModel):
    """
    +usage=Assistant message
    +docs=Defines the structure of an assistant message, including role and content
    +label=Assistant Message
    """

    role: Literal["assistant"]
    content: Union[constr(regex=r"^.[\s\S]*$"), BlobStorageReference] = Field(
        ..., description="+label=Content\n+usage=Text content for the assistant message"
    )
    name: Optional[str] = Field(
        None, description="+label=Name\n+usage=Name of the assistant message"
    )


class BasePrompt(BaseModel):
    """
    Main prompt structure
    """

    name: constr(regex=r"^[a-zA-Z][a-zA-Z0-9-_]{0,254}[a-zA-Z0-9]$") = Field(
        ...,
        description="+label=Prompt Name\n+icon=fa-desktop:#326ce5\n+message=Alphanumeric word, may contain '-' with a maximum length of 256 characters\n+sort=1000\n+usage=Name of the Prompt. This uniquely identifies it in the workspace.\n+docs=The unique name for the prompt, consisting of alphanumeric characters and dashes, max length 100.",
    )
    description: Optional[constr(min_length=1, max_length=512)] = Field(
        None,
        description="+label=Description\n+icon=fa-desktop:#326ce5\n+message=Description in a maximum of 512 characters\n+sort=2000\n+usage=Description of the prompt\n+docs=An optional description for the prompt, with a maximum length of 512 characters.\n+optional",
    )
    model_configuration: ModelConfiguration
    metadata: Dict[str, str] = Field(
        ...,
        description="+label=Metadata\n+icon=fa-info-circle:#326ce5\n+message=Additional metadata for the prompt\n+sort=7000\n+usage=Define additional metadata for the prompt\n+docs=Key-value pairs to store additional metadata related to the prompt.",
    )


class ChatPrompt(BasePrompt):
    type: Literal["chat_prompt"] = Field(
        ...,
        description='+label=Type\n+icon=fa-desktop:#326ce5\n+message=Type of the prompt\n+sort=3000\n+usage=Type of the prompt\n+docs=Type of the prompt, set to "chat_prompt"\n+default="chat_prompt"',
    )
    messages: List[Union[SystemMessage, AssistantMessage, UserMessage]] = Field(
        ...,
        description="+sort=4000\n+usage=Chat completion messages\n+label=Messages in the chat conversation\n+message=Chat completion messages\n+usage=List of messages in the chat conversation, must be non-empty\n+docs=Messages that define the chat conversation, including system, assistant, and user messages.",
    )
    variables: Optional[Dict[str, Optional[str]]] = Field(
        None,
        description="+label=Variables\n+usage=Variables for the chat completion messages to be used in the prompt messages\n+sort=5000",
    )


class SklearnFramework(BaseModel):
    """
    +docs=Scikit-learn framework for the model version
    +label=Sklearn
    """

    type: Literal["sklearn"] = Field(
        ..., description="+label=Type\n+usage=Type of the framework\n+value=sklearn"
    )
    model_filepath: Optional[str] = Field(
        None,
        description="+label=Model file path\n+usage=Relative path to the model file",
    )
    serialization_format: Optional[SklearnSerializationFormat] = None
    model_schema: Optional[SklearnModelSchema] = None


class XGBoostFramework(BaseModel):
    """
    +docs=XGBoost framework for the model version
    +label=XGBoost
    """

    type: Literal["xgboost"] = Field(
        ..., description="+label=Type\n+usage=Type of the framework\n+value=xgboost"
    )
    serialization_format: Optional[XGBoostSerializationFormat] = None
    model_filepath: Optional[str] = Field(
        None,
        description="+label=Model file path\n+usage=Relative path to the model file",
    )
    model_schema: Optional[XGBoostModelSchema] = None


class AgentApp(BaseModel):
    type: Literal["agent-app"] = Field(..., description="+value=agent-app")
    tools: List[AgentOpenAPIToolWithFQN]
    agents: List[AgentWithFQN]
    root_agent: constr(min_length=1)


class ModelVersion(BaseArtifactVersion):
    type: Literal["model-version"] = Field(
        ..., description='+label=Type\n+usage=Model Version\n+value="model-version"'
    )
    source: Union[TrueFoundryArtifactSource, ExternalArtifactSource] = Field(
        ...,
        description="+label=Model Source\n+usage=Source for the model version\n+uiType=Group",
    )
    framework: Optional[
        Union[
            TransformersFramework,
            TensorFlowFramework,
            SklearnFramework,
            PyTorchFramework,
            KerasFramework,
            XGBoostFramework,
            LightGBMFramework,
            FastAIFramework,
            H2OFramework,
            ONNXFramework,
            SpaCyFramework,
            StatsModelsFramework,
            GluonFramework,
            PaddleFramework,
        ]
    ] = Field(
        None, description="+label=Framework\n+usage=Framework for the model version"
    )
    environment: Optional[ModelVersionEnvironment] = None
    step: conint(ge=0) = Field(0, description="+label=Step")


class VersionedArtifactType(BaseModel):
    __root__: Union[
        ChatPrompt, AgentOpenAPITool, Agent, AgentApp, ModelVersion, ArtifactVersion
    ]
