xinmi/internal/application/service/wiki_structured_output.go

package service

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"unicode/utf8"

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

// Output errors can be retried with smaller inputs. Provider errors cannot:
// splitting those would multiply authentication failures or rate limits.
type wikiOutputError struct {
	reason     string
	incomplete bool
}

type wikiCandidateValidation struct {
	Code         string `json:"code"`
	ContentRunes int    `json:"content_runes,omitempty"`
	ContentShape string `json:"content_shape,omitempty"`
	Field        string `json:"field,omitempty"`
	Index        *int   `json:"index,omitempty"`
	JSONOffset   int64  `json:"json_byte_offset,omitempty"`
}

const (
	// Bump this whenever candidate quality gates or their thresholds change.
	// It is included in the Wiki checkpoint fingerprint so cached candidates
	// produced under a weaker policy cannot bypass the current gates.
	// v2: candidate quality prunes and counts instead of failing the document.
	wikiCandidateQualityVersion         = "wiki-candidate-quality-v2"
	wikiCandidateQualityMinSample       = 16
	wikiCandidateQualityMinEntitySample = 8
	wikiCandidateQualityMaxPerSource    = 128
)

// wikiCandidatePruning counts what a bounded quality pass removed. Only counts
// leave the extractor: names and model content never enter the record.
type wikiCandidatePruning struct {
	Duplicates    int `json:"duplicates"`
	Unmatched     int `json:"unmatched"`
	Capped        int `json:"capped"`
	SlugConflicts int `json:"slug_conflicts"`
	Kept          int `json:"kept"`
}

func (p wikiCandidatePruning) Dropped() int {
	return p.Duplicates + p.Unmatched + p.Capped + p.SlugConflicts
}

// wikiCandidateCap is the largest candidate set the policy accepts for a source.
func wikiCandidateCap(source string) int {
	sourceRunes := utf8.RuneCountInString(source)
	if sourceRunes > wikiCandidateQualityMaxPerSource*wikiCandidateQualityMinSample {
		return (sourceRunes + wikiCandidateQualityMaxPerSource - 1) / wikiCandidateQualityMaxPerSource
	}
	return wikiCandidateQualityMaxPerSource
}

// pruneWikiCandidateQuality applies the candidate quality policy without failing
// the document: duplicated names and any overflow above the per-source cap are
// dropped and counted. Names that cannot be matched to the source are left to the
// identity stage, which can still repair a bounded number of them instead of
// discarding a repairable candidate. Surviving candidates keep their original
// order, so the earliest source positions win.
func pruneWikiCandidateQuality(result combinedExtraction, source string) (combinedExtraction, wikiCandidatePruning) {
	var pruning wikiCandidatePruning
	pruned := combinedExtraction{
		Entities: pruneCandidateItems(result.Entities, &pruning),
		Concepts: pruneCandidateItems(result.Concepts, &pruning),
	}
	maxCandidates := wikiCandidateCap(source)
	total := len(pruned.Entities) + len(pruned.Concepts)
	if total > maxCandidates {
		pruning.Capped = total - maxCandidates
		if len(pruned.Entities) > maxCandidates {
			pruned.Entities = pruned.Entities[:maxCandidates]
			pruned.Concepts = nil
		} else {
			pruned.Concepts = pruned.Concepts[:maxCandidates-len(pruned.Entities)]
		}
	}
	pruning.Kept = len(pruned.Entities) + len(pruned.Concepts)
	pruned.CandidatePruning = pruning
	return pruned, pruning
}

func pruneCandidateItems(items []extractedItem, pruning *wikiCandidatePruning) []extractedItem {
	seen := make(map[string]struct{}, len(items))
	kept := make([]extractedItem, 0, len(items))
	for _, item := range items {
		name := strings.TrimSpace(item.Name)
		if _, exists := seen[name]; exists {
			pruning.Duplicates++
			continue
		}
		seen[name] = struct{}{}
		kept = append(kept, item)
	}
	return kept
}

