model-gateway/service/model_manage_service.go

package service

import (
	"context"
	"fmt"
	"model-gateway/consts/model"
	"model-gateway/consts/public"
	"model-gateway/dao"
	"model-gateway/model/dto"
	"model-gateway/model/entity"

	"gitea.redpowerfuture.com/red-future/common/beans"
	"gitea.redpowerfuture.com/red-future/common/db/gfdb"
	"gitea.redpowerfuture.com/red-future/common/utils"
	"github.com/gogf/gf/v2/database/gdb"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/util/gconv"
)

var ModelManage = &modelManageService{}

type modelManageService struct{}

// Create 创建模型
func (s *modelManageService) Create(ctx context.Context, req *dto.CreateModelManageReq) (res *dto.CreateModelManageRes, err error) {
	err = gfdb.DB(ctx, public.DbNameModelGateway).Transaction(ctx, func(ctx context.Context, tx gdb.TX) (err error) {
		// 1）检查是否是超管
		var isSuperAdmin bool
		isSuperAdmin, err = IsSuperAdmin(ctx)
		if err != nil {
			return
		}
		req.SystemModel = &isSuperAdmin
		// 1）如果设为会话模型，先把该用户旧会话模型取消
		err = s.CancelChatModel(ctx, req.ModelType, req.ChatModel, isSuperAdmin)
		if err != nil {
			return
		}
		// 2）模型名称唯一性：同一用户下不允许同名模型
		var user *beans.User
		user, err = utils.GetUserInfo(ctx)
		if err != nil {
			return
		}
		if !g.IsEmpty(req.ModelName) {
			var exist *entity.ModelManage
			exist, err = dao.ModelManage.GetByCreatorAndName(ctx, user.UserName, req.ModelName)
			if err != nil {
				return
			}
			if exist != nil {
				return fmt.Errorf("模型名称 [%s] 已存在，同一用户下不能重复添加同名模型", req.ModelName)
			}
		}
		// 3）插入数据
		id, err := dao.ModelManage.Insert(ctx, req)
		if err != nil {
			return
		}
		res = &dto.CreateModelManageRes{Id: id}
		return
	})
	return
}

// Update 更新模型配置
func (s *modelManageService) Update(ctx context.Context, req *dto.UpdateModelManageReq) (res *dto.GetModelManageRes, err error) {
	err = gfdb.DB(ctx, public.DbNameModelGateway).Transaction(ctx, func(ctx context.Context, tx gdb.TX) (err error) {
		var get *entity.ModelManage
		get, err = dao.ModelManage.GetNotTenantId(ctx, &dto.GetModelManageReq{
			Id: req.Id,
		})
		if err != nil {
			return
		}
		var user *beans.User
		user, err = utils.GetUserInfo(ctx)
		if err != nil {
			return
		}
		// 1）如果不是创建者，且是系统模型，则需要拷贝
		if get.Creator != user.UserName {
			if get.SystemModel != nil && *get.SystemModel {
				// 拷贝前先查同名：用户已存在同名模型则直接返回用户自己的模型，避免重复拷贝
				copyName := req.ModelName
				if g.IsEmpty(copyName) {
					copyName = get.ModelName
				}
				if !g.IsEmpty(copyName) {
					var exist *entity.ModelManage
					exist, err = dao.ModelManage.GetByCreatorAndName(ctx, user.UserName, copyName)
					if err != nil {
						return
					}
					if exist != nil {
						res, err = s.Get(ctx, &dto.GetModelManageReq{Id: exist.Id})
						return
					}
				}
				if g.IsEmpty(req.ApiKey) {
					return fmt.Errorf("模型apiKey不能为空")
				}
				d := new(dto.CreateModelManageReq)
				err = gconv.Struct(req, d)
				if err != nil {
					return
				}
				var r *dto.CreateModelManageRes
				r, err = s.Create(ctx, d)
				if err != nil {
					return err
				}
				res, err = s.Get(ctx, &dto.GetModelManageReq{
					Id: r.Id,
				})
				return
			}
			return fmt.Errorf("无权限操作")
		}

		// 1）检查是否是超管
		var isSuperAdmin bool
		isSuperAdmin, err = IsSuperAdmin(ctx)
		if err != nil {
			return
		}
		modelType := req.ModelType
		if g.IsEmpty(modelType) {
			modelType = get.ModelType
		}
		// 1）如果设为会话模型，先把该用户旧会话模型取消
		err = s.CancelChatModel(ctx, modelType, req.ChatModel, isSuperAdmin)
		if err != nil {
			return
		}
		// 2）模型名称唯一性：同一用户下不允许同名模型（排除自身）
		if !g.IsEmpty(req.ModelName) {
			var exist *entity.ModelManage
			exist, err = dao.ModelManage.GetByCreatorAndName(ctx, user.UserName, req.ModelName)
			if err != nil {
				return
			}
			if exist != nil && exist.Id != req.Id {
				return fmt.Errorf("模型名称 [%s] 已存在，同一用户下不能重复添加同名模型", req.ModelName)
			}
		}
		// 3）更新数据
		_, err = dao.ModelManage.Update(ctx, req)
		return
	})
	return
}

