WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/application/service/wiki_candidate_catalog.go

@@ -0,0 +1,254 @@
+package service
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sort"
+	"strings"
+	"unicode/utf8"
+
+	"github.com/Tencent/WeKnora/internal/agent"
+	"github.com/Tencent/WeKnora/internal/modelcontext"
+	"github.com/Tencent/WeKnora/internal/models/chat"
+	"github.com/Tencent/WeKnora/internal/types"
+)
+
+const (
+	wikiCandidateCatalogVersion = "wiki-candidate-catalog-v1"
+	wikiCitationRequestRunes    = 65536
+)
+
+type wikiCandidateEntry struct {
+	Kind string        `json:"kind"`
+	Item extractedItem `json:"item"`
+}
+
+type wikiCandidateShard struct {
+	Hash    string               `json:"hash"`
+	Entries []wikiCandidateEntry `json:"entries"`
+	XML     string               `json:"-"`
+}
+
+type wikiCitationSourceRange struct {
+	ChunkID   string `json:"chunk_id"`
+	StartRune int    `json:"start_rune"`
+	EndRune   int    `json:"end_rune"`
+	TextHash  string `json:"text_hash"`
+}
+
+type wikiCitationSourceBatch struct {
+	Hash   string                    `json:"hash"`
+	Ranges []wikiCitationSourceRange `json:"ranges"`
+	Batch  chunkBatch                `json:"-"`
+}
+
+type wikiCitationPlan struct {
+	Version string                    `json:"version"`
+	Hash    string                    `json:"hash"`
+	Shards  []wikiCandidateShard      `json:"shards"`
+	Sources []wikiCitationSourceBatch `json:"sources"`
+}
+
+func renderWikiCandidateEntries(entries []wikiCandidateEntry) string {
+	var entities, concepts []extractedItem
+	for _, entry := range entries {
+		if entry.Kind == "entity" {
+			entities = append(entities, entry.Item)
+		} else {
+			concepts = append(concepts, entry.Item)
+		}
+	}
+	return renderCandidateSlugsXML(entities, concepts)
+}
+
+// Whole structured entries, including every alias and audit field, stay in the
+// plan. Only the actual routing renderer determines the request shard budget.
+func buildWikiCandidateShards(entities, concepts []extractedItem, budget int) ([]wikiCandidateShard, error) {
+	var entries []wikiCandidateEntry
+	for _, group := range []struct {
+		kind  string
+		items []extractedItem
+	}{{"entity", entities}, {"concept", concepts}} {
+		for _, item := range group.items {
+			if strings.TrimSpace(item.Name) == "" || strings.TrimSpace(item.Slug) == "" {
+				return nil, errors.New("Wiki candidate catalog contains an empty identity")
+			}
+			entries = append(entries, wikiCandidateEntry{group.kind, item})
+		}
+	}
+	sort.Slice(entries, func(i, j int) bool {
+		a, b := entries[i], entries[j]
+		if a.Kind != b.Kind {
+			return a.Kind < b.Kind
+		}
+		if a.Item.Slug != b.Item.Slug {
+			return a.Item.Slug < b.Item.Slug
+		}
+		return wikiCheckpointHash(a) < wikiCheckpointHash(b)
+	})
+	var shards []wikiCandidateShard
+	var current []wikiCandidateEntry
+	flush := func() {
+		if len(current) == 0 {
+			return
+		}
+		xml := renderWikiCandidateEntries(current)
+		shards = append(shards, wikiCandidateShard{Hash: wikiCheckpointHash([]any{wikiCandidateCatalogVersion, budget, len(shards), current, xml}), Entries: current, XML: xml})
+		current = nil
+	}
+	runes := 0
+	for _, entry := range entries {
+		entryRunes := utf8.RuneCountInString(renderWikiCandidateEntries([]wikiCandidateEntry{entry}))
+		if entryRunes > budget {
+			return nil, fmt.Errorf("Wiki candidate %q requires %d routing characters, exceeding the %d-character single-entry budget; no identity or alias was truncated", entry.Item.Slug, entryRunes, budget)
+		}
+		if runes+entryRunes > budget {
+			flush()
+			runes = 0
+		}
+		current = append(current, entry)
+		runes += entryRunes
+	}
+	flush()
+	return shards, nil
+}
+
+func wikiCitationBatchForChunks(chunks []*types.Chunk) chunkBatch {
+	batch := chunkBatch{chunks: chunks, handles: modelcontext.NewHandleTable("c", 3, 0)}
+	for _, chunk := range chunks {
+		batch.handles.Register(chunk.ID)
+		batch.totalRuneLen += utf8.RuneCountInString(chunk.Content)
+	}
+	return batch
+}
+
+// Source ranges are exact, contiguous views of each original text chunk. The
+// budget includes XML wrappers, and one batch never aliases two ranges to the
+// same source handle. The legacy string caller keeps its original batching.
+func buildWikiCitationSources(chunks []*types.Chunk, budget int) ([]wikiCitationSourceBatch, error) {
+	var ordered []*types.Chunk
+	seen := make(map[string]bool)
+	for _, chunk := range chunks {
+		if chunk == nil || chunk.Content == "" || (chunk.ChunkType != "" && chunk.ChunkType != types.ChunkTypeText) {
+			continue
+		}
+		if chunk.ID == "" || seen[chunk.ID] {
+			return nil, errors.New("Wiki citation source has an empty or duplicate chunk identity")
+		}
+		seen[chunk.ID] = true
+		ordered = append(ordered, chunk)
+	}
+	sort.Slice(ordered, func(i, j int) bool {
+		if ordered[i].ChunkIndex != ordered[j].ChunkIndex {
+			return ordered[i].ChunkIndex < ordered[j].ChunkIndex
+		}
+		if ordered[i].StartAt != ordered[j].StartAt {
+			return ordered[i].StartAt < ordered[j].StartAt
+		}
+		return ordered[i].ID < ordered[j].ID
+	})
+	var sources []wikiCitationSourceBatch
+	var current wikiCitationSourceBatch
+	flush := func() {
+		if len(current.Ranges) == 0 {
+			return
+		}
+		current.Hash = wikiCheckpointHash([]any{wikiCandidateCatalogVersion, budget, current.Ranges, wikiCheckpointChunks(current.Batch.chunks)})
+		sources = append(sources, current)
+		current = wikiCitationSourceBatch{}
+	}
+	for _, chunk := range ordered {
+		empty := *chunk
+		empty.Content = ""
+		wrapperRunes := utf8.RuneCountInString(renderChunksXML(wikiCitationBatchForChunks([]*types.Chunk{&empty})))
+		if wrapperRunes >= budget {
+			return nil, errors.New("Wiki citation source wrapper exceeds the source request budget")
+		}
+		for _, window := range splitWikiReadWindows(chunk.Content, budget-wrapperRunes) {
+			part := *chunk
+			part.Content = window.Text
+			if len(current.Ranges) > 0 && current.Ranges[len(current.Ranges)-1].ChunkID == chunk.ID {
+				flush()
+			}
+			trial := wikiCitationBatchForChunks(append(append([]*types.Chunk(nil), current.Batch.chunks...), &part))
+			if utf8.RuneCountInString(renderChunksXML(trial)) > budget {
+				flush()
+				trial = wikiCitationBatchForChunks([]*types.Chunk{&part})
+			}
+			current.Batch = trial
+			current.Ranges = append(current.Ranges, wikiCitationSourceRange{chunk.ID, window.StartRune, window.EndRune, wikiCheckpointHash(window.Text)})
+		}
+	}
+	flush()
+	return sources, nil
+}
+
+func buildWikiCitationPlan(entities, concepts []extractedItem, chunks []*types.Chunk, lang string) (*wikiCitationPlan, error) {
+	shards, err := buildWikiCandidateShards(entities, concepts, wikiReadCandidateRunes)
+	if err != nil {
+		return nil, err
+	}
+	sources, err := buildWikiCitationSources(chunks, maxRunesPerCitationBatch)
+	if err != nil {
+		return nil, err
+	}
+	plan := &wikiCitationPlan{Version: wikiCandidateCatalogVersion, Shards: shards, Sources: sources}
+	plan.Hash = wikiCheckpointHash([]any{plan, lang, wikiReadCandidateRunes, maxRunesPerCitationBatch, wikiCitationRequestRunes, wikiCitationMergeVersion, agent.WikiChunkCitationPrompt})
+	return plan, nil
+}
+
+type wikiCitationPair struct {
+	CandidatesXML string
+	Batch         chunkBatch
+	Key           string
+}
+
+type wikiCitationProgress struct {
+	Version         string `json:"version"`
+	PlanHash        string `json:"plan_hash"`
+	CandidateShards int    `json:"candidate_shards"`
+	SourceBatches   int    `json:"source_batches"`
+	TotalPairs      int    `json:"total_pairs"`
+	CompletedPairs  int    `json:"completed_pairs"`
+	StageComplete   bool   `json:"stage_complete"`
+}
+
+// Guard the actual templated messages, including provider-repair instructions.
+// Oversize requests stop before sending; the guard never truncates source text.
+type wikiCitationChat interface{ chat.Chat }
+type wikiCitationBoundedChat struct{ wikiCitationChat }
+
+func (model wikiCitationBoundedChat) Chat(ctx context.Context, messages []chat.Message, options *chat.ChatOptions) (*types.ChatResponse, error) {
+	runes := 0
+	for _, message := range messages {
+		runes += utf8.RuneCountInString(message.Content)
+	}
+	if runes > wikiCitationRequestRunes {
+		return nil, fmt.Errorf("Wiki citation request requires %d characters, exceeding the %d-character complete-request budget; source was not truncated", runes, wikiCitationRequestRunes)
+	}
+	return model.wikiCitationChat.Chat(ctx, messages, options)
+}
+
+func (s *wikiIngestService) classifyWikiCitationPlan(ctx context.Context, model chat.Chat, plan *wikiCitationPlan, chunks []*types.Chunk, lang string, scope *wikiCheckpointScope) (map[string][]string, []newSlugFromCitation, int, error) {
+	total := len(plan.Shards) * len(plan.Sources)
+	if total > wikiCitationSlice && (scope == nil || s.checkpointDB == nil) {
+		return nil, nil, 0, errors.New("Wiki 分批引用需要持久断点存储，未启用时不能丢弃未完成批次")
+	}
+	if err := s.saveWikiCheckpoint(ctx, scope, "citation_plan", plan.Hash, plan); err != nil {
+		return nil, nil, 0, err
+	}
+	var pairs []wikiCitationPair
+	for _, shard := range plan.Shards {
+		for _, source := range plan.Sources {
+			key := wikiCheckpointHash([]any{wikiCandidateCatalogVersion, plan.Hash, shard.Hash, source.Hash, lang})
+			pairs = append(pairs, wikiCitationPair{shard.XML, source.Batch, key})
+		}
+	}
+	citations, declarations, completed, err := s.classifyWikiCitationPairs(ctx, wikiCitationBoundedChat{model}, pairs, chunks, lang, scope, "citation_pair")
+	progress := wikiCitationProgress{wikiCandidateCatalogVersion, plan.Hash, len(plan.Shards), len(plan.Sources), total, completed, err == nil && completed == total}
+	if saveErr := s.saveWikiCheckpoint(ctx, scope, "citation_progress", plan.Hash, progress); saveErr != nil {
+		return nil, nil, completed, saveErr
+	}
+	return citations, declarations, completed, err
+}