lightning-sdk 2025.8.19.post0__py3-none-any.whl → 2025.8.26__py3-none-any.whl

lightning_sdk/__init__.py

@@ -32,6 +32,6 @@ __all__ = [
     "User",
 ]
 
-__version__ = "2025.08.19.post0"
+__version__ = "2025.08.26"
 _check_version_and_prompt_upgrade(__version__)
 _set_tqdm_envvars_noninteractive()

lightning_sdk/api/llm_api.py

@@ -146,7 +146,6 @@ class LLMApi:
                     {"contentType": "text", "parts": [prompt]},
                 ],
             },
-            "max_tokens": max_completion_tokens,
             "conversation_id": conversation_id,
             "billing_project_id": billing_project_id,
             "name": name,
@@ -159,6 +158,9 @@ class LLMApi:
             "parent_message_id": kwargs.get("parent_message_id", ""),
             "tools": tools,
         }
+        if max_completion_tokens is not None:
+            body["max_completion_tokens"] = max_completion_tokens
+
         if images:
             for image in images:
                 url = image
@@ -203,7 +205,6 @@ class LLMApi:
                     {"contentType": "text", "parts": [prompt]},
                 ],
             },
-            "max_completion_tokens": max_completion_tokens,
             "conversation_id": conversation_id,
             "billing_project_id": billing_project_id,
             "name": name,
@@ -216,6 +217,9 @@ class LLMApi:
             "parent_message_id": kwargs.get("parent_message_id", ""),
             "sent_at": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="microseconds"),
         }
+        if max_completion_tokens is not None:
+            body["max_completion_tokens"] = max_completion_tokens
+
         if images:
             for image in images:
                 url = image

lightning_sdk/api/studio_api.py

