lightning-sdk 2025.7.22__py3-none-any.whl → 2025.7.31__py3-none-any.whl

lightning_sdk/lightning_cloud/openapi/models/v1_slurm_v1_status.py

@@ -42,21 +42,36 @@ class V1SlurmV1Status(object):
     """
     swagger_types = {
         'nodes': 'list[V1SlurmNode]',
+        'ssh_host': 'str',
+        'ssh_port': 'int',
+        'ssh_username': 'str',
         'version': 'str'
     }
 
     attribute_map = {
         'nodes': 'nodes',
+        'ssh_host': 'sshHost',
+        'ssh_port': 'sshPort',
+        'ssh_username': 'sshUsername',
         'version': 'version'
     }
 
-    def __init__(self, nodes: 'list[V1SlurmNode]' =None, version: 'str' =None):  # noqa: E501
+    def __init__(self, nodes: 'list[V1SlurmNode]' =None, ssh_host: 'str' =None, ssh_port: 'int' =None, ssh_username: 'str' =None, version: 'str' =None):  # noqa: E501
         """V1SlurmV1Status - a model defined in Swagger"""  # noqa: E501
         self._nodes = None
+        self._ssh_host = None
+        self._ssh_port = None
+        self._ssh_username = None
         self._version = None
         self.discriminator = None
         if nodes is not None:
             self.nodes = nodes
+        if ssh_host is not None:
+            self.ssh_host = ssh_host
+        if ssh_port is not None:
+            self.ssh_port = ssh_port
+        if ssh_username is not None:
+            self.ssh_username = ssh_username
         if version is not None:
             self.version = version
 
@@ -81,6 +96,69 @@ class V1SlurmV1Status(object):
 
         self._nodes = nodes
 
+    @property
+    def ssh_host(self) -> 'str':
+        """Gets the ssh_host of this V1SlurmV1Status.  # noqa: E501
+
+
+        :return: The ssh_host of this V1SlurmV1Status.  # noqa: E501
+        :rtype: str
+        """
+        return self._ssh_host
+
+    @ssh_host.setter
+    def ssh_host(self, ssh_host: 'str'):
+        """Sets the ssh_host of this V1SlurmV1Status.
+
+
+        :param ssh_host: The ssh_host of this V1SlurmV1Status.  # noqa: E501
+        :type: str
+        """
+
+        self._ssh_host = ssh_host
+
+    @property
+    def ssh_port(self) -> 'int':
+        """Gets the ssh_port of this V1SlurmV1Status.  # noqa: E501
+
+
+        :return: The ssh_port of this V1SlurmV1Status.  # noqa: E501
+        :rtype: int
+        """
+        return self._ssh_port
+
+    @ssh_port.setter
+    def ssh_port(self, ssh_port: 'int'):
+        """Sets the ssh_port of this V1SlurmV1Status.
+
+
+        :param ssh_port: The ssh_port of this V1SlurmV1Status.  # noqa: E501
+        :type: int
+        """
+
+        self._ssh_port = ssh_port
+
+    @property
+    def ssh_username(self) -> 'str':
+        """Gets the ssh_username of this V1SlurmV1Status.  # noqa: E501
+
+
+        :return: The ssh_username of this V1SlurmV1Status.  # noqa: E501
+        :rtype: str
+        """
+        return self._ssh_username
+
+    @ssh_username.setter
+    def ssh_username(self, ssh_username: 'str'):
+        """Sets the ssh_username of this V1SlurmV1Status.
+
+
+        :param ssh_username: The ssh_username of this V1SlurmV1Status.  # noqa: E501
+        :type: str
+        """
+
+        self._ssh_username = ssh_username
+
     @property
     def version(self) -> 'str':
         """Gets the version of this V1SlurmV1Status.  # noqa: E501

lightning_sdk/lightning_cloud/openapi/models/v1_tool_call.py

