localmodelroutex/backend/internal/api/route.go

package api

import (
	"encoding/json"
	"maps"
	"net/http"

	"git.govekar.net/LovroG05/localmodelroutex/internal/catalog"
	"git.govekar.net/LovroG05/localmodelroutex/internal/providers"
)

// The proxy handlers walk a route's targets in order. The requested model is
// tried first. When it is rate limited, has no provider, or its upstream
// fails before the client has received anything, the next backup is tried
// with the same request. A backup answers under the requested model's name,
// so the client cannot tell the difference. The request log can: each failed
// attempt gets its own row, and the row for the backup carries both the
// requested address and the address that served it, with the backup's cost.

// failure is one target's refusal. When no target answers, the last failure
// is sent to the client.
type failure struct {
	status int
	body   []byte // upstream error payload to relay; nil writes msg as a gateway error
	msg    string
}

func (f *failure) write(w http.ResponseWriter) {
	if f.body != nil {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(f.status)
		_, _ = w.Write(f.body)
		return
	}
	writeError(w, f.status, f.msg)
}

// retryable reports whether an upstream status means the request should move
// on to a backup. Client errors are not retried: a backup would get the same
// request and most likely give the same answer.
func retryable(status int) bool {
	return status == http.StatusRequestTimeout || status == http.StatusTooManyRequests || status >= http.StatusInternalServerError
}

// admit charges requests against a target's rate limits and returns its
// provider. The failure says why the target cannot serve.
func (a *API) admit(e *catalog.Entry, requests int) (providers.Provider, *failure) {
	if !a.Limiter.AllowRequests(e.Address, ptrInt(e.RPMLimit), requests) {
		return nil, &failure{status: http.StatusTooManyRequests, msg: "rate limit exceeded (RPM)"}
	}
	if !a.Limiter.AllowTokens(e.Address, ptrInt(e.TPMLimit)) {
		return nil, &failure{status: http.StatusTooManyRequests, msg: "rate limit exceeded (TPM)"}
	}
	prov, ok := a.Manager.Get(e.ProviderName)
	if !ok {
		return nil, &failure{status: http.StatusServiceUnavailable, msg: "provider unavailable"}
	}
	return prov, nil
}

// targetPayload copies the client payload for one target and applies the
// target's parameter overrides. Each target gets its own copy, so the
// overrides of a failed target do not leak into the request to its backup.
func targetPayload(payload map[string]any, e *catalog.Entry) map[string]any {
	out := maps.Clone(payload)
	maps.Copy(out, e.ParamOverrides)
	return out
}

// switchNote is appended to the log row of a failed attempt when the request
// moves on to a backup. Empty when there is no backup after the attempt.
func switchNote(next *catalog.Entry) string {
	if next == nil {
		return ""
	}
	return "; switched to backup " + next.Address
}

// nextTarget returns the target after index i, or nil.
func nextTarget(route *catalog.Route, i int) *catalog.Entry {
	if i+1 < len(route.Targets) {
		return route.Targets[i+1]
	}
	return nil
}

// reportedModel is the model id a response from e must carry. A backup
// reports the requested model's id; the requested model keeps what its
// upstream sent.
func reportedModel(route *catalog.Route, e *catalog.Entry) string {
	if route.IsFallback(e) {
		return route.Slug
	}
	return ""
}

// editJSON decodes an object body, applies edit, and encodes it again. A body
// that is not a JSON object is returned unchanged.
func editJSON(body []byte, edit func(map[string]any)) []byte {
	var obj map[string]any
	if json.Unmarshal(body, &obj) != nil || obj == nil {
		return body
	}
	edit(obj)
	out, err := json.Marshal(obj)
	if err != nil {
		return body
	}
	return out
}

// clientBody prepares an upstream chat body or chunk for the client. A missing
// "created" timestamp is added, and when model is set the reported model is
// replaced with it.
func clientBody(body []byte, created int64, model string) []byte {
	return editJSON(body, func(obj map[string]any) {
		if _, ok := obj["created"]; !ok {
			obj["created"] = created
		}
		if model != "" {
			obj["model"] = model
		}
	})
}

// withModel replaces the reported model when model is set.
func withModel(body []byte, model string) []byte {
	if model == "" {
		return body
	}
	return editJSON(body, func(obj map[string]any) { obj["model"] = model })
}