Add a new scaffold command to the CLI

cmd/replicate/main.go

@@ -37,6 +37,7 @@ func init() {
 		prediction.RootCmd,
 		training.RootCmd,
 		hardware.RootCmd,
+		cmd.ScaffoldCmd,
 	} {
 		rootCmd.AddCommand(cmd)
 		cmd.GroupID = "core"

go.mod

@@ -14,7 +14,7 @@ require (
 	github.com/getkin/kin-openapi v0.120.0
 	github.com/golangci/golangci-lint v1.55.2
 	github.com/mattn/go-isatty v0.0.20
-	github.com/replicate/replicate-go v0.11.0
+	github.com/replicate/replicate-go v0.12.0
 	github.com/schollz/progressbar/v3 v3.13.1
 	github.com/spf13/cobra v1.8.0
 	github.com/stretchr/testify v1.8.4

go.sum

@@ -522,8 +522,8 @@ github.com/quasilyte/regex/syntax v0.0.0-20210819130434-b3f0c404a727 h1:TCg2WBOl
 github.com/quasilyte/regex/syntax v0.0.0-20210819130434-b3f0c404a727/go.mod h1:rlzQ04UMyJXu/aOvhd8qT+hvDrFpiwqp8MRXDY9szc0=
 github.com/quasilyte/stdinfo v0.0.0-20220114132959-f7386bf02567 h1:M8mH9eK4OUR4lu7Gd+PU1fV2/qnDNfzT635KRSObncs=
 github.com/quasilyte/stdinfo v0.0.0-20220114132959-f7386bf02567/go.mod h1:DWNGW8A4Y+GyBgPuaQJuWiy0XYftx4Xm/y5Jqk9I6VQ=
-github.com/replicate/replicate-go v0.11.0 h1:zjjdXVvot2TqtiL8usgzdr6CEhrcI51DhygC2Gw/qjE=
-github.com/replicate/replicate-go v0.11.0/go.mod h1:k9C4+PaYa9+hMRjn4D7ZPHOCUFb8P4jhytsCqcGa2vU=
+github.com/replicate/replicate-go v0.12.0 h1:gd/hw4hCBO5G4M3Fezb3zdKYSbe9NEfRLzGoktFk3Ks=
+github.com/replicate/replicate-go v0.12.0/go.mod h1:k9C4+PaYa9+hMRjn4D7ZPHOCUFb8P4jhytsCqcGa2vU=
 github.com/rivo/uniseg v0.1.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
 github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
 github.com/rivo/uniseg v0.4.4 h1:8TfxU8dW6PdqD27gjM8MVNuicgxIjxpm4K7x4jp8sis=

go.work.sum

@@ -1,33 +1,23 @@
 cloud.google.com/go/firestore v1.6.1/go.mod h1:asNXNOzBdyVQmEU+ggO8UPodTkEVFW5Qx+rwHnAz+EY=
 github.com/armon/go-metrics v0.3.10/go.mod h1:4O98XIr/9W0sxpJ8UaYkvjk10Iff7SnFrb4QAOwNTFc=
 github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI=
-github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/charmbracelet/harmonica v0.2.0/go.mod h1:KSri/1RMQOZLbw7AHqgcBycp8pgJnQMYYT8QZRqZ1Ao=
 github.com/coreos/go-semver v0.3.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
 github.com/coreos/go-systemd/v22 v22.3.2/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
 github.com/cristalhq/acmd v0.11.1/go.mod h1:LG5oa43pE/BbxtfMoImHCQN++0Su7dzipdgBjMCBVDQ=
 github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
-github.com/frankban/quicktest v1.14.4/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0=
-github.com/go-logr/logr v1.2.4/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
 github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0=
-github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572/go.mod h1:9Pwr4B2jHnOSGXyyzV8ROjYa2ojvAY6HCGYYfMoC3Ls=
-github.com/go-test/deep v1.0.8/go.mod h1:5C2ZWiW0ErCdrYzpqxLbTX7MG14M9iiw8DgHncVwcsE=
-github.com/go-toolsmith/pkgload v1.2.2/go.mod h1:R2hxLNRKuAsiXCo2i5J6ZQPhnPMOVtU+f0arbFPWCus=
 github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
 github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
-github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
 github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/google/uuid v1.3.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/gookit/color v1.5.4/go.mod h1:pZJOeOS8DM43rXbp4AZo1n9zCU2qjpcRko0b6/QJi9w=
 github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So=
-github.com/gostaticanalysis/testutil v0.4.0/go.mod h1:bLIoPefWXrRi/ssLFWX1dx7Repi5x3CuviD3dgAZaBU=
 github.com/hashicorp/go-cleanhttp v0.5.2/go.mod h1:kO/YDlP8L1346E6Sodw+PrpBSV4/SoxCXGY6BqNFT48=
 github.com/hashicorp/go-hclog v1.2.0/go.mod h1:whpDNt7SSdeAju8AWKIWsul05p54N/39EeqMAyrmvFQ=
 github.com/hashicorp/go-immutable-radix v1.3.1/go.mod h1:0y9vanUI8NX6FsYoO3zeMjhV/C5i9g4Q3DwcSNZ4P60=
 github.com/hashicorp/go-rootcerts v1.0.2/go.mod h1:pqUvnprVnM5bf7AOirdbb01K4ccR319Vf4pU3K5EGc8=
 github.com/hashicorp/golang-lru v0.5.4/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4=
-github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
-github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
 github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
 github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I=
@@ -45,17 +35,10 @@ github.com/replicate/replicate-go v0.8.1 h1:Mza5hWR/R1akZRKwXtA/CQJ2pY4/B9fSCYX+
 github.com/replicate/replicate-go v0.8.1/go.mod h1:k9C4+PaYa9+hMRjn4D7ZPHOCUFb8P4jhytsCqcGa2vU=
 github.com/sahilm/fuzzy v0.1.0/go.mod h1:VFvziUEIMCrT6A6tw2RFIXPXXmzXbOsSHF0DOI8ZK9Y=
 github.com/shoenig/go-m1cpu v0.1.6/go.mod h1:1JJMcUBvfNwpq05QDQVAnx3gUHr9IYF7GNg9SUEw2VQ=
-github.com/ugorji/go/codec v1.2.7/go.mod h1:WGN1fab3R1fzQlVQTkfxVtIBhWDRqOviHU95kRgeqEY=
 github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
 github.com/valyala/quicktemplate v1.7.0/go.mod h1:sqKJnoaOF88V07vkO+9FL8fb9uZg/VPSJnLYn+LmLk8=
 github.com/xo/terminfo v0.0.0-20210125001918-ca9a967f8778/go.mod h1:2MuV+tbUrU1zIOPMxZ5EncGwgmMJsa+9ucAQZXxsObs=
 github.com/yusufpapurcu/wmi v1.2.3/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0=
-go-simpler.org/assert v0.6.0/go.mod h1:74Eqh5eI6vCK6Y5l3PI8ZYFXG4Sa+tkr70OIPJAUr28=
 go.opencensus.io v0.23.0/go.mod h1:XItmlyltB5F7CS4xOC1DcqMoFqwtC6OG2xF7mCv7P7E=
-go.uber.org/goleak v1.1.11/go.mod h1:cwTWslyiVhfpKIDGSZEM2HlOvcqm+tG4zioyIeLoqMQ=
 golang.org/x/crypto v0.14.0/go.mod h1:MVFd36DqK4CsrnJYDkBA3VC4m2GkXAM0PvzMCn4JQf4=
 golang.org/x/lint v0.0.0-20210508222113-6edffad5e616/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY=
-golang.org/x/net v0.14.0/go.mod h1:PpSgVXXLK0OxS0F31C1/tv6XNguvCrnXIDrFMspZIUI=
-golang.org/x/net v0.17.0/go.mod h1:NxSsAGuq816PNPmqtQdLE42eU2Fs7NoRIZrHJAlaCOE=
-gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=

internal/cmd/scaffold.go

@@ -0,0 +1,231 @@
+package cmd
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/url"
+	"os"
+	"os/exec"
+	"strings"
+
+	"github.com/replicate/replicate-go"
+	"github.com/spf13/cobra"
+)
+
+var ScaffoldCmd = &cobra.Command{
+	Use:   "scaffold <prediction-ID-or-URL> [<directory>] [--template=<template>]",
+	Short: "Create a new local development environment from a prediction",
+	Args:  cobra.RangeArgs(1, 2),
+	RunE: func(cmd *cobra.Command, args []string) error {
+		ctx := cmd.Context()
+
+		apiToken := os.Getenv("REPLICATE_API_TOKEN")
+		if apiToken == "" {
+			return fmt.Errorf("REPLICATE_API_TOKEN environment variable not set. Please set this to your Replicate API token")
+		}
+
+		client, err := replicate.NewClient(replicate.WithToken(apiToken))
+		if err != nil {
+			return fmt.Errorf("failed to create client: %w", err)
+		}
+
+		predictionId, err := parsePredictionId(args[0])
+		if err != nil {
+			return fmt.Errorf("failed to parse prediction ID: %w", err)
+		}
+		prediction, err := client.GetPrediction(ctx, predictionId)
+		if prediction == nil || err != nil {
+			return fmt.Errorf("failed to get prediction: %w", err)
+		}
+
+		var directory string
+		if len(args) == 2 {
+			directory = args[1]
+		} else {
+			directory = predictionId
+		}
+
+		template, _ := cmd.Flags().GetString("template")
+
+		switch template {
+		case "node", "nodejs", "js", "":
+			return handleNodeTemplate(ctx, prediction, directory, apiToken)
+		case "python":
+			return handlePythonTemplate(ctx, prediction, directory, apiToken)
+		default:
+			return fmt.Errorf("unsupported template: %s, expected one of: node, python", template)
+		}
+	},
+}
+
+func init() {
+	ScaffoldCmd.Flags().StringP("template", "t", "", "Starter git repo template to use. Currently supported: node, python")
+}
+
+// Parse the prediction id from a url, or return the prediction id if it's not a url
+func parsePredictionId(value string) (string, error) {
+	// Case 1: A prediction ID
+	if !strings.Contains(value, "/") {
+		return value, nil
+	}
+
+	// Case 2: A URL in the form https://replicate.com/p/{id}
+	if strings.HasPrefix(value, "replicate.com/p/") || strings.HasPrefix(value, "https://replicate.com/p/") {
+		splitUrl := strings.Split(value, "/")
+		if len(splitUrl) == 0 {
+			return "", fmt.Errorf("invalid URL format")
+		}
+		return splitUrl[len(splitUrl)-1], nil
+	}
+
+	// Case 3: A URL in the form https://api.replicate.com/v1/predictions/{id}
+	if strings.HasPrefix(value, "api.replicate.com/v1/predictions/") || strings.HasPrefix(value, "https://api.replicate.com/v1/predictions/") {
+		splitUrl := strings.Split(value, "/")
+		if len(splitUrl) == 0 {
+			return "", fmt.Errorf("invalid URL format")
+		}
+		return splitUrl[len(splitUrl)-1], nil
+	}
+
+	// Case 4: A URL in the form "https://replicate.com/*?prediction={id}"
+	if strings.Contains(value, "replicate.com") || strings.Contains(value, "https://replicate.com") {
+		parsedUrl, err := url.Parse(value)
+		if err != nil {
+			return "", fmt.Errorf("failed to parse URL: %w", err)
+		}
+		queryParams, err := url.ParseQuery(parsedUrl.RawQuery)
+		if err != nil {
+			return "", fmt.Errorf("failed to parse query parameters: %w", err)
+		}
+		predictionId := queryParams.Get("prediction")
+		if predictionId == "" {
+			return "", fmt.Errorf("no prediction ID found in URL")
+		}
+		return predictionId, nil
+	}
+
+	// If none of the above cases match, return an error
+	return "", fmt.Errorf("invalid prediction ID or URL format")
+}
+
+func execCommand(ctx context.Context, command string) error {
+	cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	err := cmd.Run()
+	if err != nil {
+		return fmt.Errorf("failed to run command: %w", err)
+	}
+	return nil
+}
+
+func handleNodeTemplate(ctx context.Context, prediction *replicate.Prediction, directory string, apiToken string) error {
+	fmt.Println("Cloning starter repo and installing dependencies...")
+
+	// 1. Clone the starter repo
+	if err := execCommand(ctx, fmt.Sprintf("git clone https://github.com/replicate/node-starter.git %s", directory)); err != nil {
+		return fmt.Errorf("failed to clone starter repo: %w", err)
+	}
+
+	// 2. Set chdir to the output path
+	if err := os.Chdir(directory); err != nil {
+		return fmt.Errorf("failed to change directory: %w", err)
+	}
+
+	// 3. Install dependencies
+	if err := execCommand(ctx, "npm install"); err != nil {
+		return fmt.Errorf("failed to install dependencies: %w", err)
+	}
+
+	// 4. Set the REPLICATE_API_TOKEN env var
+	if err := execCommand(ctx, fmt.Sprintf(`echo 'REPLICATE_API_TOKEN="%s"' >> .env`, apiToken)); err != nil {
+		return fmt.Errorf("failed to write .env file: %w", err)
+	}
+
+	// Open the template file
+	templateData, err := os.ReadFile("index.js.template")
+	if err != nil {
+		return fmt.Errorf("failed to read template file: %w", err)
+	}
+
+	fullModelString := fmt.Sprintf("%s:%s", prediction.Model, prediction.Version)
+
+	// Perform string replacement on the template file.
+	replacedData := strings.ReplaceAll(string(templateData), "{{MODEL_STRING}}", fullModelString)
+	inputs, _ := json.Marshal(prediction.Input)
+	replacedData = strings.ReplaceAll(replacedData, "{{INPUTS}}", string(inputs))
+
+	// 5. Write the populated template to {directory}/index.js
+	fmt.Println("Writing new index.js...")
+	err = os.WriteFile("index.js", []byte(replacedData), 0o644)
+	if err != nil {
+		return err
+	}
+
+	// 6. Run the example prediction
+	fmt.Println("Running example prediction...")
+	if err := execCommand(ctx, "node index.js"); err != nil {
+		return fmt.Errorf("failed to run example prediction: %w", err)
+	}
+
+	return nil
+}
+
+func handlePythonTemplate(ctx context.Context, prediction *replicate.Prediction, directory string, apiToken string) error {
+	fmt.Println("Cloning starter repo and installing dependencies...")
+
+	// 1. Clone the starter repo
+	if err := execCommand(ctx, fmt.Sprintf("git clone git@github.com:replicate/python-starter.git %s", directory)); err != nil {
+		return fmt.Errorf("failed to clone starter repo: %w", err)
+	}
+
+	// 2. Set chdir to the output path
+	err := os.Chdir(directory)
+	if err != nil {
+		return fmt.Errorf("failed to change directory: %w", err)
+	}
+
+	// 3. Create virtualenv
+	if err := execCommand(ctx, "virtualenv .venv"); err != nil {
+		return fmt.Errorf("failed to create virtualenv: %w", err)
+	}
+
+	// 4. Install dependencies
+	if err := execCommand(ctx, ".venv/bin/pip install -r requirements.txt"); err != nil {
+		return fmt.Errorf("failed to install dependencies: %w", err)
+	}
+
+	// 5. Set the REPLICATE_API_TOKEN env var
+	if err := execCommand(ctx, fmt.Sprintf(`echo 'REPLICATE_API_TOKEN="%s"' >> .env`, apiToken)); err != nil {
+		return fmt.Errorf("failed to write .env file: %w", err)
+	}
+
+	fmt.Println("Cloning starter repo, and installing dependencies...")
+
+	// Open the template file
+	templateData, err := os.ReadFile("prediction.py.template")
+	if err != nil {
+		return fmt.Errorf("failed to read template file: %w", err)
+	}
+	fullModelString := fmt.Sprintf("%s:%s", prediction.Model, prediction.Version)
+
+	replacedData := strings.ReplaceAll(string(templateData), "{{MODEL_STRING}}", fullModelString)
+	inputs, _ := json.Marshal(prediction.Input)
+	replacedData = strings.ReplaceAll(replacedData, "{{INPUTS}}", string(inputs))
+
+	// 6. Write the populated template to {directory}/prediction.py
+	fmt.Println("Writing new prediction.py...")
+	err = os.WriteFile("prediction.py", []byte(replacedData), 0o644)
+	if err != nil {
+		return err
+	}
+
+	// 7. Run the example prediction
+	fmt.Println("Running example prediction...")
+	if err := execCommand(ctx, ".venv/bin/python prediction.py"); err != nil {
+		return fmt.Errorf("failed to run example prediction: %w", err)
+	}
+
+	return nil
+}