// wikiContentShape classifies a rejected response without quoting it: only the
// shape and size leave the extractor, never the model's text.
func wikiContentShape(raw string) string {
	trimmed := strings.TrimSpace(raw)
	if trimmed == "" {
		return "empty"
	}
	if strings.HasPrefix(trimmed, "```") {
		return "code_fence"
	}
	switch trimmed[0] {
	case '{':
		return "json_object"
	case '[':
		return "json_array"
	}
	return "prose"
}

func (v wikiCandidateValidation) String() string {
	encoded, _ := json.Marshal(v)
	return string(encoded)
}

const wikiCandidateRepairVersion = "wiki-candidate-repair-v1"

const wikiCandidatePartitionVersion = "wiki-candidate-partition-v1"
const wikiCandidateProactiveSplitRunes = 16384

const wikiCandidateGenerationInstruction = `<complete_candidate_set version="` + wikiCandidatePartitionVersion + `">
Read every part of the supplied source and return the complete entities and concepts arrays. Within each array, emit each identity only once, using its type and exact outer-trimmed name under the source identity rules above. Repeated mentions of that same identity do not need repeated items. Preserve every distinct supported identity and its source-grounded distinctions; do not merge by alias, translation, similarity or shared topic, impose a candidate count cap, omit the end of the source or return only a prefix.
Keep navigation descriptions concise rather than repeating items or facts. Complete both arrays and close the JSON object, then stop. Do not repeat an item to fill the output budget, append commentary or continue generating after the closing JSON object.
</complete_candidate_set>`

const wikiCandidateRepairInstruction = `<candidate_regeneration>
Regenerate the WHOLE candidate JSON from the SAME complete source and instructions above. Do not continue, quote, infer missing values from, or rely on the previous rejected output.
Return one valid JSON object with non-null entities and concepts arrays. Every item must have a nonempty name and slug. Entity slugs start with entity/; concept slugs start with concept/; each suffix is nonempty. Preserve other fields required by the original schema. Use JSON strings and arrays of the requested types, and escape all string contents correctly.
Retain every distinct supported candidate across the entire supplied source, including conditions, negations and uncertainty. Do not shorten the source scope, invent facts or return a partial result. Return only the complete JSON object.
</candidate_regeneration>
Local validation (positions refer to cleaned JSON): `

func wikiCandidateCorrectionInstruction(validation wikiCandidateValidation) string {
	return wikiCandidateRepairInstruction + validation.String()
}

// Only fixed codes, known fields and numeric positions enter errors/prompts.
// Never forward encoding/json's error string, which can contain model values.
func parseWikiCandidateOutput(raw string) (combinedExtraction, *wikiCandidateValidation) {
	var result combinedExtraction
	if err := json.Unmarshal([]byte(cleanLLMJSON(raw)), &result); err != nil {
		validation := &wikiCandidateValidation{Code: "json_decode_failed",
			ContentRunes: utf8.RuneCountInString(raw), ContentShape: wikiContentShape(raw)}
		var syntax *json.SyntaxError
		var mismatch *json.UnmarshalTypeError
		if errors.As(err, &syntax) {
			validation.Code, validation.JSONOffset = "invalid_json_syntax", syntax.Offset
		} else if errors.As(err, &mismatch) {
			validation.Code, validation.JSONOffset = "json_type_mismatch", mismatch.Offset
		}
		return combinedExtraction{}, validation
	}
	if result.Entities == nil {
		return combinedExtraction{}, &wikiCandidateValidation{Code: "missing_entities_array", Field: "entities"}
	}
	if result.Concepts == nil {
		return combinedExtraction{}, &wikiCandidateValidation{Code: "missing_concepts_array", Field: "concepts"}
	}
	if validation := wikiCandidateFieldValidation(result); validation != nil {
		return combinedExtraction{}, validation
	}
	return result, nil
}

