xinmi/internal/application/service/wiki_ingest_checkpoint.go

package service

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/Tencent/WeKnora/internal/agent"
	apprepo "github.com/Tencent/WeKnora/internal/application/repository"
	"github.com/Tencent/WeKnora/internal/logger"
	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// A slice deliberately leaves time for persistence, publishing and queue cleanup
// inside the outer 60 minute job. Only complete validated artifacts are stored.
const wikiCitationSlice = 16
const wikiPageSlice = 24
const wikiCheckpointSliceTimeout = 20 * time.Minute

const wikiRoutingMetadataVersion = "wiki-routing-metadata-v1"
const wikiRoutingMetadataContract = "candidate-routing:slug,type,name,aliases;description,details:audit-only;citation-evidence:original-chunks;taxonomy:identity-only;identity-faithfulness:not-asserted"

var errWikiCheckpointYield = errors.New("wiki work saved; continuation required")

func wikiCheckpointError(ctx context.Context, err error) error {
	if errors.Is(context.Cause(ctx), ErrWikiRecoveryBudgetStopped) {
		return context.Cause(ctx)
	}
	if err != nil && errors.Is(context.Cause(ctx), errWikiCheckpointYield) &&
		(errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)) {
		return errWikiCheckpointYield
	}
	return err
}

func wikiCheckpointMapError(ctx context.Context, summaryErr, citationErr error) error {
	summaryErr = wikiCheckpointError(ctx, summaryErr)
	citationErr = wikiCheckpointError(ctx, citationErr)
	if summaryErr != nil && !errors.Is(summaryErr, errWikiCheckpointYield) {
		return summaryErr
	}
	if citationErr != nil && !errors.Is(citationErr, errWikiCheckpointYield) {
		return citationErr
	}
	if summaryErr != nil || citationErr != nil {
		return errWikiCheckpointYield
	}
	return nil
}

type wikiCheckpointScope struct {
	TenantID        uint64
	KnowledgeBaseID string
	KnowledgeID     string
	Fingerprint     string
}

type wikiIngestCheckpoint struct {
	TenantID        uint64 `gorm:"primaryKey;autoIncrement:false"`
	KnowledgeBaseID string `gorm:"primaryKey;size:36"`
	KnowledgeID     string `gorm:"primaryKey;size:36"`
	Fingerprint     string `gorm:"primaryKey;size:64"`
	Stage           string `gorm:"primaryKey;size:32"`
	ItemKey         string `gorm:"primaryKey;size:64"`
	Payload         string `gorm:"type:text;not null"`
	UpdatedAt       time.Time
}

func (wikiIngestCheckpoint) TableName() string { return "wiki_ingest_checkpoints" }

func wikiCheckpointHash(value any) string {
	encoded, err := json.Marshal(value)
	if err != nil {
		panic(fmt.Sprintf("invalid checkpoint signature: %v", err))
	}
	sum := sha256.Sum256(encoded)
	return hex.EncodeToString(sum[:])
}

func wikiCheckpointChunks(chunks []*types.Chunk) []any {
	values := make([]any, 0, len(chunks))
	for _, c := range chunks {
		if c != nil {
			values = append(values, []any{c.ID, c.TenantID, c.KnowledgeBaseID, c.KnowledgeID,
				c.Content, c.ContentRevision, c.ChunkIndex, c.StartAt, c.EndAt, c.ChunkType,
				c.ParentChunkID, c.ImageInfo, c.IsEnabled, c.DeletedAt.Valid})
		}
	}
	return values
}

