Skip to content

Commit fc0b718

Browse files
committed
Rename block to wait
Update type definitions to capture expectations of wait parameter in run
1 parent 7f81a4d commit fc0b718

File tree

3 files changed

+13
-12
lines changed

3 files changed

+13
-12
lines changed

index.d.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -156,7 +156,7 @@ declare module "replicate" {
156156
identifier: `${string}/${string}` | `${string}/${string}:${string}`,
157157
options: {
158158
input: object;
159-
wait?: { interval?: number };
159+
wait?: boolean | number | { mode?: "poll"; interval?: number };
160160
webhook?: string;
161161
webhook_events_filter?: WebhookEventType[];
162162
signal?: AbortSignal;
@@ -189,6 +189,7 @@ declare module "replicate" {
189189
wait(
190190
prediction: Prediction,
191191
options?: {
192+
mode?: "poll";
192193
interval?: number;
193194
},
194195
stop?: (prediction: Prediction) => Promise<boolean>

index.js

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -144,11 +144,7 @@ class Replicate {
144144
* @returns {Promise<object>} - Resolves with the output of running the model
145145
*/
146146
async run(ref, options, progress) {
147-
let { block } = options;
148-
const { wait, signal, block: _, ...data } = options;
149-
150-
// Block if `block` is explicitly true or if `wait` is explicitly true
151-
block = block || (block === undefined && wait === true);
147+
const { wait, signal, ...data } = options;
152148

153149
const identifier = ModelVersionIdentifier.parse(ref);
154150

@@ -157,13 +153,13 @@ class Replicate {
157153
prediction = await this.predictions.create({
158154
...data,
159155
version: identifier.version,
160-
block,
156+
wait: wait,
161157
});
162158
} else if (identifier.owner && identifier.name) {
163159
prediction = await this.predictions.create({
164160
...data,
165161
model: `${identifier.owner}/${identifier.name}`,
166-
block,
162+
wait: wait,
167163
});
168164
} else {
169165
throw new Error("Invalid model version identifier");

lib/predictions.js

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,11 @@ const { transformFileInputs } = require("./util");
99
* @param {object} options.input - Required. An object with the model inputs
1010
* @param {string} [options.webhook] - An HTTPS URL for receiving a webhook when the prediction has new output
1111
* @param {string[]} [options.webhook_events_filter] - You can change which events trigger webhook requests by specifying webhook events (`start`|`output`|`logs`|`completed`)
12-
* @param {boolean} [options.block] - Whether to wait until the prediction is completed before returning. Defaults to false
12+
* @param {boolean|integer} [options.wait] - Whether to wait until the prediction is completed before returning. If an integer is provided, it will wait for that many seconds. Defaults to false
1313
* @returns {Promise<object>} Resolves with the created prediction
1414
*/
1515
async function createPrediction(options) {
16-
const { model, version, input, block, ...data } = options;
16+
const { model, version, input, wait, ...data } = options;
1717

1818
if (data.webhook) {
1919
try {
@@ -25,8 +25,12 @@ async function createPrediction(options) {
2525
}
2626

2727
const headers = {};
28-
if (block) {
29-
headers["Prefer"] = "wait";
28+
if (wait) {
29+
if (typeof wait === "number") {
30+
headers["Prefer"] = `wait=${wait}`;
31+
} else {
32+
headers["Prefer"] = "wait";
33+
}
3034
}
3135

3236
let response;

0 commit comments

Comments
 (0)