lightning-sdk 0.2.19__py3-none-any.whl → 0.2.21__py3-none-any.whl

lightning_sdk/llm/llm.py

@@ -1,6 +1,6 @@
 import os
 import warnings
-from typing import Dict, Generator, List, Optional, Set, Tuple, Union
+from typing import AsyncGenerator, Dict, Generator, List, Optional, Set, Tuple, Union
 
 from lightning_sdk.api.llm_api import LLMApi
 from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
@@ -18,6 +18,7 @@ class LLM:
         self,
         name: str,
         teamspace: Optional[str] = None,
+        enable_async: Optional[bool] = False,
     ) -> None:
         """Initializes the LLM instance with teamspace information, which is required for billing purposes.
 
@@ -30,6 +31,7 @@ class LLM:
             name (str): The name of the model or resource.
             teamspace (Optional[str]): The specified teamspace for billing. If not provided, it will be resolved
                                        through the above methods.
+            enable_async (Optional[bool]): Enable async requests
 
         Raises:
             ValueError: If teamspace information cannot be resolved.
@@ -67,6 +69,7 @@ class LLM:
         self._model_provider, self._model_name = self._parse_model_name(name)
 
         self._llm_api = LLMApi()
+        self._enable_async = enable_async
 
         try:
             # check if it is a org model
@@ -174,6 +177,41 @@ class LLM:
         for line in result:
             yield line.choices[0].delta.content
 
+    async def _async_stream_text(self, output: str) -> AsyncGenerator[str, None]:
+        async for chunk in output:
+            if chunk.choices and chunk.choices[0].delta:
+                yield chunk.choices[0].delta.content
+
+    async def _async_chat(
+        self,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        max_completion_tokens: Optional[int] = 500,
+        images: Optional[Union[List[str], str]] = None,
+        conversation: Optional[str] = None,
+        metadata: Optional[Dict[str, str]] = None,
+        stream: bool = False,
+        upload_local_images: bool = False,
+    ) -> Union[str, AsyncGenerator[str, None]]:
+        conversation_id = self._conversations.get(conversation) if conversation else None
+        output = await self._llm_api.async_start_conversation(
+            prompt=prompt,
+            system_prompt=system_prompt,
+            max_completion_tokens=max_completion_tokens,
+            images=images,
+            assistant_id=self._model.id,
+            conversation_id=conversation_id,
+            billing_project_id=self._teamspace.id,
+            metadata=metadata,
+            name=conversation,
+            stream=stream,
+        )
+        if not stream:
+            if conversation and not conversation_id:
+                self._conversations[conversation] = output.conversation_id
+            return output.choices[0].delta.content
+        return self._async_stream_text(output)
+
     def chat(
         self,
         prompt: str,
@@ -198,6 +236,19 @@ class LLM:
                     self._teamspace.upload_file(file_path=image, remote_path=f"images/{os.path.basename(image)}")
 
         conversation_id = self._conversations.get(conversation) if conversation else None
+
+        if self._enable_async:
+            return self._async_chat(
+                prompt,
+                system_prompt,
+                max_completion_tokens,
+                images,
+                conversation,
+                metadata,
+                stream,
+                upload_local_images,
+            )
+
         output = self._llm_api.start_conversation(
             prompt=prompt,
             system_prompt=system_prompt,

lightning_sdk/pipeline/pipeline.py

@@ -63,7 +63,7 @@ class Pipeline:
         pipeline = None
 
         if name.startswith("pip_"):
-            pipeline = self._pipeline_api.get_pipeline_by_id(name, self._teamspace.id)
+            pipeline = self._pipeline_api.get_pipeline_by_id(self._teamspace.id, name)
 
         if pipeline:
             self._name = pipeline.name

lightning_sdk/services/license.py

@@ -9,9 +9,74 @@ from importlib import metadata
 from pathlib import Path
 from typing import Optional, Tuple
 
-from lightning_sdk.api.license_api import LicenseApi
+from lightning_sdk.api.license_api import LICENSE_SIGNING_URL, LicenseApi, generate_url_user_settings
 from lightning_sdk.lightning_cloud.login import Auth
 
+MESSAGE_NOT_AUTHENTICATED = (
+    "┌─────────────────────────────────────────────────────────────────────────┐\n"
+    "│ ⚠️  No authenticated user found or license API is not available.        │\n"
+    "│                                                                         │\n"
+    "│   Please make sure you are logged in and have a valid license.          │\n"
+    "│   If you're not logged in, you can use the following command:           │\n"
+    "│                                                                         │\n"
+    "│     lightning login                                                     │\n"
+    "└─────────────────────────────────────────────────────────────────────────┘"
+)
+MESSAGE_AUTH_NO_LICENSE = (
+    "┌──────────────────────────────────────────────────────────────────────────────────────────────┐\n"
+    "│ ⚠️ No valid license found for the authenticated user for product '{self.product_name}'.      │\n"
+    "│                                                                                              │\n"
+    "│   Please ensure you have an approved license for this product.                               │\n"
+    "│   If you believe this is an error, please contact support.                                   │\n"
+    "│                                                                                              │\n"
+    "│   You can review or update your license settings here:                                       │\n"
+    f"│     {LICENSE_SIGNING_URL:<89}│\n"
+    "└──────────────────────────────────────────────────────────────────────────────────────────────┘"
+)
+MESSAGE_GUIDE_SIGN_LICENSE = (
+    "┌───────────────────────────────────────────────────────────────────────────────────────────────────────────┐\n"
+    "│ ⚠️ {reason} │\n"
+    "│ Details: license key ({license_strats}...{license_ends}) for package {package_name:<56} │\n"
+    "│ Please make sure you have signed the license agreement and set the license key.                           │\n"
+    "│                                                                                                           │\n"
+    f"│ Sign the license agreement here (if you dont have Lightning account, you will asked to create one):       │\n"
+    f"│   {generate_url_user_settings():<102}\n"
+    "│                                                                                                           │\n"
+    "│ Once you have the license key, you may need to reinstall this package to activate it. Use the commands:   │\n"
+    "│                                                                                                           │\n"
+    "│   export LIGHTNING_LICENSE_KEY=<your_license_key>                                                         │\n"
+    "│   pip install --force-reinstall --no-deps {package_name:<63} │\n"
+    "│                                                                                                           │\n"
+    "│ For more information, please refer to the documentation.                                                  │\n"
+    "└───────────────────────────────────────────────────────────────────────────────────────────────────────────┘"
+)
+
+
+def generate_message_guide_sign_license(package_name: str, reason: str, license_key: str = "") -> str:
+    """Generate a default message for signing the license agreement."""
+    if not license_key:
+        license_key = "." * 64
+    return MESSAGE_GUIDE_SIGN_LICENSE.format(
+        package_name=package_name,
+        reason=reason.ljust(102, " "),
+        license_strats=license_key[:5],
+        license_ends=license_key[-5:],
+    )
+
+
+MESSAGE_WITH_KEY_OFFLINE = (
+    "┌──────────────────────────────────────────────────────────────────────────────────────┐\n"
+    "│ ⚠️ License key ({license_strats}...{license_ends}) is set, but the system is offline.                    │\n"
+    "│                                                                                      │\n"
+    "│ Please ensure your license key is valid and that the system is connected online.     │\n"
+    "└──────────────────────────────────────────────────────────────────────────────────────┘"
+)
+
+
+def generate_message_with_key_offline(license_key: str) -> str:
+    """Generate a message for when the license key is set but the system is offline."""
+    return MESSAGE_WITH_KEY_OFFLINE.format(license_strats=license_key[:5], license_ends=license_key[-5:])
+
 
 class LightningLicense:
     """This class is used to manage the license for the Lightning SDK."""
@@ -48,9 +113,6 @@ class LightningLicense:
         """Validate the license key."""
         if not self.is_online():
             raise ConnectionError("No internet connection.")
-        if not self.license_api:
-            # todo: this need to be exposed with public API
-            raise RuntimeError("License validation is allowed only for authenticated users.")
 
         return self.license_api.valid_license(
             license_key=self.license_key,
@@ -76,10 +138,7 @@ class LightningLicense:
             if not user_id:
                 return False
         if not self.license_api:
-            self._stream_messages(
-                "No authenticated user found or license API is not available."
-                " Please make sure you are logged in and have a valid license."
-            )
+            self._stream_messages(MESSAGE_NOT_AUTHENTICATED)
             return False
 
         licenses = self.license_api.list_user_licenses(user_id=user_id)
@@ -90,10 +149,7 @@ class LightningLicense:
                 and license_info.is_valid
             ):
                 return True
-        self._stream_messages(
-            f"No valid license found for authenticated user for product {self.product_name}."
-            " Please make sure you have a approved product license or contact support."
-        )
+        self._stream_messages(MESSAGE_AUTH_NO_LICENSE)
         return False
 
     @staticmethod
@@ -143,15 +199,13 @@ class LightningLicense:
                 else:
                     self._is_valid = self._check_user_license(user_id=user_id)
         elif self.license_key:
-            self._stream_messages(
-                f"License key ({self.license_key[:5]}...{self.license_key[-5:]}) is set but the system is offline."
-                " Please make sure you have a valid license key and the system is online."
-            )
+            self._stream_messages(generate_message_with_key_offline(self.license_key))
         else:
             self._stream_messages(
-                "License key is not set neither cannot be found in the package root or user home."
-                " Please make sure you have signed the license agreement and set the license key."
-                " For more information, please refer to the documentation.",
+                generate_message_guide_sign_license(
+                    package_name=self.product_name,
+                    reason="License key is not set neither cannot be found in the package root or user home.",
+                )
             )
 
         return self._is_valid
@@ -265,9 +319,11 @@ def check_license(
     )
     if lit_license.is_valid is False:
         stream_messages(
-            f"License key for `{lit_license.product_name}` is not valid.\n"
-            f" Key: {lit_license.license_key}\n"
-            " Please make sure you have a valid license key."
+            generate_message_guide_sign_license(
+                package_name=lit_license.product_name,
+                reason="License key is not valid or not set.",
+                license_key=lit_license.license_key,
+            )
         )
 
 

lightning_sdk/services/utilities.py

@@ -35,7 +35,9 @@ def _get_project(client: LightningClient, project_name: Optional[str] = None) ->
     raise ValueError("No valid projects found. Please reach out to lightning.ai team to create a project")
 
 
-def _get_cluster(client: LightningClient, project_id: str, cluster_id: Optional[str] = None) -> V1ProjectClusterBinding:
+def _get_cluster(
+    client: LightningClient, project_id: str, cluster_id: Optional[str] = None, allow_neoclouds: bool = False
+) -> V1ProjectClusterBinding:
     """Get a project membership for the user from the backend."""
     clusters = client.projects_service_list_project_cluster_bindings(project_id=project_id)
     if cluster_id:
@@ -47,6 +49,18 @@ def _get_cluster(client: LightningClient, project_id: str, cluster_id: Optional[
             f"Found {[c.cluster_id for c in clusters.clusters]}."
         )
 
+    # filter neoclouds out
+    if not allow_neoclouds:
+        cluster_objs = client.cluster_service_list_clusters(project_id=project_id)
+        # filter for aws or gcp cluster
+        valid_clusters = filter(
+            lambda c: c.spec.aws_v1 is not None or c.spec.google_cloud_v1 is not None, cluster_objs.clusters
+        )
+        valid_clusters = {c.id for c in valid_clusters}
+
+        if len(valid_clusters):
+            clusters.clusters = list(filter(lambda c: c.cluster_id in valid_clusters, clusters.clusters))
+
     clusters = sorted(clusters.clusters, key=lambda x: x.created_at)
     if len(clusters):
         return clusters[0]

{lightning_sdk-0.2.19.dist-info → lightning_sdk-0.2.21.dist-info}/METADATA

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