xinmi/internal/application/service/resumable_upload.go

package service

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net/textproto"
	"os"
	"path/filepath"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/Tencent/WeKnora/internal/accesscontrol"
	"github.com/Tencent/WeKnora/internal/accesspolicy"
	"github.com/Tencent/WeKnora/internal/application/repository"
	"github.com/Tencent/WeKnora/internal/resumable"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/google/uuid"
	"github.com/hibiken/asynq"
	"gorm.io/gorm"
)

const uploadTextPartBytes = 512 << 10
const uploadBinaryParseBytes = 32 << 20

type resumableKnowledgeIDKey struct{}

func resumableDocumentID(row *resumable.Upload) string {
	return uuid.NewSHA1(uuid.NameSpaceOID, []byte(fmt.Sprintf("resumable:%d:%s:%s:%d", row.TenantID, row.TargetID, row.ID, row.IngestOffset))).String()
}

type resumableKnowledgeQueue interface {
	EnsureResumableKnowledgeQueued(context.Context, *types.Knowledge, []string) (*types.Knowledge, error)
}

// EnsureResumableKnowledgeQueued repairs a crash between the native document
// INSERT and queue delivery. The same task ID fences duplicate deliveries;
// this does not delete chunks or reparse a document already being processed.
func (s *knowledgeService) EnsureResumableKnowledgeQueued(ctx context.Context, doc *types.Knowledge, tagIDs []string) (*types.Knowledge, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if doc.ParseStatus == "completed" {
		return doc, nil
	}
	if doc.ParseStatus != "pending" && doc.ParseStatus != "processing" && doc.ParseStatus != "failed" {
		return doc, nil
	}
	if active, ok := s.task.(interface{ HasActiveKnowledgeTask(string) bool }); ok && active.HasActiveKnowledgeTask(doc.ID) {
		return doc, nil
	}
	if s.taskInspector != nil {
		queued, err := s.taskInspector.HasQueuedTasksForKnowledge(ctx, doc.ID)
		if err != nil {
			return nil, err
		}
		if queued {
			return doc, nil
		}
	}
	if err := s.setAndAttachKnowledgeTags(ctx, doc.TenantID, doc.KnowledgeBaseID, doc, tagIDs); err != nil {
		return nil, err
	}
	if doc.ParseStatus == "failed" {
		return s.ReparseKnowledge(ctx, doc.ID, nil)
	}
	kb, err := s.kbService.GetKnowledgeBaseByID(ctx, doc.KnowledgeBaseID)
	if err != nil {
		return nil, err
	}
	overrides, err := doc.ProcessOverrides()
	if err != nil {
		return nil, err
	}
	eff := ResolveProcessConfig(kb, overrides)
	questionCount := eff.QuestionGenerationConfig.QuestionCount
	if questionCount <= 0 {
		questionCount = 3
	}
	payload, err := json.Marshal(types.DocumentProcessPayload{TenantID: doc.TenantID, KnowledgeID: doc.ID, KnowledgeBaseID: doc.KnowledgeBaseID, FilePath: doc.FilePath, FileName: doc.FileName, FileType: doc.FileType, EnableMultimodel: eff.EnableMultimodel, EnableQuestionGeneration: eff.QuestionGenerationConfig.Enabled, QuestionCount: questionCount, Language: types.LanguageFromContextOrDefault(ctx)})
	if err != nil {
		return nil, err
	}
	_, err = s.task.Enqueue(asynq.NewTask(types.TypeDocumentProcess, payload), documentProcessTaskOptions(s.config, asynq.TaskID("resumable:"+doc.ID), asynq.Retention(24*time.Hour))...)
	if errors.Is(err, asynq.ErrTaskIDConflict) || errors.Is(err, asynq.ErrDuplicateTask) {
		err = nil
	}
	return doc, err
}

type uploadOptions struct {
	FileName      string                           `json:"fileName"`
	TagIDs        []string                         `json:"tag_ids"`
	ProcessConfig *types.KnowledgeProcessOverrides `json:"process_config"`
}
type uploadResult struct {
	KnowledgeIDs []string                  `json:"knowledge_ids,omitempty"`
	Source       *types.DigitalHumanSource `json:"source,omitempty"`
	Selection    string                    `json:"selection,omitempty"`
}