func wikiCandidateFieldValidation(result combinedExtraction) *wikiCandidateValidation {
	for _, group := range []struct {
		field, prefix string
		items         []extractedItem
	}{{"entities", "entity/", result.Entities}, {"concepts", "concept/", result.Concepts}} {
		for index, item := range group.items {
			if strings.TrimSpace(item.Name) == "" {
				return &wikiCandidateValidation{Code: "candidate_name_empty", Field: group.field + ".name", Index: &index}
			}
			if !strings.HasPrefix(item.Slug, group.prefix) || strings.TrimSpace(strings.TrimPrefix(item.Slug, group.prefix)) == "" {
				return &wikiCandidateValidation{Code: "candidate_slug_invalid", Field: group.field + ".slug", Index: &index}
			}
		}
	}
	return nil
}

func (e *wikiOutputError) Error() string {
	return "Wiki 模型输出不完整或格式不正确：" + e.reason
}

func validateWikiResponse(response *types.ChatResponse) error {
	if response == nil {
		return &wikiOutputError{reason: "empty response"}
	}
	switch strings.ToLower(strings.TrimSpace(response.FinishReason)) {
	case "length", "max_tokens", "max_output_tokens", types.FinishReasonIncomplete:
		return &wikiOutputError{reason: "generation stopped before completion", incomplete: true}
	}
	if strings.TrimSpace(response.Content) == "" {
		// A provider that filters or drops the request returns an empty body; say
		// so instead of reporting a JSON problem that never happened.
		reason := "provider_empty_response"
		if finish := strings.TrimSpace(response.FinishReason); finish != "" {
			reason += ":" + finish
		}
		return &wikiOutputError{reason: reason}
	}
	return nil
}

const wikiPageRepairInstruction = `<complete_page_regeneration>
The previous request did not produce a complete usable page. Regenerate the WHOLE revised page from the SAME source blocks and existing page supplied above. Do not continue, reconstruct, quote, or rely on any previous partial output.
Return only the required SUMMARY line and complete page body in the requested language. Keep this one page within 12000 Unicode characters. Remove repetition, lengthy examples and stylistic padding, not distinct supported facts relevant to this page.
Preserve all material negations, conditions, exceptions, uncertainty, quantities, dates, source attribution and valid reference handles. Keep existing supported facts unless the supplied retractions explicitly remove them. Respect the same page identity, source scope, additions and retractions; do not invent facts or broaden the topic.
Use concise sentences and compact lists. Finish every sentence and section; no ellipses, placeholders or promises to continue. If a complete faithful revision cannot fit, return exactly WIKI_PAGE_REPAIR_UNSAFE instead of omitting facts or returning a partial page.
</complete_page_regeneration>`

// A rejected page gets at most one fresh request, never a continuation of the
// partial response. Base prompts remain unchanged so checkpoint scopes survive.
func (s *wikiIngestService) generateWikiPageModification(ctx context.Context, model chat.Chat, data map[string]string, continuation ...string) (string, error) {
	if err := validateWikiPageInputBudget(data); err != nil {
		return "", err
	}
	instruction := strings.Join(continuation, "\n\n")
	content, err := s.generateWithTemplateRequest(ctx, model, agent.WikiPageModifyUserPrompt, data, instruction, wikiLLMMaxAttempts)
	if err == nil {
		err = validateWikiPageBody(content, false)
	}
	var outputErr *wikiOutputError
	if err == nil || !errors.As(err, &outputErr) || ctx.Err() != nil {
		return content, err
	}
	content, err = s.generateWithTemplateRequest(ctx, model, agent.WikiPageModifyUserPrompt, data, strings.TrimSpace(instruction+"\n\n"+wikiPageRepairInstruction), 1)
	if err != nil {
		return "", fmt.Errorf("complete wiki page regeneration failed: %w", err)
	}
	if strings.Contains(content, "WIKI_PAGE_REPAIR_UNSAFE") || len([]rune(content)) > 12000 {
		return "", &wikiOutputError{reason: "complete wiki page regeneration exceeded the safe output budget"}
	}
	if err := validateWikiPageBody(content, true); err != nil {
		return "", err
	}
	return content, nil
}

// Page synthesis has a separate input budget from source reading. Reading all
// windows must never culminate in an unbounded request for a broad topic. This
// is an explicit unresolved outcome: no evidence is dropped and no model call
// is made. The caller preserves already published pages and checkpoints.
const wikiPageInputRunes = 65536

