Atlas/internal/api/gateway_openai.go

package api

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"

	"github.com/gin-gonic/gin"

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

func (s *Server) gatewayProviderID(c *gin.Context) int64 {
	if v := c.GetHeader("X-Provider-Id"); v != "" {
		id, _ := strconv.ParseInt(v, 10, 64)
		return id
	}
	return 0
}

func (s *Server) gatewayModels(c *gin.Context) {
	ctx := c.Request.Context()
	serviceType := c.DefaultQuery("service_type", string(domain.ServiceChat))
	protocol := c.DefaultQuery("protocol", domain.DefaultProtocol(domain.ServiceType(serviceType)))
	cfg, err := s.app.Router.ResolveByProtocol(ctx, serviceType, protocol, s.gatewayProviderID(c))
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	models, err := s.app.Models.ListEffective(ctx, cfg)
	if err != nil || len(models) == 0 {
		upstream, fetchErr := s.app.Models.FetchUpstream(ctx, cfg)
		if fetchErr == nil && len(upstream) > 0 {
			data := map[string]any{
				"object": "list",
				"data":   upstream,
			}
			out, _ := json.Marshal(data)
			c.Data(http.StatusOK, "application/json", out)
			return
		}
		path := provider.ModelsPath(cfg)
		res, err := s.app.Router.ProxyWithMeta(ctx, cfg, "GET", path, nil, nil, provider.CallMeta{
			SessionID: sessionID(c),
			Method:    "GET",
			Path:      path,
		})
		if err != nil {
			c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
			return
		}
		c.Data(res.StatusCode, "application/json", res.Body)
		return
	}
	data := map[string]any{
		"object": "list",
		"data":   models,
	}
	out, _ := json.Marshal(data)
	c.Data(http.StatusOK, "application/json", out)
}

func (s *Server) gatewayChatCompletions(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
	_ = json.Unmarshal(body, &req)
	model, _ := req["model"].(string)
	stream, _ := req["stream"].(bool)

	ctx := c.Request.Context()
	cfg, err := s.app.Router.ResolveByProtocol(ctx, string(domain.ServiceChat), "openai", s.gatewayProviderID(c))
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	meta := provider.CallMeta{
		SessionID: sessionID(c),
		Method:    "POST",
		Path:      "/v1/chat/completions",
		Model:     model,
	}
	if stream {
		c.Header("Content-Type", "text/event-stream")
		c.Header("Cache-Control", "no-cache")
		c.Header("Connection", "keep-alive")
		err = s.app.Router.ProxyStream(ctx, cfg, "POST", "/v1/chat/completions", body, nil, meta, c.Writer, func() { c.Writer.Flush() })
		if err != nil {
			c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		}
		return
	}
	res, err := s.app.Router.ProxyWithMeta(ctx, cfg, "POST", "/v1/chat/completions", body, 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) gatewayResponses(c *gin.Context) {
	body, err := io.ReadAll(c.Request.Body)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	chatBody, model, err := responsesToChat(body)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	ctx := c.Request.Context()
	cfg, err := s.app.Router.ResolveByProtocol(ctx, string(domain.ServiceChat), "openai", s.gatewayProviderID(c))
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	meta := provider.CallMeta{
		SessionID: sessionID(c),
		Method:    "POST",
		Path:      "/v1/responses",
		Model:     model,
	}
	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, "POST", "/v1/chat/completions", chatBody, nil, meta, c.Writer, func() { c.Writer.Flush() }); err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
	}
}

func (s *Server) gatewayProxyOpenAI(path string) gin.HandlerFunc {
	return func(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
		_ = json.Unmarshal(body, &req)
		model, _ := req["model"].(string)
		ctx := c.Request.Context()
		cfg, err := s.app.Router.ResolveByProtocol(ctx, string(domain.ServiceChat), "openai", s.gatewayProviderID(c))
		if err != nil {
			c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
			return
		}
		meta := provider.CallMeta{
			SessionID: sessionID(c),
			Method:    c.Request.Method,
			Path:      path,
			Model:     model,
		}
		res, err := s.app.Router.ProxyWithMeta(ctx, cfg, c.Request.Method, path, body, nil, meta)
		if err != nil {
			c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
			return
		}
		c.Data(res.StatusCode, "application/json", res.Body)
	}
}

func responsesToChat(body []byte) ([]byte, string, error) {
	var req map[string]any
	if err := json.Unmarshal(body, &req); err != nil {
		return nil, "", err
	}
	model, _ := req["model"].(string)
	messages := []map[string]string{}
	if input, ok := req["input"].(string); ok && input != "" {
		messages = append(messages, map[string]string{"role": "user", "content": input})
	}
	if items, ok := req["input"].([]any); ok {
		for _, item := range items {
			m, ok := item.(map[string]any)
			if !ok {
				continue
			}
			role, _ := m["role"].(string)
			if role == "" {
				role = "user"
			}
			content, _ := m["content"].(string)
			if content == "" {
				if text, ok := m["text"].(string); ok {
					content = text
				}
			}
			if content != "" {
				messages = append(messages, map[string]string{"role": role, "content": content})
			}
		}
	}
	if len(messages) == 0 {
		return nil, "", fmt.Errorf("responses input required")
	}
	out, err := json.Marshal(map[string]any{
		"model":    model,
		"messages": messages,
		"stream":   true,
	})
	return out, model, err
}

func inferImageServiceType(input map[string]any) domain.ServiceType {
	if _, ok := input["image"]; ok {
		return domain.ServiceImageToImage
	}
	return domain.ServiceTextToImage
}

func gatewayPreferWait(c *gin.Context) bool {
	return strings.Contains(strings.ToLower(c.GetHeader("Prefer")), "wait")
}