genblaze-replicate 0.2.0__tar.gz → 0.2.2__tar.gz

{genblaze_replicate-0.2.0 → genblaze_replicate-0.2.2}/PKG-INFO

@@ -1,10 +1,12 @@
 Metadata-Version: 2.4
 Name: genblaze-replicate
-Version: 0.2.0
+Version: 0.2.2
 Summary: Replicate provider adapter for genblaze
+Project-URL: Homepage, https://github.com/backblaze-labs/genblaze
 Project-URL: Documentation, https://github.com/backblaze-labs/genblaze
 Project-URL: Repository, https://github.com/backblaze-labs/genblaze
 Project-URL: Issues, https://github.com/backblaze-labs/genblaze/issues
+Author-email: Jeronimo De Leon <jdeleon@backblaze.com>
 License-Expression: MIT
 Classifier: Development Status :: 3 - Alpha
 Classifier: License :: OSI Approved :: MIT License

{genblaze_replicate-0.2.0 → genblaze_replicate-0.2.2}/genblaze_replicate/provider.py

@@ -29,6 +29,7 @@ from genblaze_core.providers import (
     route_by_media_type,
     validate_asset_url,
 )
+from genblaze_core.providers.retry import retry_after_from_response
 from genblaze_core.runnable.config import RunnableConfig
 
 from ._errors import map_replicate_error
@@ -123,6 +124,7 @@ class ReplicateProvider(BaseProvider):
             raise ProviderError(
                 f"Replicate submit failed: {exc}",
                 error_code=map_replicate_error(exc),
+                retry_after=retry_after_from_response(exc),
             ) from exc
 
     def poll(self, prediction_id: Any, config: RunnableConfig | None = None) -> bool:
@@ -137,6 +139,7 @@ class ReplicateProvider(BaseProvider):
             raise ProviderError(
                 f"Replicate poll failed: {exc}",
                 error_code=map_replicate_error(exc),
+                retry_after=retry_after_from_response(exc),
             ) from exc
 
     def fetch_output(self, prediction_id: Any, step: Step) -> Step:
@@ -176,14 +179,35 @@ class ReplicateProvider(BaseProvider):
                     error_code=ProviderErrorCode.UNKNOWN,
                 )
 
-            output = prediction.output
-            if isinstance(output, str):
-                output = [output]
-            elif output is None:
-                output = []
+            # Replicate output shapes vary by model: str (single URL), list[str]
+            # (multi-asset), dict[str, str | list] (e.g. {"video": url,
+            # "subtitles": url} from text-to-video models with side-channels),
+            # or None (no output). Normalize to list[str].
+            raw_output = prediction.output
+            urls: list[str]
+            if raw_output is None:
+                urls = []
+            elif isinstance(raw_output, str):
+                urls = [raw_output]
+            elif isinstance(raw_output, list):
+                # Nested lists happen on batch-output models; flatten one level.
+                urls = []
+                for item in raw_output:
+                    if isinstance(item, str):
+                        urls.append(item)
+                    elif isinstance(item, list):
+                        urls.extend(str(u) for u in item if isinstance(u, (str, bytes)))
+            elif isinstance(raw_output, dict):
+                # Multi-channel outputs: keep only URL-shaped string values.
+                urls = [str(v) for v in raw_output.values() if isinstance(v, str)]
+            else:
+                raise ProviderError(
+                    f"Unexpected Replicate output shape "
+                    f"({type(raw_output).__name__}): {raw_output!r}",
+                    error_code=ProviderErrorCode.SERVER_ERROR,
+                )
 
-            for url in output:
-                url_str = str(url)
+            for url_str in urls:
                 validate_asset_url(url_str)
                 path = urlparse(url_str).path
                 mime, _ = mimetypes.guess_type(path)
@@ -199,4 +223,5 @@ class ReplicateProvider(BaseProvider):
             raise ProviderError(
                 f"Replicate fetch_output failed: {exc}",
                 error_code=map_replicate_error(exc),
+                retry_after=retry_after_from_response(exc),
             ) from exc

{genblaze_replicate-0.2.0 → genblaze_replicate-0.2.2}/pyproject.toml

@@ -4,8 +4,9 @@ build-backend = "hatchling.build"
 
 [project]
 name = "genblaze-replicate"
-version = "0.2.0"
+version = "0.2.2"
 description = "Replicate provider adapter for genblaze"
+authors = [{name = "Jeronimo De Leon", email = "jdeleon@backblaze.com"}]
 readme = "README.md"
 requires-python = ">=3.11"
 license = "MIT"
@@ -20,6 +21,7 @@ dependencies = [
 ]
 
 [project.urls]
+Homepage = "https://github.com/backblaze-labs/genblaze"
 Documentation = "https://github.com/backblaze-labs/genblaze"
 Repository = "https://github.com/backblaze-labs/genblaze"
 Issues = "https://github.com/backblaze-labs/genblaze/issues"

{genblaze_replicate-0.2.0 → genblaze_replicate-0.2.2}/tests/test_replicate_provider.py

@@ -185,6 +185,54 @@ def test_fetch_output_success_with_none():
     assert len(result.assets) == 0
 
 
+def test_fetch_output_success_with_dict():
+    """Some Replicate models return dict outputs like {'video': url, 'audio': url}."""
+    provider = ReplicateProvider(api_token="test-token")
+    mock_client = MagicMock()
+    mock_client.predictions.get.return_value = FakePrediction(
+        output={
+            "video": "https://example.com/video.mp4",
+            "subtitles": "https://example.com/subs.vtt",
+        }
+    )
+    provider._client = mock_client
+
+    step = _make_step()
+    result = provider.fetch_output("pred-abc123", step)
+
+    urls = {a.url for a in result.assets}
+    assert urls == {"https://example.com/video.mp4", "https://example.com/subs.vtt"}
+
+
+def test_fetch_output_success_with_nested_list():
+    """Batch-output models return list[list[str]]; flatten one level."""
+    provider = ReplicateProvider(api_token="test-token")
+    mock_client = MagicMock()
+    mock_client.predictions.get.return_value = FakePrediction(
+        output=[["https://example.com/a.png", "https://example.com/b.png"]]
+    )
+    provider._client = mock_client
+
+    step = _make_step()
+    result = provider.fetch_output("pred-abc123", step)
+
+    assert len(result.assets) == 2
+
+
+def test_fetch_output_unknown_shape_raises_clear_error():
+    """Unknown output shapes raise ProviderError with a specific message."""
+    from genblaze_core.exceptions import ProviderError
+
+    provider = ReplicateProvider(api_token="test-token")
+    mock_client = MagicMock()
+    mock_client.predictions.get.return_value = FakePrediction(output=42)
+    provider._client = mock_client
+
+    step = _make_step()
+    with pytest.raises(ProviderError, match="Unexpected Replicate output shape"):
+        provider.fetch_output("pred-abc123", step)
+
+
 def test_fetch_output_failed():
     from genblaze_core.exceptions import ProviderError