@@ -0,0 +1,175 @@
+# coding: utf-8
+
+"""
+    external/v1/auth_service.proto
+
+    No description provided (generated by Swagger Codegen https://github.com/swagger-api/swagger-codegen)  # noqa: E501
+
+    OpenAPI spec version: version not set
+    
+    Generated by: https://github.com/swagger-api/swagger-codegen.git
+
+    NOTE
+    ----
+    standard swagger-codegen-cli for this python client has been modified
+    by custom templates. The purpose of these templates is to include
+    typing information in the API and Model code. Please refer to the
+    main grid repository for more info
+"""
+
+import pprint
+import re  # noqa: F401
+
+from typing import TYPE_CHECKING
+
+import six
+
+if TYPE_CHECKING:
+    from datetime import datetime
+    from lightning_sdk.lightning_cloud.openapi.models import *
+
+class V1ToolCall(object):
+    """NOTE: This class is auto generated by the swagger code generator program.
+
+    Do not edit the class manually.
+    """
+    """
+    Attributes:
+      swagger_types (dict): The key is attribute name
+                            and the value is attribute type.
+      attribute_map (dict): The key is attribute name
+                            and the value is json key in definition.
+    """
+    swagger_types = {
+        'function': 'V1FunctionCall',
+        'id': 'str',
+        'type': 'str'
+    }
+
+    attribute_map = {
+        'function': 'function',
+        'id': 'id',
+        'type': 'type'
+    }
+
+    def __init__(self, function: 'V1FunctionCall' =None, id: 'str' =None, type: 'str' =None):  # noqa: E501
+        """V1ToolCall - a model defined in Swagger"""  # noqa: E501
+        self._function = None
+        self._id = None
+        self._type = None
+        self.discriminator = None
+        if function is not None:
+            self.function = function
+        if id is not None:
+            self.id = id
+        if type is not None:
+            self.type = type
+
+    @property
+    def function(self) -> 'V1FunctionCall':
+        """Gets the function of this V1ToolCall.  # noqa: E501
+
+
+        :return: The function of this V1ToolCall.  # noqa: E501
+        :rtype: V1FunctionCall
+        """
+        return self._function
+
+    @function.setter
+    def function(self, function: 'V1FunctionCall'):
+        """Sets the function of this V1ToolCall.
+
+
+        :param function: The function of this V1ToolCall.  # noqa: E501
+        :type: V1FunctionCall
+        """
+
+        self._function = function
+
+    @property
+    def id(self) -> 'str':
+        """Gets the id of this V1ToolCall.  # noqa: E501
+
+
+        :return: The id of this V1ToolCall.  # noqa: E501
+        :rtype: str
+        """
+        return self._id
+
+    @id.setter
+    def id(self, id: 'str'):
+        """Sets the id of this V1ToolCall.
+
+
+        :param id: The id of this V1ToolCall.  # noqa: E501
+        :type: str
+        """
+
+        self._id = id
+
+    @property
+    def type(self) -> 'str':
+        """Gets the type of this V1ToolCall.  # noqa: E501
+
+
+        :return: The type of this V1ToolCall.  # noqa: E501
+        :rtype: str
+        """
+        return self._type
+
+    @type.setter
+    def type(self, type: 'str'):
+        """Sets the type of this V1ToolCall.
+
+
+        :param type: The type of this V1ToolCall.  # noqa: E501
+        :type: str
+        """
+
+        self._type = type
+
+    def to_dict(self) -> dict:
+        """Returns the model properties as a dict"""
+        result = {}
+
+        for attr, _ in six.iteritems(self.swagger_types):
+            value = getattr(self, attr)
+            if isinstance(value, list):
+                result[attr] = list(map(
+                    lambda x: x.to_dict() if hasattr(x, "to_dict") else x,
+                    value
+                ))
+            elif hasattr(value, "to_dict"):
+                result[attr] = value.to_dict()
+            elif isinstance(value, dict):
+                result[attr] = dict(map(
+                    lambda item: (item[0], item[1].to_dict())
+                    if hasattr(item[1], "to_dict") else item,
+                    value.items()
+                ))
+            else:
+                result[attr] = value
+        if issubclass(V1ToolCall, dict):
+            for key, value in self.items():
+                result[key] = value
+
+        return result
+
+    def to_str(self) -> str:
+        """Returns the string representation of the model"""
+        return pprint.pformat(self.to_dict())
+
+    def __repr__(self) -> str:
+        """For `print` and `pprint`"""
+        return self.to_str()
+
+    def __eq__(self, other: 'V1ToolCall') -> bool:
+        """Returns true if both objects are equal"""
+        if not isinstance(other, V1ToolCall):
+            return False
+
+        return self.__dict__ == other.__dict__
+
+    def __ne__(self, other: 'V1ToolCall') -> bool:
+        """Returns true if both objects are not equal"""
+        return not self == other

