genblaze-replicate 0.2.2__tar.gz → 0.3.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.2 → genblaze_replicate-0.3.0}/PKG-INFO +8 -2
- {genblaze_replicate-0.2.2 → genblaze_replicate-0.3.0}/genblaze_replicate/__init__.py +2 -0
- genblaze_replicate-0.3.0/genblaze_replicate/_version.py +16 -0
- genblaze_replicate-0.3.0/genblaze_replicate/provider.py +399 -0
- {genblaze_replicate-0.2.2 → genblaze_replicate-0.3.0}/pyproject.toml +9 -2
- genblaze_replicate-0.3.0/tests/test_catalog_decoupling.py +100 -0
- {genblaze_replicate-0.2.2 → genblaze_replicate-0.3.0}/tests/test_replicate_provider.py +173 -5
- genblaze_replicate-0.2.2/genblaze_replicate/provider.py +0 -227
- {genblaze_replicate-0.2.2 → genblaze_replicate-0.3.0}/.gitignore +0 -0
- {genblaze_replicate-0.2.2 → genblaze_replicate-0.3.0}/README.md +0 -0
- {genblaze_replicate-0.2.2 → genblaze_replicate-0.3.0}/genblaze_replicate/_errors.py +0 -0
- {genblaze_replicate-0.2.2 → genblaze_replicate-0.3.0}/genblaze_replicate/py.typed +0 -0
- {genblaze_replicate-0.2.2 → genblaze_replicate-0.3.0}/tests/__init__.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: genblaze-replicate
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.3.0
|
|
4
4
|
Summary: Replicate provider adapter for genblaze
|
|
5
5
|
Project-URL: Homepage, https://github.com/backblaze-labs/genblaze
|
|
6
6
|
Project-URL: Documentation, https://github.com/backblaze-labs/genblaze
|
|
@@ -8,11 +8,17 @@ Project-URL: Repository, https://github.com/backblaze-labs/genblaze
|
|
|
8
8
|
Project-URL: Issues, https://github.com/backblaze-labs/genblaze/issues
|
|
9
9
|
Author-email: Jeronimo De Leon <jdeleon@backblaze.com>
|
|
10
10
|
License-Expression: MIT
|
|
11
|
+
Keywords: ai,c2pa-ready,genai,genblaze,image,manifest,media,pipeline,provenance,replicate,video
|
|
11
12
|
Classifier: Development Status :: 3 - Alpha
|
|
12
13
|
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
17
|
+
Classifier: Topic :: Multimedia
|
|
18
|
+
Classifier: Topic :: Software Development :: Libraries
|
|
13
19
|
Classifier: Typing :: Typed
|
|
14
20
|
Requires-Python: >=3.11
|
|
15
|
-
Requires-Dist: genblaze-core<0.
|
|
21
|
+
Requires-Dist: genblaze-core<0.4,>=0.3.0
|
|
16
22
|
Requires-Dist: replicate>=0.25
|
|
17
23
|
Provides-Extra: dev
|
|
18
24
|
Requires-Dist: pytest>=7.0; extra == 'dev'
|
|
@@ -0,0 +1,16 @@
|
|
|
1
|
+
"""``genblaze-replicate`` package version — single source of truth via importlib.metadata.
|
|
2
|
+
|
|
3
|
+
Plan 5 Phase 1B — closes the version-drift class of bug across every
|
|
4
|
+
genblaze connector. Reading from ``importlib.metadata`` makes the
|
|
5
|
+
constant always equal to whatever wheel is installed; no manual edits
|
|
6
|
+
per release.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
from importlib.metadata import PackageNotFoundError, version
|
|
12
|
+
|
|
13
|
+
try:
|
|
14
|
+
__version__: str = version("genblaze-replicate")
|
|
15
|
+
except PackageNotFoundError: # pragma: no cover — editable dev installs
|
|
16
|
+
__version__ = "0.0.0+unknown"
|
|
@@ -0,0 +1,399 @@
|
|
|
1
|
+
"""ReplicateProvider — adapter for the Replicate API.
|
|
2
|
+
|
|
3
|
+
Replicate hosts thousands of models. Slugs are free-form ``owner/name``
|
|
4
|
+
strings — any of them is a valid input. The connector ships:
|
|
5
|
+
|
|
6
|
+
* ``DiscoverySupport.NATIVE`` — Replicate's API can authoritatively answer
|
|
7
|
+
"is this slug live?" via ``GET /v1/models/{owner}/{name}``.
|
|
8
|
+
* An empty registry with a permissive fallback. There are no per-slug
|
|
9
|
+
default specs because Replicate's surface is too large to enumerate.
|
|
10
|
+
* A custom ``validate_model()`` that does a per-slug ``models.get()``
|
|
11
|
+
call (cached per-process, TTL-bounded) rather than enumerating the
|
|
12
|
+
full catalog. This is what NATIVE means for Replicate in practice.
|
|
13
|
+
* A custom ``discover_models()`` that returns the *first page* of the
|
|
14
|
+
catalog as a documentation hint — enough to populate ``known()`` for
|
|
15
|
+
IDE autocomplete and conformance testing without paying the cost of
|
|
16
|
+
enumerating ~10k models on every cold start.
|
|
17
|
+
|
|
18
|
+
**Pricing**: Replicate was previously hardcoded to
|
|
19
|
+
``predict_time × 0.000225`` USD on the fallback spec. As of
|
|
20
|
+
``genblaze-core 0.3.0`` the SDK no longer ships pricing — register the
|
|
21
|
+
recipe yourself if you want compute-time cost tracking. See
|
|
22
|
+
``docs/reference/pricing-recipes.md`` for the canonical Replicate recipe.
|
|
23
|
+
|
|
24
|
+
Docs: https://replicate.com/docs/reference/http
|
|
25
|
+
"""
|
|
26
|
+
|
|
27
|
+
from __future__ import annotations
|
|
28
|
+
|
|
29
|
+
import logging
|
|
30
|
+
import mimetypes
|
|
31
|
+
import threading
|
|
32
|
+
import time
|
|
33
|
+
from typing import Any
|
|
34
|
+
from urllib.parse import urlparse
|
|
35
|
+
|
|
36
|
+
from genblaze_core.exceptions import ProviderError
|
|
37
|
+
from genblaze_core.models.asset import Asset
|
|
38
|
+
from genblaze_core.models.enums import Modality, ProviderErrorCode
|
|
39
|
+
from genblaze_core.models.step import Step
|
|
40
|
+
from genblaze_core.providers import (
|
|
41
|
+
BaseProvider,
|
|
42
|
+
DiscoveryResult,
|
|
43
|
+
DiscoverySupport,
|
|
44
|
+
ModelRegistry,
|
|
45
|
+
ModelSpec,
|
|
46
|
+
ProviderCapabilities,
|
|
47
|
+
RetryPolicy,
|
|
48
|
+
ValidationOutcome,
|
|
49
|
+
ValidationResult,
|
|
50
|
+
ValidationSource,
|
|
51
|
+
route_by_media_type,
|
|
52
|
+
validate_asset_url,
|
|
53
|
+
)
|
|
54
|
+
from genblaze_core.providers.discovery import DEFAULT_TTL_SECONDS, _DiscoveryCache
|
|
55
|
+
from genblaze_core.providers.retry import retry_after_from_response
|
|
56
|
+
from genblaze_core.runnable.config import RunnableConfig
|
|
57
|
+
|
|
58
|
+
from ._errors import map_replicate_error
|
|
59
|
+
|
|
60
|
+
logger = logging.getLogger("genblaze.replicate")
|
|
61
|
+
|
|
62
|
+
# Cap for the first-page discovery snapshot. Replicate's /v1/models
|
|
63
|
+
# returns ~25 models per page; one page is enough to seed known() for
|
|
64
|
+
# IDE/conformance purposes without enumerating the full catalog.
|
|
65
|
+
_DISCOVERY_PAGE_LIMIT: int = 25
|
|
66
|
+
|
|
67
|
+
# Replicate has no enumerable per-slug defaults. The fallback applies to
|
|
68
|
+
# every model id with route_by_media_type wiring chain inputs to the
|
|
69
|
+
# common ``image``/``video``/``audio`` keys most Replicate models accept.
|
|
70
|
+
_FALLBACK_SPEC = ModelSpec(
|
|
71
|
+
model_id="*",
|
|
72
|
+
input_mapping=route_by_media_type({"image": "image", "video": "video", "audio": "audio"}),
|
|
73
|
+
)
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
class ReplicateProvider(BaseProvider):
|
|
77
|
+
"""Provider adapter for Replicate (replicate.com)."""
|
|
78
|
+
|
|
79
|
+
name = "replicate"
|
|
80
|
+
discovery_support = DiscoverySupport.NATIVE
|
|
81
|
+
"""Replicate exposes ``GET /v1/models/{owner}/{name}`` for authoritative
|
|
82
|
+
per-slug existence checks. The catalog itself is too large to enumerate
|
|
83
|
+
(~10k+ public models), so this provider validates per-slug rather than
|
|
84
|
+
snapshotting the full catalog."""
|
|
85
|
+
|
|
86
|
+
@classmethod
|
|
87
|
+
def create_registry(cls) -> ModelRegistry:
|
|
88
|
+
# Every model resolves through the fallback spec — Replicate's
|
|
89
|
+
# catalog is enumerated only via discovery, not shipped here.
|
|
90
|
+
# The discovery cache is wired per-instance in __init__ so each
|
|
91
|
+
# provider sees its own credentials' view of the catalog.
|
|
92
|
+
return ModelRegistry(fallback=_FALLBACK_SPEC)
|
|
93
|
+
|
|
94
|
+
def get_capabilities(self) -> ProviderCapabilities:
|
|
95
|
+
"""Replicate: multi-modal generation depending on selected model."""
|
|
96
|
+
return ProviderCapabilities(
|
|
97
|
+
supported_modalities=[Modality.IMAGE, Modality.VIDEO, Modality.AUDIO],
|
|
98
|
+
supported_inputs=["text", "image", "video", "audio"],
|
|
99
|
+
accepts_chain_input=True,
|
|
100
|
+
)
|
|
101
|
+
|
|
102
|
+
def __init__(
|
|
103
|
+
self,
|
|
104
|
+
api_token: str | None = None,
|
|
105
|
+
poll_interval: float = 1.0,
|
|
106
|
+
http_timeout: float = 30.0,
|
|
107
|
+
*,
|
|
108
|
+
models: ModelRegistry | None = None,
|
|
109
|
+
retry_policy: RetryPolicy | None = None,
|
|
110
|
+
probe_cache_ttl: float | None = None,
|
|
111
|
+
probe_cache_max_entries: int | None = None,
|
|
112
|
+
):
|
|
113
|
+
# probe_cache_* are no-ops on this NATIVE-discovery provider; accepted
|
|
114
|
+
# for API uniformity so calling code can pass them to any provider
|
|
115
|
+
# without TypeError.
|
|
116
|
+
super().__init__(
|
|
117
|
+
models=models,
|
|
118
|
+
retry_policy=retry_policy,
|
|
119
|
+
probe_cache_ttl=probe_cache_ttl,
|
|
120
|
+
probe_cache_max_entries=probe_cache_max_entries,
|
|
121
|
+
)
|
|
122
|
+
self.poll_interval = poll_interval
|
|
123
|
+
self._client: Any = None
|
|
124
|
+
self._api_token = api_token
|
|
125
|
+
self._http_timeout = http_timeout
|
|
126
|
+
# Per-slug authoritative-validation cache. Stored as
|
|
127
|
+
# ``slug → (fetched_monotonic, ValidationResult)`` so we honor a TTL
|
|
128
|
+
# without re-issuing per-slug GETs on every Pipeline.run(). RLock
|
|
129
|
+
# keeps concurrent ``validate_model`` callers from racing on the
|
|
130
|
+
# cache writes; reads are fast-path lockless once populated.
|
|
131
|
+
self._validation_cache: dict[str, tuple[float, ValidationResult]] = {}
|
|
132
|
+
self._validation_lock = threading.RLock()
|
|
133
|
+
self._validation_ttl: float = DEFAULT_TTL_SECONDS
|
|
134
|
+
# Wire a discovery cache for the first-page snapshot. The fetcher
|
|
135
|
+
# closes over ``self`` so it picks up the (possibly-late-bound)
|
|
136
|
+
# ``self._client``.
|
|
137
|
+
self._models._discovery_cache = _DiscoveryCache(
|
|
138
|
+
self._fetch_first_page,
|
|
139
|
+
default_max_age_seconds=DEFAULT_TTL_SECONDS,
|
|
140
|
+
)
|
|
141
|
+
|
|
142
|
+
# --- catalog discovery & validation ------------------------------------
|
|
143
|
+
|
|
144
|
+
def _fetch_first_page(self) -> DiscoveryResult:
|
|
145
|
+
"""Fetcher backing ``discover_models`` — first page of /v1/models.
|
|
146
|
+
|
|
147
|
+
Replicate's full catalog is too large to enumerate; one page is
|
|
148
|
+
enough to seed ``known()`` for IDE / conformance purposes. The
|
|
149
|
+
authoritative per-slug check lives in ``validate_model``.
|
|
150
|
+
"""
|
|
151
|
+
try:
|
|
152
|
+
client = self._get_client()
|
|
153
|
+
page = client.models.list()
|
|
154
|
+
slugs: set[str] = set()
|
|
155
|
+
results = getattr(page, "results", None) or list(page)
|
|
156
|
+
for model in results[:_DISCOVERY_PAGE_LIMIT]:
|
|
157
|
+
owner = getattr(model, "owner", None)
|
|
158
|
+
name = getattr(model, "name", None)
|
|
159
|
+
if owner and name:
|
|
160
|
+
slugs.add(f"{owner}/{name}")
|
|
161
|
+
return DiscoveryResult.ok(slugs, source_url="https://api.replicate.com/v1/models")
|
|
162
|
+
except Exception as exc:
|
|
163
|
+
logger.warning("Replicate discover_models first-page fetch failed: %s", exc)
|
|
164
|
+
return DiscoveryResult.failed(
|
|
165
|
+
f"first-page enumeration failed: {exc}",
|
|
166
|
+
source_url="https://api.replicate.com/v1/models",
|
|
167
|
+
)
|
|
168
|
+
|
|
169
|
+
def discover_models(
|
|
170
|
+
self,
|
|
171
|
+
*,
|
|
172
|
+
max_age_seconds: float | None = ..., # type: ignore[assignment]
|
|
173
|
+
) -> DiscoveryResult:
|
|
174
|
+
"""Return the first-page snapshot of Replicate's catalog.
|
|
175
|
+
|
|
176
|
+
This is intentionally not exhaustive — Replicate hosts thousands
|
|
177
|
+
of models, and ``known()`` is documentation-grade, not a
|
|
178
|
+
contract. For authoritative existence checks, use
|
|
179
|
+
``validate_model(slug)``, which does a per-slug ``models.get()``.
|
|
180
|
+
"""
|
|
181
|
+
cache = self._models._discovery_cache
|
|
182
|
+
assert cache is not None # wired in __init__
|
|
183
|
+
if max_age_seconds is ...: # type: ignore[comparison-overlap]
|
|
184
|
+
return cache.get()
|
|
185
|
+
return cache.get(max_age_seconds=max_age_seconds)
|
|
186
|
+
|
|
187
|
+
def validate_model(self, model_id: str, *, refresh: bool = False) -> ValidationResult:
|
|
188
|
+
"""Authoritative per-slug existence check.
|
|
189
|
+
|
|
190
|
+
Replicate's per-model GET is cheap (one round-trip, no token spend)
|
|
191
|
+
and returns 404 for missing slugs. We cache the result per-slug
|
|
192
|
+
with a TTL so successive ``Pipeline.run()`` invocations don't
|
|
193
|
+
re-fetch. ``refresh=True`` evicts the cache entry first.
|
|
194
|
+
|
|
195
|
+
Falls back to the base ``validate_model`` flow when the slug is
|
|
196
|
+
user-registered or matches a registered family — those layers are
|
|
197
|
+
already authoritative without a network call.
|
|
198
|
+
"""
|
|
199
|
+
# Defer to the base class first — user-registered specs and
|
|
200
|
+
# family matches with cached discovery hits don't need a probe.
|
|
201
|
+
base_result = self._models.validate(model_id, discovery_support=self.discovery_support)
|
|
202
|
+
if (
|
|
203
|
+
base_result.outcome is ValidationOutcome.OK_AUTHORITATIVE
|
|
204
|
+
and base_result.source is not ValidationSource.DISCOVERY
|
|
205
|
+
):
|
|
206
|
+
# USER or PROBE-source authoritative answer — no network needed.
|
|
207
|
+
return base_result
|
|
208
|
+
|
|
209
|
+
if refresh:
|
|
210
|
+
with self._validation_lock:
|
|
211
|
+
self._validation_cache.pop(model_id, None)
|
|
212
|
+
else:
|
|
213
|
+
cached = self._lookup_validation_cache(model_id)
|
|
214
|
+
if cached is not None:
|
|
215
|
+
return cached
|
|
216
|
+
|
|
217
|
+
result = self._authoritative_lookup(model_id)
|
|
218
|
+
with self._validation_lock:
|
|
219
|
+
self._validation_cache[model_id] = (time.monotonic(), result)
|
|
220
|
+
return result
|
|
221
|
+
|
|
222
|
+
def _lookup_validation_cache(self, model_id: str) -> ValidationResult | None:
|
|
223
|
+
with self._validation_lock:
|
|
224
|
+
entry = self._validation_cache.get(model_id)
|
|
225
|
+
if entry is None:
|
|
226
|
+
return None
|
|
227
|
+
fetched_at, result = entry
|
|
228
|
+
if (time.monotonic() - fetched_at) > self._validation_ttl:
|
|
229
|
+
# Expired — evict and force a fresh lookup.
|
|
230
|
+
del self._validation_cache[model_id]
|
|
231
|
+
return None
|
|
232
|
+
return result
|
|
233
|
+
|
|
234
|
+
def _authoritative_lookup(self, model_id: str) -> ValidationResult:
|
|
235
|
+
"""Single-slug GET against /v1/models/{owner}/{name}.
|
|
236
|
+
|
|
237
|
+
Translates the Replicate SDK's responses into ``ValidationResult``:
|
|
238
|
+
|
|
239
|
+
* Success → ``OK_AUTHORITATIVE`` (source PROBE — same shape a
|
|
240
|
+
PARTIAL provider's family.probe would produce).
|
|
241
|
+
* 404 / NotFound → ``NOT_FOUND``.
|
|
242
|
+
* Other errors → ``UNKNOWN_PERMISSIVE`` (we couldn't verify; let
|
|
243
|
+
the upstream call surface the real error if there is one).
|
|
244
|
+
"""
|
|
245
|
+
try:
|
|
246
|
+
client = self._get_client()
|
|
247
|
+
client.models.get(model_id)
|
|
248
|
+
except Exception as exc:
|
|
249
|
+
err_code = map_replicate_error(exc)
|
|
250
|
+
if err_code is ProviderErrorCode.MODEL_ERROR:
|
|
251
|
+
return ValidationResult.not_found(
|
|
252
|
+
ValidationSource.PROBE,
|
|
253
|
+
detail=f"upstream returned not-found: {exc}",
|
|
254
|
+
)
|
|
255
|
+
logger.debug(
|
|
256
|
+
"Replicate validate_model probe inconclusive for %s: %s",
|
|
257
|
+
model_id,
|
|
258
|
+
exc,
|
|
259
|
+
)
|
|
260
|
+
return ValidationResult.unknown_permissive(detail=f"probe inconclusive: {exc}")
|
|
261
|
+
return ValidationResult.ok_authoritative(
|
|
262
|
+
ValidationSource.PROBE,
|
|
263
|
+
detail="confirmed via /v1/models/{owner}/{name}",
|
|
264
|
+
)
|
|
265
|
+
|
|
266
|
+
def _get_client(self):
|
|
267
|
+
if self._client is None:
|
|
268
|
+
try:
|
|
269
|
+
import httpx
|
|
270
|
+
import replicate
|
|
271
|
+
|
|
272
|
+
timeout = httpx.Timeout(self._http_timeout, connect=10.0)
|
|
273
|
+
if self._api_token:
|
|
274
|
+
self._client = replicate.Client(
|
|
275
|
+
api_token=self._api_token, # noqa: S106
|
|
276
|
+
timeout=timeout,
|
|
277
|
+
)
|
|
278
|
+
else:
|
|
279
|
+
self._client = replicate.Client(timeout=timeout)
|
|
280
|
+
except ImportError as exc:
|
|
281
|
+
raise ProviderError(
|
|
282
|
+
"replicate package not installed. Run: pip install replicate"
|
|
283
|
+
) from exc
|
|
284
|
+
return self._client
|
|
285
|
+
|
|
286
|
+
def submit(self, step: Step, config: RunnableConfig | None = None) -> Any:
|
|
287
|
+
client = self._get_client()
|
|
288
|
+
try:
|
|
289
|
+
input_params = self.prepare_payload(step)
|
|
290
|
+
prediction = client.predictions.create(
|
|
291
|
+
model=step.model,
|
|
292
|
+
input=input_params,
|
|
293
|
+
)
|
|
294
|
+
return prediction.id
|
|
295
|
+
except Exception as exc:
|
|
296
|
+
raise ProviderError(
|
|
297
|
+
f"Replicate submit failed: {exc}",
|
|
298
|
+
error_code=map_replicate_error(exc),
|
|
299
|
+
retry_after=retry_after_from_response(exc),
|
|
300
|
+
) from exc
|
|
301
|
+
|
|
302
|
+
def poll(self, prediction_id: Any, config: RunnableConfig | None = None) -> bool:
|
|
303
|
+
client = self._get_client()
|
|
304
|
+
try:
|
|
305
|
+
prediction = client.predictions.get(prediction_id)
|
|
306
|
+
if prediction.status in ("succeeded", "failed", "canceled"):
|
|
307
|
+
self._cache_poll_result(prediction_id, prediction)
|
|
308
|
+
return True
|
|
309
|
+
return False
|
|
310
|
+
except Exception as exc:
|
|
311
|
+
raise ProviderError(
|
|
312
|
+
f"Replicate poll failed: {exc}",
|
|
313
|
+
error_code=map_replicate_error(exc),
|
|
314
|
+
retry_after=retry_after_from_response(exc),
|
|
315
|
+
) from exc
|
|
316
|
+
|
|
317
|
+
def fetch_output(self, prediction_id: Any, step: Step) -> Step:
|
|
318
|
+
client = self._get_client()
|
|
319
|
+
try:
|
|
320
|
+
prediction = self._get_cached_poll_result(prediction_id)
|
|
321
|
+
if prediction is None:
|
|
322
|
+
prediction = client.predictions.get(prediction_id)
|
|
323
|
+
|
|
324
|
+
# Capture predict_time so registry pricing can read it.
|
|
325
|
+
metrics = getattr(prediction, "metrics", None)
|
|
326
|
+
predict_time = getattr(metrics, "predict_time", None) if metrics else None
|
|
327
|
+
|
|
328
|
+
step.provider_payload = {
|
|
329
|
+
"replicate": {
|
|
330
|
+
"prediction_id": prediction.id,
|
|
331
|
+
"model": prediction.model if hasattr(prediction, "model") else None,
|
|
332
|
+
"version": prediction.version if hasattr(prediction, "version") else None,
|
|
333
|
+
"status": prediction.status,
|
|
334
|
+
"created_at": str(prediction.created_at)
|
|
335
|
+
if hasattr(prediction, "created_at")
|
|
336
|
+
else None,
|
|
337
|
+
"predict_time": predict_time,
|
|
338
|
+
}
|
|
339
|
+
}
|
|
340
|
+
|
|
341
|
+
if prediction.status == "failed":
|
|
342
|
+
error_msg = prediction.error or "Unknown error"
|
|
343
|
+
raise ProviderError(
|
|
344
|
+
error_msg,
|
|
345
|
+
error_code=map_replicate_error(error_msg),
|
|
346
|
+
)
|
|
347
|
+
|
|
348
|
+
if prediction.status == "canceled":
|
|
349
|
+
raise ProviderError(
|
|
350
|
+
"Prediction was canceled",
|
|
351
|
+
error_code=ProviderErrorCode.UNKNOWN,
|
|
352
|
+
)
|
|
353
|
+
|
|
354
|
+
# Replicate output shapes vary by model: str (single URL), list[str]
|
|
355
|
+
# (multi-asset), dict[str, str | list] (e.g. {"video": url,
|
|
356
|
+
# "subtitles": url} from text-to-video models with side-channels),
|
|
357
|
+
# or None (no output). Normalize to list[str].
|
|
358
|
+
raw_output = prediction.output
|
|
359
|
+
urls: list[str]
|
|
360
|
+
if raw_output is None:
|
|
361
|
+
urls = []
|
|
362
|
+
elif isinstance(raw_output, str):
|
|
363
|
+
urls = [raw_output]
|
|
364
|
+
elif isinstance(raw_output, list):
|
|
365
|
+
# Nested lists happen on batch-output models; flatten one level.
|
|
366
|
+
urls = []
|
|
367
|
+
for item in raw_output:
|
|
368
|
+
if isinstance(item, str):
|
|
369
|
+
urls.append(item)
|
|
370
|
+
elif isinstance(item, list):
|
|
371
|
+
urls.extend(str(u) for u in item if isinstance(u, (str, bytes)))
|
|
372
|
+
elif isinstance(raw_output, dict):
|
|
373
|
+
# Multi-channel outputs: keep only URL-shaped string values.
|
|
374
|
+
urls = [str(v) for v in raw_output.values() if isinstance(v, str)]
|
|
375
|
+
else:
|
|
376
|
+
raise ProviderError(
|
|
377
|
+
f"Unexpected Replicate output shape "
|
|
378
|
+
f"({type(raw_output).__name__}): {raw_output!r}",
|
|
379
|
+
error_code=ProviderErrorCode.SERVER_ERROR,
|
|
380
|
+
)
|
|
381
|
+
|
|
382
|
+
for url_str in urls:
|
|
383
|
+
validate_asset_url(url_str)
|
|
384
|
+
path = urlparse(url_str).path
|
|
385
|
+
mime, _ = mimetypes.guess_type(path)
|
|
386
|
+
if mime is None:
|
|
387
|
+
mime = f"{step.modality.value}/octet-stream"
|
|
388
|
+
step.assets.append(Asset(url=url_str, media_type=mime))
|
|
389
|
+
|
|
390
|
+
self._apply_registry_pricing(step)
|
|
391
|
+
return step
|
|
392
|
+
except ProviderError:
|
|
393
|
+
raise
|
|
394
|
+
except Exception as exc:
|
|
395
|
+
raise ProviderError(
|
|
396
|
+
f"Replicate fetch_output failed: {exc}",
|
|
397
|
+
error_code=map_replicate_error(exc),
|
|
398
|
+
retry_after=retry_after_from_response(exc),
|
|
399
|
+
) from exc
|
|
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "genblaze-replicate"
|
|
7
|
-
version = "0.
|
|
7
|
+
version = "0.3.0"
|
|
8
8
|
description = "Replicate provider adapter for genblaze"
|
|
9
9
|
authors = [{name = "Jeronimo De Leon", email = "jdeleon@backblaze.com"}]
|
|
10
10
|
readme = "README.md"
|
|
@@ -14,9 +14,16 @@ classifiers = [
|
|
|
14
14
|
"Development Status :: 3 - Alpha",
|
|
15
15
|
"License :: OSI Approved :: MIT License",
|
|
16
16
|
"Typing :: Typed",
|
|
17
|
+
"Programming Language :: Python :: 3.11",
|
|
18
|
+
"Programming Language :: Python :: 3.12",
|
|
19
|
+
"Programming Language :: Python :: 3.13",
|
|
20
|
+
"Topic :: Multimedia",
|
|
21
|
+
"Topic :: Software Development :: Libraries",
|
|
17
22
|
]
|
|
23
|
+
keywords = ["genblaze", "ai", "media", "manifest", "provenance", "c2pa-ready", "genai", "pipeline", "replicate", "video", "image"]
|
|
24
|
+
|
|
18
25
|
dependencies = [
|
|
19
|
-
"genblaze-core>=0.
|
|
26
|
+
"genblaze-core>=0.3.0,<0.4",
|
|
20
27
|
"replicate>=0.25",
|
|
21
28
|
]
|
|
22
29
|
|
|
@@ -0,0 +1,100 @@
|
|
|
1
|
+
"""Catalog-decoupling regression tests for ``genblaze-replicate``.
|
|
2
|
+
|
|
3
|
+
Replicate is the proof-point for ``DiscoverySupport.NATIVE`` with an
|
|
4
|
+
empty default registry: every slug routes through the permissive
|
|
5
|
+
fallback spec, and the upstream catalog is enumerated dynamically via
|
|
6
|
+
``client.models.list()`` (wired in ``__init__`` per-instance for
|
|
7
|
+
auth-aware caches).
|
|
8
|
+
|
|
9
|
+
This file pins the contract for that minimal surface so a future
|
|
10
|
+
edit can't silently regress the post-decoupling shape.
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
from __future__ import annotations
|
|
14
|
+
|
|
15
|
+
from unittest.mock import MagicMock, patch
|
|
16
|
+
|
|
17
|
+
import pytest
|
|
18
|
+
from genblaze_core.providers import DiscoverySupport
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
@pytest.fixture(autouse=True)
|
|
22
|
+
def _patch_replicate_sdk():
|
|
23
|
+
"""Avoid importing the real replicate SDK in tests."""
|
|
24
|
+
fake = MagicMock()
|
|
25
|
+
with patch.dict("sys.modules", {"replicate": fake}):
|
|
26
|
+
yield
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
# --- DiscoverySupport declaration -----------------------------------------
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
class TestDiscoverySupportDeclaration:
|
|
33
|
+
def test_native(self) -> None:
|
|
34
|
+
from genblaze_replicate import ReplicateProvider
|
|
35
|
+
|
|
36
|
+
assert ReplicateProvider.discovery_support is DiscoverySupport.NATIVE
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
# --- Registry shape (no families, fallback only) --------------------------
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
class TestRegistryShape:
|
|
43
|
+
def test_no_provider_families_shipped(self) -> None:
|
|
44
|
+
"""Replicate is meta-vendor — every model on its hub is a valid
|
|
45
|
+
slug. The connector ships zero families and lets the upstream
|
|
46
|
+
catalog speak for itself via discovery."""
|
|
47
|
+
from genblaze_replicate import ReplicateProvider
|
|
48
|
+
|
|
49
|
+
provider = ReplicateProvider(api_token="test")
|
|
50
|
+
assert provider._models.families == ()
|
|
51
|
+
|
|
52
|
+
def test_arbitrary_slug_resolves_via_fallback(self) -> None:
|
|
53
|
+
"""Any slug — known model, brand-new release, typo — resolves
|
|
54
|
+
through the fallback. Replicate's own API gates submit-time
|
|
55
|
+
liveness."""
|
|
56
|
+
from genblaze_replicate import ReplicateProvider
|
|
57
|
+
|
|
58
|
+
provider = ReplicateProvider(api_token="test")
|
|
59
|
+
for slug in (
|
|
60
|
+
"black-forest-labs/flux-schnell",
|
|
61
|
+
"stability-ai/sdxl",
|
|
62
|
+
"owner/brand-new-model",
|
|
63
|
+
"a-typo-slug",
|
|
64
|
+
):
|
|
65
|
+
spec = provider._models.get(slug)
|
|
66
|
+
assert spec is not None, slug
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
# --- Pricing-removed contract --------------------------------------------
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
class TestPricingPhaseOut:
|
|
73
|
+
def test_fallback_spec_carries_no_pricing(self) -> None:
|
|
74
|
+
"""The connector's fallback spec ships with ``pricing=None``;
|
|
75
|
+
cost is user-registered via ``register_pricing``."""
|
|
76
|
+
from genblaze_replicate import ReplicateProvider
|
|
77
|
+
|
|
78
|
+
provider = ReplicateProvider(api_token="test")
|
|
79
|
+
spec = provider._models.get("any/slug")
|
|
80
|
+
assert spec.pricing is None
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
# --- Cross-provider isolation --------------------------------------------
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
class TestCrossProviderIsolation:
|
|
87
|
+
def test_replicate_does_not_match_other_provider_slugs(self) -> None:
|
|
88
|
+
"""Replicate has no families — ``match_family`` returns None
|
|
89
|
+
for any slug; nothing matches, including slugs that look like
|
|
90
|
+
another provider's catalog."""
|
|
91
|
+
from genblaze_replicate import ReplicateProvider
|
|
92
|
+
|
|
93
|
+
provider = ReplicateProvider(api_token="test")
|
|
94
|
+
for slug in (
|
|
95
|
+
"veo-3.0-generate-001",
|
|
96
|
+
"imagen-3.0-generate-002",
|
|
97
|
+
"tts-1",
|
|
98
|
+
"ray-2",
|
|
99
|
+
):
|
|
100
|
+
assert provider._models.match_family(slug) is None, slug
|
|
@@ -7,8 +7,9 @@ from typing import Any
|
|
|
7
7
|
from unittest.mock import MagicMock
|
|
8
8
|
|
|
9
9
|
import pytest
|
|
10
|
-
from genblaze_core.models.enums import ProviderErrorCode
|
|
10
|
+
from genblaze_core.models.enums import Modality, ProviderErrorCode
|
|
11
11
|
from genblaze_core.models.step import Step
|
|
12
|
+
from genblaze_core.providers import ModelSpec
|
|
12
13
|
from genblaze_core.testing import ProviderComplianceTests
|
|
13
14
|
from genblaze_replicate._errors import map_replicate_error
|
|
14
15
|
from genblaze_replicate.provider import ReplicateProvider
|
|
@@ -310,8 +311,12 @@ def test_fetch_output_rejects_schemeless_url():
|
|
|
310
311
|
provider.fetch_output("pred-abc123", step)
|
|
311
312
|
|
|
312
313
|
|
|
313
|
-
def
|
|
314
|
-
"""
|
|
314
|
+
def test_cost_none_by_default():
|
|
315
|
+
"""As of genblaze-core 0.3.0 the SDK no longer ships pricing for
|
|
316
|
+
Replicate. ``cost_usd`` is ``None`` unless the user has registered
|
|
317
|
+
a pricing strategy via ``provider.models.register_pricing()``.
|
|
318
|
+
See ``docs/reference/pricing-recipes.md``.
|
|
319
|
+
"""
|
|
315
320
|
provider = ReplicateProvider(api_token="test-token")
|
|
316
321
|
mock_client = MagicMock()
|
|
317
322
|
pred = FakePrediction()
|
|
@@ -320,6 +325,33 @@ def test_cost_tracked_from_metrics():
|
|
|
320
325
|
mock_client.predictions.get.return_value = pred
|
|
321
326
|
provider._client = mock_client
|
|
322
327
|
|
|
328
|
+
step = _make_step()
|
|
329
|
+
result = provider.fetch_output("pred-abc123", step)
|
|
330
|
+
assert result.cost_usd is None
|
|
331
|
+
|
|
332
|
+
|
|
333
|
+
def test_cost_tracked_with_user_registered_pricing():
|
|
334
|
+
"""User-registered compute-time pricing flows through the same path."""
|
|
335
|
+
from genblaze_core.providers import per_response_metric
|
|
336
|
+
|
|
337
|
+
def compute_time_cost(ctx):
|
|
338
|
+
payload = ctx.provider_payload.get("replicate") if ctx.provider_payload else None
|
|
339
|
+
if not isinstance(payload, dict):
|
|
340
|
+
return None
|
|
341
|
+
predict_time = payload.get("predict_time")
|
|
342
|
+
if predict_time is None:
|
|
343
|
+
return None
|
|
344
|
+
return float(predict_time) * 0.000225
|
|
345
|
+
|
|
346
|
+
provider = ReplicateProvider(api_token="test-token")
|
|
347
|
+
provider.models.register_pricing("test/model", per_response_metric(compute_time_cost))
|
|
348
|
+
mock_client = MagicMock()
|
|
349
|
+
pred = FakePrediction()
|
|
350
|
+
pred.metrics = MagicMock()
|
|
351
|
+
pred.metrics.predict_time = 10.0
|
|
352
|
+
mock_client.predictions.get.return_value = pred
|
|
353
|
+
provider._client = mock_client
|
|
354
|
+
|
|
323
355
|
step = _make_step()
|
|
324
356
|
result = provider.fetch_output("pred-abc123", step)
|
|
325
357
|
assert result.cost_usd is not None
|
|
@@ -327,8 +359,11 @@ def test_cost_tracked_from_metrics():
|
|
|
327
359
|
|
|
328
360
|
|
|
329
361
|
def test_cost_none_without_metrics():
|
|
330
|
-
"""Cost stays None when prediction has no metrics."""
|
|
362
|
+
"""Cost stays None when prediction has no metrics, even with pricing registered."""
|
|
363
|
+
from genblaze_core.providers import per_response_metric
|
|
364
|
+
|
|
331
365
|
provider = ReplicateProvider(api_token="test-token")
|
|
366
|
+
provider.models.register_pricing("test/model", per_response_metric(lambda ctx: 0.0))
|
|
332
367
|
mock_client = MagicMock()
|
|
333
368
|
pred = FakePrediction()
|
|
334
369
|
# No metrics attribute
|
|
@@ -337,7 +372,9 @@ def test_cost_none_without_metrics():
|
|
|
337
372
|
|
|
338
373
|
step = _make_step()
|
|
339
374
|
result = provider.fetch_output("pred-abc123", step)
|
|
340
|
-
|
|
375
|
+
# The strategy returns 0.0 if metrics are absent — but we registered
|
|
376
|
+
# one that returns 0.0 unconditionally. Verify the pricing path runs.
|
|
377
|
+
assert result.cost_usd == pytest.approx(0.0)
|
|
341
378
|
|
|
342
379
|
|
|
343
380
|
def test_token_not_in_provider_payload():
|
|
@@ -357,6 +394,137 @@ def test_token_not_in_provider_payload():
|
|
|
357
394
|
assert "r8_secret_token_123" not in payload_str
|
|
358
395
|
|
|
359
396
|
|
|
397
|
+
# --- Catalog-decoupling: discovery + per-slug validation ---
|
|
398
|
+
|
|
399
|
+
|
|
400
|
+
def _make_model(owner: str, name: str):
|
|
401
|
+
"""Minimal fake of a Replicate ``Model`` object."""
|
|
402
|
+
m = MagicMock()
|
|
403
|
+
m.owner = owner
|
|
404
|
+
m.name = name
|
|
405
|
+
return m
|
|
406
|
+
|
|
407
|
+
|
|
408
|
+
def test_discovery_support_native():
|
|
409
|
+
"""Replicate is the proof-point for NATIVE discovery in PR #3."""
|
|
410
|
+
from genblaze_core.providers import DiscoverySupport
|
|
411
|
+
|
|
412
|
+
assert ReplicateProvider.discovery_support is DiscoverySupport.NATIVE
|
|
413
|
+
|
|
414
|
+
|
|
415
|
+
def test_discover_models_returns_first_page():
|
|
416
|
+
"""``discover_models`` snapshots the first page of /v1/models — enough
|
|
417
|
+
to seed ``known()`` without enumerating the entire catalog."""
|
|
418
|
+
from genblaze_core.providers import DiscoveryStatus
|
|
419
|
+
|
|
420
|
+
provider = ReplicateProvider(api_token="test-token")
|
|
421
|
+
mock_client = MagicMock()
|
|
422
|
+
page = MagicMock()
|
|
423
|
+
page.results = [
|
|
424
|
+
_make_model("black-forest-labs", "flux-schnell"),
|
|
425
|
+
_make_model("stability-ai", "sdxl"),
|
|
426
|
+
_make_model("meta", "llama-3"),
|
|
427
|
+
]
|
|
428
|
+
mock_client.models.list.return_value = page
|
|
429
|
+
provider._client = mock_client
|
|
430
|
+
|
|
431
|
+
result = provider.discover_models()
|
|
432
|
+
assert result.status is DiscoveryStatus.OK
|
|
433
|
+
assert "black-forest-labs/flux-schnell" in result.slugs
|
|
434
|
+
assert "stability-ai/sdxl" in result.slugs
|
|
435
|
+
assert result.source_url == "https://api.replicate.com/v1/models"
|
|
436
|
+
|
|
437
|
+
|
|
438
|
+
def test_discover_models_failure_returns_failed():
|
|
439
|
+
"""A failed list call surfaces as ``DiscoveryStatus.FAILED`` — never
|
|
440
|
+
raises into the caller, never poisons the cache."""
|
|
441
|
+
from genblaze_core.providers import DiscoveryStatus
|
|
442
|
+
|
|
443
|
+
provider = ReplicateProvider(api_token="test-token")
|
|
444
|
+
mock_client = MagicMock()
|
|
445
|
+
mock_client.models.list.side_effect = RuntimeError("boom")
|
|
446
|
+
provider._client = mock_client
|
|
447
|
+
|
|
448
|
+
result = provider.discover_models()
|
|
449
|
+
assert result.status is DiscoveryStatus.FAILED
|
|
450
|
+
assert "boom" in (result.detail or "")
|
|
451
|
+
|
|
452
|
+
|
|
453
|
+
def test_validate_model_authoritative_via_models_get():
|
|
454
|
+
"""A live slug returns ``OK_AUTHORITATIVE`` (source PROBE) via
|
|
455
|
+
``client.models.get()`` — the cheap-existence-check that Replicate
|
|
456
|
+
supports authoritatively."""
|
|
457
|
+
from genblaze_core.providers import ValidationOutcome, ValidationSource
|
|
458
|
+
|
|
459
|
+
provider = ReplicateProvider(api_token="test-token")
|
|
460
|
+
mock_client = MagicMock()
|
|
461
|
+
mock_client.models.get.return_value = _make_model("owner", "live-model")
|
|
462
|
+
provider._client = mock_client
|
|
463
|
+
|
|
464
|
+
result = provider.validate_model("owner/live-model")
|
|
465
|
+
assert result.outcome is ValidationOutcome.OK_AUTHORITATIVE
|
|
466
|
+
assert result.source is ValidationSource.PROBE
|
|
467
|
+
|
|
468
|
+
|
|
469
|
+
def test_validate_model_not_found_via_models_get():
|
|
470
|
+
"""A 404-class error from ``client.models.get()`` surfaces as
|
|
471
|
+
``NOT_FOUND``. The Pipeline preflight phase raises on this outcome
|
|
472
|
+
before any prediction is created."""
|
|
473
|
+
from genblaze_core.providers import ValidationOutcome, ValidationSource
|
|
474
|
+
|
|
475
|
+
provider = ReplicateProvider(api_token="test-token")
|
|
476
|
+
mock_client = MagicMock()
|
|
477
|
+
mock_client.models.get.side_effect = RuntimeError("404 model not found")
|
|
478
|
+
provider._client = mock_client
|
|
479
|
+
|
|
480
|
+
result = provider.validate_model("owner/dead-model")
|
|
481
|
+
assert result.outcome is ValidationOutcome.NOT_FOUND
|
|
482
|
+
assert result.source is ValidationSource.PROBE
|
|
483
|
+
|
|
484
|
+
|
|
485
|
+
def test_validate_model_caches_result():
|
|
486
|
+
"""Per-slug validation cache memoizes the per-slug GET so successive
|
|
487
|
+
Pipeline runs don't re-fetch."""
|
|
488
|
+
provider = ReplicateProvider(api_token="test-token")
|
|
489
|
+
mock_client = MagicMock()
|
|
490
|
+
mock_client.models.get.return_value = _make_model("owner", "x")
|
|
491
|
+
provider._client = mock_client
|
|
492
|
+
|
|
493
|
+
provider.validate_model("owner/x")
|
|
494
|
+
provider.validate_model("owner/x")
|
|
495
|
+
provider.validate_model("owner/x")
|
|
496
|
+
|
|
497
|
+
assert mock_client.models.get.call_count == 1
|
|
498
|
+
|
|
499
|
+
|
|
500
|
+
def test_validate_model_refresh_evicts_cache():
|
|
501
|
+
"""``refresh=True`` forces a fresh per-slug GET."""
|
|
502
|
+
provider = ReplicateProvider(api_token="test-token")
|
|
503
|
+
mock_client = MagicMock()
|
|
504
|
+
mock_client.models.get.return_value = _make_model("owner", "x")
|
|
505
|
+
provider._client = mock_client
|
|
506
|
+
|
|
507
|
+
provider.validate_model("owner/x")
|
|
508
|
+
provider.validate_model("owner/x", refresh=True)
|
|
509
|
+
assert mock_client.models.get.call_count == 2
|
|
510
|
+
|
|
511
|
+
|
|
512
|
+
def test_validate_model_user_registered_skips_probe():
|
|
513
|
+
"""A user-registered exact spec is authoritative without a network
|
|
514
|
+
round-trip — the per-slug probe should not fire."""
|
|
515
|
+
from genblaze_core.providers import ValidationOutcome, ValidationSource
|
|
516
|
+
|
|
517
|
+
provider = ReplicateProvider(api_token="test-token")
|
|
518
|
+
provider.models.register(ModelSpec(model_id="owner/local", modality=Modality.IMAGE))
|
|
519
|
+
mock_client = MagicMock()
|
|
520
|
+
provider._client = mock_client
|
|
521
|
+
|
|
522
|
+
result = provider.validate_model("owner/local")
|
|
523
|
+
assert result.outcome is ValidationOutcome.OK_AUTHORITATIVE
|
|
524
|
+
assert result.source is ValidationSource.USER
|
|
525
|
+
mock_client.models.get.assert_not_called()
|
|
526
|
+
|
|
527
|
+
|
|
360
528
|
# --- Compliance harness ---
|
|
361
529
|
|
|
362
530
|
|
|
@@ -1,227 +0,0 @@
|
|
|
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.providers.retry import retry_after_from_response
|
|
33
|
-
from genblaze_core.runnable.config import RunnableConfig
|
|
34
|
-
|
|
35
|
-
from ._errors import map_replicate_error
|
|
36
|
-
|
|
37
|
-
# Default per-second GPU-compute rate (Nvidia A40/T4 tier).
|
|
38
|
-
_COST_PER_SEC = 0.000225
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
def _compute_time_cost(ctx: PricingContext) -> float | None:
|
|
42
|
-
payload = ctx.provider_payload.get("replicate") if ctx.provider_payload else None
|
|
43
|
-
if not isinstance(payload, dict):
|
|
44
|
-
return None
|
|
45
|
-
predict_time = payload.get("predict_time")
|
|
46
|
-
if predict_time is None:
|
|
47
|
-
return None
|
|
48
|
-
try:
|
|
49
|
-
return float(predict_time) * _COST_PER_SEC
|
|
50
|
-
except (TypeError, ValueError):
|
|
51
|
-
return None
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
# Replicate has no enumerable model list; fallback spec applies to every model.
|
|
55
|
-
_FALLBACK_SPEC = ModelSpec(
|
|
56
|
-
model_id="*",
|
|
57
|
-
pricing=per_response_metric(_compute_time_cost),
|
|
58
|
-
input_mapping=route_by_media_type({"image": "image", "video": "video", "audio": "audio"}),
|
|
59
|
-
)
|
|
60
|
-
|
|
61
|
-
|
|
62
|
-
class ReplicateProvider(BaseProvider):
|
|
63
|
-
"""Provider adapter for Replicate (replicate.com)."""
|
|
64
|
-
|
|
65
|
-
name = "replicate"
|
|
66
|
-
|
|
67
|
-
@classmethod
|
|
68
|
-
def create_registry(cls) -> ModelRegistry:
|
|
69
|
-
# No enumerable defaults — every model uses the fallback spec.
|
|
70
|
-
return ModelRegistry(defaults={}, fallback=_FALLBACK_SPEC)
|
|
71
|
-
|
|
72
|
-
def get_capabilities(self) -> ProviderCapabilities:
|
|
73
|
-
"""Replicate: multi-modal generation depending on selected model."""
|
|
74
|
-
return ProviderCapabilities(
|
|
75
|
-
supported_modalities=[Modality.IMAGE, Modality.VIDEO, Modality.AUDIO],
|
|
76
|
-
supported_inputs=["text", "image", "video", "audio"],
|
|
77
|
-
accepts_chain_input=True,
|
|
78
|
-
)
|
|
79
|
-
|
|
80
|
-
def __init__(
|
|
81
|
-
self,
|
|
82
|
-
api_token: str | None = None,
|
|
83
|
-
poll_interval: float = 1.0,
|
|
84
|
-
http_timeout: float = 30.0,
|
|
85
|
-
*,
|
|
86
|
-
models: ModelRegistry | None = None,
|
|
87
|
-
):
|
|
88
|
-
super().__init__(models=models)
|
|
89
|
-
self.poll_interval = poll_interval
|
|
90
|
-
self._client: Any = None
|
|
91
|
-
self._api_token = api_token
|
|
92
|
-
self._http_timeout = http_timeout
|
|
93
|
-
|
|
94
|
-
def _get_client(self):
|
|
95
|
-
if self._client is None:
|
|
96
|
-
try:
|
|
97
|
-
import httpx
|
|
98
|
-
import replicate
|
|
99
|
-
|
|
100
|
-
timeout = httpx.Timeout(self._http_timeout, connect=10.0)
|
|
101
|
-
if self._api_token:
|
|
102
|
-
self._client = replicate.Client(
|
|
103
|
-
api_token=self._api_token, # noqa: S106
|
|
104
|
-
timeout=timeout,
|
|
105
|
-
)
|
|
106
|
-
else:
|
|
107
|
-
self._client = replicate.Client(timeout=timeout)
|
|
108
|
-
except ImportError as exc:
|
|
109
|
-
raise ProviderError(
|
|
110
|
-
"replicate package not installed. Run: pip install replicate"
|
|
111
|
-
) from exc
|
|
112
|
-
return self._client
|
|
113
|
-
|
|
114
|
-
def submit(self, step: Step, config: RunnableConfig | None = None) -> Any:
|
|
115
|
-
client = self._get_client()
|
|
116
|
-
try:
|
|
117
|
-
input_params = self.prepare_payload(step)
|
|
118
|
-
prediction = client.predictions.create(
|
|
119
|
-
model=step.model,
|
|
120
|
-
input=input_params,
|
|
121
|
-
)
|
|
122
|
-
return prediction.id
|
|
123
|
-
except Exception as exc:
|
|
124
|
-
raise ProviderError(
|
|
125
|
-
f"Replicate submit failed: {exc}",
|
|
126
|
-
error_code=map_replicate_error(exc),
|
|
127
|
-
retry_after=retry_after_from_response(exc),
|
|
128
|
-
) from exc
|
|
129
|
-
|
|
130
|
-
def poll(self, prediction_id: Any, config: RunnableConfig | None = None) -> bool:
|
|
131
|
-
client = self._get_client()
|
|
132
|
-
try:
|
|
133
|
-
prediction = client.predictions.get(prediction_id)
|
|
134
|
-
if prediction.status in ("succeeded", "failed", "canceled"):
|
|
135
|
-
self._cache_poll_result(prediction_id, prediction)
|
|
136
|
-
return True
|
|
137
|
-
return False
|
|
138
|
-
except Exception as exc:
|
|
139
|
-
raise ProviderError(
|
|
140
|
-
f"Replicate poll failed: {exc}",
|
|
141
|
-
error_code=map_replicate_error(exc),
|
|
142
|
-
retry_after=retry_after_from_response(exc),
|
|
143
|
-
) from exc
|
|
144
|
-
|
|
145
|
-
def fetch_output(self, prediction_id: Any, step: Step) -> Step:
|
|
146
|
-
client = self._get_client()
|
|
147
|
-
try:
|
|
148
|
-
prediction = self._get_cached_poll_result(prediction_id)
|
|
149
|
-
if prediction is None:
|
|
150
|
-
prediction = client.predictions.get(prediction_id)
|
|
151
|
-
|
|
152
|
-
# Capture predict_time so registry pricing can read it.
|
|
153
|
-
metrics = getattr(prediction, "metrics", None)
|
|
154
|
-
predict_time = getattr(metrics, "predict_time", None) if metrics else None
|
|
155
|
-
|
|
156
|
-
step.provider_payload = {
|
|
157
|
-
"replicate": {
|
|
158
|
-
"prediction_id": prediction.id,
|
|
159
|
-
"model": prediction.model if hasattr(prediction, "model") else None,
|
|
160
|
-
"version": prediction.version if hasattr(prediction, "version") else None,
|
|
161
|
-
"status": prediction.status,
|
|
162
|
-
"created_at": str(prediction.created_at)
|
|
163
|
-
if hasattr(prediction, "created_at")
|
|
164
|
-
else None,
|
|
165
|
-
"predict_time": predict_time,
|
|
166
|
-
}
|
|
167
|
-
}
|
|
168
|
-
|
|
169
|
-
if prediction.status == "failed":
|
|
170
|
-
error_msg = prediction.error or "Unknown error"
|
|
171
|
-
raise ProviderError(
|
|
172
|
-
error_msg,
|
|
173
|
-
error_code=map_replicate_error(error_msg),
|
|
174
|
-
)
|
|
175
|
-
|
|
176
|
-
if prediction.status == "canceled":
|
|
177
|
-
raise ProviderError(
|
|
178
|
-
"Prediction was canceled",
|
|
179
|
-
error_code=ProviderErrorCode.UNKNOWN,
|
|
180
|
-
)
|
|
181
|
-
|
|
182
|
-
# Replicate output shapes vary by model: str (single URL), list[str]
|
|
183
|
-
# (multi-asset), dict[str, str | list] (e.g. {"video": url,
|
|
184
|
-
# "subtitles": url} from text-to-video models with side-channels),
|
|
185
|
-
# or None (no output). Normalize to list[str].
|
|
186
|
-
raw_output = prediction.output
|
|
187
|
-
urls: list[str]
|
|
188
|
-
if raw_output is None:
|
|
189
|
-
urls = []
|
|
190
|
-
elif isinstance(raw_output, str):
|
|
191
|
-
urls = [raw_output]
|
|
192
|
-
elif isinstance(raw_output, list):
|
|
193
|
-
# Nested lists happen on batch-output models; flatten one level.
|
|
194
|
-
urls = []
|
|
195
|
-
for item in raw_output:
|
|
196
|
-
if isinstance(item, str):
|
|
197
|
-
urls.append(item)
|
|
198
|
-
elif isinstance(item, list):
|
|
199
|
-
urls.extend(str(u) for u in item if isinstance(u, (str, bytes)))
|
|
200
|
-
elif isinstance(raw_output, dict):
|
|
201
|
-
# Multi-channel outputs: keep only URL-shaped string values.
|
|
202
|
-
urls = [str(v) for v in raw_output.values() if isinstance(v, str)]
|
|
203
|
-
else:
|
|
204
|
-
raise ProviderError(
|
|
205
|
-
f"Unexpected Replicate output shape "
|
|
206
|
-
f"({type(raw_output).__name__}): {raw_output!r}",
|
|
207
|
-
error_code=ProviderErrorCode.SERVER_ERROR,
|
|
208
|
-
)
|
|
209
|
-
|
|
210
|
-
for url_str in urls:
|
|
211
|
-
validate_asset_url(url_str)
|
|
212
|
-
path = urlparse(url_str).path
|
|
213
|
-
mime, _ = mimetypes.guess_type(path)
|
|
214
|
-
if mime is None:
|
|
215
|
-
mime = f"{step.modality.value}/octet-stream"
|
|
216
|
-
step.assets.append(Asset(url=url_str, media_type=mime))
|
|
217
|
-
|
|
218
|
-
self._apply_registry_pricing(step)
|
|
219
|
-
return step
|
|
220
|
-
except ProviderError:
|
|
221
|
-
raise
|
|
222
|
-
except Exception as exc:
|
|
223
|
-
raise ProviderError(
|
|
224
|
-
f"Replicate fetch_output failed: {exc}",
|
|
225
|
-
error_code=map_replicate_error(exc),
|
|
226
|
-
retry_after=retry_after_from_response(exc),
|
|
227
|
-
) from exc
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|