Add support for training endpoints

Author: mattt

index.d.ts

@@ -43,6 +43,18 @@ declare module "replicate" {
     updated: string;
   }
 
+  export interface Training {
+    id: string;
+    destination: string;
+    version: string;
+    input: any;
+    output: any;
+    webhook?: string;
+    webhook_events_filter?: WebhookEventType[];
+    created: string;
+    updated: string;
+  }
+
   export interface CollectionsGetOptions {
     collection_slug: string;
   }
@@ -74,6 +86,22 @@ declare module "replicate" {
     predictionId: string;
   }
 
+  export interface TrainingsCreateOptions {
+    destination: string;
+    version: string;
+    input: any;
+    webhook?: string;
+    webhook_events_filter?: WebhookEventType[];
+  }
+
+  export interface TrainingsGetOptions {
+    trainingId: string;
+  }
+
+  export interface TrainingsCancelOptions {
+    trainingId: string;
+  }
+
   export class Replicate {
     constructor(options: ReplicateOptions);
 
@@ -113,6 +141,12 @@ declare module "replicate" {
       get(options: PredictionsGetOptions): Promise<Prediction>;
       list(): Promise<Prediction[]>;
     };
+
+    trainings: {
+      create(options: TrainingsCreateOptions): Promise<Training>;
+      get(options: TrainingsGetOptions): Promise<Training>;
+      cancel(options: TrainingsGetOptions): Promise<Training>;
+    };
   }
 
   export default Replicate;

index.js

@@ -3,6 +3,7 @@ const axios = require('axios');
 const collections = require('./lib/collections');
 const models = require('./lib/models');
 const predictions = require('./lib/predictions');
+const trainings = require('./lib/trainings');
 const packageJSON = require('./package.json');
 
 /**
@@ -63,6 +64,12 @@ class Replicate {
       get: predictions.get.bind(this),
       list: predictions.list.bind(this),
     };
+
+    this.trainings = {
+      create: trainings.create.bind(this),
+      get: trainings.get.bind(this),
+      cancel: trainings.cancel.bind(this),
+    };
   }
 
   /**

index.test.js

@@ -131,6 +131,59 @@ describe('Replicate client', () => {
     // Add more tests for error handling, edge cases, etc.
   });
 
+  describe('trainings.create', () => {
+    test('Calls the correct API route with the correct payload', async () => {
+      client.request = jest.fn();
+      await client.trainings.create(
+        'owner',
+        'model',
+        '632231d0d49d34d5c4633bd838aee3d81d936e59a886fbf28524702003b4c532',
+        {
+          destination: 'new_owner/new_model',
+          input: {
+            text: '...'
+          }
+        }
+      );
+
+      expect(client.request).toHaveBeenCalledWith('/models/owner/model/versions/632231d0d49d34d5c4633bd838aee3d81d936e59a886fbf28524702003b4c532/trainings', {
+        method: 'POST',
+        data: {
+          destination: 'new_owner/new_model',
+          input: {
+            text: '...'
+          },
+        }
+      });
+    });
+
+    // Add more tests for error handling, edge cases, etc.
+  });
+
+  describe('trainings.get', () => {
+    test('Calls the correct API route with the correct payload', async () => {
+      client.request = jest.fn();
+      await client.trainings.get(123);
+      expect(client.request).toHaveBeenCalledWith('/trainings/123', {
+        method: 'GET',
+      });
+    });
+
+    // Add more tests for error handling, edge cases, etc.
+  });
+
+  describe('trainings.cancel', () => {
+    test('Calls the correct API route with the correct payload', async () => {
+      client.request = jest.fn();
+      await client.trainings.cancel(123);
+      expect(client.request).toHaveBeenCalledWith('/trainings/123/cancel', {
+        method: 'POST',
+      });
+    });
+
+    // Add more tests for error handling, edge cases, etc.
+  });
+
   describe('run', () => {
     test('Calls the correct API routes', async () => {
       client.request = jest.fn();

lib/trainings.js

@@ -0,0 +1,53 @@
+/**
+ * Create a new training
+ *
+ * @param {string} model_owner - Required. The username of the user or organization who owns the model
+ * @param {string} model_name - Required. The name of the model
+ * @param {string} version_id - Required. The version ID
+ * @param {object} options
+ * @param {string} options.destination - Required. The destination for the trained version in the form "{username}/{model_name}"
+ * @param {object} options.input - Required. An object with the model inputs
+ * @param {string} [options.webhook] - An HTTPS URL for receiving a webhook when the training updates
+ * @param {string[]} [options.webhook_events_filter] - You can change which events trigger webhook requests by specifying webhook events (`start`|`output`|`logs`|`completed`)
+ * @returns {Promise<object>} Resolves with the data for the created training
+ */
+async function createTraining(model_owner, model_name, version_id, options) {
+  const { ...data } = options;
+
+  const training = this.request(`/models/${model_owner}/${model_name}/versions/${version_id}/trainings`, {
+    method: 'POST',
+    data,
+  });
+
+  return training;
+}
+
+/**
+ * Fetch a training by ID
+ *
+ * @param {string} training_id - Required. The training ID
+ * @returns {Promise<object>} Resolves with the data for the training
+ */
+async function getTraining(training_id) {
+  return this.request(`/trainings/${training_id}`, {
+    method: 'GET',
+  });
+}
+
+/**
+ * Cancel a training by ID
+ *
+ * @param {string} training_id - Required. The training ID
+ * @returns {Promise<object>} Resolves with the data for the training
+ */
+async function cancelTraining(training_id) {
+  return this.request(`/trainings/${training_id}/cancel`, {
+    method: 'POST',
+  });
+}
+
+module.exports = {
+  create: createTraining,
+  get: getTraining,
+  cancel: cancelTraining,
+};