lightning_sdk/lightning_cloud/openapi/models/v1_user_features.py

@@ -56,19 +56,20 @@ class V1UserFeatures(object):
         'capacity_reservation_dry_run': 'bool',
         'chat_models': 'bool',
         'cloudspace_schedules': 'bool',
-        'cloudy_vibe_code': 'bool',
         'code_tab': 'bool',
         'collab_screen_sharing': 'bool',
         'control_center_monitoring': 'bool',
         'cost_attribution_settings': 'bool',
         'custom_app_domain': 'bool',
-        'data_connection_flushing_v2': 'bool',
         'datasets': 'bool',
         'default_one_cluster': 'bool',
         'deployment_persistent_disk': 'bool',
         'drive_v2': 'bool',
         'enterprise_compute_admin': 'bool',
         'f236': 'bool',
+        'f237': 'bool',
+        'f238': 'bool',
+        'f239': 'bool',
         'fair_share': 'bool',
         'featured_studios_admin': 'bool',
         'gcp_overprovisioning': 'bool',
@@ -143,19 +144,20 @@ class V1UserFeatures(object):
         'capacity_reservation_dry_run': 'capacityReservationDryRun',
         'chat_models': 'chatModels',
         'cloudspace_schedules': 'cloudspaceSchedules',
-        'cloudy_vibe_code': 'cloudyVibeCode',
         'code_tab': 'codeTab',
         'collab_screen_sharing': 'collabScreenSharing',
         'control_center_monitoring': 'controlCenterMonitoring',
         'cost_attribution_settings': 'costAttributionSettings',
         'custom_app_domain': 'customAppDomain',
-        'data_connection_flushing_v2': 'dataConnectionFlushingV2',
         'datasets': 'datasets',
         'default_one_cluster': 'defaultOneCluster',
         'deployment_persistent_disk': 'deploymentPersistentDisk',
         'drive_v2': 'driveV2',
         'enterprise_compute_admin': 'enterpriseComputeAdmin',
         'f236': 'f236',
+        'f237': 'f237',
+        'f238': 'f238',
+        'f239': 'f239',
         'fair_share': 'fairShare',
         'featured_studios_admin': 'featuredStudiosAdmin',
         'gcp_overprovisioning': 'gcpOverprovisioning',
@@ -214,7 +216,7 @@ class V1UserFeatures(object):
         'writable_s3_connections': 'writableS3Connections'
     }
 