func NewResumableUploadService(db *gorm.DB, bases interfaces.KnowledgeBaseService, knowledge interfaces.KnowledgeService, humans *DigitalHumanService, access *accesscontrol.Service) (*resumable.Service, error) {
	authorize := func(ctx context.Context, kind, id string) error {
		tenant, _ := types.TenantIDFromContext(ctx)
		owner, _ := types.UserIDFromContext(ctx)
		if !types.TenantRoleFromContext(ctx).HasPermission(types.TenantRoleContributor) && !types.IsSystemAdminFromContext(ctx) {
			return resumable.ErrForbidden
		}
		if kind == "digital-human" {
			if !accesspolicy.Allowed(ctx, "digital-humans.distill") || accesspolicy.Hidden(ctx, "digital-humans", id) {
				return resumable.ErrForbidden
			}
			if e := humans.CheckOwner(ctx, id); e != nil {
				return resumable.ErrForbidden
			}
			return nil
		}
		if !accesspolicy.Allowed(ctx, "knowledge.write") || accesspolicy.Hidden(ctx, "knowledge", id) {
			return resumable.ErrForbidden
		}
		kb, e := bases.GetKnowledgeBaseByID(ctx, id)
		if e != nil {
			return resumable.ErrForbidden
		}
		if kb.TenantID != tenant || (kb.CreatorID != owner && !types.TenantRoleFromContext(ctx).HasPermission(types.TenantRoleAdmin) && !types.IsSystemAdminFromContext(ctx)) {
			return resumable.ErrForbidden
		}
		if kb.Type == types.KnowledgeBaseTypeFAQ {
			return resumable.UserError("FAQ 知识库请使用 FAQ 导入")
		}
		return nil
	}
	workerContext := func(ctx context.Context, row *resumable.Upload) (context.Context, error) {
		var user types.User
		if e := db.WithContext(ctx).Where("id=? AND is_active=true", row.OwnerID).Take(&user).Error; e != nil {
			return ctx, resumable.ErrForbidden
		}
		var tenant types.Tenant
		if e := db.WithContext(ctx).Where("id=?", row.TenantID).Take(&tenant).Error; e != nil {
			return ctx, e
		}
		var member types.TenantMember
		e := db.WithContext(ctx).Where("tenant_id=? AND user_id=? AND status=?", row.TenantID, row.OwnerID, types.TenantMemberStatusActive).Take(&member).Error
		if e != nil && !user.IsSystemAdmin {
			return ctx, resumable.ErrForbidden
		}
		ctx = context.WithValue(ctx, types.TenantIDContextKey, row.TenantID)
		ctx = context.WithValue(ctx, types.TenantInfoContextKey, &tenant)
		ctx = context.WithValue(ctx, types.UserIDContextKey, row.OwnerID)
		ctx = context.WithValue(ctx, types.TenantRoleContextKey, member.Role)
		ctx = context.WithValue(ctx, types.SystemAdminContextKey, user.IsSystemAdmin)
		snapshot, e := access.Snapshot(ctx)
		if e != nil {
			return ctx, e
		}
		return accesspolicy.WithSnapshot(ctx, snapshot), nil
	}
	ingest := func(ctx context.Context, row *resumable.Upload, path string) (int64, types.JSON, bool, error) {
		if row.Kind == "digital-human" {
			source, selection, e := humans.readResumableMaterial(ctx, row.Name, path, row.Size)
			if e != nil {
				return 0, nil, false, e
			}
			result, e := json.Marshal(uploadResult{Source: &source, Selection: selection})
			return row.Size, result, true, e
		}
		return ingestResumableKnowledge(ctx, knowledge, row, path)
	}
	uploads, e := resumable.NewService(db, os.Getenv("RESUMABLE_UPLOAD_DIR"), authorize, workerContext, ingest)
	if e != nil {
		return nil, e
	}
	humans.uploadSources = func(ctx context.Context, humanID string, ids []string) ([]types.DigitalHumanSource, error) {
		if len(ids) > DigitalHumanFileCountLimit {
			return nil, digitalHumanInputError("每次最多选择99份材料")
		}
		sources := []types.DigitalHumanSource{}
		seen := map[string]bool{}
		for _, id := range ids {
			if seen[id] {
				continue
			}
			seen[id] = true
			row, e := uploads.Get(ctx, id)
			if e != nil || row.Kind != "digital-human" || row.TargetID != humanID || row.State != "completed" || row.ExpiresAt.Before(time.Now()) {
				return nil, digitalHumanInputError("材料尚未解析完成、已过期或不属于此数字人")
			}
			var result uploadResult
			if json.Unmarshal(row.Result, &result) != nil || result.Source == nil {
				return nil, digitalHumanInputError("材料解析结果无效")
			}
			if result.Selection != "full" {
				return nil, digitalHumanInputError("材料“" + row.Name + "”是旧版原文节选，请重新上传以完整蒸馏")
			}
			sources = append(sources, *result.Source)
		}
		return sources, nil
	}
	return uploads, nil
}

