105 lines
3.0 KiB
Go
105 lines
3.0 KiB
Go
package runner
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"strings"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
|
|
)
|
|
|
|
func (s *Service) uploadGeneratedAssets(ctx context.Context, result map[string]any) (map[string]any, error) {
|
|
if s.cfg.ServerMainBaseURL == "" || s.cfg.ServerMainInternalToken == "" {
|
|
return result, nil
|
|
}
|
|
data, _ := result["data"].([]any)
|
|
if len(data) == 0 {
|
|
return result, nil
|
|
}
|
|
next := map[string]any{}
|
|
for key, value := range result {
|
|
next[key] = value
|
|
}
|
|
nextData := make([]any, 0, len(data))
|
|
for index, rawItem := range data {
|
|
item, _ := rawItem.(map[string]any)
|
|
if item == nil {
|
|
nextData = append(nextData, rawItem)
|
|
continue
|
|
}
|
|
b64 := stringFromMap(item, "b64_json")
|
|
if b64 == "" {
|
|
nextData = append(nextData, rawItem)
|
|
continue
|
|
}
|
|
upload, err := s.uploadBase64Image(ctx, b64, index)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
merged := map[string]any{}
|
|
for key, value := range item {
|
|
if key != "b64_json" {
|
|
merged[key] = value
|
|
}
|
|
}
|
|
merged["upload"] = upload
|
|
if urlValue, ok := upload["url"].(string); ok && urlValue != "" {
|
|
merged["url"] = urlValue
|
|
}
|
|
nextData = append(nextData, merged)
|
|
}
|
|
next["data"] = nextData
|
|
return next, nil
|
|
}
|
|
|
|
func (s *Service) uploadBase64Image(ctx context.Context, b64 string, index int) (map[string]any, error) {
|
|
payload, err := base64.StdEncoding.DecodeString(stripDataURLPrefix(b64))
|
|
if err != nil {
|
|
return nil, &clients.ClientError{Code: "upload_decode_failed", Message: err.Error(), Retryable: false}
|
|
}
|
|
var body bytes.Buffer
|
|
writer := multipart.NewWriter(&body)
|
|
fileWriter, err := writer.CreateFormFile("file", fmt.Sprintf("gateway-result-%d.png", index+1))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if _, err := fileWriter.Write(payload); err != nil {
|
|
return nil, err
|
|
}
|
|
_ = writer.WriteField("source", "ai-gateway")
|
|
if err := writer.Close(); err != nil {
|
|
return nil, err
|
|
}
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(s.cfg.ServerMainBaseURL, "/")+"/v1/files/upload", &body)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
req.Header.Set("Content-Type", writer.FormDataContentType())
|
|
req.Header.Set("Authorization", "Bearer "+s.cfg.ServerMainInternalToken)
|
|
resp, err := http.DefaultClient.Do(req)
|
|
if err != nil {
|
|
return nil, &clients.ClientError{Code: "upload_network", Message: err.Error(), Retryable: true}
|
|
}
|
|
defer resp.Body.Close()
|
|
var decoded map[string]any
|
|
if err := json.NewDecoder(resp.Body).Decode(&decoded); err != nil {
|
|
return nil, &clients.ClientError{Code: "upload_invalid_response", Message: err.Error(), Retryable: false}
|
|
}
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
return nil, &clients.ClientError{Code: "upload_failed", Message: "server-main upload failed", StatusCode: resp.StatusCode, Retryable: resp.StatusCode >= 500}
|
|
}
|
|
return decoded, nil
|
|
}
|
|
|
|
func stripDataURLPrefix(value string) string {
|
|
if index := strings.Index(value, ","); strings.HasPrefix(value, "data:") && index >= 0 {
|
|
return value[index+1:]
|
|
}
|
|
return value
|
|
}
|