Update stream interface to always use FileOutput

replicate · Oct 9, 2024 · 57e7255 · 57e7255
1 parent 9921f4c
commit 57e7255
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 10 deletions.
diff --git a/replicate/client.py b/replicate/client.py
@@ -164,7 +164,8 @@ def run(
         self,
         ref: str,
         input: Optional[Dict[str, Any]] = None,
-        use_file_output: Optional[bool] = None,
+        *,
+        use_file_output: Optional[bool] = True,
         **params: Unpack["Predictions.CreatePredictionParams"],
     ) -> Union[Any, Iterator[Any]]:  # noqa: ANN401
         """
@@ -177,7 +178,8 @@ async def async_run(
         self,
         ref: str,
         input: Optional[Dict[str, Any]] = None,
-        use_file_output: Optional[bool] = None,
+        *,
+        use_file_output: Optional[bool] = True,
         **params: Unpack["Predictions.CreatePredictionParams"],
     ) -> Union[Any, AsyncIterator[Any]]:  # noqa: ANN401
         """
@@ -191,28 +193,30 @@ async def async_run(
     def stream(
         self,
         ref: str,
+        *,
         input: Optional[Dict[str, Any]] = None,
-        use_file_output: Optional[bool] = None,
+        use_file_output: Optional[bool] = True,
         **params: Unpack["Predictions.CreatePredictionParams"],
     ) -> Iterator["ServerSentEvent"]:
         """
         Stream a model's output.
         """
 
-        return stream(self, ref, input, use_file_output, **params)
+        return stream(self, ref, input, use_file_output=use_file_output, **params)
 
     async def async_stream(
         self,
         ref: str,
         input: Optional[Dict[str, Any]] = None,
-        use_file_output: Optional[bool] = None,
+        *,
+        use_file_output: Optional[bool] = True,
         **params: Unpack["Predictions.CreatePredictionParams"],
     ) -> AsyncIterator["ServerSentEvent"]:
         """
         Stream a model's output asynchronously.
         """
 
-        return async_stream(self, ref, input, use_file_output, **params)
+        return async_stream(self, ref, input, use_file_output=use_file_output, **params)
 
 
 # Adapted from https://github.com/encode/httpx/issues/108#issuecomment-1132753155

diff --git a/replicate/stream.py b/replicate/stream.py
@@ -71,11 +71,12 @@ def __init__(
         self,
         client: "Client",
         response: "httpx.Response",
-        use_file_output: Optional[bool] = None,
+        *,
+        use_file_output: Optional[bool] = True,
     ) -> None:
         self.client = client
         self.response = response
-        self.use_file_output = use_file_output or False
+        self.use_file_output = use_file_output or True
         content_type, _, _ = response.headers["content-type"].partition(";")
         if content_type != "text/event-stream":
             raise ValueError(
@@ -193,7 +194,8 @@ def stream(
     client: "Client",
     ref: Union["Model", "Version", "ModelVersionIdentifier", str],
     input: Optional[Dict[str, Any]] = None,
-    use_file_output: Optional[bool] = None,
+    *,
+    use_file_output: Optional[bool] = True,
     **params: Unpack["Predictions.CreatePredictionParams"],
 ) -> Iterator[ServerSentEvent]:
     """
@@ -234,7 +236,8 @@ async def async_stream(
     client: "Client",
     ref: Union["Model", "Version", "ModelVersionIdentifier", str],
     input: Optional[Dict[str, Any]] = None,
-    use_file_output: Optional[bool] = None,
+    *,
+    use_file_output: Optional[bool] = True,
     **params: Unpack["Predictions.CreatePredictionParams"],
 ) -> AsyncIterator[ServerSentEvent]:
     """