truefoundry 0.5.5rc2__py3-none-any.whl → 0.5.6rc1__py3-none-any.whl

truefoundry/ml/autogen/entities/artifacts.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  artifacts.json
-#   timestamp: 2025-01-03T09:11:54+00:00
+#   timestamp: 2025-01-28T10:54:09+00:00
 
 from __future__ import annotations
 
@@ -10,27 +10,6 @@ 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
@@ -44,25 +23,29 @@ class Method(str, Enum):
     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",
+        None, description="+sort=2\n+label=Description"
     )
     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",
+        description="+label=Metadata\n+docs=Key value pairs to store additional metadata\n+usage=Key value pairs to store additional metadata\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)",
+        description="+label=Version Alias\n+usage=The version alias 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+docs=The version alias 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 it can include '.' and '-' in between (e.g. v1.0.0, v1-prod, v3-dev, etc)",
+    )
+    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",
+    )
+    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+docs=name of the ML Repo\n+uiType=Hidden",
+        )
     )
 
 
@@ -95,16 +78,17 @@ class BlobStorageReference(BaseModel):
     )
 
 
-class ExternalArtifactSource(BaseModel):
+class ExternalBlobStorageSource(BaseModel):
     """
-    +label=External
+    +label=External Blob Storage
     """
 
     type: Literal["external"] = Field(
-        ...,
-        description="+label=Type\n+usage=Type of the source\n+value=external\n+type=External",
+        ..., description="+label=Type\n+usage=Type of the source\n+value=external"
+    )
+    uri: str = Field(
+        ..., description="+label=URI\n+usage=URI of the Blob Storage source"
     )
-    uri: str = Field(..., description="+label=URI\n+usage=URI of the Artifact source")
 
 
 class FastAIFramework(BaseModel):
@@ -415,18 +399,17 @@ class TransformersFramework(BaseModel):
     )
 
 
-class TrueFoundryArtifactSource(BaseModel):
+class TrueFoundryManagedSource(BaseModel):
     """
-    +usage=Source for the Artifact or Model
+    +usage=Source for the entity
     +label=Upload
     """
 
     type: Literal["truefoundry"] = Field(
-        ...,
-        description="+usage=Type of the source\n+type=Upload\n+value=truefoundry\n+label=Upload",
+        ..., description="+label=Type\n+usage=Type of the source\n+value=truefoundry"
     )
