wechat-api/internal/logic/agent/create_agent_logic.go

package agent

import (
	"context"
	"errors"
	"github.com/zeromicro/go-zero/core/errorx"
	"wechat-api/hook/fastgpt"
	"wechat-api/internal/svc"
	"wechat-api/internal/types"
	"wechat-api/internal/utils/dberrorhandler"

	"github.com/suyuan32/simple-admin-common/msg/errormsg"

	"github.com/zeromicro/go-zero/core/logx"
)

type CreateAgentLogic struct {
	ctx    context.Context
	svcCtx *svc.ServiceContext
	logx.Logger
}

func NewCreateAgentLogic(ctx context.Context, svcCtx *svc.ServiceContext) *CreateAgentLogic {
	return &CreateAgentLogic{
		ctx:    ctx,
		svcCtx: svcCtx,
		Logger: logx.WithContext(ctx),
	}
}

func (l *CreateAgentLogic) CreateAgent(req *types.AgentInfo) (*types.BaseMsgResp, error) {
	organizationId := l.ctx.Value("organizationId").(uint64)

	agent, err := l.svcCtx.DB.Agent.Create().
		SetNotNilName(req.Name).
		SetNotNilRole(req.Role).
		//SetNotNilStatus(req.Status).
		SetOrganizationID(organizationId).
		SetNotNilBackground(req.Background).
		SetNotNilExamples(req.Examples).
		SetNotNilModel(req.Model).
		SetNotNilAPIKey(req.ApiKey).
		SetNotNilAPIBase(req.ApiBase).
		SetNotNilType(req.Type).
		Save(l.ctx)

	if err != nil {
		return nil, dberrorhandler.DefaultEntError(l.Logger, err, req)
	}

	/**
	1. 创建知识库
	2. 创建默认集合
	3. 保存到Agent信息里
	*/
	var datasetReq fastgpt.DatasetReq
	datasetReq.Name = *req.Name
	datasetReq.Intro = *req.Name
	datasetReq.AgentModel = "gpt-3.5-turbo"
	datasetReq.VectorModel = "text-embedding-ada-002"
	datasetResp, err := fastgpt.CreateDataset(&datasetReq)
	if err != nil {
		return nil, errorx.NewInvalidArgumentError("fastgpt create dataset failed " + err.Error())
	}
	if datasetResp.Code != 200 {
		return nil, errors.New(datasetResp.Message)
	}

	var collectionReq fastgpt.CreateCollectionReq
	collectionReq.DatasetId = datasetResp.Data
	collectionReq.Name = "手动录入"
	collectionReq.Type = "virtual" // 暂时改动
	collectionResp, err := fastgpt.CreateEmptyCollection(&collectionReq)
	if err != nil {
		return nil, errors.New("create dataset failed")
	}
	if collectionResp.Code != 200 {
		return nil, errorx.NewInvalidArgumentError("fastgpt create collection failed " + err.Error())
	}

	_, err = l.svcCtx.DB.Agent.UpdateOneID(agent.ID).
		SetNotNilDatasetID(&datasetResp.Data).
		SetNotNilCollectionID(&collectionResp.Data).
		Save(l.ctx)

	if err != nil {
		return nil, errors.New("update dataset and collection failed")
	}

	return &types.BaseMsgResp{Msg: errormsg.CreateSuccess}, nil
}