LovroG05
abecc4b888
All checks were successful
ci/woodpecker/push/woodpecker Pipeline was successful
Add an optional exporter that writes one point per proxied request to the InfluxDB 3 write endpoint. Points are batched in memory and flushed by size or interval. The exporter is off until GATEWAY_INFLUX_URL is set.
415 lines
12 KiB
Go
415 lines
12 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
|
|
|
|
entry, err := a.Cache.Resolve(address)
|
|
if err != nil {
|
|
writeError(w, http.StatusNotFound, err.Error())
|
|
return
|
|
}
|
|
if !allowed(key, entry.Address) {
|
|
writeError(w, http.StatusForbidden, fmt.Sprintf("model %q is not allowed for this key", address))
|
|
return
|
|
}
|
|
if !a.Limiter.AllowRequest(entry.Address, ptrInt(entry.RPMLimit)) {
|
|
writeError(w, http.StatusTooManyRequests, "rate limit exceeded (RPM)")
|
|
return
|
|
}
|
|
if !a.Limiter.AllowTokens(entry.Address, ptrInt(entry.TPMLimit)) {
|
|
writeError(w, http.StatusTooManyRequests, "rate limit exceeded (TPM)")
|
|
return
|
|
}
|
|
|
|
prov, ok := a.Manager.Get(entry.ProviderName)
|
|
if !ok {
|
|
writeError(w, http.StatusServiceUnavailable, "provider unavailable")
|
|
return
|
|
}
|
|
cc, ok := prov.(providers.ChatCompleter)
|
|
if !ok {
|
|
writeError(w, http.StatusNotImplemented, "provider does not support chat completions")
|
|
return
|
|
}
|
|
|
|
// Snapshot the messages before we overwrite req.Model below.
|
|
messages := marshalMessages(req.Payload["messages"])
|
|
|
|
req.Model = entry.Slug
|
|
applyOverrides(req.Payload, entry.ParamOverrides)
|
|
|
|
start := time.Now()
|
|
if req.Stream {
|
|
a.streamChat(w, r, cc, req, entry, key, address, messages, start)
|
|
return
|
|
}
|
|
|
|
resp, err := cc.ChatCompletion(r.Context(), req)
|
|
if err != nil {
|
|
writeError(w, http.StatusBadGateway, err.Error())
|
|
a.logRequest(r, key, entry, address, http.StatusBadGateway, false, providers.Usage{}, err.Error(), start, nil, messages, "")
|
|
return
|
|
}
|
|
a.Limiter.RecordTokens(entry.Address, resp.Usage.TotalTokens)
|
|
respBody := resp.Body
|
|
errMsg := ""
|
|
if resp.StatusCode >= 400 {
|
|
errMsg = upstreamErrorMessage(resp.Body)
|
|
} else {
|
|
respBody = injectCreated(resp.Body)
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(resp.StatusCode)
|
|
_, _ = w.Write(respBody)
|
|
|
|
completion := extractCompletion(resp.Body)
|
|
a.logRequest(r, key, entry, address, resp.StatusCode, false, resp.Usage, errMsg, start, nil, messages, completion)
|
|
}
|
|
|
|
func (a *API) streamChat(w http.ResponseWriter, r *http.Request, cc providers.ChatCompleter, req *providers.ChatRequest, entry *catalog.Entry, key *store.VirtualKey, address string, messages json.RawMessage, start time.Time) {
|
|
stream, err := cc.ChatCompletionStream(r.Context(), req)
|
|
if err != nil {
|
|
var apiErr *providers.APIError
|
|
if errors.As(err, &apiErr) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(apiErr.StatusCode)
|
|
_, _ = w.Write(apiErr.Body)
|
|
a.logRequest(r, key, entry, address, apiErr.StatusCode, true, providers.Usage{}, upstreamErrorMessage(apiErr.Body), start, nil, messages, "")
|
|
return
|
|
}
|
|
writeError(w, http.StatusBadGateway, err.Error())
|
|
a.logRequest(r, key, entry, address, http.StatusBadGateway, true, providers.Usage{}, err.Error(), start, nil, messages, "")
|
|
return
|
|
}
|
|
defer stream.Close()
|
|
|
|
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)
|
|
|
|
var ttft *int
|
|
var reasoning, content strings.Builder
|
|
created := time.Now().Unix()
|
|
errMsg := ""
|
|
for {
|
|
ev, err := stream.Recv()
|
|
if errors.Is(err, io.EOF) {
|
|
break
|
|
}
|
|
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.
|
|
r, c := extractDelta(ev.Data)
|
|
reasoning.WriteString(r)
|
|
content.WriteString(c)
|
|
_, _ = fmt.Fprintf(w, "data: %s\n\n", injectCreatedChunk(ev.Data, created))
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
}
|
|
_, _ = 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, address, http.StatusOK, true, usage, errMsg, start, ttft, messages, completion)
|
|
}
|
|
|
|
func applyOverrides(payload, overrides map[string]any) {
|
|
maps.Copy(payload, overrides)
|
|
}
|
|
|
|
func (a *API) logRequest(r *http.Request, key *store.VirtualKey, entry *catalog.Entry, address string, 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
|
|
rec := store.RequestLog{
|
|
ProviderModelID: &entryID,
|
|
ProviderID: &providerID,
|
|
ProviderName: entry.ProviderName,
|
|
ModelName: address,
|
|
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
|
|
}
|
|
|
|
// injectCreated ensures a non-streaming response has a "created" timestamp.
|
|
// If the upstream already provided one, it is preserved.
|
|
func injectCreated(body []byte) []byte {
|
|
var obj map[string]any
|
|
if json.Unmarshal(body, &obj) == nil {
|
|
if _, ok := obj["created"]; ok {
|
|
return body // upstream already set it
|
|
}
|
|
obj["created"] = time.Now().Unix()
|
|
if out, err := json.Marshal(obj); err == nil {
|
|
return out
|
|
}
|
|
}
|
|
return body // fall back to raw body on any error
|
|
}
|
|
|
|
// injectCreatedChunk ensures a streaming SSE chunk has a "created" timestamp.
|
|
func injectCreatedChunk(data []byte, created int64) []byte {
|
|
var obj map[string]any
|
|
if json.Unmarshal(data, &obj) == nil {
|
|
if _, ok := obj["created"]; ok {
|
|
return data // upstream already set it
|
|
}
|
|
obj["created"] = created
|
|
if out, err := json.Marshal(obj); err == nil {
|
|
return out
|
|
}
|
|
}
|
|
return data // fall back to raw chunk on any error
|
|
}
|