lightning-sdk 0.2.22__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/__init__.py

@@ -1,5 +1,14 @@
+from lightning_sdk import Studio
 from lightning_sdk.pipeline.pipeline import Pipeline
 from lightning_sdk.pipeline.schedule import Schedule
-from lightning_sdk.pipeline.types import MMT, Deployment, Job
+from lightning_sdk.pipeline.steps import DeploymentReleaseStep, DeploymentStep, JobStep, MMTStep
 
-__all__ = ["Pipeline", "Job", "MMT", "Deployment", "Schedule"]
+__all__ = [
+    "Pipeline",
+    "JobStep",
+    "MMTStep",
+    "DeploymentStep",
+    "Schedule",
+    "Studio",
+    "DeploymentReleaseStep",
+]

lightning_sdk/pipeline/pipeline.py

@@ -1,3 +1,4 @@
+import os
 from typing import TYPE_CHECKING, List, Optional, Union
 
 from lightning_sdk.api import UserApi
@@ -5,9 +6,10 @@ from lightning_sdk.api.pipeline_api import PipelineApi
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.organization import Organization
 from lightning_sdk.pipeline.printer import PipelinePrinter
-from lightning_sdk.pipeline.types import MMT, Deployment, Job
+from lightning_sdk.pipeline.steps import DeploymentStep, JobStep, MMTStep, _get_studio
 from lightning_sdk.pipeline.utils import prepare_steps
 from lightning_sdk.services.utilities import _get_cluster
+from lightning_sdk.studio import Studio
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.user import User
 from lightning_sdk.utils.resolve import _resolve_org, _resolve_teamspace, _resolve_user
@@ -25,6 +27,7 @@ class Pipeline:
         user: Union[str, "User", None] = None,
         cloud_account: Optional[str] = None,
         shared_filesystem: Optional[bool] = None,
+        studio: Optional[Union[Studio, str]] = None,
     ) -> None:
         """The Lightning Pipeline can be used to create complex DAG.
 
@@ -58,10 +61,12 @@ class Pipeline:
         )
 
         self._pipeline_api = PipelineApi()
-        self._cloud_account = _get_cluster(
+        self._cloud_account = cloud_account
+        self._default_cluster = _get_cluster(
             client=self._pipeline_api._client, project_id=self._teamspace.id, cluster_id=cloud_account
         )
-        self._shared_filesystem = shared_filesystem
+        self._shared_filesystem = shared_filesystem if shared_filesystem is not None else True
+        self._studio = _get_studio(studio)
         self._is_created = False
 
         pipeline = None
@@ -75,35 +80,66 @@ class Pipeline:
         else:
             self._pipeline = None
 
-    def run(self, steps: List[Union[Job, Deployment, MMT]], schedules: Optional[List["Schedule"]] = None) -> None:
+    def run(
+        self, steps: List[Union[JobStep, DeploymentStep, MMTStep]], schedules: Optional[List["Schedule"]] = None
+    ) -> None:
         if len(steps) == 0:
             raise ValueError("The provided steps is empty")
 
-        for step_idx, step in enumerate(steps):
-            if step.name in [None, ""]:
-                raise ValueError(f"The step {step_idx} requires a name")
+        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}"
+
+            if (
+                self._studio is not None
+                and (pipeline_step.image == "" or pipeline_step.image is None)
+                and pipeline_step.studio is None
+            ):
+                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
+        )
 
-        steps = [
-            step.to_proto(self._teamspace, self._cloud_account.cluster_id or "", self._shared_filesystem)
-            for step in steps
-        ]
+        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 or False,
+            self._shared_filesystem,
             schedules,
             parent_pipeline_id,
         )
 
         printer = PipelinePrinter(
-            self._name, parent_pipeline_id is None, self._pipeline, self._teamspace, proto_steps, schedules
+            self._name,
+            parent_pipeline_id is None,
+            self._pipeline,
+            self._teamspace,
+            proto_steps,
+            schedules,
         )
         printer.print_summary()
 
@@ -124,3 +160,7 @@ class Pipeline:
         if self._pipeline:
             return self._pipeline.name
         return None
+
+    @classmethod
+    def from_env(cls) -> "Pipeline":
+        return Pipeline(name=os.getenv("LIGHTNING_PIPELINE_ID", ""))

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, V1PipelineStepType
+from lightning_sdk.lightning_cloud.openapi.models import V1Pipeline, V1PipelineStepType
+from lightning_sdk.pipeline.utils import _get_spec
 
 
 class PipelinePrinter:
@@ -14,14 +15,26 @@ class PipelinePrinter:
     }
 
     def __init__(
-        self, name: str, initial: bool, pipeline: Any, teamspace: Any, proto_steps: List[Any], schedules: List[Any]
+        self,
+        name: str,
+        initial: bool,
+        pipeline: V1Pipeline,
+        teamspace: Any,
+        proto_steps: List[Any],
+        schedules: List[Any],
     ) -> None:
         self._name = name
         self._initial = initial
         self._pipeline = pipeline
         self._teamspace = teamspace
         self._proto_steps = proto_steps
+        self._shared_filesystem = pipeline.shared_filesystem
         self._schedules = schedules
+        cluster_ids: set[str] = set()
+        for step in self._proto_steps:
+            job_spec = _get_spec(step)
+            cluster_ids.add(job_spec.cluster_id)
+        self._cluster_ids = cluster_ids
 
     def print_summary(self) -> None:
         """Prints the full, formatted summary of the created pipeline."""
@@ -32,6 +45,7 @@ class PipelinePrinter:
         self._print_steps()
         self._print_schedules()
         self._print_cloud_account()
+        self._print_shared_filesystem()
         self._print_footer()
 
     def _print(self, value: str) -> None:
@@ -75,7 +89,7 @@ class PipelinePrinter:
         team: str = self._teamspace.name
         pipeline_name: str = self._name
 
-        pipeline_url = f"{cloud_url}/{owner}/{team}/pipelines/{pipeline_name}?app_id=pipeline&section=Graph"
+        pipeline_url = f"{cloud_url}/{owner}/{team}/pipelines/{pipeline_name}?app_id=pipeline"
 
         self._print("\n" + "─" * 60)
         self._print(f"🔗 View your pipeline in the browser:\n   {pipeline_url}")
@@ -85,18 +99,24 @@ class PipelinePrinter:
         if not self._proto_steps:
             return
 
-        cluster_ids: set[str] = set()
-        for step in self._proto_steps:
-            job_spec = self._get_spec(step)
-            cluster_ids.add(job_spec.cluster_id)
-
-        self._print(f"\nCloud account{'s' if len(cluster_ids) > 1 else ''}:")
-        for cluster_id in sorted(cluster_ids):
+        self._print(f"\nCloud account{'s' if len(self._cluster_ids) > 1 else ''}:")
+        for cluster_id in sorted(self._cluster_ids):
             self._print(f"  - {cluster_id}")
 
-    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
+    def _print_shared_filesystem(self) -> None:
+        self._print(f"\nShared filesystem: {self._shared_filesystem.enabled}")
+
+        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:
+                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}")

lightning_sdk/pipeline/schedule.py

@@ -1,7 +1,8 @@
 from dataclasses import dataclass
+from typing import Optional
 
 
 @dataclass
 class Schedule:
-    name: str
     cron_expression: str
+    name: Optional[str] = None