func (s *wikiIngestService) loadWikiCheckpoint(ctx context.Context, scope *wikiCheckpointScope, stage, key string, out any) (bool, error) {
	if s.checkpointDB == nil || scope == nil {
		return false, nil
	}
	var row wikiIngestCheckpoint
	err := s.checkpointDB.WithContext(ctx).Where("tenant_id = ? AND knowledge_base_id = ? AND knowledge_id = ? AND fingerprint = ? AND stage = ? AND item_key = ?", scope.TenantID, scope.KnowledgeBaseID, scope.KnowledgeID, scope.Fingerprint, stage, key).Take(&row).Error
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return false, nil
	}
	if err != nil {
		return false, fmt.Errorf("read wiki checkpoint: %w", err)
	}
	if err := json.Unmarshal([]byte(row.Payload), out); err != nil {
		return false, fmt.Errorf("invalid wiki checkpoint: %w", err)
	}
	return true, nil
}

func (s *wikiIngestService) saveWikiCheckpoint(ctx context.Context, scope *wikiCheckpointScope, stage, key string, value any) error {
	if s.checkpointDB == nil || scope == nil {
		return nil
	}
	encoded, err := json.Marshal(value)
	if err != nil {
		return err
	}
	row := wikiIngestCheckpoint{TenantID: scope.TenantID, KnowledgeBaseID: scope.KnowledgeBaseID, KnowledgeID: scope.KnowledgeID, Fingerprint: scope.Fingerprint, Stage: stage, ItemKey: key, Payload: string(encoded), UpdatedAt: time.Now()}
	return s.checkpointDB.WithContext(ctx).Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "tenant_id"}, {Name: "knowledge_base_id"}, {Name: "knowledge_id"}, {Name: "fingerprint"}, {Name: "stage"}, {Name: "item_key"}}, DoUpdates: clause.AssignmentColumns([]string{"payload", "updated_at"})}).Create(&row).Error
}

func (s *wikiIngestService) newWikiCheckpointScope(ctx context.Context, model chat.Chat, payload WikiIngestPayload, op WikiPendingOp, chunks []*types.Chunk, content, title, language string, batch *WikiBatchContext) (*wikiCheckpointScope, error) {
	if s.checkpointDB == nil {
		return nil, nil
	}
	if payload.TenantID == 0 || payload.KnowledgeBaseID == "" || op.KnowledgeID == "" {
		return nil, errors.New("wiki checkpoint scope is incomplete")
	}
	ordered := append([]*types.Chunk(nil), chunks...)
	sort.SliceStable(ordered, func(i, j int) bool {
		if ordered[i] == nil {
			return ordered[j] != nil
		}
		if ordered[j] == nil {
			return false
		}
		return ordered[i].ID < ordered[j].ID
	})
	for _, chunk := range ordered {
		if chunk != nil && (chunk.TenantID != payload.TenantID || chunk.KnowledgeBaseID != payload.KnowledgeBaseID || chunk.KnowledgeID != op.KnowledgeID) {
			return nil, errors.New("wiki checkpoint source scope mismatch")
		}
	}
	var config any = struct{ ID, Name string }{model.GetModelID(), model.GetModelName()}
	if s.modelService != nil {
		m, err := s.modelService.GetModelByID(ctx, model.GetModelID())
		if err != nil {
			return nil, fmt.Errorf("load checkpoint model configuration: %w", err)
		}
		if m == nil {
			return nil, errors.New("checkpoint model configuration missing")
		}
		// Secrets participate only in the irreversible signature, never in rows/logs.
		config = []any{m.ID, m.TenantID, m.Name, m.Type, m.Source, m.Parameters}
	}
	// A retry releases the same pending row; a new explicit ingest operation
	// has a different row ID and must reconsider prior page contributions.
	fingerprint := wikiCheckpointHash([]any{"wiki-checkpoint-v2", wikiFullReadVersion, wikiSourceIdentityVersion, wikiCitationEvidenceVersion, wikiCitationMergeVersion, wikiRoutingMetadataVersion, wikiRoutingMetadataContract, wikiCandidateCatalogVersion, wikiCandidateQualityVersion, wikiCandidateRepairVersion, wikiCandidateRepairInstruction, wikiCandidatePartitionVersion, wikiCandidateGenerationInstruction, wikiCandidateProactiveSplitRunes, wikiReadSectionPrompt, wikiReadMergePrompt, maxContentForWiki, wikiReadCandidateRunes, wikiReadSummaryRunes, wikiPageInputRunes, op.dbID, wikiCheckpointChunks(ordered), content, title, language, config, batch.ExtractionGranularity, batch.ExtractionInstructions, batch.ContentInstructions, agent.WikiCandidateSlugPrompt, agent.WikiGranularityGuidance(string(batch.ExtractionGranularity)), agent.WikiChunkCitationPrompt, wikiCitationDeclarationRepairPrompt, agent.WikiSummaryPrompt, agent.WikiDeduplicationPrompt, agent.WikiPageModifyUserPrompt, agent.WikiPageModifySystemPrompt, agent.WikiTaxonomyPlanPrompt})
	return &wikiCheckpointScope{payload.TenantID, payload.KnowledgeBaseID, op.KnowledgeID, fingerprint}, nil
}

