localmodelroutex/backend/internal/providers/openrouter/images.go

package openrouter

import (
	"context"
	"encoding/json"
	"fmt"
	"maps"
	"net/http"
	"strings"
	"time"

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

// maxImagesPerRequest caps the OpenAI "n" parameter. Each image is one
// upstream chat completion, so this bounds the fan-out per client call.
const maxImagesPerRequest = 10

// imageOnlyFields are OpenAI images-API fields that have no meaning on the
// chat completions endpoint. They are dropped from the upstream body; the
// ones that matter (prompt, size, quality) are translated instead.
var imageOnlyFields = []string{
	"prompt", "n", "size", "quality", "response_format", "style",
	"background", "output_format", "output_compression", "moderation",
	"partial_images", "stream",
}

// ImageGeneration serves the OpenAI images API on top of OpenRouter, which
// has no /images/generations endpoint. Image models on OpenRouter generate
// through chat completions with modalities ["image","text"] and return the
// result in message.images. This translates the call there and back, the
// same way the LiteLLM proxy does, so OpenAI-style image clients (such as
// Open WebUI's image generation setting) work unchanged.
func (p *provider) ImageGeneration(ctx context.Context, req *providers.ImageRequest) (*providers.ImageResponse, error) {
	if strings.TrimSpace(req.Prompt) == "" {
		return errorResponse(http.StatusBadRequest, "invalid_request_error", "missing 'prompt'"), nil
	}
	n := req.N
	if n <= 0 {
		n = 1
	}
	if n > maxImagesPerRequest {
		return errorResponse(http.StatusBadRequest, "invalid_request_error",
			fmt.Sprintf("'n' must be at most %d", maxImagesPerRequest)), nil
	}

	chatReq := &providers.ChatRequest{Model: req.Model, Payload: chatPayload(req)}

	var (
		data     []map[string]any
		usage    providers.Usage
		costSeen bool
		cost     float64
	)
	for i := 0; i < n; i++ {
		resp, err := p.ChatCompletion(ctx, chatReq)
		if err != nil {
			return nil, err
		}
		if resp.StatusCode >= http.StatusBadRequest {
			// Relay the upstream error as is; it is already OpenAI-shaped.
			return &providers.ImageResponse{
				StatusCode: resp.StatusCode,
				Header:     resp.Header,
				Body:       resp.Body,
				Usage:      resp.Usage,
			}, nil
		}
		usage.PromptTokens += resp.Usage.PromptTokens
		usage.CompletionTokens += resp.Usage.CompletionTokens
		usage.CachedTokens += resp.Usage.CachedTokens
		usage.TotalTokens += resp.Usage.TotalTokens
		if resp.Usage.Cost != nil {
			costSeen = true
			cost += *resp.Usage.Cost
		}
		data = append(data, extractImages(resp.Body)...)
	}
	if costSeen {
		usage.Cost = &cost
	}
	if len(data) == 0 {
		return errorResponse(http.StatusBadGateway, "server_error",
			"model returned no image; it may not support image output"), nil
	}

	out := map[string]any{
		"created": time.Now().Unix(),
		"data":    data,
		"usage":   imageUsage(usage),
	}
	body, err := json.Marshal(out)
	if err != nil {
		return nil, err
	}
	return &providers.ImageResponse{
		StatusCode: http.StatusOK,
		Header:     http.Header{"Content-Type": []string{"application/json"}},
		Body:       body,
		Usage:      usage,
	}, nil
}

// chatPayload builds the upstream chat completion body. Fields the client or
// an admin override put in the payload pass through, so image_config set
// explicitly wins over the values mapped from size and quality.
func chatPayload(req *providers.ImageRequest) map[string]any {
	out := make(map[string]any, len(req.Payload)+4)
	maps.Copy(out, req.Payload)
	for _, k := range imageOnlyFields {
		delete(out, k)
	}
	out["messages"] = []any{map[string]any{"role": "user", "content": req.Prompt}}
	if _, ok := out["modalities"]; !ok {
		out["modalities"] = []any{"image", "text"}
	}

	cfg, _ := out["image_config"].(map[string]any)
	if cfg == nil {
		cfg = map[string]any{}
	}
	if _, ok := cfg["aspect_ratio"]; !ok && req.Size != "" {
		if ar := aspectRatio(req.Size); ar != "" {
			cfg["aspect_ratio"] = ar
		}
	}
	if _, ok := cfg["image_size"]; !ok && req.Quality != "" {
		if sz := imageSize(req.Quality); sz != "" {
			cfg["image_size"] = sz
		}
	}
	if len(cfg) > 0 {
		out["image_config"] = cfg
	}
	return out
}

// aspectRatio maps an OpenAI "WxH" size to the nearest OpenRouter
// image_config.aspect_ratio. Unknown or "auto" sizes map to nothing so the
// model default applies.
func aspectRatio(size string) string {
	switch size {
	case "256x256", "512x512", "1024x1024":
		return "1:1"
	case "1536x1024":
		return "3:2"
	case "1792x1024":
		return "16:9"
	case "1024x1536":
		return "2:3"
	case "1024x1792":
		return "9:16"
	}
	w, h, ok := strings.Cut(size, "x")
	if !ok {
		return ""
	}
	var wi, hi int
	if _, err := fmt.Sscanf(w+" "+h, "%d %d", &wi, &hi); err != nil || wi <= 0 || hi <= 0 {
		return ""
	}
	// Thresholds sit halfway between neighbouring supported ratios
	// (16:9, 3:2, 4:3, 1:1, 3:4, 2:3, 9:16).
	r := float64(wi) / float64(hi)
	switch {
	case r > 1.64:
		return "16:9"
	case r > 1.42:
		return "3:2"
	case r > 1.17:
		return "4:3"
	case r > 0.875:
		return "1:1"
	case r > 0.708:
		return "3:4"
	case r > 0.615:
		return "2:3"
	default:
		return "9:16"
	}
}

// imageSize maps an OpenAI quality to OpenRouter image_config.image_size.
// Only Gemini models honor it; others ignore the field.
func imageSize(quality string) string {
	switch quality {
	case "low", "standard", "auto":
		return "1K"
	case "medium":
		return "2K"
	case "high", "hd":
		return "4K"
	}
	return ""
}

// extractImages pulls message.images out of a chat completion body and
// renders each as an OpenAI images "data" item. Data URLs become b64_json;
// anything else is relayed as url. The text content, if any, is attached as
// revised_prompt so it is not lost.
func extractImages(body []byte) []map[string]any {
	var resp struct {
		Choices []struct {
			Message struct {
				Content any `json:"content"`
				Images  []struct {
					Type     string `json:"type"`
					ImageURL struct {
						URL string `json:"url"`
					} `json:"image_url"`
				} `json:"images"`
			} `json:"message"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil
	}
	var out []map[string]any
	for _, ch := range resp.Choices {
		text, _ := ch.Message.Content.(string)
		for _, img := range ch.Message.Images {
			u := img.ImageURL.URL
			if u == "" {
				continue
			}
			item := map[string]any{}
			if strings.HasPrefix(u, "data:") {
				_, b64, ok := strings.Cut(u, ",")
				if !ok || b64 == "" {
					continue
				}
				item["b64_json"] = b64
			} else {
				item["url"] = u
			}
			if text != "" {
				item["revised_prompt"] = text
			}
			out = append(out, item)
		}
	}
	return out
}

// imageUsage renders usage in the OpenAI images schema, plus OpenRouter's
// cost when reported.
func imageUsage(u providers.Usage) map[string]any {
	m := map[string]any{
		"input_tokens":  u.PromptTokens,
		"output_tokens": u.CompletionTokens,
		"total_tokens":  u.TotalTokens,
	}
	if u.Cost != nil {
		m["cost"] = *u.Cost
	}
	return m
}

// errorResponse builds an OpenAI-shaped error body.
func errorResponse(status int, typ, msg string) *providers.ImageResponse {
	body, _ := json.Marshal(map[string]any{
		"error": map[string]any{"message": msg, "type": typ},
	})
	return &providers.ImageResponse{
		StatusCode: status,
		Header:     http.Header{"Content-Type": []string{"application/json"}},
		Body:       body,
	}
}