opencatd-open/internal/controller/proxy/proxy.go

package proxy

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"opencatd-open/internal/channel"
	"opencatd-open/internal/dao"
	"opencatd-open/internal/store"
	"opencatd-open/pkg/config"
	"os"
	"strings"
	"sync"

	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
)

type Proxy struct {
	ctx        context.Context
	cfg        *config.Config
	db         *gorm.DB
	wg         *sync.WaitGroup
	httpClient *http.Client

	userDAO    *dao.UserDAO
	apiKeyDAO  *dao.ApiKeyDAO
	usageDAO   *dao.UsageDAO
	dailyDAO   *dao.DailyUsageDAO
	channelSvc *channel.Service
}

func NewProxy(ctx context.Context, cfg *config.Config, db *gorm.DB, wg *sync.WaitGroup, userDAO *dao.UserDAO, apiKeyDAO *dao.ApiKeyDAO, usageDAO *dao.UsageDAO, dailyDAO *dao.DailyUsageDAO) *Proxy {
	client := http.DefaultClient
	if os.Getenv("LOCAL_PROXY") != "" {
		proxyUrl, err := url.Parse(os.Getenv("LOCAL_PROXY"))
		if err == nil {
			tr := &http.Transport{
				Proxy: http.ProxyURL(proxyUrl),
			}
			client.Transport = tr
		}
	}

	np := &Proxy{
		ctx:        ctx,
		cfg:        cfg,
		db:         db,
		wg:         wg,
		httpClient: client,
		userDAO:    userDAO,
		apiKeyDAO:  apiKeyDAO,
		usageDAO:   usageDAO,
		dailyDAO:   dailyDAO,
	}

	return np
}

// SetChannelService sets the channel service (called after construction)
func (p *Proxy) SetChannelService(svc *channel.Service) {
	p.channelSvc = svc
}

func (p *Proxy) HandleProxy(c *gin.Context) {
	path := c.Request.URL.Path
	switch {
	case path == "/v1/chat/completions":
		// TODO: Phase 3 - implement chat completions handler
		c.JSON(http.StatusNotImplemented, gin.H{"error": "chat completions not yet implemented"})
	case strings.HasPrefix(path, "/v1/messages"):
		// TODO: Phase 3 - implement messages handler
		c.JSON(http.StatusNotImplemented, gin.H{"error": "messages not yet implemented"})
	case path == "/v1/responses":
		// TODO: Phase 3 - implement responses handler
		c.JSON(http.StatusNotImplemented, gin.H{"error": "responses not yet implemented"})
	default:
		c.JSON(http.StatusNotFound, gin.H{"error": "unknown endpoint"})
	}
}

func (p *Proxy) HandleModels(c *gin.Context) {
	// TODO: Phase 3 - implement models list
	c.JSON(http.StatusOK, gin.H{"object": "list", "data": []interface{}{}})
}

func (p *Proxy) GetDB() *gorm.DB {
	return p.db
}

// SelectChannel selects the best channel for a model
func (p *Proxy) SelectChannel(modelName string) (*store.Channel, error) {
	if p.channelSvc == nil {
		return nil, fmt.Errorf("channel service not initialized")
	}
	return p.channelSvc.SelectChannel(p.ctx, modelName)
}

// RecordSuccess records a successful request
func (p *Proxy) RecordSuccess(channelID uint64) {
	if p.channelSvc != nil {
		p.channelSvc.RecordSuccess(channelID)
	}
}

// RecordFailure records a failed request
func (p *Proxy) RecordFailure(channelID uint64) {
	if p.channelSvc != nil {
		p.channelSvc.RecordFailure(channelID)
	}
}

// SendUsagePlaceholder placeholder for usage processing
func (p *Proxy) SendUsagePlaceholder(model string, userID uint64, promptTokens, completionTokens int) {
	log.Printf("Usage: model=%s user=%d prompt=%d completion=%d", model, userID, promptTokens, completionTokens)
}

// Placeholder to keep the file compilable
var _ = json.Marshal
var _ = io.ReadAll