From 2e3f90705fcffdebfd95fa0ee6b9694eebea00f4 Mon Sep 17 00:00:00 2001
From: Ettore Di Giacinto <mudler@localai.io>
Date: Wed, 22 Jan 2025 10:34:43 +0100
Subject: [PATCH] chore(tests): use quality to specify step

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
---
 core/http/endpoints/openai/request.go | 9 +++++++++
 core/schema/openai.go                 | 5 +++--
 tests/e2e-aio/e2e_test.go             | 7 +++++--
 3 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/core/http/endpoints/openai/request.go b/core/http/endpoints/openai/request.go
index 2451f15f..4eaeec24 100644
--- a/core/http/endpoints/openai/request.go
+++ b/core/http/endpoints/openai/request.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"strconv"
 
 	"github.com/gofiber/fiber/v2"
 	"github.com/google/uuid"
@@ -296,6 +297,14 @@ func updateRequestConfig(config *config.BackendConfig, input *schema.OpenAIReque
 			}
 		}
 	}
+
+	// If a quality was defined as number, convert it to step
+	if input.Quality != "" {
+		q, err := strconv.Atoi(input.Quality)
+		if err == nil {
+			config.Step = q
+		}
+	}
 }
 
 func mergeRequestWithConfig(modelFile string, input *schema.OpenAIRequest, cm *config.BackendConfigLoader, loader *model.ModelLoader, debug bool, threads, ctx int, f16 bool) (*config.BackendConfig, *schema.OpenAIRequest, error) {
diff --git a/core/schema/openai.go b/core/schema/openai.go
index b06120ae..e445bee1 100644
--- a/core/schema/openai.go
+++ b/core/schema/openai.go
@@ -191,8 +191,9 @@ type OpenAIRequest struct {
 	Stream bool `json:"stream"`
 
 	// Image (not supported by OpenAI)
-	Mode int `json:"mode"`
-	Step int `json:"step"`
+	Mode    int    `json:"mode"`
+	Quality string `json:"quality"`
+	Step    int    `json:"step"`
 
 	// A grammar to constrain the LLM output
 	Grammar string `json:"grammar" yaml:"grammar"`
diff --git a/tests/e2e-aio/e2e_test.go b/tests/e2e-aio/e2e_test.go
index 3836f351..4d9eb4d8 100644
--- a/tests/e2e-aio/e2e_test.go
+++ b/tests/e2e-aio/e2e_test.go
@@ -123,8 +123,9 @@ var _ = Describe("E2E test", func() {
 			It("correctly", func() {
 				resp, err := client.CreateImage(context.TODO(),
 					openai.ImageRequest{
-						Prompt: "test",
-						Size:   openai.CreateImageSize256x256,
+						Prompt:  "test",
+						Quality: "1",
+						Size:    openai.CreateImageSize256x256,
 					},
 				)
 				Expect(err).ToNot(HaveOccurred())
@@ -136,6 +137,7 @@ var _ = Describe("E2E test", func() {
 					openai.ImageRequest{
 						Prompt:         "test",
 						Size:           openai.CreateImageSize256x256,
+						Quality:        "1",
 						ResponseFormat: openai.CreateImageResponseFormatURL,
 					},
 				)
@@ -148,6 +150,7 @@ var _ = Describe("E2E test", func() {
 					openai.ImageRequest{
 						Prompt:         "test",
 						Size:           openai.CreateImageSize256x256,
+						Quality:        "1",
 						ResponseFormat: openai.CreateImageResponseFormatB64JSON,
 					},
 				)