replicate-api
Version:
A typed client library for the replicate.com API
43 lines (42 loc) • 1.6 kB
JavaScript
import { getPrediction } from "./getPrediction.js";
const getSleepDuration = (elapsedTimeMillis) => {
if (elapsedTimeMillis < 10000) {
return 1000;
}
if (elapsedTimeMillis < 60000) {
return 5000;
}
return 10000;
};
/** Poll a prediction by ID.
*
* ```typescript
* const result = await pollPrediction({
* id: "ID of your prediction",
* token: "Get your token at https://replicate.com/account"
* })
* ```
* If you have a `PredictionState`, you don't have to use this function, just call `.poll()` on that object.
*
* If the timeout occurs an error is thrown.
*
* @returns A new `PredictionState`. It has a status of either "succeeded", "failed" or "canceled".
*/
export const pollPrediction = async (options, initialResult) => {
let newPrediction = initialResult || (await getPrediction({ ...options, id: options.id }));
const endAt = Date.now() + (options.timeout ?? 3600000);
while (Date.now() < endAt) {
if (newPrediction.status === "succeeded" ||
newPrediction.status === "failed" ||
newPrediction.status === "canceled") {
return newPrediction;
}
const elapsedTime = newPrediction.startedAt ? Date.now() - newPrediction.startedAt.getTime() : 0;
const sleepDuration = getSleepDuration(elapsedTime);
if (newPrediction !== initialResult) {
await new Promise(resolve => setTimeout(resolve, sleepDuration));
}
newPrediction = await getPrediction({ ...options, id: options.id });
}
throw new Error("Prediction timed out");
};