wechat-api/hook/fastgpt/search.go

package fastgpt

import (
	"errors"
)

type SearchReq struct {
	DatasetID   string  `json:"datasetId"`           //知识库ID
	Text        string  `json:"text"`                //需要测试的文本
	Limit       uint64  `json:"limit"`               //最大 tokens 数量
	Similarity  float64 `json:"similarity,optional"` //最低相关度（0~1，可选）
	SearchMode  string  `json:"searchMode"`          //搜索模式：embedding | fullTextRecall | mixedRecall
	UsingReRank bool    `json:"usingReRank"`         //使用重排
}

type SearchResp struct {
	Code       int    `json:"code"`
	StatusText string `json:"statusText"`
	Message    string `json:"message"`
	Data       struct {
		List []struct {
			ID           string `json:"id"`
			DatasetID    string `json:"datasetId"`
			CollectionID string `json:"collectionId"`
			SourceName   string `json:"sourceName"`
			Q            string `json:"q"`
			A            string `json:"a"`
			ChunkIndex   int    `json:"chunkIndex"`
			Score        []struct {
				Type  string  `json:"type"`
				Value float64 `json:"value"`
				Index uint64  `json:"index"`
			} `json:"score"`
		} `json:"list,optional"`
		Duration              string `json:"duration"`
		SearchMode            string `json:"searchMode"`
		Limit                 uint64 `json:"limit"`
		Similarity            uint64 `json:"similarity"`
		UsingReRank           bool   `json:"usingReRank"`
		UsingSimilarityFilter bool   `json:"usingSimilarityFilter"`
	} `json:"data,optional"`
}

// SearchText 搜索数据
func SearchText(params *SearchReq) (dataList *SearchResp, err error) {
	resp, err := NewResty().
		R().
		SetResult(&dataList).
		SetBody(params).
		Post("core/dataset/searchTest")

	if err != nil {
		return nil, err
	}

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

	return
}