func textUpload(name string) bool {
	switch strings.ToLower(filepath.Ext(name)) {
	case ".txt", ".md", ".markdown", ".srt", ".vtt", ".csv", ".json", ".html", ".htm":
		return true
	}
	return false
}

// readUploadPart never retains more than one bounded native parser request.
// Large UTF-8 text is split without breaking a code point; structured text is
// represented as plain text parts so partial JSON/CSV does not enter a parser.
func readUploadPart(path, name string, offset, size int64) ([]byte, int64, string, error) {
	limit := int64(uploadBinaryParseBytes)
	if textUpload(name) {
		limit = uploadTextPartBytes
	} else if size > limit {
		return nil, offset, name, resumable.UserError("文件已完成上传；此格式的原生解析上限为32 MiB，请拆分文档或转换为 UTF-8 TXT 后导入")
	}
	f, e := os.Open(path)
	if e != nil {
		return nil, offset, name, e
	}
	defer f.Close()
	if _, e = f.Seek(offset, io.SeekStart); e != nil {
		return nil, offset, name, e
	}
	data, e := io.ReadAll(io.LimitReader(f, min(limit, size-offset)))
	if e != nil {
		return nil, offset, name, e
	}
	end := offset + int64(len(data))
	if textUpload(name) {
		if end < size {
			for trim := 0; trim < utf8.UTFMax && len(data) > 0 && !utf8.Valid(data); trim++ {
				data = data[:len(data)-1]
				end--
			}
		}
		if !utf8.Valid(data) {
			return nil, offset, name, resumable.UserError("文本必须采用 UTF-8 编码")
		}
		if size > limit {
			name = strings.TrimSuffix(name, filepath.Ext(name)) + fmt.Sprintf(".part-%012d.txt", offset)
		}
	}
	if len(data) == 0 {
		return nil, offset, name, resumable.UserError("没有可读取的材料内容")
	}
	return data, end, name, nil
}