-    def __init__(self, academic_tier: 'bool' =None, add_data_v2: 'bool' =None, affiliate_links: 'bool' =None, agents_v2: 'bool' =None, ai_hub_monetization: 'bool' =None, auto_fast_load: 'bool' =None, auto_join_orgs: 'bool' =None, b2c_experience: 'bool' =None, byo_machine_type: 'bool' =None, cap_add: 'list[str]' =None, cap_drop: 'list[str]' =None, capacity_reservation_byoc: 'bool' =None, capacity_reservation_dry_run: 'bool' =None, chat_models: 'bool' =None, cloudspace_schedules: 'bool' =None, cloudy_vibe_code: 'bool' =None, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, control_center_monitoring: 'bool' =None, cost_attribution_settings: 'bool' =None, custom_app_domain: 'bool' =None, data_connection_flushing_v2: 'bool' =None, datasets: 'bool' =None, default_one_cluster: 'bool' =None, deployment_persistent_disk: 'bool' =None, drive_v2: 'bool' =None, enterprise_compute_admin: 'bool' =None, f236: 'bool' =None, fair_share: 'bool' =None, featured_studios_admin: 'bool' =None, gcp_overprovisioning: 'bool' =None, gcs_connections_optimized: 'bool' =None, gcs_folders: 'bool' =None, gcs_fuse: 'bool' =None, instant_capacity_reservation: 'bool' =None, job_artifacts_v2: 'bool' =None, kubernetes_cluster_ui: 'bool' =None, kubernetes_clusters: 'bool' =None, landing_studios: 'bool' =None, lit_logger: 'bool' =None, marketplace: 'bool' =None, mmt_fault_tolerance: 'bool' =None, mmt_strategy_selector: 'bool' =None, model_api_dashboard: 'bool' =None, model_api_dashboard_clickhouse: 'bool' =None, multicloud_folders: 'bool' =None, multiple_studio_versions: 'bool' =None, nerf_fs_nonpaying: 'bool' =None, onboarding_v2: 'bool' =None, org_level_member_permissions: 'bool' =None, org_usage_limits: 'bool' =None, persistent_disk: 'bool' =None, plugin_distributed: 'bool' =None, plugin_inference: 'bool' =None, plugin_label_studio: 'bool' =None, plugin_langflow: 'bool' =None, plugin_python_profiler: 'bool' =None, plugin_service: 'bool' =None, plugin_sweeps: 'bool' =None, pricing_updates: 'bool' =None, product_generator: 'bool' =None, product_license: 'bool' =None, project_selector: 'bool' =None, publish_pipelines: 'bool' =None, r2_data_connections: 'bool' =None, r2_uploads: 'bool' =None, reserved_machines_tab: 'bool' =None, restartable_jobs: 'bool' =None, runnable_public_studio_page: 'bool' =None, security_docs: 'bool' =None, show_dev_admin: 'bool' =None, single_wallet: 'bool' =None, slurm: 'bool' =None, specialised_studios: 'bool' =None, storage_overuse_deletion: 'bool' =None, studio_config: 'bool' =None, studio_sharing_v2: 'bool' =None, studio_version_visibility: 'bool' =None, trainium2: 'bool' =None, use_internal_data_connection_mounts: 'bool' =None, use_rclone_mounts_only: 'bool' =None, vultr: 'bool' =None, weka: 'bool' =None, writable_s3_connections: 'bool' =None):  # noqa: E501
+    def __init__(self, academic_tier: 'bool' =None, add_data_v2: 'bool' =None, affiliate_links: 'bool' =None, agents_v2: 'bool' =None, ai_hub_monetization: 'bool' =None, auto_fast_load: 'bool' =None, auto_join_orgs: 'bool' =None, b2c_experience: 'bool' =None, byo_machine_type: 'bool' =None, cap_add: 'list[str]' =None, cap_drop: 'list[str]' =None, capacity_reservation_byoc: 'bool' =None, capacity_reservation_dry_run: 'bool' =None, chat_models: 'bool' =None, cloudspace_schedules: 'bool' =None, code_tab: 'bool' =None, collab_screen_sharing: 'bool' =None, control_center_monitoring: 'bool' =None, cost_attribution_settings: 'bool' =None, custom_app_domain: 'bool' =None, datasets: 'bool' =None, default_one_cluster: 'bool' =None, deployment_persistent_disk: 'bool' =None, drive_v2: 'bool' =None, enterprise_compute_admin: 'bool' =None, f236: 'bool' =None, f237: 'bool' =None, f238: 'bool' =None, f239: 'bool' =None, fair_share: 'bool' =None, featured_studios_admin: 'bool' =None, gcp_overprovisioning: 'bool' =None, gcs_connections_optimized: 'bool' =None, gcs_folders: 'bool' =None, gcs_fuse: 'bool' =None, instant_capacity_reservation: 'bool' =None, job_artifacts_v2: 'bool' =None, kubernetes_cluster_ui: 'bool' =None, kubernetes_clusters: 'bool' =None, landing_studios: 'bool' =None, lit_logger: 'bool' =None, marketplace: 'bool' =None, mmt_fault_tolerance: 'bool' =None, mmt_strategy_selector: 'bool' =None, model_api_dashboard: 'bool' =None, model_api_dashboard_clickhouse: 'bool' =None, multicloud_folders: 'bool' =None, multiple_studio_versions: 'bool' =None, nerf_fs_nonpaying: 'bool' =None, onboarding_v2: 'bool' =None, org_level_member_permissions: 'bool' =None, org_usage_limits: 'bool' =None, persistent_disk: 'bool' =None, plugin_distributed: 'bool' =None, plugin_inference: 'bool' =None, plugin_label_studio: 'bool' =None, plugin_langflow: 'bool' =None, plugin_python_profiler: 'bool' =None, plugin_service: 'bool' =None, plugin_sweeps: 'bool' =None, pricing_updates: 'bool' =None, product_generator: 'bool' =None, product_license: 'bool' =None, project_selector: 'bool' =None, publish_pipelines: 'bool' =None, r2_data_connections: 'bool' =None, r2_uploads: 'bool' =None, reserved_machines_tab: 'bool' =None, restartable_jobs: 'bool' =None, runnable_public_studio_page: 'bool' =None, security_docs: 'bool' =None, show_dev_admin: 'bool' =None, single_wallet: 'bool' =None, slurm: 'bool' =None, specialised_studios: 'bool' =None, storage_overuse_deletion: 'bool' =None, studio_config: 'bool' =None, studio_sharing_v2: 'bool' =None, studio_version_visibility: 'bool' =None, trainium2: 'bool' =None, use_internal_data_connection_mounts: 'bool' =None, use_rclone_mounts_only: 'bool' =None, vultr: 'bool' =None, weka: 'bool' =None, writable_s3_connections: 'bool' =None):  # noqa: E501
         """V1UserFeatures - a model defined in Swagger"""  # noqa: E501
         self._academic_tier = None
         self._add_data_v2 = None
