Update request method to return response instead of JSON

Author: mattt

index.d.ts

@@ -85,7 +85,14 @@ declare module 'replicate' {
         webhook_events_filter?: WebhookEventType[];
       }
     ): Promise<object>;
-    request(route: string, parameters: any): Promise<any>;
+
+    request(route: string | URL, options: {
+      method?: string;
+      headers?: object | Headers;
+      params?: object;
+      data?: object;
+    }): Promise<Response>;
+
     paginate<T>(endpoint: () => Promise<Page<T>>): AsyncGenerator<[ T ]>;
     wait(
       prediction: Prediction,

index.js

@@ -125,43 +125,53 @@ class Replicate {
    * Make a request to the Replicate API.
    *
    * @param {string} route - REST API endpoint path
-   * @param {object} parameters - Request parameters
-   * @param {string} [parameters.method] - HTTP method. Defaults to GET
-   * @param {object} [parameters.params] - Query parameters
-   * @param {object} [parameters.data] - Body parameters
-   * @returns {Promise<object>} - Resolves with the API response data
+   * @param {object} options - Request parameters
+   * @param {string} [options.method] - HTTP method. Defaults to GET
+   * @param {object} [options.params] - Query parameters
+   * @param {object|Headers} [options.headers] - HTTP headers
+   * @param {object} [options.data] - Body parameters
+   * @returns {Promise<Response>} - Resolves with the response object
    * @throws {ApiError} If the request failed
    */
-  async request(route, parameters) {
+  async request(route, options) {
     const { auth, baseUrl, userAgent } = this;
 
-    const url = new URL(
-      route.startsWith('/') ? route.slice(1) : route,
-      baseUrl.endsWith('/') ? baseUrl : `${baseUrl}/`
-    );
+    let url;
+    if (route instanceof URL) {
+      url = route;
+    } else {
+      url = new URL(
+        route.startsWith('/') ? route.slice(1) : route,
+        baseUrl.endsWith('/') ? baseUrl : `${baseUrl}/`
+      );
+    }
 
-    const { method = 'GET', params = {}, data } = parameters;
+    const { method = 'GET', params = {}, data } = options;
 
     Object.entries(params).forEach(([key, value]) => {
       url.searchParams.append(key, value);
     });
 
-    const headers = {
-      Authorization: `Token ${auth}`,
-      'Content-Type': 'application/json',
-      'User-Agent': userAgent,
-    };
+    const headers = new Headers();
+    headers.append('Authorization', `Token ${auth}`);
+    headers.append('Content-Type', 'application/json');
+    headers.append('User-Agent', userAgent);
+    if (options.headers) {
+      options.headers.forEach((value, key) => {
+        headers.append(key, value);
+      });
+    }
 
-    const options = {
+    const init = {
       method,
       headers,
       body: data ? JSON.stringify(data) : undefined,
     };
 
-    const response = await this.fetch(url, options);
+    const response = await this.fetch(url, init);
 
     if (!response.ok) {
-      const request = new Request(url, options);
+      const request = new Request(url, init);
       const responseText = await response.text();
       throw new ApiError(
         `Request to ${url} failed with status ${response.status} ${response.statusText}: ${responseText}.`,
@@ -170,7 +180,7 @@ class Replicate {
       );
     }
 
-    return response.json();
+    return response;
   }
 
   /**
@@ -188,7 +198,7 @@ class Replicate {
     const response = await endpoint();
     yield response.results;
     if (response.next) {
-      const nextPage = () => this.request(response.next, { method: 'GET' });
+      const nextPage = () => this.request(response.next, { method: 'GET' }).then((r) => r.json());
       yield* this.paginate(nextPage);
     }
   }

lib/collections.js

@@ -5,9 +5,11 @@
  * @returns {Promise<object>} - Resolves with the collection data
  */
 async function getCollection(collection_slug) {
-  return this.request(`/collections/${collection_slug}`, {
+  const response = await this.request(`/collections/${collection_slug}`, {
     method: 'GET',
   });
+
+  return response.json();
 }
 
 /**
@@ -16,9 +18,11 @@ async function getCollection(collection_slug) {
  * @returns {Promise<object>} - Resolves with the collections data
  */
 async function listCollections() {
-  return this.request('/collections', {
+  const response = await this.request('/collections', {
     method: 'GET',
   });
+
+  return response.json();
 }
 
 module.exports = { get: getCollection, list: listCollections };

lib/models.js

@@ -6,9 +6,11 @@
  * @returns {Promise<object>} Resolves with the model data
  */
 async function getModel(model_owner, model_name) {
-  return this.request(`/models/${model_owner}/${model_name}`, {
+  const response = await this.request(`/models/${model_owner}/${model_name}`, {
     method: 'GET',
   });
+
+  return response.json();
 }
 
 /**
@@ -19,9 +21,11 @@ async function getModel(model_owner, model_name) {
  * @returns {Promise<object>} Resolves with the list of model versions
  */
 async function listModelVersions(model_owner, model_name) {
-  return this.request(`/models/${model_owner}/${model_name}/versions`, {
+  const response = await this.request(`/models/${model_owner}/${model_name}/versions`, {
     method: 'GET',
   });
+
+  return response.json();
 }
 
 /**
@@ -33,12 +37,14 @@ async function listModelVersions(model_owner, model_name) {
  * @returns {Promise<object>} Resolves with the model version data
  */
 async function getModelVersion(model_owner, model_name, version_id) {
-  return this.request(
+  const response = await this.request(
     `/models/${model_owner}/${model_name}/versions/${version_id}`,
     {
       method: 'GET',
     }
   );
+
+  return response.json();
 }
 
 module.exports = {

lib/predictions.js

@@ -23,7 +23,7 @@ async function createPrediction(options) {
     }
   }
 
-  const prediction = this.request('/predictions', {
+  const response = await this.request('/predictions', {
     method: 'POST',
     data,
   });
@@ -35,10 +35,10 @@ async function createPrediction(options) {
     const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
     await sleep(interval || 250);
 
-    return this.wait(await prediction, { maxAttempts, interval });
+    return this.wait(await response, { maxAttempts, interval });
   }
 
-  return prediction;
+  return response.json();
 }
 
 /**
@@ -48,9 +48,11 @@ async function createPrediction(options) {
  * @returns {Promise<object>} Resolves with the prediction data
  */
 async function getPrediction(prediction_id) {
-  return this.request(`/predictions/${prediction_id}`, {
+  const response = await this.request(`/predictions/${prediction_id}`, {
     method: 'GET',
   });
+
+  return response.json();
 }
 
 /**
@@ -60,9 +62,11 @@ async function getPrediction(prediction_id) {
  * @returns {Promise<object>} Resolves with the data for the training
  */
 async function cancelPrediction(prediction_id) {
-  return this.request(`/predictions/${prediction_id}/cancel`, {
+  const response = await this.request(`/predictions/${prediction_id}/cancel`, {
     method: 'POST',
   });
+
+  return response.json();
 }
 
 /**
@@ -71,9 +75,11 @@ async function cancelPrediction(prediction_id) {
  * @returns {Promise<object>} - Resolves with a page of predictions
  */
 async function listPredictions() {
-  return this.request('/predictions', {
+  const response = await this.request('/predictions', {
     method: 'GET',
   });
+
+  return response.json();
 }
 
 module.exports = {

lib/trainings.js

@@ -23,12 +23,12 @@ async function createTraining(model_owner, model_name, version_id, options) {
     }
   }
 
-  const training = this.request(`/models/${model_owner}/${model_name}/versions/${version_id}/trainings`, {
+  const response = await this.request(`/models/${model_owner}/${model_name}/versions/${version_id}/trainings`, {
     method: 'POST',
     data,
   });
 
-  return training;
+  return response.json();
 }
 
 /**
@@ -38,9 +38,11 @@ async function createTraining(model_owner, model_name, version_id, options) {
  * @returns {Promise<object>} Resolves with the data for the training
  */
 async function getTraining(training_id) {
-  return this.request(`/trainings/${training_id}`, {
+  const response = await this.request(`/trainings/${training_id}`, {
     method: 'GET',
   });
+
+  return response.json();
 }
 
 /**
@@ -50,9 +52,11 @@ async function getTraining(training_id) {
  * @returns {Promise<object>} Resolves with the data for the training
  */
 async function cancelTraining(training_id) {
-  return this.request(`/trainings/${training_id}/cancel`, {
+  const response = await this.request(`/trainings/${training_id}/cancel`, {
     method: 'POST',
   });
+
+  return response.json();
 }
 
 /**
@@ -61,9 +65,11 @@ async function cancelTraining(training_id) {
  * @returns {Promise<object>} - Resolves with a page of trainings
  */
 async function listTrainings() {
-  return this.request('/trainings', {
+  const response = await this.request('/trainings', {
     method: 'GET',
   });
+
+  return response.json();
 }
 
 module.exports = {