lightning-sdk 2025.8.6rc0__py3-none-any.whl → 2025.8.6rc2__py3-none-any.whl

lightning_sdk/lightning_cloud/utils/data_connection.py

@@ -1,7 +1,7 @@
 import os
 from time import sleep, time
 from lightning_sdk.lightning_cloud import rest_client
-from lightning_sdk.lightning_cloud.openapi import Create, V1AwsDataConnection, V1S3FolderDataConnection, V1EfsConfig, V1GcpDataConnection
+from lightning_sdk.lightning_cloud.openapi import Create, V1AwsDataConnection, V1S3FolderDataConnection, V1EfsConfig, V1GcpDataConnection, V1FilestoreDataConnection, V1DataConnectionTier
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 import urllib3
 
@@ -272,6 +272,61 @@ def add_gcs_connection(connection_name: str, bucket_name: str, create_timeout: i
 
     return
 
+def create_filestore_folder(folder_name: str, region: str, capacity_gb: int = 1024, tier: V1DataConnectionTier = V1DataConnectionTier.HDD, create_timeout: int = 600) -> None:
+    """
+    Utility function to create a Filestore folder.
+
+    Args:
+        folder_name: The name of the folder to create.
+        region: the region to create the folder in, e.g. "us-central1".
+        capacity_gb: capacity of the Filestore folder.
+        tier: Filestore folder tier, SSD or HDD.
+        create_timeout (int): The timeout for the folder creation.
+    """
+    client = rest_client.LightningClient(retry=False)
+
+    project_id = os.getenv("LIGHTNING_CLOUD_PROJECT_ID")
+    cluster_id = os.getenv("LIGHTNING_CLUSTER_ID")
+
+    # Get existing data connections
+    data_connections = client.data_connection_service_list_data_connections(project_id).data_connections
+
+    for connection in data_connections:
+        existing_folder_name = getattr(connection, 'name', None)
+        isFilestore = getattr(connection, 'filestore', None) is not None
+
+        if existing_folder_name == folder_name and isFilestore:
+            return
+
+    # If we get here, no matching folder was found, proceed with creation
+    body = Create(
+        name=folder_name,
+        create_resources=True,
+        cluster_id=cluster_id,
+        access_cluster_ids=[cluster_id],
+        filestore=V1FilestoreDataConnection(region=region, capacity_gb=capacity_gb, tier=tier, source=folder_name),
+    )
+    try:
+        connection = client.data_connection_service_create_data_connection(body, project_id)
+    except ApiException as e:
+        # Note: This function can be called in a distributed way.
+        # There is a race condition where one machine might create the entry before another machine
+        # and this request would fail with duplicated key
+        # In this case, it is fine not to raise
+        if'duplicate key value violates unique constraint' in str(e.body):
+            return
+        raise e from None
+
+    except urllib3.exceptions.HTTPError as e:
+        raise e from None
+
+    # Wait for the filesystem picks up the newly added Filestore folder
+    start = time()
+
+    while not os.path.isdir(f"/teamspace/filestore_folders/{folder_name}") and (time() - start) < create_timeout:
+        sleep(1)
+
+
 def delete_data_connection(name: str):
     """Utility to delete a data connection
 

lightning_sdk/llm/llm.py

@@ -1,10 +1,10 @@
-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.api.llm_api import LLMApi, authenticate
 from lightning_sdk.lightning_cloud.openapi.models.v1_conversation_response_chunk import V1ConversationResponseChunk
+from lightning_sdk.llm.public_assistants import PUBLIC_MODELS
 
 PUBLIC_MODEL_PROVIDERS: Dict[str, str] = {
     "openai": "OpenAI",
@@ -14,17 +14,6 @@ PUBLIC_MODEL_PROVIDERS: 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")
-        with open(json_path) as f:
-            return json.load(f)
-    except Exception as e:
-        print(f"[warning] Failed to load public_assistants.json: {e}")
-        return {}
-
-
 class LLM:
     _auth_info_cached: ClassVar[bool] = False
     _cached_auth_info: ClassVar[Dict[str, Optional[str]]] = {}
@@ -67,9 +56,9 @@ class LLM:
                     "(e.g., 'my-org/my-teamspace')."
                 ) from e
 
+        self._model_provider, self._model_name = self._parse_model_name(name)
         self._get_auth_info(teamspace_name)
 
-        self._model_provider, self._model_name = self._parse_model_name(name)
         self._enable_async = enable_async
 
         # Reuse LLMApi per teamspace (as billing is based on teamspace)
@@ -109,6 +98,7 @@ class LLM:
             if teamspace_name is None:
                 # local users with no given teamspace
                 try:
+                    authenticate(model=f"{self.provider}/{self.name}")
                     teamspace_api = TeamspaceApi()
                     user_api = UserApi()
                     authed_user = user_api._client.auth_service_get_user()
@@ -138,7 +128,7 @@ class LLM:
             }
             LLM._auth_info_cached = True
             if LLM._public_assistants is None:
-                LLM._public_assistants = _load_public_assistants()
+                LLM._public_assistants = PUBLIC_MODELS
         # Always assign to the current instance
         self._teamspace_name = LLM._cached_auth_info["teamspace_name"]
         self._teamspace_id = LLM._cached_auth_info["teamspace_id"]

lightning_sdk/llm/public_assistants.py

@@ -0,0 +1,17 @@
+PUBLIC_MODELS = {
+    "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/gpt-oss-20b": {
+        "id": "ast_01k1y2ywfh9zzexjzrva6y96ms",
+        "context_length": 128000,
+    },
+    "lightning-ai/gpt-oss-120b": {
+        "id": "ast_01k1z7ajxbz7wq62xj5bc0mrcs",
+        "context_length": 128000,
+    },
+}

lightning_sdk/services/file_endpoint.py

@@ -8,7 +8,6 @@ import urllib3
 
 from lightning_sdk.api.utils import _FileUploader
 from lightning_sdk.lightning_cloud.login import Auth
-from lightning_sdk.lightning_cloud.openapi import CommandArgumentCommandArgumentType
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 from lightning_sdk.services.utilities import _get_cluster, _get_project, _get_service_url
 from lightning_sdk.utils.resolve import _resolve_deprecated_cluster
@@ -162,7 +161,7 @@ class Client:
 class Argument:
     """A holder for the service argument."""
 
-    def __init__(self, name: str, type: CommandArgumentCommandArgumentType, **kwargs: Any) -> None:  # noqa: A002
+    def __init__(self, name: str, type: str, **kwargs: Any) -> None:  # noqa: A002
         self._name = name
         self._type = type
         self._value = None
@@ -171,12 +170,12 @@ class Argument:
     @property
     def is_text(self) -> bool:
         """Whether this argument is of type Text."""
-        return self._type == CommandArgumentCommandArgumentType.TEXT
+        return self._type == "Text"
 
     @property
     def is_file(self) -> bool:
         """Whether this argument is of type File."""
-        return self._type == CommandArgumentCommandArgumentType.FILE
+        return self._type == "File"
 
     @property
     def value(self) -> Any:

{lightning_sdk-2025.8.6rc0.dist-info → lightning_sdk-2025.8.6rc2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightning_sdk
-Version: 2025.8.6rc0
+Version: 2025.8.6rc2
 Summary: SDK to develop using Lightning AI Studios
 Author-email: Lightning-AI <justus@lightning.ai>
 License: MIT License