Ensure OutputIterator objects are converted into strings

Author: aron

replicate/use.py

@@ -1,8 +1,8 @@
 # TODO
 # - [x] Support downloading files and conversion into Path when schema is URL
 # - [x] Support list outputs
-# - [ ] Support iterator outputs
-# - [ ] Support helpers for working with ContatenateIterator
+# - [x] Support iterator outputs
+# - [x] Support helpers for working with ContatenateIterator
 # - [ ] Support reusing output URL when passing to new method
 # - [ ] Support lazy downloading of files into Path
 # - [ ] Support text streaming
@@ -187,12 +187,12 @@ class OutputIterator:
     An iterator wrapper that handles both regular iteration and string conversion.
     """
 
-    def __init__(self, iterator_factory, schema: dict, is_concatenate: bool):
+    def __init__(self, iterator_factory, schema: dict, *, is_concatenate: bool) -> None:
         self.iterator_factory = iterator_factory
         self.schema = schema
         self.is_concatenate = is_concatenate
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Any]:
         """Iterate over output items."""
         for chunk in self.iterator_factory():
             if self.is_concatenate:
@@ -230,7 +230,9 @@ def wait(self) -> Union[Any, Iterator[Any]]:
         if _has_iterator_output_type(self.schema):
             is_concatenate = _has_concatenate_iterator_output_type(self.schema)
             return OutputIterator(
-                lambda: self.prediction.output_iterator(), self.schema, is_concatenate
+                lambda: self.prediction.output_iterator(),
+                self.schema,
+                is_concatenate=is_concatenate,
             )
 
         # Process output for file downloads based on schema
@@ -299,15 +301,23 @@ def create(self, **inputs: Dict[str, Any]) -> Run:
         """
         Start a prediction with the specified inputs.
         """
+        # Process inputs to convert concatenate OutputIterators to strings
+        processed_inputs = {}
+        for key, value in inputs.items():
+            if isinstance(value, OutputIterator) and value.is_concatenate:
+                processed_inputs[key] = str(value)
+            else:
+                processed_inputs[key] = value
+
         version = self._version
 
         if version:
             prediction = self._client().predictions.create(
-                version=version, input=inputs
+                version=version, input=processed_inputs
             )
         else:
             prediction = self._client().models.predictions.create(
-                model=self._model, input=inputs
+                model=self._model, input=processed_inputs
             )
 
         return Run(prediction, self.openapi_schema)

tests/test_use.py

@@ -1,3 +1,4 @@
+import json
 import os
 from pathlib import Path
 
@@ -365,6 +366,44 @@ async def test_use_concatenate_iterator_output(use_async_client):
     output_list = list(output)
     assert output_list == ["Hello", " ", "world", "!"]
 
+    # Test that concatenate OutputIterators are stringified when passed to create()
+    # Set up a mock for the prediction creation to capture the request
+    request_body = None
+
+    def capture_request(request):
+        nonlocal request_body
+        request_body = request.read()
+        return httpx.Response(
+            201,
+            json={
+                "id": "pred456",
+                "model": "acme/hotdog-detector",
+                "version": "xyz123",
+                "urls": {
+                    "get": "https://api.replicate.com/v1/predictions/pred456",
+                    "cancel": "https://api.replicate.com/v1/predictions/pred456/cancel",
+                },
+                "created_at": "2024-01-01T00:00:00Z",
+                "source": "api",
+                "status": "processing",
+                "input": {"text_input": "Hello world!"},
+                "output": None,
+                "error": None,
+                "logs": "",
+            },
+        )
+
+    respx.post("https://api.replicate.com/v1/predictions").mock(
+        side_effect=capture_request
+    )
+
+    # Pass the OutputIterator as input to create()
+    run = hotdog_detector.create(text_input=output)
+
+    # Verify the request body contains the stringified version
+    parsed_body = json.loads(request_body)
+    assert parsed_body["input"]["text_input"] == "Hello world!"
+
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("use_async_client", [False])