func validateWikiPageInputBudget(data map[string]string) error {
	chars := 0
	for _, value := range data {
		chars += utf8.RuneCountInString(value)
	}
	if chars > wikiPageInputRunes {
		return fmt.Errorf("Wiki 页面综合已暂停：该主题输入共 %d 字符，超过单次预算 %d；请拆分来源文档或主题。原文和已完成断点保留，未截断证据或标记完成", chars, wikiPageInputRunes)
	}
	return nil
}

// wikiPageRepairUnsafeMarker is the sentinel the page prompts tell the model to emit
// when it judges the request unsafe to fulfil. It is a refusal, not malformed JSON, so
// it must be reported separately from a structurally incomplete page.
const wikiPageRepairUnsafeMarker = "WIKI_PAGE_REPAIR_UNSAFE"

func validateWikiPageBody(content string, requireSummary bool) error {
	summary, body := splitSummaryLine(content)
	runes := utf8.RuneCountInString(content)
	switch {
	case strings.Contains(content, wikiPageRepairUnsafeMarker):
		return &wikiOutputError{reason: fmt.Sprintf(
			"the model reported the page as unsafe to generate (%s, runes=%d)", wikiPageRepairUnsafeMarker, runes)}
	case strings.TrimSpace(body) == "":
		return &wikiOutputError{reason: fmt.Sprintf(
			"complete wiki page requires a nonempty summary and body (empty body, runes=%d)", runes)}
	case requireSummary && summary == "":
		return &wikiOutputError{reason: fmt.Sprintf(
			"complete wiki page requires a nonempty summary and body (missing SUMMARY line, runes=%d)", runes)}
	}
	return nil
}

// Large candidate inputs start at the first split level. Each complete invalid
// output gets one fresh correction before further splitting. The same two-level
// bound applies, and every leaf must succeed; provider errors never initiate
// output correction or splitting.
func (s *wikiIngestService) generateWikiExtraction(
	ctx context.Context, model chat.Chat, prompt string, data map[string]string,
) (combinedExtraction, error) {
	result, err := s.generateWikiFreshExtraction(ctx, model, prompt, data)
	if err != nil {
		return combinedExtraction{}, err
	}
	return normalizeWikiFreshCandidates(result)
}

