ce7cb702218693d08ea6d83464d2457c9250027c

services/vertex-batch/front.go

@@ -0,0 +1,136 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"strings"
+	"time"
+)
+
+type frontend struct {
+	client      *http.Client
+	gateway     string
+	prefix      string
+	routePrefix string
+	heartbeat   time.Duration
+}
+
+func (f frontend) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	data, err := io.ReadAll(r.Body)
+	var body map[string]json.RawMessage
+	if err != nil || json.Unmarshal(data, &body) != nil {
+		http.Error(w, "invalid JSON", http.StatusBadRequest)
+		return
+	}
+	var model string
+	var streaming bool
+	if json.Unmarshal(body["model"], &model) != nil || !strings.HasPrefix(model, f.prefix) || strings.TrimPrefix(model, f.prefix) == "" {
+		http.Error(w, "invalid batch model", http.StatusBadRequest)
+		return
+	}
+	if raw, ok := body["stream"]; ok {
+		if json.Unmarshal(raw, &streaming) != nil {
+			http.Error(w, "invalid stream flag", http.StatusBadRequest)
+			return
+		}
+	}
+	body["model"], err = json.Marshal(f.routePrefix + strings.TrimPrefix(model, f.prefix))
+	if err != nil {
+		writeError(w)
+		return
+	}
+	data, err = json.Marshal(body)
+	if err != nil {
+		writeError(w)
+		return
+	}
+	ctx, cancel := context.WithCancel(r.Context())
+	defer cancel()
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, f.gateway, bytes.NewReader(data))
+	if err != nil {
+		writeError(w)
+		return
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Authorization", r.Header.Get("Authorization"))
+	if streaming {
+		f.stream(w, req)
+		return
+	}
+	res, err := f.client.Do(req)
+	if err != nil {
+		writeError(w)
+		return
+	}
+	defer func() { _ = res.Body.Close() }()
+	w.Header().Set("Content-Type", res.Header.Get("Content-Type"))
+	w.WriteHeader(res.StatusCode)
+	_, _ = io.Copy(w, res.Body)
+}
+
+type fragment struct {
+	data []byte
+	end  bool
+}
+
+func (f frontend) receive(req *http.Request, incoming chan<- fragment) {
+	send := func(part fragment) bool {
+		select {
+		case incoming <- part:
+			return true
+		case <-req.Context().Done():
+			return false
+		}
+	}
+	res, err := f.client.Do(req)
+	if err != nil {
+		send(fragment{end: true})
+		return
+	}
+	defer func() { _ = res.Body.Close() }()
+	if res.StatusCode < http.StatusOK || res.StatusCode >= http.StatusMultipleChoices {
+		send(fragment{end: true})
+		return
+	}
+	buffer := make([]byte, 4096)
+	for {
+		n, err := res.Body.Read(buffer)
+		if !send(fragment{data: bytes.Clone(buffer[:n]), end: err != nil}) || err != nil {
+			return
+		}
+	}
+}
+
+func (f frontend) stream(w http.ResponseWriter, req *http.Request) {
+	incoming := make(chan fragment)
+	go f.receive(req, incoming)
+	w.Header().Set("Content-Type", "text/event-stream")
+	w.Header().Set("Cache-Control", "no-cache")
+	flush := http.NewResponseController(w)
+	if flush.Flush() != nil {
+		return
+	}
+	ticker := time.NewTicker(f.heartbeat)
+	defer ticker.Stop()
+	var parser events
+	for {
+		var output []byte
+		select {
+		case <-req.Context().Done():
+			return
+		case <-ticker.C:
+			output = []byte(": batch pending\n\n")
+		case part := <-incoming:
+			output = parser.feed(part.data, part.end)
+		}
+		if _, err := w.Write(output); err != nil {
+			return
+		}
+		if flush.Flush() != nil || parser.done {
+			return
+		}
+	}
+}