1

kb/controller/dataset_controller.go

@@ -0,0 +1,42 @@
+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,
+		Status:         1,
+	})
+	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
+}