Add support for /predictions/{id}/cancel endpoint (#46)

* Add cancel prediction endpoint.

* Add documentation.

* Add test for CancelPrediction

* Reorder training methods

---------

Co-authored-by: Mattt Zmuda <[email protected]>

Author: mehmettokgoz

client_test.go

@@ -610,6 +610,41 @@ func TestCreatePredictionWithModel(t *testing.T) {
 	assert.Equal(t, replicate.Starting, prediction.Status)
 }
 
+func TestCancelPrediction(t *testing.T) {
+	mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "POST", r.Method)
+		assert.Equal(t, "/predictions/ufawqhfynnddngldkgtslldrkq/cancel", r.URL.Path)
+
+		response := replicate.Prediction{
+			ID:     "ufawqhfynnddngldkgtslldrkq",
+			Status: replicate.Canceled,
+		}
+		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()
+
+	prediction, err := client.CancelPrediction(ctx, "ufawqhfynnddngldkgtslldrkq")
+	assert.NoError(t, err)
+	assert.Equal(t, "ufawqhfynnddngldkgtslldrkq", prediction.ID)
+	assert.Equal(t, replicate.Canceled, prediction.Status)
+}
+
 func TestPredictionProgress(t *testing.T) {
 	prediction := replicate.Prediction{
 		ID:        "ufawqhfynnddngldkgtslldrkq",

prediction.go

@@ -141,3 +141,13 @@ func (r *Client) GetPrediction(ctx context.Context, id string) (*Prediction, err
 	}
 	return prediction, nil
 }
+
+// CancelPrediction cancels a running prediction by its ID.
+func (r *Client) CancelPrediction(ctx context.Context, id string) (*Prediction, error) {
+	prediction := &Prediction{}
+	err := r.fetch(ctx, "POST", fmt.Sprintf("/predictions/%s/cancel", id), nil, prediction)
+	if err != nil {
+		return nil, fmt.Errorf("failed to cancel prediction: %w", err)
+	}
+	return prediction, nil
+}

training.go

@@ -31,6 +31,16 @@ func (r *Client) CreateTraining(ctx context.Context, model_owner string, model_n
 	return training, nil
 }
 
+// ListTrainings returns a list of trainings.
+func (r *Client) ListTrainings(ctx context.Context) (*Page[Training], error) {
+	response := &Page[Training]{}
+	err := r.fetch(ctx, "GET", "/trainings", nil, response)
+	if err != nil {
+		return nil, fmt.Errorf("failed to list trainings: %w", err)
+	}
+	return response, nil
+}
+
 // GetTraining sends a request to the Replicate API to get a training.
 func (r *Client) GetTraining(ctx context.Context, trainingID string) (*Training, error) {
 	training := &Training{}
@@ -52,13 +62,3 @@ func (r *Client) CancelTraining(ctx context.Context, trainingID string) (*Traini
 
 	return training, nil
 }
-
-// ListTrainings returns a list of trainings.
-func (r *Client) ListTrainings(ctx context.Context) (*Page[Training], error) {
-	response := &Page[Training]{}
-	err := r.fetch(ctx, "GET", "/trainings", nil, response)
-	if err != nil {
-		return nil, fmt.Errorf("failed to list trainings: %w", err)
-	}
-	return response, nil
-}