truefoundry 0.5.10rc2__py3-none-any.whl → 0.5.11rc1__py3-none-any.whl

truefoundry/ml/autogen/entities/artifacts.py

@@ -1,795 +0,0 @@
-# generated by datamodel-codegen:
-#   filename:  artifacts.json
-#   timestamp: 2025-02-26T17:30:50+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):
-    """
-    +sort=50
-    +uiType=Hidden
-    +label=API HTTP Method
-    +usage=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="+label=Source path\n+usage=Local file or folder path"
-    )
-    dest: Optional[str] = Field(
-        None,
-        description="+label=Destination path\n+usage=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="+sort=1\n+label=Name\n+message=The name should start with lowercase alphabets  and can contain alphanumeric and can include '-' in between\n+usage=Name of the entity",
-    )
-    description: Optional[constr(max_length=512)] = Field(
-        None, description="+sort=2\n+label=Description"
-    )
-    metadata: Dict[str, Any] = Field(
-        ...,
-        description='+label=Metadata\n+usage=Key value metadata. Should be valid JSON. For e.g. `{"business-unit": "sales", "quality": "good", "rating": 4.5}`\n+uiType=JsonInput\n+uiProps={"descriptionInline":true}',
-    )
-    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=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)\n+message=The version alias should start with `v` followed by alphanumeric and can include `.` and `-` in between (e.g. `v1.0.0`, `v1-prod`, `v3-dev`, etc)\n+uiProps={"descriptionInline":true}\n+placeholder=v1.0.0',
-    )
-    ml_repo: Optional[constr(regex=r"^[a-zA-Z][a-zA-Z0-9\-]{1,98}[a-zA-Z0-9]$")] = (
-        Field(
-            None,
-            description="+label=ML Repo\n+usage=Name of the ML Repo\n+uiType=Hidden",
-        )
-    )
-    version: Optional[conint(ge=1)] = Field(
-        None,
-        description="+label=Version of the entity\n+usage=Version of the entity\n+uiType=Hidden",
-    )
-
-
-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):
-    """
-    +label=Blob Storage Reference
-    +usage=A reference to content offloaded to blob storage
-    +docs=Defines the structure for blob storage content references, including type, path, and mime_type
-    """
-
-    type: Literal["blob-storage"] = Field(..., description="+value=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 Collaborator(BaseModel):
-    subject: str = Field(
-        ...,
-        description="+label=Subject FQN\n+usage=Fully Qualified Name of the subject. eg: user:email or team:teamname",
-    )
-    role_id: str = Field(
-        ..., description="+label=Role ID\n+usage=Role ID for the resource"
-    )
-
-
-class ExternalBlobStorageSource(BaseModel):
-    """
-    +label=External Blob Storage Source
-    """
-
-    type: Literal["external"] = Field(
-        ..., description="+label=Type\n+usage=Type of the source\n+value=external"
-    )
-    uri: str = Field(
-        ...,
-        description="+label=URI\n+usage=URI referencing a path in the blob storage bucket linked to the MLRepo",
-    )
-
-
-class FastAIFramework(BaseModel):
-    """
-    +label=FastAI
-    +icon=fastai
-    """
-
-    type: Literal["fastai"] = Field(
-        ..., description="+label=Type\n+usage=Type of the framework\n+value=fastai"
-    )
-
-
-class GluonFramework(BaseModel):
-    """
-    +label=Gluon
-    +icon=gluon
-    """
-
-    type: Literal["gluon"] = Field(
-        ..., description="+label=Type\n+usage=Type of the framework\n+value=gluon"
-    )
-
-
-class H2OFramework(BaseModel):
-    """
-    +label=H2O
-    +icon=h2o
-    """
-
-    type: Literal["h2o"] = Field(
-        ..., description="+label=Type\n+usage=Type of the framework\n+value=h2o"
-    )
-
-
-class ImageUrl(BaseModel):
-    """
-    +label=Image URL
-    +usage=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="+label=URL\n+usage=Image URL linking to the image")
-
-
-class ImageContentPart(BaseModel):
-    """
-    +label=Image Content
-    +usage=Image content for the message
-    """
-
-    type: Literal["image_url"] = Field(
-        ..., description="+value=image_url\n+usage=Type of the content part"
-    )
-    image_url: ImageUrl = Field(
-        ..., description="+label=Image URL\n+usage=Image URL linking to the image"
-    )
-
-
-class KerasFramework(BaseModel):
-    """
-    +label=Keras
-    +icon=keras
-    """
-
-    type: Literal["keras"] = Field(
-        ..., description="+label=Type\n+usage=Type of the framework\n+value=keras"
-    )
-
-
-class LightGBMFramework(BaseModel):
-    """
-    +label=LightGBM
-    +icon=lightbgm
-    """
-
-    type: Literal["lightgbm"] = Field(
-        ..., description="+label=Type\n+usage=Type of the framework\n+value=lightgbm"
-    )
-
-
-class LocalArtifactSource(BaseModel):
-    type: Literal["local"] = Field(
-        ..., description="+label=Type\n+usage=Type of the source\n+value=local"
-    )
-    paths: List[ArtifactPath] = Field(
-        ...,
-        description="+label=Paths\n+usage=Array of ArtifactPath objects representing the source and destination paths",
-    )
-
-
-class LocalModelSource(BaseModel):
-    type: Literal["local"] = Field(
-        ..., description="+label=Type\n+usage=Type of the source\n+value=local"
-    )
-    file_or_folder: str = Field(
-        ...,
-        description="+label=Path to the model file or folder\n+usage=Path to the model file or folder",
-    )
-
-
-class MLRepo(BaseModel):
-    """
-    +label=MLRepo
-    +usage=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="+value=ml-repo")
-    name: constr(regex=r"^[a-zA-Z][a-zA-Z0-9\-]{1,98}[a-zA-Z0-9]$") = Field(
-        ...,
-        description="+label=Name\n+icon=fa-desktop:#black\n+usage=Name of the ML Repo.\n+message=Alphanumeric word, may contain '-' with a maximum length of 100 characters\n+sort=1",
-    )
-    description: Optional[constr(regex=r"^.{0,500}$")] = Field(
-        None,
-        description='+label=Description\n+icon=fa-desktop:#black\n+usage=Description for the ML Repo.\n+message=Description can be maximum 500 characters\n+sort=2\n+placeholder=MLRepo to track text-classification experiments and models\n+uiProps={"descriptionInline":true}',
-    )
-    storage_integration_fqn: str = Field(
-        ...,
-        description="+label=Storage Integration\n+icon=hard-drive\n+usage=Storage Integration to store artifacts and models. A storage integration represents a connected blob storage like AWS S3 / Azure Blob Storage / Google Cloud Storage.\n+message=Select the storage-integration where you want to save your artifacts and models\n[View Docs](https://docs.truefoundry.com/docs/integrations)\n+sort=4",
-    )
-    collaborators: List[Collaborator] = Field(
-        ...,
-        description="+label=Collaborators\n+sort=5\n+usage=Users and Teams that have access to MLRepo\n+uiType=Collaborators",
-    )
-
-
-class Parameters(BaseModel):
-    """
-    +label=Parameters
-    +usage=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):
-    """
-    +label=Model Configuration
-    +icon=fa-cogs:#326ce5
-    +usage=Add a model from Gateway along with parameters to be used for chat completion
-    """
-
-    provider: str = Field(
-        ...,
-        description='+sort=1000\n+uiType=SelectModelAndProvider\n+uiProps={"isModel": false,"providerJsonKey": "model_configuration.provider","searchable": true}\n+label=Provider Name\n+usage=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='+sort=1001\n+uiType=SelectModelAndProvider\n+uiProps={"isModel": true,"providerJsonKey": "model_configuration.provider","searchable": true}\n+label=Model Name\n+usage=Name of the model to be used for generations. This model should be available in the provider',
-    )
-    parameters: Optional[Parameters] = Field(
-        None,
-        description="+label=Parameters\n+usage=Parameters to pass to the model when generating",
-    )
-    extra_parameters: Optional[Dict[str, Any]] = Field(
-        None,
-        description="+label=Extra Parameters\n+usage=Arbitrary extra parameters to pass to the model when generating",
-    )
-
-
-class ModelVersionEnvironment(BaseModel):
-    """
-    +label=Environment
-    +usage=Runtime environment describing python version and dependencies for running the model
-    +uiProps={"descriptionInline":true}
-    """
-
-    python_version: Optional[constr(regex=r"^\d+(\.\d+){1,2}([\-\.a-z0-9]+)?$")] = (
-        Field(
-            None,
-            description="+label=Python Version\n+usage=Python version to use when running the model",
-        )
-    )
-    pip_packages: Optional[List[str]] = Field(
-        None,
-        description="+label=PIP Packages\n+usage=pip dependencies needed for running the model",
-    )
-
-
-class ONNXFramework(BaseModel):
-    """
-    +label=ONNX
-    +icon=onnx
-    """
-
-    type: Literal["onnx"] = Field(
-        ..., description="+label=Type\n+usage=Type of the framework\n+value=onnx"
-    )
-
-
-class PaddleFramework(BaseModel):
-    """
-    +label=Paddle
-    +icon=paddle
-    """
-
-    type: Literal["paddle"] = Field(
-        ..., description="+label=Type\n+usage=Type of the framework\n+value=paddle"
-    )
-
-
-class PyTorchFramework(BaseModel):
-    """
-    +label=PyTorch
-    +icon=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=Scikit Learn 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\n+uiProps={"descriptionInline":true}',
-    )
-    outputs: List[Dict[str, Any]] = Field(
-        ...,
-        description='+label= Output Schema\n+usage=Schema of the output\n+uiProps={"descriptionInline":true}',
-    )
-
-
-class SklearnSerializationFormat(str, Enum):
-    """
-    +label=Scikit Learn Serialization Format
-    +usage=Serialization format used to save the sklearn model
-    """
-
-    cloudpickle = "cloudpickle"
-    joblib = "joblib"
-    pickle = "pickle"
-
-
-class SpaCyFramework(BaseModel):
-    """
-    +label=SpaCy
-    +icon=spacy
-    """
-
-    type: Literal["spacy"] = Field(
-        ..., description="+label=Type\n+usage=Type of the framework\n+value=spacy"
-    )
-
-
-class StatsModelsFramework(BaseModel):
-    """
-    +label=StatsModels
-    +icon=statsmodel
-    """
-
-    type: Literal["statsmodels"] = Field(
-        ..., description="+label=Type\n+usage=Type of the framework\n+value=statsmodels"
-    )
-
-
-class SystemMessage(BaseModel):
-    """
-    +label=System Message
-    +usage=System message for the chat
-    """
-
-    role: Literal["system"] = Field(
-        ...,
-        description="+value=system\n+label=Role\n+usage=Role of the message\n+uiType=Ignore",
-    )
-    content: Union[constr(regex=r"^.[\s\S]*$"), BlobStorageReference] = Field(
-        ...,
-        description="+label=Content\n+usage=Text content for the system message\n+uiType=TextArea",
-    )
-    name: Optional[str] = Field(
-        None, description="+label=Name\n+usage=Name of the system"
-    )
-
-
-class TensorFlowFramework(BaseModel):
-    """
-    +label=Tensorflow
-    +icon=tensorflow
-    """
-
-    type: Literal["tensorflow"] = Field(
-        ..., description="+label=Type\n+usage=Type of the framework\n+value=tensorflow"
-    )
-
-
-class TextContentPart(BaseModel):
-    """
-    +label=Text Content
-    +usage=Text content for the message
-    """
-
-    type: Literal["text"] = Field(
-        ..., description="+value=text\n+usage=Type of the content part"
-    )
-    text: Union[constr(regex=r"^.[\s\S]*$"), BlobStorageReference] = Field(
-        ..., description="+label=Text\n+usage=Text content for the message"
-    )
-
-
-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):
-    """
-    +label=Transformers
-    +icon=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+uiProps={"descriptionInline":true}',
-    )
-    base_model: Optional[str] = Field(
-        None,
-        description='+label=Base Model\n+usage=Base model Id from HuggingFace Hub. If this is a finetuned model, this points to the base model id used for finetuning.\n+uiProps={"descriptionInline":true}',
-    )
-
-
-class TrueFoundryManagedSource(BaseModel):
-    """
-    +label=TrueFoundry Managed Source
-    """
-
-    type: Literal["truefoundry"] = Field(
-        ..., description="+label=Type\n+usage=Type of the source\n+value=truefoundry"
-    )
-    uri: Optional[str] = Field(
-        None,
-        description='+label=URI\n+uiType=ModelUploadInput\n+uiProps={"hideClear":true}',
-    )
-
-
-class UserMessage(BaseModel):
-    """
-    +label=User Message
-    +usage=User message turn
-    """
-
-    role: Literal["user"] = Field(
-        ..., description="+value=user\n+usage=Role of the message\n+uiType=Ignore"
-    )
-    content: Union[
-        constr(regex=r"^.[\s\S]*$"),
-        BlobStorageReference,
-        List[Union[TextContentPart, ImageContentPart]],
-    ] = Field(
-        ...,
-        description="+label=Content\n+usage=Content of the user message. can be a mix of text and images\n+uiType=TextArea",
-    )
-    name: Optional[str] = Field(
-        None, description="+label=Name\n+usage=Name of the user this message is from"
-    )
-
-
-class XGBoostModelSchema(BaseModel):
-    """
-    +label=XGBoost Model Schema
-    +usage=Schema of the XGBoost model
-    """
-
-    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\n+uiProps={"descriptionInline":true}',
-    )
-    outputs: List[Dict[str, Any]] = Field(
-        ...,
-        description='+label= Output Schema\n+usage=Schema of the output\n+uiProps={"descriptionInline":true}',
-    )
-
-
-class XGBoostSerializationFormat(str, Enum):
-    """
-    +label=Serialization format
-    +usage=Serialization format used to save the xgboost model
-    """
-
-    cloudpickle = "cloudpickle"
-    joblib = "joblib"
-    pickle = "pickle"
-    json = "json"
-
-
-class Agent(BaseArtifactVersion):
-    type: Literal["agent"] = Field(
-        ..., description="+value=agent\n+usage=Type of the entity"
-    )
-    available_tools: List[constr(min_length=1)] = Field(
-        ...,
-        description='+sort=30\n+uiType=Hidden\n+label=Tools\n+usage=Tools available to the agent\n+uiProps={"descriptionInline":true}',
-    )
-    goal: constr(min_length=1, max_length=128) = Field(
-        ...,
-        description='+sort=10\n+uiType=TextArea\n+label=Goal\n+usage=Short form description. Will be used as `description` when this agent is used as a tool.\n+uiProps={"descriptionInline":true}',
-    )
-    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"}\n+label=Instructions\n+usage=Instructions for the agent to follow to achieve the goal\n+uiProps={"descriptionInline":true}',
-    )
-    model_id: constr(min_length=1) = Field(
-        ...,
-        description='+sort=40\n+uiType=EnabledModelSelector\n+uiProps={"searchable":true,"modelType":"chat","providerType":"openai"}\n+label=Model\n+usage=Model to use when running the agent',
-    )
-
-
-class AgentOpenAPITool(BaseArtifactVersion):
-    type: Literal["openapi-tool"] = Field(
-        ..., description="+value=openapi-tool\n+usage=Type of the agent"
-    )
-    openapi_spec: Union[BlobStorageReference, Dict[str, Any]] = Field(
-        ...,
-        description='+sort=20\n+uiType=OpenapiSchema\n+label=OpenAPI Spec\n+usage=OpenAPI Spec for the tool describing the API, endpoints and parameters. [Sample OpenAPI Spec Link](https://assets.production.truefoundry.com/sample-openapi.json)\n+uiProps={"descriptionInline":true}\n+placeholder={"openapi":"3.0.0","info":{"title":"LLM Agent Tools API","version":"1.0.0"},"servers":[{"url":"https://api.example.com/v1"}],"paths":{"/weather":{"get":{"summary":"Get current weather","description":"Fetches the current weather for a given location.","parameters":[{"name":"location","in":"query","required":true,"schema":{"type":"string"}}],"responses":{"200":{"description":"Successful response","content":{"application/json":{"schema":{"type":"object","properties":{"location":{"type":"string"},"temperature":{"type":"number"},"condition":{"type":"string"}}}}}}}}}}}',
-    )
-    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\n+label=Base URL\n+usage=HTTP endpoint where the API is hosted for the tools. E.g. `https://api.example.com/v1`\n+uiProps={"descriptionInline":true}\n+placeholder=https://api.example.com/v1',
-    )
-    path: constr(min_length=1) = Field(
-        ...,
-        description='+sort=40\n+uiType=MethodPathSelector\n+label=API Route Path\n+usage=API Route Path for the tool call HTTP request. E.g. `GET /weather`\n+uiProps={"descriptionInline":true}',
-    )
-    method: Method = Field(
-        ...,
-        description="+sort=50\n+uiType=Hidden\n+label=API HTTP Method\n+usage=HTTP Method for the tool call HTTP request",
-    )
-    headers: Optional[Dict[str, str]] = Field(
-        None,
-        description='+sort=60\n+uiType=KV\n+uiProps={"allowSecrets":true,"secretConfig":{"enableNew":true,"hideOptions":true}}\n+label=Headers\n+usage=HTTP Headers for the tool call HTTP request. E.g. `Authorization: Bearer <token>`\n+uiProps={"descriptionInline":true}',
-    )
-
-
-class AgentOpenAPIToolWithFQN(AgentOpenAPITool):
-    id: str = Field(..., description="+usage=ID of the agent tool")
-    fqn: str = Field(..., description="+usage=FQN of the agent tool")
-
-
-class AgentWithFQN(Agent):
-    id: str = Field(..., description="+usage=ID of the agent")
-    fqn: str = Field(..., description="+usage=FQN of the agent")
-
-
-class ArtifactVersion(BaseArtifactVersion):
-    """
-    +label=Artifact Version
-    +usage=Log a new Artifact Version containing files and folders with metadata
-    """
-
-    type: Literal["artifact-version"] = Field(
-        ..., description="+label=Type\n+usage=Artifact Version\n+value=artifact-version"
-    )
-    source: Union[
-        TrueFoundryManagedSource, ExternalBlobStorageSource, LocalArtifactSource
-    ] = Field(..., description="+label=Artifact Source\n+uiType=Group")
-    step: conint(ge=0) = Field(
-        0,
-        description='+label=Step\n+usage=Step/Epoch number in an iterative training loop the artifact version was created. Generally useful when logging a model version from a MLRepo Run\n+uiProps={"descriptionInline":true}',
-    )
-    run_id: Optional[str] = Field(
-        None,
-        description="+label=Run ID\n+usage=ID of the MLRepo Run that generated the artifact version\n+uiType=Hidden",
-    )
-
-
-class AssistantMessage(BaseModel):
-    """
-    +label=Assistant Message
-    +usage=Assistant message turn
-    """
-
-    role: Literal["assistant"] = Field(
-        ...,
-        description="+value=assistant\n+label=Role\n+usage=Role of the message\n+uiType=Ignore",
-    )
-    content: Union[constr(regex=r"^.[\s\S]*$"), BlobStorageReference] = Field(
-        ...,
-        description="+label=Content\n+usage=Text content for the assistant message\n+uiType=TextArea",
-    )
-    name: Optional[str] = Field(
-        None,
-        description="+label=Name\n+usage=Name of the assistant this message is from",
-    )
-
-
-class ChatPrompt(BaseArtifactVersion):
-    """
-    +label=Chat Prompt
-    +usage=Fill in the information for opening a new chat prompt
-    """
-
-    type: Literal["chat_prompt"] = Field(
-        ...,
-        description="+label=Type\n+icon=fa-desktop:#326ce5\n+sort=3000\n+usage=Type of the prompt\n+value=chat_prompt",
-    )
-    messages: List[Union[SystemMessage, AssistantMessage, UserMessage]] = Field(
-        ...,
-        description="+sort=4000\n+label=Messages in the chat conversation\n+usage=List of messages in the chat conversation, must be non-empty",
-    )
-    variables: Optional[Dict[str, Optional[str]]] = Field(
-        None,
-        description="+label=Variables\n+usage=Variables referenced in messages and that can be replaced when running generation\n+sort=5000",
-    )
-    model_configuration: ModelConfiguration
-
-
-class DataDirectory(BaseModel):
-    """
-    +label=Data Directory
-    """
-
-    type: Literal["data-dir"] = Field(..., description="+label=Type\n+value=data-dir")
-    name: constr(regex=r"^[A-Za-z0-9_\-]+$", max_length=256) = Field(
-        ...,
-        description="+sort=1\n+label=Name\n+usage=Name of the data directory\n+message=The data directory name should start with lowercase alphabets and can contain alphanumeric and can include '-' in between",
-    )
-    ml_repo: constr(regex=r"^[a-zA-Z][a-zA-Z0-9\-]{1,98}[a-zA-Z0-9]$") = Field(
-        ...,
-        description="+label=MLRepo\n+usage=Name of the ML Repo to create the data directory under\n+uiType=Hidden",
-    )
-    description: Optional[constr(max_length=512)] = Field(
-        None,
-        description="+sort=2\n+label=Description\n+usage=Description for the data directory",
-    )
-    metadata: Dict[str, Any] = Field(
-        ...,
-        description='+label=Metadata\n+usage=Key Value metadata. Should be valid JSON. For e.g. `{"business-unit": "sales", "quality": "good", "rating": 4.5}`\n+uiType=JsonInput',
-    )
-    source: Union[TrueFoundryManagedSource, ExternalBlobStorageSource] = Field(
-        ..., description="+label=Data Directory Source\n+uiType=Group"
-    )
-
-
-class SklearnFramework(BaseModel):
-    """
-    +label=Sklearn
-    +icon=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 in the model version contents\n+uiProps={"descriptionInline":true}',
-    )
-    serialization_format: Optional[SklearnSerializationFormat] = None
-    model_schema: Optional[SklearnModelSchema] = None
-
-
-class XGBoostFramework(BaseModel):
-    """
-    +label=XGBoost
-    +icon=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 in the model version contents",
-    )
-    model_schema: Optional[XGBoostModelSchema] = None
-
-
-class AgentApp(BaseModel):
-    type: Literal["agent-app"] = Field(
-        ..., description="+value=agent-app\n+usage=Type of the app"
-    )
-    tools: List[AgentOpenAPIToolWithFQN] = Field(
-        ..., description="+label=Tools\n+usage=Tools available to the Agent app"
-    )
-    agents: List[AgentWithFQN] = Field(
-        ..., description="+label=Agents\n+usage=Agents available to the Agent app"
-    )
-    root_agent: constr(min_length=1) = Field(
-        ...,
-        description="+label=Root Agent\n+usage=Root Agent for the app. This will be the first agent invoked",
-    )
-
-
-class ModelVersion(BaseArtifactVersion):
-    """
-    +label=Artifact Version
-    +usage=Log a new Model Version containing model files and folders with metadata
-    """
-
-    type: Literal["model-version"] = Field(
-        ..., description="+label=Type\n+usage=Model Version\n+value=model-version"
-    )
-    source: Union[
-        TrueFoundryManagedSource, ExternalBlobStorageSource, LocalModelSource
-    ] = Field(..., description="+label=Model Source\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 like Transformers, PyTorch, Sklearn, Xgboost etc with framework specific metadata. This will be used to infer model deployment configuration\n+uiProps={"descriptionInline":true}',
-    )
-    environment: Optional[ModelVersionEnvironment] = None
-    step: conint(ge=0) = Field(
-        0,
-        description='+label=Step\n+usage=Step/Epoch number in an iterative training loop the model version was created. Generally useful when logging a model version from a MLRepo Run\n+uiProps={"descriptionInline":true}',
-    )
-    run_id: Optional[str] = Field(
-        None,
-        description="+label=Run ID\n+usage=ID of the MLRepo Run that generated the model version\n+uiType=Hidden",
-    )
-
-
-class VersionedArtifactType(BaseModel):
-    __root__: Union[
-        ArtifactVersion, ModelVersion, ChatPrompt, AgentOpenAPITool, Agent, AgentApp
-    ]