rag/internal/logic/document.go
2024-07-23 01:55:16 +08:00

165 lines
3.9 KiB
Go

package logic
import (
"context"
"leafdev.top/leaf/rag/consts"
"leafdev.top/leaf/rag/ent"
"leafdev.top/leaf/rag/ent/document"
"leafdev.top/leaf/rag/ent/documentblock"
"leafdev.top/leaf/rag/models"
)
type DocumentLogic struct {
LibraryLogic LibraryLogic
}
func NewDocumentLogic() *DocumentLogic {
return &DocumentLogic{
LibraryLogic: *NewLibraryLogic(),
}
}
func (d DocumentLogic) ListDocument(ctx context.Context, libraryId int) ([]*ent.Document, error) {
library, err := d.LibraryLogic.ExistsByUser(ctx, libraryId)
if err != nil {
return nil, err
}
if !library {
return nil, consts.ErrLibraryNotFound
}
return orm.Document.Query().Where(
document.UserIDEQ(GetUserId(ctx)),
document.LibraryIDEQ(int64(libraryId)),
).All(ctx)
}
func (d DocumentLogic) CreateDocument(ctx context.Context, libraryId int, name string) (*ent.Document, error) {
// 检查 Library 是否存在
exists, err := d.LibraryLogic.ExistsByUser(ctx, libraryId)
if err != nil {
return nil, err
}
if !exists {
return nil, consts.ErrLibraryNotFound
}
if name == "" {
return nil, consts.ErrValidateNameMustBeNotEmpty
}
return orm.Document.Create().
SetName(name).
SetUserID(GetUserId(ctx)).
SetLibraryID(int64(libraryId)).
Save(ctx)
}
func (d DocumentLogic) ExistsByUser(ctx context.Context, libraryId int) (bool, error) {
return orm.Document.Query().Where(
document.UserIDEQ(GetUserId(ctx)),
document.LibraryIDEQ(int64(libraryId)),
).Exist(ctx)
}
func (d DocumentLogic) CreateDocumentBlock(ctx context.Context, documentId int, createDocumentModel *models.CreateDocumentBlock) (*ent.DocumentBlock, error) {
// 检查 Document 是否存在
exists, err := d.ExistsByUser(ctx, int(documentId))
if err != nil {
return nil, err
}
if !exists {
return nil, consts.ErrDocumentNotFound
}
return orm.DocumentBlock.Create().
SetOrder(createDocumentModel.Order).
SetDocumentID(int64(documentId)).
SetUserID(GetUserId(ctx)).
SetContent(createDocumentModel.Content).
Save(ctx)
}
func (d DocumentLogic) ListDocumentBlock(ctx context.Context, documentId int) ([]*ent.DocumentBlock, error) {
exists, err := d.ExistsByUser(ctx, documentId)
if err != nil {
return nil, err
}
if !exists {
return nil, consts.ErrDocumentNotFound
}
blocks := orm.DocumentBlock.Query().Where(
documentblock.DocumentID(int64(documentId)),
)
blocks = blocks.Order(ent.Asc(documentblock.FieldOrder))
allBlocks, err := blocks.All(ctx)
return allBlocks, err
}
func (d DocumentLogic) UpdateDocumentBlock(ctx context.Context, documentId int, blockId int, content string) (*ent.DocumentBlock, error) {
exists, err := d.ExistsByUser(ctx, documentId)
if err != nil {
return nil, err
}
if !exists {
return nil, consts.ErrDocumentNotFound
}
block, err := d.GetBlockById(ctx, blockId)
if err != nil {
return nil, err
}
if block.DocumentID != int64(documentId) {
return nil, consts.ErrDocumentNotFound
}
if block.UserID != GetUserId(ctx) {
return nil, consts.ErrNotYourResource
}
block.Content = content
return block.Update().Save(ctx)
}
func (d DocumentLogic) DeleteDocumentBlockById(ctx context.Context, documentId int, blockId int) error {
exists, err := d.ExistsByUser(ctx, documentId)
if err != nil {
return err
}
if !exists {
return consts.ErrDocumentNotFound
}
_, err = orm.DocumentBlock.Delete().Where(
documentblock.DocumentIDEQ(int64(documentId)),
documentblock.IDEQ(blockId),
).Exec(ctx)
return err
}
func (d DocumentLogic) GetBlockById(ctx context.Context, documentBlockId int) (*ent.DocumentBlock, error) {
return orm.DocumentBlock.Query().Where(documentblock.IDEQ(documentBlockId)).First(ctx)
}
func (d DocumentLogic) BlockExistsById(ctx context.Context, id int) (bool, error) {
return orm.Document.Query().Where(document.ID(id)).Exist(ctx)
}
func (d DocumentLogic) Exists(ctx context.Context, id int) (bool, error) {
return orm.Document.Query().Where(document.ID(id)).Exist(ctx)
}