Add support for files API endpoints

client.go

@@ -145,21 +145,11 @@ func WithRetryPolicy(maxRetries int, backoff Backoff) ClientOption {
 	}
 }
 
-// request makes an HTTP request to the Replicate API.
-func (r *Client) request(ctx context.Context, method, path string, body interface{}, out interface{}) error {
-	bodyBuffer := &bytes.Buffer{}
-	if body != nil {
-		bodyBytes, err := json.Marshal(body)
-		if err != nil {
-			return fmt.Errorf("failed to marshal request body: %w", err)
-		}
-		bodyBuffer = bytes.NewBuffer(bodyBytes)
-	}
-
+func (r *Client) newRequest(ctx context.Context, method, path string, body io.Reader) (*http.Request, error) {
 	url := constructURL(r.options.baseURL, path)
-	request, err := http.NewRequestWithContext(ctx, method, url, bodyBuffer)
+	request, err := http.NewRequestWithContext(ctx, method, url, body)
 	if err != nil {
-		return fmt.Errorf("failed to create request: %w", err)
+		return nil, fmt.Errorf("failed to create request: %w", err)
 	}
 
 	request.Header.Set("Content-Type", "application/json")
@@ -168,6 +158,10 @@ func (r *Client) request(ctx context.Context, method, path string, body interfac
 		request.Header.Set("User-Agent", *r.options.userAgent)
 	}
 
+	return request, nil
+}
+
+func (r *Client) do(ctx context.Context, request *http.Request, out interface{}) error {
 	maxRetries := r.options.retryPolicy.maxRetries
 	backoff := r.options.retryPolicy.backoff
 
@@ -187,7 +181,7 @@ func (r *Client) request(ctx context.Context, method, path string, body interfac
 
 		if response.StatusCode < 200 || response.StatusCode >= 400 {
 			apiError = unmarshalAPIError(response, responseBytes)
-			if !r.shouldRetry(response, method) {
+			if !r.shouldRetry(response, request.Method) {
 				return apiError
 			}
 
@@ -229,6 +223,25 @@ func (r *Client) request(ctx context.Context, method, path string, body interfac
 	return fmt.Errorf("request failed")
 }
 
+// fetch makes an HTTP request to Replicate's API.
+func (r *Client) fetch(ctx context.Context, method, path string, body interface{}, out interface{}) error {
+	bodyBuffer := &bytes.Buffer{}
+	if body != nil {
+		bodyBytes, err := json.Marshal(body)
+		if err != nil {
+			return fmt.Errorf("failed to marshal request body: %w", err)
+		}
+		bodyBuffer = bytes.NewBuffer(bodyBytes)
+	}
+
+	request, err := r.newRequest(ctx, method, path, bodyBuffer)
+	if err != nil {
+		return err
+	}
+
+	return r.do(ctx, request, out)
+}
+
 // shouldRetry returns true if the request should be retried.
 //
 // - GET requests should be retried if the response status code is 429 or 5xx.

client_test.go

@@ -1,12 +1,20 @@
 package replicate_test
 
 import (
+	"bytes"
 	"context"
+	"crypto/md5" // nolint:gosec
+	"crypto/sha256"
+	"encoding/hex"
 	"encoding/json"
 	"fmt"
 	"io"
+	"mime"
+	"mime/multipart"
 	"net/http"
 	"net/http/httptest"
+	"os"
+	"path/filepath"
 	"testing"
 	"time"
 
@@ -1139,7 +1147,6 @@ func TestAutomaticallyRetryGetRequests(t *testing.T) {
 
 	i := 0
 	mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-
 		status := statuses[i]
 		i++
 
@@ -1318,3 +1325,241 @@ func TestStream(t *testing.T) {
 		}
 	}
 }
+
+func TestCreateFile(t *testing.T) {
+	fileID := "file-id"
+	options := &replicate.CreateFileOptions{
+		Filename:    "hello.txt",
+		ContentType: "text/plain",
+		Metadata:    map[string]string{"foo": "bar"},
+	}
+
+	mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/files", r.URL.Path)
+		assert.Equal(t, http.MethodPost, r.Method)
+
+		_, params, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		mr := multipart.NewReader(r.Body, params["boundary"])
+		defer r.Body.Close()
+
+		part, err := mr.NextPart()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		assert.Equal(t, "form-data; name=\"content\"; filename=\"hello.txt\"", part.Header.Get("Content-Disposition"))
+		assert.Equal(t, "text/plain", part.Header.Get("Content-Type"))
+
+		content, err := io.ReadAll(part)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		etag := fmt.Sprintf("%x", md5.Sum(content)) // nolint:gosec
+		checksum := sha256.Sum256(content)
+		file := &replicate.File{
+			ID:          fileID,
+			Name:        "hello.txt",
+			ContentType: "text/plain",
+			Size:        len(content),
+			Etag:        etag,
+			Checksums:   map[string]string{"sha256": hex.EncodeToString(checksum[:])},
+			Metadata:    map[string]string{"foo": "bar"},
+			CreatedAt:   "2022-04-26T22:13:06.224088Z",
+			URLs:        map[string]string{"get": "https://api.replicate.com/v1/files/" + fileID},
+		}
+
+		responseBytes, err := json.Marshal(file)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		w.WriteHeader(http.StatusCreated)
+		w.Write(responseBytes)
+	}))
+	defer mockServer.Close()
+
+	client, err := replicate.NewClient(
+		replicate.WithToken("test-token"),
+		replicate.WithBaseURL(mockServer.URL),
+	)
+	require.NotNil(t, client)
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	t.Run("CreateFileFromBytes", func(t *testing.T) {
+		content := []byte("Hello, world!")
+		file, err := client.CreateFileFromBytes(ctx, content, options)
+		if err != nil {
+			t.Fatal(err)
+		}
+		assertCreatedFile(t, fileID, file)
+	})
+
+	t.Run("CreateFileFromBuffer", func(t *testing.T) {
+		buf := bytes.NewBufferString("Hello, world!")
+		file, err := client.CreateFileFromBuffer(ctx, buf, options)
+		if err != nil {
+			t.Fatal(err)
+		}
+		assertCreatedFile(t, fileID, file)
+	})
+
+	t.Run("CreateFileFromPath", func(t *testing.T) {
+		content := []byte("Hello, world!")
+		tmpFilePath := filepath.Join(t.TempDir(), "hello.txt")
+		if err := os.WriteFile(tmpFilePath, content, 0o644); err != nil {
+			t.Fatal(err)
+		}
+		file, err := client.CreateFileFromPath(ctx, tmpFilePath, options)
+		if err != nil {
+			t.Fatal(err)
+		}
+		assertCreatedFile(t, fileID, file)
+	})
+}
+
+func assertCreatedFile(t *testing.T, fileID string, file *replicate.File) {
+	assert.Equal(t, fileID, file.ID)
+	assert.Equal(t, "hello.txt", file.Name)
+	assert.Equal(t, "text/plain", file.ContentType)
+	assert.Equal(t, 13, file.Size)
+	assert.Equal(t, "6cd3556deb0da54bca060b4c39479839", file.Etag)
+	assert.Equal(t, "315f5bdb76d078c43b8ac0064e4a0164612b1fce77c869345bfc94c75894edd3", file.Checksums["sha256"])
+	assert.Equal(t, map[string]string{"foo": "bar"}, file.Metadata)
+	assert.Equal(t, "2022-04-26T22:13:06.224088Z", file.CreatedAt)
+	assert.Equal(t, "https://api.replicate.com/v1/files/"+fileID, file.URLs["get"])
+}
+
+func TestListFiles(t *testing.T) {
+	fileID := "file-id"
+	mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/files", r.URL.Path)
+		assert.Equal(t, http.MethodGet, r.Method)
+
+		response := replicate.Page[replicate.File]{
+			Results: []replicate.File{
+				{
+					ID:          fileID,
+					Name:        "hello.txt",
+					ContentType: "text/plain",
+					Size:        13,
+					CreatedAt:   "2022-04-26T22:13:06.224088Z",
+					URLs:        map[string]string{"get": "https://api.replicate.com/v1/files/" + fileID},
+				},
+			},
+		}
+
+		responseBytes, err := json.Marshal(response)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		w.WriteHeader(http.StatusOK)
+		w.Write(responseBytes)
+	}))
+	defer mockServer.Close()
+
+	client, err := replicate.NewClient(
+		replicate.WithToken("test-token"),
+		replicate.WithBaseURL(mockServer.URL),
+	)
+	require.NotNil(t, client)
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	files, err := client.ListFiles(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assert.Equal(t, 1, len(files.Results))
+	assert.Nil(t, files.Previous)
+	assert.Nil(t, files.Next)
+
+	file := files.Results[0]
+	assert.Equal(t, fileID, file.ID)
+	assert.Equal(t, "hello.txt", file.Name)
+	assert.Equal(t, "text/plain", file.ContentType)
+	assert.Equal(t, 13, file.Size)
+	assert.Equal(t, "2022-04-26T22:13:06.224088Z", file.CreatedAt)
+	assert.Equal(t, "https://api.replicate.com/v1/files/"+fileID, file.URLs["get"])
+}
+func TestGetFile(t *testing.T) {
+	fileID := "file-id"
+	mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/files/"+fileID, r.URL.Path)
+		assert.Equal(t, http.MethodGet, r.Method)
+
+		file := &replicate.File{
+			ID:          fileID,
+			Name:        "hello.txt",
+			ContentType: "text/plain",
+			Size:        13,
+			CreatedAt:   "2022-04-26T22:13:06.224088Z",
+			URLs:        map[string]string{"get": "https://api.replicate.com/v1/files/" + fileID},
+		}
+
+		responseBytes, err := json.Marshal(file)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		w.WriteHeader(http.StatusOK)
+		w.Write(responseBytes)
+	}))
+	defer mockServer.Close()
+
+	client, err := replicate.NewClient(
+		replicate.WithToken("test-token"),
+		replicate.WithBaseURL(mockServer.URL),
+	)
+	require.NotNil(t, client)
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	file, err := client.GetFile(ctx, fileID)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assert.Equal(t, fileID, file.ID)
+	assert.Equal(t, "hello.txt", file.Name)
+	assert.Equal(t, "text/plain", file.ContentType)
+	assert.Equal(t, 13, file.Size)
+	assert.Equal(t, "2022-04-26T22:13:06.224088Z", file.CreatedAt)
+	assert.Equal(t, "https://api.replicate.com/v1/files/"+fileID, file.URLs["get"])
+}
+
+func TestDeleteFile(t *testing.T) {
+	fileID := "file-id"
+	mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/files/"+fileID, r.URL.Path)
+		assert.Equal(t, http.MethodDelete, r.Method)
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer mockServer.Close()
+
+	client, err := replicate.NewClient(
+		replicate.WithToken("test-token"),
+		replicate.WithBaseURL(mockServer.URL),
+	)
+	require.NotNil(t, client)
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	err = client.DeleteFile(ctx, fileID)
+	assert.NoError(t, err)
+}

