xinmi/internal/application/service/wiki_full_read.go

package service

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"strings"
	"unicode/utf8"

	"github.com/Tencent/WeKnora/internal/agent"
	"github.com/Tencent/WeKnora/internal/logger"
	"github.com/Tencent/WeKnora/internal/models/chat"
)

// Each queue turn has a small amount of NEW work; validated windows survive
// retries. Bounds are explicit stops, never silent truncation or sampling.
const (
	wikiFullReadVersion       = "wiki-full-read-v1"
	wikiReadOperationsPerTurn = 4
	wikiReadMaxWindows        = 1024
	wikiReadCandidateRunes    = 24576
	wikiReadSummaryRunes      = 4096
	wikiReadSummaryFanIn      = 6
)

const wikiReadSectionPrompt = `Read the entire supplied source segment and create concise reading notes in {{.Language}}.
This is one contiguous segment of ONE document, not the whole document. Preserve its important topics, qualifications, negations, contradictions, conditions and quantities. Never infer content outside the supplied segment. The notes are a lossy navigation aid, NOT original evidence or a claim that all facts have been preserved.
Return a SUMMARY: headline followed by a Markdown body. Keep the complete response within 3000 Unicode characters: anything longer than 4096 characters is rejected outright and the whole document is parked, so condense instead of enumerating. Avoid padding and verbatim repetition. Do not invent wiki links or citations.
<source_segment>{{.Content}}</source_segment>`

const wikiReadMergePrompt = `Combine these ordered reading notes from ONE source document into concise navigation notes in {{.Language}}.
The inputs are lossy summaries, NOT original source evidence. Use only supplied information, keep distinctions and contradictions, and never claim that these notes preserve every fact or replace the original document. Do not infer connections between adjacent sections. Preserve the major topics and material qualifications.
Return a SUMMARY: headline followed by a Markdown body. Keep the complete response within 3000 Unicode characters: anything longer than 4096 characters is rejected outright and the whole document is parked, so merge and compress instead of enumerating. Do not invent wiki links or citations.
<ordered_reading_notes>{{.Content}}</ordered_reading_notes>`

type wikiReadWindow struct {
	Text      string
	StartRune int // inclusive in reconstructed source, not chunk offsets
	EndRune   int // exclusive
}

// Split on nearby paragraph boundaries while preserving EVERY byte, including
// whitespace and trailing text. Windows retain views of the existing string;
// a 19-million-character document does not need another []rune copy.
func splitWikiReadWindows(content string, maxRunes int) []wikiReadWindow {
	if content == "" || maxRunes < 1 {
		return nil
	}
	var windows []wikiReadWindow
	startByte, startRune := 0, 0
	for startByte < len(content) {
		end, count, newlineByte, newlineRunes := len(content)-startByte, 0, 0, 0
		for pos, r := range content[startByte:] {
			count++
			if r == '\n' {
				newlineByte, newlineRunes = pos+1, count
			}
			if count == maxRunes {
				_, runeBytes := utf8.DecodeRuneInString(content[startByte+pos:])
				end = pos + runeBytes
				if newlineRunes >= maxRunes-1024 && newlineRunes >= maxRunes/2 {
					end, count = newlineByte, newlineRunes
				}
				break
			}
		}
		windows = append(windows, wikiReadWindow{content[startByte : startByte+end], startRune, startRune + count})
		startByte, startRune = startByte+end, startRune+count
	}
	return windows
}

// This is machine-readable processing coverage, NOT semantic recall. In
// particular, source_windows_complete does not imply that citation verification
// or page publication has completed. Scope already binds tenant/KB/document,
// model configuration, content revision and prompts.
type wikiReadProgress struct {
	Version               string `json:"version"`
	Stage                 string `json:"stage"`
	Status                string `json:"status"`
	SourceRunes           int    `json:"source_runes"`
	TotalWindows          int    `json:"total_windows"`
	CompletedWindows      int    `json:"completed_windows"`
	ProcessedRunes        int    `json:"processed_runes"`
	SourceWindowsComplete bool   `json:"source_windows_complete"`
	StageComplete         bool   `json:"stage_complete"`
	EvidenceVerified      bool   `json:"evidence_verified"`
	Reason                string `json:"reason,omitempty"`
}

