Skip to content

Commit 42c673b

Browse files
committed
Fix the run.wait parameter
This commit fixes the use of the `run.wait` parameter to only pass it to the `prediction.create` function if it's a number or integer. We now check to see if the status of the object returned has moved out of a starting state.
1 parent 9fb5213 commit 42c673b

File tree

1 file changed

+22
-18
lines changed

1 file changed

+22
-18
lines changed

index.js

+22-18
Original file line numberDiff line numberDiff line change
@@ -147,19 +147,20 @@ class Replicate {
147147
const { wait, signal, ...data } = options;
148148

149149
const identifier = ModelVersionIdentifier.parse(ref);
150+
const isBlocking = typeof wait === "boolean" || typeof wait === "number";
150151

151152
let prediction;
152153
if (identifier.version) {
153154
prediction = await this.predictions.create({
154155
...data,
155156
version: identifier.version,
156-
wait: wait,
157+
wait: isBlocking ? wait : false,
157158
});
158159
} else if (identifier.owner && identifier.name) {
159160
prediction = await this.predictions.create({
160161
...data,
161162
model: `${identifier.owner}/${identifier.name}`,
162-
wait: wait,
163+
wait: isBlocking ? wait : false,
163164
});
164165
} else {
165166
throw new Error("Invalid model version identifier");
@@ -170,23 +171,26 @@ class Replicate {
170171
progress(prediction);
171172
}
172173

173-
prediction = await this.wait(
174-
prediction,
175-
wait || {},
176-
async (updatedPrediction) => {
177-
// Call progress callback with the updated prediction object
178-
if (progress) {
179-
progress(updatedPrediction);
174+
const isDone = isBlocking && prediction.status !== "starting";
175+
if (!isDone) {
176+
prediction = await this.wait(
177+
prediction,
178+
isBlocking ? {} : wait,
179+
async (updatedPrediction) => {
180+
// Call progress callback with the updated prediction object
181+
if (progress) {
182+
progress(updatedPrediction);
183+
}
184+
185+
// We handle the cancel later in the function.
186+
if (signal && signal.aborted) {
187+
return true; // stop polling
188+
}
189+
190+
return false; // continue polling
180191
}
181-
182-
// We handle the cancel later in the function.
183-
if (signal && signal.aborted) {
184-
return true; // stop polling
185-
}
186-
187-
return false; // continue polling
188-
}
189-
);
192+
);
193+
}
190194

191195
if (signal && signal.aborted) {
192196
prediction = await this.predictions.cancel(prediction.id);

0 commit comments

Comments
 (0)