feat: Handle SSE errors

When an error happens during LLM generation, an error is sent as a
message.
See https://platform.openai.com/docs/api-reference/responses-streaming/error

We now handle it, to transmit the error to the client

Author: paultranvan

model/rag/chat.go

@@ -303,6 +303,17 @@ func Query(inst *instance.Instance, logger logger.Logger, query QueryMessage) er
 	})
 
 	if err != nil {
+		// Send error event to client
+		errorDoc := map[string]interface{}{
+			"_id":     msg.ID,
+			"object":  "error",
+			"message": err.Error(),
+		}
+		errorPayload := couchdb.JSONDoc{
+			Type: consts.ChatEvents,
+			M:    errorDoc,
+		}
+		realtime.GetHub().Publish(inst, realtime.EventCreate, &errorPayload, nil)
 		return err
 	}
 
@@ -369,6 +380,18 @@ func foreachSSE(r io.Reader, fn func(event map[string]interface{})) error {
 		if err := json.Unmarshal(data, &event); err != nil {
 			return err
 		}
+		// Check for error event from the server
+		if errObj, ok := event["error"].(map[string]interface{}); ok {
+			message, _ := errObj["message"].(string)
+			code, _ := errObj["code"].(string)
+			if message == "" {
+				message = "unknown streaming error"
+			}
+			if code != "" {
+				return fmt.Errorf("%s: %s", code, message)
+			}
+			return errors.New(message)
+		}
 		fn(event)
 	}
 	return nil

model/rag/chat_test.go

@@ -0,0 +1,163 @@
+package rag
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestForeachSSE(t *testing.T) {
+	t.Run("normal events are passed to callback", func(t *testing.T) {
+		input := `data: {"object":"chat.completion.chunk","content":"hello"}
+
+data: {"object":"chat.completion.chunk","content":"world"}
+
+data: [DONE]
+`
+		var events []map[string]interface{}
+		err := foreachSSE(strings.NewReader(input), func(event map[string]interface{}) {
+			events = append(events, event)
+		})
+
+		require.NoError(t, err)
+		assert.Len(t, events, 2)
+		assert.Equal(t, "hello", events[0]["content"])
+		assert.Equal(t, "world", events[1]["content"])
+	})
+
+	t.Run("error event with code and message returns formatted error", func(t *testing.T) {
+		input := `data: {"error":{"message":"Error while generating answer","code":"ERROR_ANSWER_GENERATION"}}
+`
+		var events []map[string]interface{}
+		err := foreachSSE(strings.NewReader(input), func(event map[string]interface{}) {
+			events = append(events, event)
+		})
+
+		require.Error(t, err)
+		assert.Equal(t, "ERROR_ANSWER_GENERATION: Error while generating answer", err.Error())
+		assert.Empty(t, events)
+	})
+
+	t.Run("error event with only message returns error", func(t *testing.T) {
+		input := `data: {"error":{"message":"Something went wrong"}}
+`
+		var events []map[string]interface{}
+		err := foreachSSE(strings.NewReader(input), func(event map[string]interface{}) {
+			events = append(events, event)
+		})
+
+		require.Error(t, err)
+		assert.Equal(t, "Something went wrong", err.Error())
+		assert.Empty(t, events)
+	})
+
+	t.Run("error event with empty message returns unknown error", func(t *testing.T) {
+		input := `data: {"error":{"message":"","code":"SOME_CODE"}}
+`
+		var events []map[string]interface{}
+		err := foreachSSE(strings.NewReader(input), func(event map[string]interface{}) {
+			events = append(events, event)
+		})
+
+		require.Error(t, err)
+		assert.Equal(t, "SOME_CODE: unknown streaming error", err.Error())
+		assert.Empty(t, events)
+	})
+
+	t.Run("error event with no message field returns unknown error", func(t *testing.T) {
+		input := `data: {"error":{}}
+`
+		var events []map[string]interface{}
+		err := foreachSSE(strings.NewReader(input), func(event map[string]interface{}) {
+			events = append(events, event)
+		})
+
+		require.Error(t, err)
+		assert.Equal(t, "unknown streaming error", err.Error())
+		assert.Empty(t, events)
+	})
+
+	t.Run("DONE stops processing", func(t *testing.T) {
+		input := `data: {"object":"first"}
+
+data: [DONE]
+data: {"object":"should not be processed"}
+`
+		var events []map[string]interface{}
+		err := foreachSSE(strings.NewReader(input), func(event map[string]interface{}) {
+			events = append(events, event)
+		})
+
+		require.NoError(t, err)
+		assert.Len(t, events, 1)
+		assert.Equal(t, "first", events[0]["object"])
+	})
+
+	t.Run("invalid SSE format returns error", func(t *testing.T) {
+		input := `invalid line without colon
+`
+		err := foreachSSE(strings.NewReader(input), func(event map[string]interface{}) {})
+
+		require.Error(t, err)
+		assert.Equal(t, "invalid SSE response", err.Error())
+	})
+
+	t.Run("invalid JSON returns error", func(t *testing.T) {
+		input := `data: {invalid json}
+`
+		err := foreachSSE(strings.NewReader(input), func(event map[string]interface{}) {})
+
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "invalid character")
+	})
+
+	t.Run("comments are skipped", func(t *testing.T) {
+		input := `: this is a comment
+data: {"object":"event"}
+
+data: [DONE]
+`
+		var events []map[string]interface{}
+		err := foreachSSE(strings.NewReader(input), func(event map[string]interface{}) {
+			events = append(events, event)
+		})
+
+		require.NoError(t, err)
+		assert.Len(t, events, 1)
+	})
+
+	t.Run("empty lines are skipped", func(t *testing.T) {
+		input := `
+
+data: {"object":"event"}
+
+
+data: [DONE]
+`
+		var events []map[string]interface{}
+		err := foreachSSE(strings.NewReader(input), func(event map[string]interface{}) {
+			events = append(events, event)
+		})
+
+		require.NoError(t, err)
+		assert.Len(t, events, 1)
+	})
+
+	t.Run("non-data fields are skipped", func(t *testing.T) {
+		input := `event: message
+id: 123
+data: {"object":"event"}
+
+data: [DONE]
+`
+		var events []map[string]interface{}
+		err := foreachSSE(strings.NewReader(input), func(event map[string]interface{}) {
+			events = append(events, event)
+		})
+
+		require.NoError(t, err)
+		assert.Len(t, events, 1)
+	})
+}