type wikiReadRange struct {
	StartRune int    `json:"start_rune"`
	EndRune   int    `json:"end_rune"`
	TextHash  string `json:"text_hash"`
}

// Repeated text can share a model artifact, but every original occurrence
// remains independently represented by its source range in this manifest.
func (s *wikiIngestService) saveWikiReadPlan(ctx context.Context, scope *wikiCheckpointScope, windows []wikiReadWindow) error {
	ranges := make([]wikiReadRange, 0, len(windows))
	for _, window := range windows {
		ranges = append(ranges, wikiReadRange{window.StartRune, window.EndRune, wikiCheckpointHash(window.Text)})
	}
	return s.saveWikiCheckpoint(ctx, scope, "read_plan", "source", ranges)
}

func newWikiReadProgress(stage string, windows []wikiReadWindow) wikiReadProgress {
	progress := wikiReadProgress{Version: wikiFullReadVersion, Stage: stage, Status: "reading", TotalWindows: len(windows)}
	if len(windows) > 0 {
		progress.SourceRunes = windows[len(windows)-1].EndRune
	}
	return progress
}

func (s *wikiIngestService) saveWikiReadProgress(ctx context.Context, scope *wikiCheckpointScope, progress wikiReadProgress) error {
	progress.SourceWindowsComplete = progress.CompletedWindows == progress.TotalWindows
	return s.saveWikiCheckpoint(ctx, scope, "read_progress", progress.Stage, progress)
}

func (s *wikiIngestService) pauseWikiRead(ctx context.Context, scope *wikiCheckpointScope, progress wikiReadProgress, reason string) error {
	progress.Status, progress.Reason = "paused_budget", reason
	if err := s.saveWikiReadProgress(ctx, scope, progress); err != nil {
		return err
	}
	return fmt.Errorf("Wiki 全文处理已暂停（已保存断点，不会标记完成）：%s", reason)
}

func (s *wikiIngestService) checkWikiReadScope(ctx context.Context, scope *wikiCheckpointScope, progress wikiReadProgress) error {
	if s.checkpointDB == nil || scope == nil {
		return errors.New("Wiki 长文全文处理需要持久断点存储，未启用时不能截断原文继续")
	}
	if progress.TotalWindows > wikiReadMaxWindows {
		return s.pauseWikiRead(ctx, scope, progress, fmt.Sprintf("原文需要 %d 个窗口，超过单文档上限 %d，请拆分文档", progress.TotalWindows, wikiReadMaxWindows))
	}
	return nil
}

// Prior slugs are identity hints, not source evidence. Bound only these hints;
// exact-title identity resolution still runs against the wiki repository.
func boundedWikiPreviousSlugs(previous map[string]bool) map[string]bool {
	ordered := make([]string, 0, len(previous))
	for slug := range previous {
		ordered = append(ordered, slug)
	}
	sort.Strings(ordered)
	result, chars := make(map[string]bool), 0
	for _, slug := range ordered {
		chars += utf8.RuneCountInString(slug) + 3
		if chars > 8192 {
			break
		}
		result[slug] = true
	}
	return result
}

