genblaze-replicate 0.3.2__tar.gz → 0.3.3__tar.gz

{genblaze_replicate-0.3.2 → genblaze_replicate-0.3.3}/.gitignore

@@ -54,7 +54,12 @@ ehthumbs.db
 !.env.example
 *.pem
 *.key
-credentials.json
+credentials*
+*.credentials
+*_rsa
+*.p12
+*.pfx
+*.keystore
 
 # Jupyter
 .ipynb_checkpoints/

{genblaze_replicate-0.3.2 → genblaze_replicate-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: genblaze-replicate
-Version: 0.3.2
+Version: 0.3.3
 Summary: Replicate provider adapter for genblaze
 Project-URL: Homepage, https://github.com/backblaze-labs/genblaze
 Project-URL: Documentation, https://github.com/backblaze-labs/genblaze
@@ -48,7 +48,11 @@ Supports any model hosted on Replicate — including:
 - **Video** — text-to-video and image-to-video models on Replicate
 - **Audio** — music and speech models on Replicate
 
-Use the full `owner/model` slug from https://replicate.com/explore.
+Use the full `owner/model` slug from https://replicate.com/explore — official
+and community models alike. `submit()` runs each via the right Replicate
+endpoint: community models resolve to their latest published version (cached
+per-slug), official models run version-less. Pin an exact community version
+with `owner/model:<version-hash>`.
 
 ## Install
 

{genblaze_replicate-0.3.2 → genblaze_replicate-0.3.3}/README.md

@@ -21,7 +21,11 @@ Supports any model hosted on Replicate — including:
 - **Video** — text-to-video and image-to-video models on Replicate
 - **Audio** — music and speech models on Replicate
 
-Use the full `owner/model` slug from https://replicate.com/explore.
+Use the full `owner/model` slug from https://replicate.com/explore — official
+and community models alike. `submit()` runs each via the right Replicate
+endpoint: community models resolve to their latest published version (cached
+per-slug), official models run version-less. Pin an exact community version
+with `owner/model:<version-hash>`.
 
 ## Install
 

{genblaze_replicate-0.3.2 → genblaze_replicate-0.3.3}/genblaze_replicate/provider.py

@@ -72,6 +72,12 @@ _FALLBACK_SPEC = ModelSpec(
     input_mapping=route_by_media_type({"image": "image", "video": "video", "audio": "audio"}),
 )
 
+# Sentinel distinguishing "slug not yet resolved" from "slug resolved to no
+# version" in ``_version_cache``. The latter caches ``None`` — a real result
+# meaning "official model, submit via the ``model=`` path" — so a plain
+# ``dict.get(slug)`` can't tell a miss from a cached official model.
+_VERSION_UNRESOLVED: Any = object()
+
 
 class ReplicateProvider(BaseProvider):
     """Provider adapter for Replicate (replicate.com)."""
@@ -131,6 +137,19 @@ class ReplicateProvider(BaseProvider):
         self._validation_cache: dict[str, tuple[float, ValidationResult]] = {}
         self._validation_lock = threading.RLock()
         self._validation_ttl: float = DEFAULT_TTL_SECONDS
+        # Per-slug submit-time resolution cache: ``owner/name`` → version
+        # hash, or ``None`` for official/versionless models that run via the
+        # ``model=`` path (see ``_resolve_version``). Deliberately *not*
+        # TTL-bounded like ``_validation_cache``: genblaze runs a process per
+        # pipeline, so pinning the slug→version resolved at first sight gives
+        # version consistency across a run (call ``validate_model(refresh=
+        # True)`` to re-resolve a newly-published *latest* version). Kept as a
+        # small purpose-built dict rather than reusing the base ``_probe_cache``
+        # machinery — the value is a bare hash, not a ``LiveProbeResult``,
+        # resolution is a cheap idempotent token-free GET, and it is seeded for
+        # free from ``validate_model``'s probe. Guarded by the same lock as the
+        # validation cache since both are populated from one ``models.get()``.
+        self._version_cache: dict[str, str | None] = {}
         # Wire a discovery cache for the first-page snapshot. The fetcher
         # closes over ``self`` so it picks up the (possibly-late-bound)
         # ``self._client``.
@@ -207,8 +226,14 @@ class ReplicateProvider(BaseProvider):
             return base_result
 
         if refresh:
+            slug, _ = self._split_model_id(model_id)
             with self._validation_lock:
                 self._validation_cache.pop(model_id, None)
+                # Keep the version cache in lockstep — a caller explicitly
+                # asking to refresh a slug's validation almost always wants
+                # a fresh "latest_version" too (e.g. the model owner just
+                # published a new version).
+                self._version_cache.pop(slug, None)
         else:
             cached = self._lookup_validation_cache(model_id)
             if cached is not None:
@@ -242,9 +267,13 @@ class ReplicateProvider(BaseProvider):
         * Other errors → ``UNKNOWN_PERMISSIVE`` (we couldn't verify; let
           the upstream call surface the real error if there is one).
         """
+        # Strip any inline ``:version`` — Replicate's model-existence GET
+        # (/v1/models/{owner}/{name}) only accepts the bare slug; a compound
+        # ``owner/name:hash`` key 404s even when the model exists.
+        slug, _ = self._split_model_id(model_id)
         try:
             client = self._get_client()
-            client.models.get(model_id)
+            model = client.models.get(slug)
         except Exception as exc:
             err_code = map_replicate_error(exc)
             if err_code is ProviderErrorCode.MODEL_ERROR:
@@ -258,11 +287,78 @@ class ReplicateProvider(BaseProvider):
                 exc,
             )
             return ValidationResult.unknown_permissive(detail=f"probe inconclusive: {exc}")
+        # This probe already fetched the model object that ``submit()`` needs
+        # to resolve a version — seed the version cache so a Pipeline run
+        # (preflight validate_model, then submit) never fetches it twice.
+        self._cache_model_version(slug, model)
         return ValidationResult.ok_authoritative(
             ValidationSource.PROBE,
             detail="confirmed via /v1/models/{owner}/{name}",
         )
 
+    @staticmethod
+    def _split_model_id(model_id: str) -> tuple[str, str | None]:
+        """Split ``owner/name[:version-hash]`` into ``(slug, inline_version)``.
+
+        Replicate slugs never contain a bare ``:`` themselves, so the first
+        colon (if any) unambiguously separates an inline-pinned version hash.
+        Shared by every call site that talks to ``/v1/models/{owner}/{name}``
+        (which rejects the compound form) or that needs a version hash.
+        """
+        slug, _, inline_version = model_id.partition(":")
+        return slug, (inline_version or None)
+
+    def _cache_model_version(self, slug: str, model: Any) -> str | None:
+        """Extract and cache ``model.latest_version.id`` for ``slug``.
+
+        Shared by ``_authoritative_lookup`` (validate_model's per-slug probe)
+        and ``_resolve_version`` (submit-time resolution) so a single
+        ``models.get()`` round-trip seeds both caches. ``slug`` must already
+        be stripped of any inline version (see ``_split_model_id``). Returns —
+        and caches — the version hash, or ``None`` for official/versionless
+        models. Caching the ``None`` result too means a fan-out of submits over
+        an official model probes once rather than re-resolving per prediction.
+        """
+        latest_version = getattr(model, "latest_version", None)
+        version_id = getattr(latest_version, "id", None) or None
+        with self._validation_lock:
+            self._version_cache[slug] = version_id
+        return version_id
+
+    def _resolve_version(self, model_id: str) -> str | None:
+        """Resolve ``model_id`` to a version hash for prediction submission.
+
+        Replicate has two kinds of runnable models submitted via two different
+        endpoints, so resolution returns either a hash or ``None`` (#109):
+
+        * **Community models** carry published versions and 404 on the
+          ``predictions.create(model=<slug>)`` path — they must be run via
+          ``predictions.create(version=<hash>)``. Returns their hash.
+        * **Official models** (e.g. ``black-forest-labs/flux-schnell``) expose
+          *no* version and can only be run via the ``model=<slug>`` path.
+          Returns ``None`` to tell ``submit()`` to take that fallback.
+
+        Resolution order:
+
+        1. Inline version (``owner/name:hash``) — already a hash; no network.
+        2. Cached resolution from a prior ``submit()``/``validate_model()`` for
+           this slug (``None`` is a valid cached "official model" result, hence
+           the ``_VERSION_UNRESOLVED`` sentinel to detect a genuine miss).
+        3. Otherwise ``client.models.get(slug)`` — one round-trip, cached
+           per-slug (hash or ``None``) thereafter.
+        """
+        slug, inline_version = self._split_model_id(model_id)
+        if inline_version:
+            return inline_version
+
+        cached = self._version_cache.get(slug, _VERSION_UNRESOLVED)
+        if cached is not _VERSION_UNRESOLVED:
+            return cached
+
+        client = self._get_client()
+        model = client.models.get(slug)
+        return self._cache_model_version(slug, model)
+
     def _get_client(self):
         if self._client is None:
             try:
@@ -287,11 +383,25 @@ class ReplicateProvider(BaseProvider):
         client = self._get_client()
         try:
             input_params = self.prepare_payload(step)
-            prediction = client.predictions.create(
-                model=step.model,
-                input=input_params,
-            )
+            version = self._resolve_version(step.model)
+            if version is not None:
+                # Community model — run the resolved version via /v1/predictions.
+                prediction = client.predictions.create(
+                    version=version,
+                    input=input_params,
+                )
+            else:
+                # Official/versionless model — only runnable via the model=
+                # path (/v1/models/{owner}/{name}/predictions). Pass the bare
+                # slug; that endpoint rejects a compound ``owner/name:hash``.
+                slug, _ = self._split_model_id(step.model)
+                prediction = client.predictions.create(
+                    model=slug,
+                    input=input_params,
+                )
             return prediction.id
+        except ProviderError:
+            raise
         except Exception as exc:
             raise ProviderError(
                 f"Replicate submit failed: {exc}",

{genblaze_replicate-0.3.2 → genblaze_replicate-0.3.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "genblaze-replicate"
-version = "0.3.2"
+version = "0.3.3"
 description = "Replicate provider adapter for genblaze"
 authors = [{name = "Jeronimo De Leon", email = "jdeleon@backblaze.com"}]
 readme = "README.md"

{genblaze_replicate-0.3.2 → genblaze_replicate-0.3.3}/tests/test_replicate_provider.py

@@ -34,6 +34,16 @@ def _make_step(**kwargs) -> Step:
     return Step(**defaults)
 
 
+def _make_model_with_version(version_id: str | None):
+    """A fake ``replicate.model.Model`` exposing a concrete ``latest_version``."""
+    model = MagicMock()
+    if version_id is None:
+        model.latest_version = None
+    else:
+        model.latest_version = MagicMock(id=version_id)
+    return model
+
+
 # --- map_replicate_error tests ---
 
 
@@ -76,6 +86,7 @@ def test_map_error_accepts_exception():
 def test_submit_sends_prompt_and_params():
     provider = ReplicateProvider(api_token="test-token")
     mock_client = MagicMock()
+    mock_client.models.get.return_value = _make_model_with_version("v-resolved")
     mock_client.predictions.create.return_value = FakePrediction()
     provider._client = mock_client
 
@@ -83,7 +94,7 @@ def test_submit_sends_prompt_and_params():
     result = provider.submit(step)
 
     mock_client.predictions.create.assert_called_once_with(
-        model="test/model",
+        version="v-resolved",
         input={"width": 1024, "prompt": "a cat"},
     )
     assert result == "pred-abc123"
@@ -92,6 +103,7 @@ def test_submit_sends_prompt_and_params():
 def test_submit_includes_negative_prompt():
     provider = ReplicateProvider(api_token="test-token")
     mock_client = MagicMock()
+    mock_client.models.get.return_value = _make_model_with_version("v-resolved")
     mock_client.predictions.create.return_value = FakePrediction()
     provider._client = mock_client
 
@@ -102,6 +114,168 @@ def test_submit_includes_negative_prompt():
     assert call_input["negative_prompt"] == "blurry"
 
 
+# --- version resolution (community-model 404 fix, #109) ---
+
+
+def test_submit_resolves_community_model_slug_to_version():
+    """A bare ``owner/name`` slug (the common case — community models) must
+    resolve to a version hash and submit via ``version=``, not ``model=``.
+    ``predictions.create(model=...)`` only resolves official models and 404s
+    for the vast majority of Replicate's public catalog (#109)."""
+    provider = ReplicateProvider(api_token="test-token")
+    mock_client = MagicMock()
+    mock_client.models.get.return_value = _make_model_with_version("abc123hash")
+    mock_client.predictions.create.return_value = FakePrediction()
+    provider._client = mock_client
+
+    step = _make_step(model="sczhou/codeformer")
+    provider.submit(step)
+
+    mock_client.models.get.assert_called_once_with("sczhou/codeformer")
+    mock_client.predictions.create.assert_called_once_with(
+        version="abc123hash",
+        input={"prompt": "a cat"},
+    )
+
+
+def test_submit_uses_inline_version_without_network_call():
+    """``owner/name:hash`` carries the version inline — no ``models.get()``
+    round-trip needed to resolve it."""
+    provider = ReplicateProvider(api_token="test-token")
+    mock_client = MagicMock()
+    mock_client.predictions.create.return_value = FakePrediction()
+    provider._client = mock_client
+
+    step = _make_step(model="sczhou/codeformer:deadbeef1234")
+    provider.submit(step)
+
+    mock_client.models.get.assert_not_called()
+    mock_client.predictions.create.assert_called_once_with(
+        version="deadbeef1234",
+        input={"prompt": "a cat"},
+    )
+
+
+def test_submit_caches_resolved_version_across_calls():
+    """The version resolution is cached per-slug so repeated submits (e.g.
+    across a batch run) don't re-issue ``models.get()`` per prediction."""
+    provider = ReplicateProvider(api_token="test-token")
+    mock_client = MagicMock()
+    mock_client.models.get.return_value = _make_model_with_version("cached-version")
+    mock_client.predictions.create.return_value = FakePrediction()
+    provider._client = mock_client
+
+    provider.submit(_make_step(model="sczhou/codeformer"))
+    provider.submit(_make_step(model="sczhou/codeformer"))
+
+    assert mock_client.models.get.call_count == 1
+    assert mock_client.predictions.create.call_count == 2
+
+
+def test_submit_falls_back_to_model_path_for_official_versionless_model():
+    """Official Replicate models (e.g. ``black-forest-labs/flux-schnell``)
+    expose no version and are only runnable via ``predictions.create(model=…)``
+    (POST /v1/models/{owner}/{name}/predictions). When ``models.get`` reports
+    no ``latest_version``, submit() must fall back to the ``model=`` path with
+    the bare slug rather than raise — raising would break every official model,
+    including several this connector documents as supported (#109)."""
+    provider = ReplicateProvider(api_token="test-token")
+    mock_client = MagicMock()
+    mock_client.models.get.return_value = _make_model_with_version(None)
+    mock_client.predictions.create.return_value = FakePrediction()
+    provider._client = mock_client
+
+    result = provider.submit(_make_step(model="black-forest-labs/flux-schnell"))
+
+    mock_client.predictions.create.assert_called_once_with(
+        model="black-forest-labs/flux-schnell",
+        input={"prompt": "a cat"},
+    )
+    # A versionless official model must NOT be submitted with a version= kwarg.
+    assert "version" not in mock_client.predictions.create.call_args.kwargs
+    assert result == "pred-abc123"
+
+
+def test_submit_caches_versionless_resolution_across_calls():
+    """The "official / no version" decision is cached per-slug just like a
+    resolved hash — a fan-out of submits over an official model issues one
+    ``models.get`` probe, not one per prediction."""
+    provider = ReplicateProvider(api_token="test-token")
+    mock_client = MagicMock()
+    mock_client.models.get.return_value = _make_model_with_version(None)
+    mock_client.predictions.create.return_value = FakePrediction()
+    provider._client = mock_client
+
+    provider.submit(_make_step(model="black-forest-labs/flux-schnell"))
+    provider.submit(_make_step(model="black-forest-labs/flux-schnell"))
+
+    assert mock_client.models.get.call_count == 1
+    assert mock_client.predictions.create.call_count == 2
+
+
+def test_validate_model_seeds_version_cache_for_submit():
+    """``validate_model``'s per-slug probe and ``submit``'s version resolution
+    share one ``models.get()`` round-trip per slug — the pipeline preflight
+    phase already fetches the model, so submit() shouldn't fetch it again."""
+    provider = ReplicateProvider(api_token="test-token")
+    mock_client = MagicMock()
+    mock_client.models.get.return_value = _make_model_with_version("seeded-version")
+    mock_client.predictions.create.return_value = FakePrediction()
+    provider._client = mock_client
+
+    provider.validate_model("sczhou/codeformer")
+    provider.submit(_make_step(model="sczhou/codeformer"))
+
+    assert mock_client.models.get.call_count == 1
+    mock_client.predictions.create.assert_called_once_with(
+        version="seeded-version",
+        input={"prompt": "a cat"},
+    )
+
+
+def test_validate_model_strips_inline_version_before_probe():
+    """A ``Pipeline.run()`` preflight validates ``step.model`` verbatim,
+    including any inline ``:version`` suffix. Replicate's model-existence
+    GET (``/v1/models/{owner}/{name}``) only accepts the bare slug — a
+    compound ``owner/name:hash`` key 404s even when the model exists, which
+    would otherwise abort the run with a false "model not found" before
+    ``submit()`` ever runs. ``validate_model`` must probe the bare slug."""
+    from genblaze_core.providers import ValidationOutcome
+
+    provider = ReplicateProvider(api_token="test-token")
+    mock_client = MagicMock()
+    mock_client.models.get.return_value = _make_model_with_version("resolved-hash")
+    provider._client = mock_client
+
+    result = provider.validate_model("sczhou/codeformer:deadbeef1234")
+
+    mock_client.models.get.assert_called_once_with("sczhou/codeformer")
+    assert result.outcome is ValidationOutcome.OK_AUTHORITATIVE
+
+
+def test_validate_model_refresh_also_evicts_version_cache():
+    """``refresh=True`` is a request to re-check everything about a slug —
+    it must also drop any cached version resolution, not just the
+    validation-outcome cache, so a subsequent ``submit()`` re-resolves
+    ``latest_version`` instead of silently reusing a stale hash."""
+    provider = ReplicateProvider(api_token="test-token")
+    mock_client = MagicMock()
+    mock_client.models.get.return_value = _make_model_with_version("v1")
+    mock_client.predictions.create.return_value = FakePrediction()
+    provider._client = mock_client
+
+    provider.validate_model("sczhou/codeformer")
+    provider.submit(_make_step(model="sczhou/codeformer"))
+    assert mock_client.predictions.create.call_args.kwargs["version"] == "v1"
+
+    # Model owner publishes a new version; caller refreshes.
+    mock_client.models.get.return_value = _make_model_with_version("v2")
+    provider.validate_model("sczhou/codeformer", refresh=True)
+    provider.submit(_make_step(model="sczhou/codeformer"))
+
+    assert mock_client.predictions.create.call_args.kwargs["version"] == "v2"
+
+
 # --- poll tests ---
 
 
@@ -563,6 +737,10 @@ class TestReplicateCompliance(ProviderComplianceTests):
     def make_provider(self):
         provider = ReplicateProvider(api_token="test-token")
         mock_client = MagicMock()
+        # submit() resolves a version via models.get().latest_version.id
+        # before creating a prediction — stub it explicitly rather than
+        # relying on MagicMock's auto-vivified (but opaque) attribute chain.
+        mock_client.models.get.return_value = _make_model_with_version("compliance-version")
         mock_client.predictions.create.return_value = FakePrediction()
         mock_client.predictions.get.return_value = FakePrediction()
         provider._client = mock_client