Rename block parameter to wait

Signed-off-by: Mattt Zmuda <[email protected]>

Author: mattt

replicate/deployment.py

@@ -426,7 +426,7 @@ def create(
                 client=self._client,
                 file_encoding_strategy=file_encoding_strategy,
             )
-        headers = _create_prediction_headers(block=params.pop("block", None))
+        headers = _create_prediction_headers(wait=params.pop("wait", None))
         body = _create_prediction_body(version=None, input=input, **params)
 
         resp = self._client._request(
@@ -454,7 +454,7 @@ async def async_create(
                 client=self._client,
                 file_encoding_strategy=file_encoding_strategy,
             )
-        headers = _create_prediction_headers(block=params.pop("block", None))
+        headers = _create_prediction_headers(wait=params.pop("wait", None))
         body = _create_prediction_body(version=None, input=input, **params)
 
         resp = await self._client._async_request(

replicate/model.py

@@ -401,7 +401,7 @@ def create(
                 client=self._client,
                 file_encoding_strategy=file_encoding_strategy,
             )
-        headers = _create_prediction_headers(block=params.pop("block", None))
+        headers = _create_prediction_headers(wait=params.pop("wait", None))
         body = _create_prediction_body(version=None, input=input, **params)
 
         resp = self._client._request(
@@ -432,7 +432,7 @@ async def async_create(
                 client=self._client,
                 file_encoding_strategy=file_encoding_strategy,
             )
-        headers = _create_prediction_headers(block=params.pop("block", None))
+        headers = _create_prediction_headers(wait=params.pop("wait", None))
         body = _create_prediction_body(version=None, input=input, **params)
 
         resp = await self._client._async_request(

replicate/prediction.py

@@ -383,8 +383,14 @@ class CreatePredictionParams(TypedDict):
         stream: NotRequired[bool]
         """Enable streaming of prediction output."""
 
-        block: NotRequired[bool]
-        """Wait until the prediction is completed before returning."""
+        wait: NotRequired[Union[int, bool]]
+        """
+        Wait until the prediction is completed before returning.
+
+        If `True`, wait a predetermined number of seconds until the prediction
+        is completed before returning.
+        If an `int`, wait for the specified number of seconds.
+        """
 
         file_encoding_strategy: NotRequired[FileEncodingStrategy]
         """The strategy to use for encoding files in the prediction input."""
@@ -466,7 +472,7 @@ def create(  # type: ignore
                 client=self._client,
                 file_encoding_strategy=file_encoding_strategy,
             )
-        headers = _create_prediction_headers(block=params.pop("block", None))
+        headers = _create_prediction_headers(wait=params.pop("wait", None))
         body = _create_prediction_body(
             version,
             input,
@@ -559,7 +565,7 @@ async def async_create(  # type: ignore
                 client=self._client,
                 file_encoding_strategy=file_encoding_strategy,
             )
-        headers = _create_prediction_headers(block=params.pop("block", None))
+        headers = _create_prediction_headers(wait=params.pop("wait", None))
         body = _create_prediction_body(
             version,
             input,
@@ -612,13 +618,15 @@ async def async_cancel(self, id: str) -> Prediction:
 
 def _create_prediction_headers(
     *,
-    block: Optional[bool] = None,
+    wait: Optional[Union[int, bool]] = None,
 ) -> Dict[str, Any]:
     headers = {}
 
-    if block:
-        headers["Prefer"] = "wait"
-
+    if wait:
+        if isinstance(wait, bool):
+            headers["Prefer"] = "wait"
+        elif isinstance(wait, int):
+            headers["Prefer"] = f"wait={wait}"
     return headers