New streaming api

go.mod

@@ -5,7 +5,9 @@ go 1.21
 toolchain go1.22.0
 
 require (
+	github.com/launchdarkly/eventsource v1.7.1
 	github.com/stretchr/testify v1.8.4
+	github.com/vincent-petithory/dataurl v1.0.0
 	golang.org/x/sync v0.6.0
 )
 

go.sum

@@ -1,20 +1,33 @@
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
+github.com/launchdarkly/eventsource v1.7.1 h1:StoRQeiPyrcQIXjlQ7b5jWMzHW4p+GGczN2r2oBhujg=
+github.com/launchdarkly/eventsource v1.7.1/go.mod h1:LHxSeb4OnqznNZxCSXbFghxS/CjIQfzHovNoAqbO/Wk=
+github.com/launchdarkly/go-test-helpers/v2 v2.2.0 h1:L3kGILP/6ewikhzhdNkHy1b5y4zs50LueWenVF0sBbs=
+github.com/launchdarkly/go-test-helpers/v2 v2.2.0/go.mod h1:L7+th5govYp5oKU9iN7To5PgznBuIjBPn+ejqKR0avw=
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e h1:fD57ERR4JtEqsWbfPhv4DMiApHyliiK5xCTNVSPiaAs=
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
+github.com/stretchr/testify v1.6.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
 github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+github.com/vincent-petithory/dataurl v1.0.0 h1:cXw+kPto8NLuJtlMsI152irrVw9fRDX8AbShPRpg2CI=
+github.com/vincent-petithory/dataurl v1.0.0/go.mod h1:FHafX5vmDzyP+1CQATJn7WFKc9CvnvxyvZy6I1MrG/U=
 golang.org/x/sync v0.6.0 h1:5BMeUDZ7vkXGfEr1x9B4bRcTH4lpkTkpdh0T/J+qjbQ=
 golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f h1:BLraFXnmrev5lT+xlilqcH8XK9/i0At2xKjWk4p6zsU=
 gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

stream.go

@@ -8,8 +8,12 @@
 	"fmt"
 	"io"
 	"net/http"
+	"strings"
 	"unicode/utf8"
 
+	"github.com/launchdarkly/eventsource"
+	"github.com/replicate/replicate-go/streaming"
+	"github.com/vincent-petithory/dataurl"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -140,6 +144,198 @@
 	return sseChan, errChan
 }
 
