truefoundry 0.5.9__py3-none-any.whl → 0.5.10__py3-none-any.whl

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,11 +37,14 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
             timeout=3,
             params={"id": integration_id, "type": "blob-storage"},
             session=session,
+            headers={
+                "User-Agent": get_user_agent(),
+            },
         )
 
         try:
-            data = request_handling(response)
-            assert isinstance(data, dict)
+            result = request_handling(response)
+            assert isinstance(result, dict)
         except HttpRequestException as he:
             raise MlFoundryException(
                 f"Failed to get storage integration from id: {integration_id}. Error: {he.message}",
@@ -51,13 +55,10 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
                 f"Failed to get storage integration from id: {integration_id}. Error: {str(e)}"
             ) from None
 
+        data = result.get("data", result.get("providerAccounts"))
         # TODO (chiragjn): Parse this using Pydantic
-        if (
-            data.get("providerIntegrations")
-            and len(data["providerIntegrations"]) > 0
-            and data["providerIntegrations"][0]
-        ):
-            return data["providerIntegrations"][0]
+        if data and len(data) > 0 and data[0]:
+            return data[0]
         else:
             raise MlFoundryException(
                 f"Invalid storage integration id: {integration_id}"

truefoundry/ml/log_types/artifacts/artifact.py

@@ -3,11 +3,10 @@ import datetime
 import json
 import os
 import tempfile
-import typing
 import uuid
 import warnings
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, NamedTuple, Optional, Union
 
 from truefoundry.common.warnings import TrueFoundryDeprecationWarning
 from truefoundry.ml.artifact.truefoundry_artifact_repo import (
@@ -24,16 +23,11 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ExperimentsApi,
     ExternalBlobStorageSource,
     FinalizeArtifactVersionRequestDto,
-    Manifest,
     MlfoundryArtifactsApi,
     NotifyArtifactVersionFailureDto,
     TrueFoundryManagedSource,
     UpdateArtifactVersionRequestDto,
 )
-from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
-    InternalMetadata as InternalMetadataDto,
-)
-from truefoundry.ml.autogen.entities.artifacts import ChatPrompt
 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 (
@@ -41,6 +35,7 @@ from truefoundry.ml.log_types.artifacts.utils import (
     _validate_artifact_metadata,
     _validate_description,
     calculate_total_size,
+    get_autogen_type,
 )
 from truefoundry.ml.logger import logger
 from truefoundry.ml.session import _get_api_client
@@ -404,13 +399,18 @@ class ArtifactVersion:
         else:
             manifest = None
         try:
+            _manifest_cls = get_autogen_type(
+                UpdateArtifactVersionRequestDto, "manifest"
+            )
             _artifact_version = (
                 self._mlfoundry_artifacts_api.update_artifact_version_post(
                     update_artifact_version_request_dto=UpdateArtifactVersionRequestDto(
                         id=self._artifact_version.id,
                         description=self.description,
                         artifact_metadata=self.metadata,
-                        manifest=manifest,
+                        manifest=_manifest_cls.from_dict(manifest.to_dict())
+                        if manifest
+                        else None,
                     )
                 )
             )
@@ -423,66 +423,6 @@ class ArtifactVersion:
             self._set_mutable_attrs()
 
 
-class ChatPromptVersion(ArtifactVersion):
-    def __init__(self, artifact_version: ArtifactVersionDto, artifact: ArtifactDto):
-        if artifact.type != ArtifactType.CHAT_PROMPT:
-            raise ValueError(
-                f"{artifact_version.fqn!r} is not a chat prompt type artifact"
-            )
-        super().__init__(
-            artifact_version=artifact_version,
-            artifact=artifact,
-        )
-        self._chat_prompt = ChatPrompt.parse_obj(
-            artifact_version.internal_metadata.to_dict()
-        )
-
-    @classmethod
-    def from_fqn(cls, fqn: str) -> "ChatPromptVersion":
-        api_client = _get_api_client()
-        mlfoundry_artifacts_api = MlfoundryArtifactsApi(api_client=api_client)
-        _artifact_version = mlfoundry_artifacts_api.get_artifact_version_by_fqn_get(
-            fqn=fqn
-        )
-        artifact_version = _artifact_version.artifact_version
-        _artifact = mlfoundry_artifacts_api.get_artifact_by_id_get(
-            id=artifact_version.artifact_id
-        )
-        return cls(
-            artifact_version=_artifact_version.artifact_version,
-            artifact=_artifact.artifact,
-        )
-
-    @property
-    def model(self) -> str:
-        return self._chat_prompt.model_configuration.model
-
-    @property
-    def provider(self) -> str:
-        return self._chat_prompt.model_configuration.provider
-
-    @property
-    def messages(self) -> List[Dict[str, Any]]:
-        return [message.dict() for message in self._chat_prompt.messages]
-
-    @property
-    def parameters(self) -> Dict[str, Any]:
-        _parameters = self._chat_prompt.model_configuration.parameters
-        return _parameters.dict(exclude_unset=True) if _parameters else {}
-
-    @property
-    def extra_parameters(self) -> Dict[str, Any]:
-        _extra_parameters = self._chat_prompt.model_configuration.extra_parameters
-        return _extra_parameters if _extra_parameters else {}
-
-    @property
-    def variables(self) -> Dict[str, Any]:
-        return self._chat_prompt.variables or {}
-
-    def update(self):
-        raise NotImplementedError("ChatPromptVersion does not support update operation")
-
-
 def _log_artifact_version_helper(
     run: Optional["MlFoundryRun"],
     name: str,
@@ -525,6 +465,7 @@ def _log_artifact_version_helper(
     )
     version_id = _create_artifact_response.id
     artifact_storage_root = _create_artifact_response.artifact_storage_root
+    total_size = 0
     if isinstance(artifact_dir, tempfile.TemporaryDirectory):
         # Source is of type TrueFoundryManagedSource
         source = TrueFoundryManagedSource(type="truefoundry", uri=artifact_storage_root)
@@ -556,13 +497,6 @@ def _log_artifact_version_helper(
             raise MlFoundryException("Failed to log Artifact") from e
         finally:
             artifact_dir.cleanup()
-
-        # Note: Here we call from_dict instead of directly passing in init and relying on it
-        # to convert because the complicated union of types generates a custom type to handle casting
-        # Check the source of `InternalMetadataDto` to see the generated code
-        internal_metadata_dto = InternalMetadataDto.from_dict(
-            internal_metadata.dict() if internal_metadata is not None else {}
-        )
     elif isinstance(artifact_dir, BlobStorageDirectory):
         source = ExternalBlobStorageSource(type="external", uri=artifact_dir.uri)
     else:
@@ -570,7 +504,7 @@ def _log_artifact_version_helper(
 
     artifact_manifest = None
     if artifact_type == ArtifactType.ARTIFACT:
-        _source_cls = typing.get_type_hints(ArtifactVersionManifest)["source"]
+        _source_cls = get_autogen_type(ArtifactVersionManifest, "source")
         artifact_manifest = ArtifactVersionManifest(
             name=name,
             ml_repo=ml_repo,
@@ -579,15 +513,18 @@ def _log_artifact_version_helper(
             source=_source_cls.from_dict(source.dict()),
             step=step,
         )
+    _manifest_cls = get_autogen_type(FinalizeArtifactVersionRequestDto, "manifest")
     finalize_artifact_version_request_dto = FinalizeArtifactVersionRequestDto(
         id=version_id,
         run_uuid=run.run_id if run else None,
-        internal_metadata=internal_metadata_dto if internal_metadata else None,
+        internal_metadata=internal_metadata.dict()
+        if internal_metadata is not None
+        else {},
         artifact_metadata=metadata,
         data_path=INTERNAL_METADATA_PATH if internal_metadata else None,
         step=step,
         artifact_size=total_size,
-        manifest=Manifest.from_dict(artifact_manifest.to_dict())
+        manifest=_manifest_cls.from_dict(artifact_manifest.to_dict())
         if artifact_manifest
         else None,
     )

truefoundry/ml/log_types/artifacts/model.py

@@ -4,7 +4,6 @@ import json
 import logging
 import os.path
 import tempfile
-import typing
 import uuid
 import warnings
 from pathlib import Path
@@ -23,7 +22,6 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ExternalBlobStorageSource,
     FinalizeArtifactVersionRequestDto,
     Framework,
-    Manifest,
     MlfoundryArtifactsApi,
     ModelDto,
     ModelVersionDto,
@@ -46,6 +44,7 @@ from truefoundry.ml.log_types.artifacts.utils import (
     _validate_artifact_metadata,
     _validate_description,
     calculate_total_size,
+    get_autogen_type,
 )
 from truefoundry.ml.model_framework import (
     ModelFrameworkType,
@@ -618,8 +617,7 @@ def _log_model_version(  # noqa: C901
     else:
         raise MlFoundryException("Invalid model_file_or_folder provided")
 
-    _source_cls = typing.get_type_hints(ModelVersionManifest)["source"]
-
+    _source_cls = get_autogen_type(ModelVersionManifest, "source")
     # Auto fetch the framework & environment details if not provided
     framework = _ModelFramework.to_model_framework_type(framework)
     if framework and isinstance(model_file_or_folder, str):
@@ -639,6 +637,7 @@ def _log_model_version(  # noqa: C901
         environment=environment,
         step=step,
     )
+    _manifest_cls = get_autogen_type(FinalizeArtifactVersionRequestDto, "manifest")
     artifact_version_response = mlfoundry_artifacts_api.finalize_artifact_version_post(
         finalize_artifact_version_request_dto=FinalizeArtifactVersionRequestDto(
             id=version_id,
@@ -647,7 +646,7 @@ def _log_model_version(  # noqa: C901
             artifact_metadata=metadata,
             internal_metadata=None,
             step=model_manifest.step,
-            manifest=Manifest.from_dict(model_manifest.to_dict()),
+            manifest=_manifest_cls.from_dict(model_manifest.to_dict()),
         )
     )
     return ModelVersion.from_fqn(fqn=artifact_version_response.artifact_version.fqn)

truefoundry/ml/log_types/artifacts/utils.py

@@ -2,11 +2,13 @@ import json
 import logging
 import os
 import posixpath
+import typing
 from pathlib import Path, PureWindowsPath
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Type, Union
 
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.log_types.artifacts.constants import DESCRIPTION_MAX_LENGTH
+from truefoundry.pydantic_v1 import BaseModel
 
 logger = logging.getLogger(__name__)
 
@@ -206,9 +208,23 @@ def calculate_total_size(
     Tells about the size of the artifact
 
     Args:
-        directory (str): directory path
+        paths (str): list of paths to include in total size calculation
 
     Returns:
         total size of the artifact
     """
     return sum(os.stat(os.path.realpath(file_path)).st_size for file_path in paths)
+
+
+def get_autogen_type(parent_type: Type[BaseModel], field_name: str) -> Type[BaseModel]:
+    type_ = typing.get_type_hints(parent_type)[field_name]
+    origin = typing.get_origin(type_)
+
+    if origin is None:
+        return type_
+    elif origin is typing.Union:
+        args = typing.get_args(type_)
+        if len(args) == 2 and args[-1] is type(None):
+            return args[0]
+
+    raise NotImplementedError(f"Cannot extract main type from {type_}")

truefoundry/ml/mlfoundry_api.py

@@ -52,7 +52,6 @@ from truefoundry.ml.log_types.artifacts.artifact import (
     ArtifactPath,
     ArtifactVersion,
     BlobStorageDirectory,
-    ChatPromptVersion,
 )
 from truefoundry.ml.log_types.artifacts.dataset import DataDirectory
 from truefoundry.ml.log_types.artifacts.general_artifact import _log_artifact_version
@@ -1556,45 +1555,6 @@ class MlFoundry:
             if not datasets or not page_token:
                 done = True
 
-    def get_chat_prompt_version_by_fqn(self, fqn: str) -> ChatPromptVersion:
-        """
-        Get the prompt by prompt version FQN
-        Args:
-            fqn (str): Fully qualified name of the chat prompt version.
-        Returns:
-            ChatPromptVersion: An instance of the ChatPromptVersion class with the fetched prompt details.
-
-        Examples:
-            ```python
-            from truefoundry.ml import get_client
-            client = get_client()
-
-            # Get the chat prompt by FQN
-            fqn = 'chat_prompt:truefoundry/prompt-demo/demo-prompt-1:2'
-            chat_prompt = client.get_chat_prompt_version_by_fqn(fqn=fqn)
-
-            # Check the chat prompt values
-            print(chat_prompt)
-
-            # Use with OpenAI client
-            from openai import OpenAI
-            openai_client = OpenAI(api_key="api_key")
-
-            # Make the OpenAI call
-            response = openai_client.chat.completions.create(
-                model=chat_prompt.model,
-                messages=chat_prompt.messages,
-                **chat_prompt.parameters,
-                **chat_prompt.extra_parameters
-            )
-
-            # Extract the content from the response
-            response_content = response.choices[0].message.content
-            print(response_content)
-            ```
-        """
-        return ChatPromptVersion.from_fqn(fqn=fqn)
-
 
 def get_client() -> MlFoundry:
     """Initializes and returns the mlfoundry client.

truefoundry/ml/session.py

@@ -5,14 +5,17 @@ from typing import TYPE_CHECKING, Dict, Optional
 
 from truefoundry.common.request_utils import urllib3_retry
 from truefoundry.common.session import Session
-from truefoundry.common.utils import get_tfy_servers_config, relogin_error_message
+from truefoundry.common.utils import (
+    get_tfy_servers_config,
+    get_user_agent,
+    relogin_error_message,
+)
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ApiClient,
     Configuration,
 )
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.logger import logger
-from truefoundry.version import __version__
 
 if TYPE_CHECKING:
     from truefoundry.ml.mlfoundry_run import MlFoundryRun
@@ -99,7 +102,7 @@ class MLFoundryServerApiClient(ApiClient):
         )
         configuration.retries = urllib3_retry(retries=2)
         api_client = cls(session=session, configuration=configuration)
-        api_client.user_agent = f"truefoundry-cli/{__version__}"
+        api_client.user_agent = get_user_agent()
         return api_client
 
     def _ensure_session(self):

{truefoundry-0.5.9.dist-info → truefoundry-0.5.10.dist-info}/METADATA

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: truefoundry
-Version: 0.5.9
+Version: 0.5.10
 Summary: TrueFoundry CLI
 Author-email: TrueFoundry Team <abhishek@truefoundry.com>
-Requires-Python: <3.13,>=3.8.1
+Requires-Python: <3.14,>=3.8.1
 Requires-Dist: aenum<4.0.0,>=3.0.0
 Requires-Dist: click<9.0.0,>=7.0.0
 Requires-Dist: coolname<2.0.0,>=1.1.0
@@ -14,8 +14,7 @@ Requires-Dist: gitpython<4.0.0,>=3.1.43
 Requires-Dist: importlib-metadata<9.0.0,>=4.11.3
 Requires-Dist: importlib-resources<7.0.0,>=5.2.0
 Requires-Dist: mako<2.0.0,>=1.1.6
-Requires-Dist: numpy<2.0.0,>=1.23.0; python_version < '3.12'
-Requires-Dist: numpy<2.0.0,>=1.26.0; python_version >= '3.12'
+Requires-Dist: numpy<3.0.0,>=1.23.0
 Requires-Dist: openai<2.0.0,>=1.16.2
 Requires-Dist: packaging<25.0,>=20.0
 Requires-Dist: pydantic<3.0.0,>=1.8.2
@@ -30,8 +29,6 @@ Requires-Dist: requests<3.0.0,>=2.18.0
 Requires-Dist: requirements-parser<0.12.0,>=0.11.0
 Requires-Dist: rich-click<2.0.0,>=1.2.1
 Requires-Dist: rich<14.0.0,>=13.7.1
-Requires-Dist: scipy<2.0.0,>=1.12.0; python_version >= '3.12'
-Requires-Dist: scipy<2.0.0,>=1.5.0; python_version < '3.12'
 Requires-Dist: tqdm<5.0.0,>=4.0.0
 Requires-Dist: typing-extensions>=4.0
 Requires-Dist: urllib3<3,>=1.26.18