truefoundry 0.3.2__tar.gz → 0.4.0.dev0__tar.gz

{truefoundry-0.3.2 → truefoundry-0.4.0.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: truefoundry
-Version: 0.3.2
+Version: 0.4.0.dev0
 Summary: Truefoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com
@@ -10,37 +10,44 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Provides-Extra: ml
 Provides-Extra: workflow
-Requires-Dist: GitPython (>=3.1.43,<3.2.0)
+Requires-Dist: GitPython (>=3.1.43,<4.0.0)
 Requires-Dist: Mako (>=1.1.6,<2.0.0)
-Requires-Dist: PyJWT (>=2.0.0,<3.0.0)
+Requires-Dist: PyJWT (>=2.4.0,<3.0.0)
 Requires-Dist: PyYAML (>=6.0.0,<7.0.0)
+Requires-Dist: aenum (>=3.0.0,<4.0.0)
 Requires-Dist: click (>=7.0.0,<9.0.0)
+Requires-Dist: coolname (>=1.1.0,<2.0.0)
 Requires-Dist: docker (>=6.1.2,<8.0.0)
 Requires-Dist: fastapi (>=0.56.0,<0.200.0)
 Requires-Dist: filelock (>=3.8.0,<4.0.0)
 Requires-Dist: flytekit (==1.12.2) ; extra == "workflow"
-Requires-Dist: gitignorefile (>=1.1.2,<1.2.0)
-Requires-Dist: importlib-metadata (>=6.0.1,<8.0.0)
-Requires-Dist: importlib-resources (>=5.2.0,<6.0.0)
-Requires-Dist: mlfoundry (==0.11.5) ; extra == "ml"
+Requires-Dist: gitignorefile (>=1.1.2,<2.0.0)
+Requires-Dist: importlib-metadata (>=4.11.3,<9.0.0)
+Requires-Dist: importlib-resources (>=5.2.0,<7.0.0)
+Requires-Dist: numpy (>=1.23.0,<2.0.0) ; python_version < "3.12"
+Requires-Dist: numpy (>=1.26.0,<2.0.0) ; python_version >= "3.12"
 Requires-Dist: openai (>=1.16.2,<2.0.0)
 Requires-Dist: packaging (>=20.0,<25.0)
-Requires-Dist: pydantic (>=1.10.0,<3)
+Requires-Dist: pandas (>=1.0.0,<3.0.0) ; python_version < "3.10"
+Requires-Dist: pandas (>=1.4.0,<3.0.0) ; python_version >= "3.10"
+Requires-Dist: psutil (>=5.9.0,<7.0.0)
+Requires-Dist: pydantic (>=1.8.2,<3.0.0)
 Requires-Dist: pygments (>=2.12.0,<3.0.0)
 Requires-Dist: python-dateutil (>=2.8.2,<3.0.0)
-Requires-Dist: python-dotenv (>=1.0.1,<1.1.0)
+Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
 Requires-Dist: python-socketio[client] (>=5.5.2,<6.0.0)
 Requires-Dist: questionary (>=1.10.0,<2.0.0)
 Requires-Dist: requests (>=2.31.0,<3.0.0)
-Requires-Dist: requirements-parser (>=0.10.2,<0.11.0)
+Requires-Dist: requirements-parser (>=0.11.0,<0.12.0)
 Requires-Dist: rich (>=13.7.1,<14.0.0)
 Requires-Dist: rich-click (>=1.2.1,<2.0.0)
+Requires-Dist: scipy (>=1.12.0,<2.0.0) ; python_version >= "3.12"
+Requires-Dist: scipy (>=1.5.0,<2.0.0) ; python_version < "3.12"
 Requires-Dist: tqdm (>=4.0.0,<5.0.0)
 Requires-Dist: typing-extensions (>=3.10.0)
 Requires-Dist: urllib3 (>=1.26.18,<3)
-Requires-Dist: uvicorn (>=0.13.0,<0.40.0)
+Requires-Dist: uvicorn (>=0.13.0,<1.0.0)
 Requires-Dist: yq (>=3.1.0,<4.0.0)
 Description-Content-Type: text/markdown
 

truefoundry-0.4.0.dev0/pyproject.toml

@@ -0,0 +1,154 @@
+[tool.poetry]
+name = "truefoundry"
+version = "0.4.0.dev0"
+description = "Truefoundry CLI"
+authors = ["Abhishek Choudhary <abhishek@truefoundry.com>"]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = ">=3.8.1,<3.13"
+aenum = ">=3.0.0,<4.0.0"
+click = ">=7.0.0,<9.0.0"
+coolname = ">=1.1.0,<2.0.0"
+docker = ">=6.1.2,<8.0.0"
+fastapi = ">=0.56.0,<0.200.0"
+filelock = ">=3.8.0,<4.0.0"
+flytekit = { version = "1.12.2", optional = true }
+gitignorefile = ">=1.1.2,<2.0.0"
+GitPython = ">=3.1.43,<4.0.0"
+importlib-metadata = ">=4.11.3,<9.0.0"
+importlib-resources = ">=5.2.0,<7.0.0"
+Mako = ">=1.1.6,<2.0.0"
+openai = ">=1.16.2,<2.0.0"
+packaging = ">=20.0,<25.0"
+psutil = ">=5.9.0,<7.0.0"
+pydantic = ">=1.8.2,<3.0.0"
+pygments = ">=2.12.0,<3.0.0"
+PyJWT = ">=2.4.0,<3.0.0"
+python-dateutil = ">=2.8.2,<3.0.0"
+python-dotenv = ">=1.0.1,<2.0.0"
+python-socketio = { version = ">=5.5.2,<6.0.0", extras = ["client"] }
+PyYAML = ">=6.0.0,<7.0.0"
+questionary = ">=1.10.0,<2.0.0"
+requests = ">=2.31.0,<3.0.0"
+requirements-parser = ">=0.11.0,<0.12.0"
+rich = ">=13.7.1,<14.0.0"
+rich-click = ">=1.2.1,<2.0.0"
+tqdm = ">=4.0.0,<5.0.0"
+typing-extensions = ">=3.10.0"
+urllib3 = ">=1.26.18,<3"
+uvicorn = ">=0.13.0,<1.0.0"
+yq = ">=3.1.0,<4.0.0"
+# Check and try to eliminate libs below this comment
+scipy = [
+    # This is split into two because poetry will select a version which has no wheel on python >= 3.12 :/
+    { version = ">=1.5.0,<2.0.0", python = "<3.12" },
+    { version = ">=1.12.0,<2.0.0", python = ">=3.12" },
+]
+numpy = [
+    # This is split into two because poetry will select a version which has no wheel on python >= 3.12 :/
+    { version = ">=1.23.0,<2.0.0", python = "<3.12" },
+    { version = ">=1.26.0,<2.0.0", python = ">=3.12" },
+]
+pandas = [
+    # This is split into two because poetry will select a version which has no wheel on python >= 3.10 :/
+    { version = ">=1.0.0,<3.0.0", python = "<3.10" },
+    { version = ">=1.4.0,<3.0.0", python = ">=3.10" },
+]
+
+
+[tool.poetry.extras]
+workflow = ["flytekit"]
+
+[tool.poetry.group.dev.dependencies]
+Pillow = ">=9.1.1,<11.0.0"
+ipykernel = ">=6.16.0"
+ipython = ">=7.10.0"
+langchain = ">=0.2.0,<0.3.0"
+matplotlib = ">=3.0.0,<4.0.0"
+mypy = "*"
+plotly = ">=5.8.2,<6.0.0"
+pre-commit = ">=3.5.0,<4.0.0"
+pytest = ">=8.2.2,<9.0.0"
+ruff = ">=0.4.10"
+seaborn = ">=0.13.0,<0.14.0"
+scikit-learn = ">=1.3.0,<2.0.0"
+shap =  [
+    { version = ">=0.40.0", python = "<3.12" },
+    { version = ">=0.45.0", python = ">=3.12" },
+]
+
+[tool.poetry-dynamic-versioning]
+enable = false
+
+[build-system]
+requires = ["poetry-core>=1.0.0", "poetry-dynamic-versioning>=1.0.0,<2.0.0"]
+build-backend = "poetry_dynamic_versioning.backend"
+
+[tool.poetry.plugins."console_scripts"]
+tfy = "truefoundry.cli.__main__:main"
+truefoundry = "truefoundry.cli.__main__:main"
+
+[tool.ruff]
+lint.select = [
+    "E",  # pycodestyle errors
+    "W",  # pycodestyle warnings
+    "F",  # pyflakes
+    "C",  # flake8-comprehensions
+    "B",  # flake8-bugbear
+    "Q",  # flake8-quotes
+    "I",
+]
+lint.ignore = ["E501"]
+exclude = ["venv", "repos"]
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["F401"]
+"truefoundry/ml/autogen/client/*" = ["B028", "B904", "C901", "C409", "F841"]
+"truefoundry/ml/autogen/client/models/*" = ["F821"]
+"truefoundry/ml/autogen/client/api_client.py" = ["E721"]
+
+[tool.mypy]
+plugins = [
+  "pydantic.mypy"
+]
+install_types = true
+non_interactive = true
+
+allow_redefinition = false
+allow_untyped_globals = false
+check_untyped_defs = true
+disallow_any_generics = true
+disallow_incomplete_defs = false
+disallow_subclassing_any = true
+disallow_untyped_calls = false
+disallow_untyped_decorators = false
+disallow_untyped_defs = false
+follow_imports = "silent"
+force_uppercase_builtins = true
+no_implicit_optional = true
+no_implicit_reexport = true
+strict_equality = true
+strict_optional = true
+warn_no_return = true
+warn_redundant_casts = true
+warn_return_any = true
+warn_unreachable = true
+warn_unused_ignores = true
+
+packages = [
+    "truefoundry.ml",
+]
+exclude = [
+    "^truefoundry/ml/autogen",
+    "^truefoundry/deploy/auto_gen",
+    "^truefoundry/pydantic_v1.py$",
+]
+disable_error_code = [
+    "import-untyped",
+]
+
+[tool.pydantic-mypy]
+init_forbid_extra = true
+init_typed = true
+warn_required_dynamic_aliases = true

{truefoundry-0.3.2 → truefoundry-0.4.0.dev0}/truefoundry/cli/__main__.py

@@ -1,21 +1,9 @@
 import sys
 
-import click
+import rich_click as click
 
 from truefoundry.deploy.cli.cli import create_truefoundry_cli
-
-MLFOUNDRY_INSTALLED = True
-
-
-try:
-    from mlfoundry.cli.commands import download
-except ImportError:
-    MLFOUNDRY_INSTALLED = False
-
-
-@click.group()
-def ml():
-    """MlFoundry CLI"""
+from truefoundry.ml.cli.cli import get_ml_cli
 
 
 def main():
@@ -25,9 +13,7 @@ def main():
     # If it is another kind of object, it will be printed and the system exit status will be one (i.e., failure).
     try:
         cli = create_truefoundry_cli()
-        if MLFOUNDRY_INSTALLED:
-            ml.add_command(download)
-            cli.add_command(ml)
+        cli.add_command(get_ml_cli())
     except Exception as e:
         raise click.exceptions.UsageError(message=str(e)) from e
     sys.exit(cli())

truefoundry-0.4.0.dev0/truefoundry/common/request_utils.py

@@ -0,0 +1,56 @@
+import requests
+from requests.adapters import HTTPAdapter
+from urllib3 import Retry
+
+
+def urllib3_retry(
+    retries=5,
+    backoff_factor=0.3,
+    status_forcelist=(408, 429, 500, 502, 503, 504),
+    method_whitelist=frozenset({"GET", "POST"}),
+):
+    retry = Retry(
+        total=retries,
+        read=retries,
+        connect=retries,
+        status=retries,
+        backoff_factor=backoff_factor,
+        allowed_methods=method_whitelist,
+        status_forcelist=status_forcelist,
+        respect_retry_after_header=True,
+    )
+    return retry
+
+
+def requests_retry_session(
+    retries=5,
+    backoff_factor=0.3,
+    status_forcelist=(408, 429, 500, 502, 503, 504),
+    method_whitelist=frozenset({"GET", "POST"}),
+    session=None,
+):
+    """
+    Returns a `requests` session with retry capabilities for certain HTTP status codes.
+
+    Args:
+        retries (int): The number of retries for HTTP requests.
+        backoff_factor (float): The backoff factor for exponential backoff during retries.
+        status_forcelist (tuple): A tuple of HTTP status codes that should trigger a retry.
+        method_whitelist (frozenset): The set of HTTP methods that should be retried.
+        session (requests.Session, optional): An optional existing requests session to use.
+
+    Returns:
+        requests.Session: A session with retry capabilities.
+    """
+    # Implementation taken from https://www.peterbe.com/plog/best-practice-with-retries-with-requests
+    session = session or requests.Session()
+    retry = urllib3_retry(
+        retries=retries,
+        backoff_factor=backoff_factor,
+        status_forcelist=status_forcelist,
+        method_whitelist=method_whitelist,
+    )
+    adapter = HTTPAdapter(max_retries=retry)
+    session.mount("http://", adapter)
+    session.mount("https://", adapter)
+    return session

{truefoundry-0.3.2 → truefoundry-0.4.0.dev0}/truefoundry/deploy/auto_gen/models.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  application.json
-#   timestamp: 2024-09-10T13:19:47+00:00
+#   timestamp: 2024-09-19T11:05:06+00:00
 
 from __future__ import annotations
 
@@ -444,6 +444,30 @@ class Image(BaseModel):
     )
 
 
