New streaming api (#91)

This adds functions StreamPredictionText() and StreamPredictionFiles() that stream output from models that return iterators of strings or files, respectively.

Author: philandstuff

go.mod

@@ -6,6 +6,7 @@ toolchain go1.22.0
 
 require (
 	github.com/stretchr/testify v1.8.4
+	github.com/vincent-petithory/dataurl v1.0.0
 	golang.org/x/sync v0.6.0
 )
 

go.sum

@@ -11,6 +11,8 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 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=

identifier_test.go

@@ -3,8 +3,9 @@ package replicate_test
 import (
 	"testing"
 
-	"github.com/replicate/replicate-go"
 	"github.com/stretchr/testify/assert"
+
+	"github.com/replicate/replicate-go"
 )
 
 func TestValidWithVersion(t *testing.T) {

internal/sse/decoder.go

@@ -0,0 +1,70 @@
+package sse
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+	"strings"
+)
+
+type Event struct {
+	Type string
+	ID   string
+	Data string
+}
+
+type Decoder struct {
+	r *bufio.Reader
+}
+
+func NewDecoder(r io.Reader) *Decoder {
+	return &Decoder{r: bufio.NewReader(r)}
+}
+
+var (
+	eventField = []byte("event:")
+	dataField  = []byte("data:")
+	idField    = []byte("id:")
+	retryField = []byte("retry:")
+	space      = []byte{' '}
+)
+
+func buildEvent(t, id string, data *strings.Builder) Event {
+	return Event{
+		Type: t,
+		ID:   id,
+		Data: data.String(),
+	}
+}
+
+func (d *Decoder) Next() (Event, error) {
+	var t, id string
+	var data strings.Builder
+	for {
+		line, err := d.r.ReadBytes('\n')
+		if err == io.EOF {
+			return buildEvent(t, id, &data), io.ErrUnexpectedEOF
+		}
+		if err != nil {
+			return buildEvent(t, id, &data), err
+		}
+
+		switch {
+		case line[0] == '\n':
+			// a blank line finishes the event, so we return it
+			return buildEvent(t, id, &data), nil
+		case bytes.HasPrefix(line, eventField):
+			t = string(bytes.TrimPrefix(line[6:len(line)-1], space))
+		case bytes.HasPrefix(line, dataField):
+			// strings.Builder.Write() always returns nil error, so we don't
+			// need to handle it
+			data.Write(bytes.TrimPrefix(line[5:], space))
+		case bytes.HasPrefix(line, idField):
+			id = string(bytes.TrimPrefix(line[3:len(line)-1], space))
+		case bytes.HasPrefix(line, retryField):
+		default:
+			// ignore the line
+		}
+
+	}
+}

internal/sse/decoder_test.go

@@ -0,0 +1,133 @@
+package sse_test
+
+import (
+	"fmt"
+	"io"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/replicate/replicate-go/internal/sse"
+)
+
+func TestDecodeOneEventNoSpace(t *testing.T) {
+	input := `event:output
+id:123abc
+data:giraffe
+
+`
+	d := sse.NewDecoder(strings.NewReader(input))
+
+	e, err := d.Next()
+
+	require.NoError(t, err)
+
+	assert.Equal(t, "output", e.Type)
+	assert.Equal(t, "123abc", e.ID)
+	assert.Equal(t, "giraffe\n", e.Data)
+}
+
+func TestDecodeOneEventWithSpace(t *testing.T) {
+	input := `event: output
+id: 123abc
+data:   giraffe
+
+`
+	d := sse.NewDecoder(strings.NewReader(input))
+
+	e, err := d.Next()
+
+	require.NoError(t, err)
+
+	assert.Equal(t, "output", e.Type)
+	assert.Equal(t, "123abc", e.ID)
+	// only one space should be trimmed
+	assert.Equal(t, "  giraffe\n", e.Data)
+}
+
+func TestDecodeOneEventMultipleData(t *testing.T) {
+	input := `event:output
+data:giraffe
+data:rhino
+data:wombat
+
+`
+	d := sse.NewDecoder(strings.NewReader(input))
+
+	e, err := d.Next()
+
+	require.NoError(t, err)
+
+	assert.Equal(t, "output", e.Type)
+	assert.Equal(t, "giraffe\nrhino\nwombat\n", e.Data)
+}
+
+func TestDecodeOneEventHugeData(t *testing.T) {
+	// this test is mainly to make sure we're not constrained by the
+	// bufio.Reader buffer size
+	input := fmt.Sprintf(`event:output
+data:%s
+
+`, strings.Repeat("0123456789abcdef", 1_000_000))
+	d := sse.NewDecoder(strings.NewReader(input))
+
+	e, err := d.Next()
+
+	require.NoError(t, err)
+
+	assert.Equal(t, "output", e.Type)
+	// 16_000_000 data bytes and the terminal LF character
+	assert.Equal(t, 16_000_001, len(e.Data))
+}
+
+func TestDecodeManyEvents(t *testing.T) {
+	input := `event:output
+id:alpha1
+data:giraffe
+
+event:output
+id:bravo2
+data:rhino
+
+event:output
+id:gamma3
+data:pine marten
+
+`
+	d := sse.NewDecoder(strings.NewReader(input))
+
+	e, err := d.Next()
+
+	require.NoError(t, err)
+
+	assert.Equal(t, "output", e.Type)
+	assert.Equal(t, "alpha1", e.ID)
+	assert.Equal(t, "giraffe\n", e.Data)
+
+	e, err = d.Next()
+
+	require.NoError(t, err)
+
+	assert.Equal(t, "output", e.Type)
+	assert.Equal(t, "bravo2", e.ID)
+	assert.Equal(t, "rhino\n", e.Data)
+
+	e, err = d.Next()
+
+	require.NoError(t, err)
+
+	assert.Equal(t, "output", e.Type)
+	assert.Equal(t, "gamma3", e.ID)
+	assert.Equal(t, "pine marten\n", e.Data)
+}
+
+func TestDecodeEarlyEOF(t *testing.T) {
+	input := ``
+	d := sse.NewDecoder(strings.NewReader(input))
+
+	_, err := d.Next()
+
+	assert.ErrorIs(t, err, io.ErrUnexpectedEOF)
+}

internal/sse/stream.go

@@ -0,0 +1,124 @@
+package sse
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"time"
+)
+
+// Backoff is a copy of replicate.Backoff to avoid import cycles
+type Backoff interface {
+	NextDelay(retries int) time.Duration
+}
+
+type Streamer struct {
+	c          *http.Client
+	url        string
+	maxRetries int
+	backoff    Backoff
+
+	attempt     int
+	lastEventID string
+
+	decoder       *Decoder
+	currentStream io.ReadCloser
+}
+
+func NewStreamer(c *http.Client, url string, maxRetries int, backoff Backoff) *Streamer {
+	return &Streamer{
+		c:          c,
+		url:        url,
+		maxRetries: maxRetries,
+		backoff:    backoff,
+	}
+}
+
+var ErrMaximumRetries = errors.New("Exceeded maximum retries")
+
+// connect (re-)establishes the connection to the SSE server. It only returns an
+// error if it cannot recover through retries.
+func (s *Streamer) connect(ctx context.Context) error {
+	for {
+		if s.attempt > s.maxRetries {
+			return ErrMaximumRetries
+		}
+
+		delay := 0 * time.Second
+		if s.attempt > 0 {
+			// delay on connection retry
+			delay = s.backoff.NextDelay(s.attempt - 1)
+		}
+		s.attempt++
+		reconnectDelay := time.NewTimer(delay)
+		// once we only support go 1.23+, we can use time.After() here and simplify
+		defer reconnectDelay.Stop()
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-reconnectDelay.C:
+		}
+
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.url, nil)
+		if err != nil {
+			return err
+		}
+		req.Header.Set("Accept", "text/event-stream")
+
+		if s.lastEventID != "" {
+			req.Header.Set("Last-Event-ID", s.lastEventID)
+		}
+
+		//nolint:bodyclose
+		resp, err := s.c.Do(req)
+		if err != nil {
+			// try again
+			continue
+		}
+
+		if resp.StatusCode != http.StatusOK {
+			return fmt.Errorf("received invalid status code: %d", resp.StatusCode)
+		}
+
+		if s.currentStream != nil {
+			err = s.currentStream.Close()
+			if err != nil {
+				return err
+			}
+		}
+		s.currentStream = resp.Body
+		s.decoder = NewDecoder(s.currentStream)
+		return nil
+	}
+}
+
+func (s *Streamer) NextEvent(ctx context.Context) (*Event, error) {
+	if s.decoder == nil {
+		if err := s.connect(ctx); err != nil {
+			return nil, err
+		}
+	}
+	for {
+		e, err := s.decoder.Next()
+		if err != nil {
+			if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
+				if err = s.connect(ctx); err != nil {
+					return nil, err
+				}
+				continue
+			}
+			return nil, err
+		}
+		s.lastEventID = e.ID
+		return &e, nil
+	}
+}
+
+func (s *Streamer) Close() error {
+	if s.currentStream != nil {
+		return s.currentStream.Close()
+	}
+	return nil
+}