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) }