fix linter issues

Author: dtmeadows

src/replicate/_client.py

@@ -8,6 +8,8 @@
 
 import httpx
 
+from replicate.types.prediction_create_params import PredictionCreateParamsWithoutVersion
+
 from . import _exceptions
 from ._qs import Querystring
 from .types import PredictionOutput, PredictionCreateParams
@@ -135,7 +137,7 @@ def run(
         ref: str,
         *,
         wait: Union[int, bool, NotGiven] = NOT_GIVEN,
-        **params: Unpack[PredictionCreateParams],
+        **params: Unpack[PredictionCreateParamsWithoutVersion],
     ) -> PredictionOutput | FileOutput | Iterable[FileOutput] | Dict[str, FileOutput]:
         """Run a model and wait for its output."""
         from .lib._predictions import run

src/replicate/lib/_predictions.py

@@ -3,6 +3,8 @@
 from typing import TYPE_CHECKING, Dict, Union, Iterable
 from typing_extensions import Unpack
 
+from replicate.types.prediction_create_params import PredictionCreateParamsWithoutVersion
+
 from ..types import PredictionOutput, PredictionCreateParams
 from .._types import NOT_GIVEN, NotGiven
 from .._utils import is_given
@@ -21,7 +23,7 @@ def run(
     *,
     wait: Union[int, bool, NotGiven] = NOT_GIVEN,
     # use_file_output: Optional[bool] = True,
-    **params: Unpack[PredictionCreateParams],
+    **params: Unpack[PredictionCreateParamsWithoutVersion],
 ) -> PredictionOutput | FileOutput | Iterable[FileOutput] | Dict[str, FileOutput]:
     from ._files import transform_output
 
@@ -40,7 +42,8 @@ def run(
             params.setdefault("prefer", f"wait={wait}")
 
     # TODO: support more ref types
-    prediction = client.predictions.create(version=ref, **params)
+    params_with_version: PredictionCreateParams = {**params, "version": ref}
+    prediction = client.predictions.create(**params_with_version)
 
     # Currently the "Prefer: wait" interface will return a prediction with a status
     # of "processing" rather than a terminal state because it returns before the
@@ -91,7 +94,8 @@ async def async_run(
             params.setdefault("prefer", f"wait={wait}")
 
     # TODO: support more ref types
-    prediction = await client.predictions.create(version=ref, **params)
+    params_with_version: PredictionCreateParams = {**params, "version": ref}
+    prediction = await client.predictions.create(**params_with_version)
 
     # Currently the "Prefer: wait" interface will return a prediction with a status
     # of "processing" rather than a terminal state because it returns before the

src/replicate/resources/predictions.py

@@ -54,10 +54,10 @@ def with_streaming_response(self) -> PredictionsResourceWithStreamingResponse:
 
     def wait(self, prediction_id: str) -> Prediction:
         """Wait for prediction to finish."""
-        prediction = self.retrieve(prediction_id)
+        prediction = self.get(prediction_id)
         while prediction.status not in PREDICTION_TERMINAL_STATES:
             self._sleep(self._client.poll_interval)
-            prediction = self.retrieve(prediction.id)
+            prediction = self.get(prediction.id)
         return prediction
 
     def create(
@@ -473,10 +473,10 @@ def with_streaming_response(self) -> AsyncPredictionsResourceWithStreamingRespon
 
     async def wait(self, prediction_id: str) -> Prediction:
         """Wait for prediction to finish."""
-        prediction = await self.retrieve(prediction_id)
+        prediction = await self.get(prediction_id)
         while prediction.status not in PREDICTION_TERMINAL_STATES:
             await self._sleep(self._client.poll_interval)
-            prediction = await self.retrieve(prediction.id)
+            prediction = await self.get(prediction.id)
         return prediction
 
     async def create(

src/replicate/types/__init__.py

@@ -15,7 +15,10 @@
 from .deployment_create_params import DeploymentCreateParams as DeploymentCreateParams
 from .deployment_list_response import DeploymentListResponse as DeploymentListResponse
 from .deployment_update_params import DeploymentUpdateParams as DeploymentUpdateParams
-from .prediction_create_params import PredictionCreateParams as PredictionCreateParams
+from .prediction_create_params import (
+    PredictionCreateParams as PredictionCreateParams,
+    PredictionCreateParamsWithoutVersion as PredictionCreateParamsWithoutVersion,
+)
 from .training_cancel_response import TrainingCancelResponse as TrainingCancelResponse
 from .deployment_create_response import DeploymentCreateResponse as DeploymentCreateResponse
 from .deployment_update_response import DeploymentUpdateResponse as DeploymentUpdateResponse

src/replicate/types/prediction_create_params.py

@@ -7,10 +7,10 @@
 
 from .._utils import PropertyInfo
 
-__all__ = ["PredictionCreateParams"]
+__all__ = ["PredictionCreateParams", "PredictionCreateParamsWithoutVersion"]
 
 
-class PredictionCreateParams(TypedDict, total=False):
+class PredictionCreateParamsWithoutVersion(TypedDict, total=False):
     input: Required[object]
     """The model's input as a JSON object.
 
@@ -36,9 +36,6 @@ class PredictionCreateParams(TypedDict, total=False):
     - you don't need to use the file again (Replicate will not store it)
     """
 
-    version: Required[str]
-    """The ID of the model version that you want to run."""
-
     stream: bool
     """**This field is deprecated.**
 
@@ -94,3 +91,8 @@ class PredictionCreateParams(TypedDict, total=False):
     """
 
     prefer: Annotated[str, PropertyInfo(alias="Prefer")]
+
+
+class PredictionCreateParams(PredictionCreateParamsWithoutVersion):
+    version: Required[str]
+    """The ID of the model version that you want to run."""