+class JobAlert(BaseModel):
+    """
+    +docs=Describes the configuration for the job alerts
+    +label=Alert
+    """
+
+    notification_channel: str = Field(
+        ...,
+        description='+label=Notification Channel\n+usage=Specify the notification channel to send alerts to\n+uiType=IntegrationSelect\n+uiProps={"integrationType":"notification-channel"}\n+sort=660',
+    )
+    on_start: bool = Field(
+        False,
+        description="+label=On Start\n+usage=Send an alert when the job starts\n+sort=670",
+    )
+    on_completion: bool = Field(
+        False,
+        description="+label=On Completion\n+usage=Send an alert when the job completes\n+sort=680",
+    )
+    on_failure: bool = Field(
+        True,
+        description="+label=On Failure\n+usage=Send an alert when the job fails\n+sort=690",
+    )
+
+
 class KafkaMetricConfig(BaseModel):
     type: Literal["kafka"] = Field(..., description="+value=kafka")
     lag_threshold: conint(ge=1) = Field(
@@ -1350,6 +1374,10 @@ class Job(BaseModel):
         description="+label=Environment Variables\n+usage=Configure environment variables to be injected in the service either as plain text or secrets. [Docs](https://docs.truefoundry.com/docs/env-variables)\n+icon=fa-globe\n+sort=500",
     )
     resources: Optional[Resources] = None
