diff --git a/main.go b/main.go index 9e7b8e0..62fe04a 100644 --- a/main.go +++ b/main.go @@ -135,20 +135,20 @@ func main() { defer stream.Close() + flusher, ok := c.Writer.(http.Flusher) + if !ok { + c.JSON(http.StatusInternalServerError, gin.H{"error": "streaming not supported"}) + return + } + header := c.Writer.Header() header.Set("Content-Type", "text/event-stream; charset=utf-8") header.Set("Cache-Control", "no-cache") header.Set("Connection", "keep-alive") - if _, ok := c.Writer.(http.Flusher); !ok { - c.JSON(http.StatusInternalServerError, gin.H{"error": "streaming not supported"}) - return - } - if flusher, ok := c.Writer.(http.Flusher); ok { - flusher.Flush() - } - c.Status(http.StatusOK) + c.Writer.WriteHeaderNow() + flusher.Flush() replacer := strings.NewReplacer("\r\n", "\n", "\r", "\n") @@ -160,25 +160,30 @@ func main() { if data, ok := payload["data"].(string); ok { sanitized["data"] = replacer.Replace(data) } - bytes, err := json.Marshal(sanitized) + + body, err := json.Marshal(sanitized) if err != nil { log.Printf("failed to marshal sse payload: %v", err) return } + if event != "" { - fmt.Fprintf(w, "event: %s\n", event) + if _, err := fmt.Fprintf(w, "event: %s\n", event); err != nil { + log.Printf("failed to write sse event: %v", err) + return + } } - fmt.Fprintf(w, "data: %s\n\n", bytes) + + if _, err := fmt.Fprintf(w, "data: %s\n\n", body); err != nil { + log.Printf("failed to write sse data: %v", err) + return + } + + flusher.Flush() } - sentDone := false - c.Stream(func(w io.Writer) bool { - if sentDone { - return false - } - - for stream.Next() { + if stream.Next() { chunk := stream.Current() var builder strings.Builder for _, choice := range chunk.Choices { @@ -186,19 +191,21 @@ func main() { builder.WriteString(choice.Delta.Content) } } + if builder.Len() == 0 { - continue + return true } + writeSSE(w, "chunk", gin.H{"data": builder.String()}) return true } - sentDone = true if err := stream.Err(); err != nil { writeSSE(w, "error", gin.H{"message": err.Error()}) } + writeSSE(w, "done", gin.H{"data": ""}) - return true + return false }) })