Skip to content

Commit

Permalink
Emit done and error events to SSE channel (#34)
Browse files Browse the repository at this point in the history
* Add String() method to SSEEvent

* Emit done event

* Don't send stream error events to error channel
  • Loading branch information
mattt committed Dec 4, 2023
1 parent c5d8c36 commit 1db43b6
Showing 1 changed file with 11 additions and 6 deletions.
17 changes: 11 additions & 6 deletions stream.go
Expand Up @@ -58,6 +58,15 @@ func (e *SSEEvent) decode(b []byte) error {
return nil
}

func (e *SSEEvent) String() string {
switch e.Type {
case "output":
return e.Data
default:
return ""
}
}

func (r *Client) Stream(ctx context.Context, identifier string, input PredictionInput, webhook *Webhook) (<-chan SSEEvent, <-chan error) {
sseChan := make(chan SSEEvent, 64)
errChan := make(chan error, 64)
Expand Down Expand Up @@ -177,14 +186,10 @@ func (r *Client) streamPrediction(ctx context.Context, prediction *Prediction, l
errChan <- err
}

switch event.Type {
case "error":
errChan <- unmarshalAPIError(nil, []byte(event.Data))
case "done":
sseChan <- event
if event.Type == "done" {
close(done)
return
default:
sseChan <- event
}
}
}
Expand Down

0 comments on commit 1db43b6

Please sign in to comment.