初始提交：边缘AI算力机统一AI通讯层

internal/handler/sse.go

@@ -0,0 +1,139 @@
+package handler
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+
+	"github.com/edgeai/gateway/internal/adapter"
+	"github.com/edgeai/gateway/pkg/api"
+)
+
+// SSEWriter writes Server-Sent Events to an HTTP response.
+type SSEWriter struct {
+	w       http.ResponseWriter
+	flusher http.Flusher
+}
+
+// NewSSEWriter creates a new SSEWriter. Returns nil if streaming is not supported.
+func NewSSEWriter(w http.ResponseWriter) *SSEWriter {
+	flusher, ok := w.(http.Flusher)
+	if !ok {
+		return nil
+	}
+
+	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")
+
+	return &SSEWriter{w: w, flusher: flusher}
+}
+
+// WriteChunk writes a single SSE data event.
+func (s *SSEWriter) WriteChunk(data any) error {
+	jsonData, err := json.Marshal(data)
+	if err != nil {
+		return fmt.Errorf("marshal sse data: %w", err)
+	}
+	fmt.Fprintf(s.w, "data: %s\n\n", jsonData)
+	s.flusher.Flush()
+	return nil
+}
+
+// WriteDone writes the [DONE] marker.
+func (s *SSEWriter) WriteDone() {
+	fmt.Fprintf(s.w, "data: [DONE]\n\n")
+	s.flusher.Flush()
+}
+
+// StreamChatCompletion streams chunks from an adapter to the client in OpenAI SSE format.
+func StreamChatCompletion(sse *SSEWriter, ch <-chan adapter.StreamChunk, requestID, taskID, model string) (int, int, error) {
+	inputTokens := 0
+	outputTokens := 0
+
+	for chunk := range ch {
+		if chunk.Error != nil {
+			return inputTokens, outputTokens, chunk.Error
+		}
+
+		if chunk.Done {
+			if chunk.InputTokens > 0 {
+				inputTokens = chunk.InputTokens
+			}
+			if chunk.OutputTokens > 0 {
+				outputTokens = chunk.OutputTokens
+			}
+
+			// Write final chunk with finish_reason
+			sseChunk := map[string]any{
+				"id":      requestID,
+				"object":  "chat.completion.chunk",
+				"model":   model,
+				"choices": []map[string]any{
+					{
+						"index":         0,
+						"delta":         map[string]any{},
+						"finish_reason": chunk.FinishReason,
+					},
+				},
+			}
+			if inputTokens > 0 || outputTokens > 0 {
+				sseChunk["usage"] = map[string]int{
+					"input_tokens":  inputTokens,
+					"output_tokens": outputTokens,
+					"total_tokens":  inputTokens + outputTokens,
+				}
+			}
+			sse.WriteChunk(sseChunk)
+			sse.WriteDone()
+			return inputTokens, outputTokens, nil
+		}
+
+		// Write content delta
+		sseChunk := map[string]any{
+			"id":      requestID,
+			"object":  "chat.completion.chunk",
+			"model":   model,
+			"choices": []map[string]any{
+				{
+					"index": 0,
+					"delta": map[string]any{
+						"content": chunk.Delta,
+					},
+					"finish_reason": nil,
+				},
+			},
+		}
+		sse.WriteChunk(sseChunk)
+	}
+
+	return inputTokens, outputTokens, nil
+}
+
+// BuildChatResponse creates a non-streaming ChatResponse from adapter result.
+func BuildChatResponse(requestID, taskID, logicalModel string, resp *adapter.ChatResponse) api.ChatResponse {
+	return api.ChatResponse{
+		RequestID:   requestID,
+		TaskID:      taskID,
+		Status:      "completed",
+		Model:       logicalModel,
+		Choices: []api.Choice{
+			{
+				Index: 0,
+				Message: &api.Message{
+					Role:    "assistant",
+					Content: resp.Content,
+				},
+				FinishReason: resp.FinishReason,
+			},
+		},
+		LogicalModel: logicalModel,
+		ActualModel:  resp.ActualModel,
+		Usage: &api.Usage{
+			InputTokens:  resp.InputTokens,
+			OutputTokens: resp.OutputTokens,
+			TotalTokens:  resp.InputTokens + resp.OutputTokens,
+		},
+	}
+}