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.
- genblaze_replicate-0.2.0/.gitignore +76 -0
- genblaze_replicate-0.2.0/PKG-INFO +128 -0
- genblaze_replicate-0.2.0/README.md +110 -0
- genblaze_replicate-0.2.0/genblaze_replicate/__init__.py +5 -0
- genblaze_replicate-0.2.0/genblaze_replicate/_errors.py +21 -0
- genblaze_replicate-0.2.0/genblaze_replicate/provider.py +202 -0
- genblaze_replicate-0.2.0/genblaze_replicate/py.typed +0 -0
- genblaze_replicate-0.2.0/pyproject.toml +39 -0
- genblaze_replicate-0.2.0/tests/__init__.py +0 -0
- genblaze_replicate-0.2.0/tests/test_replicate_provider.py +333 -0
|
@@ -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,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")
|