lightning-sdk 0.2.23__py3-none-any.whl → 0.2.24rc0__py3-none-any.whl

lightning_sdk/llm/llm.py

@@ -1,19 +1,37 @@
+import json
 import os
-import warnings
-from typing import AsyncGenerator, Dict, Generator, List, Optional, Set, Tuple, Union
+from typing import AsyncGenerator, ClassVar, Dict, Generator, List, Optional, Tuple, Union
 
 from lightning_sdk.api.llm_api import LLMApi
-from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
-from lightning_sdk.lightning_cloud.openapi import V1Assistant
 from lightning_sdk.lightning_cloud.openapi.models.v1_conversation_response_chunk import V1ConversationResponseChunk
-from lightning_sdk.lightning_cloud.openapi.rest import ApiException
-from lightning_sdk.organization import Organization
-from lightning_sdk.owner import Owner
-from lightning_sdk.teamspace import Teamspace
-from lightning_sdk.utils.resolve import _get_authed_user, _resolve_org, _resolve_teamspace
+
+PUBLIC_MODEL_PROVIDERS: Dict[str, str] = {
+    "openai": "OpenAI",
+    "anthropic": "Anthropic",
+    "google": "Google",
+}
+
+
+def _load_public_assistants() -> Dict[str, str]:
+    """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]]] = {}
+    _llm_api_cache: ClassVar[Dict[Optional[str], LLMApi]] = {}
+    _public_assistants: ClassVar[Optional[Dict[str, str]]] = None
+
+    def __new__(cls, name: str, teamspace: Optional[str] = None, enable_async: Optional[bool] = False) -> "LLM":
+        return super().__new__(cls)
+
     def __init__(
         self,
         name: str,
@@ -36,67 +54,18 @@ class LLM:
         Raises:
             ValueError: If teamspace information cannot be resolved.
         """
-        menu = _TeamspacesMenu()
-        user = _get_authed_user()
-        possible_teamspaces = menu._get_possible_teamspaces(user)
-        if teamspace is None:
-            # get current teamspace
-            self._teamspace = _resolve_teamspace(teamspace=None, org=None, user=None)
-        else:
-            self._teamspace = Teamspace(**menu._get_teamspace_from_name(teamspace, possible_teamspaces))
-
-        if self._teamspace is None:
-            # select the first available teamspace
-            first_teamspace = next(iter(possible_teamspaces.values()), None)
-
-            if first_teamspace:
-                self._teamspace = Teamspace(
-                    name=first_teamspace["name"],
-                    org=first_teamspace["org"],
-                    user=first_teamspace["user"],
-                )
-                warnings.warn(
-                    f"No teamspace given. Using teamspace: {self._teamspace.name}.",
-                    UserWarning,
-                    stacklevel=2,
-                )
-
-        if self._teamspace is None:
-            raise ValueError("Teamspace is required for billing but could not be resolved. ")
-
-        self._user = user
+        # TODO support user input teamspace
+        self._get_auth_info()
 
         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
-            self._org = _resolve_org(self._model_provider)
+        # Reuse LLMApi per teamspace (as billing is based on teamspace)
+        if teamspace not in LLM._llm_api_cache:
+            LLM._llm_api_cache[teamspace] = LLMApi()
+        self._llm_api = LLM._llm_api_cache[teamspace]
 
-            try:
-                # check if user has access to the org
-                self._org_models = self._build_model_lookup(self._get_org_models())
-            except ApiException:
-                warnings.warn(
-                    f"User is not authenticated to access the model in organization: '{self._model_provider}'.\n"
-                    " Proceeding with appropriate org models, user models, or public models.",
-                    UserWarning,
-                    stacklevel=2,
-                )
-                self._model_provider = None
-                raise
-        except ApiException:
-            if isinstance(self._teamspace.owner, Organization):
-                self._org = self._teamspace.owner
-            else:
-                self._org = None
-            self._org_models = self._build_model_lookup(self._get_org_models())
-
-        self._public_models = self._build_model_lookup(self._get_public_models())
-        self._user_models = self._build_model_lookup(self._get_user_models())
-        self._model = self._get_model()
+        self._model_id = self._get_model_id()
         self._conversations = {}
 
     @property
@@ -107,9 +76,33 @@ class LLM:
     def provider(self) -> str:
         return self._model_provider
 