-    uri: str = Field(
-        ...,
+    uri: Optional[str] = Field(
+        None,
         description='+label=URI\n+uiType=ModelUploadInput\n+uiProps={"hideClear":true}',
     )
 
@@ -480,14 +463,26 @@ class XGBoostSerializationFormat(str, Enum):
     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"
+class Agent(BaseArtifactVersion):
+    type: Literal["agent"] = Field(..., description="+value=agent")
+    available_tools: List[constr(min_length=1)] = Field(
+        ..., description="+sort=30\n+uiType=Hidden"
     )
-    description: constr(min_length=1, max_length=128) = Field(
+    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 AgentOpenAPITool(BaseArtifactVersion):
+    type: Literal["openapi-tool"] = Field(..., description="+value=openapi-tool")
     openapi_spec: Union[BlobStorageReference, Dict[str, Any]] = Field(
         ..., description="+sort=20\n+uiType=OpenapiSchema"
     )
@@ -510,12 +505,17 @@ class AgentOpenAPIToolWithFQN(AgentOpenAPITool):
     fqn: str
 
 
+class AgentWithFQN(Agent):
+    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(
+    source: Union[TrueFoundryManagedSource, ExternalBlobStorageSource] = Field(
         ...,
         description="+label=Artifact Source\n+usage=Source for the Artifact version\n+uiType=Group",
     )
@@ -538,38 +538,50 @@ class AssistantMessage(BaseModel):
     )
 
 
-class BasePrompt(BaseModel):
+class ChatPrompt(BaseArtifactVersion):
+    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",
+    )
+    model_configuration: ModelConfiguration
+
+
+class DataDirectory(BaseModel):
     """
-    Main prompt structure
+    +label=Data Directory
     """
 
-    name: constr(regex=r"^[a-zA-Z][a-zA-Z0-9-_]{0,254}[a-zA-Z0-9]$") = Field(
+    type: Literal["data-dir"] = 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='+label=Type\n+usage=Data Directory\n+docs=Data Directory\n+value="data-dir"',
     )
-    description: Optional[constr(min_length=1, max_length=512)] = Field(
+    description: Optional[constr(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",
+        description="+label=Description\n+usage=Description of the data directory\n+docs=Description of the data directory",
     )
-    model_configuration: ModelConfiguration
-    metadata: Dict[str, str] = Field(
+    metadata: Dict[str, Any] = 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.",
+        description="+label=Metadata\n+docs=Metadata for the data directory\n+usage=Metadata for the data directory\n+uiType=JsonInput",
     )
-
-
-class ChatPrompt(BasePrompt):
-    type: Literal["chat_prompt"] = Field(
+    name: constr(regex=r"^[A-Za-z0-9_\-]+$", max_length=256) = 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"',
+        description="+sort=1\n+label=Name\n+usage=Name of the data directory\n+docs=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",
     )
-    messages: List[Union[SystemMessage, AssistantMessage, UserMessage]] = Field(
+    ml_repo: constr(regex=r"^[a-zA-Z][a-zA-Z0-9\-]{1,98}[a-zA-Z0-9]$") = 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.",
+        description="+label=ML Repo\n+usage=name of the ML Repo of the data directory\n+docs=name of the ML Repo of the data directory\n+uiType=Hidden",
     )
-    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",
+    source: Union[TrueFoundryManagedSource, ExternalBlobStorageSource] = Field(
+        ...,
+        description="+label=Data Directory Source\n+usage=Source for the Data Directory\n+docs=Source for the Data Directory\n+uiType=Group",
     )
 
 
@@ -618,7 +630,7 @@ class ModelVersion(BaseArtifactVersion):
     type: Literal["model-version"] = Field(
         ..., description='+label=Type\n+usage=Model Version\n+value="model-version"'
     )
-    source: Union[TrueFoundryArtifactSource, ExternalArtifactSource] = Field(
+    source: Union[TrueFoundryManagedSource, ExternalBlobStorageSource] = Field(
         ...,
         description="+label=Model Source\n+usage=Source for the model version\n+uiType=Group",
     )
@@ -648,5 +660,11 @@ class ModelVersion(BaseArtifactVersion):
 
 class VersionedArtifactType(BaseModel):
     __root__: Union[
-        ChatPrompt, AgentOpenAPITool, Agent, AgentApp, ModelVersion, ArtifactVersion
+        ChatPrompt,
+        AgentOpenAPITool,
+        Agent,
+        AgentApp,
+        ModelVersion,
+        ArtifactVersion,
+        DataDirectory,
     ]

truefoundry/ml/log_types/artifacts/artifact.py

@@ -21,12 +21,13 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactVersionManifest,
     CreateArtifactVersionRequestDto,
     DeleteArtifactVersionsRequestDto,
-    ExternalArtifactSource,
+    ExperimentsApi,
+    ExternalBlobStorageSource,
     FinalizeArtifactVersionRequestDto,
     Manifest,
     MlfoundryArtifactsApi,
     NotifyArtifactVersionFailureDto,
-    TrueFoundryArtifactSource,
+    TrueFoundryManagedSource,
     UpdateArtifactVersionRequestDto,
 )
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
@@ -488,29 +489,36 @@ def _log_artifact_version_helper(
     artifact_type: ArtifactType,
     artifact_dir: Union[tempfile.TemporaryDirectory, BlobStorageDirectory],
     dest_to_src_map: Dict[str, str],
-    mlfoundry_artifacts_api: Optional[MlfoundryArtifactsApi] = None,
-    ml_repo_id: Optional[str] = None,
+    ml_repo: Optional[str] = None,
     description: Optional[str] = None,
     internal_metadata: Optional[BaseModel] = None,
     metadata: Optional[Dict[str, Any]] = None,
     step: int = 0,
     progress: Optional[bool] = None,
 ) -> ArtifactVersion:
-    if (run and mlfoundry_artifacts_api) or (not run and not mlfoundry_artifacts_api):
-        raise MlFoundryException(
-            "Exactly one of run, mlfoundry_artifacts_api should be passed"
-        )
-    if mlfoundry_artifacts_api and not ml_repo_id:
-        raise MlFoundryException(
-            "If mlfoundry_artifacts_api is passed, ml_repo_id must also be passed"
-        )
+    if not run and not ml_repo:
+        raise MlFoundryException("Exactly one of run, ml_repo should be passed")
+
     if run:
         mlfoundry_artifacts_api = run._mlfoundry_artifacts_api
+        experiments_api = run._experiments_api
+        ml_repo = experiments_api.get_experiment_get(
+            experiment_id=str(run._experiment_id)
+        ).experiment.name
+        ml_repo_id = run._experiment_id
+    else:
+        assert ml_repo is not None
+        api_client = _get_api_client()
+        mlfoundry_artifacts_api = MlfoundryArtifactsApi(api_client=api_client)
+        experiments_api = ExperimentsApi(api_client=api_client)
+        ml_repo_id = experiments_api.get_experiment_by_name_get(
+            experiment_name=ml_repo
+        ).experiment.experiment_id
 
     assert mlfoundry_artifacts_api is not None
     _create_artifact_response = mlfoundry_artifacts_api.create_artifact_version_post(
         create_artifact_version_request_dto=CreateArtifactVersionRequestDto(
-            experiment_id=int(run._experiment_id if run else ml_repo_id),
+            experiment_id=int(ml_repo_id),
             name=name,
             artifact_type=artifact_type,
         )
@@ -518,10 +526,8 @@ def _log_artifact_version_helper(
     version_id = _create_artifact_response.id
     artifact_storage_root = _create_artifact_response.artifact_storage_root
     if isinstance(artifact_dir, tempfile.TemporaryDirectory):
-        # Source is of type TrueFoundryArtifactSource
-        source = TrueFoundryArtifactSource(
-            type="truefoundry", uri=artifact_storage_root
-        )
+        # Source is of type TrueFoundryManagedSource
+        source = TrueFoundryManagedSource(type="truefoundry", uri=artifact_storage_root)
         artifacts_repo = MlFoundryArtifactsRepository(
             artifact_identifier=ArtifactIdentifier(
                 artifact_version_id=uuid.UUID(version_id),
@@ -558,7 +564,7 @@ def _log_artifact_version_helper(
             internal_metadata.dict() if internal_metadata is not None else {}
         )
     elif isinstance(artifact_dir, BlobStorageDirectory):
-        source = ExternalArtifactSource(type="external", uri=artifact_dir.uri)
+        source = ExternalBlobStorageSource(type="external", uri=artifact_dir.uri)
     else:
         raise MlFoundryException("Invalid artifact_dir provided")
 
@@ -566,6 +572,8 @@ def _log_artifact_version_helper(
     if artifact_type == ArtifactType.ARTIFACT:
         _source_cls = typing.get_type_hints(ArtifactVersionManifest)["source"]
         artifact_manifest = ArtifactVersionManifest(
+            name=name,
+            ml_repo=ml_repo,
             description=description,
             metadata=metadata or {},
             source=_source_cls.from_dict(source.dict()),

truefoundry/ml/log_types/artifacts/general_artifact.py

@@ -4,7 +4,6 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactType,
-    MlfoundryArtifactsApi,
 )
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.log_types.artifacts.artifact import (
@@ -27,21 +26,14 @@ def _log_artifact_version(
     run: Optional["MlFoundryRun"],
     name: str,
     artifact_paths: List[Union[ArtifactPath, Tuple[str, Optional[str]], Tuple[str]]],
-    mlfoundry_artifacts_api: Optional[MlfoundryArtifactsApi] = None,
-    ml_repo_id: Optional[str] = None,
+    ml_repo: Optional[str] = None,
     description: Optional[str] = None,
     metadata: Optional[Dict[str, Any]] = None,
     step: Optional[int] = 0,
     progress: Optional[bool] = None,
 ) -> ArtifactVersion:
-    if (run and mlfoundry_artifacts_api) or (not run and not mlfoundry_artifacts_api):
-        raise MlFoundryException(
-            "Exactly one of run, mlfoundry_artifacts_api should be passed"
-        )
-    if mlfoundry_artifacts_api and not ml_repo_id:
-        raise MlFoundryException(
-            "If mlfoundry_artifacts_api is passed, ml_repo_id must also be passed"
-        )
+    if not run and not ml_repo:
+        raise MlFoundryException("Exactly one of run, ml_repo should be passed")
     for i, artifact_path in enumerate(artifact_paths):
         if isinstance(artifact_path, ArtifactPath):
             continue
@@ -79,12 +71,11 @@ def _log_artifact_version(
 
     return _log_artifact_version_helper(
         run=run,
-        ml_repo_id=ml_repo_id,
+        ml_repo=ml_repo,
         name=name,
         artifact_type=ArtifactType.ARTIFACT,
         artifact_dir=temp_dir,
         dest_to_src_map=temp_dest_to_src_map,
-        mlfoundry_artifacts_api=mlfoundry_artifacts_api,
         description=description,
         internal_metadata=None,
         metadata=metadata,

truefoundry/ml/log_types/artifacts/model.py

@@ -19,7 +19,8 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactType,
     CreateArtifactVersionRequestDto,
     DeleteArtifactVersionsRequestDto,
-    ExternalArtifactSource,
+    ExperimentsApi,
+    ExternalBlobStorageSource,
     FinalizeArtifactVersionRequestDto,
     Framework,
     Manifest,
@@ -29,7 +30,7 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ModelVersionEnvironment,
     ModelVersionManifest,
     NotifyArtifactVersionFailureDto,
-    TrueFoundryArtifactSource,
+    TrueFoundryManagedSource,
     UpdateModelVersionRequestDto,
 )
 from truefoundry.ml.autogen.models import infer_signature as _infer_signature
@@ -512,8 +513,7 @@ def _log_model_version(  # noqa: C901
     run: Optional["MlFoundryRun"],
     name: str,
     model_file_or_folder: Union[str, BlobStorageDirectory],
-    mlfoundry_artifacts_api: Optional[MlfoundryArtifactsApi] = None,
-    ml_repo_id: Optional[str] = None,
+    ml_repo: Optional[str] = None,
     description: Optional[str] = None,
     metadata: Optional[Dict[str, Any]] = None,
     step: Optional[int] = 0,
@@ -521,24 +521,29 @@ def _log_model_version(  # noqa: C901
     framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
     environment: Optional[ModelVersionEnvironment] = None,
 ) -> ModelVersion:
-    if (run and mlfoundry_artifacts_api) or (not run and not mlfoundry_artifacts_api):
-        raise MlFoundryException(
-            "Exactly one of run, mlfoundry_artifacts_api should be passed"
-        )
-
     if not isinstance(model_file_or_folder, (str, BlobStorageDirectory)):
         raise MlFoundryException(
             "model_file_or_folder should be of type str or BlobStorageDirectory"
         )
 
-    if mlfoundry_artifacts_api and not ml_repo_id:
-        raise MlFoundryException(
-            "If mlfoundry_artifacts_api is passed, ml_repo_id must also be passed"
-        )
+    if not run and not ml_repo:
+        raise MlFoundryException("Exactly one of run, ml_repo should be passed")
+
     if run:
         mlfoundry_artifacts_api = run._mlfoundry_artifacts_api
-
-    assert mlfoundry_artifacts_api is not None
+        experiments_api = run._experiments_api
+        ml_repo = experiments_api.get_experiment_get(
+            experiment_id=str(run._experiment_id)
+        ).experiment.name
+        ml_repo_id = run._experiment_id
+    else:
+        assert ml_repo is not None
+        api_client = _get_api_client()
+        mlfoundry_artifacts_api = MlfoundryArtifactsApi(api_client=api_client)
+        experiments_api = ExperimentsApi(api_client=api_client)
+        ml_repo_id = experiments_api.get_experiment_by_name_get(
+            experiment_name=ml_repo
+        ).experiment.experiment_id
 
     step = step or 0
     total_size = None
@@ -568,7 +573,7 @@ def _log_model_version(  # noqa: C901
     _create_artifact_version_response = (
         mlfoundry_artifacts_api.create_artifact_version_post(
             create_artifact_version_request_dto=CreateArtifactVersionRequestDto(
-                experiment_id=int(run._experiment_id if run else ml_repo_id),
+                experiment_id=int(ml_repo_id),
                 artifact_type=ArtifactType.MODEL,
                 name=name,
             )
@@ -577,10 +582,8 @@ def _log_model_version(  # noqa: C901
     version_id = _create_artifact_version_response.id
     artifact_storage_root = _create_artifact_version_response.artifact_storage_root
     if isinstance(model_file_or_folder, str):
-        # Source is of type TrueFoundryArtifactSource
-        source = TrueFoundryArtifactSource(
-            type="truefoundry", uri=artifact_storage_root
-        )
+        # Source is of type TrueFoundryManagedSource
+        source = TrueFoundryManagedSource(type="truefoundry", uri=artifact_storage_root)
         artifacts_repo = MlFoundryArtifactsRepository(
             artifact_identifier=ArtifactIdentifier(
                 artifact_version_id=uuid.UUID(version_id)
@@ -609,7 +612,9 @@ def _log_model_version(  # noqa: C901
         finally:
             temp_dir.cleanup()
     elif isinstance(model_file_or_folder, BlobStorageDirectory):
-        source = ExternalArtifactSource(type="external", uri=model_file_or_folder.uri)
+        source = ExternalBlobStorageSource(
+            type="external", uri=model_file_or_folder.uri
+        )
     else:
         raise MlFoundryException("Invalid model_file_or_folder provided")
 
@@ -625,6 +630,8 @@ def _log_model_version(  # noqa: C901
     auto_update_environment_details(environment=environment, framework=framework)
 
     model_manifest = ModelVersionManifest(
+        name=name,
+        ml_repo=ml_repo,
         description=description,
         metadata=metadata,
         source=_source_cls.from_dict(source.dict()),

truefoundry/ml/mlfoundry_api.py

@@ -1209,11 +1209,9 @@ class MlFoundry:
             raise MlFoundryException(
                 "artifact_paths cannot be empty, atleast one artifact_path must be passed"
             )
-        ml_repo_id = self._get_ml_repo_id(ml_repo=ml_repo)
         artifact_version = _log_artifact_version(
             run=None,
-            mlfoundry_artifacts_api=self._mlfoundry_artifacts_api,
-            ml_repo_id=ml_repo_id,
+            ml_repo=ml_repo,
             name=name,
             artifact_paths=artifact_paths,
             description=description,
@@ -1345,11 +1343,9 @@ class MlFoundry:
             ```
 
         """
-        ml_repo_id = self._get_ml_repo_id(ml_repo=ml_repo)
         model_version = _log_model_version(
             run=None,
-            mlfoundry_artifacts_api=self._mlfoundry_artifacts_api,
-            ml_repo_id=ml_repo_id,
+            ml_repo=ml_repo,
             name=name,
             model_file_or_folder=model_file_or_folder,
             description=description,

truefoundry/workflow/task.py

@@ -17,7 +17,7 @@ def task(
     Callable[..., FuncOut],
 ]:
     """
-    This is the decorator used to run flyte tasks in Truefoundry.
+    This is the decorator used to run flyte tasks in TrueFoundry.
 
     Tasks are the building blocks of workflow. They represent users code. Tasks have the following properties
 

{truefoundry-0.5.5rc2.dist-info → truefoundry-0.5.6rc1.dist-info}/METADATA

@@ -1,7 +1,7 @@
 Metadata-Version: 2.3
 Name: truefoundry
-Version: 0.5.5rc2
-Summary: Truefoundry CLI
+Version: 0.5.6rc1
+Summary: TrueFoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com
 Requires-Python: >=3.8.1,<3.13