func mergeWikiWindowCandidates(all *combinedExtraction, part combinedExtraction) error {
	kept, dropped, err := pruneIncomingWindowCandidates(*all, part)
	if err != nil {
		return err
	}
	if dropped > 0 {
		// Counted, content-free evidence that an identity conflict or an
		// unsupported alias removed candidates instead of the whole window.
		all.CandidatePruning.SlugConflicts += dropped
		logger.Infof(context.Background(), "wiki window merge pruning: policy=%s slug_conflicts=%d kept=%d",
			wikiSourceIdentityRepairVersion, all.CandidatePruning.SlugConflicts, len(kept.Entities)+len(kept.Concepts))
	}
	part = kept
	combined := combinedExtraction{
		Entities: append(append([]extractedItem(nil), all.Entities...), part.Entities...),
		Concepts: append(append([]extractedItem(nil), all.Concepts...), part.Concepts...),
	}
	if err := validateWikiIdentityCollisions(combined); err != nil {
		return err
	}
	merge := func(dst, incoming []extractedItem) []extractedItem {
		index := make(map[string]int, len(dst))
		for i, item := range dst {
			index[item.Slug] = i
		}
		for _, item := range incoming {
			if i, ok := index[item.Slug]; ok {
				// Strict source identity never turns a display-name
				// difference into a generated alias. Both candidates have
				// already passed exact outer-trimmed title validation.
				merged := dst[i]
				if len([]rune(item.Description)) > len([]rune(merged.Description)) {
					merged.Description = item.Description
				}
				if len([]rune(item.Details)) > len([]rune(merged.Details)) {
					merged.Details = item.Details
				}
				for _, chunkID := range item.SourceChunks {
					merged.SourceChunks = appendUniqueString(merged.SourceChunks, chunkID)
				}
				dst[i] = merged
			} else {
				index[item.Slug] = len(dst)
				dst = append(dst, item)
			}
		}
		return dst
	}
	all.Entities = merge(all.Entities, part.Entities)
	all.Concepts = merge(all.Concepts, part.Concepts)
	return nil
}

func (s *wikiIngestService) extractWikiFullCandidates(ctx context.Context, model chat.Chat, kbID, content, lang string, previous map[string]bool, batch *WikiBatchContext, scope *wikiCheckpointScope) ([]extractedItem, []extractedItem, map[string]extractedItem, error) {
	windows := splitWikiReadWindows(content, maxContentForWiki)
	if len(windows) <= 1 {
		return s.extractCandidateSlugs(ctx, model, kbID, content, lang, boundedWikiPreviousSlugs(previous), batch)
	}
	progress := newWikiReadProgress("candidates", windows)
	if err := s.checkWikiReadScope(ctx, scope, progress); err != nil {
		return nil, nil, nil, err
	}
	if err := s.saveWikiReadPlan(ctx, scope, windows); err != nil {
		return nil, nil, nil, err
	}
	merged, operations := combinedExtraction{}, 0
	for windowIndex, window := range windows {
		key := wikiCheckpointHash([]any{wikiFullReadVersion, window.Text})
		var part combinedExtraction
		hit, err := s.loadWikiCheckpoint(ctx, scope, "read_candidate", key, &part)
		if err == nil && hit {
			err = validateWikiFreshCandidateCheckpoint(part, window.Text)
		}
		if err != nil {
			return nil, nil, nil, fmt.Errorf("Wiki 候选读取第 %d/%d 个原文窗口失败：%w", windowIndex+1, len(windows), err)
		}
		if !hit {
			if operations >= wikiReadOperationsPerTurn {
				if err := s.saveWikiReadProgress(ctx, scope, progress); err != nil {
					return nil, nil, nil, err
				}
				return nil, nil, nil, errWikiCheckpointYield
			}
			if ctx.Err() != nil {
				return nil, nil, nil, wikiCheckpointError(ctx, ctx.Err())
			}
			operations++
			part, err = s.extractFreshCandidateSlugs(ctx, model, window.Text, lang, boundedWikiPreviousSlugs(previous), batch)
			if err == nil {
				err = validateWikiFreshCandidateCheckpoint(part, window.Text)
			}
			if err == nil {
				err = s.saveWikiCheckpoint(ctx, scope, "read_candidate", key, part)
			}
			if err != nil {
				return nil, nil, nil, wikiCheckpointError(ctx, fmt.Errorf("Wiki 候选抽取第 %d/%d 个原文窗口失败：%w", windowIndex+1, len(windows), err))
			}
		}
		merged.Entities = append(merged.Entities, part.Entities...)
		merged.Concepts = append(merged.Concepts, part.Concepts...)
		progress.CompletedWindows++
		progress.ProcessedRunes = window.EndRune
		if err := s.saveWikiReadProgress(ctx, scope, progress); err != nil {
			return nil, nil, nil, err
		}
	}
	// Window checkpoints preserve fresh proposals. Claims begin only after the
	// complete set has deterministic, collision-free routing identifiers.
	merged, err := normalizeWikiFreshCandidates(merged)
	if err == nil {
		merged, err = s.resolveWikiSourceIdentities(ctx, kbID, merged, batch)
	}
	if err != nil {
		return nil, nil, nil, err
	}
	progress.StageComplete, progress.Status = true, "complete"
	if err := s.saveWikiReadProgress(ctx, scope, progress); err != nil {
		return nil, nil, nil, err
	}
	items := make(map[string]extractedItem)
	for _, item := range append(append([]extractedItem(nil), merged.Entities...), merged.Concepts...) {
		items[item.Slug] = item
	}
	return merged.Entities, merged.Concepts, items, nil
}