-    @property
-    def owner(self) -> Optional[Owner]:
-        return self._teamspace.owner
+    def _get_auth_info(self) -> None:
+        if not LLM._auth_info_cached:
+            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._cached_auth_info = {
+                "teamspace_name": teamspace_name,
+                "teamspace_id": os.environ.get("LIGHTNING_CLOUD_PROJECT_ID", None),
+                "user_name": os.environ.get("LIGHTNING_USERNAME", ""),
+                "user_id": os.environ.get("LIGHTNING_USER_ID", None),
+                "org_name": os.environ.get("LIGHTNING_ORG", ""),
+                "cloud_url": os.environ.get("LIGHTNING_CLOUD_URL", None),
+            }
+            LLM._auth_info_cached = True
+            if LLM._public_assistants is None:
+                LLM._public_assistants = _load_public_assistants()
+        # Always assign to the current instance
+        self._teamspace_name = LLM._cached_auth_info["teamspace_name"]
+        self._teamspace_id = LLM._cached_auth_info["teamspace_id"]
+        self._user_name = LLM._cached_auth_info["user_name"]
+        self._user_id = LLM._cached_auth_info["user_id"]
+        self._org_name = LLM._cached_auth_info["org_name"]
+        self._cloud_url = LLM._cached_auth_info["cloud_url"]
+        self._org = None
 
     def _parse_model_name(self, name: str) -> Tuple[str, str]:
         parts = name.split("/")
@@ -117,50 +110,60 @@ class LLM:
             # a user model or a org model
             return None, parts[0]
         if len(parts) == 2:
-            return parts[0], parts[1]
+            return parts[0].lower(), parts[1]
         raise ValueError(
             f"Model name must be in the format `organization/model_name` or `model_name`, but got '{name}'."
         )
 
-    def _build_model_lookup(self, endpoints: List[str]) -> Dict[str, Set[str]]:
-        result = {}
-        for endpoint in endpoints:
-            result.setdefault(endpoint.model, []).append(endpoint)
-        return result
-
-    def _get_public_models(self) -> List[str]:
-        return self._llm_api.get_public_models()
-
-    def _get_org_models(self) -> List[str]:
-        return self._llm_api.get_org_models(self._org.id) if self._org else []
-
-    def _get_user_models(self) -> List[str]:
-        return self._llm_api.get_user_models(self._user.id) if self._user else []
-
-    def _get_model(self) -> V1Assistant:
-        # TODO how to handle multiple models with same model type? For now, just use the first one
-        if self._model_name in self._public_models:
-            return self._public_models.get(self._model_name)[0]
-        if self._model_name in self._org_models:
-            return self._org_models.get(self._model_name)[0]
-        if self._model_name in self._user_models:
-            return self._user_models.get(self._model_name)[0]
-
-        available_models = []
-        if self._public_models:
-            available_models.append(f"Public Models: {', '.join(self._public_models.keys())}")
-
-        if self._org and self._org_models:
-            available_models.append(f"Org ({self._org.name}) Models: {', '.join(self._org_models.keys())}")
+    # returns the assistant ID
+    def _get_model_id(self) -> str:
+        if self._model_provider in PUBLIC_MODEL_PROVIDERS:
+            # if prod
+            if (
+                self._cloud_url == "https://lightning.ai"
+                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}"]
+            try:
+                return self._llm_api.get_assistant(
+                    model_provider=PUBLIC_MODEL_PROVIDERS[self._model_provider],
+                    model_name=self._model_name,
+                    user_name="",
+                    org_name="",
+                )
+            except Exception as e:
+                raise ValueError(
+                    f"Public model '{self._model_provider}/{self._model_name}' not found. "
+                    "Please check the model name or provider."
+                ) from e
 
-        if self._user and self._user_models:
-            available_models.append(f"User ({self._user.name}) Models: {', '.join(self._user_models.keys())}")
+        # Try organization model
+        try:
+            return self._llm_api.get_assistant(
+                model_provider="",
+                model_name=self._model_name,
+                user_name="",
+                org_name=self._model_provider,
+            )
+        except Exception:
+            pass
 
-        available_models_str = "\n".join(available_models)
-        raise ValueError(f"Model '{self._model_name}' not found. \nAvailable models: \n{available_models_str}")
+        # Try user model
+        try:
+            return self._llm_api.get_assistant(
+                model_provider="",
+                model_name=self._model_name,
+                user_name=self._model_provider,
+                org_name="",
+            )
+        except Exception as user_error:
+            raise ValueError(
+                f"Model '{self._model_provider}/{self._model_name}' not found as either an org or user model.\n"
+            ) from user_error
 
     def _get_conversations(self) -> None:
-        conversations = self._llm_api.list_conversations(assistant_id=self._model.id)
+        conversations = self._llm_api.list_conversations(assistant_id=self._model_id)
         for conversation in conversations:
             if conversation.name and conversation.name not in self._conversations:
                 self._conversations[conversation.name] = conversation.id
