truefoundry 0.5.2__py3-none-any.whl → 0.5.3__py3-none-any.whl

truefoundry/__init__.py

@@ -1,3 +1,12 @@
+from truefoundry.common.warnings import (
+    suppress_truefoundry_deprecation_warnings,
+    surface_truefoundry_deprecation_warnings,
+)
 from truefoundry.deploy.core import login, logout
 
-__all__ = ["login", "logout"]
+surface_truefoundry_deprecation_warnings()
+__all__ = [
+    "login",
+    "logout",
+    "suppress_truefoundry_deprecation_warnings",
+]

truefoundry/autodeploy/agents/base.py

@@ -17,12 +17,12 @@ from openai.types.chat.chat_completion_user_message_param import (
     ChatCompletionUserMessageParam,
 )
 from openai.types.shared_params.function_definition import FunctionDefinition
-from pydantic import BaseModel, ValidationError
 
 from truefoundry.autodeploy.constants import AUTODEPLOY_MODEL_NAME
 from truefoundry.autodeploy.logger import logger
 from truefoundry.autodeploy.tools import Event, RequestEvent, Tool
 from truefoundry.autodeploy.utils.pydantic_compat import model_dump, model_json_schema
+from truefoundry.pydantic_v1 import BaseModel, ValidationError
 
 
 def llm(

truefoundry/autodeploy/agents/developer.py

@@ -4,7 +4,6 @@ from typing import Dict, Optional
 
 import docker
 from openai import OpenAI
-from pydantic import Field
 
 from truefoundry.autodeploy.agents.base import Agent
 from truefoundry.autodeploy.agents.project_identifier import ProjectIdentifier
@@ -18,6 +17,7 @@ from truefoundry.autodeploy.tools import (
     RequestEvent,
     ResponseEvent,
 )
+from truefoundry.pydantic_v1 import Field
 
 
 class Developer(Agent):

truefoundry/autodeploy/agents/project_identifier.py

@@ -4,7 +4,6 @@ import enum
 from typing import List, Optional
 
 from openai import OpenAI
-from pydantic import Field
 from rich.console import Console
 
 from truefoundry.autodeploy.agents.base import Agent
@@ -16,6 +15,7 @@ from truefoundry.autodeploy.tools import (
     RequestEvent,
     ResponseEvent,
 )
+from truefoundry.pydantic_v1 import Field
 
 
 class ComponentType(str, enum.Enum):

truefoundry/autodeploy/agents/tester.py

@@ -4,7 +4,6 @@ from typing import Dict, Optional
 
 import docker
 from openai import OpenAI
-from pydantic import Field
 from rich.console import Console
 
 from truefoundry.autodeploy.agents.base import Agent
@@ -16,6 +15,7 @@ from truefoundry.autodeploy.tools import (
     ResponseEvent,
     SendRequest,
 )
+from truefoundry.pydantic_v1 import Field
 
 
 class Tester(Agent):

truefoundry/autodeploy/cli.py

@@ -32,8 +32,8 @@ from truefoundry.autodeploy.tools.commit import CommitConfirmation
 from truefoundry.autodeploy.tools.docker_run import DockerRun, DockerRunLog
 from truefoundry.autodeploy.utils.client import get_git_binary
 from truefoundry.cli.const import COMMAND_CLS
+from truefoundry.common.session import Session
 from truefoundry.deploy import Build, DockerFileBuild, Job, LocalSource, Port, Service
-from truefoundry.deploy.lib.auth.servicefoundry_session import ServiceFoundrySession
 from truefoundry.deploy.lib.clients.servicefoundry_client import (
     ServiceFoundryServiceClient,
 )
@@ -45,7 +45,7 @@ def _get_openai_client() -> OpenAI:
             api_key=AUTODEPLOY_OPENAI_API_KEY, base_url=AUTODEPLOY_OPENAI_BASE_URL
         )
     try:
