import os from dataclasses import dataclass from typing import ( Any, AsyncGenerator, ClassVar, Coroutine, Dict, Generator, List, Literal, Optional, Tuple, Union, cast, ) from lightning_sdk.api import TeamspaceApi, UserApi 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 from lightning_sdk.utils.resolve import _resolve_org, _resolve_teamspace _PROD_CLOUD_URL = "https://platform.lightning.ai" PUBLIC_MODEL_PROVIDERS: Dict[str, str] = { "openai": "OpenAI", "anthropic": "Anthropic", "google": "Google", "lightning-ai": "lightning-ai", } @dataclass class ModelMetadata: name: str provider: str status: str context_length: int max_completion_tokens: Optional[int] prompt_price: float completion_price: float capabilities: Dict[str, bool] throughput: float time_to_first_token: float def __str__(self) -> str: """Return a user-friendly string representation of the model metadata. Returns: str: A formatted multi-line string containing model information including name, provider, status, context length, pricing, performance metrics, and key capabilities (images and files support). """ return f""" Model: {self.name} Provider: {self.provider} Status: {self.status} Context Length: {self.context_length:,} tokens Pricing: ${self.prompt_price:.2e}/prompt token, ${self.completion_price:.2e}/completion token Performance: {self.throughput:.1f} tokens/sec, {self.time_to_first_token:.1f}ms TTFT Capabilities: Images={self.capabilities.get("images", False)}, Files={self.capabilities.get("files", False)} """.strip() 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, Dict[str, Any]]]] = 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, teamspace: Optional[str] = None, enable_async: Optional[bool] = False, ) -> None: """Initializes the LLM instance with teamspace information, which is required for billing purposes. Args: name (str): The name of the model or resource. teamspace (Optional[str]): The specified teamspace for billing. If not provided, it will use the available default teamspace. enable_async (Optional[bool]): Enable async requests Raises: ValueError: If teamspace information cannot be resolved. """ teamspace_name = None teamspace_owner = None if teamspace: if "/" in teamspace: try: teamspace_owner, teamspace_name = teamspace.split("/", maxsplit=1) except ValueError as e: raise ValueError( f"Invalid teamspace format: '{teamspace}'. " "Teamspace should be specified as '{org}/{teamspace_name}' or '{org}'" "(e.g., 'my-org/my-teamspace', 'my-org')." ) from e else: # org is given teamspace_name = teamspace self._model_provider, self._model_name = self._parse_model_name(name) self._get_auth_info(teamspace_owner, teamspace_name) self._enable_async = enable_async # 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] self._context_length: Optional[int] = None self._model_id = self._get_model_id() self._conversations: Dict[str, str] = {} self._metadata: Optional[ModelMetadata] = None @property def name(self) -> str: """The model name portion of the identifier (everything after the first ``/``).""" return self._model_name @property def provider(self) -> str: """The provider prefix of the model identifier (e.g. ``openai``, ``anthropic``).""" return self._model_provider @property def metadata(self) -> ModelMetadata: """Lazily fetched metadata for this model, including pricing and capabilities.""" if self._metadata is None: model = self._llm_api.get_model_metadata(self._teamspace_id, self._model_name) abilities = ( model.abilities or type( "obj", (object,), {"can_receive_images": False, "can_receive_files": False, "can_call_hub_deployment": False}, )() ) self._metadata = ModelMetadata( name=self._model_name, provider=self._model_provider, status=model.status, context_length=int(model.context_length), max_completion_tokens=int(model.max_completion_tokens) if model.max_completion_tokens != "0" else None, prompt_price=model.prompt_token_price, completion_price=model.completion_token_price, capabilities={ "images": abilities.can_receive_images, "files": abilities.can_receive_files, "hub_deployment": abilities.can_call_hub_deployment, }, throughput=model.throughput, time_to_first_token=model.time_to_first_token, ) return self._metadata @property def context_length(self) -> Optional[int]: """Context length for the current model.""" if self._context_length is None: try: self._context_length = self.metadata.context_length except Exception as e: raise ValueError(f"Cannot access context length: {e}") from e return self._context_length def get_context_length(self, model: Optional[str] = None) -> Optional[int]: """Return the context-window length for the given model. Args: model: Full ``provider/name`` model identifier to query. Defaults to this instance's model when ``None``. Returns: Optional[int]: Context length in tokens, or ``None`` if unavailable. Raises: ValueError: If the model metadata cannot be retrieved. """ model_key = model or f"{self._model_provider}/{self._model_name}" context_info = (self._public_assistants or {}).get(model_key) if context_info and "context_length" in context_info: return int(context_info["context_length"]) try: metadata_name = model.split("/", maxsplit=1)[-1] if model else self._model_name temp_metadata = self._llm_api.get_model_metadata(self._teamspace_id, metadata_name) return int(temp_metadata.context_length) except Exception as e: raise ValueError(f"Cannot access context length of model '{model}': {e}") from e def _get_auth_info(self, teamspace_owner: Optional[str] = None, teamspace_name: Optional[str] = None) -> None: """Resolve and cache teamspace / user auth information for billing. On the first call the method resolves the teamspace from the given owner/name combination (or falls back to environment variables and the default teamspace) and caches the result on the class so subsequent ``LLM`` instances reuse it. Args: teamspace_owner: Username or organisation that owns the target teamspace. teamspace_name: Name of the target teamspace (or organisation name when ``teamspace_owner`` is ``None``). Raises: ValueError: If the teamspace cannot be resolved or no teamspace information is available at all. """ if not LLM._auth_info_cached: if teamspace_owner and teamspace_name: # org with specific teamspace try: t = _resolve_teamspace(teamspace=teamspace_name, org=None, user=teamspace_owner) except Exception: try: t = _resolve_teamspace(teamspace=teamspace_name, org=teamspace_owner, user=None) except Exception as err: raise ValueError( f"Teamspace {teamspace_owner}/{teamspace_name} not found." "Please verify owner name (username or organization) and the teamspace name are correct." ) from err if t is None: raise ValueError(f"Teamspace {teamspace_owner}/{teamspace_name} not found.") from None os.environ["LIGHTNING_TEAMSPACE"] = t.name os.environ["LIGHTNING_CLOUD_PROJECT_ID"] = t.id elif teamspace_name and teamspace_owner is None: # if only org name is given, use the default teamspace try: org = _resolve_org(teamspace_name) if org is None: raise ValueError(f"Organization {teamspace_name} not found.") teamspace_api = TeamspaceApi() teamspaces = teamspace_api.list_teamspaces(org.id) or [] teamspace = teamspaces[0] except Exception as err: raise ValueError( f"Organization {teamspace_name} not found. Please verify the organization name." ) from err os.environ["LIGHTNING_CLOUD_PROJECT_ID"] = teamspace.id os.environ["LIGHTNING_TEAMSPACE"] = teamspace.name else: if teamspace_name is None: # studio users teamspace_name = os.environ.get("LIGHTNING_TEAMSPACE", None) 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() teamspaces = teamspace_api.list_teamspaces(owner_id=authed_user.id) or [] default_teamspace = teamspaces[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://platform.lightning.ai/sign-up " "and re-run your script, or set environment variable LIGHTNING_TEAMSPACE=." ) from err # 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), "user_name": os.environ.get("LIGHTNING_USERNAME", ""), "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 = PUBLIC_MODELS # Always assign to the current instance self._teamspace_name = LLM._cached_auth_info["teamspace_name"] teamspace_id = LLM._cached_auth_info["teamspace_id"] if teamspace_id is None: raise ValueError("Teamspace ID is missing from the resolved authentication information.") self._teamspace_id = teamspace_id self._user_name = LLM._cached_auth_info["user_name"] self._org_name = LLM._cached_auth_info["org_name"] self._cloud_url = LLM._cached_auth_info["cloud_url"] self._org = None @staticmethod def _parse_model_name(name: str) -> Tuple[str, str]: """Parses the model name into provider and model name. >>> LLM._parse_model_name("openai/v1/gpt-3.5-turbo") ('openai', 'v1/gpt-3.5-turbo') """ if "/" not in name: raise ValueError( f"Invalid model name format: '{name}'. " "Model name must be in the format `provider/model_name`." "(e.g., 'lightning-ai/gpt-oss-20b')" ) provider, model_name = name.split("/", maxsplit=1) return provider.lower(), model_name def _get_model_id(self) -> str: """Resolve the assistant ID for the current model. Tries public model providers first, then organisation models, then user models. Returns: str: The assistant ID used by the LLM API. Raises: ValueError: If no matching model is found for the configured provider and name. """ if self._model_provider in PUBLIC_MODEL_PROVIDERS: # if prod if ( self._cloud_url == _PROD_CLOUD_URL and LLM._public_assistants and f"{self._model_provider}/{self._model_name}" in LLM._public_assistants ): 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], 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._model_provider == "lightning-ai": # Try model provider model try: return self._llm_api.get_assistant( model_provider=self._model_provider, model_name=self._model_name, user_name="", org_name="", ) except Exception: pass # 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 # 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: """Fetch remote conversations for this model and merge them into the local cache.""" 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 def _stream_chat_response( self, result: Generator[V1ConversationResponseChunk, None, None], conversation: Optional[str] = None, full_response: bool = False, ) -> Generator[str, None, None]: """Yield text chunks (or full response objects) from a streaming conversation result. Args: result: Raw streaming generator from the LLM API. conversation: Conversation name; used to cache the conversation ID from the first chunk. full_response: When ``True``, yield the raw ``V1ConversationResponseChunk`` objects instead of the extracted text content. Returns: Generator[str, None, None]: Stream of response text chunks (or raw chunk objects when ``full_response=True``). """ first_line = next(result, None) if first_line: if conversation and first_line.conversation_id: self._conversations[conversation] = first_line.conversation_id if full_response: yield first_line else: yield first_line.choices[0].delta.content for line in result: if full_response: yield line else: yield line.choices[0].delta.content async def _async_stream_text( self, output: AsyncGenerator[V1ConversationResponseChunk, None], full_response: bool = False ) -> AsyncGenerator[Any, None]: async for chunk in output: if chunk.choices and chunk.choices[0].delta: if full_response: yield chunk else: yield chunk.choices[0].delta.content async def _async_chat( self, prompt: str, system_prompt: Optional[str] = None, max_completion_tokens: Optional[int] = None, images: Optional[Union[List[str], str]] = None, conversation: Optional[str] = None, metadata: Optional[Dict[str, str]] = None, stream: bool = False, full_response: bool = False, reasoning_effort: Optional[Literal["none", "low", "medium", "high"]] = None, **kwargs: Any, ) -> Union[str, V1ConversationResponseChunk, AsyncGenerator[Any, None]]: normalized_images = [images] if isinstance(images, str) else images 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=normalized_images, assistant_id=self._model_id, conversation_id=conversation_id, billing_project_id=self._teamspace_id, metadata=metadata, name=conversation, stream=stream, **kwargs, ) if not stream: response = cast(V1ConversationResponseChunk, output) if conversation and not conversation_id: self._conversations[conversation] = response.conversation_id if full_response: return response return response.choices[0].delta.content return self._async_stream_text(cast(AsyncGenerator[V1ConversationResponseChunk, None], output), full_response) def chat( self, prompt: str, system_prompt: Optional[str] = None, max_completion_tokens: Optional[int] = None, images: Optional[Union[List[str], str]] = None, conversation: Optional[str] = None, metadata: Optional[Dict[str, str]] = None, stream: bool = False, full_response: bool = False, tools: Optional[List[Dict[str, Any]]] = None, reasoning_effort: Optional[Literal["none", "low", "medium", "high"]] = None, **kwargs: Any, ) -> Union[ V1ConversationResponseChunk, Generator[V1ConversationResponseChunk, None, None], str, Generator[str, None, None], Coroutine[Any, Any, Union[str, V1ConversationResponseChunk, AsyncGenerator[Any, None]]], ]: if reasoning_effort is not None and reasoning_effort not in ["none", "low", "medium", "high"]: raise ValueError("reasoning_effort must be 'none', 'low', 'medium', 'high', or None") if conversation and conversation not in self._conversations: self._get_conversations() normalized_images = [images] if isinstance(images, str) else images if normalized_images: for image in normalized_images: if not isinstance(image, str): raise NotImplementedError(f"Image type {type(image)} are not supported yet.") conversation_id = self._conversations.get(conversation) if conversation else None if self._enable_async: return self._async_chat( prompt, system_prompt, max_completion_tokens, normalized_images, conversation, metadata, stream, full_response, reasoning_effort, **kwargs, ) output = self._llm_api.start_conversation( prompt=prompt, system_prompt=system_prompt, max_completion_tokens=max_completion_tokens, images=normalized_images, assistant_id=self._model_id, conversation_id=conversation_id, billing_project_id=self._teamspace_id, metadata=metadata, name=conversation, stream=stream, tools=tools, reasoning_effort=reasoning_effort, **kwargs, ) if not stream: response = cast(V1ConversationResponseChunk, output) if conversation and not conversation_id: self._conversations[conversation] = response.conversation_id if full_response: return response return response.choices[0].delta.content return self._stream_chat_response( cast(Generator[V1ConversationResponseChunk, None, None], output), conversation=conversation, full_response=full_response, ) def list_conversations(self) -> List[str]: """Return the names of all conversations for this model. Returns: List[str]: Conversation names tracked for the current model instance. """ self._get_conversations() return list(self._conversations.keys()) def _get_conversation_messages(self, conversation_id: str) -> List[Any]: """Fetch the raw message list for a conversation by its ID. Args: conversation_id: The API-level ID of the conversation to retrieve. Returns: List[Any]: The conversation messages as returned by the LLM API. """ return self._llm_api.get_conversation(assistant_id=self._model_id, conversation_id=conversation_id) def get_history(self, conversation: str) -> List[Dict[str, Any]]: """Return the message history of a named conversation as a list of role/content dicts. Args: conversation: The display name of the conversation to retrieve. Returns: Optional[List[Dict]]: A list of ``{"role": ..., "content": ...}`` dicts for each message in the conversation, in chronological order. Raises: ValueError: If no conversation with the given name exists. """ if conversation not in self._conversations: self._get_conversations() if conversation not in self._conversations: raise ValueError( f"Conversation '{conversation}' not found. \nAvailable conversations: {self._conversations.keys()}" ) messages = self._get_conversation_messages(self._conversations[conversation]) history = [] for message in messages: if message.author.role == "user": history.append({"role": "user", "content": message.content[0].parts[0]}) elif message.author.role == "assistant": history.append({"role": "assistant", "content": message.content[0].parts[0]}) return history def reset_conversation(self, conversation: str) -> None: """Clear the message history for a named conversation. The conversation is deleted on the server and removed from the local cache. Sending a new message under the same name will start a fresh conversation. Args: conversation: The display name of the conversation to reset. """ if conversation not in self._conversations: self._get_conversations() if conversation in self._conversations: self._llm_api.reset_conversation( assistant_id=self._model_id, conversation_id=self._conversations[conversation], ) del self._conversations[conversation]