feat(api): update via SDK Studio (#8)

Author: stainless-app[bot]

README.md

@@ -31,6 +31,8 @@ from luma_ai import LumaAI
 
 client = LumaAI(
     auth_token="My Auth Token",
+    # or 'production' | 'staging' | 'localhost'; defaults to "production".
+    environment="production_api",
 )
 
 generation = client.generations.create(
@@ -49,6 +51,8 @@ from luma_ai import AsyncLumaAI
 
 client = AsyncLumaAI(
     auth_token="My Auth Token",
+    # or 'production' | 'staging' | 'localhost'; defaults to "production".
+    environment="production_api",
 )
 
 

src/luma_ai/__init__.py

@@ -3,7 +3,18 @@
 from . import types
 from ._types import NOT_GIVEN, NoneType, NotGiven, Transport, ProxiesTypes
 from ._utils import file_from_path
-from ._client import Client, LumaAI, Stream, Timeout, Transport, AsyncClient, AsyncLumaAI, AsyncStream, RequestOptions
+from ._client import (
+    ENVIRONMENTS,
+    Client,
+    LumaAI,
+    Stream,
+    Timeout,
+    Transport,
+    AsyncClient,
+    AsyncLumaAI,
+    AsyncStream,
+    RequestOptions,
+)
 from ._models import BaseModel
 from ._version import __title__, __version__
 from ._response import APIResponse as APIResponse, AsyncAPIResponse as AsyncAPIResponse
@@ -58,6 +69,7 @@
     "AsyncStream",
     "LumaAI",
     "AsyncLumaAI",
+    "ENVIRONMENTS",
     "file_from_path",
     "BaseModel",
     "DEFAULT_TIMEOUT",

src/luma_ai/_client.py

@@ -3,8 +3,8 @@
 from __future__ import annotations
 
 import os
-from typing import Any, Union, Mapping
-from typing_extensions import Self, override
+from typing import Any, Dict, Union, Mapping, cast
+from typing_extensions import Self, Literal, override
 
 import httpx
 
@@ -33,6 +33,7 @@
 )
 
 __all__ = [
+    "ENVIRONMENTS",
     "Timeout",
     "Transport",
     "ProxiesTypes",
@@ -44,6 +45,13 @@
     "AsyncClient",
 ]
 
+ENVIRONMENTS: Dict[str, str] = {
+    "production": "http://api.lumalabs.ai/dream-machine/v1alpha",
+    "production_api": "http://internal-api.virginia.labs.lumalabs.ai/dream-machine/v1alpha",
+    "staging": "http://internal-api.sandbox.labs.lumalabs.ai/dream-machine/v1alpha",
+    "localhost": "http://localhost:9600/dream-machine/v1alpha",
+}
+
 
 class LumaAI(SyncAPIClient):
     ping: resources.PingResource
@@ -54,11 +62,14 @@ class LumaAI(SyncAPIClient):
     # client options
     auth_token: str
 
+    _environment: Literal["production", "production_api", "staging", "localhost"] | NotGiven
+
     def __init__(
         self,
         *,
         auth_token: str,
-        base_url: str | httpx.URL | None = None,
+        environment: Literal["production", "production_api", "staging", "localhost"] | NotGiven = NOT_GIVEN,
+        base_url: str | httpx.URL | None | NotGiven = NOT_GIVEN,
         timeout: Union[float, Timeout, None, NotGiven] = NOT_GIVEN,
         max_retries: int = DEFAULT_MAX_RETRIES,
         default_headers: Mapping[str, str] | None = None,
@@ -80,10 +91,31 @@ def __init__(
         """Construct a new synchronous luma_ai client instance."""
         self.auth_token = auth_token
 
-        if base_url is None:
-            base_url = os.environ.get("LUMA_AI_BASE_URL")
-        if base_url is None:
-            base_url = f"http://api.lumalabs.ai/dream-machine/v1alpha"
+        self._environment = environment
+
+        base_url_env = os.environ.get("LUMA_AI_BASE_URL")
+        if is_given(base_url) and base_url is not None:
+            # cast required because mypy doesn't understand the type narrowing
+            base_url = cast("str | httpx.URL", base_url)  # pyright: ignore[reportUnnecessaryCast]
+        elif is_given(environment):
+            if base_url_env and base_url is not None:
+                raise ValueError(
+                    "Ambiguous URL; The `LUMA_AI_BASE_URL` env var and the `environment` argument are given. If you want to use the environment, you must pass base_url=None",
+                )
+
+            try:
+                base_url = ENVIRONMENTS[environment]
+            except KeyError as exc:
+                raise ValueError(f"Unknown environment: {environment}") from exc
+        elif base_url_env is not None:
+            base_url = base_url_env
+        else:
+            self._environment = environment = "production"
+
+            try:
+                base_url = ENVIRONMENTS[environment]
+            except KeyError as exc:
+                raise ValueError(f"Unknown environment: {environment}") from exc
 
         super().__init__(
             version=__version__,
@@ -125,6 +157,7 @@ def copy(
         self,
         *,
         auth_token: str | None = None,
+        environment: Literal["production", "production_api", "staging", "localhost"] | None = None,
         base_url: str | httpx.URL | None = None,
         timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
         http_client: httpx.Client | None = None,
@@ -160,6 +193,7 @@ def copy(
         return self.__class__(
             auth_token=auth_token or self.auth_token,
             base_url=base_url or self.base_url,
+            environment=environment or self._environment,
             timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
             http_client=http_client,
             max_retries=max_retries if is_given(max_retries) else self.max_retries,
@@ -215,11 +249,14 @@ class AsyncLumaAI(AsyncAPIClient):
     # client options
     auth_token: str
 
+    _environment: Literal["production", "production_api", "staging", "localhost"] | NotGiven
+
     def __init__(
         self,
         *,
         auth_token: str,
-        base_url: str | httpx.URL | None = None,
+        environment: Literal["production", "production_api", "staging", "localhost"] | NotGiven = NOT_GIVEN,
+        base_url: str | httpx.URL | None | NotGiven = NOT_GIVEN,
         timeout: Union[float, Timeout, None, NotGiven] = NOT_GIVEN,
         max_retries: int = DEFAULT_MAX_RETRIES,
         default_headers: Mapping[str, str] | None = None,
@@ -241,10 +278,31 @@ def __init__(
         """Construct a new async luma_ai client instance."""
         self.auth_token = auth_token
 
-        if base_url is None:
-            base_url = os.environ.get("LUMA_AI_BASE_URL")
-        if base_url is None:
-            base_url = f"http://api.lumalabs.ai/dream-machine/v1alpha"
+        self._environment = environment
+
+        base_url_env = os.environ.get("LUMA_AI_BASE_URL")
+        if is_given(base_url) and base_url is not None:
+            # cast required because mypy doesn't understand the type narrowing
+            base_url = cast("str | httpx.URL", base_url)  # pyright: ignore[reportUnnecessaryCast]
+        elif is_given(environment):
+            if base_url_env and base_url is not None:
+                raise ValueError(
+                    "Ambiguous URL; The `LUMA_AI_BASE_URL` env var and the `environment` argument are given. If you want to use the environment, you must pass base_url=None",
+                )
+
+            try:
+                base_url = ENVIRONMENTS[environment]
+            except KeyError as exc:
+                raise ValueError(f"Unknown environment: {environment}") from exc
+        elif base_url_env is not None:
+            base_url = base_url_env
+        else:
+            self._environment = environment = "production"
+
+            try:
+                base_url = ENVIRONMENTS[environment]
+            except KeyError as exc:
+                raise ValueError(f"Unknown environment: {environment}") from exc
 
         super().__init__(
             version=__version__,
@@ -286,6 +344,7 @@ def copy(
         self,
         *,
         auth_token: str | None = None,
+        environment: Literal["production", "production_api", "staging", "localhost"] | None = None,
         base_url: str | httpx.URL | None = None,
         timeout: float | Timeout | None | NotGiven = NOT_GIVEN,
         http_client: httpx.AsyncClient | None = None,
@@ -321,6 +380,7 @@ def copy(
         return self.__class__(
             auth_token=auth_token or self.auth_token,
             base_url=base_url or self.base_url,
+            environment=environment or self._environment,
             timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
             http_client=http_client,
             max_retries=max_retries if is_given(max_retries) else self.max_retries,

tests/test_client.py

@@ -552,6 +552,16 @@ def test_base_url_env(self) -> None:
             client = LumaAI(auth_token=auth_token, _strict_response_validation=True)
             assert client.base_url == "http://localhost:5000/from/env/"
 
+        # explicit environment arg requires explicitness
+        with update_env(LUMA_AI_BASE_URL="http://localhost:5000/from/env"):
+            with pytest.raises(ValueError, match=r"you must pass base_url=None"):
+                LumaAI(auth_token=auth_token, _strict_response_validation=True, environment="production")
+
+            client = LumaAI(
+                base_url=None, auth_token=auth_token, _strict_response_validation=True, environment="production"
+            )
+            assert str(client.base_url).startswith("http://api.lumalabs.ai/dream-machine/v1alpha")
+
     @pytest.mark.parametrize(
         "client",
         [
@@ -1265,6 +1275,16 @@ def test_base_url_env(self) -> None:
             client = AsyncLumaAI(auth_token=auth_token, _strict_response_validation=True)
             assert client.base_url == "http://localhost:5000/from/env/"
 
+        # explicit environment arg requires explicitness
+        with update_env(LUMA_AI_BASE_URL="http://localhost:5000/from/env"):
+            with pytest.raises(ValueError, match=r"you must pass base_url=None"):
+                AsyncLumaAI(auth_token=auth_token, _strict_response_validation=True, environment="production")
+
+            client = AsyncLumaAI(
+                base_url=None, auth_token=auth_token, _strict_response_validation=True, environment="production"
+            )
+            assert str(client.base_url).startswith("http://api.lumalabs.ai/dream-machine/v1alpha")
+
     @pytest.mark.parametrize(
         "client",
         [