forked from triggerdotdev/trigger.dev
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathpredictions.ts
More file actions
101 lines (94 loc) · 2.65 KB
/
Copy pathpredictions.ts
File metadata and controls
101 lines (94 loc) · 2.65 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
import { IntegrationTaskKey } from "@trigger.dev/sdk";
import ReplicateClient, { Page, Prediction } from "replicate";
import { ReplicateRunTask } from "./index";
import { CallbackTimeout, ReplicateReturnType } from "./types";
import { callbackProperties, createPredictionProperties } from "./utils";
export class Predictions {
constructor(private runTask: ReplicateRunTask) {}
/** Cancel a prediction. */
cancel(key: IntegrationTaskKey, params: { id: string }): ReplicateReturnType<Prediction> {
return this.runTask(
key,
(client) => {
return client.predictions.cancel(params.id);
},
{
name: "Cancel Prediction",
params,
properties: [{ label: "Prediction ID", text: params.id }],
}
);
}
/** Create a new prediction. */
create(
key: IntegrationTaskKey,
params: Parameters<ReplicateClient["predictions"]["create"]>[0]
): ReplicateReturnType<Prediction> {
return this.runTask(
key,
(client) => {
return client.predictions.create(params);
},
{
name: "Create Prediction",
params,
properties: createPredictionProperties(params),
}
);
}
/** Create a new prediction and await the result. */
createAndAwait(
key: IntegrationTaskKey,
params: Omit<
Parameters<ReplicateClient["predictions"]["create"]>[0],
"webhook" | "webhook_events_filter"
>,
options: CallbackTimeout = { timeoutInSeconds: 3600 }
): ReplicateReturnType<Prediction> {
return this.runTask(
key,
(client, task) => {
return client.predictions.create({
...params,
webhook: task.callbackUrl ?? "",
webhook_events_filter: ["completed"],
});
},
{
name: "Create And Await Prediction",
params,
properties: [...createPredictionProperties(params), ...callbackProperties(options)],
callback: {
enabled: true,
timeoutInSeconds: options.timeoutInSeconds,
},
}
);
}
/** Fetch a prediction. */
get(key: IntegrationTaskKey, params: { id: string }): ReplicateReturnType<Prediction> {
return this.runTask(
key,
(client) => {
return client.predictions.get(params.id);
},
{
name: "Get Prediction",
params,
properties: [{ label: "Prediction ID", text: params.id }],
}
);
}
/** List all predictions. */
list(key: IntegrationTaskKey): ReplicateReturnType<Page<Prediction>> {
return this.runTask(
key,
(client) => {
return client.predictions.list();
},
{
name: "List Predictions",
}
);
}
}