rag/internal/logic/document.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, err := orm.DocumentBlock.Query().Where(
		documentblock.DocumentID(int64(documentId)),
	).All(ctx)

	return blocks, 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)
}