xinmi/internal/application/service/wiki_candidate_catalog.go

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
}