llmgw/internal/control/adminapi/requestbody.go

package adminapi

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

type bodyFinder interface {
	RequestBodyByID(ctx context.Context, id string) (store.RequestBody, error)
}

// Replayer sends a kept request again to a model: the router's Replay.
type Replayer func(ctx context.Context, path, model string, body []byte) (int, []byte, time.Duration, error)

// SetReplayer gives the API the router to replay kept requests through.
func (a *API) SetReplayer(f Replayer) { a.replay = f }

// ReplayView is a kept request sent again, beside what it got the first time.
type ReplayView struct {
	Original store.RequestBody `json:"original"`
	Model    string            `json:"model"`
	Status   int               `json:"status"`
	Body     json.RawMessage   `json:"body"`
	TookMs   int64             `json:"took_ms"`
}

func (a *API) keptBody(w http.ResponseWriter, r *http.Request) (store.RequestBody, bool) {
	f, ok := a.usage.(bodyFinder)
	if !ok {
		writeError(w, http.StatusNotImplemented, "this control plane keeps no requests")
		return store.RequestBody{}, false
	}
	b, err := f.RequestBodyByID(r.Context(), r.PathValue("id"))
	if errors.Is(err, store.ErrNotFound) {
		writeError(w, http.StatusNotFound, "no request with that id is kept: its project keeps none "+
			"(llmgw keys requests --project P --retain-hours N), it fell outside the share kept, or it has expired")
		return store.RequestBody{}, false
	}
	if err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return store.RequestBody{}, false
	}
	return b, true
}

func (a *API) requestBody(w http.ResponseWriter, r *http.Request) {
	if b, ok := a.keptBody(w, r); ok {
		writeJSON(w, http.StatusOK, b)
	}
}

func (a *API) replayRequest(w http.ResponseWriter, r *http.Request) {
	if a.replay == nil {
		writeError(w, http.StatusNotImplemented, "this control plane does not route requests")
		return
	}
	var req struct {
		Model string `json:"model"`
	}
	if err := json.NewDecoder(io.LimitReader(r.Body, 1<<12)).Decode(&req); err != nil || req.Model == "" {
		writeError(w, http.StatusBadRequest, "say which model: {\"model\": \"...\"}")
		return
	}
	b, ok := a.keptBody(w, r)
	if !ok {
		return
	}
	status, body, took, err := a.replay(r.Context(), b.Path, req.Model, []byte(b.Request))
	if err != nil {
		writeError(w, http.StatusBadGateway, err.Error())
		return
	}
	v := ReplayView{Original: b, Model: req.Model, Status: status, TookMs: took.Milliseconds()}
	if json.Valid(body) {
		v.Body = body
	} else {
		v.Body, _ = json.Marshal(string(body))
	}
	writeJSON(w, http.StatusOK, v)
}