-        session = ServiceFoundrySession()
+        session = Session.new()
         resp = requests.get(
             f"{AUTODEPLOY_TFY_BASE_URL}/api/svc/v1/llm-gateway/access-details",
             headers={

truefoundry/autodeploy/tools/base.py

@@ -2,9 +2,10 @@ from __future__ import annotations
 
 from typing import Any, ClassVar, Generator, Optional, Protocol, Type
 
-from pydantic import BaseModel
 from rich.console import Console
 
+from truefoundry.pydantic_v1 import BaseModel
+
 
 class Event(BaseModel):
     def render(self, _: Console) -> Optional[Any]: ...

truefoundry/autodeploy/tools/commit.py

@@ -4,7 +4,6 @@ import os
 import tempfile
 from typing import Any, Generator, Optional
 
-from pydantic import Field
 from rich.console import Console
 from rich.padding import Padding
 from rich.prompt import Confirm
@@ -20,6 +19,7 @@ from truefoundry.autodeploy.tools.base import (
 from truefoundry.autodeploy.utils.client import get_git_binary
 from truefoundry.autodeploy.utils.diff import LLMDiff
 from truefoundry.autodeploy.utils.pydantic_compat import model_dump
+from truefoundry.pydantic_v1 import Field
 
 
 class CommitConfirmation(Event):

truefoundry/autodeploy/tools/docker_build.py

@@ -6,7 +6,6 @@ from typing import Any, Generator, Iterable, Optional
 
 import docker
 from docker.models.images import BuildError, json_stream
-from pydantic import Field
 from rich.console import Console, ConsoleOptions, RenderResult
 from rich.padding import Padding
 from rich.text import Text
@@ -19,6 +18,7 @@ from truefoundry.autodeploy.tools.base import (
     Rule,
     Tool,
 )
+from truefoundry.pydantic_v1 import Field
 
 
 class DockerBuildLog(Event):

truefoundry/autodeploy/tools/docker_run.py

@@ -5,7 +5,6 @@ from typing import Any, Dict, Generator, List, Optional
 
 import docker
 import requests
-from pydantic import Field
 from rich.console import Console, ConsoleOptions, RenderResult
 from rich.padding import Padding
 from rich.text import Text
@@ -18,6 +17,7 @@ from truefoundry.autodeploy.tools.base import (
     Rule,
     Tool,
 )
+from truefoundry.pydantic_v1 import Field
 
 
 class DockerRunLog(Event):

truefoundry/autodeploy/tools/file_type_counts.py

@@ -5,7 +5,6 @@ from collections import Counter
 from typing import Any, Dict, Generator
 
 import gitignorefile
-from pydantic import Field
 from rich.console import Console
 from rich.table import Table
 
@@ -17,6 +16,7 @@ from truefoundry.autodeploy.tools.base import (
     Rule,
     Tool,
 )
+from truefoundry.pydantic_v1 import Field
 
 
 class ShowFileCount(Event):

truefoundry/autodeploy/tools/list_files.py

@@ -5,7 +5,6 @@ from fnmatch import fnmatch
 from typing import Any, Generator, List, Optional
 
 import gitignorefile
-from pydantic import Field
 
 from truefoundry.autodeploy.tools.base import (
     Event,
@@ -15,6 +14,7 @@ from truefoundry.autodeploy.tools.base import (
     Rule,
     Tool,
 )
+from truefoundry.pydantic_v1 import Field
 
 
 class ListFiles(Tool):

truefoundry/autodeploy/tools/read_file.py

@@ -3,8 +3,6 @@ from __future__ import annotations
 import os
 from typing import Any, Generator, List, Optional
 
-from pydantic import BaseModel, Field
-
 from truefoundry.autodeploy.constants import MAX_FILE_SIZE_READ
 from truefoundry.autodeploy.tools.base import (
     Event,
@@ -14,6 +12,7 @@ from truefoundry.autodeploy.tools.base import (
     Rule,
     Tool,
 )
+from truefoundry.pydantic_v1 import BaseModel, Field
 
 
 class Line(BaseModel):

truefoundry/autodeploy/tools/send_request.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 from typing import Any, Generator, Optional
 
 import requests
-from pydantic import Field
 
 from truefoundry.autodeploy.tools.base import (
     Event,
@@ -13,6 +12,7 @@ from truefoundry.autodeploy.tools.base import (
     Rule,
     Tool,
 )
+from truefoundry.pydantic_v1 import Field
 
 
 class SendRequest(Tool):

truefoundry/autodeploy/tools/write_file.py

@@ -4,12 +4,12 @@ from __future__ import annotations
 import os
 from typing import Optional
 
-from pydantic import BaseModel, Field
 from rich.console import Console
 from rich.prompt import Confirm
 
 from truefoundry.autodeploy.tools.base import Tool
 from truefoundry.autodeploy.utils.diff import Diff
+from truefoundry.pydantic_v1 import BaseModel, Field
 
 INTERACTIVE_SESSION = True
 

truefoundry/cli/util.py

@@ -5,13 +5,13 @@ from typing import Dict
 
 import rich_click as click
 from packaging.version import parse as parse_version
-from requests.exceptions import ConnectionError
+from requests.exceptions import ConnectionError, Timeout
 from rich.padding import Padding
 from rich.panel import Panel
 from rich.table import Table
 
 from truefoundry.cli.console import console
-from truefoundry.common.exceptions import BadRequestException
+from truefoundry.common.exceptions import HttpRequestException
 from truefoundry.common.utils import is_debug_env_set
 
 
@@ -25,12 +25,21 @@ def setup_rich_click():
 def handle_exception(exception):
     if is_debug_env_set():
         console.print_exception(show_locals=True)
-    if isinstance(exception, BadRequestException):
+    if isinstance(exception, HttpRequestException):
         print_dict_as_table_panel(
             {"Status Code": str(exception.status_code), "Error": exception.message},
             title="Command Failed",
             border_style="red",
         )
+    elif isinstance(exception, Timeout):
+        loc = ""
+        if exception.request:
+            loc = f" at {exception.request.url}"
+        print_dict_as_table_panel(
+            {"Error": f"Request to TrueFoundry{loc} timed out."},
+            title="Command Failed",
+            border_style="red",
+        )
     elif isinstance(exception, ConnectionError):
         loc = ""
         if exception.request:

truefoundry/common/auth_service_client.py

@@ -7,7 +7,10 @@ import requests
 from truefoundry.common.constants import VERSION_PREFIX
 from truefoundry.common.entities import DeviceCode, Token
 from truefoundry.common.exceptions import BadRequestException
-from truefoundry.common.request_utils import request_handling, requests_retry_session
+from truefoundry.common.request_utils import (
+    request_handling,
+    requests_retry_session,
+)
 from truefoundry.common.utils import poll_for_function, relogin_error_message
 from truefoundry.logger import logger
 
@@ -16,14 +19,13 @@ class AuthServiceClient(ABC):
     def __init__(self, tenant_name: str):
         self._tenant_name = tenant_name
 
-    # TODO (chiragjn): Rename base_url to tfy_host
     @classmethod
-    def from_base_url(cls, base_url: str) -> "AuthServiceClient":
+    def from_tfy_host(cls, tfy_host: str) -> "AuthServiceClient":
         from truefoundry.common.servicefoundry_client import (
             ServiceFoundryServiceClient,
         )
 
-        client = ServiceFoundryServiceClient(base_url=base_url)
+        client = ServiceFoundryServiceClient(tfy_host=tfy_host)
         if client.python_sdk_config.use_sfy_server_auth_apis:
             return ServiceFoundryServerAuthServiceClient(
                 tenant_name=client.tenant_info.tenant_name, url=client._api_server_url
@@ -159,6 +161,7 @@ class AuthServerServiceClient(AuthServiceClient):
         session = requests_retry_session(retries=2)
         response = session.post(url, json=data)
         response_data = request_handling(response)
+        assert isinstance(response_data, dict)
         # TODO: temporary cleanup of incorrect attributes
         return DeviceCode.parse_obj(
             {

truefoundry/common/constants.py

@@ -38,6 +38,9 @@ class TrueFoundrySdkEnv(BaseSettings):
     )
     TFY_INTERNAL_SIGNED_URL_SERVER_MAX_TIMEOUT: int = 5  # default: 5 seconds
     TFY_INTERNAL_SIGNED_URL_SERVER_DEFAULT_TTL: int = 3600  # default: 1 hour
+    TFY_INTERNAL_MULTIPART_UPLOAD_FINALIZE_SIGNED_URL_TIMEOUT: int = (
+        24 * 60 * 60
+    )  # default: 24 hour
     TFY_INTERNAL_SIGNED_URL_REQUEST_TIMEOUT: int = 3600  # default: 1 hour
     TFY_INTERNAL_SIGNED_URL_CLIENT_LOG_LEVEL: str = "WARNING"
 

truefoundry/common/credential_provider.py

@@ -17,10 +17,9 @@ class CredentialProvider(ABC):
     @abstractmethod
     def token(self) -> Token: ...
 
-    # TODO (chiragjn): Rename base_url to tfy_host
     @property
     @abstractmethod
-    def base_url(self) -> str: ...
+    def tfy_host(self) -> str: ...
 
     @staticmethod
     @abstractmethod
@@ -38,7 +37,7 @@ class EnvCredentialProvider(CredentialProvider):
                 f"Value of {TFY_API_KEY_ENV_KEY} env var should be non-empty string"
             )
         self._host = resolve_tfy_host()
-        self._auth_service = AuthServiceClient.from_base_url(base_url=self._host)
+        self._auth_service = AuthServiceClient.from_tfy_host(tfy_host=self._host)
         self._token: Token = Token(access_token=api_key, refresh_token=None)  # type: ignore[call-arg]
 
     @staticmethod
@@ -51,12 +50,12 @@ class EnvCredentialProvider(CredentialProvider):
             if self._token.is_going_to_be_expired():
                 logger.info("Refreshing access token")
                 self._token = self._auth_service.refresh_token(
-                    self._token, self.base_url
+                    self._token, self.tfy_host
                 )
             return self._token
 
     @property
-    def base_url(self) -> str:
+    def tfy_host(self) -> str:
         return self._host
 
 
@@ -70,7 +69,7 @@ class FileCredentialProvider(CredentialProvider):
             self._token = self._last_cred_file_content.to_token()
             self._host = self._last_cred_file_content.host
 
-        self._auth_service = AuthServiceClient.from_base_url(base_url=self._host)
+        self._auth_service = AuthServiceClient.from_tfy_host(tfy_host=self._host)
 
     @staticmethod
     def can_provide() -> bool:
@@ -91,7 +90,7 @@ class FileCredentialProvider(CredentialProvider):
 
                 if new_cred_file_content == self._last_cred_file_content:
                     self._token = self._auth_service.refresh_token(
-                        self._token, self.base_url
+                        self._token, self.tfy_host
                     )
                     self._last_cred_file_content = CredentialsFileContent(
                         host=self._host,
@@ -115,5 +114,5 @@ class FileCredentialProvider(CredentialProvider):
                 )
 
     @property
-    def base_url(self) -> str:
+    def tfy_host(self) -> str:
         return self._host

truefoundry/common/exceptions.py

@@ -1,12 +1,16 @@
 from typing import Optional
 
 
-class BadRequestException(Exception):
-    def __init__(self, status_code: int, message: Optional[str] = None):
-        super().__init__()
-        self.status_code = status_code
-        self.message = message
+# TODO (chiragjn): We need to establish uniform exception handling across codebase
+class HttpRequestException(Exception):
+    def __init__(self, message: str, status_code: Optional[int] = None):
+        self.message = str(message)
         self.status_code = status_code
+        super().__init__(message)
+
+    def __str__(self) -> str:
+        return self.message or ""
+
 
-    def __str__(self):
-        return self.message
+class BadRequestException(HttpRequestException):
+    pass

truefoundry/common/request_utils.py

@@ -1,30 +1,40 @@
-from typing import Optional
+import json
+from contextlib import contextmanager
+from typing import Any, Dict, Optional, Type, Union
 
 import requests
 from requests import Response
 from requests.adapters import HTTPAdapter
 from urllib3 import Retry
 
-from truefoundry.common.exceptions import BadRequestException
+from truefoundry.common.exceptions import BadRequestException, HttpRequestException
 
 
-def request_handling(response: Response):
-    try:
-        status_code = response.status_code
-    except Exception as e:
-        raise Exception("Unknown error occurred. Couldn't get status code.") from e
+# TODO (chiragjn): Rename this to json_response_handling
+def request_handling(response: Response) -> Optional[Any]:
+    status_code = response.status_code
     if 200 <= status_code <= 299:
         if response.content == b"":
+            # TODO (chiragjn): Do we really need this empty check?
             return None
-        return response.json()
-    if 400 <= status_code <= 499:
         try:
-            message = str(response.json()["message"])
-        except Exception:
-            message = response.text
-        raise BadRequestException(status_code=response.status_code, message=message)
+            return response.json()
+        except json.JSONDecodeError as je:
+            raise ValueError(
+                f"Failed to parse response as json. Response: {response.text}",
+            ) from je
+
+    try:
+        message = str(response.json()["message"])
+    except Exception:
+        message = response.text
+
+    if 400 <= status_code <= 499:
+        raise BadRequestException(message=message, status_code=response.status_code)
     if 500 <= status_code <= 599:
-        raise Exception(response.content)
+        raise HttpRequestException(message=message, status_code=response.status_code)
+
+    raise HttpRequestException("Unknown error occurred", status_code=status_code)
 
 
 def urllib3_retry(
@@ -82,3 +92,75 @@ def requests_retry_session(
     session.mount("http://", adapter)
     session.mount("https://", adapter)
     return session
+
+
+def http_request(
+    *,
+    method: str,
+    url: str,
+    token: Optional[str] = None,
+    timeout=None,
+    headers: Optional[Dict[str, str]] = None,
+    session: Optional[requests.Session] = None,
+    exception_class: Type[HttpRequestException] = HttpRequestException,
+    **kwargs,
+) -> requests.Response:
+    session = session or requests.Session()
+    headers = headers or {}
+    if token is not None:
+        headers["Authorization"] = f"Bearer {token}"
+    try:
+        response = session.request(
+            method=method, url=url, headers=headers, timeout=timeout, **kwargs
+        )
+        return response
+    except requests.exceptions.ConnectionError as ce:
+        raise exception_class("Failed to connect to TrueFoundry") from ce
+    except requests.exceptions.Timeout as te:
+        raise exception_class(f"Request to {url} timed out") from te
+    except Exception as e:
+        raise exception_class(f"Request to {url} failed with error {str(e)}") from e
+
+
+@contextmanager
+def cloud_storage_http_request(
+    *,
+    method: str,
+    url: str,
+    session: Optional[requests.Session] = None,
+    timeout: Optional[Union[int, float]] = None,
+    exception_class: Type[HttpRequestException] = HttpRequestException,
+    **kwargs,
+):
+    """
+    Performs an HTTP PUT/GET request using Python's `requests` module with automatic retry.
+    """
+    # Note: This does not support auth and is only meant to be used for pre-signed URLs.
+    session = session or requests_retry_session(retries=5, backoff_factor=0.5)
+    headers = kwargs.get("headers", {}) or {}
+    if "blob.core.windows.net" in url:
+        headers.update({"x-ms-blob-type": "BlockBlob"})
+    if method.lower() not in ("put", "get"):
+        raise ValueError(f"Illegal http method: {method}")
+    yield http_request(
+        method=method,
+        url=url,
+        session=session,
+        timeout=timeout,
+        headers=headers,
+        exception_class=exception_class,
+        **kwargs,
+    )
+
+
+# TODO: Try and eliminate this
+def augmented_raise_for_status(
+    response, exception_class: Type[HttpRequestException] = HttpRequestException
+):
+    try:
+        response.raise_for_status()
+    except requests.exceptions.HTTPError as he:
+        raise exception_class(
+            message=f"Request failed with status code {he.response.status_code}. Response: {he.response.text}",
+            status_code=he.response.status_code,
+        ) from he

truefoundry/common/servicefoundry_client.py

@@ -13,7 +13,10 @@ from truefoundry.common.entities import (
     PythonSDKConfig,
     TenantInfo,
 )
-from truefoundry.common.request_utils import request_handling, requests_retry_session
+from truefoundry.common.request_utils import (
+    request_handling,
+    requests_retry_session,
+)
 from truefoundry.common.utils import (
     get_tfy_servers_config,
     timed_lru_cache,
@@ -25,19 +28,8 @@ from truefoundry.version import __version__
 def check_min_cli_version(fn):
     @functools.wraps(fn)
     def inner(*args, **kwargs):
-        if __version__ != "0.0.0":
-            client: "ServiceFoundryServiceClient" = args[0]
-            # "0.0.0" indicates dev version
-            # noinspection PyProtectedMember
-            min_cli_version_required = client._min_cli_version_required
-            if version.parse(__version__) < version.parse(min_cli_version_required):
-                raise Exception(
-                    "You are using an outdated version of `truefoundry`.\n"
-                    f"Run `pip install truefoundry>={min_cli_version_required}` to install the supported version.",
-                )
-        else:
-            logger.debug("Ignoring minimum cli version check")
-
+        client: "ServiceFoundryServiceClient" = args[0]
+        client.check_min_cli_version()
         return fn(*args, **kwargs)
 
     return inner
@@ -51,34 +43,33 @@ def session_with_retries(
 
 @timed_lru_cache(seconds=30 * 60)
 def _cached_get_tenant_info(api_server_url: str, tenant_host: str) -> TenantInfo:
-    res = session_with_retries().get(
+    response = session_with_retries().get(
         url=f"{api_server_url}/{VERSION_PREFIX}/tenant-id",
         params={"hostName": tenant_host},
     )
-    res = request_handling(res)
-    return TenantInfo.parse_obj(res)
+    response_data = request_handling(response)
+    return TenantInfo.parse_obj(response_data)
 
 
 @timed_lru_cache(seconds=30 * 60)
 def _cached_get_python_sdk_config(api_server_url: str) -> PythonSDKConfig:
-    res = session_with_retries().get(
+    response = session_with_retries().get(
         url=f"{api_server_url}/{VERSION_PREFIX}/min-cli-version"
     )
-    res = request_handling(res)
-    return PythonSDKConfig.parse_obj(res)
+    response_data = request_handling(response)
+    return PythonSDKConfig.parse_obj(response_data)
 
 
 class ServiceFoundryServiceClient:
-    # TODO (chiragjn): Rename base_url to tfy_host
-    def __init__(self, base_url: str):
-        self._base_url = base_url.strip("/")
-        tfy_servers_config = get_tfy_servers_config(self._base_url)
+    def __init__(self, tfy_host: str):
+        self._tfy_host = tfy_host.strip("/")
+        tfy_servers_config = get_tfy_servers_config(self._tfy_host)
         self._tenant_host = tfy_servers_config.tenant_host
         self._api_server_url = tfy_servers_config.servicefoundry_server_url
 
     @property
-    def base_url(self) -> str:
-        return self._base_url
+    def tfy_host(self) -> str:
+        return self._tfy_host
 
     @property
     def tenant_info(self) -> TenantInfo:
@@ -95,6 +86,17 @@ class ServiceFoundryServiceClient:
 
     @functools.cached_property
     def _min_cli_version_required(self) -> str:
-        return _cached_get_python_sdk_config(
-            self._api_server_url
-        ).truefoundry_cli_min_version
+        return self.python_sdk_config.truefoundry_cli_min_version
+
+    def check_min_cli_version(self) -> None:
+        if __version__ != "0.0.0":
+            # "0.0.0" indicates dev version
+            # noinspection PyProtectedMember
+            min_cli_version_required = self._min_cli_version_required
+            if version.parse(__version__) < version.parse(min_cli_version_required):
+                raise Exception(
+                    "You are using an outdated version of `truefoundry`.\n"
+                    f"Run `pip install truefoundry>={min_cli_version_required}` to install the supported version.",
+                )
+        else:
+            logger.debug("Ignoring minimum cli version check")