Allow replicate.predictions.create to accept version or model options

Author: mattt

index.d.ts

@@ -75,6 +75,13 @@ declare module "replicate" {
     results: T[];
   }
 
+  export interface ServerSentEvent {
+    event: string;
+    data: string;
+    id?: string;
+    retry?: number;
+  }
+
   export default class Replicate {
     constructor(options?: {
       auth?: string;
@@ -103,6 +110,16 @@ declare module "replicate" {
       progress?: (prediction: Prediction) => void
     ): Promise<object>;
 
+    stream(
+      identifier: `${string}/${string}` | `${string}/${string}:${string}`,
+      options: {
+        input: object;
+        webhook?: string;
+        webhook_events_filter?: WebhookEventType[];
+        signal?: AbortSignal;
+      }
+    ): AsyncGenerator<ServerSentEvent>;
+
     request(
       route: string | URL,
       options: {
@@ -186,13 +203,15 @@ declare module "replicate" {
     };
 
     predictions: {
-      create(options: {
-        version: string;
-        input: object;
-        stream?: boolean;
-        webhook?: string;
-        webhook_events_filter?: WebhookEventType[];
-      }): Promise<Prediction>;
+      create(
+        options: {
+          version: string;
+          input: object;
+          stream?: boolean;
+          webhook?: string;
+          webhook_events_filter?: WebhookEventType[];
+        } & ({ version: string } | { model: string })
+      ): Promise<Prediction>;
       get(prediction_id: string): Promise<Prediction>;
       cancel(prediction_id: string): Promise<Prediction>;
       list(): Promise<Page<Prediction>>;

lib/predictions.js

@@ -2,15 +2,16 @@
  * Create a new prediction
  *
  * @param {object} options
- * @param {string} options.version - Required. The model version
+ * @param {string} options.model - The model.
+ * @param {string} options.version - The model version.
  * @param {object} options.input - Required. An object with the model inputs
  * @param {string} [options.webhook] - An HTTPS URL for receiving a webhook when the prediction has new output
  * @param {string[]} [options.webhook_events_filter] - You can change which events trigger webhook requests by specifying webhook events (`start`|`output`|`logs`|`completed`)
  * @param {boolean} [options.stream] - Whether to stream the prediction output. Defaults to false
  * @returns {Promise<object>} Resolves with the created prediction
  */
 async function createPrediction(options) {
-  const { stream, ...data } = options;
+  const { model, version, stream, ...data } = options;
 
   if (data.webhook) {
     try {
@@ -21,10 +22,20 @@ async function createPrediction(options) {
     }
   }
 
-  const response = await this.request("/predictions", {
-    method: "POST",
-    data: { ...data, stream },
-  });
+  let response;
+  if (version) {
+    response = await this.request("/predictions", {
+      method: "POST",
+      data: { ...data, stream, version },
+    });
+  } else if (model) {
+    response = await this.request(`/models/${model}/predictions`, {
+      method: "POST",
+      data: { ...data, stream },
+    });
+  } else {
+    throw new Error("Either model or version must be specified");
+  }
 
   return response.json();
 }