diff --git a/CHANGELOG.md b/CHANGELOG.md index 5f4968d..ca1da45 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,18 @@ notes for each version. ### Added +- `client.models.run()` (and `AsyncModels.run`) now accept three optional + keyword-only params matching Comfy Router's own query params on this route: + `model_provider` (run the model on a specific alternate provider instead of + its current default), `strict_mode` (only meaningful with `model_provider`; + `False`, the default, translates the request body to/from the alternate + provider's own schema, `True` sends it through unmodified), and + `fallback_provider` (Router's own retry against the model's other + registered provider on a failure attributable to Router or the provider + tried; defaults ON, pass `False` to opt out). All three default to `None` + and are omitted from the request entirely when unset, so an existing caller + sees no change. See the README's "`models.run`" section. + - Every exception `client.models.run()` raises **for a failed call** now carries the `Idempotency-Key` it was made under, on `.idempotency_key` — the typed `RouterError` buckets, a `RouterError` whose `error_type` this version diff --git a/README.md b/README.md index ac04659..730e18c 100644 --- a/README.md +++ b/README.md @@ -422,6 +422,36 @@ Three things follow from that, and they are the whole contract of this method: partner's own API documents as the request body is what you pass here, so you can move between the partner's API and Router by changing the host. +**`model_provider`, `strict_mode` and `fallback_provider`** are three optional +keyword-only params, all `None` by default: + +```python +result = client.models.run( + "openai/gpt-image-2", + {"prompt": "a red circle"}, + model_provider="fal", # run this model on fal instead of its default provider + strict_mode=False, # (default) translate the body to/from fal's own schema + fallback_provider=True, # (default) retry once on Router's or fal's own failure +) +``` + +- **`model_provider`** picks a specific alternate provider for `model` instead + of its current default. Omitted (`None`, the default) is byte-for-byte + today's default-provider behavior. +- **`strict_mode`** is only meaningful together with `model_provider`. + `False` (Router's own default) translates `arguments` from `model`'s native + contract into the alternate provider's real schema; `True` sends + `arguments` through unmodified, so it must already be that provider's own + native shape. +- **`fallback_provider`** controls Router's own retry: on a failure + attributable to Router's own side or to the specific provider tried — never + to `arguments` or your account — Router retries once against the model's + other registered provider. This defaults ON (Router's own behavior when the + param is omitted); pass `False` to opt out. + +All three default to `None`, which omits the corresponding query param +entirely — a caller who never passes them gets exactly today's request. + `run` returns when the generation is **complete**. There is no submit step and nothing to poll: where the platform has to submit-and-poll an upstream provider, that happens server side inside this one call. The value you get back diff --git a/spec/router-openapi.yaml b/spec/router-openapi.yaml index fca146e..3dbf9fb 100644 --- a/spec/router-openapi.yaml +++ b/spec/router-openapi.yaml @@ -90,6 +90,9 @@ paths: - $ref: '#/components/parameters/RouterProvider' - $ref: '#/components/parameters/RouterModel' - $ref: '#/components/parameters/RouterIdempotencyKey' + - $ref: '#/components/parameters/ModelProvider' + - $ref: '#/components/parameters/StrictMode' + - $ref: '#/components/parameters/FallbackProvider' requestBody: required: true description: The partner model's native JSON input, forwarded to the provider unchanged. @@ -476,6 +479,22 @@ components: schema: $ref: '#/components/schemas/RouterErrorResponse' parameters: + FallbackProvider: + name: fallback_provider + in: query + required: false + description: 'Controls whether Router retries this call against the model''s OTHER registered provider when the FIRST attempt fails for a reason attributable to Router''s own side or to the specific provider tried - never for a reason attributable to the request itself (an unretried failure is refused exactly as it always was). Omitted, or any value other than `false`: fallback is ON (the default) and Router uses the one alternate the model has today. `false`: fallback is OFF - a failure is refused, never retried. A response where a fallback actually ran carries the `X-Comfy-Router-Fallback-Provider` header, naming the provider that ultimately served it, and no case retries a generation that may already have been submitted to a provider.' + schema: + type: string + example: fal + ModelProvider: + name: model_provider + in: query + required: false + description: Selects an alternate provider for this model, instead of its current default. Omitted, or `default`, is byte-for-byte today's behavior. A value naming a real provider that does not serve this model is refused `404 provider_not_available`; an unrecognized value is refused `422 validation_error`. + schema: + type: string + example: fal RouterCatalogCursor: name: cursor in: query @@ -516,6 +535,14 @@ components: description: Lowercase provider segment of the canonical `{provider}/{model}[/{variant}]` model ID - the partner whose model is being run. schema: $ref: '#/components/schemas/RouterProviderSegment' + StrictMode: + name: strict_mode + in: query + required: false + description: 'Only meaningful together with `model_provider`. `false` (the default): the request body must be this model''s own native contract, translated to the alternate provider''s real schema - any native field that cannot be expressed exactly is dropped and disclosed via the response''s `notes`, never silently. `true`: the body must already be the alternate provider''s own real schema, passed through unmodified in both directions - no translation, so no `notes`.' + schema: + type: boolean + default: false headers: CommittedSpendCurrentHeader: description: The USD cents the caller currently has committed to calls still in flight, not counting the refused call. Present alongside `X-Committed-Spend-Limit`. diff --git a/src/comfy_low/transport.py b/src/comfy_low/transport.py index 8b564d7..b17a920 100644 --- a/src/comfy_low/transport.py +++ b/src/comfy_low/transport.py @@ -39,7 +39,7 @@ from importlib.metadata import PackageNotFoundError from importlib.metadata import version as _pkg_version from typing import Any, BinaryIO -from urllib.parse import parse_qs, quote, urlsplit, urlunsplit +from urllib.parse import parse_qs, quote, urlencode, urlsplit, urlunsplit import httpx @@ -143,6 +143,10 @@ def model_run_request( model: str, arguments: Mapping[str, Any], idempotency_key: str | None, + *, + model_provider: str | None = None, + strict_mode: bool | None = None, + fallback_provider: bool | None = None, ) -> tuple[str, dict[str, Any], dict[str, str]]: """Sans-IO ``(path, json_body, headers)`` for one model run. @@ -160,11 +164,32 @@ def model_run_request( ``arguments`` is copied into a plain dict so any ``Mapping`` is accepted and the caller's object is never handed to the JSON encoder directly. + + ``model_provider``, ``strict_mode`` and ``fallback_provider`` are Router's + three query params on this route (``spec/router-openapi.yaml``'s + ``ModelProvider``, ``StrictMode``, ``FallbackProvider`` parameters) — + ``None`` (the default for all three) omits the param entirely rather than + sending an empty or ``"None"`` value, which is what lets a caller who never + heard of them get byte-for-byte today's request. ``fallback_provider`` is + the one with an inverted sense on the wire: Router defaults it ON, so this + only ever sends the query param when the caller passes ``False`` — sending + ``fallback_provider=true`` explicitly would be a no-op byte string, never + a distinct request the stub or a real deployment could tell apart from + omitting it. """ provider, name = parse_model_id(model) path = _MODEL_RUN_PATH_TEMPLATE.format( provider=quote(provider, safe=""), model=quote(name, safe="") ) + query: dict[str, str] = {} + if model_provider is not None: + query["model_provider"] = model_provider + if strict_mode is not None: + query["strict_mode"] = "true" if strict_mode else "false" + if fallback_provider is False: + query["fallback_provider"] = "false" + if query: + path = f"{path}?{urlencode(query)}" body: dict[str, Any] = dict(arguments) headers = {"Idempotency-Key": idempotency_key} if idempotency_key else {} return path, body, headers @@ -765,6 +790,9 @@ def post_model_run( *, idempotency_key: str | None = None, timeout: Any = MODEL_RUN_TIMEOUT, + model_provider: str | None = None, + strict_mode: bool | None = None, + fallback_provider: bool | None = None, ) -> dict[str, Any]: """POST ``{router_base_url}/v2/models/{provider}/{model}`` — awaited server-side. @@ -786,10 +814,21 @@ def post_model_run( ``spec/router-openapi.yaml``, hand-bound — see :data:`_MODEL_RUN_PATH_TEMPLATE`. + ``model_provider``, ``strict_mode`` and ``fallback_provider`` are + Router's three query params — see :func:`model_run_request`, which + builds them into the path; ``None`` (every default) omits all three. + Raises ``TypeError``/``ValueError`` from :func:`parse_model_id` before any request when ``model`` is not a ``{provider}/{model}`` id. """ - path, body, headers = model_run_request(model, arguments, idempotency_key) + path, body, headers = model_run_request( + model, + arguments, + idempotency_key, + model_provider=model_provider, + strict_mode=strict_mode, + fallback_provider=fallback_provider, + ) url = self._p.router_base_url + path resp = self.raw_request("POST", url, headers=headers, json=body, timeout=timeout) return self._p.parse_or_raise(resp, (200, 201)) @@ -1102,9 +1141,19 @@ async def post_model_run( *, idempotency_key: str | None = None, timeout: Any = MODEL_RUN_TIMEOUT, + model_provider: str | None = None, + strict_mode: bool | None = None, + fallback_provider: bool | None = None, ) -> dict[str, Any]: """Async :meth:`ComfyLow.post_model_run`.""" - path, body, headers = model_run_request(model, arguments, idempotency_key) + path, body, headers = model_run_request( + model, + arguments, + idempotency_key, + model_provider=model_provider, + strict_mode=strict_mode, + fallback_provider=fallback_provider, + ) url = self._p.router_base_url + path resp = await self.raw_request("POST", url, headers=headers, json=body, timeout=timeout) return self._p.parse_or_raise(resp, (200, 201)) diff --git a/src/comfy_sdk/models.py b/src/comfy_sdk/models.py index d79b9d9..ca45871 100644 --- a/src/comfy_sdk/models.py +++ b/src/comfy_sdk/models.py @@ -131,6 +131,9 @@ def run( *, idempotency_key: str | None = None, timeout: float | httpx.Timeout | None = MODEL_RUN_TIMEOUT, + model_provider: str | None = None, + strict_mode: bool | None = None, + fallback_provider: bool | None = None, ) -> dict[str, Any]: """Run ``model`` with ``arguments`` and return the completed result. @@ -223,6 +226,22 @@ def run( generation instead of collecting the first. Every exception *this* method raises carries a real key, but an ``except ComfyError`` that also catches errors from other surfaces can hand you one that does not. + + ``model_provider`` requests a specific alternate provider for + ``model`` instead of its current default (``?model_provider=`` on the + wire); omitted, or ``None``, is byte-for-byte today's default-provider + behavior. ``strict_mode`` is only meaningful together with + ``model_provider``: ``False`` (Router's own default) translates + ``arguments`` from ``model``'s native contract to the alternate + provider's real schema; ``True`` sends ``arguments`` through + unmodified, so it must already be that provider's own native shape. + ``fallback_provider`` is Router's own retry: on a failure attributable + to Router's own side or to the specific provider tried — never to + ``arguments`` or your account — Router retries once against the + model's other registered provider, and this defaults ON (``None`` or + ``True``); pass ``False`` to opt out. All three default to ``None``, + which omits the corresponding query param entirely rather than + sending an explicit "off" value Router would have to special-case. """ low = cast(ComfyLow, self._low) # Minted once, outside the loop: reusing this exact value on every @@ -250,7 +269,15 @@ def run( with translating(idempotency_key=key): while True: try: - return low.post_model_run(model, payload, idempotency_key=key, timeout=timeout) + return low.post_model_run( + model, + payload, + idempotency_key=key, + timeout=timeout, + model_provider=model_provider, + strict_mode=strict_mode, + fallback_provider=fallback_provider, + ) except _CANDIDATE_FAILURES as exc: delay = retrier.delay_before_retry(exc) if delay is None: @@ -272,14 +299,18 @@ async def run( *, idempotency_key: str | None = None, timeout: float | httpx.Timeout | None = MODEL_RUN_TIMEOUT, + model_provider: str | None = None, + strict_mode: bool | None = None, + fallback_provider: bool | None = None, ) -> dict[str, Any]: """Awaitable :meth:`Models.run` — same arguments, same result shape. This *is* the async form of ``run``: awaiting it on ``AsyncComfy`` is the whole difference from the sync client — including the model-id - rule, the retry policy, the one-key-per-call rule, and the - ``.idempotency_key`` every exception it raises carries for the replay. - See :meth:`Models.run`. + rule, the retry policy, the one-key-per-call rule, the + ``.idempotency_key`` every exception it raises carries for the replay, + and ``model_provider``/``strict_mode``/``fallback_provider``. See + :meth:`Models.run`. """ low = cast(AsyncComfyLow, self._low) key = ( @@ -300,7 +331,13 @@ async def run( while True: try: return await low.post_model_run( - model, payload, idempotency_key=key, timeout=timeout + model, + payload, + idempotency_key=key, + timeout=timeout, + model_provider=model_provider, + strict_mode=strict_mode, + fallback_provider=fallback_provider, ) except _CANDIDATE_FAILURES as exc: delay = retrier.delay_before_retry(exc) diff --git a/tests/conftest.py b/tests/conftest.py index 799351a..bf55b31 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -22,7 +22,7 @@ from dataclasses import dataclass, field from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Any -from urllib.parse import unquote +from urllib.parse import parse_qs, unquote import pytest @@ -197,9 +197,13 @@ class ServerState: # asserts the id the caller passed rather than a particular encoding of it. last_model_run_provider: str | None = None last_model_run_model: str | None = None - # ...and the raw, still-encoded request path, for the tests that are about - # the encoding itself. + # ...and the raw, still-encoded request path (query string stripped), for + # the tests that are about the encoding itself. last_model_run_path: str | None = None + # The last model run's query string, decoded to one value per key (see + # `_post_model_run`'s own note on why `parse_qs`'s list form is not kept). + # `None` before any run; `{}` after a run that carried none. + last_model_run_query: dict[str, str] | None = None # Every Idempotency-Key seen on a model run, in arrival order (`None` # records a run that arrived without the header at all). model_run_idempotency_keys: list[str | None] = field(default_factory=list) @@ -498,10 +502,14 @@ def do_POST(self) -> None: # `/api/v2` paths above (a different host in production; the same # stub here, with `COMFY_ROUTER_BASE_URL` pointed at it). The two # segments are the model id, so they are matched rather than - # compared to a fixed string. - m = re.match(r"/v2/models/([^/]+)/([^/]+)$", self.path) + # compared to a fixed string. Split off the query string before + # matching: `model_provider`/`strict_mode`/`fallback_provider` ride + # in it, and `[^/]+` would otherwise swallow a `?...` suffix into + # the model segment rather than leaving it for `_post_model_run`. + path_only, _, query_string = self.path.partition("?") + m = re.match(r"/v2/models/([^/]+)/([^/]+)$", path_only) if m: - self._post_model_run(m.group(1), m.group(2)) + self._post_model_run(m.group(1), m.group(2), query_string) return m = re.match(r"/api/v2/jobs/([^/]+)/cancel$", self.path) if m: @@ -527,7 +535,7 @@ def _post_from_hash(self) -> None: else: self._err(404, "blob_not_found", "no such blob") - def _post_model_run(self, provider: str, model: str) -> None: + def _post_model_run(self, provider: str, model: str, query_string: str) -> None: state.model_run_count += 1 # Decoded, because the SDK percent-encodes each segment and a real # origin server decodes it before routing — asserting the encoded @@ -536,7 +544,12 @@ def _post_model_run(self, provider: str, model: str) -> None: # for the tests that are about the encoding. state.last_model_run_provider = unquote(provider) state.last_model_run_model = unquote(model) - state.last_model_run_path = self.path + state.last_model_run_path = self.path.partition("?")[0] + # `parse_qs` drops a key with no value at all, which never happens + # here — model_run_request only ever adds a key with a real value — + # so a plain single-valued dict is the faithful, easy-to-assert + # shape rather than parse_qs's list-per-key one. + state.last_model_run_query = {k: v[0] for k, v in parse_qs(query_string).items()} state.last_model_run_body = json.loads(self._read_body() or b"{}") key = self.headers.get("Idempotency-Key") state.model_run_idempotency_keys.append(key) diff --git a/tests/test_models_run.py b/tests/test_models_run.py index 9dbb2c5..b265a2a 100644 --- a/tests/test_models_run.py +++ b/tests/test_models_run.py @@ -216,6 +216,99 @@ def test_run_accepts_any_mapping_and_does_not_alias_the_callers_object(server) - assert caller_args == {"prompt": "a dog"} +# --- model_provider / strict_mode / fallback_provider -------------------- +# +# Comfy Router's three query params on this route (`spec/router-openapi.yaml`'s +# `ModelProvider`, `StrictMode`, `FallbackProvider` parameters). All default to +# `None` on `run`, which must omit the corresponding query param entirely — a +# caller who never heard of them gets byte-for-byte today's plain request, +# which `test_run_addresses_the_model_by_path_and_sends_the_native_body` above +# already pins for the fully-omitted case. + + +def test_omitting_all_three_sends_no_query_string_at_all(server) -> None: + with Comfy() as client: + client.models.run(MODEL, ARGS) + assert server.state.last_model_run_query == {} + + +def test_model_provider_is_sent_verbatim(server) -> None: + with Comfy() as client: + client.models.run(MODEL, ARGS, model_provider="fal") + assert server.state.last_model_run_query == {"model_provider": "fal"} + + +def test_strict_mode_true_and_false_are_both_sent_explicitly(server) -> None: + with Comfy() as client: + client.models.run(MODEL, ARGS, model_provider="fal", strict_mode=True) + assert server.state.last_model_run_query == {"model_provider": "fal", "strict_mode": "true"} + + with Comfy() as client: + client.models.run(MODEL, ARGS, model_provider="fal", strict_mode=False) + assert server.state.last_model_run_query == {"model_provider": "fal", "strict_mode": "false"} + + +def test_fallback_provider_false_opts_out_explicitly(server) -> None: + # The one param with an inverted sense: Router defaults it ON, so `False` + # is the only value this SDK ever has a reason to put on the wire. + with Comfy() as client: + client.models.run(MODEL, ARGS, fallback_provider=False) + assert server.state.last_model_run_query == {"fallback_provider": "false"} + + +def test_fallback_provider_true_is_the_same_as_omitting_it(server) -> None: + # `True` is Router's own default, so sending it explicitly would be a + # distinct wire value nothing downstream could tell apart from omitting it + # — asserted here so this stays a deliberate no-op rather than a forgotten + # branch. + with Comfy() as client: + client.models.run(MODEL, ARGS, fallback_provider=True) + assert server.state.last_model_run_query == {} + + +async def test_the_async_client_sends_the_same_three_params(server) -> None: + async with AsyncComfy() as client: + await client.models.run( + MODEL, ARGS, model_provider="fal", strict_mode=True, fallback_provider=False + ) + assert server.state.last_model_run_query == { + "model_provider": "fal", + "strict_mode": "true", + "fallback_provider": "false", + } + + +def test_the_sans_io_builder_agrees_with_the_wire_for_all_three(server) -> None: + # The one place the query-string shape is decided, asserted directly + # against the same real HTTP round trip the tests above exercise — + # mirrors test_the_sans_io_request_builder_agrees_with_the_wire above. + path, _body, _headers = model_run_request( + MODEL, ARGS, None, model_provider="fal", strict_mode=True, fallback_provider=False + ) + assert path == ( + "/v2/models/acme/flux-dev?model_provider=fal&strict_mode=true&fallback_provider=false" + ) + + with Comfy() as client: + client.models.run( + MODEL, ARGS, model_provider="fal", strict_mode=True, fallback_provider=False + ) + assert server.state.last_model_run_path == "/v2/models/acme/flux-dev" + assert server.state.last_model_run_query == { + "model_provider": "fal", + "strict_mode": "true", + "fallback_provider": "false", + } + + +def test_a_run_with_these_params_still_returns_the_native_result(server) -> None: + # The three params change what Router does server side, never the shape of + # what `run` hands back — still the provider's own payload, unwrapped. + with Comfy() as client: + result = client.models.run(MODEL, ARGS, model_provider="fal") + assert result == server.state.model_run_result + + # --- which host the run is addressed to --------------------------------- @@ -495,6 +588,9 @@ def post_model_run( *, idempotency_key: str | None = None, timeout: Any = None, + model_provider: str | None = None, + strict_mode: bool | None = None, + fallback_provider: bool | None = None, ) -> dict[str, Any]: self.keys.append(idempotency_key) raise self._exc @@ -508,6 +604,9 @@ async def post_model_run( # type: ignore[override] *, idempotency_key: str | None = None, timeout: Any = None, + model_provider: str | None = None, + strict_mode: bool | None = None, + fallback_provider: bool | None = None, ) -> dict[str, Any]: self.keys.append(idempotency_key) raise self._exc diff --git a/tests/test_models_run_retry.py b/tests/test_models_run_retry.py index 7f646b5..d7a2ed0 100644 --- a/tests/test_models_run_retry.py +++ b/tests/test_models_run_retry.py @@ -108,6 +108,9 @@ def post_model_run( *, idempotency_key: str | None = None, timeout: Any = None, + model_provider: str | None = None, + strict_mode: bool | None = None, + fallback_provider: bool | None = None, ) -> dict[str, Any]: return self._attempt(arguments, idempotency_key) @@ -120,6 +123,9 @@ async def post_model_run( # type: ignore[override] *, idempotency_key: str | None = None, timeout: Any = None, + model_provider: str | None = None, + strict_mode: bool | None = None, + fallback_provider: bool | None = None, ) -> dict[str, Any]: return self._attempt(arguments, idempotency_key)