genblaze-replicate 0.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,76 @@
1
+ # Python
2
+ __pycache__/
3
+ *.py[cod]
4
+ *$py.class
5
+ *.so
6
+
7
+ # Distribution / packaging
8
+ dist/
9
+ build/
10
+ *.egg-info/
11
+ *.egg
12
+ wheels/
13
+ /MANIFEST
14
+
15
+ # Virtual environments
16
+ .venv/
17
+ venv/
18
+ env/
19
+ ENV/
20
+
21
+ # IDE / editors
22
+ .vscode/
23
+ .idea/
24
+ *.swp
25
+ *.swo
26
+ *~
27
+ .project
28
+ .settings/
29
+
30
+ # Testing / coverage
31
+ .pytest_cache/
32
+ .coverage
33
+ .coverage.*
34
+ htmlcov/
35
+ coverage.xml
36
+ *.cover
37
+
38
+ # Type checking
39
+ .mypy_cache/
40
+ .pytype/
41
+ .pyre/
42
+
43
+ # Linting
44
+ .ruff_cache/
45
+
46
+ # OS files
47
+ .DS_Store
48
+ Thumbs.db
49
+ ehthumbs.db
50
+
51
+ # Environment / secrets
52
+ .env
53
+ .env.*
54
+ !.env.example
55
+ *.pem
56
+ *.key
57
+ credentials.json
58
+
59
+ # Jupyter
60
+ .ipynb_checkpoints/
61
+
62
+ # MkDocs build output
63
+ site/
64
+
65
+ # Claude Code — local/ephemeral
66
+ .claude/settings.local.json
67
+ .claude/worktrees/
68
+ CLAUDE.local.md
69
+
70
+ # Hypothesis test cache
71
+ .hypothesis/
72
+
73
+ # Temp files
74
+ *.tmp
75
+ *.bak
76
+ *.log
@@ -0,0 +1,128 @@
1
+ Metadata-Version: 2.4
2
+ Name: genblaze-replicate
3
+ Version: 0.2.0
4
+ Summary: Replicate provider adapter for genblaze
5
+ Project-URL: Documentation, https://github.com/backblaze-labs/genblaze
6
+ Project-URL: Repository, https://github.com/backblaze-labs/genblaze
7
+ Project-URL: Issues, https://github.com/backblaze-labs/genblaze/issues
8
+ License-Expression: MIT
9
+ Classifier: Development Status :: 3 - Alpha
10
+ Classifier: License :: OSI Approved :: MIT License
11
+ Classifier: Typing :: Typed
12
+ Requires-Python: >=3.11
13
+ Requires-Dist: genblaze-core<0.3,>=0.2.0
14
+ Requires-Dist: replicate>=0.25
15
+ Provides-Extra: dev
16
+ Requires-Dist: pytest>=7.0; extra == 'dev'
17
+ Description-Content-Type: text/markdown
18
+
19
+ <!-- last_verified: 2026-04-22 -->
20
+ # genblaze-replicate
21
+
22
+ **[Replicate](https://replicate.com) provider adapter for [genblaze](https://github.com/backblaze-labs/genblaze) — run any Replicate-hosted image, video, or audio model through a unified AI pipeline with SHA-256 provenance.**
23
+
24
+ `genblaze-replicate` wraps Replicate's hosted model catalog (FLUX, SDXL, Stable Diffusion, video models, music models — anything on Replicate) as a genblaze provider. Compose Replicate calls into multi-step pipelines, persist outputs to [Backblaze B2](https://www.backblaze.com/cloud-storage?utm_source=github&utm_medium=referral&utm_campaign=ai_artifacts&utm_content=genblaze) or any S3-compatible store, and emit a tamper-evident provenance manifest for every run.
25
+
26
+ ## Why genblaze-replicate
27
+
28
+ - **Any Replicate model, one API** — Pass the `owner/model` slug; genblaze handles submission, polling, asset download, and manifest capture.
29
+ - **Provenance by default** — SHA-256 hash of every output, canonical manifest with prompt + params + timestamps.
30
+ - **Composable** — Chain Replicate outputs into downstream steps (image → video, upscale, transcode).
31
+ - **Production-ready** — Timeouts, retries, moderation hooks, step caching, OpenTelemetry tracing.
32
+ - **Durable storage** — Drop the `genblaze-s3` sink in for B2/S3/R2/MinIO persistence of assets + manifests.
33
+
34
+ ## Models
35
+
36
+ Supports any model hosted on Replicate — including:
37
+
38
+ - **Image** — `black-forest-labs/flux-schnell`, `black-forest-labs/flux-dev`, `stability-ai/sdxl`, `stability-ai/stable-diffusion-3`
39
+ - **Video** — text-to-video and image-to-video models on Replicate
40
+ - **Audio** — music and speech models on Replicate
41
+
42
+ Use the full `owner/model` slug from https://replicate.com/explore.
43
+
44
+ ## Install
45
+
46
+ ```bash
47
+ pip install genblaze-replicate
48
+ ```
49
+
50
+ Registers the `replicate` provider via entry points; [`genblaze-core`](https://pypi.org/project/genblaze-core/) discovers it automatically.
51
+
52
+ ## Quickstart — FLUX Schnell (text-to-image)
53
+
54
+ ```bash
55
+ pip install genblaze-core genblaze-replicate
56
+ export REPLICATE_API_TOKEN="r8_..."
57
+ ```
58
+
59
+ ```python
60
+ from genblaze_core import Modality, Pipeline
61
+ from genblaze_replicate import ReplicateProvider
62
+
63
+ run, manifest = (
64
+ Pipeline("flux-demo")
65
+ .step(
66
+ ReplicateProvider(),
67
+ model="black-forest-labs/flux-schnell",
68
+ prompt="a photorealistic golden retriever puppy in a field of wildflowers, "
69
+ "golden hour, shallow depth of field",
70
+ modality=Modality.IMAGE,
71
+ num_outputs=1,
72
+ aspect_ratio="1:1",
73
+ )
74
+ .run()
75
+ )
76
+
77
+ print(run.steps[0].assets[0].url)
78
+ print(manifest.canonical_hash)
79
+ assert manifest.verify()
80
+ ```
81
+
82
+ ## Quickstart — persist to Backblaze B2
83
+
84
+ ```bash
85
+ pip install genblaze-core genblaze-replicate genblaze-s3
86
+ export REPLICATE_API_TOKEN="r8_..."
87
+ export B2_KEY_ID="..." B2_APP_KEY="..."
88
+ ```
89
+
90
+ ```python
91
+ from genblaze_core import KeyStrategy, ObjectStorageSink, Pipeline
92
+ from genblaze_replicate import ReplicateProvider
93
+ from genblaze_s3 import S3StorageBackend
94
+
95
+ storage = ObjectStorageSink(
96
+ S3StorageBackend.for_backblaze("my-bucket"),
97
+ key_strategy=KeyStrategy.CONTENT_ADDRESSABLE,
98
+ )
99
+
100
+ result = (
101
+ Pipeline("flux-b2")
102
+ .step(ReplicateProvider(), model="black-forest-labs/flux-schnell",
103
+ prompt="cyberpunk tokyo street at night, neon reflections on wet pavement")
104
+ .run(sink=storage, timeout=120)
105
+ )
106
+
107
+ print(result.run.steps[0].assets[0].url) # durable B2 URL
108
+ ```
109
+
110
+ ## Credentials
111
+
112
+ | Env var | Where to get it |
113
+ |---|---|
114
+ | `REPLICATE_API_TOKEN` | https://replicate.com/account/api-tokens |
115
+
116
+ ## Documentation
117
+
118
+ - **Main repo**: https://github.com/backblaze-labs/genblaze
119
+ - **Runnable example**: [`replicate_flux_pipeline.py`](https://github.com/backblaze-labs/genblaze/blob/main/examples/replicate_flux_pipeline.py)
120
+
121
+ ## Related packages
122
+
123
+ - [`genblaze-core`](https://pypi.org/project/genblaze-core/) — the pipeline SDK
124
+ - [`genblaze-s3`](https://pypi.org/project/genblaze-s3/) — durable storage on [Backblaze B2](https://www.backblaze.com/cloud-storage?utm_source=github&utm_medium=referral&utm_campaign=ai_artifacts&utm_content=genblaze) and other S3-compatible backends
125
+
126
+ ## License
127
+
128
+ MIT
@@ -0,0 +1,110 @@
1
+ <!-- last_verified: 2026-04-22 -->
2
+ # genblaze-replicate
3
+
4
+ **[Replicate](https://replicate.com) provider adapter for [genblaze](https://github.com/backblaze-labs/genblaze) — run any Replicate-hosted image, video, or audio model through a unified AI pipeline with SHA-256 provenance.**
5
+
6
+ `genblaze-replicate` wraps Replicate's hosted model catalog (FLUX, SDXL, Stable Diffusion, video models, music models — anything on Replicate) as a genblaze provider. Compose Replicate calls into multi-step pipelines, persist outputs to [Backblaze B2](https://www.backblaze.com/cloud-storage?utm_source=github&utm_medium=referral&utm_campaign=ai_artifacts&utm_content=genblaze) or any S3-compatible store, and emit a tamper-evident provenance manifest for every run.
7
+
8
+ ## Why genblaze-replicate
9
+
10
+ - **Any Replicate model, one API** — Pass the `owner/model` slug; genblaze handles submission, polling, asset download, and manifest capture.
11
+ - **Provenance by default** — SHA-256 hash of every output, canonical manifest with prompt + params + timestamps.
12
+ - **Composable** — Chain Replicate outputs into downstream steps (image → video, upscale, transcode).
13
+ - **Production-ready** — Timeouts, retries, moderation hooks, step caching, OpenTelemetry tracing.
14
+ - **Durable storage** — Drop the `genblaze-s3` sink in for B2/S3/R2/MinIO persistence of assets + manifests.
15
+
16
+ ## Models
17
+
18
+ Supports any model hosted on Replicate — including:
19
+
20
+ - **Image** — `black-forest-labs/flux-schnell`, `black-forest-labs/flux-dev`, `stability-ai/sdxl`, `stability-ai/stable-diffusion-3`
21
+ - **Video** — text-to-video and image-to-video models on Replicate
22
+ - **Audio** — music and speech models on Replicate
23
+
24
+ Use the full `owner/model` slug from https://replicate.com/explore.
25
+
26
+ ## Install
27
+
28
+ ```bash
29
+ pip install genblaze-replicate
30
+ ```
31
+
32
+ Registers the `replicate` provider via entry points; [`genblaze-core`](https://pypi.org/project/genblaze-core/) discovers it automatically.
33
+
34
+ ## Quickstart — FLUX Schnell (text-to-image)
35
+
36
+ ```bash
37
+ pip install genblaze-core genblaze-replicate
38
+ export REPLICATE_API_TOKEN="r8_..."
39
+ ```
40
+
41
+ ```python
42
+ from genblaze_core import Modality, Pipeline
43
+ from genblaze_replicate import ReplicateProvider
44
+
45
+ run, manifest = (
46
+ Pipeline("flux-demo")
47
+ .step(
48
+ ReplicateProvider(),
49
+ model="black-forest-labs/flux-schnell",
50
+ prompt="a photorealistic golden retriever puppy in a field of wildflowers, "
51
+ "golden hour, shallow depth of field",
52
+ modality=Modality.IMAGE,
53
+ num_outputs=1,
54
+ aspect_ratio="1:1",
55
+ )
56
+ .run()
57
+ )
58
+
59
+ print(run.steps[0].assets[0].url)
60
+ print(manifest.canonical_hash)
61
+ assert manifest.verify()
62
+ ```
63
+
64
+ ## Quickstart — persist to Backblaze B2
65
+
66
+ ```bash
67
+ pip install genblaze-core genblaze-replicate genblaze-s3
68
+ export REPLICATE_API_TOKEN="r8_..."
69
+ export B2_KEY_ID="..." B2_APP_KEY="..."
70
+ ```
71
+
72
+ ```python
73
+ from genblaze_core import KeyStrategy, ObjectStorageSink, Pipeline
74
+ from genblaze_replicate import ReplicateProvider
75
+ from genblaze_s3 import S3StorageBackend
76
+
77
+ storage = ObjectStorageSink(
78
+ S3StorageBackend.for_backblaze("my-bucket"),
79
+ key_strategy=KeyStrategy.CONTENT_ADDRESSABLE,
80
+ )
81
+
82
+ result = (
83
+ Pipeline("flux-b2")
84
+ .step(ReplicateProvider(), model="black-forest-labs/flux-schnell",
85
+ prompt="cyberpunk tokyo street at night, neon reflections on wet pavement")
86
+ .run(sink=storage, timeout=120)
87
+ )
88
+
89
+ print(result.run.steps[0].assets[0].url) # durable B2 URL
90
+ ```
91
+
92
+ ## Credentials
93
+
94
+ | Env var | Where to get it |
95
+ |---|---|
96
+ | `REPLICATE_API_TOKEN` | https://replicate.com/account/api-tokens |
97
+
98
+ ## Documentation
99
+
100
+ - **Main repo**: https://github.com/backblaze-labs/genblaze
101
+ - **Runnable example**: [`replicate_flux_pipeline.py`](https://github.com/backblaze-labs/genblaze/blob/main/examples/replicate_flux_pipeline.py)
102
+
103
+ ## Related packages
104
+
105
+ - [`genblaze-core`](https://pypi.org/project/genblaze-core/) — the pipeline SDK
106
+ - [`genblaze-s3`](https://pypi.org/project/genblaze-s3/) — durable storage on [Backblaze B2](https://www.backblaze.com/cloud-storage?utm_source=github&utm_medium=referral&utm_campaign=ai_artifacts&utm_content=genblaze) and other S3-compatible backends
107
+
108
+ ## License
109
+
110
+ MIT
@@ -0,0 +1,5 @@
1
+ """Replicate provider adapter for genblaze."""
2
+
3
+ from genblaze_replicate.provider import ReplicateProvider
4
+
5
+ __all__ = ["ReplicateProvider"]
@@ -0,0 +1,21 @@
1
+ """Shared Replicate error mapping — used by provider.py."""
2
+
3
+ from genblaze_core.models.enums import ProviderErrorCode
4
+
5
+
6
+ def map_replicate_error(exc_or_msg: Exception | str) -> ProviderErrorCode:
7
+ """Map a Replicate API exception (or error string) to a ProviderErrorCode."""
8
+ msg = (str(exc_or_msg)).lower()
9
+ if "rate" in msg or "429" in msg:
10
+ return ProviderErrorCode.RATE_LIMIT
11
+ if "auth" in msg or "401" in msg or "403" in msg or "token" in msg:
12
+ return ProviderErrorCode.AUTH_FAILURE
13
+ if "invalid" in msg or "400" in msg or "validation" in msg:
14
+ return ProviderErrorCode.INVALID_INPUT
15
+ if "timeout" in msg or "timed out" in msg:
16
+ return ProviderErrorCode.TIMEOUT
17
+ if "model" in msg and ("not found" in msg or "error" in msg or "does not exist" in msg):
18
+ return ProviderErrorCode.MODEL_ERROR
19
+ if "500" in msg or "502" in msg or "503" in msg:
20
+ return ProviderErrorCode.SERVER_ERROR
21
+ return ProviderErrorCode.UNKNOWN
@@ -0,0 +1,202 @@
1
+ """ReplicateProvider — adapter for the Replicate API.
2
+
3
+ Models are free-form strings (any Replicate model slug). Default pricing is
4
+ compute-time-based (``predict_time`` × rate). Override per-model or wildcard::
5
+
6
+ from genblaze_replicate import ReplicateProvider
7
+ reg = ReplicateProvider.models_default().fork()
8
+ reg.register_pricing("owner/my-model", per_response_metric(...))
9
+ provider = ReplicateProvider(models=reg)
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ import mimetypes
15
+ from typing import Any
16
+ from urllib.parse import urlparse
17
+
18
+ from genblaze_core.exceptions import ProviderError
19
+ from genblaze_core.models.asset import Asset
20
+ from genblaze_core.models.enums import Modality, ProviderErrorCode
21
+ from genblaze_core.models.step import Step
22
+ from genblaze_core.providers import (
23
+ BaseProvider,
24
+ ModelRegistry,
25
+ ModelSpec,
26
+ PricingContext,
27
+ ProviderCapabilities,
28
+ per_response_metric,
29
+ route_by_media_type,
30
+ validate_asset_url,
31
+ )
32
+ from genblaze_core.runnable.config import RunnableConfig
33
+
34
+ from ._errors import map_replicate_error
35
+
36
+ # Default per-second GPU-compute rate (Nvidia A40/T4 tier).
37
+ _COST_PER_SEC = 0.000225
38
+
39
+
40
+ def _compute_time_cost(ctx: PricingContext) -> float | None:
41
+ payload = ctx.provider_payload.get("replicate") if ctx.provider_payload else None
42
+ if not isinstance(payload, dict):
43
+ return None
44
+ predict_time = payload.get("predict_time")
45
+ if predict_time is None:
46
+ return None
47
+ try:
48
+ return float(predict_time) * _COST_PER_SEC
49
+ except (TypeError, ValueError):
50
+ return None
51
+
52
+
53
+ # Replicate has no enumerable model list; fallback spec applies to every model.
54
+ _FALLBACK_SPEC = ModelSpec(
55
+ model_id="*",
56
+ pricing=per_response_metric(_compute_time_cost),
57
+ input_mapping=route_by_media_type({"image": "image", "video": "video", "audio": "audio"}),
58
+ )
59
+
60
+
61
+ class ReplicateProvider(BaseProvider):
62
+ """Provider adapter for Replicate (replicate.com)."""
63
+
64
+ name = "replicate"
65
+
66
+ @classmethod
67
+ def create_registry(cls) -> ModelRegistry:
68
+ # No enumerable defaults — every model uses the fallback spec.
69
+ return ModelRegistry(defaults={}, fallback=_FALLBACK_SPEC)
70
+
71
+ def get_capabilities(self) -> ProviderCapabilities:
72
+ """Replicate: multi-modal generation depending on selected model."""
73
+ return ProviderCapabilities(
74
+ supported_modalities=[Modality.IMAGE, Modality.VIDEO, Modality.AUDIO],
75
+ supported_inputs=["text", "image", "video", "audio"],
76
+ accepts_chain_input=True,
77
+ )
78
+
79
+ def __init__(
80
+ self,
81
+ api_token: str | None = None,
82
+ poll_interval: float = 1.0,
83
+ http_timeout: float = 30.0,
84
+ *,
85
+ models: ModelRegistry | None = None,
86
+ ):
87
+ super().__init__(models=models)
88
+ self.poll_interval = poll_interval
89
+ self._client: Any = None
90
+ self._api_token = api_token
91
+ self._http_timeout = http_timeout
92
+
93
+ def _get_client(self):
94
+ if self._client is None:
95
+ try:
96
+ import httpx
97
+ import replicate
98
+
99
+ timeout = httpx.Timeout(self._http_timeout, connect=10.0)
100
+ if self._api_token:
101
+ self._client = replicate.Client(
102
+ api_token=self._api_token, # noqa: S106
103
+ timeout=timeout,
104
+ )
105
+ else:
106
+ self._client = replicate.Client(timeout=timeout)
107
+ except ImportError as exc:
108
+ raise ProviderError(
109
+ "replicate package not installed. Run: pip install replicate"
110
+ ) from exc
111
+ return self._client
112
+
113
+ def submit(self, step: Step, config: RunnableConfig | None = None) -> Any:
114
+ client = self._get_client()
115
+ try:
116
+ input_params = self.prepare_payload(step)
117
+ prediction = client.predictions.create(
118
+ model=step.model,
119
+ input=input_params,
120
+ )
121
+ return prediction.id
122
+ except Exception as exc:
123
+ raise ProviderError(
124
+ f"Replicate submit failed: {exc}",
125
+ error_code=map_replicate_error(exc),
126
+ ) from exc
127
+
128
+ def poll(self, prediction_id: Any, config: RunnableConfig | None = None) -> bool:
129
+ client = self._get_client()
130
+ try:
131
+ prediction = client.predictions.get(prediction_id)
132
+ if prediction.status in ("succeeded", "failed", "canceled"):
133
+ self._cache_poll_result(prediction_id, prediction)
134
+ return True
135
+ return False
136
+ except Exception as exc:
137
+ raise ProviderError(
138
+ f"Replicate poll failed: {exc}",
139
+ error_code=map_replicate_error(exc),
140
+ ) from exc
141
+
142
+ def fetch_output(self, prediction_id: Any, step: Step) -> Step:
143
+ client = self._get_client()
144
+ try:
145
+ prediction = self._get_cached_poll_result(prediction_id)
146
+ if prediction is None:
147
+ prediction = client.predictions.get(prediction_id)
148
+
149
+ # Capture predict_time so registry pricing can read it.
150
+ metrics = getattr(prediction, "metrics", None)
151
+ predict_time = getattr(metrics, "predict_time", None) if metrics else None
152
+
153
+ step.provider_payload = {
154
+ "replicate": {
155
+ "prediction_id": prediction.id,
156
+ "model": prediction.model if hasattr(prediction, "model") else None,
157
+ "version": prediction.version if hasattr(prediction, "version") else None,
158
+ "status": prediction.status,
159
+ "created_at": str(prediction.created_at)
160
+ if hasattr(prediction, "created_at")
161
+ else None,
162
+ "predict_time": predict_time,
163
+ }
164
+ }
165
+
166
+ if prediction.status == "failed":
167
+ error_msg = prediction.error or "Unknown error"
168
+ raise ProviderError(
169
+ error_msg,
170
+ error_code=map_replicate_error(error_msg),
171
+ )
172
+
173
+ if prediction.status == "canceled":
174
+ raise ProviderError(
175
+ "Prediction was canceled",
176
+ error_code=ProviderErrorCode.UNKNOWN,
177
+ )
178
+
179
+ output = prediction.output
180
+ if isinstance(output, str):
181
+ output = [output]
182
+ elif output is None:
183
+ output = []
184
+
185
+ for url in output:
186
+ url_str = str(url)
187
+ validate_asset_url(url_str)
188
+ path = urlparse(url_str).path
189
+ mime, _ = mimetypes.guess_type(path)
190
+ if mime is None:
191
+ mime = f"{step.modality.value}/octet-stream"
192
+ step.assets.append(Asset(url=url_str, media_type=mime))
193
+
194
+ self._apply_registry_pricing(step)
195
+ return step
196
+ except ProviderError:
197
+ raise
198
+ except Exception as exc:
199
+ raise ProviderError(
200
+ f"Replicate fetch_output failed: {exc}",
201
+ error_code=map_replicate_error(exc),
202
+ ) from exc
File without changes
@@ -0,0 +1,39 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "genblaze-replicate"
7
+ version = "0.2.0"
8
+ description = "Replicate provider adapter for genblaze"
9
+ readme = "README.md"
10
+ requires-python = ">=3.11"
11
+ license = "MIT"
12
+ classifiers = [
13
+ "Development Status :: 3 - Alpha",
14
+ "License :: OSI Approved :: MIT License",
15
+ "Typing :: Typed",
16
+ ]
17
+ dependencies = [
18
+ "genblaze-core>=0.2.0,<0.3",
19
+ "replicate>=0.25",
20
+ ]
21
+
22
+ [project.urls]
23
+ Documentation = "https://github.com/backblaze-labs/genblaze"
24
+ Repository = "https://github.com/backblaze-labs/genblaze"
25
+ Issues = "https://github.com/backblaze-labs/genblaze/issues"
26
+
27
+ [project.optional-dependencies]
28
+ dev = [
29
+ "pytest>=7.0",
30
+ ]
31
+
32
+ [project.entry-points."genblaze.providers"]
33
+ replicate = "genblaze_replicate:ReplicateProvider"
34
+
35
+ [tool.hatch.build.targets.wheel]
36
+ packages = ["genblaze_replicate"]
37
+
38
+ [tool.pytest.ini_options]
39
+ testpaths = ["tests"]
File without changes
@@ -0,0 +1,333 @@
1
+ """Tests for ReplicateProvider with mocked Replicate API."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from dataclasses import dataclass, field
6
+ from typing import Any
7
+ from unittest.mock import MagicMock
8
+
9
+ import pytest
10
+ from genblaze_core.models.enums import ProviderErrorCode
11
+ from genblaze_core.models.step import Step
12
+ from genblaze_core.testing import ProviderComplianceTests
13
+ from genblaze_replicate._errors import map_replicate_error
14
+ from genblaze_replicate.provider import ReplicateProvider
15
+
16
+
17
+ @dataclass
18
+ class FakePrediction:
19
+ """Minimal mock of a Replicate prediction object."""
20
+
21
+ id: str = "pred-abc123"
22
+ status: str = "succeeded"
23
+ output: Any = field(default_factory=lambda: ["https://example.com/out.png"])
24
+ error: str | None = None
25
+ model: str = "test/model"
26
+ version: str | None = "v1"
27
+ created_at: str = "2026-01-01T00:00:00Z"
28
+
29
+
30
+ def _make_step(**kwargs) -> Step:
31
+ defaults = {"provider": "replicate", "model": "test/model", "prompt": "a cat"}
32
+ defaults.update(kwargs)
33
+ return Step(**defaults)
34
+
35
+
36
+ # --- map_replicate_error tests ---
37
+
38
+
39
+ def test_map_error_timeout():
40
+ assert map_replicate_error("Request timed out") == ProviderErrorCode.TIMEOUT
41
+
42
+
43
+ def test_map_error_rate_limit():
44
+ assert map_replicate_error("rate limit exceeded 429") == ProviderErrorCode.RATE_LIMIT
45
+
46
+
47
+ def test_map_error_auth():
48
+ assert map_replicate_error("401 unauthorized token") == ProviderErrorCode.AUTH_FAILURE
49
+
50
+
51
+ def test_map_error_invalid_input():
52
+ assert map_replicate_error("invalid input: bad format") == ProviderErrorCode.INVALID_INPUT
53
+
54
+
55
+ def test_map_error_model_not_found():
56
+ assert map_replicate_error("model not found") == ProviderErrorCode.MODEL_ERROR
57
+
58
+
59
+ def test_map_error_server_error():
60
+ assert map_replicate_error("500 internal server error") == ProviderErrorCode.SERVER_ERROR
61
+
62
+
63
+ def test_map_error_unknown():
64
+ assert map_replicate_error("something went wrong") == ProviderErrorCode.UNKNOWN
65
+
66
+
67
+ def test_map_error_accepts_exception():
68
+ """map_replicate_error should accept Exception objects too."""
69
+ assert map_replicate_error(RuntimeError("timed out")) == ProviderErrorCode.TIMEOUT
70
+
71
+
72
+ # --- submit tests ---
73
+
74
+
75
+ def test_submit_sends_prompt_and_params():
76
+ provider = ReplicateProvider(api_token="test-token")
77
+ mock_client = MagicMock()
78
+ mock_client.predictions.create.return_value = FakePrediction()
79
+ provider._client = mock_client
80
+
81
+ step = _make_step(params={"width": 1024})
82
+ result = provider.submit(step)
83
+
84
+ mock_client.predictions.create.assert_called_once_with(
85
+ model="test/model",
86
+ input={"width": 1024, "prompt": "a cat"},
87
+ )
88
+ assert result == "pred-abc123"
89
+
90
+
91
+ def test_submit_includes_negative_prompt():
92
+ provider = ReplicateProvider(api_token="test-token")
93
+ mock_client = MagicMock()
94
+ mock_client.predictions.create.return_value = FakePrediction()
95
+ provider._client = mock_client
96
+
97
+ step = _make_step(negative_prompt="blurry")
98
+ provider.submit(step)
99
+
100
+ call_input = mock_client.predictions.create.call_args.kwargs["input"]
101
+ assert call_input["negative_prompt"] == "blurry"
102
+
103
+
104
+ # --- poll tests ---
105
+
106
+
107
+ def test_poll_returns_true_on_success():
108
+ provider = ReplicateProvider(api_token="test-token", poll_interval=0.0)
109
+ mock_client = MagicMock()
110
+ mock_client.predictions.get.return_value = FakePrediction(status="succeeded")
111
+ provider._client = mock_client
112
+
113
+ assert provider.poll("pred-abc123") is True
114
+
115
+
116
+ def test_poll_returns_true_on_failed():
117
+ provider = ReplicateProvider(api_token="test-token", poll_interval=0.0)
118
+ mock_client = MagicMock()
119
+ mock_client.predictions.get.return_value = FakePrediction(status="failed")
120
+ provider._client = mock_client
121
+
122
+ assert provider.poll("pred-abc123") is True
123
+
124
+
125
+ def test_poll_returns_true_on_canceled():
126
+ provider = ReplicateProvider(api_token="test-token", poll_interval=0.0)
127
+ mock_client = MagicMock()
128
+ mock_client.predictions.get.return_value = FakePrediction(status="canceled")
129
+ provider._client = mock_client
130
+
131
+ assert provider.poll("pred-abc123") is True
132
+
133
+
134
+ def test_poll_returns_false_on_processing():
135
+ provider = ReplicateProvider(api_token="test-token", poll_interval=0.0)
136
+ mock_client = MagicMock()
137
+ mock_client.predictions.get.return_value = FakePrediction(status="processing")
138
+ provider._client = mock_client
139
+
140
+ assert provider.poll("pred-abc123") is False
141
+
142
+
143
+ # --- fetch_output tests ---
144
+
145
+
146
+ def test_fetch_output_success_with_list():
147
+ provider = ReplicateProvider(api_token="test-token")
148
+ mock_client = MagicMock()
149
+ mock_client.predictions.get.return_value = FakePrediction(
150
+ output=["https://example.com/a.png", "https://example.com/b.png"]
151
+ )
152
+ provider._client = mock_client
153
+
154
+ step = _make_step()
155
+ result = provider.fetch_output("pred-abc123", step)
156
+
157
+ assert len(result.assets) == 2
158
+ assert result.assets[0].url == "https://example.com/a.png"
159
+
160
+
161
+ def test_fetch_output_success_with_string():
162
+ provider = ReplicateProvider(api_token="test-token")
163
+ mock_client = MagicMock()
164
+ mock_client.predictions.get.return_value = FakePrediction(
165
+ output="https://example.com/single.png"
166
+ )
167
+ provider._client = mock_client
168
+
169
+ step = _make_step()
170
+ result = provider.fetch_output("pred-abc123", step)
171
+
172
+ assert len(result.assets) == 1
173
+ assert result.assets[0].url == "https://example.com/single.png"
174
+
175
+
176
+ def test_fetch_output_success_with_none():
177
+ provider = ReplicateProvider(api_token="test-token")
178
+ mock_client = MagicMock()
179
+ mock_client.predictions.get.return_value = FakePrediction(output=None)
180
+ provider._client = mock_client
181
+
182
+ step = _make_step()
183
+ result = provider.fetch_output("pred-abc123", step)
184
+
185
+ assert len(result.assets) == 0
186
+
187
+
188
+ def test_fetch_output_failed():
189
+ from genblaze_core.exceptions import ProviderError
190
+
191
+ provider = ReplicateProvider(api_token="test-token")
192
+ mock_client = MagicMock()
193
+ mock_client.predictions.get.return_value = FakePrediction(
194
+ status="failed", error="Model crashed"
195
+ )
196
+ provider._client = mock_client
197
+
198
+ step = _make_step()
199
+ with pytest.raises(ProviderError, match="Model crashed"):
200
+ provider.fetch_output("pred-abc123", step)
201
+
202
+
203
+ def test_fetch_output_canceled():
204
+ from genblaze_core.exceptions import ProviderError
205
+
206
+ provider = ReplicateProvider(api_token="test-token")
207
+ mock_client = MagicMock()
208
+ mock_client.predictions.get.return_value = FakePrediction(status="canceled")
209
+ provider._client = mock_client
210
+
211
+ step = _make_step()
212
+ with pytest.raises(ProviderError, match="canceled"):
213
+ provider.fetch_output("pred-abc123", step)
214
+
215
+
216
+ # --- Token safety ---
217
+
218
+
219
+ def test_fetch_output_rejects_non_https_url():
220
+ """Asset URLs must be HTTPS to prevent SSRF via malicious provider output."""
221
+ from genblaze_core.exceptions import ProviderError
222
+
223
+ provider = ReplicateProvider(api_token="test-token")
224
+ mock_client = MagicMock()
225
+ mock_client.predictions.get.return_value = FakePrediction(output=["file:///etc/passwd"])
226
+ provider._client = mock_client
227
+
228
+ step = _make_step()
229
+ with pytest.raises(ProviderError, match="Unsafe asset URL"):
230
+ provider.fetch_output("pred-abc123", step)
231
+
232
+
233
+ def test_fetch_output_rejects_http_url():
234
+ """Plain HTTP URLs rejected — could target internal metadata endpoints."""
235
+ from genblaze_core.exceptions import ProviderError
236
+
237
+ provider = ReplicateProvider(api_token="test-token")
238
+ mock_client = MagicMock()
239
+ mock_client.predictions.get.return_value = FakePrediction(
240
+ output=["http://169.254.169.254/latest/meta-data/"]
241
+ )
242
+ provider._client = mock_client
243
+
244
+ step = _make_step()
245
+ with pytest.raises(ProviderError, match="Unsafe asset URL"):
246
+ provider.fetch_output("pred-abc123", step)
247
+
248
+
249
+ def test_fetch_output_rejects_schemeless_url():
250
+ """Scheme-less URLs rejected — urlparse sets empty scheme for these."""
251
+ from genblaze_core.exceptions import ProviderError
252
+
253
+ provider = ReplicateProvider(api_token="test-token")
254
+ mock_client = MagicMock()
255
+ mock_client.predictions.get.return_value = FakePrediction(
256
+ output=["//169.254.169.254/latest/meta-data/"]
257
+ )
258
+ provider._client = mock_client
259
+
260
+ step = _make_step()
261
+ with pytest.raises(ProviderError, match="Unsafe asset URL"):
262
+ provider.fetch_output("pred-abc123", step)
263
+
264
+
265
+ def test_cost_tracked_from_metrics():
266
+ """Cost is computed from prediction.metrics.predict_time."""
267
+ provider = ReplicateProvider(api_token="test-token")
268
+ mock_client = MagicMock()
269
+ pred = FakePrediction()
270
+ pred.metrics = MagicMock()
271
+ pred.metrics.predict_time = 10.0
272
+ mock_client.predictions.get.return_value = pred
273
+ provider._client = mock_client
274
+
275
+ step = _make_step()
276
+ result = provider.fetch_output("pred-abc123", step)
277
+ assert result.cost_usd is not None
278
+ assert result.cost_usd == pytest.approx(10.0 * 0.000225)
279
+
280
+
281
+ def test_cost_none_without_metrics():
282
+ """Cost stays None when prediction has no metrics."""
283
+ provider = ReplicateProvider(api_token="test-token")
284
+ mock_client = MagicMock()
285
+ pred = FakePrediction()
286
+ # No metrics attribute
287
+ mock_client.predictions.get.return_value = pred
288
+ provider._client = mock_client
289
+
290
+ step = _make_step()
291
+ result = provider.fetch_output("pred-abc123", step)
292
+ assert result.cost_usd is None
293
+
294
+
295
+ def test_token_not_in_provider_payload():
296
+ """API token must never leak into the provider_payload stored in manifests."""
297
+ provider = ReplicateProvider(api_token="r8_secret_token_123")
298
+ mock_client = MagicMock()
299
+ mock_client.predictions.get.return_value = FakePrediction()
300
+ provider._client = mock_client
301
+
302
+ step = _make_step()
303
+ result = provider.fetch_output("pred-abc123", step)
304
+
305
+ # Serialize the entire payload and check token is absent
306
+ import json
307
+
308
+ payload_str = json.dumps(result.provider_payload)
309
+ assert "r8_secret_token_123" not in payload_str
310
+
311
+
312
+ # --- Compliance harness ---
313
+
314
+
315
+ class TestReplicateCompliance(ProviderComplianceTests):
316
+ """Verify ReplicateProvider satisfies the genblaze provider contract."""
317
+
318
+ # ReplicateProvider populates cost_usd from prediction.metrics.predict_time,
319
+ # but the FakePrediction compliance fixture doesn't expose that metric.
320
+ # A dedicated test (test_cost_tracked_from_predict_time) covers the real
321
+ # code path; compliance waives the check here.
322
+ expects_cost = False
323
+
324
+ def make_provider(self):
325
+ provider = ReplicateProvider(api_token="test-token")
326
+ mock_client = MagicMock()
327
+ mock_client.predictions.create.return_value = FakePrediction()
328
+ mock_client.predictions.get.return_value = FakePrediction()
329
+ provider._client = mock_client
330
+ return provider
331
+
332
+ def make_step(self):
333
+ return Step(provider="replicate", model="test/model", prompt="a cat")