255 lines
9.3 KiB
Go
255 lines
9.3 KiB
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
|
||
|
|
}
|