Luis Pater
e3e60f914b
- Added `disable-image-generation` configuration flag to disable the `image_generation` tool globally. - Updated payload handling to remove `image_generation` tools from request payload arrays when the flag is enabled. - Modified OpenAI image handlers (`ImagesGenerations`, `ImagesEdits`) to return 404 when the feature is disabled. - Enhanced configuration diff logging to track changes for the `disable-image-generation` flag. - Added accompanying unit tests for the new feature in payload helpers and image handler logic.
122 lines
4.3 KiB
Go
122 lines
4.3 KiB
Go
package openai
|
|
|
|
import (
|
|
"bytes"
|
|
"io"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
|
|
sdkconfig "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
func performImagesEndpointRequest(t *testing.T, endpointPath string, contentType string, body io.Reader, handler gin.HandlerFunc) *httptest.ResponseRecorder {
|
|
t.Helper()
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
router := gin.New()
|
|
router.POST(endpointPath, handler)
|
|
|
|
req := httptest.NewRequest(http.MethodPost, endpointPath, body)
|
|
if contentType != "" {
|
|
req.Header.Set("Content-Type", contentType)
|
|
}
|
|
resp := httptest.NewRecorder()
|
|
router.ServeHTTP(resp, req)
|
|
return resp
|
|
}
|
|
|
|
func assertUnsupportedImagesModelResponse(t *testing.T, resp *httptest.ResponseRecorder, model string) {
|
|
t.Helper()
|
|
|
|
if resp.Code != http.StatusBadRequest {
|
|
t.Fatalf("status = %d, want %d: %s", resp.Code, http.StatusBadRequest, resp.Body.String())
|
|
}
|
|
|
|
message := gjson.GetBytes(resp.Body.Bytes(), "error.message").String()
|
|
expectedMessage := "Model " + model + " is not supported on " + imagesGenerationsPath + " or " + imagesEditsPath + ". Use " + defaultImagesToolModel + "."
|
|
if message != expectedMessage {
|
|
t.Fatalf("error message = %q, want %q", message, expectedMessage)
|
|
}
|
|
if errorType := gjson.GetBytes(resp.Body.Bytes(), "error.type").String(); errorType != "invalid_request_error" {
|
|
t.Fatalf("error type = %q, want invalid_request_error", errorType)
|
|
}
|
|
}
|
|
|
|
func TestImagesModelValidationAllowsGPTImage2WithOptionalPrefix(t *testing.T) {
|
|
for _, model := range []string{"gpt-image-2", "codex/gpt-image-2"} {
|
|
if !isSupportedImagesModel(model) {
|
|
t.Fatalf("expected %s to be supported", model)
|
|
}
|
|
}
|
|
if isSupportedImagesModel("gpt-5.4-mini") {
|
|
t.Fatal("expected gpt-5.4-mini to be rejected")
|
|
}
|
|
}
|
|
|
|
func TestImagesGenerationsRejectsUnsupportedModel(t *testing.T) {
|
|
handler := &OpenAIAPIHandler{}
|
|
body := strings.NewReader(`{"model":"gpt-5.4-mini","prompt":"draw a square"}`)
|
|
|
|
resp := performImagesEndpointRequest(t, imagesGenerationsPath, "application/json", body, handler.ImagesGenerations)
|
|
|
|
assertUnsupportedImagesModelResponse(t, resp, "gpt-5.4-mini")
|
|
}
|
|
|
|
func TestImagesEditsJSONRejectsUnsupportedModel(t *testing.T) {
|
|
handler := &OpenAIAPIHandler{}
|
|
body := strings.NewReader(`{"model":"gpt-5.4-mini","prompt":"edit this","images":[{"image_url":"data:image/png;base64,AA=="}]}`)
|
|
|
|
resp := performImagesEndpointRequest(t, imagesEditsPath, "application/json", body, handler.ImagesEdits)
|
|
|
|
assertUnsupportedImagesModelResponse(t, resp, "gpt-5.4-mini")
|
|
}
|
|
|
|
func TestImagesEditsMultipartRejectsUnsupportedModel(t *testing.T) {
|
|
handler := &OpenAIAPIHandler{}
|
|
var body bytes.Buffer
|
|
writer := multipart.NewWriter(&body)
|
|
if err := writer.WriteField("model", "gpt-5.4-mini"); err != nil {
|
|
t.Fatalf("write model field: %v", err)
|
|
}
|
|
if err := writer.WriteField("prompt", "edit this"); err != nil {
|
|
t.Fatalf("write prompt field: %v", err)
|
|
}
|
|
if errClose := writer.Close(); errClose != nil {
|
|
t.Fatalf("close multipart writer: %v", errClose)
|
|
}
|
|
|
|
resp := performImagesEndpointRequest(t, imagesEditsPath, writer.FormDataContentType(), &body, handler.ImagesEdits)
|
|
|
|
assertUnsupportedImagesModelResponse(t, resp, "gpt-5.4-mini")
|
|
}
|
|
|
|
func TestImagesGenerations_DisableImageGeneration_Returns404(t *testing.T) {
|
|
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{DisableImageGeneration: true}, nil)
|
|
handler := NewOpenAIAPIHandler(base)
|
|
body := strings.NewReader(`{"prompt":"draw a square"}`)
|
|
|
|
resp := performImagesEndpointRequest(t, imagesGenerationsPath, "application/json", body, handler.ImagesGenerations)
|
|
|
|
if resp.Code != http.StatusNotFound {
|
|
t.Fatalf("status = %d, want %d: %s", resp.Code, http.StatusNotFound, resp.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestImagesEdits_DisableImageGeneration_Returns404(t *testing.T) {
|
|
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{DisableImageGeneration: true}, nil)
|
|
handler := NewOpenAIAPIHandler(base)
|
|
body := strings.NewReader(`{"prompt":"edit this","images":[{"image_url":"data:image/png;base64,AA=="}]}`)
|
|
|
|
resp := performImagesEndpointRequest(t, imagesEditsPath, "application/json", body, handler.ImagesEdits)
|
|
|
|
if resp.Code != http.StatusNotFound {
|
|
t.Fatalf("status = %d, want %d: %s", resp.Code, http.StatusNotFound, resp.Body.String())
|
|
}
|
|
}
|