rony-chat-bot/internal/streaming/sse.go

package streaming

import (
	"encoding/json"
	"fmt"
	"net/http"
)

type Event struct {
	Type    string `json:"type"`
	Content string `json:"content,omitempty"`
}

func WriteEvent(w http.ResponseWriter, eventType string, payload any) error {
	data, err := json.Marshal(payload)
	if err != nil {
		return fmt.Errorf("marshal event: %w", err)
	}
	if _, err := fmt.Fprintf(w, "event: %s\ndata: %s\n\n", eventType, data); err != nil {
		return err
	}
	if f, ok := w.(http.Flusher); ok {
		f.Flush()
	}
	return nil
}

func WriteChunk(w http.ResponseWriter, content string) error {
	return WriteEvent(w, "chunk", Event{Type: "chunk", Content: content})
}

func WriteStart(w http.ResponseWriter, conversationID string) error {
	return WriteEvent(w, "start", map[string]any{
		"type":            "start",
		"conversation_id": conversationID,
	})
}

func WriteSources(w http.ResponseWriter, sources []string) error {
	return WriteEvent(w, "sources", map[string]any{
		"type":     "sources",
		"documents": sources,
	})
}

// WriteCompaction surfaces a context-compaction event so the client can
// render a "summary injected" hint. Always emitted right after the
// `sources` event (before the streamed chunks), and only when compaction
// actually fired for this request.
func WriteCompaction(w http.ResponseWriter, payload map[string]any) error {
	out := map[string]any{"type": "compaction"}
	for k, v := range payload {
		out[k] = v
	}
	return WriteEvent(w, "compaction", out)
}

type Usage struct {
	InputTokens  int `json:"input_tokens"`
	OutputTokens int `json:"output_tokens"`
}

func WriteDone(w http.ResponseWriter, usage Usage) error {
	return WriteEvent(w, "done", map[string]any{
		"type":  "done",
		"usage": usage,
	})
}

func WriteError(w http.ResponseWriter, errMsg string) error {
	return WriteEvent(w, "error", map[string]string{
		"type":  "error",
		"error": errMsg,
	})
}