llmgw/internal/control/adminapi/federation.go

package adminapi

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"strings"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/federation"
)

// Federation: other llmgw clusters whose models this one offers, as
// PEER/MODEL, and the fallbacks that send overflow to them.

// SetFederation gives the API the peers to list and change.
func (a *API) SetFederation(f *federation.Service) { a.federation = f }

// FederationView is this cluster's id and its peers.
type FederationView struct {
	// Cluster is this cluster's own id, which its peers see in the hop
	// header.
	Cluster string              `json:"cluster"`
	Peers   []federation.Status `json:"peers"`
}

// FallbacksRequest replaces a model's fallbacks; empty takes them all away.
type FallbacksRequest struct {
	Fallbacks []string `json:"fallbacks"`
}

// FallbacksView is a model's fallbacks, as they are now.
type FallbacksView struct {
	Model     string   `json:"model"`
	Fallbacks []string `json:"fallbacks"`
}

type fallbackSetter interface {
	SetFallbacks(ctx context.Context, name string, fallbacks []string, by string) (catalog.Model, error)
}

func (a *API) registerFederation(mux Mux) {
	need := func(w http.ResponseWriter) bool {
		if a.federation == nil {
			writeError(w, http.StatusNotImplemented, "this control plane has no store to keep federation peers in")
			return false
		}
		return true
	}
	mux.HandleFunc("GET /api/v1/federation", func(w http.ResponseWriter, _ *http.Request) {
		if !need(w) {
			return
		}
		writeJSON(w, http.StatusOK, FederationView{Cluster: a.federation.Self(), Peers: a.federation.List()})
	})
	mux.HandleFunc("POST /api/v1/federation", func(w http.ResponseWriter, r *http.Request) {
		if !need(w) {
			return
		}
		var req federation.AddRequest
		if err := json.NewDecoder(io.LimitReader(r.Body, 1<<14)).Decode(&req); err != nil {
			writeError(w, http.StatusBadRequest, "the request is not valid JSON: "+err.Error())
			return
		}
		st, err := a.federation.Add(r.Context(), req, actorOf(r))
		if err != nil {
			writeError(w, http.StatusBadRequest, err.Error())
			return
		}
		writeJSON(w, http.StatusCreated, st)
	})
	mux.HandleFunc("DELETE /api/v1/federation/{peer}", func(w http.ResponseWriter, r *http.Request) {
		if !need(w) {
			return
		}
		name := r.PathValue("peer")
		if err := a.federation.Remove(r.Context(), name); err != nil {
			writeError(w, http.StatusBadRequest, err.Error())
			return
		}
		writeJSON(w, http.StatusOK, Removed{Removed: name})
	})
	mux.HandleFunc("POST /api/v1/federation/{peer}/sync", func(w http.ResponseWriter, r *http.Request) {
		if !need(w) {
			return
		}
		st, err := a.federation.Sync(r.Context(), r.PathValue("peer"))
		if err != nil {
			writeError(w, http.StatusNotFound, err.Error())
			return
		}
		writeJSON(w, http.StatusOK, st)
	})
	mux.HandleFunc("PUT /api/v1/models/{name}/fallbacks", func(w http.ResponseWriter, r *http.Request) {
		setter, ok := a.models.(fallbackSetter)
		if !ok {
			writeError(w, http.StatusNotImplemented, "this control plane has no store to keep the setting in")
			return
		}
		var req FallbacksRequest
		if err := json.NewDecoder(io.LimitReader(r.Body, 1<<14)).Decode(&req); err != nil {
			writeError(w, http.StatusBadRequest, "the request is not valid JSON: "+err.Error())
			return
		}
		var fbs []string
		for _, f := range req.Fallbacks {
			if f = strings.TrimSpace(f); f != "" {
				fbs = append(fbs, f)
			}
		}
		m, err := setter.SetFallbacks(r.Context(), r.PathValue("name"), fbs, actorOf(r))
		if err != nil {
			writeError(w, http.StatusBadRequest, err.Error())
			return
		}
		out := m.Fallbacks
		if out == nil {
			out = []string{}
		}
		writeJSON(w, http.StatusOK, FallbacksView{Model: m.Name, Fallbacks: out})
	})
}