58 lines
1.3 KiB
Go
58 lines
1.3 KiB
Go
package server
|
||
|
||
import (
|
||
"encoding/json"
|
||
"fmt"
|
||
"net/http"
|
||
|
||
"github.com/daniel/cursor-adapter/internal/types"
|
||
)
|
||
|
||
// SSEWriter 封裝 http.ResponseWriter 用於 SSE streaming。
|
||
type SSEWriter struct {
|
||
w http.ResponseWriter
|
||
flush http.Flusher
|
||
}
|
||
|
||
// NewSSEWriter 建立 SSEWriter，設定必要的 headers。
|
||
func NewSSEWriter(w http.ResponseWriter) *SSEWriter {
|
||
w.Header().Set("Content-Type", "text/event-stream")
|
||
w.Header().Set("Cache-Control", "no-cache")
|
||
w.Header().Set("Connection", "keep-alive")
|
||
w.Header().Set("X-Accel-Buffering", "no")
|
||
|
||
flusher, _ := w.(http.Flusher)
|
||
|
||
return &SSEWriter{w: w, flush: flusher}
|
||
}
|
||
|
||
// WriteChunk 寫入一個 SSE chunk。
|
||
func (s *SSEWriter) WriteChunk(chunk types.ChatCompletionChunk) error {
|
||
data, err := json.Marshal(chunk)
|
||
if err != nil {
|
||
return fmt.Errorf("marshal chunk: %w", err)
|
||
}
|
||
fmt.Fprintf(s.w, "data: %s\n\n", data)
|
||
if s.flush != nil {
|
||
s.flush.Flush()
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// WriteDone 寫入 SSE 結束標記。
|
||
func (s *SSEWriter) WriteDone() {
|
||
fmt.Fprint(s.w, "data: [DONE]\n\n")
|
||
if s.flush != nil {
|
||
s.flush.Flush()
|
||
}
|
||
}
|
||
|
||
// WriteError 寫入 SSE 格式的錯誤。
|
||
func (s *SSEWriter) WriteError(errMsg string) {
|
||
stopReason := "stop"
|
||
chunk := types.NewChatCompletionChunk("error", 0, "", types.Delta{Content: &errMsg})
|
||
chunk.Choices[0].FinishReason = &stopReason
|
||
s.WriteChunk(chunk)
|
||
s.WriteDone()
|
||
}
|