rag-local/kb/service/document_service.go

package service

import (
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"

	"rag-local/common"
	"rag-local/kb/consts"
	"rag-local/kb/dao"
	"rag-local/kb/model/domain"
	"rag-local/kb/model/entity"

	"github.com/gogf/gf/v2/errors/gerror"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/net/ghttp"
)

var DocumentService = &documentService{}

type documentService struct{}

// Upload 保存上传文件到 workspace/{datasetId}/{yyyymmdd}/{uuid}.ext，落库并提交解析任务
func (s *documentService) Upload(ctx context.Context, datasetId int64, file *ghttp.UploadFile) (*entity.Document, error) {
	if file == nil {
		return nil, gerror.New("请选择文件")
	}
	f, err := file.Open()
	if err != nil {
		return nil, err
	}
	defer func() { _ = f.Close() }()
	data, err := io.ReadAll(f)
	if err != nil {
		return nil, err
	}
	filename := file.Filename
	dataset, err := dao.Dataset.GetOne(ctx, datasetId)
	if err != nil {
		return nil, err
	}
	if dataset == nil {
		return nil, gerror.New("数据集不存在")
	}
	ext := strings.TrimPrefix(strings.ToLower(filepath.Ext(filename)), ".")
	if !isSupportedExt(ext) {
		return nil, gerror.Newf("不支持的文档类型: %s（支持 %s）", ext, strings.Join(common.SupportedExts(), "/"))
	}
	if len(data) == 0 {
		return nil, gerror.New("文件内容为空")
	}

	relDir := filepath.Join(fmt.Sprintf("%d", datasetId), time.Now().Format("20060102"))
	relPath := filepath.Join(relDir, fmt.Sprintf("%s.%s", common.RandomToken(16), ext))
	absPath := filepath.Join("workspace", relPath)
	if err := os.MkdirAll(filepath.Dir(absPath), 0o755); err != nil {
		return nil, err
	}
	if err := os.WriteFile(absPath, data, 0o644); err != nil {
		return nil, err
	}

	docId, err := dao.Document.Insert(ctx, &entity.Document{
		DatasetId: datasetId,
		Filename:  filename,
		FilePath:  relPath,
		FileSize:  int64(len(data)),
		FileType:  ext,
		Status:    consts.DocumentStatusPending,
	})
	if err != nil {
		_ = os.Remove(absPath)
		return nil, err
	}
	if _, err := dao.ParseTask.Insert(ctx, docId, datasetId, consts.TaskTypeParse); err != nil {
		_ = os.Remove(absPath)
		return nil, err
	}
	return dao.Document.GetOne(ctx, docId)
}

func isSupportedExt(ext string) bool {
	for _, e := range common.SupportedExts() {
		if e == ext {
			return true
		}
	}
	return false
}

// List 文档列表
func (s *documentService) List(ctx context.Context, datasetId int64, page, pageSize int) ([]*entity.Document, int, error) {
	return dao.Document.List(ctx, datasetId, page, pageSize)
}

// Detail 文档详情（含原始全文）。存量文档 content 为空时，现场从源文件解析并回填落库
func (s *documentService) Detail(ctx context.Context, id int64) (*entity.Document, error) {
	doc, err := dao.Document.GetOne(ctx, id)
	if err != nil {
		return nil, err
	}
	if doc == nil {
		return nil, gerror.New("文档不存在")
	}
	if doc.Content == "" && doc.FilePath != "" {
		text, err := common.ParseFile(filepath.Join("workspace", doc.FilePath))
		if err != nil {
			return nil, gerror.Wrap(err, "解析源文件失败")
		}
		if err := dao.Document.UpdateFields(ctx, id, g.Map{"content": text}); err != nil {
			return nil, err
		}
		doc.Content = text
	}
	return doc, nil
}

// Reembed 文档全部分块重新向量化（分块文本不变，仅重算向量），返回处理的分块数
func (s *documentService) Reembed(ctx context.Context, id int64) (int, error) {
	doc, err := dao.Document.GetOne(ctx, id)
	if err != nil {
		return 0, err
	}
	if doc == nil {
		return 0, gerror.New("文档不存在")
	}
	cfgId, err := dao.Dataset.GetEmbeddingCfgId(ctx, doc.DatasetId)
	if err != nil {
		return 0, err
	}
	if cfgId <= 0 {
		return 0, gerror.New("数据集未绑定向量模型，无法向量化")
	}
	em, err := BuildEmbedder(ctx, cfgId)
	if err != nil {
		return 0, err
	}
	chunks, _, err := dao.Chunk.ListByDocument(ctx, id, 1, 100000)
	if err != nil {
		return 0, err
	}
	if len(chunks) == 0 {
		return 0, gerror.New("该文档无分块（解析失败或分块已删除），无法重新向量化，请先重新解析")
	}
	for start := 0; start < len(chunks); start += consts.EmbedBatchSize {
		end := min(start+consts.EmbedBatchSize, len(chunks))
		texts := make([]string, 0, end-start)
		for _, c := range chunks[start:end] {
			texts = append(texts, c.Content)
		}
		vecs, err := em.EmbedStrings(ctx, texts)
		if err != nil {
			return 0, gerror.Wrap(err, "向量化失败")
		}
		for j, c := range chunks[start:end] {
			if err := dao.Chunk.UpdateVec(ctx, c.Id, domain.VecJsonF64(vecs[j])); err != nil {
				return 0, err
			}
		}
	}
	return len(chunks), nil
}

// Delete 删除文档：先删文件与索引数据，再删记录
func (s *documentService) Delete(ctx context.Context, id int64) error {
	doc, err := dao.Document.GetOne(ctx, id)
	if err != nil {
		return err
	}
	if doc == nil {
		return nil
	}
	// 先取分块 id 清理知识图谱数据（分块删除后无法再映射）
	chunks, _, err := dao.Chunk.ListByDocument(ctx, id, 1, 100000)
	if err != nil {
		return err
	}
	chunkIds := make([]int64, 0, len(chunks))
	for _, c := range chunks {
		chunkIds = append(chunkIds, c.Id)
	}
	if err := dao.KgRelation.DeleteByChunkIds(ctx, chunkIds); err != nil {
		return err
	}
	if err := dao.KgEntity.DeleteByChunkIds(ctx, chunkIds); err != nil {
		return err
	}
	if err := ChunkService.DeleteByDocument(ctx, id); err != nil {
		return err
	}
	if err := dao.ParseTask.DeleteByDocument(ctx, id); err != nil {
		return err
	}
	if doc.FilePath != "" {
		_ = os.Remove(filepath.Join("workspace", doc.FilePath))
	}
	return dao.Document.Delete(ctx, id)
}