LovroG05
8bb62a1c20
All checks were successful
ci/woodpecker/push/woodpecker Pipeline was successful
Any model can name another model as its backup. The gateway tries the backup when the model is offline, rate limited, or its upstream fails before the client has received anything. The response reports the requested model, so the client sees no switch. The request log records each failed attempt and the address that served the request, and the cost uses the backup's prices.
412 lines
14 KiB
Go
412 lines
14 KiB
Go
package api
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"maps"
|
|
"net"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
chimw "github.com/go-chi/chi/v5/middleware"
|
|
"github.com/google/uuid"
|
|
|
|
"git.govekar.net/LovroG05/localmodelroutex/internal/billing"
|
|
"git.govekar.net/LovroG05/localmodelroutex/internal/catalog"
|
|
"git.govekar.net/LovroG05/localmodelroutex/internal/providers"
|
|
"git.govekar.net/LovroG05/localmodelroutex/internal/store"
|
|
)
|
|
|
|
const maxRequestBody = 10 << 20 // 10 MiB
|
|
|
|
func (a *API) chatCompletions(w http.ResponseWriter, r *http.Request) {
|
|
body, err := io.ReadAll(io.LimitReader(r.Body, maxRequestBody))
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, "failed to read request body")
|
|
return
|
|
}
|
|
req, err := providers.ParseChatRequest(body)
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, err.Error())
|
|
return
|
|
}
|
|
if req.Model == "" {
|
|
writeError(w, http.StatusBadRequest, "missing 'model'")
|
|
return
|
|
}
|
|
|
|
key := keyFrom(r.Context())
|
|
address := req.Model
|
|
|
|
route, err := a.Cache.Route(address)
|
|
if err != nil {
|
|
writeError(w, http.StatusNotFound, err.Error())
|
|
return
|
|
}
|
|
if !allowed(key, route.Primary) {
|
|
writeError(w, http.StatusForbidden, fmt.Sprintf("model %q is not allowed for this key", address))
|
|
return
|
|
}
|
|
|
|
messages := marshalMessages(req.Payload["messages"])
|
|
|
|
var last *failure
|
|
for i, entry := range route.Targets {
|
|
if i > 0 && r.Context().Err() != nil {
|
|
break // the client is gone; a backup cannot help
|
|
}
|
|
prov, fail := a.admit(entry, 1)
|
|
if fail != nil {
|
|
last = fail
|
|
continue
|
|
}
|
|
cc, ok := prov.(providers.ChatCompleter)
|
|
if !ok {
|
|
last = &failure{status: http.StatusNotImplemented, msg: "provider does not support chat completions"}
|
|
continue
|
|
}
|
|
attempt := &providers.ChatRequest{Model: entry.Slug, Stream: req.Stream, Payload: targetPayload(req.Payload, entry)}
|
|
next := nextTarget(route, i)
|
|
start := time.Now()
|
|
if req.Stream {
|
|
fail = a.streamChat(w, r, cc, attempt, entry, route, next, key, messages, start)
|
|
} else {
|
|
fail = a.completeChat(w, r, cc, attempt, entry, route, next, key, messages, start)
|
|
}
|
|
if fail == nil {
|
|
return
|
|
}
|
|
last = fail
|
|
}
|
|
last.write(w)
|
|
}
|
|
|
|
// completeChat serves a non-streaming completion from one target. It returns
|
|
// nil once the client has an answer, or the failure when the request should
|
|
// move on to next. An upstream error with no backup behind it is relayed.
|
|
func (a *API) completeChat(w http.ResponseWriter, r *http.Request, cc providers.ChatCompleter, req *providers.ChatRequest, entry *catalog.Entry, route *catalog.Route, next *catalog.Entry, key *store.VirtualKey, messages json.RawMessage, start time.Time) *failure {
|
|
resp, err := cc.ChatCompletion(r.Context(), req)
|
|
if err != nil {
|
|
a.logRequest(r, key, entry, route, http.StatusBadGateway, false, providers.Usage{}, err.Error()+switchNote(next), start, nil, messages, "")
|
|
return &failure{status: http.StatusBadGateway, msg: err.Error()}
|
|
}
|
|
a.Limiter.RecordTokens(entry.Address, resp.Usage.TotalTokens)
|
|
if resp.StatusCode >= 400 {
|
|
errMsg := upstreamErrorMessage(resp.Body)
|
|
if next != nil && retryable(resp.StatusCode) {
|
|
a.logRequest(r, key, entry, route, resp.StatusCode, false, resp.Usage, errMsg+switchNote(next), start, nil, messages, "")
|
|
return &failure{status: resp.StatusCode, body: resp.Body}
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(resp.StatusCode)
|
|
_, _ = w.Write(resp.Body)
|
|
a.logRequest(r, key, entry, route, resp.StatusCode, false, resp.Usage, errMsg, start, nil, messages, "")
|
|
return nil
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(resp.StatusCode)
|
|
_, _ = w.Write(clientBody(resp.Body, time.Now().Unix(), reportedModel(route, entry)))
|
|
|
|
a.logRequest(r, key, entry, route, resp.StatusCode, false, resp.Usage, "", start, nil, messages, extractCompletion(resp.Body))
|
|
return nil
|
|
}
|
|
|
|
// streamChat serves a streaming completion from one target. The first chunk is
|
|
// read before the status line goes out, so an upstream that fails at once can
|
|
// still hand the request to a backup. After the first byte reaches the
|
|
// client, a failure can only be recorded in the log.
|
|
func (a *API) streamChat(w http.ResponseWriter, r *http.Request, cc providers.ChatCompleter, req *providers.ChatRequest, entry *catalog.Entry, route *catalog.Route, next *catalog.Entry, key *store.VirtualKey, messages json.RawMessage, start time.Time) *failure {
|
|
stream, err := cc.ChatCompletionStream(r.Context(), req)
|
|
if err != nil {
|
|
var apiErr *providers.APIError
|
|
if errors.As(err, &apiErr) {
|
|
errMsg := upstreamErrorMessage(apiErr.Body)
|
|
if next != nil && retryable(apiErr.StatusCode) {
|
|
a.logRequest(r, key, entry, route, apiErr.StatusCode, true, providers.Usage{}, errMsg+switchNote(next), start, nil, messages, "")
|
|
return &failure{status: apiErr.StatusCode, body: apiErr.Body}
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(apiErr.StatusCode)
|
|
_, _ = w.Write(apiErr.Body)
|
|
a.logRequest(r, key, entry, route, apiErr.StatusCode, true, providers.Usage{}, errMsg, start, nil, messages, "")
|
|
return nil
|
|
}
|
|
a.logRequest(r, key, entry, route, http.StatusBadGateway, true, providers.Usage{}, err.Error()+switchNote(next), start, nil, messages, "")
|
|
return &failure{status: http.StatusBadGateway, msg: err.Error()}
|
|
}
|
|
defer stream.Close()
|
|
|
|
ev, err := stream.Recv()
|
|
if err != nil && !errors.Is(err, io.EOF) {
|
|
errMsg := "stream interrupted: " + err.Error()
|
|
a.logRequest(r, key, entry, route, http.StatusBadGateway, true, providers.Usage{}, errMsg+switchNote(next), start, nil, messages, "")
|
|
return &failure{status: http.StatusBadGateway, msg: errMsg}
|
|
}
|
|
|
|
flusher, _ := w.(http.Flusher)
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.Header().Set("Cache-Control", "no-cache")
|
|
w.Header().Set("Connection", "keep-alive")
|
|
w.WriteHeader(http.StatusOK)
|
|
|
|
model := reportedModel(route, entry)
|
|
var ttft *int
|
|
var reasoning, content strings.Builder
|
|
created := time.Now().Unix()
|
|
errMsg := ""
|
|
for !errors.Is(err, io.EOF) {
|
|
if err != nil {
|
|
// The status line is already sent, so the error can only be
|
|
// recorded in the log row and signalled by closing the stream.
|
|
if r.Context().Err() != nil {
|
|
errMsg = "client disconnected"
|
|
} else {
|
|
errMsg = "stream interrupted: " + err.Error()
|
|
}
|
|
break
|
|
}
|
|
if ttft == nil {
|
|
ms := int(time.Since(start).Milliseconds())
|
|
ttft = &ms
|
|
}
|
|
// Accumulate reasoning and content deltas separately so the recorded
|
|
// completion wraps reasoning in think tags, matching extractCompletion.
|
|
rs, c := extractDelta(ev.Data)
|
|
reasoning.WriteString(rs)
|
|
content.WriteString(c)
|
|
_, _ = fmt.Fprintf(w, "data: %s\n\n", clientBody(ev.Data, created, model))
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
ev, err = stream.Recv()
|
|
}
|
|
_, _ = fmt.Fprint(w, "data: [DONE]\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
|
|
completion := content.String()
|
|
if reasoning.Len() > 0 {
|
|
completion = "<think>\n" + reasoning.String() + "\n</think>\n\n" + completion
|
|
}
|
|
usage := stream.Usage()
|
|
a.Limiter.RecordTokens(entry.Address, usage.TotalTokens)
|
|
a.logRequest(r, key, entry, route, http.StatusOK, true, usage, errMsg, start, ttft, messages, completion)
|
|
return nil
|
|
}
|
|
|
|
// logRequest records one attempt. ModelName is the address the client sent;
|
|
// when a backup served the request, ServedModel names it and Cost uses its
|
|
// prices, so the switch is visible in the log but not to the client.
|
|
func (a *API) logRequest(r *http.Request, key *store.VirtualKey, entry *catalog.Entry, route *catalog.Route, status int, streamed bool, usage providers.Usage, errMsg string, start time.Time, ttft *int, messages json.RawMessage, completion string) {
|
|
entryID := entry.ID
|
|
providerID := entry.ProviderID
|
|
served := ""
|
|
if route.IsFallback(entry) {
|
|
served = entry.Address
|
|
}
|
|
rec := store.RequestLog{
|
|
ProviderModelID: &entryID,
|
|
ProviderID: &providerID,
|
|
ProviderName: entry.ProviderName,
|
|
ModelName: route.Address,
|
|
ServedModel: served,
|
|
StatusCode: status,
|
|
Streamed: streamed,
|
|
Error: errMsg,
|
|
PromptTokens: usage.PromptTokens,
|
|
CompletionTokens: usage.CompletionTokens,
|
|
CachedTokens: usage.CachedTokens,
|
|
TotalTokens: usage.TotalTokens,
|
|
Cost: billing.Cost(entry, usage),
|
|
LatencyMS: int(time.Since(start).Milliseconds()),
|
|
TTFTMs: ttft,
|
|
RequestID: chimw.GetReqID(r.Context()),
|
|
ClientIP: clientIP(r),
|
|
}
|
|
if key != nil {
|
|
rec.KeyPrefix = key.KeyPrefix
|
|
rec.UserID = key.UserID
|
|
// Only set VirtualKeyID for real keys — synthetic dashboard keys have a nil UUID.
|
|
if key.ID != uuid.Nil {
|
|
rec.VirtualKeyID = &key.ID
|
|
}
|
|
}
|
|
|
|
if a.Stats != nil {
|
|
a.Stats.Record(rec)
|
|
}
|
|
|
|
// The log write must outlive the request, so it gets its own context.
|
|
go func() { //nolint:gosec // G118: intentional, see above
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
id, err := a.Logs.InsertRequestLog(ctx, rec)
|
|
if err != nil {
|
|
slog.Warn("failed to write request log", "err", err)
|
|
return
|
|
}
|
|
if err := a.Logs.InsertLogContent(ctx, id, messages, completion); err != nil {
|
|
slog.Warn("failed to write request log content", "err", err)
|
|
}
|
|
if key != nil && key.ID != uuid.Nil {
|
|
_ = a.Logs.TouchVirtualKey(ctx, key.ID)
|
|
}
|
|
}()
|
|
}
|
|
|
|
// marshalMessages extracts the messages array from the request payload as raw
|
|
// JSON for the content log. Inline image data is replaced by a short marker so
|
|
// a single multimodal request cannot write megabytes of base64 into the log.
|
|
func marshalMessages(v any) json.RawMessage {
|
|
arr, ok := v.([]any)
|
|
if !ok {
|
|
return json.RawMessage("[]")
|
|
}
|
|
b, err := json.Marshal(redactImages(arr))
|
|
if err != nil {
|
|
return json.RawMessage("[]")
|
|
}
|
|
return b
|
|
}
|
|
|
|
// redactImages returns a copy of the messages with data: image URLs replaced.
|
|
func redactImages(messages []any) []any {
|
|
out := make([]any, len(messages))
|
|
for i, raw := range messages {
|
|
m, ok := raw.(map[string]any)
|
|
if !ok {
|
|
out[i] = raw
|
|
continue
|
|
}
|
|
parts, ok := m["content"].([]any)
|
|
if !ok {
|
|
out[i] = m
|
|
continue
|
|
}
|
|
cp := make(map[string]any, len(m))
|
|
maps.Copy(cp, m)
|
|
newParts := make([]any, len(parts))
|
|
for j, p := range parts {
|
|
part, ok := p.(map[string]any)
|
|
if !ok || part["type"] != "image_url" {
|
|
newParts[j] = p
|
|
continue
|
|
}
|
|
iu, _ := part["image_url"].(map[string]any)
|
|
url, _ := iu["url"].(string)
|
|
if !strings.HasPrefix(url, "data:") {
|
|
newParts[j] = p
|
|
continue
|
|
}
|
|
mediaType := url
|
|
if i := strings.IndexAny(url, ";,"); i > 0 {
|
|
mediaType = url[:i]
|
|
}
|
|
newParts[j] = map[string]any{
|
|
"type": "image_url",
|
|
"image_url": map[string]any{"url": fmt.Sprintf("[inline image %s, %d bytes omitted]", strings.TrimPrefix(mediaType, "data:"), len(url))},
|
|
}
|
|
}
|
|
cp["content"] = newParts
|
|
out[i] = cp
|
|
}
|
|
return out
|
|
}
|
|
|
|
// upstreamErrorMessage pulls the message out of an OpenAI-style error body,
|
|
// falling back to a truncated copy of the raw body.
|
|
func upstreamErrorMessage(body []byte) string {
|
|
var e struct {
|
|
Error struct {
|
|
Message string `json:"message"`
|
|
} `json:"error"`
|
|
}
|
|
if json.Unmarshal(body, &e) == nil && e.Error.Message != "" {
|
|
return e.Error.Message
|
|
}
|
|
msg := strings.TrimSpace(string(body))
|
|
if len(msg) > 500 {
|
|
msg = msg[:500] + "…"
|
|
}
|
|
return msg
|
|
}
|
|
|
|
// extractCompletion pulls the first choice's message content from a non-streaming
|
|
// response body. Reasoning (reasoning_content, reasoning, or reasoning_details) is
|
|
// wrapped in think tags and prepended when present.
|
|
func extractCompletion(body []byte) string {
|
|
var resp struct {
|
|
Choices []struct {
|
|
Message struct {
|
|
ReasoningContent string `json:"reasoning_content"`
|
|
Reasoning string `json:"reasoning"`
|
|
ReasoningDetails []reasoningDetail `json:"reasoning_details"`
|
|
Content string `json:"content"`
|
|
} `json:"message"`
|
|
} `json:"choices"`
|
|
}
|
|
if err := json.Unmarshal(body, &resp); err != nil || len(resp.Choices) == 0 {
|
|
return ""
|
|
}
|
|
m := resp.Choices[0].Message
|
|
r := reasoningText(m.ReasoningContent, m.Reasoning, m.ReasoningDetails)
|
|
if r != "" {
|
|
return "<think>\n" + r + "\n</think>\n\n" + m.Content
|
|
}
|
|
return m.Content
|
|
}
|
|
|
|
// extractDelta pulls the reasoning and content deltas from a streaming SSE chunk.
|
|
// reasoning_content, reasoning, and reasoning_details are all recognized.
|
|
func extractDelta(data []byte) (reasoning, content string) {
|
|
var chunk struct {
|
|
Choices []struct {
|
|
Delta struct {
|
|
ReasoningContent string `json:"reasoning_content"`
|
|
Reasoning string `json:"reasoning"`
|
|
ReasoningDetails []reasoningDetail `json:"reasoning_details"`
|
|
Content string `json:"content"`
|
|
} `json:"delta"`
|
|
} `json:"choices"`
|
|
}
|
|
if err := json.Unmarshal(data, &chunk); err != nil || len(chunk.Choices) == 0 {
|
|
return "", ""
|
|
}
|
|
d := chunk.Choices[0].Delta
|
|
return reasoningText(d.ReasoningContent, d.Reasoning, d.ReasoningDetails), d.Content
|
|
}
|
|
|
|
// reasoningDetail is one entry in a reasoning_details array.
|
|
type reasoningDetail struct {
|
|
Type string `json:"type"`
|
|
Text string `json:"text"`
|
|
Summary string `json:"summary"`
|
|
}
|
|
|
|
// reasoningText joins reasoning fragments a provider may emit: reasoning_content
|
|
// (llama.cpp), reasoning (OpenRouter), and reasoning_details (structured blocks).
|
|
func reasoningText(reasoningContent, reasoning string, details []reasoningDetail) string {
|
|
var b strings.Builder
|
|
b.WriteString(reasoningContent)
|
|
b.WriteString(reasoning)
|
|
for _, d := range details {
|
|
b.WriteString(d.Text)
|
|
b.WriteString(d.Summary)
|
|
}
|
|
return b.String()
|
|
}
|
|
|
|
func clientIP(r *http.Request) string {
|
|
host, _, err := net.SplitHostPort(r.RemoteAddr)
|
|
if err != nil {
|
|
return r.RemoteAddr
|
|
}
|
|
return host
|
|
}
|