localmodelroutex/backend/internal/api/images.go

package api

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"slices"
	"time"

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

// maxImagesPerRequest caps the OpenAI "n" parameter per client call. Each
// image may be a separate upstream request, so this bounds the fan-out.
const maxImagesPerRequest = 10

// imageGenerations serves POST /v1/images/generations in OpenAI schema. The
// model must carry the "images" capability in the catalog, which the sync
// sets from the provider's reported output modalities. The provider does any
// translation its upstream needs (OpenRouter, for example, has no images
// endpoint and generates through chat completions).
func (a *API) imageGenerations(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.ParseImageRequest(body)
	if err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	if req.Model == "" {
		writeError(w, http.StatusBadRequest, "missing 'model'")
		return
	}
	if req.Prompt == "" {
		writeError(w, http.StatusBadRequest, "missing 'prompt'")
		return
	}
	if req.N > maxImagesPerRequest {
		writeError(w, http.StatusBadRequest, fmt.Sprintf("'n' must be at most %d", maxImagesPerRequest))
		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
	}
	// Log the prompt as a one-message conversation so it shows in the request
	// log like a chat turn.
	messages, _ := json.Marshal([]map[string]any{{"role": "user", "content": req.Prompt}})

	var last *failure
	for i, entry := range route.Targets {
		if i > 0 && r.Context().Err() != nil {
			break
		}
		if !slices.Contains(entry.Capabilities, string(providers.CapabilityImages)) {
			last = &failure{status: http.StatusBadRequest, msg: fmt.Sprintf("model %q does not generate images", address)}
			continue
		}
		// Each requested image is one upstream call, so charge all of them
		// against the RPM limit up front.
		prov, fail := a.admit(entry, max(req.N, 1))
		if fail != nil {
			last = fail
			continue
		}
		gen, ok := prov.(providers.ImageGenerator)
		if !ok {
			last = &failure{status: http.StatusNotImplemented, msg: "provider does not support image generation"}
			continue
		}
		attempt := *req
		attempt.Model = entry.Slug
		attempt.Payload = targetPayload(req.Payload, entry)
		next := nextTarget(route, i)
		start := time.Now()
		resp, err := gen.ImageGeneration(r.Context(), &attempt)
		if err != nil {
			a.logRequest(r, key, entry, route, http.StatusBadGateway, false, providers.Usage{}, err.Error()+switchNote(next), start, nil, messages, "")
			last = &failure{status: http.StatusBadGateway, msg: err.Error()}
			continue
		}
		a.Limiter.RecordTokens(entry.Address, resp.Usage.TotalTokens)

		errMsg := ""
		completion := ""
		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, "")
				last = &failure{status: resp.StatusCode, body: resp.Body}
				continue
			}
		} else {
			completion = imageCompletionSummary(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, completion)
		return
	}
	last.write(w)
}

// imageCompletionSummary describes an images response for the content log
// without storing the image bytes.
func imageCompletionSummary(body []byte) string {
	var resp struct {
		Data []struct {
			B64           string `json:"b64_json"`
			URL           string `json:"url"`
			RevisedPrompt string `json:"revised_prompt"`
		} `json:"data"`
	}
	if err := json.Unmarshal(body, &resp); err != nil || len(resp.Data) == 0 {
		return ""
	}
	var bytes int
	for _, d := range resp.Data {
		bytes += len(d.B64) / 4 * 3
	}
	s := fmt.Sprintf("[%d image(s) generated, %d bytes omitted]", len(resp.Data), bytes)
	if rp := resp.Data[0].RevisedPrompt; rp != "" {
		s += "\n\n" + rp
	}
	return s
}