xinmi/internal/application/service/wiki_citation_evidence.go

package service

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"strings"

	"github.com/Tencent/WeKnora/internal/types"
)

const wikiCitationEvidenceVersion = "wiki-citation-evidence-v2"

var errWikiCitationEvidence = errors.New("wiki source evidence is incomplete or invalid")

// resolveWikiAdditionEvidence is the publication gate for entity/concept
// additions. A candidate's description/details are model output, not source
// evidence. Every declared citation must still resolve inside the exact source
// scope before Reduce mutates a page or calls the editor model. Rejecting the
// whole update preserves the old page and prevents partial citations from
// silently masquerading as a complete evidence set.
func (s *wikiIngestService) resolveWikiAdditionEvidence(
	ctx context.Context, tenantID uint64, kbID string, additions []SlugUpdate,
) (map[string]*types.Chunk, error) {
	out := make(map[string]*types.Chunk)
	if len(additions) == 0 {
		return out, nil
	}
	if tenantID == 0 || strings.TrimSpace(kbID) == "" || s.chunkRepo == nil {
		return nil, fmt.Errorf("%w: missing source scope or chunk repository", errWikiCitationEvidence)
	}

	owners := make(map[string]string)
	var ids []string
	for _, add := range additions {
		refID, _, _ := strings.Cut(add.SourceRef, "|")
		if strings.TrimSpace(add.KnowledgeID) == "" || refID != add.KnowledgeID || len(add.SourceChunks) == 0 {
			return nil, fmt.Errorf("%w: page %q has an unverified candidate from document %q", errWikiCitationEvidence, add.Slug, add.KnowledgeID)
		}
		for _, id := range add.SourceChunks {
			if id == "" || strings.TrimSpace(id) != id {
				return nil, fmt.Errorf("%w: document %q has a blank or malformed chunk reference", errWikiCitationEvidence, add.KnowledgeID)
			}
			if owner, ok := owners[id]; ok {
				if owner != add.KnowledgeID {
					return nil, fmt.Errorf("%w: chunk %q is attributed to multiple documents", errWikiCitationEvidence, id)
				}
				continue
			}
			owners[id] = add.KnowledgeID
			ids = append(ids, id)
		}
	}

	const lookupBatchSize = 128
	for start := 0; start < len(ids); start += lookupBatchSize {
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		end := min(start+lookupBatchSize, len(ids))
		chunks, err := s.chunkRepo.ListChunksByID(ctx, tenantID, ids[start:end])
		if err != nil {
			return nil, fmt.Errorf("%w: citation lookup failed: %w", errWikiCitationEvidence, err)
		}
		requested := make(map[string]bool, end-start)
		for _, id := range ids[start:end] {
			requested[id] = true
		}
		for _, chunk := range chunks {
			if chunk == nil || !requested[chunk.ID] {
				return nil, fmt.Errorf("%w: citation lookup returned an unexpected chunk", errWikiCitationEvidence)
			}
			if chunk.TenantID != tenantID || chunk.KnowledgeBaseID != kbID || chunk.KnowledgeID != owners[chunk.ID] {
				return nil, fmt.Errorf("%w: chunk %q does not belong to the cited source", errWikiCitationEvidence, chunk.ID)
			}
			if !chunk.IsEnabled || chunk.DeletedAt.Valid || strings.TrimSpace(chunk.Content) == "" {
				return nil, fmt.Errorf("%w: chunk %q is unavailable or has no source text", errWikiCitationEvidence, chunk.ID)
			}
			if _, duplicate := out[chunk.ID]; duplicate {
				return nil, fmt.Errorf("%w: duplicate chunk %q returned by citation lookup", errWikiCitationEvidence, chunk.ID)
			}
			// Keep an immutable row snapshot across a potentially long editor call.
			copy := *chunk
			out[chunk.ID] = &copy
		}
		for _, id := range ids[start:end] {
			if out[id] == nil {
				return nil, fmt.Errorf("%w: cited chunk %q was not found", errWikiCitationEvidence, id)
			}
		}
	}
	return out, nil
}

// wikiEvidenceSourceHash covers source text, identity, revision, ordering and
// enabled/deleted state. Database result order is not part of the signature.
func wikiEvidenceSourceHash(chunks []*types.Chunk) string {
	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
	})
	return wikiCheckpointHash(wikiCheckpointChunks(ordered))
}

