72 lines
1.8 KiB
Go
72 lines
1.8 KiB
Go
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) Detail(ctx context.Context, req *dto.GetDocumentDetailReq) (res *dto.GetDocumentDetailRes, err error) {
|
|
doc, err := service.DocumentService.Detail(ctx, req.Id)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &dto.GetDocumentDetailRes{Document: doc}, 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) {
|
|
n, err := service.DocumentService.Reembed(ctx, req.Id)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &dto.ReembedDocumentRes{Count: n}, nil
|
|
}
|