+    alerts: Optional[List[JobAlert]] = Field(
+        None,
+        description="+label=Alerts\n+usage=Configure alerts to be sent when the job starts/fails/completes\n+icon=fa-bell\n+sort=650",
+    )
     retries: conint(ge=0, le=10) = Field(
         0,
         description="+label=Retries\n+usage=Specify the maximum number of attempts to retry a job before it is marked as failed.\n+icon=fa-repeat\n+sort=700",

{truefoundry-0.3.2 → truefoundry-0.4.0.dev0}/truefoundry/deploy/cli/cli.py

@@ -24,7 +24,7 @@ from truefoundry.version import __version__
 click.rich_click.USE_RICH_MARKUP = True
 
 
-def create_truefoundry_cli():
+def create_truefoundry_cli() -> click.MultiCommand:
     """Generates CLI by combining all subcommands into a main CLI and returns in
 
     Returns:

{truefoundry-0.3.2 → truefoundry-0.4.0.dev0}/truefoundry/deploy/lib/auth/credential_provider.py

@@ -31,10 +31,6 @@ class CredentialProvider(ABC):
 
 class EnvCredentialProvider(CredentialProvider):
     def __init__(self) -> None:
-        from truefoundry.deploy.lib.clients.servicefoundry_client import (
-            ServiceFoundryServiceClient,
-        )
-
         logger.debug("Using env var credential provider")
         api_key = os.getenv(API_KEY_ENV_NAME)
         if not api_key:
@@ -45,13 +41,7 @@ class EnvCredentialProvider(CredentialProvider):
         base_url = resolve_base_url().strip("/")
         self._host = base_url
         self._auth_service = AuthServiceClient.from_base_url(base_url=base_url)
-
-        servicefoundry_client = ServiceFoundryServiceClient(
-            init_session=False, base_url=base_url
-        )
-        self._token: Token = servicefoundry_client.get_token_from_api_key(
-            api_key=api_key
-        )
+        self._token: Token = Token(access_token=api_key, refresh_token=None)
 
     @staticmethod
     def can_provide() -> bool:
@@ -123,7 +113,7 @@ class FileCredentialProvider(CredentialProvider):
                     return self.token
 
                 raise Exception(
-                    "Credentials on disk changed while mlfoundry was running."
+                    "Credentials on disk changed while truefoundry was running."
                 )
 
     @property

{truefoundry-0.3.2 → truefoundry-0.4.0.dev0}/truefoundry/deploy/lib/clients/servicefoundry_client.py

@@ -28,7 +28,6 @@ from truefoundry.deploy.lib.model.entity import (
     JobRun,
     PythonSDKConfig,
     TenantInfo,
-    Token,
     TriggerJobResult,
     Workspace,
     WorkspaceResources,
@@ -712,14 +711,6 @@ class ServiceFoundryServiceClient:
         response = request_handling(res)
         return parse_obj_as(List[Deployment], response)
 
-    @check_min_cli_version
-    def get_token_from_api_key(self, api_key: str) -> Token:
-        url = f"{self._api_server_url}/{VERSION_PREFIX}/oauth/api-key/token"
-        data = {"apiKey": api_key}
-        res = requests.get(url, params=data)
-        res = request_handling(res)
-        return Token.parse_obj(res)
-
     @check_min_cli_version
     def apply(self, manifest: Dict[str, Any]) -> None:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/apply"

{truefoundry-0.3.2 → truefoundry-0.4.0.dev0}/truefoundry/deploy/lib/exceptions.py

@@ -6,13 +6,11 @@ class BadRequestException(Exception):
         super().__init__()
         self.status_code = status_code
         self.message = message
+        self.status_code = status_code
 
     def __str__(self):
         return self.message
 
-    def __repr__(self):
-        return self.message
-
 
 class ConfigurationException(Exception):
     def __init__(self, message: Optional[str] = None):
@@ -21,6 +19,3 @@ class ConfigurationException(Exception):
 
     def __str__(self):
         return self.message
-
-    def __repr__(self):
-        return self.message

{truefoundry-0.3.2 → truefoundry-0.4.0.dev0}/truefoundry/deploy/lib/session.py

@@ -7,9 +7,6 @@ from truefoundry.deploy.io.output_callback import OutputCallBack
 from truefoundry.deploy.lib.auth.auth_service_client import AuthServiceClient
 from truefoundry.deploy.lib.auth.credential_file_manager import CredentialsFileManager
 from truefoundry.deploy.lib.auth.credential_provider import EnvCredentialProvider
-from truefoundry.deploy.lib.clients.servicefoundry_client import (
-    ServiceFoundryServiceClient,
-)
 from truefoundry.deploy.lib.clients.utils import resolve_base_url
 from truefoundry.deploy.lib.const import (
     API_KEY_ENV_NAME,
@@ -85,12 +82,7 @@ def login(
             return False
 
         if api_key:
-            servicefoundry_client = ServiceFoundryServiceClient(
-                init_session=False, base_url=host
-            )
-            token = _login_with_api_key(
-                api_key=api_key, servicefoundry_client=servicefoundry_client
-            )
+            token = Token(access_token=api_key, refresh_token=None)
         else:
             auth_service = AuthServiceClient.from_base_url(base_url=host)
             # interactive login
@@ -122,13 +114,6 @@ def logout(
             output_hook.print_line(PROMPT_ALREADY_LOGGED_OUT)
 
 
-def _login_with_api_key(
-    api_key: str, servicefoundry_client: ServiceFoundryServiceClient
-) -> Token:
-    logger.debug("Logging in with api key")
-    return servicefoundry_client.get_token_from_api_key(api_key=api_key)
-
-
 def _login_with_device_code(
     base_url: str,
     auth_service: AuthServiceClient,

{truefoundry-0.3.2/truefoundry → truefoundry-0.4.0.dev0/truefoundry/deploy}/python_deploy_codegen.py

@@ -16,6 +16,7 @@ def generate_code(
     workspace_fqn: str,
 ):
     symbols = ",".join(symbols_to_import)
+    application_type = application_type.replace(" ", "").replace("-", "_")
     code = f"""\
 import logging
 from truefoundry.deploy import (

{truefoundry-0.3.2 → truefoundry-0.4.0.dev0}/truefoundry/langchain/truefoundry_chat.py

@@ -10,8 +10,8 @@ from langchain.schema.messages import (
     SystemMessage,
 )
 
+from truefoundry.common.request_utils import requests_retry_session
 from truefoundry.langchain.utils import (
-    requests_retry_session,
     validate_tfy_environment,
 )
 from truefoundry.logger import logger

{truefoundry-0.3.2 → truefoundry-0.4.0.dev0}/truefoundry/langchain/truefoundry_embeddings.py

@@ -6,8 +6,8 @@ import tqdm
 from langchain.embeddings.base import Embeddings
 from langchain.pydantic_v1 import BaseModel, Extra, Field, root_validator
 
+from truefoundry.common.request_utils import requests_retry_session
 from truefoundry.langchain.utils import (
-    requests_retry_session,
     validate_tfy_environment,
 )
 from truefoundry.logger import logger

{truefoundry-0.3.2 → truefoundry-0.4.0.dev0}/truefoundry/langchain/truefoundry_llm.py

@@ -3,8 +3,8 @@ from typing import Any, Dict, List, Optional
 from langchain.llms.base import LLM
 from langchain.pydantic_v1 import Extra, Field, root_validator
 
+from truefoundry.common.request_utils import requests_retry_session
 from truefoundry.langchain.utils import (
-    requests_retry_session,
     validate_tfy_environment,
 )
 from truefoundry.logger import logger

truefoundry-0.4.0.dev0/truefoundry/langchain/utils.py

@@ -0,0 +1,44 @@
+import os
+from typing import Dict, List, Optional
+from urllib.parse import urljoin
+
+from langchain.pydantic_v1 import BaseModel
+
+from truefoundry.deploy.lib.auth.servicefoundry_session import ServiceFoundrySession
+
+
+class ModelParameters(BaseModel):
+    temperature: Optional[float]
+    maximum_length: Optional[int]
+    top_p: Optional[float]
+    top_k: Optional[int]
+    repetition_penalty: Optional[float]
+    frequency_penalty: Optional[float]
+    presence_penalty: Optional[float]
+    stop_sequences: Optional[List[str]]
+
+
+def validate_tfy_environment(values: Dict):
+    gateway_url = values["tfy_llm_gateway_url"] or os.getenv("TFY_LLM_GATEWAY_URL")
+    api_key = values["tfy_api_key"] or os.getenv("TFY_API_KEY")
+
+    if gateway_url and api_key:
+        values["tfy_llm_gateway_url"] = gateway_url
+        values["tfy_api_key"] = api_key
+        return values
+
+    sfy_session = ServiceFoundrySession()
+    if not sfy_session:
+        raise Exception(
+            "Unauthenticated: Please login using truefoundry login --host <https://example-domain.com>"
+        )
+
+    if not gateway_url:
+        gateway_url = urljoin(sfy_session.base_url, "/api/llm")
+
+    if not api_key:
+        api_key = sfy_session.access_token
+
+    values["tfy_llm_gateway_url"] = gateway_url
+    values["tfy_api_key"] = api_key
+    return values

truefoundry-0.4.0.dev0/truefoundry/ml/__init__.py

@@ -0,0 +1,46 @@
+from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
+    FeatureDto as Feature,
+)
+from truefoundry.ml.enums import (
+    DataSlice,
+    FileFormat,
+    ModelFramework,
+    ModelType,
+    ViewType,
+)
+from truefoundry.ml.exceptions import MlFoundryException
+from truefoundry.ml.log_types import Image, Plot
+from truefoundry.ml.log_types.artifacts.artifact import ArtifactPath, ArtifactVersion
+from truefoundry.ml.log_types.artifacts.dataset import DataDirectory, DataDirectoryPath
+from truefoundry.ml.log_types.artifacts.model import (
+    ModelVersion,
+)
+from truefoundry.ml.log_types.artifacts.model_extras import CustomMetric, ModelSchema
+from truefoundry.ml.logger import init_logger
+from truefoundry.ml.login import login
+from truefoundry.ml.mlfoundry_api import get_client
+from truefoundry.ml.mlfoundry_run import MlFoundryRun
+
+__all__ = [
+    "ArtifactPath",
+    "ArtifactVersion",
+    "CustomMetric",
+    "DataDirectory",
+    "DataDirectoryPath",
+    "DataSlice",
+    "Feature",
+    "FileFormat",
+    "Image",
+    "MlFoundryRun",
+    "MlFoundryException",
+    "ModelFramework",
+    "ModelSchema",
+    "ModelType",
+    "ModelVersion",
+    "Plot",
+    "ViewType",
+    "get_client",
+    "login",
+]
+
+init_logger()