wechat-api/hook/fastgpt/dataset.go

package fastgpt

import (
	"errors"
	"time"
)

type DatasetResp struct {
	Code       int    `json:"code"`
	StatusText string `json:"statusText"`
	Message    string `json:"message"`
	Data       string `json:"data"`
}

type DatasetReq struct {
	ParentID    string `json:"parentId,optional"`
	Type        string `json:"type,optional"`
	Avatar      string `json:"avatar,optional"`
	Name        string `json:"name"`
	Intro       string `json:"intro,optional"`
	VectorModel string `json:"vectorModel"`
	AgentModel  string `json:"agentModel"`
}

type DatasetInfo struct {
	ID          string `json:"_id"`
	ParentID    string `json:"parentId,optional"`
	TeamID      string `json:"teamId,optional"`
	TmbID       string `json:"tmbId,optional"`
	Type        string `json:"type,optional"`
	Status      string `json:"status,optional"`
	Avatar      string `json:"avatar,optional"`
	Name        string `json:"name,optional"`
	VectorModel struct {
		Model            string `json:"model"`
		Name             string `json:"name"`
		Avatar           string `json:"avatar"`
		CharsPointsPrice uint64 `json:"charsPointsPrice"`
		DefaultToken     uint64 `json:"defaultToken"`
		MaxToken         uint64 `json:"maxToken"`
		Weight           uint64 `json:"weight"`
		DefaultConfig    struct {
		} `json:"defaultConfig"`
		DbConfig struct {
		} `json:"dbConfig"`
		QueryConfig struct {
		} `json:"queryConfig"`
	} `json:"vectorModel,optional"`
	AgentModel struct {
		Model                   string  `json:"model"`
		Name                    string  `json:"name"`
		Avatar                  string  `json:"avatar"`
		MaxContext              uint64  `json:"maxContext"`
		MaxResponse             uint64  `json:"maxResponse"`
		QuoteMaxToken           uint64  `json:"quoteMaxToken"`
		MaxTemperature          float64 `json:"maxTemperature"`
		CharsPointsPrice        uint64  `json:"charsPointsPrice"`
		Censor                  bool    `json:"censor"`
		Vision                  bool    `json:"vision"`
		DatasetProcess          bool    `json:"datasetProcess"`
		UsedInClassify          bool    `json:"usedInClassify"`
		UsedInExtractFields     bool    `json:"usedInExtractFields"`
		UsedInToolCall          bool    `json:"usedInToolCall"`
		UsedInQueryExtension    bool    `json:"usedInQueryExtension"`
		ToolChoice              bool    `json:"toolChoice"`
		FunctionCall            bool    `json:"functionCall"`
		CustomCQPrompt          string  `json:"customCQPrompt"`
		CustomExtractPrompt     string  `json:"customExtractPrompt"`
		DefaultSystemChatPrompt string  `json:"defaultSystemChatPrompt"`
		DefaultConfig           struct {
		} `json:"defaultConfig"`
	} `json:"agentModel,optional"`
	Intro      string `json:"intro,optional"`
	Permission struct {
		Value        int64 `json:"value"`
		IsOwner      bool  `json:"isOwner"`
		HasManagePer bool  `json:"hasManagePer"`
		HasWritePer  bool  `json:"hasWritePer"`
		HasReadPer   bool  `json:"hasReadPer"`
	} `json:"permission,optional"`
	DefaultPermission uint64    `json:"defaultPermission.optional"`
	UpdateTime        time.Time `json:"updateTime,optional"`
	V                 uint64    `json:"__v,optional"`
}

type DatasetList struct {
	Code       int           `json:"code"`
	StatusText string        `json:"statusText"`
	Message    string        `json:"message"`
	Data       []DatasetInfo `json:"data"`
}

type DatasetDetail struct {
	Code       int         `json:"code"`
	StatusText string      `json:"statusText"`
	Message    string      `json:"message"`
	Data       DatasetInfo `json:"data"`
}

// GetDatasetDetail 获取数据集详情
func GetDatasetDetail(id string) (dataset *DatasetDetail, err error) {
	resp, err := NewResty().
		R().
		SetResult(&dataset).
		SetQueryParam("id", id).
		Get("core/dataset/detail")

	if err != nil {
		return nil, err
	}

	if resp.IsError() {
		return nil, errors.New(resp.String())
	}
	return
}

// GetDatasetList 获取数据集列表
func GetDatasetList(parentId string) (datasetList *DatasetList, err error) {
	resp, err := NewResty().
		R().
		SetResult(&datasetList).
		SetBody(map[string]string{"parentId": parentId}).
		Post("core/dataset/list")

	if err != nil {
		return nil, err
	}

	if resp.IsError() {
		return nil, errors.New(resp.String())
	}

	return
}

// CreateDataset 创建数据集
func CreateDataset(data *DatasetReq) (dataset *DatasetResp, err error) {
	data.Type = "dataset"
	if data.Avatar == "" {
		data.Avatar = "/icon/logo.svg"
	}

	resp, err := NewResty().
		R().
		SetResult(&dataset).
		SetBody(*data).
		Post("core/dataset/create")

	if err != nil {
		return nil, err
	}

	if resp.IsError() {
		return nil, errors.New(resp.String())
	}

	return
}

// CreateDatasetFolder 创建数据集目录
func CreateDatasetFolder(data *DatasetReq) (*DatasetResp, error) {
	data.Type = "folder"
	data.Avatar = "/imgs/files/folder.svg"
	return CreateDataset(data)
}

// DeleteDataset 删除数据集
func DeleteDataset(id string) (dataset *DatasetResp, err error) {
	resp, err := NewResty().
		R().
		SetResult(&dataset).
		SetQueryParam("id", id).
		Delete("core/dataset/delete")

	if err != nil {
		return nil, err
	}

	if resp.IsError() {
		return nil, errors.New(resp.String())
	}

	return
}