+func (r *Client) StreamPredictionFiles(ctx context.Context, prediction *Prediction) (<-chan streaming.File, error) {
+	url := prediction.URLs["stream"]
+	if url == "" {
+		return nil, errors.New("streaming not supported or not enabled for this prediction")
+	}
+
+	ch := make(chan streaming.File)
+
+	go r.streamFilesTo(ctx, ch, url, "")
+	return ch, nil
+}
+
+func (r *Client) StreamPredictionText(ctx context.Context, prediction *Prediction) (io.Reader, error) {
+	url := prediction.URLs["stream"]
+	if url == "" {
+		return nil, errors.New("streaming not supported or not enabled for this prediction")
+	}
+
+	reader, writer := io.Pipe()
+
+	go r.streamTextTo(ctx, writer, url, "")
+	return reader, nil
+}
+
+func (r *Client) streamTextTo(ctx context.Context, writer *io.PipeWriter, url string, lastEventID string) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		writer.CloseWithError(fmt.Errorf("failed to create request: %w", err))
+		return
+	}
+	req.Header.Set("Accept", "text/event-stream")
+	req.Header.Set("Cache-Control", "no-cache")
+	req.Header.Set("Connection", "keep-alive")
+
+	if lastEventID != "" {
+		req.Header.Set("Last-Event-ID", lastEventID)
+	}
+
+	resp, err := r.c.Do(req)
+	if err != nil {
+		writer.CloseWithError(fmt.Errorf("failed to send request: %w", err))
+		return
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		writer.CloseWithError(fmt.Errorf("received invalid status code: %d", resp.StatusCode))
+		return
+	}
+	defer resp.Body.Close()
+	decoder := eventsource.NewDecoder(resp.Body)
+	for {
+		event, err := decoder.Decode()
+		if err != nil {
+			if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
+				// retry (TODO: backoff policy?)
+				r.streamTextTo(ctx, writer, url, lastEventID)
+				return
+			}
+			writer.CloseWithError(fmt.Errorf("Failed to get token: %w", err))
+			return
+		}
+		lastEventID = event.Id()
+		switch event.Event() {
+		case SSETypeOutput:
+			io.WriteString(writer, event.Data())
+		case SSETypeDone:
+			writer.Close()
+			return
+		case SSETypeLogs:
+			// TODO
+		default:
+			writer.CloseWithError(fmt.Errorf("unknown event type %s", event.Event()))
+			return
+		}
+	}
+}
+
+type dataURL struct {
+	url string
+}
+
+var _ streaming.File = &dataURL{}
+
+func (d *dataURL) Body(_ context.Context) (io.ReadCloser, error) {
+	data, err := dataurl.DecodeString(d.url)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return io.NopCloser(bytes.NewReader(data.Data)), nil
+}
+
+type httpURL struct {
+	c   *http.Client
+	url string
+}
+
+var _ streaming.File = &httpURL{}
+
+func (h *httpURL) Body(ctx context.Context) (io.ReadCloser, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, h.url, nil)
+	if err != nil {
+		return nil, err
+	}
+	resp, err := h.c.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	return resp.Body, nil
+}
+
+type errWrapper struct {
+	err error
+}
+
+var _ streaming.File = &errWrapper{}
+
+func fileError(err error) streaming.File {
+	return &errWrapper{err: err}
+}
+
+func (e *errWrapper) Body(_ context.Context) (io.ReadCloser, error) {
+	return nil, e.err
+}
+
+func (e *errWrapper) Close() error {
+	return nil
+}
+
+func (r *Client) streamFilesTo(ctx context.Context, out chan<- streaming.File, url string, lastEventID string) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		out <- fileError(err)
+		close(out)
+		return
+	}
+	req.Header.Set("Accept", "text/event-stream")
+	req.Header.Set("Cache-Control", "no-cache")
+	req.Header.Set("Connection", "keep-alive")
+
+	if lastEventID != "" {
+		req.Header.Set("Last-Event-ID", lastEventID)
+	}
+
+	resp, err := r.c.Do(req)
+	if err != nil {
+		out <- fileError(fmt.Errorf("failed to send request: %w", err))
+		close(out)
+		return
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		out <- fileError(fmt.Errorf("received invalid status code: %d", resp.StatusCode))
+		close(out)
+		return
+	}
+	defer resp.Body.Close()
+	decoder := eventsource.NewDecoder(resp.Body)
+	for {
+		event, err := decoder.Decode()
+		if err != nil {
+			if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
+				// retry (TODO: backoff policy?)
+				r.streamFilesTo(ctx, out, url, lastEventID)
+				return
+			}
+			out <- fileError(fmt.Errorf("Failed to get token: %w", err))
+			close(out)
+			return
+		}
+		lastEventID = event.Id()
+		switch event.Event() {
+		case SSETypeOutput:
+			if strings.HasPrefix(event.Data(), "data:") {
+				out <- &dataURL{url: event.Data()}
+			} else if strings.HasPrefix(event.Data(), "http") {
+				out <- &httpURL{c: r.c, url: event.Data()}
+			}
+		case SSETypeDone:
+			close(out)
+			return
+		case SSETypeLogs:
+			// TODO
+		default:
+			out <- fileError(fmt.Errorf("unknown event type %s", event.Event()))
+			return
+		}
+	}
+
+}
+
 func (r *Client) streamPrediction(ctx context.Context, prediction *Prediction, lastEvent *SSEEvent, sseChan chan SSEEvent, errChan chan error) {
 	url := prediction.URLs["stream"]
 	if url == "" {