collection.go

@@ -34,7 +34,7 @@ func (c *Collection) UnmarshalJSON(data []byte) error {
 // ListCollections returns a list of all collections.
 func (r *Client) ListCollections(ctx context.Context) (*Page[Collection], error) {
 	response := &Page[Collection]{}
-	err := r.request(ctx, "GET", "/collections", nil, response)
+	err := r.fetch(ctx, "GET", "/collections", nil, response)
 	if err != nil {
 		return nil, fmt.Errorf("failed to list collections: %w", err)
 	}
@@ -44,7 +44,7 @@ func (r *Client) ListCollections(ctx context.Context) (*Page[Collection], error)
 // GetCollection returns a collection by slug.
 func (r *Client) GetCollection(ctx context.Context, slug string) (*Collection, error) {
 	collection := &Collection{}
-	err := r.request(ctx, "GET", fmt.Sprintf("/collections/%s", slug), nil, collection)
+	err := r.fetch(ctx, "GET", fmt.Sprintf("/collections/%s", slug), nil, collection)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get collection: %w", err)
 	}

deployment.go

@@ -22,7 +22,7 @@ func (r *Client) CreatePredictionWithDeployment(ctx context.Context, deployment_
 
 	prediction := &Prediction{}
 	path := fmt.Sprintf("/deployments/%s/%s/predictions", deployment_owner, deployment_name)
-	err := r.request(ctx, "POST", path, data, prediction)
+	err := r.fetch(ctx, "POST", path, data, prediction)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create prediction: %w", err)
 	}