aiproxy/core/model/groupmcp.go

package model

import (
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/labring/aiproxy/core/common"
	log "github.com/sirupsen/logrus"
	"gorm.io/gorm"
)

type GroupMCPStatus int

const (
	GroupMCPStatusEnabled GroupMCPStatus = iota + 1
	GroupMCPStatusDisabled
)

const (
	ErrGroupMCPNotFound = "group mcp"
)

type GroupMCPType string

const (
	GroupMCPTypeProxySSE        GroupMCPType = "mcp_proxy_sse"
	GroupMCPTypeProxyStreamable GroupMCPType = "mcp_proxy_streamable"
	GroupMCPTypeOpenAPI         GroupMCPType = "mcp_openapi"
)

type GroupMCPProxyConfig struct {
	URL     string            `json:"url"`
	Querys  map[string]string `json:"querys"`
	Headers map[string]string `json:"headers"`
}

type GroupMCP struct {
	ID            string               `gorm:"primaryKey"                    json:"id"`
	GroupID       string               `gorm:"primaryKey"                    json:"group_id"`
	Group         *Group               `gorm:"foreignKey:GroupID"            json:"-"`
	Status        GroupMCPStatus       `gorm:"index;default:1"               json:"status"`
	CreatedAt     time.Time            `gorm:"index,autoCreateTime"          json:"created_at"`
	UpdateAt      time.Time            `gorm:"index,autoUpdateTime"          json:"update_at"`
	Name          string               `                                     json:"name"`
	Type          GroupMCPType         `gorm:"index"                         json:"type"`
	Description   string               `                                     json:"description"`
	ProxyConfig   *GroupMCPProxyConfig `gorm:"serializer:fastjson;type:text" json:"proxy_config,omitempty"`
	OpenAPIConfig *MCPOpenAPIConfig    `gorm:"serializer:fastjson;type:text" json:"openapi_config,omitempty"`
}

func (g *GroupMCP) BeforeSave(_ *gorm.DB) (err error) {
	if g.GroupID == "" {
		return errors.New("group id is empty")
	}

	if err := validateMCPID(g.ID); err != nil {
		return err
	}

	if g.UpdateAt.IsZero() {
		g.UpdateAt = time.Now()
	}

	if g.Status == 0 {
		g.Status = GroupMCPStatusEnabled
	}

	if g.OpenAPIConfig != nil {
		config := g.OpenAPIConfig
		if config.OpenAPISpec != "" {
			return validateHTTPURL(config.OpenAPISpec)
		}

		if config.OpenAPIContent != "" {
			return nil
		}

		return errors.New("openapi spec and content is empty")
	}

	if g.ProxyConfig != nil {
		config := g.ProxyConfig
		return validateHTTPURL(config.URL)
	}

	return err
}

// CreateGroupMCP creates a new GroupMCP
func CreateGroupMCP(mcp *GroupMCP) error {
	err := DB.Create(mcp).Error
	if err != nil && errors.Is(err, gorm.ErrDuplicatedKey) {
		return errors.New("group mcp already exists")
	}

	return err
}

// UpdateGroupMCP updates an existing GroupMCP
func UpdateGroupMCP(mcp *GroupMCP) (err error) {
	defer func() {
		if err == nil {
			if err := CacheDeleteGroupMCP(mcp.GroupID, mcp.ID); err != nil {
				log.Error("cache delete group mcp error: " + err.Error())
			}
		}
	}()

	selects := []string{
		"name",
		"proxy_config",
		"openapi_config",
		"description",
	}
	if mcp.Type != "" {
		selects = append(selects, "type")
	}

	if mcp.Status != 0 {
		selects = append(selects, "status")
	}

	result := DB.
		Select(selects).
		Where("id = ? AND group_id = ?", mcp.ID, mcp.GroupID).
		Updates(mcp)

	return HandleUpdateResult(result, ErrGroupMCPNotFound)
}