func validateWikiCandidateCheckpoint(result combinedExtraction, source string) error {
	if err := validateWikiSourceIdentities(result, source); err != nil {
		return fmt.Errorf("invalid persisted wiki candidate: %w", err)
	}
	return nil
}

// pruneAndValidateWikiCandidates keeps a document moving when two model-written
// names claim one slug, and still refuses a set that cannot be published at all.
func pruneAndValidateWikiCandidates(ctx context.Context, result combinedExtraction, source string) (combinedExtraction, error) {
	pruned, dropped, err := pruneCandidateSetCollisions(result)
	if err != nil {
		return combinedExtraction{}, err
	}
	if dropped > 0 {
		logger.Infof(ctx, "wiki candidate set pruning: policy=%s slug_conflicts=%d kept=%d",
			wikiSourceIdentityVersion, pruned.CandidatePruning.SlugConflicts, len(pruned.Entities)+len(pruned.Concepts))
	}
	if err := validateWikiCandidateCheckpoint(pruned, source); err != nil {
		return combinedExtraction{}, err
	}
	return pruned, nil
}

// pruneCandidateSetCollisions applies the candidate quality policy to a complete
// candidate set before it is validated or persisted: two model-written names that
// claim one routing slug would otherwise abort a document that has already read
// its whole source. The first identity keeps the slug, later claimants are dropped
// and counted, and a set that prunes to nothing still fails closed.
func pruneCandidateSetCollisions(result combinedExtraction) (combinedExtraction, int, error) {
	pruned, dropped := pruneIdentityCollisions(result)
	if dropped == 0 {
		return result, 0, nil
	}
	if len(pruned.Entities)+len(pruned.Concepts) == 0 {
		return combinedExtraction{}, dropped, wikiIdentityFailure("wiki_identity_slug_conflict", types.WikiPageTypeEntity, -1, nil)
	}
	pruned.CandidatePruning.SlugConflicts += dropped
	return pruned, dropped, nil
}

func validateWikiFreshCandidateCheckpoint(result combinedExtraction, source string) error {
	if err := validateWikiFreshSourceIdentities(result, source); err != nil {
		return fmt.Errorf("invalid persisted wiki candidate: %w", err)
	}
	return nil
}

// Field paths identify invalid model output without logging source text or
// candidate names. Use the same checks before deduplication and on cache reads.
func validateWikiCandidates(result combinedExtraction) error {
	if validation := wikiCandidateFieldValidation(result); validation != nil {
		group, field, _ := strings.Cut(validation.Field, ".")
		return fmt.Errorf("%s[%d].%s: %s", group, *validation.Index, field, validation.Code)
	}
	return nil
}

func wikiCandidateSlugSet(candidates string) map[string]bool {
	allowed := make(map[string]bool)
	for _, line := range strings.Split(candidates, "\n") {
		slug, _, _ := strings.Cut(strings.TrimPrefix(strings.TrimSpace(line), "- slug: "), ",")
		if strings.HasPrefix(slug, "entity/") || strings.HasPrefix(slug, "concept/") {
			allowed[slug] = true
		}
	}
	return allowed
}