// wikiSummaryCompressionInput states the overshoot and asks for compression, so the
// retry cannot simply repeat the same over-long answer.
func wikiSummaryCompressionInput(input string, producedRunes int) string {
	return fmt.Sprintf("%s\n\n<compression_request>Your previous answer for this segment was %d characters, over the %d character budget, so it was rejected. Rewrite it now within 3000 characters: keep the SUMMARY: headline, keep one line per topic, drop examples, enumerations and repeated wording.</compression_request>",
		input, producedRunes, wikiReadSummaryRunes)
}

func validateWikiReadSummary(content string) error {
	if err := validateWikiPageBody(content, true); err != nil {
		return err
	}
	if utf8.RuneCountInString(content) > wikiReadSummaryRunes {
		return &wikiOutputError{reason: "全文阅读摘要超过 4096 字符预算，拒绝截断或当作完整结果"}
	}
	return nil
}

type wikiReadSummaryNode struct {
	Content            string
	StartRune, EndRune int
}

func (s *wikiIngestService) generateWikiFullSummary(ctx context.Context, model chat.Chat, content, lang, slugs string, batch *WikiBatchContext, scope *wikiCheckpointScope) (string, error) {
	windows := splitWikiReadWindows(content, maxContentForWiki)
	if len(windows) <= 1 && utf8.RuneCountInString(slugs) <= wikiReadCandidateRunes {
		return s.generateWithTemplate(ctx, model, agent.WikiSummaryPrompt, map[string]string{
			"Content": content, "Language": lang, "ExtractedSlugs": slugs,
			"CustomInstructions": batch.ContentInstructions, "InstructionScope": "wiki_content",
		})
	}
	if len(windows) == 0 {
		return "", errors.New("Wiki 全文摘要缺少原文，不能用候选目录代替原文")
	}
	progress := newWikiReadProgress("summary", windows)
	if err := s.checkWikiReadScope(ctx, scope, progress); err != nil {
		return "", err
	}
	if err := s.saveWikiReadPlan(ctx, scope, windows); err != nil {
		return "", err
	}
	operations := 0
	generate := func(stage, prompt, input string) (string, error) {
		key := wikiCheckpointHash([]any{wikiFullReadVersion, prompt, input})
		var pausedReason string
		if hit, err := s.loadWikiCheckpoint(ctx, scope, "read_pause", key, &pausedReason); err != nil {
			return "", err
		} else if hit {
			return "", s.pauseWikiRead(ctx, scope, progress, pausedReason)
		}
		var output string
		hit, err := s.loadWikiCheckpoint(ctx, scope, stage, key, &output)
		if err != nil {
			return "", err
		}
		if hit {
			return output, validateWikiReadSummary(output)
		}
		if operations >= wikiReadOperationsPerTurn {
			return "", errWikiCheckpointYield
		}
		if ctx.Err() != nil {
			return "", wikiCheckpointError(ctx, ctx.Err())
		}
		operations++
		output, err = s.generateWithTemplate(ctx, model, prompt, map[string]string{
			"Content": input, "Language": lang, "CustomInstructions": batch.ContentInstructions,
			"InstructionScope": "wiki_content",
		})
		if err == nil {
			err = validateWikiReadSummary(output)
		}
		if err != nil && utf8.RuneCountInString(output) > wikiReadSummaryRunes {
			// Policy: a summary that overshoots the character budget gets ONE bounded
			// compression retry, because the observed failure is a model ignoring an
			// explicit budget rather than an unsolvable segment. Only a second
			// overshoot parks the document, so a model that cannot follow the budget
			// still stops instead of looping.
			if operations < wikiReadOperationsPerTurn && ctx.Err() == nil {
				operations++
				repaired, repairErr := s.generateWithTemplate(ctx, model, prompt, map[string]string{
					"Content":            wikiSummaryCompressionInput(input, utf8.RuneCountInString(output)),
					"Language":           lang,
					"CustomInstructions": batch.ContentInstructions,
					"InstructionScope":   "wiki_content",
				})
				if repairErr == nil {
					repairErr = validateWikiReadSummary(repaired)
				}
				if repairErr == nil {
					output, err = repaired, nil
				}
			}
		}
		if err != nil && utf8.RuneCountInString(output) > wikiReadSummaryRunes {
			reason := "模型摘要超过单窗口输出预算，需要调整模型或分段策略，已保留成功窗口"
			if saveErr := s.saveWikiCheckpoint(ctx, scope, "read_pause", key, reason); saveErr != nil {
				return "", saveErr
			}
			return "", s.pauseWikiRead(ctx, scope, progress, reason)
		}
		if err == nil {
			err = s.saveWikiCheckpoint(ctx, scope, stage, key, output)
		}
		return output, wikiCheckpointError(ctx, err)
	}
	var nodes []wikiReadSummaryNode
	for _, window := range windows {
		output, err := generate("read_summary", wikiReadSectionPrompt, window.Text)
		if err != nil {
			return "", err
		}
		nodes = append(nodes, wikiReadSummaryNode{output, window.StartRune, window.EndRune})
		progress.CompletedWindows++
		progress.ProcessedRunes = window.EndRune
		if err := s.saveWikiReadProgress(ctx, scope, progress); err != nil {
			return "", err
		}
	}
	for len(nodes) > 1 {
		var next []wikiReadSummaryNode
		for start := 0; start < len(nodes); start += wikiReadSummaryFanIn {
			end := min(start+wikiReadSummaryFanIn, len(nodes))
			if end-start == 1 {
				next = append(next, nodes[start])
				continue
			}
			var input strings.Builder
			for _, node := range nodes[start:end] {
				fmt.Fprintf(&input, "<notes source_runes=\"%d:%d\">\n%s\n</notes>\n", node.StartRune, node.EndRune, node.Content)
			}
			output, err := generate("read_reduce", wikiReadMergePrompt, input.String())
			if err != nil {
				return "", err
			}
			next = append(next, wikiReadSummaryNode{output, nodes[start].StartRune, nodes[end-1].EndRune})
		}
		nodes = next
	}
	progress.StageComplete, progress.Status = true, "complete"
	if err := s.saveWikiReadProgress(ctx, scope, progress); err != nil {
		return "", err
	}
	// The visible provenance note prevents this compressed overview from being
	// presented as an exhaustive evidence store. Verbatim cited chunks remain
	// the independent source of facts for entity/concept pages.
	return nodes[0].Content + fmt.Sprintf("\n\n> 阅读范围：原文 %d 个窗口、%d 字符均已分段处理。此页为分层摘要，不等同完整原文或全部事实核验；具体结论应回查原文引用。", len(windows), progress.SourceRunes), nil
}