feat: support NOT_GIVEN type

Author: etang-ai21

ai21/__init__.py

@@ -13,6 +13,7 @@
 from ai21.logger import setup_logger
 from ai21.services.sagemaker import SageMaker
 from ai21.version import VERSION
+from ai21.types import NOT_GIVEN
 
 __version__ = VERSION
 setup_logger()
@@ -63,4 +64,5 @@ def __getattr__(name: str) -> Any:
     "AI21SageMakerClient",
     "BedrockModelID",
     "SageMaker",
+    "NOT_GIVEN",
 ]

ai21/clients/studio/resources/studio_completion.py

@@ -1,8 +1,11 @@
+from __future__ import annotations
+
 from typing import Optional, List
 
 from ai21.clients.common.completion_base import Completion
 from ai21.clients.studio.resources.studio_resource import StudioResource
 from ai21.models import Penalty, CompletionsResponse
+from ai21.types import NOT_GIVEN
 
 
 class StudioCompletion(StudioResource, Completion):
@@ -11,7 +14,7 @@ def create(
         model: str,
         prompt: str,
         *,
-        max_tokens: Optional[int] = None,
+        max_tokens: int | NOT_GIVEN = None,
         num_results: Optional[int] = 1,
         min_tokens: Optional[int] = 0,
         temperature: Optional[float] = 0.7,

ai21/types.py

@@ -0,0 +1,30 @@
+from typing_extensions import Literal
+
+
+# Sentinel class used until PEP 0661 is accepted
+class NotGiven:
+    """
+    A sentinel singleton class used to distinguish omitted keyword arguments
+    from those passed in with the value None (which may have different behavior).
+
+    For example:
+
+    ```py
+    def get(timeout: Union[int, NotGiven, None] = NotGiven()) -> Response:
+        ...
+
+
+    get(timeout=1)  # 1s timeout
+    get(timeout=None)  # No timeout
+    get()  # Default timeout behavior, which may not be statically known at the method definition.
+    ```
+    """
+
+    def __bool__(self) -> Literal[False]:
+        return False
+
+    def __repr__(self) -> str:
+        return "NOT_GIVEN"
+
+
+NOT_GIVEN = NotGiven()

poetry.lock

@@ -60,6 +60,7 @@ requests = "^2.31.0"
 ai21-tokenizer = "^0.3.9"
 boto3 = { version = "^1.28.82", optional = true }
 dataclasses-json = "^0.6.3"
+typing-extensions = "^4.9.0"
 
 
 [tool.poetry.group.dev.dependencies]