49 lines
1.2 KiB
Go
49 lines
1.2 KiB
Go
package controller
|
|
|
|
import (
|
|
"context"
|
|
|
|
"rag-local/kb/model/dto"
|
|
"rag-local/kb/model/entity"
|
|
"rag-local/kb/service"
|
|
)
|
|
|
|
type dataset struct{}
|
|
|
|
var Dataset = &dataset{}
|
|
|
|
func (c *dataset) List(ctx context.Context, _ *dto.ListDatasetReq) (*dto.ListDatasetRes, error) {
|
|
list, err := service.DatasetService.List(ctx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &dto.ListDatasetRes{List: list}, nil
|
|
}
|
|
|
|
func (c *dataset) Save(ctx context.Context, req *dto.SaveDatasetReq) (*dto.SaveDatasetRes, error) {
|
|
id, err := service.DatasetService.Save(ctx, &entity.Dataset{
|
|
Id: req.Id,
|
|
Name: req.Name,
|
|
Description: req.Description,
|
|
EmbeddingCfgId: req.EmbeddingCfgId,
|
|
ChunkSize: req.ChunkSize,
|
|
ChunkOverlap: req.ChunkOverlap,
|
|
ReactRounds: req.ReactRounds,
|
|
VecTopK: req.VecTopK,
|
|
FtsTopK: req.FtsTopK,
|
|
RerankTopK: req.RerankTopK,
|
|
RecallTopK: req.RecallTopK,
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &dto.SaveDatasetRes{Id: id}, nil
|
|
}
|
|
|
|
func (c *dataset) Delete(ctx context.Context, req *dto.DeleteDatasetReq) (*dto.DeleteDatasetRes, error) {
|
|
if err := service.DatasetService.Delete(ctx, req.Id); err != nil {
|
|
return nil, err
|
|
}
|
|
return &dto.DeleteDatasetRes{}, nil
|
|
}
|