func validateWikiCitationCheckpoint(result *citationBatchResult, batch chunkBatch, candidates map[string]bool) error {
	if err := validateWikiCitationResult(result, batch, candidates); err != nil {
		return fmt.Errorf("invalid persisted citation: %w", err)
	}
	return nil
}

// Unknown namespaced slugs may be an omitted declaration, not a missing
// source. Fresh output has a narrowly bounded declaration-only repair path;
// persisted results still fail closed. Do not include model text in errors.
type wikiUnknownCitationSlugsError struct{ slugs []string }

func (e *wikiUnknownCitationSlugsError) Error() string {
	return fmt.Sprintf("citation has %d undeclared slug(s)", len(e.slugs))
}

func validateWikiCitationResult(result *citationBatchResult, batch chunkBatch, candidates map[string]bool) error {
	if result == nil || result.Citations == nil {
		return errors.New("invalid citation schema")
	}
	allowed := make(map[string]bool, len(batch.chunks))
	for _, chunk := range batch.chunks {
		if chunk != nil {
			allowed[chunk.ID] = true
		}
	}
	validateIDs := func(ids []string) error {
		for _, id := range ids {
			if !allowed[id] {
				return errors.New("citation references a different source chunk")
			}
		}
		return nil
	}
	allowedSlugs := make(map[string]bool, len(candidates)+len(result.NewSlugs))
	for slug := range candidates {
		allowedSlugs[slug] = true
	}
	for index, item := range result.NewSlugs {
		if strings.TrimSpace(item.Name) == "" || (item.Type != "entity" && item.Type != "concept") ||
			!strings.HasPrefix(item.Slug, item.Type+"/") || strings.TrimPrefix(item.Slug, item.Type+"/") == "" {
			return errors.New("citation has invalid new slug")
		}
		allowedSlugs[item.Slug] = true
		if err := validateIDs(item.SourceChunks); err != nil {
			return err
		}
		if len(item.SourceChunks) == 0 && len(result.Citations[item.Slug]) == 0 {
			return errors.New("citation new slug has no source evidence")
		}
		ids := append(append([]string(nil), item.SourceChunks...), result.Citations[item.Slug]...)
		if err := validateIDs(ids); err != nil {
			return err
		}
		var sources []string
		for _, chunk := range batch.chunks {
			if chunk == nil {
				continue
			}
			for _, id := range ids {
				if chunk.ID == id {
					sources = append(sources, chunk.Content)
					break
				}
			}
		}
		if err := validateWikiSourceItemIdentity(extractedItem{Name: item.Name, Aliases: item.Aliases}, item.Type, index, sources); err != nil {
			return err
		}
	}
	if err := validateWikiCitationIdentityMerge(nil, nil, result.NewSlugs); err != nil {
		return err
	}
	var unknown []string
	for slug, ids := range result.Citations {
		if err := validateIDs(ids); err != nil {
			return err
		}
		if !allowedSlugs[slug] {
			// An invented short handle, wrong namespace or empty evidence
			// cannot be repaired by manufacturing a new declaration.
			kind, name, namespaced := strings.Cut(slug, "/")
			if !namespaced || (kind != "entity" && kind != "concept") || strings.TrimSpace(name) == "" || len(ids) == 0 {
				return errors.New("citation has unknown or invalid reference")
			}
			unknown = append(unknown, slug)
		}
	}
	if len(unknown) > 0 {
		sort.Strings(unknown)
		return &wikiUnknownCitationSlugsError{slugs: unknown}
	}
	return nil
}

type wikiMapCheckpoint struct {
	Result  docIngestResult
	Updates []SlugUpdate
}

