1

kb/controller/document_controller.go

@@ -0,0 +1,62 @@
+package controller
+
+import (
+	"context"
+	"io"
+
+	"rag-local/kb/model/dto"
+	"rag-local/kb/service"
+
+	"github.com/gogf/gf/v2/errors/gerror"
+)
+
+type document struct{}
+
+var Document = &document{}
+
+func (c *document) Upload(ctx context.Context, req *dto.UploadDocumentReq) (*dto.UploadDocumentRes, error) {
+	if req.File == nil {
+		return nil, gerror.New("请选择文件")
+	}
+	f, err := req.File.Open()
+	if err != nil {
+		return nil, err
+	}
+	defer func() { _ = f.Close() }()
+	data, err := io.ReadAll(f)
+	if err != nil {
+		return nil, err
+	}
+	doc, err := service.DocumentService.Upload(ctx, req.DatasetId, req.File.Filename, data)
+	if err != nil {
+		return nil, err
+	}
+	return &dto.UploadDocumentRes{Id: doc.Id}, nil
+}
+
+func (c *document) List(ctx context.Context, req *dto.ListDocumentReq) (*dto.ListDocumentRes, error) {
+	list, total, err := service.DocumentService.List(ctx, req.DatasetId, req.Page, req.PageSize)
+	if err != nil {
+		return nil, err
+	}
+	return &dto.ListDocumentRes{
+		List:     list,
+		Total:    total,
+		Page:     req.Page,
+		PageSize: req.PageSize,
+	}, nil
+}
+
+func (c *document) Delete(ctx context.Context, req *dto.DeleteDocumentReq) (*dto.DeleteDocumentRes, error) {
+	if err := service.DocumentService.Delete(ctx, req.Id); err != nil {
+		return nil, err
+	}
+	return &dto.DeleteDocumentRes{}, nil
+}
+
+func (c *document) Reembed(ctx context.Context, req *dto.ReembedDocumentReq) (*dto.ReembedDocumentRes, error) {
+	if err := service.DocumentService.Reembed(ctx, req.Id); err != nil {
+		return nil, err
+	}
+	return &dto.ReembedDocumentRes{}, nil
+}