feat: add sub2api compatibility

handler/ai.go

@@ -9,8 +9,10 @@ import (
 	"mime"
 	"mime/multipart"
 	"net/http"
+	"net/textproto"
 	"strings"
 
+	"github.com/basketikun/infinite-canvas/model"
 	"github.com/basketikun/infinite-canvas/service"
 )
 
@@ -84,6 +86,12 @@ func proxyAIRequest(w http.ResponseWriter, r *http.Request, path string) {
 		Fail(w, "AI 接口请求失败")
 		return
 	}
+	body, contentType, err = adaptAIRequestForChannel(body, contentType, path, channel)
+	if err != nil {
+		log.Printf("AI proxy adapt request failed: model=%s path=%s err=%v", modelName, path, err)
+		Fail(w, "AI 接口请求失败")
+		return
+	}
 	path = resolveAIProxyPath(channel.BaseURL, modelName, path)
 	request, err := http.NewRequest(http.MethodPost, service.BuildModelChannelURL(channel, path), bytes.NewReader(body))
 	if err != nil {
@@ -140,6 +148,90 @@ func copyAIResponse(w http.ResponseWriter, request *http.Request, onFailure func
 	_, _ = io.Copy(w, response.Body)
 }
 
+func adaptAIRequestForChannel(body []byte, contentType string, path string, channel model.ModelChannel) ([]byte, string, error) {
+	if path != "/images/generations" && path != "/images/edits" {
+		return body, contentType, nil
+	}
+	format := strings.ToLower(strings.TrimSpace(channel.RequestOptions.ImageResponseFormat))
+	if format == "" {
+		if strings.EqualFold(strings.TrimSpace(channel.Compatibility), "sub2api") {
+			format = "url"
+		} else {
+			format = "b64_json"
+		}
+	}
+	if format != "url" && format != "b64_json" {
+		return body, contentType, nil
+	}
+	if strings.HasPrefix(contentType, "multipart/form-data") {
+		return setMultipartField(body, contentType, "response_format", format)
+	}
+	var payload map[string]any
+	if err := json.Unmarshal(body, &payload); err != nil {
+		return nil, "", err
+	}
+	payload["response_format"] = format
+	nextBody, err := json.Marshal(payload)
+	if err != nil {
+		return nil, "", err
+	}
+	return nextBody, contentType, nil
+}
+
+func setMultipartField(body []byte, contentType string, fieldName string, value string) ([]byte, string, error) {
+	_, params, err := mime.ParseMediaType(contentType)
+	if err != nil {
+		return nil, "", err
+	}
+	reader := multipart.NewReader(bytes.NewReader(body), params["boundary"])
+	var buffer bytes.Buffer
+	writer := multipart.NewWriter(&buffer)
+	written := false
+	for {
+		part, err := reader.NextPart()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return nil, "", err
+		}
+		if part.FormName() == fieldName && part.FileName() == "" {
+			if err := writer.WriteField(fieldName, value); err != nil {
+				_ = part.Close()
+				return nil, "", err
+			}
+			written = true
+			_ = part.Close()
+			continue
+		}
+		header := make(textproto.MIMEHeader, len(part.Header))
+		for key, values := range part.Header {
+			copied := make([]string, len(values))
+			copy(copied, values)
+			header[key] = copied
+		}
+		target, err := writer.CreatePart(header)
+		if err != nil {
+			_ = part.Close()
+			return nil, "", err
+		}
+		if _, err := io.Copy(target, part); err != nil {
+			_ = part.Close()
+			return nil, "", err
+		}
+		_ = part.Close()
+	}
+	if !written {
+		if err := writer.WriteField(fieldName, value); err != nil {
+			return nil, "", err
+		}
+	}
+	if err := writer.Close(); err != nil {
+		return nil, "", err
+	}
+	return buffer.Bytes(), writer.FormDataContentType(), nil
+}
+
 func readAIRequest(r *http.Request) ([]byte, string, string, error) {
 	contentType := r.Header.Get("Content-Type")
 	body, err := io.ReadAll(r.Body)

handler/ai_test.go

@@ -1,8 +1,13 @@
 package handler
 
 import (
+	"bytes"
+	"encoding/json"
+	"mime/multipart"
 	"strings"
 	"testing"
+
+	"github.com/basketikun/infinite-canvas/model"
 )
 
 func TestAIUpstreamErrorDetail(t *testing.T) {
@@ -25,3 +30,51 @@ func TestSafeUpstreamTextTruncates(t *testing.T) {
 		t.Fatalf("truncated rune length = %d", len([]rune(got)))
 	}
 }
+
+func TestAdaptAIRequestForSub2APIImagesUsesURLResponseFormat(t *testing.T) {
+	body, contentType, err := adaptAIRequestForChannel([]byte(`{"model":"gpt-image-1","prompt":"cat"}`), "application/json", "/images/generations", model.ModelChannel{Compatibility: "sub2api"})
+	if err != nil {
+		t.Fatalf("adaptAIRequestForChannel returned error: %v", err)
+	}
+	if contentType != "application/json" {
+		t.Fatalf("contentType = %q", contentType)
+	}
+	var payload map[string]any
+	if err := json.Unmarshal(body, &payload); err != nil {
+		t.Fatalf("invalid json body: %v", err)
+	}
+	if got := payload["response_format"]; got != "url" {
+		t.Fatalf("response_format = %#v, want url", got)
+	}
+}
+
+func TestAdaptAIRequestForSub2APIImageEditsUpdatesMultipartResponseFormat(t *testing.T) {
+	var input bytes.Buffer
+	writer := multipart.NewWriter(&input)
+	_ = writer.WriteField("model", "gpt-image-1")
+	_ = writer.WriteField("response_format", "b64_json")
+	part, err := writer.CreateFormFile("image", "input.png")
+	if err != nil {
+		t.Fatalf("CreateFormFile: %v", err)
+	}
+	_, _ = part.Write([]byte("png"))
+	if err := writer.Close(); err != nil {
+		t.Fatalf("Close multipart writer: %v", err)
+	}
+
+	body, contentType, err := adaptAIRequestForChannel(input.Bytes(), writer.FormDataContentType(), "/images/edits", model.ModelChannel{Compatibility: "sub2api"})
+	if err != nil {
+		t.Fatalf("adaptAIRequestForChannel returned error: %v", err)
+	}
+	reader := multipart.NewReader(bytes.NewReader(body), strings.TrimPrefix(contentType, "multipart/form-data; boundary="))
+	form, err := reader.ReadForm(1024)
+	if err != nil {
+		t.Fatalf("ReadForm: %v", err)
+	}
+	if got := form.Value["response_format"]; len(got) != 1 || got[0] != "url" {
+		t.Fatalf("response_format = %#v, want [url]", got)
+	}
+	if files := form.File["image"]; len(files) != 1 || files[0].Filename != "input.png" {
+		t.Fatalf("image file = %#v", files)
+	}
+}