@@ -231,19 +233,20 @@ class V1UserFeatures(object):
         self._capacity_reservation_dry_run = None
         self._chat_models = None
         self._cloudspace_schedules = None
-        self._cloudy_vibe_code = None
         self._code_tab = None
         self._collab_screen_sharing = None
         self._control_center_monitoring = None
         self._cost_attribution_settings = None
         self._custom_app_domain = None
-        self._data_connection_flushing_v2 = None
         self._datasets = None
         self._default_one_cluster = None
         self._deployment_persistent_disk = None
         self._drive_v2 = None
         self._enterprise_compute_admin = None
         self._f236 = None
+        self._f237 = None
+        self._f238 = None
+        self._f239 = None
         self._fair_share = None
         self._featured_studios_admin = None
         self._gcp_overprovisioning = None
@@ -331,8 +334,6 @@ class V1UserFeatures(object):
             self.chat_models = chat_models
         if cloudspace_schedules is not None:
             self.cloudspace_schedules = cloudspace_schedules
-        if cloudy_vibe_code is not None:
-            self.cloudy_vibe_code = cloudy_vibe_code
         if code_tab is not None:
             self.code_tab = code_tab
         if collab_screen_sharing is not None:
@@ -343,8 +344,6 @@ class V1UserFeatures(object):
             self.cost_attribution_settings = cost_attribution_settings
         if custom_app_domain is not None:
             self.custom_app_domain = custom_app_domain
