# generated by datamodel-codegen:
#   filename:  artifacts.json
#   timestamp: 2025-03-03T13:22:20+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 Method(str, Enum):
    """
    HTTP Method for the tool call HTTP request
    """

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


class ArtifactPath(BaseModel):
    src: str = Field(..., description="Local file or folder path")
    dest: Optional[str] = Field(
        None,
        description="Relative path where the file or folder will be uploaded to in the artifact",
    )


class BaseArtifactVersion(BaseModel):
    name: Optional[constr(regex=r"^[A-Za-z0-9_\-]+$", max_length=256)] = Field(
        None, description="Name of the entity"
    )
    description: Optional[constr(max_length=512)] = Field(None, description="")
    metadata: Dict[str, Any] = Field(
        ...,
        description='Key value metadata. Should be valid JSON. For e.g. `{"business-unit": "sales", "quality": "good", "rating": 4.5}`',
    )
    version_alias: Optional[
        constr(regex=r"^v[a-zA-Z0-9.-]*([a-zA-Z0-9]+)$", max_length=128)
    ] = Field(
        None,
        description="Version alias is alternate, ideally human readable, version string to reference an artifact version. It should start with `v` followed by alphanumeric and it can include `.` and `-` in between (e.g. `v1.0.0`, `v1-prod`, `v3-dev`, etc)",
    )
    ml_repo: Optional[constr(regex=r"^[a-zA-Z][a-zA-Z0-9\-]{1,98}[a-zA-Z0-9]$")] = (
        Field(None, description="Name of the ML Repo")
    )
    version: Optional[conint(ge=1)] = Field(None, description="Version of the entity")


