Atlas/internal/api/gateway_replicate.go

package api

import (
	"encoding/json"
	"io"
	"net/http"

	"github.com/gin-gonic/gin"

	"github.com/rose_cat707/Atlas/internal/domain"
	"github.com/rose_cat707/Atlas/internal/provider"
)

func (s *Server) gatewayCreatePrediction(c *gin.Context) {
	body, err := io.ReadAll(c.Request.Body)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	var req map[string]any
	if err := json.Unmarshal(body, &req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	model, _ := req["model"].(string)
	input, _ := req["input"].(map[string]any)
	if input == nil {
		input = map[string]any{}
	}
	params, _ := json.Marshal(input)
	serviceType := inferImageServiceType(input)
	mode := domain.ModeAsync
	if gatewayPreferWait(c) {
		mode = domain.ModeSync
	}
	ctx := c.Request.Context()
	result, err := s.app.Invoke.Run(ctx, domain.InvokeRequest{
		SessionID:   sessionID(c),
		ServiceType: serviceType,
		ProviderID:  s.gatewayProviderID(c),
		Model:       model,
		Params:      params,
		Mode:        mode,
	})
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	if mode == domain.ModeSync {
		var output any
		_ = json.Unmarshal(result.Output, &output)
		c.JSON(http.StatusOK, gin.H{
			"id":     "sync",
			"status": "succeeded",
			"output": output,
		})
		return
	}
	task, err := s.app.Store.GetTask(ctx, sessionID(c), *result.TaskID)
	if err != nil {
		c.JSON(http.StatusOK, gin.H{"id": result.TaskID, "status": "starting"})
		return
	}
	c.JSON(http.StatusOK, gin.H{
		"id":     task.ExternalID,
		"status": task.Status,
	})
}

func (s *Server) gatewayListPredictions(c *gin.Context) {
	ctx := c.Request.Context()
	cfg, err := s.app.Router.ResolveByProtocol(ctx, string(domain.ServiceTextToImage), "replicate", s.gatewayProviderID(c))
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	meta := provider.CallMeta{
		SessionID: sessionID(c),
		Method:    "GET",
		Path:      "/v1/predictions",
	}
	res, err := s.app.Router.ProxyWithMeta(ctx, cfg, "GET", "/v1/predictions", nil, nil, meta)
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	c.Data(res.StatusCode, "application/json", res.Body)
}

func (s *Server) gatewayGetPrediction(c *gin.Context) {
	ctx := c.Request.Context()
	cfg, err := s.app.Router.ResolveByProtocol(ctx, string(domain.ServiceTextToImage), "replicate", s.gatewayProviderID(c))
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	path := "/v1/predictions/" + c.Param("id")
	meta := provider.CallMeta{
		SessionID: sessionID(c),
		Method:    "GET",
		Path:      "/v1/predictions/:id",
	}
	res, err := s.app.Router.ProxyWithMeta(ctx, cfg, "GET", path, nil, nil, meta)
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	c.Data(res.StatusCode, "application/json", res.Body)
}

func (s *Server) gatewayCancelPrediction(c *gin.Context) {
	ctx := c.Request.Context()
	id := c.Param("id")
	cfg, err := s.app.Router.ResolveByProtocol(ctx, string(domain.ServiceTextToImage), "replicate", s.gatewayProviderID(c))
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	path := "/v1/predictions/" + id + "/cancel"
	meta := provider.CallMeta{
		SessionID: sessionID(c),
		Method:    "POST",
		Path:      "/v1/predictions/:id/cancel",
	}
	res, err := s.app.Router.ProxyWithMeta(ctx, cfg, "POST", path, nil, nil, meta)
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	c.Data(res.StatusCode, "application/json", res.Body)
}

func (s *Server) gatewayStreamPrediction(c *gin.Context) {
	ctx := c.Request.Context()
	id := c.Param("id")
	cfg, err := s.app.Router.ResolveByProtocol(ctx, string(domain.ServiceTextToImage), "replicate", s.gatewayProviderID(c))
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	path := "/v1/predictions/" + id + "/stream"
	meta := provider.CallMeta{
		SessionID: sessionID(c),
		Method:    "GET",
		Path:      "/v1/predictions/:id/stream",
	}
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	if err := s.app.Router.ProxyStream(ctx, cfg, "GET", path, nil, nil, meta, c.Writer, func() { c.Writer.Flush() }); err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
	}
}