-        if data_connection_flushing_v2 is not None:
-            self.data_connection_flushing_v2 = data_connection_flushing_v2
         if datasets is not None:
             self.datasets = datasets
         if default_one_cluster is not None:
@@ -357,6 +356,12 @@ class V1UserFeatures(object):
             self.enterprise_compute_admin = enterprise_compute_admin
         if f236 is not None:
             self.f236 = f236
+        if f237 is not None:
+            self.f237 = f237
+        if f238 is not None:
+            self.f238 = f238
+        if f239 is not None:
+            self.f239 = f239
         if fair_share is not None:
             self.fair_share = fair_share
         if featured_studios_admin is not None:
@@ -785,27 +790,6 @@ class V1UserFeatures(object):
 
         self._cloudspace_schedules = cloudspace_schedules
 
-    @property
-    def cloudy_vibe_code(self) -> 'bool':
-        """Gets the cloudy_vibe_code of this V1UserFeatures.  # noqa: E501
-
-
-        :return: The cloudy_vibe_code of this V1UserFeatures.  # noqa: E501
-        :rtype: bool
-        """
-        return self._cloudy_vibe_code
-
-    @cloudy_vibe_code.setter
-    def cloudy_vibe_code(self, cloudy_vibe_code: 'bool'):
-        """Sets the cloudy_vibe_code of this V1UserFeatures.
-
-
-        :param cloudy_vibe_code: The cloudy_vibe_code of this V1UserFeatures.  # noqa: E501
-        :type: bool
-        """
-
-        self._cloudy_vibe_code = cloudy_vibe_code
-
     @property
     def code_tab(self) -> 'bool':
         """Gets the code_tab of this V1UserFeatures.  # noqa: E501
@@ -911,27 +895,6 @@ class V1UserFeatures(object):
 
         self._custom_app_domain = custom_app_domain
 
-    @property
-    def data_connection_flushing_v2(self) -> 'bool':
-        """Gets the data_connection_flushing_v2 of this V1UserFeatures.  # noqa: E501
-
-
-        :return: The data_connection_flushing_v2 of this V1UserFeatures.  # noqa: E501
-        :rtype: bool
-        """
-        return self._data_connection_flushing_v2
-
-    @data_connection_flushing_v2.setter
-    def data_connection_flushing_v2(self, data_connection_flushing_v2: 'bool'):
-        """Sets the data_connection_flushing_v2 of this V1UserFeatures.
-
-
-        :param data_connection_flushing_v2: The data_connection_flushing_v2 of this V1UserFeatures.  # noqa: E501
-        :type: bool
-        """
-
-        self._data_connection_flushing_v2 = data_connection_flushing_v2
-
     @property
     def datasets(self) -> 'bool':
         """Gets the datasets of this V1UserFeatures.  # noqa: E501
@@ -1058,6 +1021,69 @@ class V1UserFeatures(object):
 
         self._f236 = f236
 
+    @property
+    def f237(self) -> 'bool':
+        """Gets the f237 of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The f237 of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._f237
+
+    @f237.setter
+    def f237(self, f237: 'bool'):
+        """Sets the f237 of this V1UserFeatures.
+
+
+        :param f237: The f237 of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._f237 = f237
+
+    @property
+    def f238(self) -> 'bool':
+        """Gets the f238 of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The f238 of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._f238
+
+    @f238.setter
+    def f238(self, f238: 'bool'):
+        """Sets the f238 of this V1UserFeatures.
+
+
+        :param f238: The f238 of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._f238 = f238
+
+    @property
+    def f239(self) -> 'bool':
+        """Gets the f239 of this V1UserFeatures.  # noqa: E501
+
+
+        :return: The f239 of this V1UserFeatures.  # noqa: E501
+        :rtype: bool
+        """
+        return self._f239
+
+    @f239.setter
+    def f239(self, f239: 'bool'):
+        """Sets the f239 of this V1UserFeatures.
+
+
+        :param f239: The f239 of this V1UserFeatures.  # noqa: E501
+        :type: bool
+        """
+
+        self._f239 = f239
+
     @property
     def fair_share(self) -> 'bool':
         """Gets the fair_share of this V1UserFeatures.  # noqa: E501

lightning_sdk/lightning_cloud/openapi/models/v1_volume_state.py

@@ -40,6 +40,7 @@ class V1VolumeState(object):
     UNSPECIFIED = "VOLUME_STATE_UNSPECIFIED"
     ATTACHED = "VOLUME_STATE_ATTACHED"
     DETACHED = "VOLUME_STATE_DETACHED"
+    OVERPROVISIONED = "VOLUME_STATE_OVERPROVISIONED"
     """
     Attributes:
       swagger_types (dict): The key is attribute name