class MimeType(str, Enum):
    """
    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):
    """
    A reference to content offloaded to blob storage
    """

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


class Collaborator(BaseModel):
    subject: str = Field(
        ...,
        description="Fully Qualified Name of the subject. eg: user:email or team:teamname",
    )
    role_id: str = Field(..., description="Role ID for the resource")


class ExternalBlobStorageSource(BaseModel):
    type: Literal["external"] = Field(..., description="Type of the source")
    uri: constr(regex=r"^(s3|gs|wasbs)://.*") = Field(
        ...,
        description="URI referencing a path in the blob storage bucket linked to the MLRepo",
    )


class FastAIFramework(BaseModel):
    type: Literal["fastai"] = Field(..., description="Type of the framework")


class GluonFramework(BaseModel):
    type: Literal["gluon"] = Field(..., description="Type of the framework")


class H2OFramework(BaseModel):
    type: Literal["h2o"] = Field(..., description="Type of the framework")


class ImageUrl(BaseModel):
    """
    Image URL linking to the image
    """

    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="Image URL linking to the image")


class ImageContentPart(BaseModel):
    """
    Image content for the message
    """

    type: Literal["image_url"] = Field(..., description="Type of the content part")
    image_url: ImageUrl = Field(..., description="Image URL linking to the image")


class KerasFramework(BaseModel):
    type: Literal["keras"] = Field(..., description="Type of the framework")


class LightGBMFramework(BaseModel):
    type: Literal["lightgbm"] = Field(..., description="Type of the framework")


class LocalArtifactSource(BaseModel):
    type: Literal["local"] = Field(..., description="Type of the source")
    paths: List[ArtifactPath] = Field(
        ...,
        description="Array of ArtifactPath objects representing the source and destination paths",
    )


class LocalModelSource(BaseModel):
    type: Literal["local"] = Field(..., description="Type of the source")
    file_or_folder: str = Field(..., description="Path to the model file or folder")


class MLRepo(BaseModel):
    """
    MLRepo is a repository ML training runs that log params, metrics, plots, images and versioned entities like artifacts, models, prompts, tools, agents
    """

    type: Literal["ml-repo"] = Field(..., description="")
    name: constr(regex=r"^[a-zA-Z][a-zA-Z0-9\-]{1,98}[a-zA-Z0-9]$") = Field(
        ..., description="Name of the ML Repo."
    )
    description: Optional[constr(regex=r"^.{0,500}$")] = Field(
        None, description="Description for the ML Repo."
    )
    storage_integration_fqn: str = Field(
        ...,
        description="Storage Integration to store artifacts and models. A storage integration represents a connected blob storage like AWS S3 / Azure Blob Storage / Google Cloud Storage.",
    )
    collaborators: List[Collaborator] = Field(
        ..., description="Users and Teams that have access to MLRepo"
    )


class Parameters(BaseModel):
    """
    Parameters to pass to the model when generating
    """

    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):
    """
    Add a model from Gateway along with parameters to be used for chat completion
    """

    provider: str = Field(
        ...,
        description="Name of the provider, must be one of the integration providers configured for the Gateway",
    )
    model: constr(regex=r"^[a-zA-Z][a-zA-Z0-9\-]{1,30}[a-zA-Z0-9]$") = Field(
        ...,
        description="Name of the model to be used for generations. This model should be available in the provider",
    )
    parameters: Optional[Parameters] = Field(
        None, description="Parameters to pass to the model when generating"
    )
    extra_parameters: Optional[Dict[str, Any]] = Field(
        None,
        description="Arbitrary extra parameters to pass to the model when generating",
    )


class ModelVersionEnvironment(BaseModel):
    """
    Runtime environment describing python version and dependencies for running the model
    """

    python_version: Optional[constr(regex=r"^\d+(\.\d+){1,2}([\-\.a-z0-9]+)?$")] = (
        Field(None, description="Python version to use when running the model")
    )
    pip_packages: Optional[List[str]] = Field(
        None, description="pip dependencies needed for running the model"
    )


class ONNXFramework(BaseModel):
    type: Literal["onnx"] = Field(..., description="Type of the framework")


class PaddleFramework(BaseModel):
    type: Literal["paddle"] = Field(..., description="Type of the framework")


class PyTorchFramework(BaseModel):
    type: Literal["pytorch"] = Field(..., description="Type of the framework")


class InferMethodName(str, Enum):
    """
    Name of the method used for inference
    """

    predict = "predict"
    predict_proba = "predict_proba"


class SklearnModelSchema(BaseModel):
    infer_method_name: InferMethodName = Field(
        ..., description="Name of the method used for inference"
    )
    inputs: List[Dict[str, Any]] = Field(..., description="Schema of the input")
    outputs: List[Dict[str, Any]] = Field(..., description="Schema of the output")


class SklearnSerializationFormat(str, Enum):
    """
    Serialization format used to save the sklearn model
    """

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


class SpaCyFramework(BaseModel):
    type: Literal["spacy"] = Field(..., description="Type of the framework")


class StatsModelsFramework(BaseModel):
    type: Literal["statsmodels"] = Field(..., description="Type of the framework")


class SystemMessage(BaseModel):
    """
    System message for the chat
    """

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


class TensorFlowFramework(BaseModel):
    type: Literal["tensorflow"] = Field(..., description="Type of the framework")


class TextContentPart(BaseModel):
    """
    Text content for the message
    """

    type: Literal["text"] = Field(..., description="Type of the content part")
    text: Union[constr(regex=r"^.[\s\S]*$"), BlobStorageReference] = Field(
        ..., description="Text content for the message"
    )


class LibraryName(str, Enum):
    """
    Name of the library for the framework
    """

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


class TransformersFramework(BaseModel):
    type: Literal["transformers"] = Field(..., description="Type of the framework")
    library_name: Optional[LibraryName] = Field(
        "transformers", description="Name of the library for the framework"
    )
    pipeline_tag: Optional[str] = Field(
        None,
        description="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",
    )
    base_model: Optional[str] = Field(
        None,
        description="Base model Id from HuggingFace Hub. If this is a finetuned model, this points to the base model id used for finetuning.",
    )


class TrueFoundryManagedSource(BaseModel):
    type: Literal["truefoundry"] = Field(..., description="Type of the source")
    uri: Optional[constr(regex=r"^(s3|gs|wasbs)://.*")] = Field(None, description="")


class UserMessage(BaseModel):
    """
    User message turn
    """

    role: Literal["user"] = Field(..., description="Role of the message")
    content: Union[
        constr(regex=r"^.[\s\S]*$"),
        BlobStorageReference,
        List[Union[TextContentPart, ImageContentPart]],
    ] = Field(
        ..., description="Content of the user message. can be a mix of text and images"
    )
    name: Optional[str] = Field(
        None, description="Name of the user this message is from"
    )


class XGBoostModelSchema(BaseModel):
    """
    Schema of the XGBoost model
    """

    infer_method_name: Literal["predict"] = Field(
        ..., description="Name of the method used for inference"
    )
    inputs: List[Dict[str, Any]] = Field(..., description="Schema of the input")
    outputs: List[Dict[str, Any]] = Field(..., description="Schema of the output")


class XGBoostSerializationFormat(str, Enum):
    """
    Serialization format used to save the xgboost model
    """

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


class Agent(BaseArtifactVersion):
    type: Literal["agent"] = Field(..., description="Type of the entity")
    available_tools: List[constr(min_length=1)] = Field(
        ..., description="Tools available to the agent"
    )
    goal: constr(min_length=1, max_length=128) = Field(
        ...,
        description="Short form description. Will be used as `description` when this agent is used as a tool.",
    )
    instruction: constr(min_length=1, max_length=2620) = Field(
        ..., description="Instructions for the agent to follow to achieve the goal"
    )
    model_id: constr(min_length=1) = Field(
        ..., description="Model to use when running the agent"
    )


class AgentOpenAPITool(BaseArtifactVersion):
    type: Literal["openapi-tool"] = Field(..., description="Type of the agent")
    openapi_spec: Union[BlobStorageReference, Dict[str, Any]] = Field(
        ...,
        description="OpenAPI Spec for the tool describing the API, endpoints and parameters. [Sample OpenAPI Spec Link](https://assets.production.truefoundry.com/sample-openapi.json)",
    )
    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="HTTP endpoint where the API is hosted for the tools. E.g. `https://api.example.com/v1`",
    )
    path: constr(min_length=1) = Field(
        ...,
        description="API Route Path for the tool call HTTP request. E.g. `GET /weather`",
    )
    method: Method = Field(
        ..., description="HTTP Method for the tool call HTTP request"
    )
    headers: Optional[Dict[str, str]] = Field(
        None,
        description="HTTP Headers for the tool call HTTP request. E.g. `Authorization: Bearer <token>`",
    )


class AgentOpenAPIToolWithFQN(AgentOpenAPITool):
    id: str = Field(..., description="ID of the agent tool")
    fqn: str = Field(..., description="FQN of the agent tool")


class AgentWithFQN(Agent):
    id: str = Field(..., description="ID of the agent")
    fqn: str = Field(..., description="FQN of the agent")


class ArtifactVersion(BaseArtifactVersion):
    """
    Log a new Artifact Version containing files and folders with metadata
    """

    type: Literal["artifact-version"] = Field(..., description="Artifact Version")
    source: Union[
        TrueFoundryManagedSource, ExternalBlobStorageSource, LocalArtifactSource
    ] = Field(..., description="")
    step: conint(ge=0) = Field(
        0,
        description="Step/Epoch number in an iterative training loop the artifact version was created. Generally useful when logging a model version from a MLRepo Run",
    )
    run_id: Optional[str] = Field(
        None, description="ID of the MLRepo Run that generated the artifact version"
    )


class AssistantMessage(BaseModel):
    """
    Assistant message turn
    """

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


class ChatPrompt(BaseArtifactVersion):
    """
    Fill in the information for opening a new chat prompt
    """

    type: Literal["chat_prompt"] = Field(..., description="Type of the prompt")
    messages: List[Union[SystemMessage, AssistantMessage, UserMessage]] = Field(
        ..., description="List of messages in the chat conversation, must be non-empty"
    )
    variables: Optional[Dict[str, Optional[str]]] = Field(
        None,
        description="Variables referenced in messages and that can be replaced when running generation",
    )
    model_configuration: ModelConfiguration


class DataDirectory(BaseModel):
    type: Literal["data-dir"] = Field(..., description="")
    name: constr(regex=r"^[A-Za-z0-9_\-]+$", max_length=256) = Field(
        ..., description="Name of the data directory"
    )
    ml_repo: constr(regex=r"^[a-zA-Z][a-zA-Z0-9\-]{1,98}[a-zA-Z0-9]$") = Field(
        ..., description="Name of the ML Repo to create the data directory under"
    )
    description: Optional[constr(max_length=512)] = Field(
        None, description="Description for the data directory"
    )
    metadata: Dict[str, Any] = Field(
        ...,
        description='Key Value metadata. Should be valid JSON. For e.g. `{"business-unit": "sales", "quality": "good", "rating": 4.5}`',
    )
    source: Union[TrueFoundryManagedSource, ExternalBlobStorageSource] = Field(
        ..., description=""
    )


class SklearnFramework(BaseModel):
    type: Literal["sklearn"] = Field(..., description="Type of the framework")
    model_filepath: Optional[str] = Field(
        None,
        description="Relative path to the model file in the model version contents",
    )
    serialization_format: Optional[SklearnSerializationFormat] = None
    model_schema: Optional[SklearnModelSchema] = None


class XGBoostFramework(BaseModel):
    type: Literal["xgboost"] = Field(..., description="Type of the framework")
    serialization_format: Optional[XGBoostSerializationFormat] = None
    model_filepath: Optional[str] = Field(
        None,
        description="Relative path to the model file in the model version contents",
    )
    model_schema: Optional[XGBoostModelSchema] = None


class AgentApp(BaseModel):
    type: Literal["agent-app"] = Field(..., description="Type of the app")
    tools: List[AgentOpenAPIToolWithFQN] = Field(
        ..., description="Tools available to the Agent app"
    )
    agents: List[AgentWithFQN] = Field(
        ..., description="Agents available to the Agent app"
    )
    root_agent: constr(min_length=1) = Field(
        ..., description="Root Agent for the app. This will be the first agent invoked"
    )


class ModelVersion(BaseArtifactVersion):
    """
    Log a new Model Version containing model files and folders with metadata
    """

    type: Literal["model-version"] = Field(..., description="Model Version")
    source: Union[
        TrueFoundryManagedSource, ExternalBlobStorageSource, LocalModelSource
    ] = Field(..., description="")
    framework: Optional[
        Union[
            TransformersFramework,
            TensorFlowFramework,
            SklearnFramework,
            PyTorchFramework,
            KerasFramework,
            XGBoostFramework,
            LightGBMFramework,
            FastAIFramework,
            H2OFramework,
            ONNXFramework,
            SpaCyFramework,
            StatsModelsFramework,
            GluonFramework,
            PaddleFramework,
        ]
    ] = Field(
        None,
        description="Framework for the model version like Transformers, PyTorch, Sklearn, Xgboost etc with framework specific metadata. This will be used to infer model deployment configuration",
    )
    environment: Optional[ModelVersionEnvironment] = None
    step: conint(ge=0) = Field(
        0,
        description="Step/Epoch number in an iterative training loop the model version was created. Generally useful when logging a model version from a MLRepo Run",
    )
    run_id: Optional[str] = Field(
        None, description="ID of the MLRepo Run that generated the model version"
    )


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