func ingestResumableKnowledge(ctx context.Context, knowledge interfaces.KnowledgeService, row *resumable.Upload, path string) (int64, types.JSON, bool, error) {
	var options uploadOptions
	if e := json.Unmarshal(row.Options, &options); e != nil {
		return row.IngestOffset, nil, false, e
	}
	data, end, name, e := readUploadPart(path, row.Name, row.IngestOffset, row.Size)
	if e != nil {
		return row.IngestOffset, nil, false, e
	}
	documentID := resumableDocumentID(row)
	ctx = context.WithValue(ctx, resumableKnowledgeIDKey{}, documentID)
	doc, lookupErr := knowledge.GetKnowledgeByID(ctx, documentID)
	if lookupErr != nil && !errors.Is(lookupErr, repository.ErrKnowledgeNotFound) && !errors.Is(lookupErr, gorm.ErrRecordNotFound) {
		return row.IngestOffset, nil, false, lookupErr
	}
	if lookupErr == nil && doc != nil {
		if doc.TenantID != row.TenantID || doc.KnowledgeBaseID != row.TargetID {
			return row.IngestOffset, nil, false, resumable.ErrForbidden
		}
		if queue, ok := knowledge.(resumableKnowledgeQueue); ok {
			doc, e = queue.EnsureResumableKnowledgeQueued(ctx, doc, options.TagIDs)
			if e != nil {
				return row.IngestOffset, nil, false, e
			}
		}
	}
	var body bytes.Buffer
	writer := multipart.NewWriter(&body)
	header := textproto.MIMEHeader{}
	header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="file"; filename="%s"`, strings.NewReplacer("\\", "_", "\"", "_").Replace(name)))
	part, e := writer.CreatePart(header)
	if e != nil {
		return row.IngestOffset, nil, false, e
	}
	_, e = part.Write(data)
	if e != nil {
		return row.IngestOffset, nil, false, e
	}
	if e = writer.Close(); e != nil {
		return row.IngestOffset, nil, false, e
	}
	form, e := multipart.NewReader(&body, writer.Boundary()).ReadForm(64 << 10)
	if e != nil {
		return row.IngestOffset, nil, false, e
	}
	defer form.RemoveAll()
	customName := options.FileName
	if row.Size > uploadTextPartBytes && textUpload(row.Name) {
		folder, _ := types.SplitKnowledgeRelativePath(customName)
		customName = name
		if folder != "" {
			customName = folder + "/" + name
		}
	}
	if doc == nil {
		doc, e = knowledge.CreateKnowledgeFromFile(ctx, row.TargetID, form.File["file"][0], map[string]string{"resumable_upload_id": row.ID, "resumable_upload_offset": fmt.Sprint(row.IngestOffset)}, nil, customName, options.TagIDs, "web", options.ProcessConfig)
	}
	var duplicate *types.DuplicateKnowledgeError
	if errors.As(e, &duplicate) {
		doc = duplicate.Knowledge
		e = nil
		if doc != nil && doc.ID == documentID {
			if queue, ok := knowledge.(resumableKnowledgeQueue); ok {
				doc, e = queue.EnsureResumableKnowledgeQueued(ctx, doc, options.TagIDs)
			}
		}
	}
	if e != nil {
		return row.IngestOffset, nil, false, e
	}
	if doc == nil {
		return row.IngestOffset, nil, false, fmt.Errorf("原生入库未返回文档")
	}
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	for doc.ParseStatus != "completed" {
		if doc.ParseStatus == "failed" || doc.ParseStatus == "cancelled" || doc.ParseStatus == "deleting" {
			return row.IngestOffset, nil, false, resumable.UserError("原生文档解析失败，请检查文档处理记录和解析器或模型连接后重试")
		}
		select {
		case <-ctx.Done():
			return row.IngestOffset, nil, false, ctx.Err()
		case <-ticker.C:
		}
		doc, e = knowledge.GetKnowledgeByID(ctx, doc.ID)
		if e != nil {
			return row.IngestOffset, nil, false, e
		}
	}
	var result uploadResult
	if len(row.Result) > 0 && json.Unmarshal(row.Result, &result) != nil {
		return row.IngestOffset, nil, false, fmt.Errorf("上传进度记录无效")
	}
	found := false
	for _, id := range result.KnowledgeIDs {
		if id == doc.ID {
			found = true
		}
	}
	if !found {
		result.KnowledgeIDs = append(result.KnowledgeIDs, doc.ID)
	}
	encoded, e := json.Marshal(result)
	return end, encoded, end == row.Size, e
}

func (s *DigitalHumanService) readResumableMaterial(ctx context.Context, name, path string, size int64) (types.DigitalHumanSource, string, error) {
	if !DigitalHumanUploadSupported(name) {
		return types.DigitalHumanSource{}, "", digitalHumanInputError("不支持的数字人材料格式")
	}
	if textUpload(name) {
		f, e := os.Open(path)
		if e != nil {
			return types.DigitalHumanSource{}, "", e
		}
		defer f.Close()
		source, selection, e := readFullMaterial(name, f, size)
		return source, selection, e
	}
	if size > DigitalHumanFileLimit {
		return types.DigitalHumanSource{}, "", digitalHumanInputError("文件已完成上传；此格式解析上限为20 MiB，请拆分文档或转换为 UTF-8 TXT")
	}
	data, e := os.ReadFile(path)
	if e != nil {
		return types.DigitalHumanSource{}, "", e
	}
	source, e := s.readDigitalHumanUpload(ctx, DigitalHumanUpload{Name: name, Data: data})
	if e != nil {
		return source, "", e
	}
	return source, "full", nil
}

// Refuse oversized material explicitly: an excerpt must never be presented as
// a complete distillation source. The byte bound also limits memory allocation.
func readFullMaterial(name string, reader io.ReaderAt, size int64) (types.DigitalHumanSource, string, error) {
	source := types.DigitalHumanSource{Name: name}
	if size < 0 || size > int64(digitalHumanSourceLimit*utf8.UTFMax+3) {
		return source, "", digitalHumanInputError("材料超过完整蒸馏的单份24万字上限，请拆分后上传：" + name)
	}
	data := make([]byte, size)
	if _, err := io.ReadFull(io.NewSectionReader(reader, 0, size), data); err != nil {
		return source, "", err
	}
	if !utf8.Valid(data) {
		return source, "", digitalHumanInputError("文本材料须采用 UTF-8 编码")
	}
	source.Text = strings.TrimSpace(strings.TrimPrefix(string(data), "\ufeff"))
	if source.Text == "" {
		return source, "", digitalHumanInputError("材料没有可读取的文字")
	}
	if utf8.RuneCountInString(source.Text) > digitalHumanSourceLimit {
		return source, "", digitalHumanInputError("材料超过完整蒸馏的单份24万字上限，请拆分后上传：" + name)
	}
	return source, "full", nil
}