lightning_sdk/llm/llm.py

@@ -2,6 +2,7 @@ import json
 import os
 from typing import Any, AsyncGenerator, ClassVar, Dict, Generator, List, Optional, Tuple, Union
 
+from lightning_sdk.api import TeamspaceApi, UserApi
 from lightning_sdk.api.llm_api import LLMApi
 from lightning_sdk.lightning_cloud.openapi.models.v1_conversation_response_chunk import V1ConversationResponseChunk
 
@@ -13,7 +14,7 @@ PUBLIC_MODEL_PROVIDERS: Dict[str, str] = {
 }
 
 
-def _load_public_assistants() -> Dict[str, str]:
+def _load_public_assistants() -> Dict[str, Dict[str, Any]]:
     """Load public assistants from a JSON file."""
     try:
         json_path = os.path.join(os.path.dirname(__file__), "public_assistants.json")
@@ -28,7 +29,7 @@ class LLM:
     _auth_info_cached: ClassVar[bool] = False
     _cached_auth_info: ClassVar[Dict[str, Optional[str]]] = {}
     _llm_api_cache: ClassVar[Dict[Optional[str], LLMApi]] = {}
-    _public_assistants: ClassVar[Optional[Dict[str, str]]] = None
+    _public_assistants: ClassVar[Optional[Dict[str, Dict[str, Any]]]] = None
 
     def __new__(cls, name: str, teamspace: Optional[str] = None, enable_async: Optional[bool] = False) -> "LLM":
         return super().__new__(cls)
@@ -76,6 +77,7 @@ class LLM:
             LLM._llm_api_cache[teamspace] = LLMApi()
         self._llm_api = LLM._llm_api_cache[teamspace]
 
+        self._context_length = None
         self._model_id = self._get_model_id()
         self._conversations = {}
 
@@ -87,16 +89,45 @@ class LLM:
     def provider(self) -> str:
         return self._model_provider
 
+    def context_length(self, model: Optional[str] = None) -> Optional[int]:
+        if model is None:
+            return self._context_length
+
+        context_info = self._public_assistants.get(model)
+        if context_info is None or "context_length" not in context_info:
+            raise ValueError(f"Cannot access context length of model '{model}'.")
+
+        return int(context_info["context_length"])
+
     def _get_auth_info(self, teamspace_name: Optional[str] = None) -> None:
         # TODO: Validate user input teamspace name
         if not LLM._auth_info_cached:
             if teamspace_name is None:
+                # studio users
                 teamspace_name = os.environ.get("LIGHTNING_TEAMSPACE", None)
+
             if teamspace_name is None:
