wechat-api/hook/fastgpt/collection.go

package fastgpt

import (
	"encoding/json"
	"errors"
	"fmt"
	"time"
)

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

type CreateCollectionReq struct {
	DatasetId string      `json:"datasetId"`
	ParentId  string      `json:"parentId,optional"`
	Name      string      `json:"name"`
	Type      string      `json:"type"`
	Metadata  interface{} `json:"metadata,optional"`
}

type CreateTextCollectionReq struct {
	DatasetId     string      `json:"datasetId"`
	ParentId      string      `json:"parentId,optional"`
	Name          string      `json:"name"`
	Text          string      `json:"text"`
	Metadata      interface{} `json:"metadata,optional"`
	TrainingType  string      `json:"trainingType,optional"`
	ChunkSize     int         `json:"chunkSize,optional"`
	ChunkSplitter string      `json:"chunkSplitter,optional"`
	QaPrompt      string      `json:"qaPrompt,optional"`
}

type CreateTextCollectionResp struct {
	Code       int    `json:"code"`
	StatusText string `json:"statusText"`
	Message    string `json:"message"`
	Data       struct {
		CollectionID string `json:"collectionId"`
		Results      struct {
			InsertLen int   `json:"insertLen"`
			OverToken []any `json:"overToken"`
			Repeat    []any `json:"repeat"`
			Error     []any `json:"error"`
		} `json:"results"`
	} `json:"data"`
}

type CreateLinkCollectionReq struct {
	DatasetId string `json:"datasetId"`
	ParentId  string `json:"parentId,optional"`
	Name      string `json:"name"`
	Link      string `json:"link"`

	Metadata      interface{} `json:"metadata,optional"`
	TrainingType  string      `json:"trainingType,optional"`
	ChunkSize     int         `json:"chunkSize,optional"`
	ChunkSplitter string      `json:"chunkSplitter,optional"`
	QaPrompt      string      `json:"qaPrompt,optional"`
}

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

type CreateFileCollectionReq struct {
	DatasetId string `json:"datasetId"`
	ParentId  string `json:"parentId,optional"`

	Metadata      interface{} `json:"metadata,optional"`
	TrainingType  string      `json:"trainingType,optional"`
	ChunkSize     int         `json:"chunkSize,optional"`
	ChunkSplitter string      `json:"chunkSplitter,optional"`
	QaPrompt      string      `json:"qaPrompt,optional"`
}

type CreateFileCollectionResp struct {
	Code       int    `json:"code"`
	StatusText string `json:"statusText,optional"`
	Message    string `json:"message,optional"`
	Data       struct {
		CollectionID string `json:"collectionId,optional"`
		Results      struct {
			InsertLen int   `json:"insertLen"`
			OverToken []any `json:"overToken"`
			Repeat    []any `json:"repeat"`
			Error     []any `json:"error"`
		} `json:"results,optional"`
	} `json:"data,optional"`
}

type GetCollectionListReq struct {
	PageNum    int    `json:"pageNum,optional"`
	PageSize   int    `json:"pageSize,optional"`
	DatasetId  string `json:"datasetId,optional"`
	ParentId   string `json:"parentId,optional"`
	SearchText string `json:"searchText,optional"`
}

type CollectionInfo struct {
	ID        string      `json:"_id"`
	ParentID  interface{} `json:"parentId"`
	TeamID    string      `json:"teamId"`
	TmbID     string      `json:"tmbId"`
	DatasetID struct {
		ID          string      `json:"_id"`
		ParentID    interface{} `json:"parentId"`
		TeamID      string      `json:"teamId"`
		TmbID       string      `json:"tmbId"`
		Type        string      `json:"type"`
		Status      string      `json:"status"`
		Avatar      string      `json:"avatar"`
		Name        string      `json:"name"`
		VectorModel string      `json:"vectorModel"`
		AgentModel  string      `json:"agentModel"`
		Intro       string      `json:"intro"`
		Permission  string      `json:"permission"`
		UpdateTime  time.Time   `json:"updateTime"`
	} `json:"datasetId"`
	Type          string    `json:"type"`
	Name          string    `json:"name"`
	TrainingType  string    `json:"trainingType"`
	ChunkSize     int       `json:"chunkSize"`
	ChunkSplitter string    `json:"chunkSplitter"`
	QaPrompt      string    `json:"qaPrompt"`
	RawTextLength int       `json:"rawTextLength"`
	HashRawText   string    `json:"hashRawText"`
	CreateTime    time.Time `json:"createTime"`
	UpdateTime    time.Time `json:"updateTime"`
	CanWrite      bool      `json:"canWrite"`
	SourceName    string    `json:"sourceName"`
}

