clean up helpers to match underlying api changes

Author: dtmeadows

pyproject.toml

@@ -14,7 +14,6 @@ dependencies = [
     "anyio>=3.5.0, <5",
     "distro>=1.7.0, <2",
     "sniffio",
-    "asyncio>=3.4.3",
 ]
 requires-python = ">= 3.8"
 classifiers = [

src/replicate/lib/_files.py

@@ -10,8 +10,6 @@
 
 import httpx
 
-from replicate.types.prediction_output import PredictionOutput
-
 from .._utils import is_mapping, is_sequence
 
 # Use TYPE_CHECKING to avoid circular imports
@@ -218,7 +216,7 @@ def __repr__(self) -> str:
         return f'{self.__class__.__name__}("{self.url}")'
 
 
-def transform_output(value: PredictionOutput, client: "Replicate | AsyncReplicate") -> Any:
+def transform_output(value: object, client: "Replicate | AsyncReplicate") -> Any:
     """
     Transform the output of a prediction to a `FileOutput` object if it's a URL.
     """

src/replicate/lib/_predictions.py

@@ -10,7 +10,7 @@
 from replicate.types.prediction import Prediction
 from replicate.types.prediction_create_params import PredictionCreateParamsWithoutVersion
 
-from ..types import PredictionOutput, PredictionCreateParams
+from ..types import PredictionCreateParams
 from .._types import NOT_GIVEN, NotGiven
 from .._utils import is_given
 from ._models import Model, Version, ModelVersionIdentifier, resolve_reference
@@ -29,7 +29,7 @@ def run(
     use_file_output: Optional[bool] = True,
     file_encoding_strategy: Optional["FileEncodingStrategy"] = None,
     **params: Unpack[PredictionCreateParamsWithoutVersion],
-) -> PredictionOutput | FileOutput | Iterable[FileOutput] | Dict[str, FileOutput]:
+) -> object | FileOutput | Iterable[FileOutput] | Dict[str, FileOutput]:
     from ._files import transform_output
 
     if is_given(wait) and "prefer" in params:
@@ -91,7 +91,7 @@ def run(
 
     # Return an iterator for the completed prediction when needed.
     if version and _has_output_iterator_array_type(version) and prediction.output is not None:
-        return (transform_output(chunk, client) for chunk in prediction.output)
+        return (transform_output(chunk, client) for chunk in prediction.output) # type: ignore
 
     if use_file_output:
         return transform_output(prediction.output, client)  # type: ignore[no-any-return]
@@ -107,7 +107,7 @@ async def async_run(
     wait: Union[int, bool, NotGiven] = NOT_GIVEN,
     use_file_output: Optional[bool] = True,
     **params: Unpack[PredictionCreateParamsWithoutVersion],
-) -> PredictionOutput | FileOutput | Iterable[FileOutput] | Dict[str, FileOutput]:
+) -> object | FileOutput | Iterable[FileOutput] | Dict[str, FileOutput]:
     from ._files import transform_output
 
     if is_given(wait) and "prefer" in params:
@@ -176,7 +176,7 @@ async def async_run(
 
     # Return an iterator for completed output if the model has an output iterator array type.
     if version and _has_output_iterator_array_type(version) and prediction.output is not None:
-        return (transform_output(chunk, client) async for chunk in _make_async_iterator(prediction.output))
+        return (transform_output(chunk, client) async for chunk in _make_async_iterator(prediction.output)) # type: ignore
     if use_file_output:
         return transform_output(prediction.output, client)  # type: ignore[no-any-return]
 
@@ -207,7 +207,7 @@ def output_iterator(prediction: Prediction, client: Replicate) -> Iterator[Any]:
         yield from new_output
         previous_output = output
         time.sleep(client.poll_interval)
-        prediction = client.predictions.get(prediction.id)
+        prediction = client.predictions.get(prediction_id=prediction.id)
 
     if prediction.status == "failed":
         raise ModelError(prediction=prediction)

src/replicate/resources/predictions.py

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