-                raise ValueError(
-                    "Teamspace name must be provided either through "
-                    "the environment variable LIGHTNING_TEAMSPACE or as an argument - LLM(..., teamspace=...)"
-                )
+                # local users with no given teamspace
+                try:
+                    teamspace_api = TeamspaceApi()
+                    user_api = UserApi()
+                    authed_user = user_api._client.auth_service_get_user()
+                    default_teamspace = teamspace_api.list_teamspaces(owner_id=authed_user.id)[0]
+                    teamspace_name = default_teamspace.name
+                    teamspace_id = default_teamspace.id
+                    os.environ["LIGHTNING_CLOUD_PROJECT_ID"] = teamspace_id
+                    os.environ["LIGHTNING_TEAMSPACE"] = teamspace_name
+                except Exception as err:
+                    # throw an appropriate error that guides users to login through the platform
+                    raise ValueError(
+                        "Teamspace information is missing. "
+                        "If this is your first time using LitAI, please log in at https://lightning.ai/sign-up "
+                        "and re-run your script, or set the environment variable LIGHTNING_TEAMSPACE=<your-teamspace>."
+                    ) from err
+
+            # TODO: when teamspace_name is given, we don't know the teamspace_id yet
+            # TODO: if LIGHTNING_CLOUD_PROJECT_ID does not exist, we have to get the id from the teamspace name
+
             LLM._cached_auth_info = {
                 "teamspace_name": teamspace_name,
                 "teamspace_id": os.environ.get("LIGHTNING_CLOUD_PROJECT_ID", None),
@@ -137,7 +168,10 @@ class LLM:
                 and LLM._public_assistants
                 and f"{self._model_provider}/{self._model_name}" in LLM._public_assistants
             ):
-                return LLM._public_assistants[f"{self._model_provider}/{self._model_name}"]
+                self._context_length = int(
+                    LLM._public_assistants[f"{self._model_provider}/{self._model_name}"]["context_length"]
+                )
+                return LLM._public_assistants[f"{self._model_provider}/{self._model_name}"]["id"]
             try:
                 return self._llm_api.get_assistant(
                     model_provider=PUBLIC_MODEL_PROVIDERS[self._model_provider],

lightning_sdk/llm/public_assistants.json

@@ -1,10 +1,34 @@
 {
-    "openai/gpt-4o": "ast_01jdjds71fs8gt47jexzed4czs",
-    "openai/gpt-4": "ast_01jd38ze6tjbrcd4942nhz41zn",
-    "openai/o3-mini": "ast_01jz3t13fhnjhh11t1k8b5gyp1",
-    "anthropic/claude-3-5-sonnet-20240620": "ast_01jd3923a6p98rqwh3dpj686pq",
-    "google/gemini-2.5-pro": "ast_01jz3tdb1fhey798k95pv61v57",
-    "google/gemini-2.5-flash": "ast_01jz3thxskg4fcdk4xhkjkym5a",
-    "google/gemini-2.5-flash-lite-preview-06-17": "ast_01jz3thxskg4fcdk4xhkjkym5b",
-    "lightning-ai/llama4-maverick": "ast_01k073vsqs66tenpns02cf5jnq"
+    "openai/gpt-4o": {
+        "id": "ast_01jdjds71fs8gt47jexzed4czs",
+        "context_length": 128000
+    },
+    "openai/gpt-4": {
+        "id": "ast_01jd38ze6tjbrcd4942nhz41zn",
+        "context_length": 8192
+    },
+    "openai/o3-mini": {
+        "id": "ast_01jz3t13fhnjhh11t1k8b5gyp1",
+        "context_length": 128000
+    },
+    "anthropic/claude-3-5-sonnet-20240620": {
+        "id": "ast_01jd3923a6p98rqwh3dpj686pq",
+        "context_length": 200000
+    },
+    "google/gemini-2.5-pro": {
+        "id": "ast_01jz3tdb1fhey798k95pv61v57",
+        "context_length": 1048576
+    },
+    "google/gemini-2.5-flash": {
+        "id": "ast_01jz3thxskg4fcdk4xhkjkym5a",
+        "context_length": 8000
+    },
+    "google/gemini-2.5-flash-lite-preview-06-17": {
+        "id": "ast_01jz3thxskg4fcdk4xhkjkym5b",
+        "context_length": 8000
+    },
+    "lightning-ai/llama4-maverick": {
+        "id": "ast_01k0wgg56tm8mv9n12aq2mnxas",
+        "context_length": 100000
+    }
 }