// Retain raw routing proposals across leaf/window boundaries until every
// source-scoped identity has been checked and the complete set is available.
func (s *wikiIngestService) generateWikiFreshExtraction(
	ctx context.Context, model chat.Chat, prompt string, data map[string]string,
) (combinedExtraction, error) {
	candidateOnly := prompt == agent.WikiCandidateSlugPrompt
	generationInstruction := ""
	if candidateOnly {
		generationInstruction = wikiCandidateGenerationInstruction
	}
	var extract func(map[string]string, int) (combinedExtraction, error)
	split := func(input map[string]string, content []rune, depth int) (combinedExtraction, error) {
		// Split at a nearby paragraph boundary without losing any source
		// character. Proactive and recovery splits share the same depth budget.
		middle := len(content) / 2
		for i := middle; i > middle-500; i-- {
			if content[i] == '\n' {
				middle = i + 1
				break
			}
		}
		merged := combinedExtraction{Entities: []extractedItem{}, Concepts: []extractedItem{}}
		for index, part := range []string{string(content[:middle]), string(content[middle:])} {
			partInput := make(map[string]string, len(input))
			for key, value := range input {
				partInput[key] = value
			}
			partInput["Content"] = part
			parsed, partErr := extract(partInput, depth+1)
			if partErr != nil {
				return combinedExtraction{}, fmt.Errorf("extraction partition %d/2: %w", index+1, partErr)
			}
			merged.Entities = append(merged.Entities, parsed.Entities...)
			merged.Concepts = append(merged.Concepts, parsed.Concepts...)
		}
		return merged, nil
	}
	extract = func(input map[string]string, depth int) (combinedExtraction, error) {
		if err := ctx.Err(); err != nil {
			return combinedExtraction{}, err
		}
		content := []rune(input["Content"])
		if candidateOnly && depth == 0 && len(content) > wikiCandidateProactiveSplitRunes {
			return split(input, content, depth)
		}
		var result combinedExtraction
		repairInstruction := generationInstruction
		raw, err := s.generateWithTemplateRequest(ctx, model, prompt, input, repairInstruction, wikiLLMMaxAttempts)
		if err == nil {
			var validation *wikiCandidateValidation
			result, validation = parseWikiCandidateOutput(raw)
			if validation != nil {
				err = &wikiOutputError{reason: validation.String()}
				if ctx.Err() != nil {
					return combinedExtraction{}, ctx.Err()
				}
				repairInstruction = wikiCandidateCorrectionInstruction(*validation)
				if generationInstruction != "" {
					repairInstruction = generationInstruction + "\n\n" + repairInstruction
				}
				raw, err = s.generateWithTemplateRequest(ctx, model, prompt, input, repairInstruction, 1)
				if err == nil {
					result, validation = parseWikiCandidateOutput(raw)
					if validation != nil {
						err = &wikiOutputError{reason: validation.String()}
					}
				}
			}
		}
		if err == nil {
			if candidateOnly {
				original := result
				pruned, pruning := pruneWikiCandidateQuality(result, input["Content"])
				if pruning.Dropped() > 0 {
					// Counted, content-free evidence that a bounded policy removed
					// candidates instead of discarding the document.
					logger.Infof(ctx, "wiki candidate pruning: policy=%s duplicates=%d unmatched=%d capped=%d kept=%d dropped=%d",
						wikiCandidateQualityVersion, pruning.Duplicates, pruning.Unmatched, pruning.Capped, pruning.Kept, pruning.Dropped())
				}
				result = pruned
				// An empty candidate set is legitimate for a window that holds nothing
				// new. Pruning must however never turn candidates into nothing: if
				// everything was dropped, the identity failure that caused it is
				// reported instead of an empty extraction.
				if len(result.Entities)+len(result.Concepts) == 0 && len(original.Entities)+len(original.Concepts) > 0 {
					if identityErr := validateWikiFreshSourceIdentities(original, input["Content"]); identityErr != nil {
						return combinedExtraction{}, identityErr
					}
					return combinedExtraction{}, wikiCandidateQualityRejected("candidate_set_empty_after_pruning")
				}
			}
			if identityErr := validateWikiFreshSourceIdentities(result, input["Content"]); identityErr != nil {
				corrected, repairErr := s.repairWikiFreshSourceIdentities(ctx, model, prompt, input, result)
				if repairErr == nil {
					return corrected, nil
				}
				for _, stop := range []error{ErrWikiRecoveryBudgetStopped, context.DeadlineExceeded, context.Canceled} {
					if errors.Is(repairErr, stop) || errors.Is(context.Cause(ctx), stop) {
						return combinedExtraction{}, errors.Join(identityErr, stop)
					}
				}
				// Partial degradation: identities that cannot be supported by this
				// source are dropped and counted rather than failing the document.
				dropped, droppedCount := dropRejectedWikiSourceIdentities(result, input["Content"])
				dropped.CandidatePruning.Unmatched += droppedCount
				captureWikiIdentityFailure(ctx, model.GetModelID(), prompt, input, repairInstruction, raw, identityErr)
				if len(dropped.Entities)+len(dropped.Concepts) == 0 {
					return combinedExtraction{}, identityErr
				}
				logger.Infof(ctx, "wiki identity pruning: policy=%s code=%s dropped=%d kept=%d",
					wikiSourceIdentityRepairVersion, wikiIdentityFailureCode(identityErr), droppedCount,
					len(dropped.Entities)+len(dropped.Concepts))
				return dropped, nil
			}
			return result, nil
		}
		var outputErr *wikiOutputError
		if !errors.As(err, &outputErr) || ctx.Err() != nil || len(content) < 2000 {
			return combinedExtraction{}, err
		}
		if depth >= 2 || (depth == 1 && (!outputErr.incomplete || len(content) < 8000)) {
			return combinedExtraction{}, err
		}
		return split(input, content, depth)
	}
	return extract(data, 0)
}