444 lines
15 KiB
Go
444 lines
15 KiB
Go
package runner
|
|
|
|
import (
|
|
"context"
|
|
"encoding/base64"
|
|
"errors"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
|
)
|
|
|
|
func TestHydrateProviderRequestAssetsConvertsStrictBase64Field(t *testing.T) {
|
|
storageDir := t.TempDir()
|
|
fileName := "gateway-request-asset-test.png"
|
|
if err := os.WriteFile(filepath.Join(storageDir, fileName), []byte("image bytes"), 0o644); err != nil {
|
|
t.Fatalf("write request asset: %v", err)
|
|
}
|
|
service := &Service{cfg: config.Config{LocalUploadedStorageDir: storageDir}}
|
|
body := map[string]any{
|
|
"model": "demo",
|
|
"b64_json": map[string]any{
|
|
"assetRef": map[string]any{
|
|
"sha256": "sha-test",
|
|
"contentType": "image/png",
|
|
"url": "/static/uploaded/" + fileName,
|
|
"storageProvider": "local_static",
|
|
},
|
|
"url": "/static/uploaded/" + fileName,
|
|
},
|
|
}
|
|
|
|
hydrated, err := service.hydrateProviderRequestAssets(context.Background(), body, store.RuntimeModelCandidate{})
|
|
if err != nil {
|
|
t.Fatalf("hydrate request assets: %v", err)
|
|
}
|
|
if got, want := stringFromAny(hydrated["b64_json"]), base64.StdEncoding.EncodeToString([]byte("image bytes")); got != want {
|
|
t.Fatalf("unexpected hydrated base64: got %q want %q", got, want)
|
|
}
|
|
}
|
|
|
|
func TestHydrateProviderRequestAssetsConvertsBase64ArrayField(t *testing.T) {
|
|
storageDir := t.TempDir()
|
|
fileName := "gateway-request-asset-array.png"
|
|
if err := os.WriteFile(filepath.Join(storageDir, fileName), []byte("array image bytes"), 0o644); err != nil {
|
|
t.Fatalf("write request asset: %v", err)
|
|
}
|
|
service := &Service{cfg: config.Config{LocalUploadedStorageDir: storageDir}}
|
|
body := map[string]any{
|
|
"binary_data_base64": []any{
|
|
map[string]any{
|
|
"assetRef": map[string]any{
|
|
"sha256": "sha-array-image",
|
|
"contentType": "image/png",
|
|
"url": "/static/uploaded/" + fileName,
|
|
"storageProvider": "local_static",
|
|
},
|
|
"url": "/static/uploaded/" + fileName,
|
|
},
|
|
},
|
|
}
|
|
|
|
hydrated, err := service.hydrateProviderRequestAssets(context.Background(), body, store.RuntimeModelCandidate{})
|
|
if err != nil {
|
|
t.Fatalf("hydrate request assets: %v", err)
|
|
}
|
|
values := hydrated["binary_data_base64"].([]any)
|
|
if got, want := stringFromAny(values[0]), base64.StdEncoding.EncodeToString([]byte("array image bytes")); got != want {
|
|
t.Fatalf("unexpected hydrated array base64: got %q want %q", got, want)
|
|
}
|
|
}
|
|
|
|
func TestHydrateProviderRequestAssetsConvertsGeminiInlineDataAssetToRawBase64(t *testing.T) {
|
|
storageDir := t.TempDir()
|
|
fileName := "gateway-request-asset-gemini-inline.png"
|
|
payload := []byte("gemini inline image bytes")
|
|
if err := os.WriteFile(filepath.Join(storageDir, fileName), payload, 0o644); err != nil {
|
|
t.Fatalf("write request asset: %v", err)
|
|
}
|
|
service := &Service{cfg: config.Config{LocalUploadedStorageDir: storageDir}}
|
|
body := map[string]any{
|
|
"contents": []any{
|
|
map[string]any{
|
|
"role": "user",
|
|
"parts": []any{
|
|
map[string]any{
|
|
"inlineData": map[string]any{
|
|
"mimeType": "image/png",
|
|
"data": map[string]any{
|
|
"assetRef": map[string]any{
|
|
"sha256": "sha-gemini-inline",
|
|
"contentType": "image/png",
|
|
"url": "/static/uploaded/" + fileName,
|
|
"storageProvider": "local_static",
|
|
},
|
|
"url": "/static/uploaded/" + fileName,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
hydrated, err := service.hydrateProviderRequestAssets(context.Background(), body, store.RuntimeModelCandidate{Provider: "gemini"})
|
|
if err != nil {
|
|
t.Fatalf("hydrate request assets: %v", err)
|
|
}
|
|
contents := hydrated["contents"].([]any)
|
|
content := contents[0].(map[string]any)
|
|
parts := content["parts"].([]any)
|
|
part := parts[0].(map[string]any)
|
|
inlineData := part["inlineData"].(map[string]any)
|
|
if got, want := stringFromAny(inlineData["data"]), base64.StdEncoding.EncodeToString(payload); got != want {
|
|
t.Fatalf("unexpected hydrated inlineData base64: got %q want %q", got, want)
|
|
}
|
|
}
|
|
|
|
func TestHydrateProviderRequestAssetsConvertsVolcesImageURLAssetToDataURL(t *testing.T) {
|
|
storageDir := t.TempDir()
|
|
fileName := "gateway-request-asset-chat-image.png"
|
|
if err := os.WriteFile(filepath.Join(storageDir, fileName), []byte("chat image bytes"), 0o644); err != nil {
|
|
t.Fatalf("write request asset: %v", err)
|
|
}
|
|
service := &Service{cfg: config.Config{LocalUploadedStorageDir: storageDir}}
|
|
body := map[string]any{
|
|
"messages": []any{
|
|
map[string]any{
|
|
"role": "user",
|
|
"content": []any{
|
|
map[string]any{"type": "text", "text": "describe it"},
|
|
map[string]any{
|
|
"type": "image_url",
|
|
"image_url": map[string]any{
|
|
"url": map[string]any{
|
|
"assetRef": map[string]any{
|
|
"sha256": "sha-chat-image",
|
|
"contentType": "image/png",
|
|
"url": "/static/uploaded/" + fileName,
|
|
"storageProvider": "local_static",
|
|
},
|
|
"url": "/static/uploaded/" + fileName,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
hydrated, err := service.hydrateProviderRequestAssets(context.Background(), body, store.RuntimeModelCandidate{Provider: "volces"})
|
|
if err != nil {
|
|
t.Fatalf("hydrate request assets: %v", err)
|
|
}
|
|
messages := hydrated["messages"].([]any)
|
|
message := messages[0].(map[string]any)
|
|
content := message["content"].([]any)
|
|
imagePart := content[1].(map[string]any)
|
|
imageURL := imagePart["image_url"].(map[string]any)
|
|
if got, want := stringFromAny(imageURL["url"]), "data:image/png;base64,"+base64.StdEncoding.EncodeToString([]byte("chat image bytes")); got != want {
|
|
t.Fatalf("unexpected hydrated image data url: got %q want %q", got, want)
|
|
}
|
|
}
|
|
|
|
func TestHydrateProviderRequestAssetsUsesImageCapabilityBase64ForTopLevelImageAsset(t *testing.T) {
|
|
storageDir := t.TempDir()
|
|
fileName := "gateway-request-asset-edit-image.png"
|
|
payload := []byte("edit image bytes")
|
|
if err := os.WriteFile(filepath.Join(storageDir, fileName), payload, 0o644); err != nil {
|
|
t.Fatalf("write request asset: %v", err)
|
|
}
|
|
service := &Service{cfg: config.Config{LocalUploadedStorageDir: storageDir}}
|
|
body := map[string]any{
|
|
"image": map[string]any{
|
|
"assetRef": map[string]any{
|
|
"sha256": "sha-edit-image",
|
|
"contentType": "image/png",
|
|
"url": "/static/uploaded/" + fileName,
|
|
"storageProvider": "local_static",
|
|
},
|
|
"url": "/static/uploaded/" + fileName,
|
|
},
|
|
}
|
|
|
|
hydrated, err := service.hydrateProviderRequestAssets(context.Background(), body, store.RuntimeModelCandidate{
|
|
ModelType: "image_edit",
|
|
Capabilities: map[string]any{
|
|
"image_edit": map[string]any{
|
|
"support_url_input": false,
|
|
"support_base64_input": true,
|
|
},
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("hydrate request assets: %v", err)
|
|
}
|
|
if got, want := stringFromAny(hydrated["image"]), "data:image/png;base64,"+base64.StdEncoding.EncodeToString(payload); got != want {
|
|
t.Fatalf("unexpected hydrated image data url: got %q want %q", got, want)
|
|
}
|
|
}
|
|
|
|
func TestHydrateProviderRequestAssetsImageCapabilityOverridesProviderDataURLDefault(t *testing.T) {
|
|
service := &Service{}
|
|
body := map[string]any{
|
|
"messages": []any{
|
|
map[string]any{
|
|
"role": "user",
|
|
"content": []any{
|
|
map[string]any{
|
|
"type": "image_url",
|
|
"image_url": map[string]any{
|
|
"url": map[string]any{
|
|
"assetRef": map[string]any{
|
|
"sha256": "sha-url-only-image",
|
|
"contentType": "image/png",
|
|
"url": "https://cdn.example.com/request.png",
|
|
"storageProvider": "remote",
|
|
},
|
|
"url": "https://cdn.example.com/request.png",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
hydrated, err := service.hydrateProviderRequestAssets(context.Background(), body, store.RuntimeModelCandidate{
|
|
Provider: "volces",
|
|
ModelType: "image_edit",
|
|
Capabilities: map[string]any{
|
|
"image_edit": map[string]any{
|
|
"support_url_input": true,
|
|
"support_base64_input": false,
|
|
},
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("hydrate request assets: %v", err)
|
|
}
|
|
messages := hydrated["messages"].([]any)
|
|
message := messages[0].(map[string]any)
|
|
content := message["content"].([]any)
|
|
imagePart := content[0].(map[string]any)
|
|
imageURL := imagePart["image_url"].(map[string]any)
|
|
if got, want := stringFromAny(imageURL["url"]), "https://cdn.example.com/request.png"; got != want {
|
|
t.Fatalf("image capability should keep URL despite provider default, got %q want %q", got, want)
|
|
}
|
|
}
|
|
|
|
func TestHydrateProviderRequestAssetsUsesBase64ForLocalAssetEvenWhenModelSupportsURL(t *testing.T) {
|
|
storageDir := t.TempDir()
|
|
fileName := "gateway-request-asset-local-image.png"
|
|
payload := []byte("local image bytes")
|
|
if err := os.WriteFile(filepath.Join(storageDir, fileName), payload, 0o644); err != nil {
|
|
t.Fatalf("write request asset: %v", err)
|
|
}
|
|
service := &Service{cfg: config.Config{LocalUploadedStorageDir: storageDir}}
|
|
body := map[string]any{
|
|
"image": map[string]any{
|
|
"assetRef": map[string]any{
|
|
"sha256": "sha-local-image",
|
|
"contentType": "image/png",
|
|
"url": "/static/uploaded/" + fileName,
|
|
"storageProvider": "local_static",
|
|
},
|
|
"url": "/static/uploaded/" + fileName,
|
|
},
|
|
}
|
|
|
|
hydrated, err := service.hydrateProviderRequestAssets(context.Background(), body, store.RuntimeModelCandidate{
|
|
ModelType: "image_edit",
|
|
Capabilities: map[string]any{
|
|
"image_edit": map[string]any{
|
|
"support_url_input": true,
|
|
"support_base64_input": true,
|
|
},
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("hydrate request assets: %v", err)
|
|
}
|
|
if got, want := stringFromAny(hydrated["image"]), "data:image/png;base64,"+base64.StdEncoding.EncodeToString(payload); got != want {
|
|
t.Fatalf("local asset should fall back to base64, got %q want %q", got, want)
|
|
}
|
|
}
|
|
|
|
func TestHydrateProviderRequestAssetsConvertsPlainImageURLWhenModelRequiresBase64(t *testing.T) {
|
|
payload := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 0, 0, 0, 0}
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "image/png")
|
|
_, _ = w.Write(payload)
|
|
}))
|
|
defer server.Close()
|
|
service := &Service{}
|
|
body := map[string]any{"image": server.URL + "/source.png"}
|
|
|
|
hydrated, err := service.hydrateProviderRequestAssets(context.Background(), body, store.RuntimeModelCandidate{
|
|
ModelType: "image_edit",
|
|
Capabilities: map[string]any{
|
|
"image_edit": map[string]any{
|
|
"support_url_input": false,
|
|
"support_base64_input": true,
|
|
},
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("hydrate request assets: %v", err)
|
|
}
|
|
if got, want := stringFromAny(hydrated["image"]), "data:image/png;base64,"+base64.StdEncoding.EncodeToString(payload); got != want {
|
|
t.Fatalf("plain image URL should be converted to data URL, got %q want %q", got, want)
|
|
}
|
|
}
|
|
|
|
func TestHydrateProviderRequestAssetsConvertsPrivateImageURLEvenWhenModelSupportsURL(t *testing.T) {
|
|
payload := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 1, 2, 3, 4}
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "image/png")
|
|
_, _ = w.Write(payload)
|
|
}))
|
|
defer server.Close()
|
|
service := &Service{}
|
|
body := map[string]any{"image": server.URL + "/source.png"}
|
|
|
|
hydrated, err := service.hydrateProviderRequestAssets(context.Background(), body, store.RuntimeModelCandidate{
|
|
ModelType: "image_edit",
|
|
Capabilities: map[string]any{
|
|
"image_edit": map[string]any{
|
|
"support_url_input": true,
|
|
"support_base64_input": true,
|
|
},
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("hydrate request assets: %v", err)
|
|
}
|
|
if got, want := stringFromAny(hydrated["image"]), "data:image/png;base64,"+base64.StdEncoding.EncodeToString(payload); got != want {
|
|
t.Fatalf("private image URL should fall back to data URL, got %q want %q", got, want)
|
|
}
|
|
}
|
|
|
|
func TestHydrateProviderRequestAssetsKeepsImageURLAssetAsURLForProviderURLDefault(t *testing.T) {
|
|
service := &Service{}
|
|
body := map[string]any{
|
|
"messages": []any{
|
|
map[string]any{
|
|
"role": "user",
|
|
"content": []any{
|
|
map[string]any{
|
|
"type": "image_url",
|
|
"image_url": map[string]any{
|
|
"url": map[string]any{
|
|
"assetRef": map[string]any{
|
|
"sha256": "sha-url-image",
|
|
"contentType": "image/png",
|
|
"url": "https://cdn.example.com/request.png",
|
|
"storageProvider": "remote",
|
|
},
|
|
"url": "https://cdn.example.com/request.png",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
hydrated, err := service.hydrateProviderRequestAssets(context.Background(), body, store.RuntimeModelCandidate{Provider: "jimeng"})
|
|
if err != nil {
|
|
t.Fatalf("hydrate request assets: %v", err)
|
|
}
|
|
messages := hydrated["messages"].([]any)
|
|
message := messages[0].(map[string]any)
|
|
content := message["content"].([]any)
|
|
imagePart := content[0].(map[string]any)
|
|
imageURL := imagePart["image_url"].(map[string]any)
|
|
if got, want := stringFromAny(imageURL["url"]), "https://cdn.example.com/request.png"; got != want {
|
|
t.Fatalf("unexpected hydrated image URL: got %q want %q", got, want)
|
|
}
|
|
}
|
|
|
|
func TestHydrateProviderRequestAssetsReturnsExpiredError(t *testing.T) {
|
|
expiredAt := time.Now().Add(-time.Minute).UTC().Format(time.RFC3339)
|
|
service := &Service{}
|
|
body := map[string]any{
|
|
"image_url": map[string]any{
|
|
"assetRef": map[string]any{
|
|
"sha256": "sha-expired",
|
|
"contentType": "image/png",
|
|
"url": "/static/uploaded/gateway-request-asset-expired.png",
|
|
"storageProvider": "local_static",
|
|
"expiresAt": expiredAt,
|
|
},
|
|
"url": "/static/uploaded/gateway-request-asset-expired.png",
|
|
},
|
|
}
|
|
|
|
_, err := service.hydrateProviderRequestAssets(context.Background(), body, store.RuntimeModelCandidate{})
|
|
if err == nil {
|
|
t.Fatal("expected expired request asset error")
|
|
}
|
|
var clientErr *clients.ClientError
|
|
if !errors.As(err, &clientErr) || clientErr.Code != "request_asset_expired" {
|
|
t.Fatalf("expected request_asset_expired, got %T %v", err, err)
|
|
}
|
|
}
|
|
|
|
func TestStringFromAnyReadsRequestAssetWrapperURL(t *testing.T) {
|
|
wrapper := map[string]any{
|
|
"assetRef": map[string]any{"url": "https://cdn.example/request.png"},
|
|
}
|
|
if got := stringFromAny(wrapper); got != "https://cdn.example/request.png" {
|
|
t.Fatalf("expected wrapper URL, got %q", got)
|
|
}
|
|
}
|
|
|
|
func TestSlimTaskRequestSnapshotKeepsMessageRefs(t *testing.T) {
|
|
service := &Service{}
|
|
task := store.GatewayTask{Request: map[string]any{
|
|
"conversationId": "conv-1",
|
|
"messageRefs": []any{map[string]any{"messageId": "msg-1", "position": 0}},
|
|
"newMessageCount": 1,
|
|
}}
|
|
body := map[string]any{
|
|
"model": "demo",
|
|
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
|
|
}
|
|
|
|
snapshot := service.slimTaskRequestSnapshot(task, body)
|
|
|
|
if snapshot["messages"] != nil {
|
|
t.Fatalf("snapshot should not persist restored messages: %+v", snapshot)
|
|
}
|
|
if snapshot["messageRefs"] == nil || snapshot["newMessageCount"] != 1 {
|
|
t.Fatalf("snapshot should keep message refs and new count: %+v", snapshot)
|
|
}
|
|
}
|