136 lines
3.1 KiB
Go
136 lines
3.1 KiB
Go
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
|
|
}
|
|
}
|
|
}
|