@@ -191,7 +194,6 @@ class LLM:
         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(
@@ -199,9 +201,9 @@ class LLM:
             system_prompt=system_prompt,
             max_completion_tokens=max_completion_tokens,
             images=images,
-            assistant_id=self._model.id,
+            assistant_id=self._model_id,
             conversation_id=conversation_id,
-            billing_project_id=self._teamspace.id,
+            billing_project_id=self._teamspace_id,
             metadata=metadata,
             name=conversation,
             stream=stream,
@@ -221,7 +223,6 @@ class LLM:
         conversation: Optional[str] = None,
         metadata: Optional[Dict[str, str]] = None,
         stream: bool = False,
-        upload_local_images: bool = False,
     ) -> Union[str, Generator[str, None, None]]:
         if conversation and conversation not in self._conversations:
             self._get_conversations()
@@ -232,8 +233,6 @@ class LLM:
             for image in images:
                 if not isinstance(image, str):
                     raise NotImplementedError(f"Image type {type(image)} are not supported yet.")
-                if not image.startswith("http") and upload_local_images:
-                    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
 
@@ -246,7 +245,6 @@ class LLM:
                 conversation,
                 metadata,
                 stream,
-                upload_local_images,
             )
 
         output = self._llm_api.start_conversation(
@@ -254,9 +252,9 @@ class LLM:
             system_prompt=system_prompt,
             max_completion_tokens=max_completion_tokens,
             images=images,
-            assistant_id=self._model.id,
+            assistant_id=self._model_id,
             conversation_id=conversation_id,
-            billing_project_id=self._teamspace.id,
+            billing_project_id=self._teamspace_id,
             metadata=metadata,
             name=conversation,
             stream=stream,
@@ -272,7 +270,7 @@ class LLM:
         return list(self._conversations.keys())
 
     def _get_conversation_messages(self, conversation_id: str) -> Optional[str]:
-        return self._llm_api.get_conversation(assistant_id=self._model.id, conversation_id=conversation_id)
+        return self._llm_api.get_conversation(assistant_id=self._model_id, conversation_id=conversation_id)
 
     def get_history(self, conversation: str) -> Optional[List[Dict]]:
         if conversation not in self._conversations:
@@ -297,7 +295,7 @@ class LLM:
             self._get_conversations()
         if conversation in self._conversations:
             self._llm_api.reset_conversation(
-                assistant_id=self._model.id,
+                assistant_id=self._model_id,
                 conversation_id=self._conversations[conversation],
             )
             del self._conversations[conversation]

lightning_sdk/llm/public_assistants.json

@@ -0,0 +1,8 @@
+{
+    "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"
+}

lightning_sdk/pipeline/pipeline.py

@@ -86,6 +86,12 @@ class Pipeline:
         if len(steps) == 0:
             raise ValueError("The provided steps is empty")
 
+        provided_cloud_account = None
+        if self._cloud_account:
+            provided_cloud_account = self._cloud_account
+        elif self._default_cluster:
+            provided_cloud_account = self._default_cluster.cluster_id
+
         for step_idx, pipeline_step in enumerate(steps):
             if pipeline_step.name in [None, ""]:
                 pipeline_step.name = f"step-{step_idx}"
@@ -98,20 +104,29 @@ class Pipeline:
                 pipeline_step.cloud_account = self._studio.cloud_account
                 pipeline_step.studio = self._studio
 
+            if not pipeline_step.cloud_account and isinstance(provided_cloud_account, str):
+                pipeline_step.cloud_account = provided_cloud_account
+
         cluster_ids = set(step.cloud_account for step in steps if step.cloud_account not in ["", None])  # noqa: C401
 
-        cloud_account = list(cluster_ids)[0] if len(cluster_ids) == 1 and self._cloud_account is None else ""  # noqa: RUF015
+        cloud_account = (
+            list(cluster_ids)[0] if len(cluster_ids) == 1 and self._cloud_account is None else ""  # noqa: RUF015
+        )
 
         steps = [step.to_proto(self._teamspace, cloud_account, self._shared_filesystem) for step in steps]
 
         proto_steps = prepare_steps(steps)
         schedules = schedules or []
 
+        for schedule_idx, schedule in enumerate(schedules):
+            if schedule.name is None:
+                schedule.name = f"schedule-{schedule_idx}"
+
         parent_pipeline_id = None if self._pipeline is None else self._pipeline.id
 
         self._pipeline = self._pipeline_api.create_pipeline(
             self._name,
-            self._teamspace.id,
+            self._teamspace,
             proto_steps,
             self._shared_filesystem,
             schedules,

lightning_sdk/pipeline/printer.py

@@ -1,7 +1,8 @@
 import os
 from typing import Any, ClassVar, Dict, List
 
-from lightning_sdk.lightning_cloud.openapi.models import V1JobSpec, V1Pipeline, V1PipelineStepType
+from lightning_sdk.lightning_cloud.openapi.models import V1Pipeline, V1PipelineStepType
+from lightning_sdk.pipeline.utils import _get_spec
 
 
 class PipelinePrinter:
@@ -31,7 +32,7 @@ class PipelinePrinter:
         self._schedules = schedules
         cluster_ids: set[str] = set()
         for step in self._proto_steps:
-            job_spec = self._get_spec(step)
+            job_spec = _get_spec(step)
             cluster_ids.add(job_spec.cluster_id)
         self._cluster_ids = cluster_ids
 
@@ -107,15 +108,15 @@ class PipelinePrinter:
 
         if self._shared_filesystem.enabled and len(self._cluster_ids) == 1:
             shared_path = ""
+            cluster_id = list(self._cluster_ids)[0]  # noqa: RUF015
             if self._pipeline.shared_filesystem.s3_folder:
-                cluster_id = list(self._cluster_ids)[0]  # noqa: RUF015
                 shared_path = f"/teamspace/s3_folders/pipelines-{cluster_id}"
+            if self._pipeline.shared_filesystem.gcs_folder:
+                shared_path = f"/teamspace/gcs_folders/pipelines-{cluster_id}"
+            if self._pipeline.shared_filesystem.efs:
+                shared_path = f"/teamspace/efs_connections/pipelines-{cluster_id}"
+            if self._pipeline.shared_filesystem.filestore:
+                shared_path = f"/teamspace/gcs_connections/pipelines-{cluster_id}"
+
             if shared_path:
                 self._print(f"  - {shared_path}")
-
-    def _get_spec(self, step: Any) -> V1JobSpec:
-        if step.type == V1PipelineStepType.DEPLOYMENT:
-            return step.deployment.spec
-        if step.type == V1PipelineStepType.MMT:
-            return step.mmt.spec
-        return step.job.spec

lightning_sdk/pipeline/steps.py

@@ -69,12 +69,15 @@ class DeploymentStep:
 
         self.machine = machine or Machine.CPU
         self.image = image
+        autoscaling_metric_name = (
+            ("CPU" if self.machine.is_cpu() else "GPU") if isinstance(self.machine, Machine) else "CPU"
+        )
         self.autoscale = autoscale or AutoScaleConfig(
             min_replicas=0,
             max_replicas=1,
             target_metrics=[
                 AutoScalingMetric(
-                    name="CPU" if self.machine.is_cpu() else "GPU",
+                    name=autoscaling_metric_name,
                     target=80,
                 )
             ],

lightning_sdk/pipeline/utils.py

@@ -1,6 +1,11 @@
-from typing import List, Literal, Optional, Union
-
-from lightning_sdk.lightning_cloud.openapi.models import V1PipelineStep, V1SharedFilesystem
+from typing import Any, List, Literal, Optional, Union
+
+from lightning_sdk.lightning_cloud.openapi.models import (
+    V1JobSpec,
+    V1PipelineStep,
+    V1PipelineStepType,
+    V1SharedFilesystem,
+)
 from lightning_sdk.studio import Studio
 
 DEFAULT = "DEFAULT"
@@ -21,7 +26,7 @@ def prepare_steps(steps: List["V1PipelineStep"]) -> List["V1PipelineStep"]:
         else:
             raise ValueError(f"A step with the name {current_step.name} already exists.")
 
-    if steps[0].wait_for != DEFAULT:
+    if steps[0].wait_for not in [None, DEFAULT, []]:
         raise ValueError("The first step isn't allowed to receive `wait_for=...`.")
 
     steps[0].wait_for = []
@@ -89,3 +94,23 @@ def _to_wait_for(wait_for: Optional[Union[str, List[str]]]) -> Optional[Union[Li
         return []
 
     return wait_for if isinstance(wait_for, list) else [wait_for]
+
+
+def _get_cloud_account(steps: List[V1PipelineStep]) -> Optional[str]:
+    if len(steps) == 0:
+        return None
+
+    cluster_ids: set[str] = set()
+    for step in steps:
+        job_spec = _get_spec(step)
+        cluster_ids.add(job_spec.cluster_id)
+
+    return sorted(cluster_ids)[0]
+
+
+def _get_spec(step: Any) -> V1JobSpec:
+    if step.type == V1PipelineStepType.DEPLOYMENT:
+        return step.deployment.spec
+    if step.type == V1PipelineStepType.MMT:
+        return step.mmt.spec
+    return step.job.spec

lightning_sdk/studio.py

@@ -88,6 +88,9 @@ class Studio:
 
         self._plugins = {}
 
+        if self._teamspace is None:
+            raise ValueError("Couldn't resolve teamspace from the provided name, org, or user")
+
         if provider is not None:
             if isinstance(provider, str) and provider in Provider.__members__:
                 provider = Provider(provider)

{lightning_sdk-0.2.23.dist-info → lightning_sdk-0.2.24rc0.dist-info}/METADATA

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