truefoundry 0.3.3__py3-none-any.whl → 0.4.0.dev1__py3-none-any.whl

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/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/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/deploy/cli/util.py

@@ -36,7 +36,9 @@ def handle_exception(exception):
         )
     elif isinstance(exception, ConnectionError):
         print_dict_as_table_panel(
-            {"Error": "Couldn't connect to TrueFoundry"},
+            {
+                "Error": "Couldn't connect to TrueFoundry. Please make sure that the provided `--host` is correct."
+            },
             title="Command Failed",
             border_style="red",
         )

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/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/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/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/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/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/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/langchain/utils.py

@@ -2,10 +2,7 @@ import os
 from typing import Dict, List, Optional
 from urllib.parse import urljoin
 
-import requests
 from langchain.pydantic_v1 import BaseModel
-from requests.adapters import HTTPAdapter
-from urllib3.util.retry import Retry
 
 from truefoundry.deploy.lib.auth.servicefoundry_session import ServiceFoundrySession
 
@@ -45,41 +42,3 @@ def validate_tfy_environment(values: Dict):
     values["tfy_llm_gateway_url"] = gateway_url
     values["tfy_api_key"] = api_key
     return values
-
-
-def requests_retry_session(
-    retries=5,
-    backoff_factor=0.3,
-    status_forcelist=(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 = 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,
-    )
-    adapter = HTTPAdapter(max_retries=retry)
-    session.mount("http://", adapter)
-    session.mount("https://", adapter)
-    return session

truefoundry/ml/__init__.py

@@ -1,6 +1,46 @@
-try:
-    from mlfoundry import *  # noqa: F403
-except ImportError:
-    print(
-        "To use metadata/artifact logging features, please run 'pip install truefoundry[ml]'."
-    )
+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()