type CollectionListResp struct {
	Code       int    `json:"code"`
	StatusText string `json:"statusText"`
	Message    string `json:"message"`
	Data       struct {
		PageNum  int `json:"pageNum"`
		PageSize int `json:"pageSize"`
		Data     []struct {
			ID             string    `json:"_id"`
			ParentID       string    `json:"parentId"`
			TmbID          string    `json:"tmbId"`
			Type           string    `json:"type"`
			Name           string    `json:"name"`
			UpdateTime     time.Time `json:"updateTime"`
			DataAmount     int       `json:"dataAmount"`
			TrainingAmount int       `json:"trainingAmount"`
			CanWrite       bool      `json:"canWrite"`
			RawLink        string    `json:"rawLink,omitempty"`
		} `json:"data"`
		Total int `json:"total"`
	} `json:"data"`
}

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

type UpdateCollectionReq struct {
	ID       string `json:"id"`
	ParentId string `json:"parentId"`
	Name     string `json:"name"`
}

// GetCollectionDetail 获取合集详情
func GetCollectionDetail(id string) (collection *CollectionDetailResp, err error) {
	resp, err := NewResty().
		R().
		SetResult(&collection).
		SetQueryParam("id", id).
		Get("core/dataset/collection/detail")

	if err != nil {
		return nil, err
	}

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

	return
}

// GetCollectionList 获取合集列表
func GetCollectionList(params *GetCollectionListReq) (collectionList *CollectionListResp, err error) {
	resp, err := NewResty().
		R().
		SetResult(&collectionList).
		SetBody(params).
		Post("core/dataset/collection/list")

	if err != nil {
		return nil, err
	}

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

	return
}

// CreateEmptyCollection 创建空合集
func CreateEmptyCollection(data *CreateCollectionReq) (collection *CollectionResp, err error) {
	resp, err := NewResty().
		R().
		SetResult(&collection).
		SetBody(data).
		Post("core/dataset/collection/create")

	if err != nil {
		return nil, err
	}

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

	return
}

// CreateTextCollection 创建文本合集
func CreateTextCollection(data *CreateTextCollectionReq) (collection *CreateTextCollectionResp, err error) {
	resp, err := NewResty().
		R().
		SetResult(&collection).
		SetBody(*data).
		Post("core/dataset/collection/create/text")

	if err != nil {
		fmt.Printf(err.Error())
		return nil, err
	}

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

	return
}

// CreateLinkCollection 创建链接合集
func CreateLinkCollection(data *CreateLinkCollectionReq) (collection *CreateLinkCollectionResp, err error) {
	resp, err := NewResty().
		R().
		SetResult(&collection).
		SetBody(*data).
		Post("core/dataset/collection/create/link")

	if err != nil {
		fmt.Printf(err.Error())
		return nil, err
	}

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

	return
}

// CreateFileCollection 创建文件合集
func CreateFileCollection(data *map[string]string, filePath string) (collection *CreateFileCollectionResp, err error) {
	dataString, _ := json.Marshal(*data)
	resp, err := NewResty().
		R().
		SetResult(&collection).
		SetFormData(map[string]string{"data": string(dataString)}).
		SetFile("file", filePath).
		Post("core/dataset/collection/create/localFile")

	fmt.Printf("%v, %v", err, resp.String())
	if err != nil {
		return nil, err
	}

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

	return
}

// UpdateCollection 修改合集
func UpdateCollection(data *UpdateCollectionReq) (collection *CollectionResp, err error) {
	resp, err := NewResty().
		R().
		SetResult(&collection).
		SetBody(*data).
		Put("core/dataset/collection/update")

	if err != nil {
		return nil, err
	}

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

	return
}

// DeleteCollection 删除合集
func DeleteCollection(id string) (collection *CollectionResp, err error) {
	resp, err := NewResty().
		R().
		SetResult(&collection).
		SetQueryParam("id", id).
		Delete("core/dataset/collection/delete")

	if err != nil {
		return nil, err
	}

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

	return
}