103 lines
2.4 KiB
Go
103 lines
2.4 KiB
Go
|
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"}}`)
|
||
|
|
}
|