iop/apps/edge/internal/openai/stream.go

package openai

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"go.uber.org/zap"

	edgeservice "iop/apps/edge/internal/service"
)

func (s *Server) streamChatCompletion(w http.ResponseWriter, r *http.Request, req chatCompletionRequest, handle edgeservice.RunResult, outputPolicy strictOutputPolicy) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		writeError(w, http.StatusInternalServerError, "streaming_not_supported", "response writer does not support streaming")
		return
	}
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")

	created := time.Now().Unix()
	id := "chatcmpl-" + handle.Dispatch().RunID
	model := responseModel(req.Model, handle.Dispatch().Target)
	writeSSE(w, flusher, chatCompletionChunk{
		ID:      id,
		Object:  "chat.completion.chunk",
		Created: created,
		Model:   model,
		Choices: []chatCompletionChunkChoice{{
			Index: 0,
			Delta: chatDelta{Role: "assistant"},
		}},
	})

	if outputPolicy.Strict && outputPolicy.StreamBuffer {
		s.streamBufferedChatCompletion(w, r, handle, id, created, model, flusher, outputPolicy)
		return
	}

	var contentBuilder strings.Builder
	var reasoningBuilder strings.Builder
	defer func() {
		s.logger.Info("openai chat completion stream closed",
			zap.String("run_id", handle.Dispatch().RunID),
			zap.Bool("strict_output", outputPolicy.Strict),
			zap.Bool("strict_stream_buffer", outputPolicy.StreamBuffer),
			zap.Int("content_len", contentBuilder.Len()),
			zap.Int("reasoning_len", reasoningBuilder.Len()),
			zap.String("content_preview", previewString(contentBuilder.String(), 1000)),
			zap.String("reasoning_preview", previewString(reasoningBuilder.String(), 1000)),
		)
	}()

	stream := handle.Stream()
	for {
		select {
		case <-r.Context().Done():
			return
		case nodeEvent := <-stream.NodeEvents:
			if edgeservice.IsNodeDisconnected(nodeEvent) {
				writeSSEError(w, flusher, "node disconnected")
				return
			}
		case event := <-stream.Events:
			if event == nil {
				continue
			}
			switch event.GetType() {
			case "delta":
				if event.GetDelta() == "" {
					continue
				}
				contentBuilder.WriteString(event.GetDelta())
				writeSSE(w, flusher, chatCompletionChunk{
					ID:      id,
					Object:  "chat.completion.chunk",
					Created: created,
					Model:   model,
					Choices: []chatCompletionChunkChoice{{
						Index: 0,
						Delta: chatDelta{Content: event.GetDelta()},
					}},
				})
			case "reasoning_delta":
				if event.GetDelta() == "" {
					continue
				}
				reasoningBuilder.WriteString(event.GetDelta())
				if outputPolicy.Strict {
					continue
				}
				writeSSE(w, flusher, chatCompletionChunk{
					ID:      id,
					Object:  "chat.completion.chunk",
					Created: created,
					Model:   model,
					Choices: []chatCompletionChunkChoice{{
						Index: 0,
						Delta: chatDelta{ReasoningContent: event.GetDelta()},
					}},
				})
			case "complete":
				writeSSE(w, flusher, chatCompletionChunk{
					ID:      id,
					Object:  "chat.completion.chunk",
					Created: created,
					Model:   model,
					Choices: []chatCompletionChunkChoice{{
						Index:        0,
						Delta:        chatDelta{},
						FinishReason: "stop",
					}},
				})
				fmt.Fprint(w, "data: [DONE]\n\n")
				flusher.Flush()
				return
			case "error", "cancelled":
				msg := event.GetError()
				if msg == "" {
					msg = event.GetMessage()
				}
				if msg == "" {
					msg = "run failed"
				}
				writeSSEError(w, flusher, msg)
				return
			}
		case <-time.After(handle.WaitTimeout()):
			writeSSEError(w, flusher, "run timed out")
			return
		}
	}
}

func (s *Server) streamBufferedChatCompletion(w http.ResponseWriter, r *http.Request, handle edgeservice.RunResult, id string, created int64, model string, flusher http.Flusher, outputPolicy strictOutputPolicy) {
	text, reasoning, _, err := collectRunResult(r.Context(), handle.Stream(), handle.WaitTimeout())
	if err != nil {
		writeSSEError(w, flusher, err.Error())
		return
	}
	text, reasoning, normalized := normalizeCompletionOutput(outputPolicy, text, reasoning)
	s.logger.Info("openai chat completion stream closed",
		zap.String("run_id", handle.Dispatch().RunID),
		zap.Bool("strict_output", outputPolicy.Strict),
		zap.Bool("strict_stream_buffer", outputPolicy.StreamBuffer),
		zap.String("xml_completion_tool", outputPolicy.XMLCompletionTool),
		zap.Bool("normalized", normalized),
		zap.Int("content_len", len(text)),
		zap.Int("reasoning_len", len(reasoning)),
		zap.String("content_preview", previewString(text, 1000)),
		zap.String("reasoning_preview", previewString(reasoning, 1000)),
	)
	if text != "" {
		writeSSE(w, flusher, chatCompletionChunk{
			ID:      id,
			Object:  "chat.completion.chunk",
			Created: created,
			Model:   model,
			Choices: []chatCompletionChunkChoice{{
				Index: 0,
				Delta: chatDelta{Content: text},
			}},
		})
	}
	writeSSE(w, flusher, chatCompletionChunk{
		ID:      id,
		Object:  "chat.completion.chunk",
		Created: created,
		Model:   model,
		Choices: []chatCompletionChunkChoice{{
			Index:        0,
			Delta:        chatDelta{},
			FinishReason: "stop",
		}},
	})
	fmt.Fprint(w, "data: [DONE]\n\n")
	flusher.Flush()
}

func writeSSE(w http.ResponseWriter, flusher http.Flusher, v any) {
	b, err := json.Marshal(v)
	if err != nil {
		return
	}
	fmt.Fprintf(w, "data: %s\n\n", b)
	flusher.Flush()
}

func writeSSEError(w http.ResponseWriter, flusher http.Flusher, message string) {
	writeSSE(w, flusher, errorResponse{Error: errorBody{Type: "run_error", Message: message}})
	fmt.Fprint(w, "data: [DONE]\n\n")
	flusher.Flush()
}