ce7cb702218693d08ea6d83464d2457c9250027c

services/vertex-batch/http.go

@@ -0,0 +1,103 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"log/slog"
+	"net/http"
+	"sync"
+	"time"
+)
+
+type batchRunner interface {
+	run(context.Context, string, json.RawMessage) (json.RawMessage, error)
+}
+
+type handler struct {
+	runner    batchRunner
+	heartbeat time.Duration
+}
+
+func (h handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if r.Method == http.MethodGet && r.URL.Path == "/healthz" {
+		w.WriteHeader(http.StatusOK)
+		return
+	}
+	if r.Method != http.MethodPost || r.URL.Path != "/v1/chat/completions" {
+		http.NotFound(w, r)
+		return
+	}
+	model := r.Header.Get("x-batch-model")
+	streaming := r.Header.Get("x-batch-stream")
+	if model == "" || (streaming != "true" && streaming != "false") {
+		http.Error(w, "invalid batch model or stream headers", http.StatusBadRequest)
+		return
+	}
+	native, err := io.ReadAll(r.Body)
+	if err != nil || !json.Valid(native) {
+		http.Error(w, "invalid Gemini request", http.StatusBadRequest)
+		return
+	}
+	ctx, cancel := context.WithCancel(r.Context())
+	type outcome struct {
+		body json.RawMessage
+		err  error
+	}
+	done := make(chan outcome, 1)
+	var workers sync.WaitGroup
+	workers.Go(func() {
+		body, err := h.runner.run(ctx, model, native)
+		done <- outcome{body: body, err: err}
+	})
+	defer workers.Wait()
+	defer cancel()
+	if streaming != "true" {
+		select {
+		case <-ctx.Done():
+			return
+		case result := <-done:
+			if result.err != nil {
+				writeError(w)
+				return
+			}
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = w.Write(result.body)
+			return
+		}
+	}
+	w.Header().Set("Content-Type", "text/event-stream")
+	w.Header().Set("Cache-Control", "no-cache")
+	flush := http.NewResponseController(w)
+	if err := flush.Flush(); err != nil {
+		return
+	}
+	ticker := time.NewTicker(h.heartbeat)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			if _, err := io.WriteString(w, ": batch pending\n\n"); err != nil {
+				return
+			}
+			if err := flush.Flush(); err != nil {
+				return
+			}
+		case result := <-done:
+			if result.err != nil {
+				slog.Error("batch failed", "kind", "cloud")
+				return
+			}
+			_, _ = io.WriteString(w, "data: "+string(result.body)+"\n\n")
+			return
+		}
+	}
+}
+
+func writeError(w http.ResponseWriter) {
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(http.StatusBadGateway)
+	_, _ = io.WriteString(w, `{"error":{"code":502,"status":"INTERNAL","message":"Batch did not produce a result"}}`)
+}