@@ -11,7 +11,7 @@ from tqdm import tqdm
 
 from lightning_sdk.api.utils import (
     _create_app,
-    _download_studio_files,
+    _download_teamspace_files,
     _DummyBody,
     _DummyResponse,
     _FileUploader,
@@ -25,6 +25,7 @@ from lightning_sdk.constants import _LIGHTNING_DEBUG
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
     CloudspaceIdRunsBody,
+    CloudspacesIdBody,
     Externalv1LightningappInstance,
     IdCodeconfigBody,
     IdExecuteBody1,
@@ -40,6 +41,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1CloudSpaceState,
     V1ClusterAccelerator,
     V1EndpointType,
+    V1EnvVar,
     V1GetCloudSpaceInstanceStatusResponse,
     V1GetLongRunningCommandInCloudSpaceResponse,
     V1LoginRequest,
@@ -205,6 +207,32 @@ class StudioApi:
             instance_id = code_status.in_use.cloud_space_instance_id
             print(f"Studio started | {teamspace_id=} {studio_id=} {instance_id=}")
 
+    def start_studio_async(
+        self,
+        studio_id: str,
+        teamspace_id: str,
+        machine: Union[Machine, str],
+        interruptible: bool = False,
+        max_runtime: Optional[int] = None,
+    ) -> None:
+        """Start an existing Studio without blocking."""
+        # need to go via kwargs for typing compatibility since autogenerated apis accept None but aren't typed with None
+        optional_kwargs_compute_body = {}
+
+        if max_runtime is not None:
+            optional_kwargs_compute_body["requested_run_duration_seconds"] = str(max_runtime)
+        self._client.cloud_space_service_start_cloud_space_instance(
+            IdStartBody(
+                compute_config=V1UserRequestedComputeConfig(
+                    name=_machine_to_compute_name(machine),
+                    spot=interruptible,
+                    **optional_kwargs_compute_body,
+                )
+            ),
+            teamspace_id,
+            studio_id,
+        )
+
     def stop_studio(self, studio_id: str, teamspace_id: str) -> None:
         """Stop an existing Studio."""
         self.stop_keeping_alive(teamspace_id=teamspace_id, studio_id=studio_id)
@@ -289,6 +317,79 @@ class StudioApi:
                 break
             time.sleep(1)
 
+    def switch_studio_machine_with_progress(
+        self,
+        studio_id: str,
+        teamspace_id: str,
+        machine: Union[Machine, str],
+        interruptible: bool,
+        progress: Any,  # StudioProgressTracker - avoid circular import
+    ) -> None:
+        """Switches given Studio to a new machine type with progress tracking."""
+        progress.update_progress(10, "Requesting machine switch...")
+
+        self._request_switch(
+            studio_id=studio_id, teamspace_id=teamspace_id, machine=machine, interruptible=interruptible
+        )
+
+        progress.update_progress(20, "Waiting for machine allocation...")
+
+        # Wait until it's time to switch
+        requested_was_found = False
+        startup_status = None
+        base_progress = 20
+        max_wait_progress = 60
+        wait_counter = 0
+
+        while True:
+            status = self.get_studio_status(studio_id, teamspace_id)
+            requested_machine = status.requested
+
+            if requested_machine is not None:
+                requested_was_found = True
+                startup_status = requested_machine.startup_status
+
+            # if the requested machine was found in the past, use the in_use status instead.
+            # it might be that it either was cancelled or it actually is ready.
+            # Either way, since we're actually blocking below for the in use startup status
+            # it's safe to switch at this point
+            elif requested_was_found:
+                in_use_machine = status.in_use
+                if in_use_machine is not None:
+                    startup_status = in_use_machine.startup_status
+
+            if startup_status and startup_status.initial_restore_finished:
+                break
+
+            # Update progress gradually while waiting
+            wait_counter += 1
+            current_progress = min(base_progress + (wait_counter * 2), max_wait_progress)
+            progress.update_progress(current_progress, "Allocating new machine...")
+            time.sleep(1)
+
+        progress.update_progress(70, "Starting machine switch...")
+        self._client.cloud_space_service_switch_cloud_space_instance(teamspace_id, studio_id)
+
+        progress.update_progress(80, "Configuring new machine...")
+
+        # Wait until the new machine is ready to use
+        switch_counter = 0
+        while True:
+            in_use = self.get_studio_status(studio_id, teamspace_id).in_use
+            if in_use is None:
+                continue
+            startup_status = in_use.startup_status
+            if startup_status and startup_status.top_up_restore_finished:
+                break
+
+            # Update progress while waiting for machine to be ready
+            switch_counter += 1
+            current_progress = min(80 + switch_counter, 95)
+            progress.update_progress(current_progress, "Finalizing machine setup...")
+            time.sleep(1)
+
+        progress.complete("Machine switch completed successfully")
+
     def get_machine(self, studio_id: str, teamspace_id: str, cloud_account_id: str, org_id: str) -> Machine:
         """Get the current machine type the given Studio is running on."""
         response: V1CloudSpaceInstanceConfig = self._client.cloud_space_service_get_cloud_space_instance_config(
@@ -578,7 +679,7 @@ class StudioApi:
         if prefix.endswith("/") is False:
             prefix = prefix + "/"
 
-        _download_studio_files(
+        _download_teamspace_files(
             client=self._client,
             teamspace_id=teamspace_id,
             cluster_id=cloud_account,
@@ -841,3 +942,68 @@ class StudioApi:
             plugin_type=plugin_type,
             **other_arguments,
         )
+
+    def _update_cloudspace(self, studio: V1CloudSpace, teamspace_id: str, key: str, value: Any) -> None:
+        body = CloudspacesIdBody(
+            code_url=studio.code_url,
+            data_connection_mounts=studio.data_connection_mounts,
+            description=studio.description,
+            display_name=studio.display_name,
+            env=studio.env,
+            featured=studio.featured,
+            hide_files=studio.hide_files,
+            is_cloudspace_private=studio.is_cloudspace_private,
+            is_code_private=studio.is_code_private,
+            is_favorite=studio.is_favorite,
+            is_published=studio.is_published,
+            license=studio.license,
+            license_url=studio.license_url,
+            message=studio.message,
+            multi_user_edit=studio.multi_user_edit,
+            operating_cost=studio.operating_cost,
+            paper_authors=studio.paper_authors,
+            paper_org=studio.paper_org,
+            paper_org_avatar_url=studio.paper_org_avatar_url,
+            paper_url=studio.paper_url,
+            switch_to_default_machine_on_idle=studio.switch_to_default_machine_on_idle,
+            tags=studio.tags,
+            thumbnail_file_type=studio.thumbnail_file_type,
+            user_metadata=studio.user_metadata,
+        )
+
+        setattr(body, key, value)
+
+        self._client.cloud_space_service_update_cloud_space(
+            id=studio.id,
+            project_id=teamspace_id,
+            body=body,
+        )
+
+    def set_env(
+        self,
+        studio: V1CloudSpace,
+        teamspace_id: str,
+        new_env: Dict[str, str],
+        partial: bool = True,
+    ) -> None:
+        """Set the environment variables for the Studio.
+
+        Args:
+            new_env: The new environment variables to set.
+            partial: Whether to only set the environment variables that are provided.
+                If False, existing environment variables that are not in new_env will be removed.
+                If True, existing environment variables that are not in new_env will be kept.
+        """
+        updated_env_dict = {}
+        if partial:
+            updated_env_dict = {env.name: env.value for env in studio.env}
+            updated_env_dict.update(new_env)
+        else:
+            updated_env_dict = new_env
+
+        updated_env = [V1EnvVar(name=key, value=value) for key, value in updated_env_dict.items()]
+
+        self._update_cloudspace(studio, teamspace_id, "env", updated_env)
+
+    def get_env(self, studio: V1CloudSpace) -> Dict[str, str]:
+        return {env.name: env.value for env in studio.env}

lightning_sdk/api/teamspace_api.py

@@ -1,6 +1,5 @@
 import os
-import tempfile
-import zipfile
+import re
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 
@@ -9,6 +8,7 @@ from tqdm.auto import tqdm
 
 from lightning_sdk.api.utils import (
     _download_model_files,
+    _download_teamspace_files,
     _DummyBody,
     _FileUploader,
     _get_model_version,
@@ -22,6 +22,8 @@ from lightning_sdk.lightning_cloud.openapi import (
     ModelsStoreApi,
     ProjectIdAgentsBody,
     ProjectIdModelsBody,
+    ProjectIdSecretsBody,
+    SecretsIdBody,
     V1Assistant,
     V1CloudSpace,
     V1ClusterAccelerator,
@@ -34,6 +36,8 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1Project,
     V1ProjectClusterBinding,
     V1PromptSuggestion,
+    V1Secret,
+    V1SecretType,
     V1UpstreamOpenAI,
 )
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
@@ -421,39 +425,65 @@ class TeamspaceApi:
         # TODO: Update this endpoint to permit basic auth
         auth = Auth()
         auth.authenticate()
-        token = self._client.auth_service_login(V1LoginRequest(auth.api_key)).token
 
-        query_params = {
-            "clusterId": cloud_account,
-            "prefix": _resolve_teamspace_remote_path(path),
-            "token": token,
-        }
+        prefix = _resolve_teamspace_remote_path(path)
 
-        r = requests.get(
-            f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/download",
-            params=query_params,
-            stream=True,
+        # ensure we only download as a directory and not the entire prefix
+        if prefix.endswith("/") is False:
+            prefix = prefix + "/"
+
+        _download_teamspace_files(
+            client=self._client,
+            teamspace_id=teamspace_id,
+            cluster_id=cloud_account,
+            prefix=prefix,
+            download_dir=Path(target_path),
+            progress_bar=progress_bar,
         )
 
-        if progress_bar:
-            pbar = tqdm(
-                desc=f"Downloading {os.path.split(path)[1]}",
-                unit="B",
-                unit_scale=True,
-                unit_divisor=1000,
-            )
+    def get_secrets(self, teamspace_id: str) -> Dict[str, str]:
+        """Get all secrets for a teamspace."""
+        secrets = self._get_secrets(teamspace_id)
+        # this returns encrypted values for security. It doesn't make sense to show them,
+        # so we just return a placeholder
+        # not a security issue to replace in the client as we get the encrypted values from the server.
+        return {secret.name: "***REDACTED***" for secret in secrets if secret.type == V1SecretType.UNSPECIFIED}
 
-            pbar_update = pbar.update
-        else:
-            pbar_update = lambda x: None
+    def set_secret(self, teamspace_id: str, key: str, value: str) -> None:
+        """Set a secret for a teamspace.
 
-        if target_path:
-            os.makedirs(target_path, exist_ok=True)
+        This will replace the existing secret if it exists and create a new one if it doesn't.
+        """
+        secrets = self._get_secrets(teamspace_id)
+        for secret in secrets:
+            if secret.name == key:
+                return self._update_secret(teamspace_id, secret.id, value)
+        return self._create_secret(teamspace_id, key, value)
+
+    def _get_secrets(self, teamspace_id: str) -> List[V1Secret]:
+        return self._client.secret_service_list_secrets(project_id=teamspace_id).secrets
+
+    def _update_secret(self, teamspace_id: str, secret_id: str, value: str) -> None:
+        self._client.secret_service_update_secret(
+            body=SecretsIdBody(value=value),
+            project_id=teamspace_id,
+            id=secret_id,
+        )
 
-        with tempfile.TemporaryFile() as f:
-            for chunk in r.iter_content(chunk_size=4096 * 8):
-                f.write(chunk)
-                pbar_update(len(chunk))
+    def _create_secret(
+        self,
+        teamspace_id: str,
+        key: str,
+        value: str,
+    ) -> None:
+        self._client.secret_service_create_secret(
+            body=ProjectIdSecretsBody(name=key, value=value, type=V1SecretType.UNSPECIFIED), project_id=teamspace_id
+        )
 
-            with zipfile.ZipFile(f) as z:
-                z.extractall(target_path)
+    def verify_secret_name(self, name: str) -> bool:
+        """Verify if a secret name is valid.
+
+        A valid secret name starts with a letter or underscore, followed by letters, digits, or underscores.
+        """
+        pattern = r"^[A-Za-z_][A-Za-z0-9_]*$"
+        return re.match(pattern, name) is not None

lightning_sdk/api/user_api.py

@@ -1,15 +1,20 @@
-from typing import List, Union
+import re
+from typing import Dict, List, Union
 
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
+    SecretsIdBody1,
     V1CloudSpace,
+    V1CreateUserSecretRequest,
     V1GetUserResponse,
     V1ListCloudSpacesResponse,
     V1Membership,
     V1Organization,
     V1SearchUser,
+    V1Secret,
     V1UserFeatures,
 )
+from lightning_sdk.lightning_cloud.openapi.models.v1_secret_type import V1SecretType
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 
 
@@ -69,3 +74,46 @@ class UserApi:
     def _get_feature_flags(self) -> V1UserFeatures:
         resp: V1GetUserResponse = self._client.auth_service_get_user()
         return resp.features
+
+    def get_secrets(self) -> Dict[str, str]:
+        """Get all secrets for the current user."""
+        secrets = self._get_secrets()
+        # this returns encrypted values for security. It doesn't make sense to show them,
+        # so we just return a placeholder
+        # not a security issue to replace in the client as we get the encrypted values from the server.
+        return {secret.name: "***REDACTED***" for secret in secrets if secret.type == V1SecretType.UNSPECIFIED}
+
+    def set_secret(self, key: str, value: str) -> None:
+        """Set a secret for the current user.
+
+        This will replace the existing secret if it exists and create a new one if it doesn't.
+        """
+        secrets = self._get_secrets()
+        for secret in secrets:
+            if secret.name == key:
+                return self._update_secret(secret.id, value)
+        return self._create_secret(key, value)
+
+    def _get_secrets(self) -> List[V1Secret]:
+        return self._client.secret_service_list_user_secrets().secrets
+
+    def _update_secret(self, secret_id: str, value: str) -> None:
+        self._client.secret_service_update_user_secret(
+            body=SecretsIdBody1(value=value),
+            id=secret_id,
+        )
+
+    def _create_secret(
+        self,
+        key: str,
+        value: str,
+    ) -> None:
+        self._client.secret_service_create_user_secret(body=V1CreateUserSecretRequest(name=key, value=value))
+
+    def verify_secret_name(self, name: str) -> bool:
+        """Verify if a secret name is valid.
+
+        A valid secret name starts with a letter or underscore, followed by letters, digits, or underscores.
+        """
+        pattern = r"^[A-Za-z_][A-Za-z0-9_]*$"
+        return re.match(pattern, name) is not None

lightning_sdk/api/utils.py

@@ -586,7 +586,7 @@ def _download_model_files(
         return response.filepaths
 
 
-def _download_studio_files(
+def _download_teamspace_files(
     client: LightningClient,
     teamspace_id: str,
     cluster_id: str,

lightning_sdk/cli/config/set.py

@@ -1,8 +1,8 @@
 import click
 
-from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
+from lightning_sdk.cli.utils.save_to_config import save_teamspace_to_config
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.machine import CloudProvider
-from lightning_sdk.organization import Organization
 from lightning_sdk.studio import Studio
 from lightning_sdk.utils.config import Config, DefaultConfigKeys
 from lightning_sdk.utils.resolve import _resolve_org, _resolve_user
@@ -59,23 +59,11 @@ def set_studio(studio_name: str) -> None:
 @click.argument("teamspace_name")
 def set_teamspace(teamspace_name: str) -> None:
     """Set the default teamspace name in the config."""
-    config = Config()
-
-    teamspace_resolved = resolve_teamspace_owner_name_format(teamspace_name)
+    menu = TeamspacesMenu()
+    teamspace_resolved = menu(teamspace=teamspace_name)
 
-    if teamspace_resolved is None:
-        # TODO: make this a generic CLI error
-        raise ValueError(
-            f"Could not resolve teamspace: '{teamspace_name}'. "
-            "Teamspace should be specified as 'owner/name'. Does the teamspace exist?"
-        )
-
-    setattr(config, DefaultConfigKeys.teamspace_name, teamspace_resolved.name)
-    setattr(config, DefaultConfigKeys.teamspace_owner, teamspace_resolved.owner.name)
-    if isinstance(teamspace_resolved.owner, Organization):
-        setattr(config, DefaultConfigKeys.teamspace_owner_type, "organization")
-    else:
-        setattr(config, DefaultConfigKeys.teamspace_owner_type, "user")
+    # explicit user action, so overwrite the config
+    save_teamspace_to_config(teamspace_resolved, overwrite=True)
 
 
 @set_value.command("cloud-account")

lightning_sdk/cli/legacy/create.py

@@ -7,10 +7,8 @@ import click
 from rich.console import Console
 
 from lightning_sdk import Machine, Studio
-from lightning_sdk.api.cloud_account_api import CloudAccountApi
-from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.machine import CloudProvider
-from lightning_sdk.utils.resolve import _resolve_deprecated_provider
 
 _MACHINE_VALUES = tuple(
     [machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine) and machine._include_in_cli]
@@ -80,16 +78,8 @@ def studio(
 
     NAME: the name of the studio to create. If already present within teamspace, will add a random suffix.
     """
-    menu = _TeamspacesMenu()
-    teamspace_resolved = menu._resolve_teamspace(teamspace)
-
-    cloud_provider = str(_resolve_deprecated_provider(cloud_provider, provider))
-
-    if cloud_provider is not None:
-        cloud_account_api = CloudAccountApi()
-        cloud_account = cloud_account_api.resolve_cloud_account(
-            teamspace_resolved.id, cloud_account, cloud_provider, teamspace_resolved.default_cloud_account
-        )
+    menu = TeamspacesMenu()
+    teamspace_resolved = menu(teamspace)
 
     # default cloud account to current studios cloud account if run from studio
     # else it will fall back to teamspace default in the backend
@@ -107,11 +97,19 @@ def studio(
         console.print(f"Studio with name {name} already exists. Using {new_name} instead.")
         name = new_name
 
-    studio = Studio(name=name, teamspace=teamspace_resolved, cloud_account=cloud_account, create_ok=True)
+    studio = Studio(
+        name=name,
+        teamspace=teamspace_resolved,
+        cloud_account=cloud_account,
+        create_ok=True,
+        cloud_provider=cloud_provider,
+        provider=provider,
+    )
 
     console.print(f"Created Studio {studio.name}.")
 
     if start is not None:
         start_machine = getattr(Machine, start, start)
+        Studio.show_progress = True
         studio.start(start_machine)
         console.print(f"Started Studio {studio.name} on machine {start}")

lightning_sdk/cli/legacy/delete.py

@@ -5,7 +5,7 @@ from rich.console import Console
 
 from lightning_sdk.cli.legacy.exceptions import StudioCliError
 from lightning_sdk.cli.legacy.job_and_mmt_action import _JobAndMMTAction
-from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.lit_container import LitContainer
 from lightning_sdk.studio import Studio
@@ -30,8 +30,8 @@ def delete() -> None:
 def container(name: str, teamspace: Optional[str] = None) -> None:
     """Delete the docker container NAME."""
     api = LitContainer()
-    menu = _TeamspacesMenu()
-    resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace=teamspace)
     try:
         api.delete_container(name, resolved_teamspace.name, resolved_teamspace.owner.name)
         Console().print(f"Container {name} deleted successfully.")

lightning_sdk/cli/legacy/deploy/_auth.py

@@ -10,7 +10,7 @@ from rich.prompt import Confirm
 
 from lightning_sdk import Teamspace
 from lightning_sdk.api import UserApi
-from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.lightning_cloud import env
 from lightning_sdk.lightning_cloud.login import Auth, AuthServer
 from lightning_sdk.lightning_cloud.openapi import V1CloudSpace
@@ -74,13 +74,13 @@ def authenticate(mode: _AuthMode, shall_confirm: bool = True) -> None:
 def select_teamspace(teamspace: Optional[str], org: Optional[str], user: Optional[str]) -> Teamspace:
     if teamspace is None:
         user = _get_authed_user()
-        menu = _TeamspacesMenu()
+        menu = TeamspacesMenu()
         possible_teamspaces = menu._get_possible_teamspaces(user)
         if len(possible_teamspaces) == 1:
             name = next(iter(possible_teamspaces.values()))["name"]
             return Teamspace(name=name, org=org, user=user)
 
-        return menu._resolve_teamspace(teamspace)
+        return menu(teamspace)
 
     return _resolve_teamspace(teamspace=teamspace, org=org, user=user)
 
@@ -180,7 +180,7 @@ class _Onboarding:
             return select_teamspace(teamspace, org, user)
 
         # Run only when user hasn't completed onboarding yet.
-        menu = _TeamspacesMenu()
+        menu = TeamspacesMenu()
         self._wait_user_onboarding()
         # Onboarding has been completed - user already selected organization if they could
         possible_teamspaces = menu._get_possible_teamspaces(self.user)

lightning_sdk/cli/legacy/download.py

@@ -11,7 +11,7 @@ from lightning_sdk.api.license_api import LicenseApi
 from lightning_sdk.api.lit_container_api import LitContainerApi
 from lightning_sdk.cli.legacy.exceptions import StudioCliError
 from lightning_sdk.cli.legacy.studios_menu import _StudiosMenu
-from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.models import download_model
 from lightning_sdk.studio import Studio
 from lightning_sdk.utils.resolve import _get_authed_user, skip_studio_init
@@ -110,8 +110,8 @@ def folder(
         path = _expand_remote_path(path)
         resolved_downloader = _resolve_studio(studio)
     elif teamspace:
-        menu = _TeamspacesMenu()
-        resolved_downloader = menu._resolve_teamspace(teamspace)
+        menu = TeamspacesMenu()
+        resolved_downloader = menu(teamspace)
     else:
         raise ValueError("Either --studio or --teamspace must be provided")
 
@@ -173,8 +173,8 @@ def file(path: str = "", studio: Optional[str] = None, teamspace: Optional[str]
     if studio:
         resolved_downloader = _resolve_studio(studio)
     elif teamspace:
-        menu = _TeamspacesMenu()
-        resolved_downloader = menu._resolve_teamspace(teamspace)
+        menu = TeamspacesMenu()
+        resolved_downloader = menu(teamspace)
     else:
         raise ValueError("Either --studio or --teamspace must be provided")
 
@@ -214,8 +214,8 @@ def download_container(
     CONTAINER: The name of the container to download.
     """
     console = Console()
-    menu = _TeamspacesMenu()
-    resolved_teamspace = menu._resolve_teamspace(teamspace)
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace)
     with console.status("Downloading container..."):
         api = LitContainerApi()
         api.download_container(container, resolved_teamspace, tag, cloud_account)