truefoundry 0.5.10rc2__py3-none-any.whl → 0.6.1__py3-none-any.whl

truefoundry/ml/_autogen/entities/artifacts.py

@@ -0,0 +1,558 @@
+# 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
+    ]

truefoundry/ml/artifact/truefoundry_artifact_repo.py

@@ -43,7 +43,7 @@ from truefoundry.common.storage_provider_utils import (
     decide_file_parts,
     s3_compatible_multipart_upload,
 )
-from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
+from truefoundry.ml._autogen.client import (  # type: ignore[attr-defined]
     ApiClient,
     CreateMultiPartUploadForDatasetRequestDto,
     CreateMultiPartUploadRequestDto,

truefoundry/ml/cli/commands/model_init.py

@@ -6,7 +6,7 @@ import rich_click as click
 from truefoundry.cli.console import console
 from truefoundry.cli.const import COMMAND_CLS
 from truefoundry.cli.util import handle_exception_wrapper
-from truefoundry.ml.autogen.client import ModelServer  # type: ignore[attr-defined]
+from truefoundry.ml._autogen.client import ModelServer  # type: ignore[attr-defined]
 from truefoundry.ml.cli.utils import (
     AppName,
     NonEmptyString,

truefoundry/ml/clients/servicefoundry_client.py

@@ -13,6 +13,7 @@ from truefoundry.common.request_utils import (
 from truefoundry.common.servicefoundry_client import (
     ServiceFoundryServiceClient as BaseServiceFoundryServiceClient,
 )
+from truefoundry.common.utils import get_user_agent
 from truefoundry.ml.exceptions import MlFoundryException
 
 
@@ -36,6 +37,9 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
             timeout=3,
             params={"id": integration_id, "type": "blob-storage"},
             session=session,
+            headers={
+                "User-Agent": get_user_agent(),
+            },
         )
 
         try:

truefoundry/ml/entities.py

@@ -1,6 +1,6 @@
 from typing import Optional, Union
 
-from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
+from truefoundry.ml._autogen.client import (  # type: ignore[attr-defined]
     FileInfoDto,
     MetricDto,
 )

truefoundry/ml/log_types/artifacts/artifact.py

@@ -9,15 +9,11 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, NamedTuple, Optional, Union
 
 from truefoundry.common.warnings import TrueFoundryDeprecationWarning
-from truefoundry.ml.artifact.truefoundry_artifact_repo import (
-    ArtifactIdentifier,
-    MlFoundryArtifactsRepository,
-)
-from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
+from truefoundry.ml._autogen.client import (  # type: ignore[attr-defined]
     ArtifactDto,
+    ArtifactManifest,
     ArtifactType,
     ArtifactVersionDto,
-    ArtifactVersionManifest,
     CreateArtifactVersionRequestDto,
     DeleteArtifactVersionsRequestDto,
     ExperimentsApi,
@@ -28,6 +24,10 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     TrueFoundryManagedSource,
     UpdateArtifactVersionRequestDto,
 )
+from truefoundry.ml.artifact.truefoundry_artifact_repo import (
+    ArtifactIdentifier,
+    MlFoundryArtifactsRepository,
+)
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.log_types.artifacts.constants import INTERNAL_METADATA_PATH
 from truefoundry.ml.log_types.artifacts.utils import (
@@ -392,7 +392,7 @@ class ArtifactVersion:
         self._ensure_not_deleted()
         if self._artifact_version.manifest:
             manifest = self._artifact_version.manifest.actual_instance
-            assert isinstance(manifest, ArtifactVersionManifest)
+            assert isinstance(manifest, ArtifactManifest)
             manifest.description = self.description
             manifest.metadata = self.metadata
             manifest.version_alias = self.version_alias
@@ -504,8 +504,8 @@ def _log_artifact_version_helper(
 
     artifact_manifest = None
     if artifact_type == ArtifactType.ARTIFACT:
-        _source_cls = get_autogen_type(ArtifactVersionManifest, "source")
-        artifact_manifest = ArtifactVersionManifest(
+        _source_cls = get_autogen_type(ArtifactManifest, "source")
+        artifact_manifest = ArtifactManifest(
             name=name,
             ml_repo=ml_repo,
             description=description,

truefoundry/ml/log_types/artifacts/dataset.py

@@ -5,16 +5,16 @@ import tempfile
 from typing import Any, Dict, Iterator, List, NamedTuple, Optional, Tuple, Union
 from warnings import warn
 
-from truefoundry.ml.artifact.truefoundry_artifact_repo import (
-    ArtifactIdentifier,
-    MlFoundryArtifactsRepository,
-)
-from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
+from truefoundry.ml._autogen.client import (  # type: ignore[attr-defined]
     DatasetDto,
     DeleteDatasetRequestDto,
     MlfoundryArtifactsApi,
     UpdateDatasetRequestDto,
 )
+from truefoundry.ml.artifact.truefoundry_artifact_repo import (
+    ArtifactIdentifier,
+    MlFoundryArtifactsRepository,
+)
 from truefoundry.ml.entities import FileInfo
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.log_types.artifacts.utils import (

truefoundry/ml/log_types/artifacts/general_artifact.py

@@ -2,7 +2,7 @@ import collections
 import tempfile
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
-from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
+from truefoundry.ml._autogen.client import (  # type: ignore[attr-defined]
     ArtifactType,
 )
 from truefoundry.ml.exceptions import MlFoundryException