func UpdateGroupMCPStatus(id, groupID string, status GroupMCPStatus) (err error) {
	defer func() {
		if err == nil {
			if err := CacheDeleteGroupMCP(groupID, id); err != nil {
				log.Error("cache delete group mcp error: " + err.Error())
			}
		}
	}()

	result := DB.Model(&GroupMCP{}).
		Where("id = ? AND group_id = ?", id, groupID).
		Update("status", status)

	return HandleUpdateResult(result, ErrGroupMCPNotFound)
}

func GetAllGroupMCPs(status GroupMCPStatus) ([]GroupMCP, error) {
	var mcps []GroupMCP

	tx := DB.Model(&GroupMCP{})
	if status != 0 {
		tx = tx.Where("status = ?", status)
	}

	err := tx.Find(&mcps).Error

	return mcps, err
}

// DeleteGroupMCP deletes a GroupMCP by ID and GroupID
func DeleteGroupMCP(id, groupID string) (err error) {
	defer func() {
		if err == nil {
			if err := CacheDeleteGroupMCP(groupID, id); err != nil {
				log.Error("cache delete group mcp error: " + err.Error())
			}
		}
	}()

	if id == "" || groupID == "" {
		return errors.New("group mcp id or group id is empty")
	}

	result := DB.Where("id = ? AND group_id = ?", id, groupID).Delete(&GroupMCP{})

	return HandleUpdateResult(result, ErrGroupMCPNotFound)
}

// GetGroupMCPByID retrieves a GroupMCP by ID and GroupID
func GetGroupMCPByID(id, groupID string) (GroupMCP, error) {
	var mcp GroupMCP
	if id == "" || groupID == "" {
		return mcp, errors.New("group mcp id or group id is empty")
	}

	err := DB.Where("id = ? AND group_id = ?", id, groupID).First(&mcp).Error

	return mcp, HandleNotFound(err, ErrGroupMCPNotFound)
}

// GetGroupMCPs retrieves GroupMCPs with pagination and filtering
func GetGroupMCPs(
	groupID string,
	page, perPage int,
	id string,
	mcpType GroupMCPType,
	keyword string,
	status GroupMCPStatus,
) (mcps []GroupMCP, total int64, err error) {
	if groupID == "" {
		return nil, 0, errors.New("group id is empty")
	}

	tx := DB.Model(&GroupMCP{}).Where("group_id = ?", groupID)

	if id != "" {
		tx = tx.Where("id = ?", id)
	}

	if status != 0 {
		tx = tx.Where("status = ?", status)
	}

	if mcpType != "" {
		tx = tx.Where("type = ?", mcpType)
	}

	if keyword != "" {
		var (
			conditions []string
			values     []any
		)

		if id == "" {
			if !common.UsingSQLite {
				conditions = append(conditions, "id ILIKE ?")
				values = append(values, "%"+keyword+"%")
			} else {
				conditions = append(conditions, "id LIKE ?")
				values = append(values, "%"+keyword+"%")
			}
		}

		if !common.UsingSQLite {
			conditions = append(conditions, "name ILIKE ?")
			values = append(values, "%"+keyword+"%")
		} else {
			conditions = append(conditions, "name LIKE ?")
			values = append(values, "%"+keyword+"%")
		}

		if !common.UsingSQLite {
			conditions = append(conditions, "description ILIKE ?")
			values = append(values, "%"+keyword+"%")
		} else {
			conditions = append(conditions, "description LIKE ?")
			values = append(values, "%"+keyword+"%")
		}

		if len(conditions) > 0 {
			tx = tx.Where(fmt.Sprintf("(%s)", strings.Join(conditions, " OR ")), values...)
		}
	}

	err = tx.Count(&total).Error
	if err != nil {
		return nil, 0, err
	}

	if total <= 0 {
		return nil, 0, nil
	}

	limit, offset := toLimitOffset(page, perPage)
	err = tx.
		Limit(limit).
		Offset(offset).
		Find(&mcps).
		Error

	return mcps, total, err
}