func (s *modelManageService) CancelChatModel(ctx context.Context, modelType model.ModelType, chatModel *bool, isSuperAdmin bool) (err error) {
	if !g.IsEmpty(chatModel) && *chatModel {
		if !g.IsEmpty(modelType) && *modelType == *model.ModelTypeInference.Code {
			if isSuperAdmin {
				return fmt.Errorf("超级管理员不能设置会话模型")
			}
			// 2）获取该用户信息
			var user *beans.User
			user, err = utils.GetUserInfo(ctx)
			if err != nil {
				return
			}
			// 3）取消该用户之前的会话模型
			var get *entity.ModelManage
			get, err = dao.ModelManage.Get(ctx, &dto.GetModelManage{
				Creator:   user.UserName,
				ChatModel: chatModel,
			})
			if err != nil {
				return
			}
			if g.IsEmpty(get) {
				return
			}
			_, err = dao.ModelManage.Update(ctx, &dto.UpdateModelManageReq{
				Id:        get.Id,
				ChatModel: gconv.PtrBool(false),
			})
			if err != nil {
				return
			}
		} else {
			return fmt.Errorf("只有推理模型可以设置成会话模型")
		}
	}
	return
}

// Delete 删除模型
func (s *modelManageService) Delete(ctx context.Context, req *dto.DeleteModelManageReq) error {
	_, err := dao.ModelManage.Delete(ctx, req)
	return err
}

func (s *modelManageService) Get(ctx context.Context, req *dto.GetModelManageReq) (res *dto.GetModelManageRes, err error) {
	get, err := dao.ModelManage.GetNotTenantId(ctx, req)
	if err != nil {
		return
	}
	res = new(dto.GetModelManageRes)
	err = gconv.Struct(get, &res.ModelManage)
	return
}

func (s *modelManageService) GetChatModel(ctx context.Context, req *dto.GetChatModelReq) (res *dto.GetChatModelRes, err error) {
	user, err := utils.GetUserInfo(ctx)
	if err != nil {
		return
	}
	get, err := dao.ModelManage.Get(ctx, &dto.GetModelManage{
		Creator:   user.UserName,
		ChatModel: gconv.PtrBool(true),
	})
	if err != nil {
		return
	}
	res = &dto.GetChatModelRes{
		ModelManage: get,
	}
	return
}

// List 获取模型列表
func (s *modelManageService) List(ctx context.Context, req *dto.ListModelManageReq) (res *dto.ListModelManageRes, err error) {
	if req.IsSameType && !g.IsEmpty(req.Id) {
		var get *entity.ModelManage
		get, err = dao.ModelManage.GetNotTenantId(ctx, &dto.GetModelManageReq{
			Id: req.Id,
		})
		if err != nil {
			return
		}
		req.ModelType = get.ModelType
	}
	var user *beans.User
	user, err = utils.GetUserInfo(ctx)
	if err != nil {
		return
	}
	req.Creator = user.UserName
	list, total, err := dao.ModelManage.ListNotTenantId(ctx, req)
	if err != nil {
		return
	}
	res = &dto.ListModelManageRes{
		Total: total,
	}
	err = gconv.Struct(list, &res.List)
	return
}

func (s *modelManageService) CheckChatModel(ctx context.Context, req *dto.CheckChatModelReq) (res *dto.CheckChatModelRes, err error) {
	get, err := s.GetChatModel(ctx, &dto.GetChatModelReq{})
	if err != nil {
		return
	}
	res = &dto.CheckChatModelRes{
		IsChatModel: !g.IsEmpty(get),
	}
	return
}

// GetModelType 获取模型类型
func (s *modelManageService) GetModelType(ctx context.Context, req *dto.ModelTypeReq) (res *dto.ModelTypeRes, err error) {
	res = &dto.ModelTypeRes{
		List: model.GetTypeTreeList(),
	}
	return res, nil
}

// GetModelSupplier 获取运营商列表
func (s *modelManageService) GetModelSupplier(ctx context.Context, req *dto.ModelSupplierReq) (res *dto.ModelSupplierRes, err error) {
	return &dto.ModelSupplierRes{
		List: model.GetSupplierOptionList(),
	}, nil
}