func validateWikiMapCheckpoint(saved wikiMapCheckpoint, scope *wikiCheckpointScope, chunks []*types.Chunk) error {
	if scope == nil || saved.Result.KnowledgeID != scope.KnowledgeID || strings.TrimSpace(saved.Result.Summary) == "" || len(saved.Updates) == 0 {
		return errors.New("invalid persisted wiki map")
	}
	allowed := make(map[string]bool, len(chunks))
	for _, chunk := range chunks {
		if chunk != nil {
			allowed[chunk.ID] = true
		}
	}
	written := make(map[string]bool)
	for updateIndex, update := range saved.Updates {
		if update.KnowledgeID != scope.KnowledgeID || update.Slug == "" {
			return errors.New("persisted wiki update scope mismatch")
		}
		switch update.Type {
		case "retract", "retractStale":
		case types.WikiPageTypeSummary, types.WikiPageTypeEntity, types.WikiPageTypeConcept:
			if update.SourceRef != scope.KnowledgeID || !strings.HasPrefix(update.Slug, update.Type+"/") || strings.TrimPrefix(update.Slug, update.Type+"/") == "" {
				return errors.New("persisted wiki source reference mismatch")
			}
			if update.Type == types.WikiPageTypeSummary {
				if update.Slug != "summary/"+slugify(scope.KnowledgeID) || strings.TrimSpace(update.SummaryBody) == "" || strings.TrimSpace(update.DocTitle) == "" {
					return errors.New("persisted wiki summary incomplete")
				}
			} else if update.Item.Slug != update.Slug || strings.TrimSpace(update.Item.Name) == "" {
				return errors.New("persisted wiki item mismatch")
			}
			written[update.Slug] = true
		default:
			return errors.New("persisted wiki update type invalid")
		}
		for _, id := range append(append([]string(nil), update.SourceChunks...), update.Item.SourceChunks...) {
			if !allowed[id] {
				return errors.New("persisted wiki map references a different source chunk")
			}
		}
		if update.Type == types.WikiPageTypeEntity || update.Type == types.WikiPageTypeConcept {
			var sources []string
			for _, chunk := range chunks {
				if chunk == nil {
					continue
				}
				for _, id := range update.SourceChunks {
					if chunk.ID == id {
						sources = append(sources, chunk.Content)
						break
					}
				}
			}
			if err := validateWikiSourceItemIdentity(update.Item, update.Type, updateIndex, sources); err != nil {
				return err
			}
		}
	}
	if err := validateWikiPageIdentityUpdates(nil, saved.Updates); err != nil {
		return err
	}
	seen := make(map[string]bool, len(saved.Result.Pages))
	for _, page := range saved.Result.Pages {
		if strings.TrimSpace(page.Title) == "" || !written[page.Slug] || seen[page.Slug] {
			return errors.New("persisted wiki map page mismatch")
		}
		seen[page.Slug] = true
	}
	if len(seen) != len(written) || !written["summary/"+slugify(scope.KnowledgeID)] {
		return errors.New("persisted wiki map coverage incomplete")
	}
	return nil
}

type wikiPageCheckpoint struct {
	Applied  bool
	Absent   bool
	PageHash string
}

func wikiCheckpointPageHash(page *types.WikiPage) string {
	return wikiCheckpointHash([]any{page.ID, page.TenantID, page.KnowledgeBaseID, page.Slug,
		page.PageType, page.Title, page.Content, page.Summary, page.Aliases, page.SourceRefs, page.ChunkRefs})
}

func wikiPageHasSource(page *types.WikiPage, knowledgeID string) bool {
	if page != nil {
		for _, ref := range page.SourceRefs {
			if ref == knowledgeID || strings.HasPrefix(ref, knowledgeID+"|") {
				return true
			}
		}
	}
	return false
}