func (s *wikiIngestService) verifyWikiEvidenceKnowledgeLive(ctx context.Context, kbID, knowledgeID string) error {
	if s.knowledgeSvc == nil {
		return fmt.Errorf("%w: source knowledge service is unavailable", errWikiCitationEvidence)
	}
	gone, err := s.isKnowledgeGone(ctx, kbID, knowledgeID)
	if err != nil {
		return fmt.Errorf("%w: %w", errWikiCitationEvidence, err)
	}
	if gone {
		return fmt.Errorf("%w: source document %q was deleted during page generation", errWikiCitationEvidence, knowledgeID)
	}
	return nil
}

// Called immediately before saving an edited page. If a user changed or
// removed evidence while the model was working, keep the previously published
// page and let a fresh source scope rebuild it. Never publish stale evidence.
func (s *wikiIngestService) verifyWikiAdditionEvidenceUnchanged(
	ctx context.Context, tenantID uint64, kbID string, additions []SlugUpdate, before map[string]*types.Chunk,
) error {
	if len(additions) == 0 {
		return nil
	}
	seen := make(map[string]bool)
	for _, add := range additions {
		if seen[add.KnowledgeID] {
			continue
		}
		if err := s.verifyWikiEvidenceKnowledgeLive(ctx, kbID, add.KnowledgeID); err != nil {
			return err
		}
		seen[add.KnowledgeID] = true
	}
	after, err := s.resolveWikiAdditionEvidence(ctx, tenantID, kbID, additions)
	if err != nil {
		return err
	}
	if len(before) != len(after) {
		return fmt.Errorf("%w: source evidence changed during page generation", errWikiCitationEvidence)
	}
	for id, chunk := range before {
		if chunk == nil || after[id] == nil || wikiEvidenceSourceHash([]*types.Chunk{chunk}) != wikiEvidenceSourceHash([]*types.Chunk{after[id]}) {
			return fmt.Errorf("%w: cited chunk %q changed during page generation", errWikiCitationEvidence, id)
		}
	}
	return nil
}

// Summaries have no per-topic citations, so a checkpointed summary must compare
// the entire source-chunk set captured during Map, including additions/removals.
func (s *wikiIngestService) verifyWikiSummaryEvidenceUnchanged(
	ctx context.Context, tenantID uint64, kbID string, update SlugUpdate, batch *WikiBatchContext,
) error {
	if err := s.verifyWikiEvidenceKnowledgeLive(ctx, kbID, update.KnowledgeID); err != nil {
		return err
	}
	if batch == nil {
		return nil
	}
	scopeValue, hasScope := batch.checkpointScopes.Load(update.KnowledgeID)
	if !hasScope {
		return nil
	}
	scope, ok := scopeValue.(*wikiCheckpointScope)
	if !ok || scope == nil || scope.TenantID != tenantID || scope.KnowledgeBaseID != kbID || scope.KnowledgeID != update.KnowledgeID {
		return fmt.Errorf("%w: summary source scope is invalid", errWikiCitationEvidence)
	}
	baseline, ok := batch.checkpointSourceHashes.Load(update.KnowledgeID)
	if !ok || s.chunkRepo == nil {
		return fmt.Errorf("%w: summary source snapshot is missing", errWikiCitationEvidence)
	}
	chunks, err := s.chunkRepo.ListChunksByKnowledgeID(ctx, tenantID, update.KnowledgeID)
	if err != nil {
		return fmt.Errorf("%w: %w", errWikiCitationEvidence, err)
	}
	for _, chunk := range chunks {
		if chunk == nil || chunk.TenantID != tenantID || chunk.KnowledgeBaseID != kbID || chunk.KnowledgeID != update.KnowledgeID {
			return fmt.Errorf("%w: summary source lookup returned an invalid chunk", errWikiCitationEvidence)
		}
	}
	if hash, ok := baseline.(string); !ok || hash == "" || hash != wikiEvidenceSourceHash(chunks) {
		return fmt.Errorf("%w: source document changed during summary generation", errWikiCitationEvidence)
	}
	return nil
}