func (s *wikiIngestService) checkpointPageScope(batch *WikiBatchContext, updates []SlugUpdate) *wikiCheckpointScope {
	if s.checkpointDB == nil || batch == nil || len(updates) == 0 {
		return nil
	}
	var scope *wikiCheckpointScope
	for _, update := range updates {
		value, ok := batch.checkpointScopes.Load(update.KnowledgeID)
		if !ok {
			return nil
		}
		current := value.(*wikiCheckpointScope)
		if scope != nil && *scope != *current {
			return nil
		}
		scope = current
	}
	return scope
}

// Called under the existing per-slug lock. Successful pages survive the next
// bounded slice. Source document changes get a new scope before this can run.
func (s *wikiIngestService) reduceSlugUpdates(ctx context.Context, model chat.Chat, kbID, slug string, updates []SlugUpdate, tenantID uint64, batch *WikiBatchContext, spans map[string]*Span) (bool, string, bool, error) {
	scope := s.checkpointPageScope(batch, updates)
	if scope == nil {
		return s.reduceSlugUpdatesUncheckpointed(ctx, model, kbID, slug, updates, tenantID, batch, spans)
	}
	if scope.TenantID != tenantID || scope.KnowledgeBaseID != kbID {
		return false, "", false, errors.New("wiki page checkpoint scope mismatch")
	}
	expectsSource := false
	for _, update := range updates {
		if update.Type != "retract" && update.Type != "retractStale" {
			expectsSource = true
		}
	}
	key := wikiCheckpointHash([]any{slug, updates})
	var saved wikiPageCheckpoint
	hit, err := s.loadWikiCheckpoint(ctx, scope, "page", key, &saved)
	if err != nil {
		return false, "", false, err
	}
	if hit && saved.Applied {
		page, pageErr := s.wikiService.GetPageBySlug(ctx, kbID, slug)
		if pageErr == nil && page != nil {
			if err := validateWikiPageIdentityUpdates(page, updates); err != nil {
				return false, "", true, err
			}
		}
		if (pageErr == nil || errors.Is(pageErr, apprepo.ErrWikiPageNotFound)) && page == nil && !expectsSource && saved.Absent {
			return false, "", false, nil
		}
		if pageErr == nil && page != nil && page.TenantID == tenantID && page.KnowledgeBaseID == kbID &&
			(page.Status == types.WikiPageStatusDraft || page.Status == types.WikiPageStatusPublished) && saved.PageHash == wikiCheckpointPageHash(page) {
			if wikiPageHasSource(page, scope.KnowledgeID) == expectsSource {
				return false, "", false, nil
			}
		}
	}
	if ctx.Err() != nil {
		return false, "", false, wikiCheckpointError(ctx, ctx.Err())
	}
	if batch.checkpointPages.Add(1) > wikiPageSlice {
		return false, "", false, errWikiCheckpointYield
	}
	changed, affected, failed, err := s.reduceSlugUpdatesUncheckpointed(ctx, model, kbID, slug, updates, tenantID, batch, spans)
	err = wikiCheckpointError(ctx, err)
	if err == nil && !failed {
		page, pageErr := s.wikiService.GetPageBySlug(ctx, kbID, slug)
		if pageErr != nil && !errors.Is(pageErr, apprepo.ErrWikiPageNotFound) {
			return changed, affected, failed, pageErr
		}
		if expectsSource && (page == nil || !wikiPageHasSource(page, scope.KnowledgeID) || strings.TrimSpace(page.Content) == "") {
			return changed, affected, failed, errors.New("wiki page contribution was not persisted")
		}
		if page != nil && (page.TenantID != tenantID || page.KnowledgeBaseID != kbID || page.Slug != slug || wikiPageHasSource(page, scope.KnowledgeID) != expectsSource) {
			return changed, affected, failed, errors.New("wiki page checkpoint state mismatch")
		}
		receipt := wikiPageCheckpoint{Applied: true, Absent: page == nil}
		if page != nil {
			receipt.PageHash = wikiCheckpointPageHash(page)
		}
		err = s.saveWikiCheckpoint(ctx, scope, "page", key, receipt)
	}
	return changed, affected, failed, err
}