Files

3319 lines
122 KiB
Go
Raw Permalink Normal View History

package service
import (
"context"
"encoding/json"
"errors"
"fmt"
"regexp"
"sort"
"strings"
"sync"
"sync/atomic"
"text/template"
"time"
"unicode/utf8"
"github.com/Tencent/WeKnora/internal/agent"
apprepo "github.com/Tencent/WeKnora/internal/application/repository"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/models/chat"
"github.com/Tencent/WeKnora/internal/models/requestbudget"
"github.com/Tencent/WeKnora/internal/searchutil"
"github.com/Tencent/WeKnora/internal/tracing/langfuse"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
"github.com/google/uuid"
"github.com/hibiken/asynq"
"github.com/redis/go-redis/v9"
"golang.org/x/sync/singleflight"
"gorm.io/gorm"
)
// ErrWikiIngestConcurrent is returned by the wiki ingest handler in Lite mode
// when another batch is already running for the same KB (the in-process
// liteLocks guard is held). The asynq/sync executor's RetryDelayFunc uses
// errors.Is on this sentinel to apply a short, fixed retry delay instead of
// exponential backoff, so the deferred batch retries promptly once the active
// one releases. Standard (Redis) mode no longer takes an exclusive per-KB
// lock (Phase 3) and never returns this.
var ErrWikiIngestConcurrent = errors.New("concurrent wiki task active")
const (
// maxContentForWiki limits ONE original-source window, never the total
// document. Long documents continue through durable full-read windows.
maxContentForWiki = 32768
// --- Phase 3: concurrent per-KB batches (standard/Redis mode) ---------
//
// Phase 3 removed the exclusive per-KB Redis "batch in progress" lock
// (wiki:active:<kbID>). Standard mode now allows concurrent batches for
// one KB, kept safe by row claiming (claimPendingList) plus per-slug
// reduce locks (withSlugLock). Lite mode still serializes per KB via the
// in-process liteLocks map.
//
// wikiClaimStaleAfter is how long a claimed-but-undrained ingest row
// waits before another worker may re-claim it. It MUST exceed the asynq
// task Timeout for wiki:ingest (60m) so a still-running batch's rows are
// never stolen mid-flight — only genuinely crashed/abandoned claims are
// recovered. This preserves the pre-claim crash behaviour where a dead
// batch's undeleted rows were simply re-peeked by the next trigger.
wikiClaimStaleAfter = 90 * time.Minute
// wikiSlugLockPrefix guards read-modify-write on a single shared wiki
// page (entity/concept/summary/index) so two concurrent batches for the
// same KB can't lost-update the same slug. Key: wiki:slug:{kbID}:{slug}.
wikiSlugLockPrefix = "wiki:slug:"
// wikiIdentityClaimPrefix reserves the slug chosen for one normalized
// (page type, display title) identity while concurrent map phases are
// still running. The existing per-slug lock cannot help when two models
// emit different slugs for the same title, because those reducers lock
// different keys. A short-lived identity claim makes both batches use
// the same slug before summaries and page updates are materialized.
wikiIdentityClaimPrefix = "wiki:identity:"
wikiIdentityClaimTTL = 2 * time.Hour
// wikiIdentityClaimScript atomically GET-or-SET (or overwrite when
// ARGV[3] is "1"). A valid existing slug wins and has its TTL refreshed.
// Missing or corrupt values are replaced so callers cannot diverge on a
// dirty key. KEYS[1]=claim key; ARGV = proposed slug, TTL seconds,
// authoritative ("0"/"1"), required slug prefix.
wikiIdentityClaimScript = `
local proposed = ARGV[1]
local ttl = tonumber(ARGV[2])
local prefix = ARGV[4]
if ARGV[3] == '1' then
redis.call('SET', KEYS[1], proposed, 'EX', ttl)
return proposed
end
local existing = redis.call('GET', KEYS[1])
if type(existing) == 'string' and string.sub(existing, 1, #prefix) == prefix then
redis.call('EXPIRE', KEYS[1], ttl)
return existing
end
redis.call('SET', KEYS[1], proposed, 'EX', ttl)
return proposed
`
// wikiSlugLockTTL bounds the per-slug lock so a crashed reducer can't
// wedge a hot page forever. Active reducers renew this token-owned lease.
wikiSlugLockTTL = 5 * time.Minute
// wikiSlugLockWait / Poll bound how long a reduce goroutine blocks
// waiting for a contended slug before falling back to a best-effort
// skip (matching the pre-existing reduce-failure semantics).
wikiSlugLockWait = 2 * time.Minute
wikiSlugLockPoll = 50 * time.Millisecond
// --- Phase 4: per-KB in-flight cap (standard/Redis mode) --------------
//
// Since Phase 3 lets many batches for one KB run at once, a single KB's
// bulk import could otherwise occupy the entire wiki pool and starve
// other KBs. Each running batch reserves a slot in a per-KB Redis sorted
// set (score = expiry). ZREMRANGEBYSCORE purges slots left behind by a
// crashed worker, so the cap is self-healing without an explicit lock.
//
// wikiInflightPrefix keys the per-KB slot set: wiki:inflight:{kbID}.
wikiInflightPrefix = "wiki:inflight:"
// wikiInflightDefault is the fallback cap when WikiConfig.IngestMaxInflight
// is unset. 4 leaves half of the default 8-worker wiki pool for other
// KBs while still giving one KB solid parallelism.
wikiInflightDefault = 4
// wikiInflightTTL is how long a reserved slot survives without renewal.
// MUST exceed wikiInflightRenew by a comfortable margin so a single
// missed renew (GC pause, Redis blip) doesn't drop a live batch's slot.
wikiInflightTTL = 90 * time.Second
// wikiInflightRenew is how often a running batch bumps its slot's expiry.
wikiInflightRenew = 30 * time.Second
// wikiInflightBackoff delays the follow-up trigger when a batch is turned
// away by the cap, so it retries after a slot plausibly frees.
wikiInflightBackoff = 10 * time.Second
// wikiIngestDelay is how long to wait after a document is added before
// the batch task fires. Debounces rapid uploads.
wikiIngestDelay = 30 * time.Second
// wikiFollowUpDelay is the normal light debounce before a follow-up
// batch drains the remaining pending rows. Standard mode already fans a
// KB's backlog across concurrent claiming batches, so this is just a
// short breather, not a lock-release wait.
wikiFollowUpDelay = 5 * time.Second
// wikiRateLimitBackoff is the (much longer) follow-up delay used when a
// batch's failures were caused by upstream rate limiting (HTTP 429 /
// quota). Retrying a rate-limited document at the normal 5s cadence just
// throws more requests at an already-saturated rpm budget — every retry
// re-issues extract + classify + summary calls for the doc — which keeps
// the limiter tripped and drags the whole KB's ingest out. Backing off
// gives the per-minute window time to reset before we try again.
wikiRateLimitBackoff = 60 * time.Second
// wikiMaxDocsPerBatch limits how many documents a single batch processes.
// Prevents unbounded execution time. Remaining ops stay in
// task_pending_ops and are picked up by the follow-up task.
wikiMaxDocsPerBatch = 5
// wikiMaxFailRetries is the maximum number of times a single document op
// may be re-attempted via requeueFailedOps before it is permanently
// archived to task_dead_letters. 5 retries ≈ five full batch cycles
// (each with a ~30 s delay), giving transient LLM errors a fair chance
// to recover without letting a persistently-broken doc clog the queue
// indefinitely.
wikiMaxFailRetries = 5
// wikiIngestMaxRetry controls asynq retry budget for wiki:ingest tasks.
// Keep this moderate: lock conflicts already retry every 15s via
// asynqRetryDelayFunc, and follow-up/retract paths fire quickly.
wikiIngestMaxRetry = 10
// wikiDeletedKeyPrefix is the Redis key prefix for "recently deleted
// knowledge" tombstones. Key: wiki:deleted:{kbID}:{knowledgeID}. Written
// by cleanupWikiOnKnowledgeDelete so that any wiki_ingest task still in
// flight (or queued) for this knowledge can fast-path skip without
// hitting the DB. TTL > wikiIngestDelay so it's guaranteed to outlast
// any in-flight ingest.
wikiDeletedKeyPrefix = "wiki:deleted:"
// wikiDeletedTTL bounds how long we remember a deletion. Must comfortably
// exceed the longest plausible ingest run (LLM extraction + reduce).
wikiDeletedTTL = 1 * time.Hour
// wikiLLMMaxAttempts is the total attempt count (initial + retries) for
// every LLM call routed through generateWithTemplate. 3 was chosen to
// absorb transient 504/timeouts from upstream gateways without
// materially prolonging task runtime when the remote is genuinely down.
wikiLLMMaxAttempts = 3
// A single stalled call must leave room for bounded retries within the
// checkpoint slice. The chat transport otherwise inherits that full slice.
wikiLLMCallTimeout = 5 * time.Minute
// wikiLLMMaxTokens is the completion-token budget for every LLM call
// routed through generateWithTemplate. Combined wiki extraction emits a
// single large JSON document (entities + concepts + details). When
// MaxTokens is left at 0, OpenAI-compatible clients omit max_tokens and
// providers such as DeepSeek apply a default of 8192 — long extracts are
// truncated mid-JSON with finish_reason=length, and parse fails with
// "unexpected end of JSON input" (EXTRACT_FAILED). Raising the budget to
// 32768 matches verified complete outputs for large Chinese policy docs;
// shorter replies still stop early via finish_reason=stop. See #2604.
wikiLLMMaxTokens = 65536
// wikiLLMBackoffBase is the base delay for the exponential backoff
// between retry attempts. The nth retry waits base << (n-1) — so with
// a 2s base we wait 2s, 4s, 8s between attempts.
wikiLLMBackoffBase = 2 * time.Second
// wikiTaskType is the task_type stamp used in task_pending_ops and
// task_dead_letters rows for this pipeline. Stable across the lifetime
// of any pending op so the follow-up consumer can pull it back.
wikiTaskType = "wiki:ingest"
// wikiTaskScope is the scope used by both pending ops and dead letters.
// Wiki ingest is per-KB, so every op is scoped to a knowledge_base.
wikiTaskScope = types.TaskScopeKnowledgeBase
// --- Finalize lane (Phase 1: debounced KB-global convergence) ---------
//
// The KB-global convergence steps (index-intro rebuild, dead-link
// cleanup, cross-link injection) used to run at the tail of EVERY
// 5-doc ingest batch. On a bulk import that meant O(docs/batchSize)
// index rebuilds — each an LLM call plus an index-page write. They are
// now deferred to a debounced per-KB wiki:finalize task: each ingest
// batch records what changed into a separate lane of task_pending_ops
// (task_type=wiki:finalize) and schedules a coalesced trigger, so a
// burst of N documents rebuilds the index ONCE at the end.
// wikiFinalizeTaskType is the task_pending_ops lane for finalize work.
// Keyed by (task_type, scope, scope_id) like the ingest lane, but with a
// distinct task_type so PeekBatch never mixes the two.
wikiFinalizeTaskType = "wiki:finalize"
// wikiFinalizeOpSlug rows carry one affected page slug (+ its fresh title
// when this batch wrote it) for the dead-link / cross-link passes.
wikiFinalizeOpSlug = "slug"
// wikiFinalizeOpChange rows carry a doc-level add/remove change entry for
// the index-intro change description.
wikiFinalizeOpChange = "change"
// wikiFinalizeOpFolderPrune rows carry folders that may have become empty
// after a document retract. Keeping this in the durable finalize lane lets
// us wait until every ingest op for the KB has settled before deleting the
// directories; taxonomy planning creates folders before reduce writes the
// corresponding pages, so pruning any earlier can invalidate in-flight
// folder assignments.
wikiFinalizeOpFolderPrune = "folder_prune"
wikiFinalizeAdded = "added"
wikiFinalizeRemoved = "removed"
// wikiFinalizeDelay debounces the finalize trigger so multiple ingest
// batches within the window coalesce into a single index rebuild.
wikiFinalizeDelay = 20 * time.Second
// wikiFinalizeMaxRows caps how many finalize rows one run drains. A run
// that hits the cap reschedules itself for the remainder.
wikiFinalizeMaxRows = 5000
// wikiFinalizeLockTTL / Renew guard against two finalize runs for the
// same KB writing the index page concurrently (belt-and-suspenders on
// top of the asynq.TaskID coalescing).
wikiFinalizeLockTTL = 60 * time.Second
wikiFinalizeLockRenew = 20 * time.Second
// Folder cleanup is maintenance, not user-blocking work. When an ingest is
// still active, retry slowly so pruning never competes with the primary wiki
// pipeline for worker capacity.
wikiFolderPruneRetryDelay = 1 * time.Minute
// wikiIngestCleanupTimeout bounds detached tail cleanup after the asynq
// task context has been cancelled or hit its timeout.
wikiIngestCleanupTimeout = 10 * time.Second
)
// wikiFinalizeChange is a doc-level add/remove entry for the index-intro
// change description, persisted as a finalize-lane row payload.
type wikiFinalizeChange struct {
Action string `json:"action"` // wikiFinalizeAdded | wikiFinalizeRemoved
DocTitle string `json:"doc_title,omitempty"`
DocSummary string `json:"doc_summary,omitempty"`
}
// wikiFinalizeRow is the JSON payload of a task_pending_ops row in the
// finalize lane. Exactly one of {Slug, Change, FolderIDs} is set,
// distinguished by the row's Op column.
type wikiFinalizeRow struct {
types.WikiRecoveryIdentity
Slug string `json:"slug,omitempty"`
Title string `json:"title,omitempty"`
Change *wikiFinalizeChange `json:"change,omitempty"`
FolderIDs []string `json:"folder_ids,omitempty"`
}
// WikiDeletedTombstoneKey returns the Redis key used to mark a knowledge as
// recently deleted, so wiki_ingest tasks in flight can short-circuit. Exposed
// so knowledgeService.cleanupWikiOnKnowledgeDelete can write the same key
// without duplicating the format string.
func WikiDeletedTombstoneKey(kbID, knowledgeID string) string {
return wikiDeletedKeyPrefix + kbID + ":" + knowledgeID
}
// WikiIngestPayload is the asynq task payload for wiki ingest batch trigger.
// The actual document IDs are stored in the task_pending_ops table; this
// payload only carries the trigger metadata so the worker can resolve
// the queue tuple (task_type, scope, scope_id) and process whatever rows
// are queued under it.
type WikiIngestPayload struct {
types.TracingContext
TenantID uint64 `json:"tenant_id"`
KnowledgeBaseID string `json:"knowledge_base_id"`
Language string `json:"language,omitempty"`
}
// WikiRetractPayload is the asynq task payload for wiki content retraction
type WikiRetractPayload struct {
types.TracingContext
TenantID uint64 `json:"tenant_id"`
KnowledgeBaseID string `json:"knowledge_base_id"`
KnowledgeID string `json:"knowledge_id"`
DocTitle string `json:"doc_title"`
DocSummary string `json:"doc_summary,omitempty"` // one-line summary of the deleted document
Language string `json:"language,omitempty"`
PageSlugs []string `json:"page_slugs"`
FolderIDs []string `json:"folder_ids,omitempty"`
}
const (
WikiOpIngest = "ingest"
WikiOpRetract = "retract"
)
// WikiPendingOp represents a single operation queued in task_pending_ops
// under task_type="wiki:ingest". The struct is the JSON payload of the
// task_pending_ops row; the surrounding (task_type, scope, scope_id,
// dedup_key) fields live as separate columns and are not serialized
// here.
//
// dbID is the auto-increment primary key of the task_pending_ops row
// the op was loaded from. PeekBatch fills it; consumers carry it
// through Map/Reduce so DeleteByIDs (after consume) and IncrFailCount
// (after failure) can address the right row. It is intentionally
// unexported and excluded from JSON so the persisted payload does not
// duplicate the column.
type WikiPendingOp struct {
types.WikiRecoveryIdentity
Op string `json:"op"`
KnowledgeID string `json:"knowledge_id"`
// Ingest fields
Language string `json:"language,omitempty"`
// Retract fields
DocTitle string `json:"doc_title,omitempty"`
DocSummary string `json:"doc_summary,omitempty"`
PageSlugs []string `json:"page_slugs,omitempty"`
FolderIDs []string `json:"folder_ids,omitempty"`
// dbID is set by peekPendingList from task_pending_ops.id. Zero in
// constructions made outside the queue (e.g. legacy tests).
dbID int64 `json:"-"`
recoveryConflict bool `json:"-"`
}
// wikiIngestService handles the LLM-powered wiki generation pipeline.
//
// Durable state lives in two places:
// - task_pending_ops (rows tagged task_type="wiki:ingest", scope=
// "knowledge_base"): the per-document op queue. Replaces the
// legacy Redis wiki:pending:<kbID> list, which was vulnerable to
// 24h TTL eviction at 4w-document scale.
// - task_dead_letters: in-batch failures that exhausted
// wikiMaxFailRetries land here. The asynq dead-letter middleware
// also writes asynq-level archived rows here uniformly across
// every task type.
//
// Redis is still used for the per-KB active-batch lock
// (wiki:active:<kbID>) and the delete tombstone (wiki:deleted:<...>),
// both of which are correctness-critical short-lived flags rather
// than data the system should survive without.
type wikiIngestService struct {
wikiService interfaces.WikiPageService
kbService interfaces.KnowledgeBaseService
knowledgeSvc interfaces.KnowledgeService
knowledgeRepo interfaces.KnowledgeRepository
chunkRepo interfaces.ChunkRepository
modelService interfaces.ModelService
checkpointDB *gorm.DB
task interfaces.TaskEnqueuer
audit interfaces.AuditLogService
pendingRepo interfaces.TaskPendingOpsRepository
deadLetterRepo interfaces.TaskDeadLetterRepository
redisClient *redis.Client // nil in Lite mode (no Redis)
// spanTracker lets per-document map work surface as a
// postprocess.wiki subspan in the knowledge trace tree. Async
// batch design means we look up the parent attempt by knowledge
// id at run-time (LatestAttempt) rather than carrying it in the
// asynq payload, which is per-KB and would otherwise be ambiguous
// for the 5-docs-per-batch fan-out.
spanTracker SpanTracker
// liteLocks provides per-KB mutual exclusion in Lite mode (no Redis).
// Keys are kbID strings; values are unused (presence = locked).
liteLocks sync.Map
// liteFinalizeLocks is the Lite-mode counterpart of the Redis
// wiki:finalize:active:<kbID> lock, keeping two finalize runs for the
// same KB from overlapping when there is no Redis.
liteFinalizeLocks sync.Map
// llmRequests coalesces byte-identical concurrent prompts within this
// process. Keys include tenant and model to preserve isolation.
llmRequests singleflight.Group
// promptWarmups serializes only the first request for a reusable Wiki page
// prefix. Other prefixes and already-warmed cohorts stay parallel.
promptWarmups sync.Map
}
type wikiPromptWarmup struct {
done chan struct{}
once sync.Once
}
// NewWikiIngestService creates a new wiki ingest service
func NewWikiIngestService(
wikiService interfaces.WikiPageService,
kbService interfaces.KnowledgeBaseService,
knowledgeSvc interfaces.KnowledgeService,
knowledgeRepo interfaces.KnowledgeRepository,
chunkRepo interfaces.ChunkRepository,
modelService interfaces.ModelService,
task interfaces.TaskEnqueuer,
audit interfaces.AuditLogService,
pendingRepo interfaces.TaskPendingOpsRepository,
deadLetterRepo interfaces.TaskDeadLetterRepository,
redisClient *redis.Client,
spanTracker SpanTracker,
checkpointDB *gorm.DB,
) interfaces.TaskHandler {
svc := &wikiIngestService{
wikiService: wikiService,
kbService: kbService,
knowledgeSvc: knowledgeSvc,
knowledgeRepo: knowledgeRepo,
chunkRepo: chunkRepo,
modelService: modelService,
task: task,
audit: audit,
pendingRepo: pendingRepo,
deadLetterRepo: deadLetterRepo,
redisClient: redisClient,
spanTracker: spanTracker,
checkpointDB: checkpointDB,
}
return svc
}
// tracker returns a non-nil span tracker so callers don't have to
// nil-check on every Begin/End. Matches the noopSpanTracker pattern
// used elsewhere (see knowledgeService.tracker, KnowledgePostProcessService.tracker).
func (s *wikiIngestService) tracker() SpanTracker {
if s.spanTracker == nil {
return noopSpanTracker{}
}
return s.spanTracker
}
// beginWikiSubspan opens a postprocess.wiki subspan for this document
// under the knowledge's most recent attempt. Returns nil when there is
// no parse attempt to attach to (e.g. a wiki ingest fired from a manual
// reparse path that never went through the tracker) — callers must
// pair every begin with a tolerant end / fail / skip below.
//
// Lookups are by `LatestAttempt(knowledgeID)` because the asynq task
// payload (WikiIngestPayload) is KB-scoped and carries no per-doc
// attempt — see the type's comment for the batch architecture.
func (s *wikiIngestService) beginWikiSubspan(ctx context.Context, knowledgeID string, input types.JSONMap) *Span {
if knowledgeID == "" {
return nil
}
attempt := s.tracker().LatestAttempt(ctx, knowledgeID)
if attempt <= 0 {
return nil
}
parent := s.tracker().LookupStage(ctx, knowledgeID, attempt, types.StagePostProcess)
if parent == nil {
return nil
}
return s.tracker().BeginSubSpan(ctx, parent, "postprocess.wiki", types.SpanKindSubSpan, input)
}
// EnqueueWikiIngest queues a document for wiki ingestion. The returned bool
// reports whether the durable pending op was persisted. A trigger enqueue
// error may therefore be returned together with true; callers can retry only
// the KB-scoped trigger without appending a duplicate operation.
//
// Architecture: each upload inserts one row into task_pending_ops
// (task_type="wiki:ingest", scope="knowledge_base", scope_id=kbID,
// dedup_key=knowledgeID), then schedules a debounced asynq trigger task.
// When the trigger fires, the worker peeks a batch from
// task_pending_ops, processes it, deletes consumed rows, and (if more
// remain) schedules a follow-up. Multiple debounced triggers within the
// 30s window all coalesce: the first one to acquire the per-KB active
// lock drains the batch; subsequent ones see an empty queue and exit.
//
// Lite mode (no Redis) still works as long as Postgres is reachable —
// the queue lives in PG, only the active-batch lock is Redis-only and
// has a process-local fallback (liteLocks) inside the worker.
func enqueueWikiPendingOp(
ctx context.Context,
pendingRepo interfaces.TaskPendingOpsRepository,
op *types.TaskPendingOp,
) (bool, error) {
if pendingRepo == nil {
return true, nil
}
if guard, ok := pendingRepo.(interfaces.TaskPendingOpsKnowledgeBaseGuard); ok {
return guard.EnqueueIfKnowledgeBaseActive(ctx, op)
}
if err := pendingRepo.Enqueue(ctx, op); err != nil {
return false, err
}
return true, nil
}
func EnqueueWikiIngest(
ctx context.Context,
task interfaces.TaskEnqueuer,
pendingRepo interfaces.TaskPendingOpsRepository,
tenantID uint64,
kbID, knowledgeID string,
) (bool, error) {
pendingOp, err := newWikiIngestPendingOp(ctx, tenantID, kbID, knowledgeID)
// Persist the pending op. A re-ingest of the same knowledge id while
// a previous op is still queued simply appends another row; the
// peekPendingList consumer collapses by dedup_key (== knowledge_id),
// keeping the LATEST op for each knowledge — matching the legacy
// "RPush + reverse-dedupe" semantics.
if err != nil {
logger.Warnf(ctx, "wiki ingest: failed to marshal pending op for %s: %v", knowledgeID, err)
return false, err
}
accepted, err := enqueueWikiPendingOp(ctx, pendingRepo, pendingOp)
if err != nil {
logger.Warnf(ctx, "wiki ingest: failed to enqueue pending op for %s: %v", knowledgeID, err)
return false, fmt.Errorf("enqueue wiki ingest pending op: %w", err)
}
if !accepted {
logger.Infof(ctx, "wiki ingest: skip enqueue for deleted KB %s", kbID)
return false, nil
}
if err := enqueueWikiIngestTrigger(ctx, task, tenantID, kbID); err != nil {
return true, err
}
return true, nil
}
func newWikiIngestPendingOp(
ctx context.Context,
tenantID uint64,
kbID, knowledgeID string,
) (*types.TaskPendingOp, error) {
lang := types.LanguageFromContextOrDefault(ctx)
op := WikiPendingOp{
Op: WikiOpIngest,
KnowledgeID: knowledgeID,
Language: lang,
}
payloadBytes, err := json.Marshal(op)
if err != nil {
return nil, fmt.Errorf("marshal wiki ingest pending op: %w", err)
}
return &types.TaskPendingOp{
TenantID: tenantID,
TaskType: wikiTaskType,
Scope: wikiTaskScope,
ScopeID: kbID,
Op: WikiOpIngest,
DedupKey: knowledgeID,
Payload: payloadBytes,
}, nil
}
func enqueueWikiIngestTrigger(
ctx context.Context,
task interfaces.TaskEnqueuer,
tenantID uint64,
kbID string,
) error {
lang := types.LanguageFromContextOrDefault(ctx)
trigger := WikiIngestPayload{
TenantID: tenantID,
KnowledgeBaseID: kbID,
Language: lang,
}
langfuse.InjectTracing(ctx, &trigger)
triggerBytes, err := json.Marshal(trigger)
if err != nil {
logger.Warnf(ctx, "wiki ingest: failed to marshal trigger task: %v", err)
return fmt.Errorf("marshal wiki ingest trigger: %w", err)
}
if task == nil {
return errors.New("enqueue wiki ingest trigger: task enqueuer is nil")
}
t := asynq.NewTask(types.TypeWikiIngest, triggerBytes,
asynq.Queue(types.QueueWiki),
asynq.MaxRetry(wikiIngestMaxRetry),
asynq.Timeout(60*time.Minute),
asynq.ProcessIn(wikiIngestDelay),
)
if _, err := task.Enqueue(t); err != nil {
logger.Warnf(ctx, "wiki ingest: failed to enqueue trigger task: %v", err)
return fmt.Errorf("enqueue wiki ingest trigger: %w", err)
}
return nil
}
// EnqueueWikiRetract queues a wiki retraction op (a delete cleanup).
// Identical persistence model as EnqueueWikiIngest — the op rides in
// task_pending_ops and an asynq trigger fires shortly after to
// process the batch. Retracts use a slightly shorter ProcessIn delay
// because there is no "user upload arriving in waves" pattern to
// debounce against — a deletion fires once and we want the cleanup
// to land promptly.
func EnqueueWikiRetract(
ctx context.Context,
task interfaces.TaskEnqueuer,
pendingRepo interfaces.TaskPendingOpsRepository,
payload WikiRetractPayload,
) {
op := WikiPendingOp{
Op: WikiOpRetract,
KnowledgeID: payload.KnowledgeID,
DocTitle: payload.DocTitle,
DocSummary: payload.DocSummary,
PageSlugs: payload.PageSlugs,
FolderIDs: payload.FolderIDs,
Language: payload.Language,
}
payloadBytes, err := json.Marshal(op)
if err != nil {
logger.Warnf(ctx, "wiki retract: failed to marshal pending op: %v", err)
return
}
accepted, err := enqueueWikiPendingOp(ctx, pendingRepo, &types.TaskPendingOp{
TenantID: payload.TenantID,
TaskType: wikiTaskType,
Scope: wikiTaskScope,
ScopeID: payload.KnowledgeBaseID,
Op: WikiOpRetract,
DedupKey: payload.KnowledgeID,
Payload: payloadBytes,
})
if err != nil {
logger.Warnf(ctx, "wiki retract: failed to enqueue pending op: %v", err)
return
}
if !accepted {
logger.Infof(ctx, "wiki retract: skip enqueue for deleted KB %s", payload.KnowledgeBaseID)
return
}
trigger := WikiIngestPayload{
TenantID: payload.TenantID,
KnowledgeBaseID: payload.KnowledgeBaseID,
Language: payload.Language,
}
langfuse.InjectTracing(ctx, &trigger)
triggerBytes, _ := json.Marshal(trigger)
t := asynq.NewTask(types.TypeWikiIngest, triggerBytes,
asynq.Queue(types.QueueWiki),
asynq.MaxRetry(wikiIngestMaxRetry),
asynq.Timeout(60*time.Minute),
asynq.ProcessIn(5*time.Second), // Retract can trigger the batch quickly
)
if _, err := task.Enqueue(t); err != nil {
logger.Warnf(ctx, "wiki retract: failed to enqueue trigger task: %v", err)
}
}
// Handle implements interfaces.TaskHandler for asynq task processing. The
// wiki service owns two task types, dispatched here by type so a single
// registered handler covers both the ingest batch and the debounced
// KB-global finalize pass.
func (s *wikiIngestService) Handle(ctx context.Context, t *asynq.Task) error {
switch t.Type() {
case types.TypeWikiFinalize:
return s.ProcessWikiFinalize(ctx, t)
default:
return s.ProcessWikiIngest(ctx, t)
}
}
func wikiIngestCleanupContext(ctx context.Context) (context.Context, context.CancelFunc) {
if ctx == nil {
ctx = context.Background()
}
return context.WithTimeout(context.WithoutCancel(ctx), wikiIngestCleanupTimeout)
}
func (s *wikiIngestService) clearDeletedKnowledgeBasePendingOps(ctx context.Context, kbID string) error {
cleaner, ok := s.pendingRepo.(interfaces.TaskPendingOpsScopeCleaner)
if !ok || kbID == "" {
return nil
}
cleanupCtx, cancel := wikiIngestCleanupContext(ctx)
defer cancel()
return cleaner.DeleteByScope(cleanupCtx, types.TaskScopeKnowledgeBase, kbID)
}
func (s *wikiIngestService) enqueueFinalizeRow(ctx context.Context, op *types.TaskPendingOp) bool {
accepted, err := enqueueWikiPendingOp(ctx, s.pendingRepo, op)
if err != nil {
logger.Warnf(ctx, "wiki finalize: enqueue %s row failed: %v", op.Op, err)
return false
}
return accepted
}
// enqueueFinalize persists this batch's KB-global convergence work into the
// finalize lane of task_pending_ops and schedules a debounced trigger. One
// "slug" row per affected page (carrying its fresh title when this batch wrote
// it, for cross-linking) plus one "change" row per doc added/removed (for the
// index-intro change description).
func (s *wikiIngestService) enqueueFinalize(
ctx context.Context,
payload WikiIngestPayload,
affectedSlugs []string,
freshTitleBySlug map[string]string,
changes []wikiFinalizeChange,
folderIDs []string,
) error {
if s.pendingRepo == nil {
return nil
}
acceptedAny := false
recovery := wikiRecoveryIdentityFromContext(ctx)
if recovery.Present() {
return s.enqueueControlledWikiFinalize(ctx, payload, affectedSlugs, freshTitleBySlug, changes)
}
dedupPrefix := ""
if recovery.Present() {
dedupPrefix = "recovery:" + recovery.Key() + ":"
}
for _, slug := range affectedSlugs {
row := wikiFinalizeRow{WikiRecoveryIdentity: recovery, Slug: slug, Title: freshTitleBySlug[slug]}
b, err := json.Marshal(row)
if err != nil {
continue
}
if s.enqueueFinalizeRow(ctx, &types.TaskPendingOp{
TenantID: payload.TenantID,
TaskType: wikiFinalizeTaskType,
Scope: wikiTaskScope,
ScopeID: payload.KnowledgeBaseID,
Op: wikiFinalizeOpSlug,
DedupKey: dedupPrefix + slug,
Payload: b,
}) {
acceptedAny = true
}
}
for i := range changes {
row := wikiFinalizeRow{WikiRecoveryIdentity: recovery, Change: &changes[i]}
b, err := json.Marshal(row)
if err != nil {
continue
}
if s.enqueueFinalizeRow(ctx, &types.TaskPendingOp{
TenantID: payload.TenantID,
TaskType: wikiFinalizeTaskType,
Scope: wikiTaskScope,
ScopeID: payload.KnowledgeBaseID,
Op: wikiFinalizeOpChange,
DedupKey: "",
Payload: b,
}) {
acceptedAny = true
}
}
if len(folderIDs) > 0 {
row := wikiFinalizeRow{WikiRecoveryIdentity: recovery, FolderIDs: uniqueWikiFolderIDs(folderIDs)}
if b, err := json.Marshal(row); err == nil {
if s.enqueueFinalizeRow(ctx, &types.TaskPendingOp{
TenantID: payload.TenantID,
TaskType: wikiFinalizeTaskType,
Scope: wikiTaskScope,
ScopeID: payload.KnowledgeBaseID,
Op: wikiFinalizeOpFolderPrune,
DedupKey: "",
Payload: b,
}) {
acceptedAny = true
}
}
}
if !acceptedAny {
return nil
}
s.scheduleFinalize(ctx, payload)
return nil
}
func uniqueWikiFolderIDs(values []string) []string {
seen := make(map[string]struct{}, len(values))
out := make([]string, 0, len(values))
for _, value := range values {
value = strings.TrimSpace(value)
if value == "" {
continue
}
if _, ok := seen[value]; ok {
continue
}
seen[value] = struct{}{}
out = append(out, value)
}
return out
}
// scheduleFinalize enqueues a debounced, coalesced KB-global finalize trigger.
// asynq.TaskID ("wiki-finalize-<kbID>") makes concurrent schedules within the
// debounce window collapse into one pending task; the conflict error is the
// expected coalescing signal, not a failure. In Lite mode the sync executor
// ignores TaskID, so finalize simply runs once per batch (acceptable at the
// small scale Lite mode targets).
func (s *wikiIngestService) scheduleFinalize(ctx context.Context, payload WikiIngestPayload) {
langfuse.InjectTracing(ctx, &payload)
b, _ := json.Marshal(payload)
t := asynq.NewTask(types.TypeWikiFinalize, b,
asynq.Queue(types.QueueWiki),
asynq.MaxRetry(wikiIngestMaxRetry),
asynq.Timeout(30*time.Minute),
asynq.ProcessIn(wikiFinalizeDelay),
asynq.TaskID("wiki-finalize-"+payload.KnowledgeBaseID),
)
if _, err := s.task.Enqueue(t); err != nil {
if errors.Is(err, asynq.ErrTaskIDConflict) || errors.Is(err, asynq.ErrDuplicateTask) {
return // a finalize is already scheduled/running for this KB — coalesced
}
logger.Warnf(ctx, "wiki finalize: schedule trigger failed: %v", err)
}
}
// scheduleFinalizeRetry is used when folder pruning is waiting for ingest
// rows to drain. It deliberately has no stable TaskID: the currently-running
// finalize task still owns that ID until it returns, so reusing it here would
// coalesce the only retry away. Duplicate retries are harmless because the
// durable prune row is deleted exactly once and an empty lane is a no-op.
func (s *wikiIngestService) scheduleFinalizeRetry(ctx context.Context, payload WikiIngestPayload) {
langfuse.InjectTracing(ctx, &payload)
b, _ := json.Marshal(payload)
t := asynq.NewTask(types.TypeWikiFinalize, b,
asynq.Queue(types.QueueWiki),
asynq.MaxRetry(wikiIngestMaxRetry),
asynq.Timeout(30*time.Minute),
asynq.ProcessIn(wikiFolderPruneRetryDelay),
)
if _, err := s.task.Enqueue(t); err != nil {
logger.Warnf(ctx, "wiki finalize: schedule deferred folder prune failed: %v", err)
}
}
// peekPendingList loads up to `limit` ops from task_pending_ops for
// this KB, ordered FIFO. Rows are NOT removed; callers must
// DeleteByIDs once they have been consumed (or IncrFailCount + leave
// them in place for the next pass).
//
// peekedIDs returns the DB ids of every row included in the peek
// (NOT just the ones that survived dedup) so trimPendingList can
// delete them all in one statement at the end of the batch — this
// matches the legacy "LTrim peekedCount entries" semantics, where
// duplicates collapsed by the consumer were also drained from the
// list once their canonical sibling had been processed.
func (s *wikiIngestService) peekPendingList(ctx context.Context, kbID string, limit int) (ops []WikiPendingOp, peekedIDs []int64, err error) {
if s.pendingRepo == nil {
return nil, nil, nil
}
if limit <= 0 {
limit = wikiMaxDocsPerBatch
}
rows, err := s.pendingRepo.PeekBatch(ctx, wikiTaskType, wikiTaskScope, kbID, limit)
if err != nil {
// Surface the error so the caller can distinguish a transient DB
// failure from a genuinely empty queue: the former must trigger an
// asynq retry, whereas returning "no rows" here would ack the task
// as a false success and strand the pending list.
return nil, nil, err
}
ops, peekedIDs = s.decodePendingRows(ctx, rows)
return ops, peekedIDs, nil
}
// claimPendingList is the standard-mode (Redis) counterpart of
// peekPendingList: it atomically CLAIMS up to `limit` ops (marks
// claimed_at) so concurrent batches for the same KB — allowed since Phase 3
// removed the exclusive per-KB lock — pull DISJOINT documents instead of
// double-processing. Stale claims (older than wikiClaimStaleAfter, i.e. from
// a crashed worker) are recovered. Dedup / peekedIDs semantics match
// peekPendingList; the returned peekedIDs are the claimed rows that the
// caller must DeleteByIDs on success or ReleaseByIDs to retry.
func (s *wikiIngestService) claimPendingList(ctx context.Context, kbID string, limit int) (ops []WikiPendingOp, peekedIDs []int64, err error) {
if s.pendingRepo == nil {
return nil, nil, nil
}
if limit <= 0 {
limit = wikiMaxDocsPerBatch
}
rows, err := s.pendingRepo.ClaimBatch(ctx, wikiTaskType, wikiTaskScope, kbID, limit,
time.Now().Add(-wikiClaimStaleAfter))
if err != nil {
// A claim failure is transient (DB blip). Propagate it so the batch
// returns an error and asynq retries, instead of acking the trigger
// as a false "no pending ops" success and stranding the queue.
return nil, nil, err
}
ops, peekedIDs = s.decodePendingRows(ctx, rows)
return ops, peekedIDs, nil
}
// wikiInflightReserveScript atomically enforces the per-KB in-flight cap.
// It purges expired slots (crash recovery), counts live slots, and adds the
// caller's slot iff under the cap. Returning 1 = reserved, 0 = at cap.
// Keeping purge+count+add in one Lua call closes the check-then-act race two
// concurrent reservers would otherwise have.
const wikiInflightReserveScript = `
local now = tonumber(ARGV[1])
local expiry = tonumber(ARGV[2])
local maxInflight = tonumber(ARGV[3])
local token = ARGV[4]
local ttl = tonumber(ARGV[5])
redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, now)
if redis.call('ZCARD', KEYS[1]) >= maxInflight then
return 0
end
redis.call('ZADD', KEYS[1], expiry, token)
redis.call('PEXPIRE', KEYS[1], ttl)
return 1
`
// reserveInflightSlot claims one of KB's concurrent-batch slots (Phase 4,
// standard/Redis mode). Returns (release, true) when granted — release() MUST
// run when the batch finishes; (nil, false) when the KB is already at
// maxInflight, so the caller should reschedule and bail. A background renew
// keeps the slot alive for the batch's duration; a crashed batch's slot simply
// expires (wikiInflightTTL) and is purged by the next reserver. Lite mode has
// no shared-pool contention (liteLocks already serialize per KB), so it always
// grants a no-op slot. Fails OPEN on a Redis error: a blip must not halt wiki
// generation, and the pool size still bounds total work.
func (s *wikiIngestService) reserveInflightSlot(ctx context.Context, kbID string, maxInflight int) (func(), bool) {
if s.redisClient == nil || maxInflight <= 0 {
return func() {}, true
}
key := wikiInflightPrefix + kbID
token := uuid.New().String()
now := time.Now()
res, err := s.redisClient.Eval(ctx, wikiInflightReserveScript,
[]string{key},
now.UnixMilli(),
now.Add(wikiInflightTTL).UnixMilli(),
maxInflight,
token,
wikiInflightTTL.Milliseconds(),
).Int()
if err != nil {
logger.Warnf(ctx, "wiki ingest: inflight reserve failed for KB %s: %v (running uncapped)", kbID, err)
return func() {}, true
}
if res == 0 {
return nil, false
}
renewCtx, cancel := context.WithCancel(context.Background())
go func() {
ticker := time.NewTicker(wikiInflightRenew)
defer ticker.Stop()
for {
select {
case <-renewCtx.Done():
return
case <-ticker.C:
exp := float64(time.Now().Add(wikiInflightTTL).UnixMilli())
s.redisClient.ZAdd(context.Background(), key, redis.Z{Score: exp, Member: token})
s.redisClient.PExpire(context.Background(), key, wikiInflightTTL)
}
}
}()
return func() {
cancel()
s.redisClient.ZRem(context.Background(), key, token)
}, true
}
// scheduleCappedRetry enqueues a single coalesced follow-up trigger after a
// batch was turned away by the in-flight cap. asynq.TaskID collapses all
// turned-away triggers for one KB into a single pending retry (no thundering
// herd), and the running batches that hold the slots also chain their own
// follow-ups on completion, so the turned-away rows are guaranteed to drain
// once a slot frees.
func (s *wikiIngestService) scheduleCappedRetry(ctx context.Context, payload WikiIngestPayload) {
langfuse.InjectTracing(ctx, &payload)
b, _ := json.Marshal(payload)
t := asynq.NewTask(types.TypeWikiIngest, b,
asynq.Queue(types.QueueWiki),
asynq.MaxRetry(wikiIngestMaxRetry),
asynq.Timeout(60*time.Minute),
asynq.ProcessIn(wikiInflightBackoff),
asynq.TaskID("wiki-ingest-capped-"+payload.KnowledgeBaseID),
)
if _, err := s.task.Enqueue(t); err != nil {
if errors.Is(err, asynq.ErrTaskIDConflict) || errors.Is(err, asynq.ErrDuplicateTask) {
return // a capped retry is already pending for this KB — coalesced
}
logger.Warnf(ctx, "wiki ingest: capped-retry enqueue failed: %v", err)
}
}
// scheduleStaleClaimRecheck arms a single, far-future safety-net trigger for a
// KB that still has pending rows but yielded nothing to claim (every eligible
// row is held by a FRESH claim). Normally a running batch drains those rows and
// chains its own fast follow-up on completion; this net exists only for the
// case where the claim holder CRASHED, leaving claimed_at stamped so no worker
// can re-claim until wikiClaimStaleAfter elapses — and where nothing else would
// ever re-trigger the KB afterwards.
//
// The delay is set past the stale threshold so that when the net fires the
// abandoned claims are guaranteed eligible again. asynq.TaskID coalesces all
// rechecks for one KB into a single pending net (no thundering herd from
// concurrent no-op batches). If PendingCount is already zero the KB has fully
// drained and no net is needed. Returns true if a net is (or already was)
// scheduled.
func (s *wikiIngestService) scheduleStaleClaimRecheck(ctx context.Context, payload WikiIngestPayload) bool {
if s.pendingRepo == nil {
return false
}
count, err := s.pendingRepo.PendingCount(ctx, wikiTaskType, wikiTaskScope, payload.KnowledgeBaseID)
if err != nil || count == 0 {
return false
}
logger.Infof(ctx, "wiki ingest: %d rows for KB %s held by fresh claims, arming stale-claim recheck", count, payload.KnowledgeBaseID)
langfuse.InjectTracing(ctx, &payload)
b, _ := json.Marshal(payload)
t := asynq.NewTask(types.TypeWikiIngest, b,
asynq.Queue(types.QueueWiki),
asynq.MaxRetry(wikiIngestMaxRetry),
asynq.Timeout(60*time.Minute),
asynq.ProcessIn(wikiClaimStaleAfter+wikiFollowUpDelay),
asynq.TaskID("wiki-ingest-recheck-"+payload.KnowledgeBaseID),
)
if _, err := s.task.Enqueue(t); err != nil {
if errors.Is(err, asynq.ErrTaskIDConflict) || errors.Is(err, asynq.ErrDuplicateTask) {
return true // a recheck is already armed for this KB — coalesced
}
logger.Warnf(ctx, "wiki ingest: stale-claim recheck enqueue failed: %v", err)
return false
}
return true
}
// decodePendingRows converts raw task_pending_ops rows into WikiPendingOps,
// applying the last-write-wins dedup by knowledge_id. peekedIDs carries the
// db ids of EVERY row (including dedup-collapsed ones) so the caller can
// drain them all at trim time. Shared by peekPendingList (no claim) and
// claimPendingList (claimed rows).
func (s *wikiIngestService) decodePendingRows(ctx context.Context, rows []*types.TaskPendingOp) (ops []WikiPendingOp, peekedIDs []int64) {
if len(rows) == 0 {
return nil, nil
}
all := make([]WikiPendingOp, 0, len(rows))
peekedIDs = make([]int64, 0, len(rows))
for _, r := range rows {
peekedIDs = append(peekedIDs, r.ID)
var op WikiPendingOp
if len(r.Payload) > 0 {
if err := json.Unmarshal(r.Payload, &op); err != nil {
logger.Warnf(ctx, "wiki ingest: failed to unmarshal pending op id=%d: %v", r.ID, err)
continue
}
} else {
// Defensive: if payload was lost, fall back to column data
// so the row is still drainable (otherwise it would loop
// on every batch as un-deletable).
op = WikiPendingOp{
Op: r.Op,
KnowledgeID: r.DedupKey,
}
}
op.dbID = r.ID
all = append(all, op)
}
// Deduplicate by KnowledgeID, keeping only the *last* operation for
// each document. Optimizes out redundant sequences (e.g., upload
// then immediate delete: [ingest, retract] → [retract]). The
// non-canonical rows still get drained at trim time — their dbIDs
// are in peekedIDs.
seen := make(map[string]bool)
recoveryDocs := make(map[string]bool)
docCounts := make(map[string]int)
for _, op := range all {
docCounts[op.KnowledgeID]++
if op.WikiRecoveryIdentity.Present() {
recoveryDocs[op.KnowledgeID] = true
}
}
reversedUnique := make([]WikiPendingOp, 0, len(all))
for i := len(all) - 1; i >= 0; i-- {
op := all[i]
if recoveryDocs[op.KnowledgeID] {
op.recoveryConflict = docCounts[op.KnowledgeID] > 1
reversedUnique = append(reversedUnique, op)
continue
}
if op.KnowledgeID == "" {
// No dedup key — keep verbatim (rare; edge case for
// future ops without a knowledge anchor).
reversedUnique = append(reversedUnique, op)
continue
}
if seen[op.KnowledgeID] {
continue
}
seen[op.KnowledgeID] = true
reversedUnique = append(reversedUnique, op)
}
ops = make([]WikiPendingOp, 0, len(reversedUnique))
for i := len(reversedUnique) - 1; i >= 0; i-- {
ops = append(ops, reversedUnique[i])
}
return ops, peekedIDs
}
// trimPendingList deletes consumed rows from task_pending_ops. Empty
// input is a no-op so callers can invoke unconditionally at the end
// of a batch.
func (s *wikiIngestService) trimPendingList(ctx context.Context, ids []int64) error {
if s.pendingRepo == nil || len(ids) == 0 {
return nil
}
if err := s.pendingRepo.DeleteByIDs(ctx, ids); err != nil {
logger.Warnf(ctx, "wiki ingest: failed to trim %d pending rows: %v", len(ids), err)
return err
}
return nil
}
// finalizeWikiSubtask releases this knowledge's slot in the finalizing
// counter once its wiki op reaches a terminal state (mapped successfully
// or dead-lettered). The matching +1 is seeded by
// KnowledgePostProcess.SetFinalizing when willSpawnWiki is true. Callers
// must only invoke this for ingest ops — retract ops are for deleted
// knowledge that has no counter to drain.
//
// Safe to call on a row that is already completed or whose counter is
// already zero: FinalizeSubtask guards both the decrement (count > 0) and
// the promote (parse_status = finalizing AND count = 0), so an op enqueued
// before this accounting shipped is a harmless no-op.
func (s *wikiIngestService) finalizeWikiSubtask(ctx context.Context, knowledgeID string) {
// Wiki is only finalized when its op reaches a terminal state, so this is
// always an intended drain (retErr=nil, final=true). Detached context: the
// wiki batch worker may be mid-shutdown or have a cancelled ctx when this
// runs; a swallowed failure would strand the parent in "finalizing".
finalizeSubtaskDetached(ctx, s.knowledgeRepo, knowledgeID, "wiki", nil, false, true)
}
// requeueFailedOps records in-batch failures.
//
// For each failed op:
//
// - IncrFailCount on the source row. The repo returns the new total,
// so a single round trip handles both bookkeeping and retry-budget
// check.
// - If the count is <= wikiMaxFailRetries: leave the row in place.
// The next follow-up batch's PeekBatch will pick it up naturally
// (rows are ordered by id ASC and we never moved/touched it).
// - If the count exceeds the retry cap: archive the op into
// task_dead_letters and DeleteByIDs to remove it from the queue.
// Settlement failures are returned so the caller does not mark claims
// settled while rows are still claimed or undeleted.
func (s *wikiIngestService) requeueFailedOps(ctx context.Context, payload WikiIngestPayload, ops []WikiPendingOp) error {
if s.pendingRepo == nil || len(ops) == 0 {
return nil
}
var settleErrs []error
for _, op := range ops {
if op.dbID == 0 {
// Op was never persisted (synthetic / test) — nothing to
// retry against.
continue
}
count, err := s.pendingRepo.IncrFailCount(ctx, op.dbID)
if err != nil {
logger.Warnf(ctx, "wiki ingest: failed to increment fail count for %s (id=%d): %v", op.KnowledgeID, op.dbID, err)
settleErrs = append(settleErrs, fmt.Errorf("increment fail count id=%d: %w", op.dbID, err))
// Without a fresh count we can't tell whether to drop. Be
// conservative: leave the row in place; the next PeekBatch
// will see it again and we'll try once more.
continue
}
if count <= wikiMaxFailRetries {
// Release the claim so the row is immediately eligible for the
// next trigger's ClaimBatch instead of waiting out
// wikiClaimStaleAfter. No-op in Lite mode (row was peeked, never
// claimed). ReleaseByIDs preserves fail_count, so the retry
// budget still counts down.
if err := s.pendingRepo.ReleaseByIDs(ctx, []int64{op.dbID}); err != nil {
logger.Warnf(ctx, "wiki ingest: failed to release claim for retry id=%d: %v", op.dbID, err)
settleErrs = append(settleErrs, fmt.Errorf("release retry claim id=%d: %w", op.dbID, err))
}
logger.Infof(ctx, "wiki ingest: re-queued failed op %s (%s) for retry (attempt %d/%d)", op.KnowledgeID, op.DocTitle, count, wikiMaxFailRetries)
continue
}
// Preserve the original operation until its recovery record is durable.
if s.deadLetterRepo == nil {
settleErrs = append(settleErrs, fmt.Errorf("archive dead letter id=%d: repository unavailable", op.dbID))
continue
}
logger.Warnf(ctx, "wiki ingest: dropping op %s (%s) after %d failures (limit %d)", op.KnowledgeID, op.DocTitle, count, wikiMaxFailRetries)
if s.deadLetterRepo != nil {
payloadBytes, _ := json.Marshal(op)
if dlErr := s.deadLetterRepo.Insert(ctx, &types.TaskDeadLetter{
TenantID: payload.TenantID,
TaskType: wikiTaskType,
Scope: wikiTaskScope,
ScopeID: payload.KnowledgeBaseID,
RelatedID: op.KnowledgeID,
Payload: payloadBytes,
LastError: fmt.Sprintf("exceeded wikiMaxFailRetries=%d (in-batch retries)", wikiMaxFailRetries),
FailCount: count,
}); dlErr != nil {
logger.Warnf(ctx, "wiki ingest: failed to archive op %s to dead letters: %v", op.KnowledgeID, dlErr)
settleErrs = append(settleErrs, fmt.Errorf("archive dead letter id=%d: %w", op.dbID, dlErr))
continue
}
}
if err := s.pendingRepo.DeleteByIDs(ctx, []int64{op.dbID}); err != nil {
logger.Warnf(ctx, "wiki ingest: failed to drop dead-lettered row id=%d: %v", op.dbID, err)
settleErrs = append(settleErrs, fmt.Errorf("drop dead-lettered row id=%d: %w", op.dbID, err))
continue
}
if op.Op == WikiOpIngest {
s.finalizeWikiSubtask(ctx, op.KnowledgeID)
}
}
return errors.Join(settleErrs...)
}
// docIngestResult captures per-document info for batch post-processing.
type wikiIngestPageRef struct {
Slug string
Title string
}
type docIngestResult struct {
types.WikiRecoveryIdentity
KnowledgeID string
DocTitle string
Summary string // one-line summary of the document (from summary page)
// Pages records the wiki pages this document touched, carrying both
// the slug used for link/retract bookkeeping and its human-readable title.
Pages []wikiIngestPageRef
// MapStats are the per-doc map-phase metrics captured at the moment
// mapOneDocument finishes. Surfaced into the postprocess.wiki span's
// output so the trace viewer can show "what the map phase produced"
// even though the span itself stays open until the batch's reduce +
// cleanup phases complete (so the user-visible duration covers the
// whole pipeline for this doc, not just LLM extraction).
MapStats types.JSONMap
// WikiSpan is the postprocess.wiki subspan opened at the start of
// mapOneDocument. ProcessWikiIngest holds it open across the reduce
// + cleanup phases and closes it once this doc's pages have all
// been materialised — see the EndSpan call near the end of
// ProcessWikiIngest. nil when no parent attempt was found, in which
// case the tracker helpers are all no-ops anyway.
WikiSpan *Span
}
// WikiBatchContext holds shared data across Map and Reduce phases.
//
// Historically this carried a fully materialized `AllPages` slice plus
// pre-built SlugTitleMap / SummaryContentByKnowledgeID lookup tables.
// At 4w-document scale that meant the very first thing every batch
// did was load 100K+ wiki_pages rows (content TEXT included) into Go
// memory — and then walk them several more times for cleanDeadLinks /
// injectCrossLinks / getExistingPageSlugsForKnowledge.
//
// We now lazy-load via fetchers backed by lightweight projections
// (ListBySlugs / ListSummariesByKnowledgeIDs). Each fetcher caches
// results keyed by its input so repeat lookups within a batch are
// free; the cache is per-batch and goroutine-local-via-mutex (sync.Map
// would also work but mutex keeps the surface small).
type WikiBatchContext struct {
checkpointScopes sync.Map
checkpointSourceHashes sync.Map
checkpointPages atomic.Int32
checkpointEvidenceSteps atomic.Int32
// SlugTitle resolves a slug to its current title (or "" if missing).
// Backed by ListBySlugs; cache is populated as callers ask, so we
// only pay for the slugs we actually look at.
SlugTitle func(ctx context.Context, slug string) string
// SlugTitleMany batches a slug-set into a single ListBySlugs query
// and returns the resolved titles map. Convenient when a caller
// already has the full slug list; results are still cached.
SlugTitleMany func(ctx context.Context, slugs []string) map[string]string
// SummaryContentByKnowledgeID returns the surviving summary page's
// content for the given knowledge id (or "" if no summary page
// exists / was archived). Backed by ListSummariesByKnowledgeIDs;
// cache is populated lazily as well.
SummaryContentByKnowledgeID func(ctx context.Context, kid string) string
// ExtractionGranularity drives Pass 0 (candidate slug extraction)
// aggressiveness. Resolved once per batch from the KnowledgeBase's
// WikiConfig so every doc in the batch sees the same scope rules.
// Already Normalize()'d — consumers can assume it is one of the
// three valid values.
ExtractionGranularity types.WikiExtractionGranularity
// ContentInstructions and ExtractionInstructions are KB-scoped business
// guidance. Stable citation, merge, taxonomy and JSON rules remain in the
// system templates and cannot be replaced by these fields.
ContentInstructions string
ExtractionInstructions string
// PlannedFolderID holds the per-slug wiki_folders.id assigned by the batch
// taxonomy planning pass (planBatchTaxonomy + folder resolution), keyed by
// page slug. Reduce applies it only to pages that aren't already filed
// (FolderID == ""), so the whole batch lands on one coherent tree without
// churning user-curated placements. The folders themselves are created
// sequentially before reduce, so the parallel reduce phase only assigns
// pre-resolved ids and never races on folder creation. Read-only during
// reduce.
PlannedFolderID map[string]string
// identityClaims is the Lite-mode and Redis-error fallback for identity
// reservations. Map workers in one batch run concurrently even though Lite
// serializes batches, so they still need to converge before Reduce groups
// updates by slug. The map is batch-scoped and disappears with the batch.
identityClaims sync.Map
// identityPages memoizes exact title lookups for this batch so concurrent
// map workers probing the same (page type, normalized title) share one DB
// round-trip. Values are []*types.WikiPageLite, including empty slices
// for confirmed misses.
identityPages sync.Map
}
// SlugUpdate represents a single update operation for a specific slug
type SlugUpdate struct {
types.WikiRecoveryIdentity
Slug string
Type string // "entity", "concept", "summary", "retract", "retractStale"
Item extractedItem // For entity/concept
DocTitle string
KnowledgeID string
SourceRef string
// Language is the already-resolved, human-readable language name the
// Reduce phase interpolates into the editor prompt (e.g. "Chinese
// (Simplified)"), NOT a locale code. Map resolves it once per document
// so every page derived from that document shares one language.
Language string
SummaryBody string // For summary
SummaryLine string // For summary
RetractDocContent string // For retract / retractStale
// SourceChunks lists the chunk IDs (within KnowledgeID) that substantively
// support this update. Mirrors Item.SourceChunks for convenience — the
// Reduce phase reads from here to avoid an extra field hop.
SourceChunks []string
// DocSummary is the document-level summary body produced by
// WikiSummaryPrompt (everything after the SUMMARY: ... headline, falling
// back to the raw output if no headline could be parsed out). Carried
// here so the Reduce phase can frame cited chunks with a rich
// <source_context> block that tells the editor model what the document
// is about AND what kind of document it is (resume vs announcement vs
// product page). The one-line headline alone was too terse to keep the
// editor grounded on longer / multi-topic source documents.
DocSummary string
}
func previewText(s string, maxRunes int) string {
s = strings.TrimSpace(s)
s = strings.ReplaceAll(s, "\n", " ")
s = strings.ReplaceAll(s, "\t", " ")
for strings.Contains(s, " ") {
s = strings.ReplaceAll(s, " ", " ")
}
r := []rune(s)
if maxRunes <= 0 || len(r) <= maxRunes {
return s
}
return string(r[:maxRunes]) + "...(truncated)"
}
func previewStringSlice(items []string, limit int) string {
if len(items) == 0 {
return "[]"
}
if limit <= 0 {
limit = 1
}
n := len(items)
if n > limit {
items = items[:limit]
}
out := make([]string, 0, len(items))
for _, it := range items {
out = append(out, previewText(it, 48))
}
if n > limit {
return fmt.Sprintf("[%s ...(+%d)]", strings.Join(out, ", "), n-limit)
}
return fmt.Sprintf("[%s]", strings.Join(out, ", "))
}
// previewExtractedItems returns a JSON-friendly preview of the first
// `limit` extracted entities or concepts so the trace viewer's
// postprocess.wiki.extract span shows actual names/slugs/descriptions
// instead of bare counts. Each item is trimmed to a small fixed
// budget — these end up serialised into the spans table's JSONB
// output column, so the cumulative size matters more than per-item
// fidelity.
func previewExtractedItems(items []extractedItem, limit int) []map[string]string {
if limit <= 0 {
limit = 1
}
n := len(items)
if n > limit {
items = items[:limit]
}
out := make([]map[string]string, 0, len(items))
for _, it := range items {
out = append(out, map[string]string{
"name": previewText(it.Name, 60),
"slug": it.Slug,
"description": previewText(it.Description, 120),
})
}
return out
}
// topCitedSlugs returns the top `limit` slugs by chunk-citation count.
// Used by postprocess.wiki.classify so the trace surfaces which
// candidate slugs the citation pass attached the most chunks to —
// useful when triaging "this LLM run extracted weird things" without
// having to open and diff full chunk lists.
func topCitedSlugs(citations map[string][]string, limit int) []map[string]any {
if len(citations) == 0 {
return nil
}
type entry struct {
slug string
count int
}
entries := make([]entry, 0, len(citations))
for slug, ids := range citations {
entries = append(entries, entry{slug: slug, count: len(ids)})
}
sort.Slice(entries, func(i, j int) bool {
if entries[i].count != entries[j].count {
return entries[i].count > entries[j].count
}
return entries[i].slug < entries[j].slug
})
if limit > 0 && len(entries) > limit {
entries = entries[:limit]
}
out := make([]map[string]any, 0, len(entries))
for _, e := range entries {
out = append(out, map[string]any{
"slug": e.slug,
"chunks": e.count,
})
}
return out
}
// previewNewSlugs returns a JSON-friendly preview of the first
// `limit` slugs that the citation pass discovered (i.e. did not appear
// in pass-0's candidate list). Surfacing these makes "the citation
// LLM kept inventing entries" trivially diagnosable from the trace
// viewer.
func previewNewSlugs(items []newSlugFromCitation, limit int) []map[string]string {
if limit <= 0 {
limit = 1
}
n := len(items)
if n > limit {
items = items[:limit]
}
out := make([]map[string]string, 0, len(items))
for _, it := range items {
out = append(out, map[string]string{
"name": previewText(it.Name, 60),
"slug": it.Slug,
"type": it.Type,
"chunks": fmt.Sprintf("%d", len(it.SourceChunks)),
})
}
return out
}
// wikiLinkRE matches `[[slug]]` and `[[slug|display text]]` references
// inside wiki page content. The slug capture group rejects whitespace and
// the closing-bracket / pipe characters so we don't accidentally swallow
// adjacent text. Display text (group 2) is optional.
var wikiLinkRE = regexp.MustCompile(`\[\[([^\[\]\|\s]+)(?:\|([^\]]+))?\]\]`)
// sanitizeDeadSummaryLinks rewrites the summary pages produced by THIS
// batch to fix `[[slug]]` / `[[slug|display]]` references that point
// at slugs whose entity/concept page generation failed in reduce.
//
// Background: WikiSummaryPrompt instructs the LLM to embed wiki links
// for every extracted slug it knows about, but slug extraction happens
// during map (parallel with summary generation) and the actual page
// creation happens later in reduce. When reduce's WikiPageModifyUserPrompt
// fails on an entity/concept slug the page never gets written — and
// the already-persisted summary is left holding a `[[entity/foo|name]]`
// link that 404s.
//
// We pass the batch's affected-slug set + the SlugTitleMany fetcher
// to the resolver so that LLM-mangled slugs (e.g. extra pinyin hyphens
// in "shang-hai-tower" vs "shanghai-tower") are healed in place rather
// than stripped to plain text — preserving cross-link information
// whenever the display text or surface form unambiguously identifies a
// live page.
//
// Pure text replacement, no LLM call. Scoped to the doc-summary slugs
// in this batch (`summary/<slugify(knowledgeID)>`), keeping the work
// proportional to batch size.
func (s *wikiIngestService) sanitizeDeadSummaryLinks(
ctx context.Context,
kbID string,
docResults []*docIngestResult,
failedSlugs map[string]struct{},
batchCtx *WikiBatchContext,
) {
if len(failedSlugs) == 0 || len(docResults) == 0 {
return
}
// Build a (live-slug-set, title->slug) pair the resolver can consult.
// We seed liveSlugs from batchCtx (the slugs that DID make it into
// pages this batch) and expand it lazily as needed via SlugTitleMany.
// titleToSlug is filled with the same successful pages' titles so the
// display-text reverse lookup works on first try.
for _, r := range docResults {
if r == nil || r.KnowledgeID == "" {
continue
}
summarySlug := "summary/" + slugify(r.KnowledgeID)
page, err := s.wikiService.GetPageBySlug(ctx, kbID, summarySlug)
if err != nil || page == nil {
continue
}
// Collect the slugs this summary actually links to (so the
// resolver has a non-empty pool of candidates), plus all the
// successfully-written sibling pages from the same doc. These
// two sets together cover the LLM-vs-actual mismatch cases
// without paying for a full ListAll scan.
candidateSlugs := make(map[string]struct{}, len(page.OutLinks)+len(r.Pages))
for _, slug := range page.OutLinks {
candidateSlugs[slug] = struct{}{}
}
for _, ref := range r.Pages {
if _, bad := failedSlugs[ref.Slug]; bad {
continue
}
candidateSlugs[ref.Slug] = struct{}{}
}
liveSlugs, titleToSlug := s.resolveLiveSlugs(ctx, batchCtx, candidateSlugs)
newContent, changed := stripDeadWikiLinks(page.Content, failedSlugs, liveSlugs, titleToSlug)
if !changed {
continue
}
page.Content = newContent
if err := s.wikiService.UpdateAutoLinkedContent(ctx, page); err != nil {
logger.Warnf(ctx, "wiki ingest: failed to sanitize dead links in summary %s: %v", summarySlug, err)
continue
}
logger.Infof(ctx, "wiki ingest: sanitized dead [[slug]] refs in summary %s", summarySlug)
}
}
// resolveLiveSlugs builds the (liveSlugs, titleToSlug) pair that
// stripDeadWikiLinks / cleanDeadLinks pass into resolveDeadSlug.
//
// We start from a caller-supplied candidate set (typically the page's
// own out-links + this batch's freshly-written slugs) and ask the
// batch's SlugTitleMany fetcher to resolve them in one batched query.
// The fetcher already filters out archived / system pages, so missing
// entries naturally translate to "not live" without an extra check.
//
// titleToSlug is keyed by the page's exact title only — we don't have
// aliases in the lite projection. That's an acceptable trade-off: the
// reported breakage pattern is "slug munged, display = title", not
// "slug munged, display = alias", so display-by-title carries the
// majority of the rescue value at a fraction of the storage cost.
func (s *wikiIngestService) resolveLiveSlugs(
ctx context.Context,
batchCtx *WikiBatchContext,
candidates map[string]struct{},
) (map[string]struct{}, map[string]string) {
if len(candidates) == 0 || batchCtx == nil || batchCtx.SlugTitleMany == nil {
return nil, nil
}
slugList := make([]string, 0, len(candidates))
for s := range candidates {
slugList = append(slugList, s)
}
titles := batchCtx.SlugTitleMany(ctx, slugList)
live := make(map[string]struct{}, len(titles))
titleToSlug := make(map[string]string, len(titles))
for slug, title := range titles {
live[slug] = struct{}{}
if title != "" {
titleToSlug[title] = slug
}
}
return live, titleToSlug
}
// stripDeadWikiLinks rewrites `[[slug]]` / `[[slug|display]]` references
// whose `slug` falls into the dead set. The handling depends on whether
// the dead slug can be repaired:
//
// - If the resolver maps the dead slug to a live one (typically via
// display-text reverse lookup or hyphen-normalized equality —
// see resolveDeadSlug), the link is REWRITTEN with the corrected
// slug. Display text is preserved.
// - If no live candidate is close enough, the link is STRIPPED to
// plain text (display text when present; otherwise a humanized
// last-segment of the slug). This is the original behaviour.
//
// The resolver is optional: when liveSlugs / titleToSlug are nil or
// empty, every dead slug falls through to the strip path. This keeps
// backward compatibility for tests / call sites that don't yet wire
// the resolution data.
func stripDeadWikiLinks(
content string,
deadSlugs map[string]struct{},
liveSlugs map[string]struct{},
titleToSlug map[string]string,
) (string, bool) {
if len(deadSlugs) == 0 || content == "" {
return content, false
}
changed := false
out := wikiLinkRE.ReplaceAllStringFunc(content, func(match string) string {
sub := wikiLinkRE.FindStringSubmatch(match)
if len(sub) < 2 {
return match
}
slug := sub[1]
if _, dead := deadSlugs[slug]; !dead {
return match
}
display := ""
if len(sub) >= 3 {
display = strings.TrimSpace(sub[2])
}
// (1) Try fuzzy resolve before falling back to strip. The
// resolver consults display-text reverse lookup, hyphen-
// normalized equality, and bigram similarity in that order;
// returns "" only when no candidate is safe.
if resolved, ok := resolveDeadSlug(slug, display, liveSlugs, titleToSlug); ok && resolved != slug {
changed = true
if display != "" {
return "[[" + resolved + "|" + display + "]]"
}
return "[[" + resolved + "]]"
}
// (2) Strip — best-effort plain text. Prefer the LLM-supplied
// display text; otherwise humanize the slug's last path segment
// so the prose stays readable.
changed = true
if display != "" {
return display
}
parts := strings.Split(slug, "/")
label := parts[len(parts)-1]
label = strings.ReplaceAll(label, "-", " ")
return label
})
return out, changed
}
// cleanDeadLinks rewrites `[[slug]]` references in the batch's affected
// pages whose targets no longer exist (or were archived). Pure text
// cleanup — no LLM call.
//
// Scope is intentionally limited to the slugs touched by this batch:
// at 4w-document scale the legacy "scan every page in the KB" path was
// the dominant tail in the post-batch phase, and the long-tail
// historical dead links are better handled by the lint AutoFix pipeline
// (which runs out-of-band and can afford a full table walk).
//
// For each affected page:
//
// 1. Pull its lite projection (out_links + status) via the batch's
// SlugTitle fetcher (one IN query for the whole affected set,
// amortized via the batchCtx cache).
// 2. Probe the union of out-link targets through ExistsSlugs to
// classify them as live vs dead.
// 3. For each dead link, try resolveDeadSlug first; rewrite if a
// safe candidate exists, otherwise strip to plain text.
// 4. Persist the rewritten content via UpdateAutoLinkedContent so
// the version counter stays unchanged (this is a maintenance
// pass, not a user-visible edit).
func (s *wikiIngestService) cleanDeadLinks(ctx context.Context, kbID string, affectedSlugs []string, batchCtx *WikiBatchContext) error {
if len(affectedSlugs) == 0 {
return nil
}
// (1) Load the affected pages' content + out-links in one go.
// We need the full WikiPage rows here (not just lite projections)
// because we're going to rewrite content; the lite path saves
// nothing once we're touching content anyway.
cleaned := 0
for _, slug := range affectedSlugs {
page, err := s.wikiService.GetPageBySlug(ctx, kbID, slug)
if err != nil || page == nil {
if wikiRecoveryFromContext(ctx) != nil {
return fmt.Errorf("%w: cannot load page %s for link cleanup: %v", ErrWikiRecoveryBudgetStopped, slug, err)
}
continue
}
if page.Status == types.WikiPageStatusArchived {
continue
}
if page.PageType == types.WikiPageTypeIndex {
continue
}
if len(page.OutLinks) == 0 {
continue
}
// (2) Classify out-links as live vs dead via one batched
// ExistsSlugs query. Empty slug list → no-op.
liveMap, err := s.wikiService.ExistsSlugs(ctx, kbID, []string(page.OutLinks))
if err != nil {
logger.Warnf(ctx, "wiki: ExistsSlugs failed during dead-link cleanup for %s: %v", slug, err)
if wikiRecoveryFromContext(ctx) != nil {
return err
}
continue
}
deadSlugs := make(map[string]struct{})
liveSlugs := make(map[string]struct{}, len(liveMap))
for outSlug, alive := range liveMap {
if alive {
liveSlugs[outSlug] = struct{}{}
} else {
deadSlugs[outSlug] = struct{}{}
}
}
if len(deadSlugs) == 0 {
continue
}
// (3) Build the title->slug reverse-lookup map for fuzzy
// resolve. We pull titles for the live slugs only — those
// are the candidates a dead reference could be remapped to.
titles := batchCtx.SlugTitleMany(ctx, []string(page.OutLinks))
titleToSlug := make(map[string]string, len(titles))
for s, t := range titles {
if t != "" {
titleToSlug[t] = s
}
}
newContent, changed := stripDeadWikiLinks(page.Content, deadSlugs, liveSlugs, titleToSlug)
if !changed {
continue
}
// (4) Persist. UpdateAutoLinkedContent skips the version bump
// because dead-link cleanup is a machine-only edit.
page.Content = newContent
if err := s.wikiService.UpdateAutoLinkedContent(ctx, page); err != nil {
logger.Warnf(ctx, "wiki: failed to clean dead links in page %s: %v", page.Slug, err)
if wikiRecoveryFromContext(ctx) != nil {
return err
}
continue
}
cleaned++
}
if cleaned > 0 {
logger.Infof(ctx, "wiki: cleaned dead links in %d pages", cleaned)
}
return nil
}
// injectCrossLinks scans the batch's affected pages and injects
// `[[wiki-links]]` for mentions of other wiki page titles / aliases
// in the content. Pure text replacement, no LLM call.
//
// Scope is intentionally limited to two slug sets:
//
// 1. The affected pages themselves — we only rewrite their content.
// 2. The candidate refs come from (a) the affected pages' existing
// out-links (already known to be relevant via prior linkification
// or manual edits) plus (b) the batch's freshly-written sibling
// slugs supplied via `linkRefs` from the caller.
//
// At 4w-document scale this is the difference between loading 100K+
// pages just to find link candidates vs O(batch-size) lookups. We
// trade off some long-tail recall (a brand new entity in this batch
// won't be linkified into pages from previous batches until they get
// re-edited), but lint AutoFix is the right place for that.
//
// linkifyContent does the actual matching work, including code-block /
// existing-link / word-boundary exclusions.
func (s *wikiIngestService) injectCrossLinks(
ctx context.Context,
kbID string,
affectedSlugs []string,
freshRefs []linkRef,
batchCtx *WikiBatchContext,
) error {
if len(affectedSlugs) == 0 {
return nil
}
updated := 0
for _, slug := range affectedSlugs {
page, err := s.wikiService.GetPageBySlug(ctx, kbID, slug)
if err != nil || page == nil {
if wikiRecoveryFromContext(ctx) != nil {
return fmt.Errorf("%w: cannot load page %s for cross-links: %v", ErrWikiRecoveryBudgetStopped, slug, err)
}
continue
}
if page.PageType == types.WikiPageTypeIndex {
continue
}
// Build the per-page candidate ref set: the existing out-links
// (resolved via the batch's title fetcher to skip archived /
// system pages) plus the freshly-written sibling slugs from
// this batch.
var refs []linkRef
if len(page.OutLinks) > 0 {
titles := batchCtx.SlugTitleMany(ctx, []string(page.OutLinks))
for outSlug, title := range titles {
if title == "" || outSlug == slug {
continue
}
refs = append(refs, linkRef{slug: outSlug, matchText: title})
}
}
for _, fr := range freshRefs {
if fr.slug == slug {
continue
}
refs = append(refs, fr)
}
if len(refs) == 0 {
continue
}
newContent, changed := linkifyContent(page.Content, refs, page.Slug)
if !changed {
continue
}
page.Content = newContent
if err := s.wikiService.UpdateAutoLinkedContent(ctx, page); err != nil {
logger.Warnf(ctx, "wiki ingest: cross-link injection failed for %s: %v", page.Slug, err)
if wikiRecoveryFromContext(ctx) != nil {
return err
}
continue
}
updated++
}
if updated > 0 {
logger.Infof(ctx, "wiki ingest: injected cross-links in %d pages", updated)
}
return nil
}
// collectLinkRefs flattens (title + aliases) of all non-system pages into a
// single linkRef slice suitable for linkifyContent.
func collectLinkRefs(pages []*types.WikiPage) []linkRef {
refs := make([]linkRef, 0, len(pages)*2)
for _, p := range pages {
if p.PageType == types.WikiPageTypeIndex {
continue
}
if p.Title != "" {
refs = append(refs, linkRef{slug: p.Slug, matchText: p.Title})
}
for _, alias := range p.Aliases {
if alias != "" {
refs = append(refs, linkRef{slug: p.Slug, matchText: alias})
}
}
}
return refs
}
// wikiTaxonomyPromptMaxPaths caps how many existing folders are rendered into a
// planning prompt as the set to reuse. Reached only for pathologically large
// taxonomies; the similarity preprocessing keeps the fed set well under it.
const wikiTaxonomyPromptMaxPaths = 150
// wikiTaxonomyFolderPoolMax bounds the existing folders pulled from the DB as the
// candidate pool for similarity selection. Distinct folders are few even for
// large KBs, so this only guards against a degenerate taxonomy.
const wikiTaxonomyFolderPoolMax = 400
// wikiTaxonomyFeedAllMaxFolders is the folder count at or below which the whole
// folder set is fed to the planner as-is: a healthy navigation directory is
// small, so feeding everything gives perfect reuse recall with no embedding cost
// (similarity preprocessing only earns its keep once folders are numerous).
const wikiTaxonomyFeedAllMaxFolders = 60
// wikiTaxonomyRelevantTopK is how many nearest existing deeper folders each item
// contributes to the reuse set when similarity preprocessing kicks in.
const wikiTaxonomyRelevantTopK = 3
// wikiTaxonomyPlanChunkSize caps how many items go into a single planning call.
// Larger batches are split into chunks; folders assigned by earlier chunks are
// fed forward as "existing folders" so later chunks converge onto the same tree.
const wikiTaxonomyPlanChunkSize = 60
const wikiTaxonomyEmptyTreeHint = "(none yet — this knowledge base has no folders, design a fresh directory)"
type wikiTaxonomyNode struct {
children map[string]*wikiTaxonomyNode
}
func insertWikiTaxonomyPath(root *wikiTaxonomyNode, path []string) {
if root == nil || len(path) == 0 {
return
}
cur := root
for _, part := range path {
part = strings.TrimSpace(part)
if part == "" {
continue
}
if cur.children == nil {
cur.children = make(map[string]*wikiTaxonomyNode)
}
child := cur.children[part]
if child == nil {
child = &wikiTaxonomyNode{}
cur.children[part] = child
}
cur = child
}
}
func appendWikiTaxonomyNode(buf *strings.Builder, label string, node *wikiTaxonomyNode, depth int) {
if label != "" {
fmt.Fprintf(buf, "%s%s\n", strings.Repeat(" ", depth), label)
}
if node == nil || len(node.children) == 0 {
return
}
keys := make([]string, 0, len(node.children))
for k := range node.children {
keys = append(keys, k)
}
sort.Strings(keys)
for _, k := range keys {
appendWikiTaxonomyNode(buf, k, node.children[k], depth+1)
}
}
// formatExistingTaxonomyForPrompt renders distinct category_path values as an
// indented folder tree for LLM extraction prompts.
func formatExistingTaxonomyForPrompt(paths [][]string) string {
if len(paths) == 0 {
return ""
}
root := &wikiTaxonomyNode{}
for _, path := range paths {
insertWikiTaxonomyPath(root, path)
}
if len(root.children) == 0 {
return ""
}
var buf strings.Builder
keys := make([]string, 0, len(root.children))
for k := range root.children {
keys = append(keys, k)
}
sort.Strings(keys)
for _, k := range keys {
appendWikiTaxonomyNode(&buf, k, root.children[k], 0)
}
return strings.TrimSpace(buf.String())
}
// getExistingPageSlugsForKnowledge returns all page slugs that currently
// reference a given knowledge ID in their source_refs. Used to snapshot
// state before re-ingest so the reduce phase can reconcile additions vs
// retractions.
//
// Backed by idx_wiki_pages_source_refs (GIN jsonb_path_ops, migration
// 000041) and the legacy text-index fallback for "kid|title" entries.
// We project to slugs only — no need to load full row content for a
// per-doc snapshot.
//
// Index/log slugs (wiki-intrinsic system pages) never carry real
// source_refs in practice, but we filter them out explicitly here as
// a defense-in-depth measure: an old buggy ingest that mistakenly
// stamped a system page with a knowledge ref would otherwise show up
// in the reparse "old set" and confuse the reduce stage.
func (s *wikiIngestService) getExistingPageSlugsForKnowledge(ctx context.Context, kbID, knowledgeID string) map[string]bool {
slugs, err := s.wikiService.ListSlugsBySourceRef(ctx, kbID, knowledgeID)
if err != nil {
logger.Warnf(ctx, "wiki ingest: ListSlugsBySourceRef(%s) failed: %v", knowledgeID, err)
return nil
}
if len(slugs) == 0 {
return nil
}
out := make(map[string]bool, len(slugs))
for _, slug := range slugs {
// Defense-in-depth: skip wiki-intrinsic slugs that never have
// real source refs.
if slug == "index" {
continue
}
out[slug] = true
}
return out
}
// retractStalePages handles pages that were previously linked to this document
// but are no longer produced by the updated extraction.
// - Single-source stale pages → deleted
// - Multi-source stale pages → LLM retract to clean content synchronously
// Build set of newly affected slugs (including summary)
// Stale = was in old set but not in new set
// Remove this doc's source ref
// No other sources → delete the page
// Multi-source → remove ref, queue retract
// extractedItem represents a single extracted entity or concept.
//
// SourceChunks holds the stable chunk IDs (from the source document) that
// substantively discuss this item. The Reduce phase requires these scoped,
// available chunks as verbatim evidence. Description/Details are navigation
// metadata and cannot substitute for missing source evidence.
type extractedItem struct {
Name string `json:"name"`
Slug string `json:"slug"`
Aliases []string `json:"aliases"`
Description string `json:"description"`
Details string `json:"details"`
SourceChunks []string `json:"source_chunks,omitempty"`
}
// combinedExtraction represents the parsed result of the combined entity+concept extraction
type combinedExtraction struct {
Entities []extractedItem `json:"entities"`
Concepts []extractedItem `json:"concepts"`
// CandidatePruning reports what the bounded quality pass removed. It is
// counted, never serialized: only counts leave the extractor, never names or
// model content.
CandidatePruning wikiCandidatePruning `json:"-"`
}
// rebuildIndexPage refreshes the LLM-generated intro that sits on the
// index wiki_pages row.
//
// History: the index page used to store "intro + full directory listing" as
// a single multi-MB markdown blob in content. Every ingest batch rewrote
// the whole column, which on KBs with tens of thousands of pages caused
// O(N) TOAST writes per batch. The directory was lifted out into the
// structured GET /wiki/index endpoint (see wikiPageService.GetIndexView),
// and this method now only maintains the intro.
//
// Intro lifecycle:
// - First time (empty or legacy placeholder): generate from all document
// summaries via WikiIndexIntroPrompt.
// - Subsequent calls with a change description: incremental update via
// WikiIndexIntroUpdatePrompt so the intro reflects what just landed.
// - No change description: keep the existing intro untouched.
//
// The new intro is written to both Content and Summary so readers that
// still fall back to Summary (older clients, legacy migrations) stay in
// sync with the column the view actually renders.
// indexIntroSummaryCap caps how many summary pages we feed into the
// LLM when generating the wiki index intro from scratch. A 4w-document
// KB would otherwise blow the context window every batch, and the
// intro is a "set the scene" artifact where the most-recently-touched
// documents carry disproportionately more signal anyway. We pick the
// top-N most-recently-updated summaries and add a "showing N of M"
// hint to the prompt so the LLM can be honest about its sample.
const indexIntroSummaryCap = 200
// rebuildIndexPage refreshes the LLM-generated intro on the index
// page. Two paths:
//
// - First-time generation (no existing intro, or only the legacy
// placeholder): the LLM gets a CAPPED window of the most recent
// summary pages (most-recently-updated wins). Compare with the
// legacy path which loaded ALL summaries — at 4w-document scale
// that produced multi-MB prompts that simply broke the context
// window and silently fell back to a hardcoded intro.
// - Incremental update: the LLM gets only the existing intro plus
// the change description for THIS batch. Document summaries are
// intentionally NOT included — at scale the change-description
// alone is enough signal for "what landed?", and excluding the
// full summary set keeps the prompt size bounded regardless of
// KB size.
//
// The intro is written to both Content and Summary so legacy readers
// that fall through to Summary stay in sync.
func (s *wikiIngestService) rebuildIndexPage(ctx context.Context, chatModel chat.Chat, payload WikiIngestPayload,
changeDesc, lang, customInstructions string,
) error {
indexPage, indexErr := s.wikiService.GetIndex(ctx, payload.KnowledgeBaseID)
if indexErr != nil && wikiRecoveryFromContext(ctx) != nil {
return indexErr
}
if indexPage == nil {
if wikiRecoveryFromContext(ctx) != nil {
return fmt.Errorf("%w: missing recovery index page", ErrWikiRecoveryBudgetStopped)
}
return nil
}
// The intro lives on both Content and Summary. Prefer Content since
// that's what the new index view returns; fall back to Summary for
// rows written before this refactor so the incremental-update prompt
// has something to work with.
existingIntro := strings.TrimSpace(indexPage.Content)
if existingIntro == "" {
existingIntro = strings.TrimSpace(indexPage.Summary)
}
// Detect the legacy "intro + directory" payload. Such rows embed the
// fence-separated "## Summary" sections right after the intro, so we
// clip everything from the first directory heading onward to keep the
// intro length bounded when we feed it back into the update prompt.
if idx := strings.Index(existingIntro, "\n## "); idx >= 0 {
existingIntro = strings.TrimSpace(existingIntro[:idx])
}
var intro string
switch {
case existingIntro == "" || existingIntro == "Wiki index - table of contents":
// First-time generation: pull the top-N most-recent summary
// pages via the lite projection. CountByType lets us tell the
// LLM "showing N of M" so it can frame the intro honestly when
// the KB is bigger than what we're sampling.
recentSummaries, listErr := s.wikiService.ListByTypeRecent(ctx, payload.KnowledgeBaseID, types.WikiPageTypeSummary, indexIntroSummaryCap)
if listErr != nil {
return listErr
}
var docSummaries strings.Builder
for _, e := range recentSummaries {
fmt.Fprintf(&docSummaries, "<document>\n<title>%s</title>\n<summary>%s</summary>\n</document>\n\n", e.Title, e.Summary)
}
// Best-effort total count for the framing hint. CountByType
// counts every page type; we need just summary, so we read
// directly. A failure here doesn't block intro generation.
totalSummaries := int64(len(recentSummaries))
if counts, cntErr := s.wikiService.CountByType(ctx, payload.KnowledgeBaseID); cntErr == nil {
if t, ok := counts[types.WikiPageTypeSummary]; ok {
totalSummaries = t
}
}
framing := ""
if int(totalSummaries) > len(recentSummaries) && len(recentSummaries) > 0 {
framing = fmt.Sprintf("(showing %d most recent of %d total documents)\n\n", len(recentSummaries), totalSummaries)
}
if docSummaries.Len() == 0 {
docSummaries.WriteString("(no documents yet)")
}
generatedIntro, genErr := s.generateWithTemplate(ctx, chatModel, agent.WikiIndexIntroPrompt, map[string]string{
"DocumentSummaries": framing + docSummaries.String(),
"Language": lang,
"CustomInstructions": customInstructions,
"InstructionScope": "wiki_content",
})
if genErr != nil {
if requestbudget.FromContext(ctx) != nil {
return genErr
}
intro = "# Wiki Index\n\nThis wiki contains knowledge extracted from uploaded documents.\n"
} else {
intro = strings.TrimSpace(generatedIntro)
}
case changeDesc != "":
// Incremental update: only the existing intro + this batch's
// change description go into the prompt. We deliberately stop
// passing the full DocumentSummaries set here — at 4w docs it
// would re-flood the context every batch, and the
// change-description block already encodes the "what just
// changed" signal the prompt is asking for.
updatedIntro, genErr := s.generateWithTemplate(ctx, chatModel, agent.WikiIndexIntroUpdatePrompt, map[string]string{
"ExistingIntro": existingIntro,
"ChangeDescription": changeDesc,
"DocumentSummaries": "",
"Language": lang,
"CustomInstructions": customInstructions,
"InstructionScope": "wiki_content",
})
if genErr != nil {
if requestbudget.FromContext(ctx) != nil {
return genErr
}
intro = existingIntro // keep existing on error
} else {
intro = strings.TrimSpace(updatedIntro)
}
default:
// No change description and an existing intro: leave it as-is so
// we don't bump the version for a no-op.
intro = existingIntro
}
// Defensive: some LLM outputs occasionally bleed into a directory-
// like section even when the intro prompt doesn't ask for one. If
// the freshly-generated intro starts to look like a legacy payload,
// clip it at the first "\n## " just like we did on the read path
// above. This keeps indexPage.Content a bounded intro-only blob.
if idx := strings.Index(intro, "\n## "); idx >= 0 {
intro = strings.TrimSpace(intro[:idx])
}
indexPage.Content = intro
indexPage.Summary = intro
_, err := s.wikiService.UpdatePage(ctx, indexPage)
return err
}
// splitSummaryLine extracts the "SUMMARY: ..." line from LLM output.
// Returns (summary, content). If no SUMMARY line found, summary is empty.
func splitSummaryLine(raw string) (summary string, content string) {
raw = strings.TrimSpace(raw)
if strings.HasPrefix(raw, "SUMMARY:") || strings.HasPrefix(raw, "SUMMARY:") {
idx := strings.IndexByte(raw, '\n')
if idx < 0 {
// Only one line
return strings.TrimSpace(strings.TrimPrefix(strings.TrimPrefix(raw, "SUMMARY:"), "SUMMARY:")), ""
}
summaryLine := raw[:idx]
summaryLine = strings.TrimPrefix(summaryLine, "SUMMARY:")
summaryLine = strings.TrimPrefix(summaryLine, "SUMMARY:")
return strings.TrimSpace(summaryLine), strings.TrimSpace(raw[idx+1:])
}
return "", raw
}
// publishDraftPages transitions draft pages to published status after ingest completes.
// This ensures users don't see half-built pages during the ingest process.
func (s *wikiIngestService) publishDraftPages(ctx context.Context, kbID string, slugs []string) map[string]error {
failed := make(map[string]error)
for _, slug := range slugs {
if wikiRecoveryFromContext(ctx) != nil {
if err := requestbudget.Check(ctx); err != nil {
failed[slug] = err
return failed
}
}
page, err := s.wikiService.GetPageBySlug(ctx, kbID, slug)
if err != nil {
failed[slug] = err
continue
}
if page == nil {
failed[slug] = apprepo.ErrWikiPageNotFound
continue
}
if page.Status == types.WikiPageStatusDraft {
page.Status = types.WikiPageStatusPublished
if err := s.wikiService.UpdatePageMeta(ctx, page); err != nil {
logger.Warnf(ctx, "wiki ingest: failed to publish page %s: %v", slug, err)
failed[slug] = err
}
}
}
return failed
}
// writeDedupCandidateGroup renders one new item together with its own
// similarity-candidate existing pages, nested under a <candidates> element.
// This per-item grouping is what constrains the dedup model to local
// decisions (see the grouping rationale in deduplicateExtractedBatch). The
// candidate pages keep their aliases so the model still has the acronym /
// translation signal it needs to accept a legitimate merge.
func writeDedupCandidateGroup(
buf *strings.Builder, item extractedItem, itemType string, candidates []*types.WikiPageLite,
) {
fmt.Fprintf(buf, " <item slug=%q type=%q>\n", item.Slug, itemType)
fmt.Fprintf(buf, " <name>%s</name>\n", xmlEscape(item.Name))
for _, alias := range item.Aliases {
if alias == "" {
continue
}
fmt.Fprintf(buf, " <alias>%s</alias>\n", xmlEscape(alias))
}
buf.WriteString(" <candidates>\n")
for _, p := range candidates {
if p == nil {
continue
}
fmt.Fprintf(buf, " <page slug=%q type=%q>\n", p.Slug, p.PageType)
fmt.Fprintf(buf, " <name>%s</name>\n", xmlEscape(p.Title))
for _, alias := range []string(p.Aliases) {
if alias == "" {
continue
}
fmt.Fprintf(buf, " <alias>%s</alias>\n", xmlEscape(alias))
}
buf.WriteString(" </page>\n")
}
buf.WriteString(" </candidates>\n")
buf.WriteString(" </item>\n")
}
// xmlEscape escapes the minimal set of characters that can break XML text
// content. Slugs are ASCII-only so they don't need escaping when used as
// attribute values.
func xmlEscape(s string) string {
s = strings.ReplaceAll(s, "&", "&amp;")
s = strings.ReplaceAll(s, "<", "&lt;")
s = strings.ReplaceAll(s, ">", "&gt;")
return s
}
// deduplicateExtractedBatch deduplicates both entities and concepts against
// existing wiki pages in a single LLM call. Pre-filters candidates via the
// pg_trgm trigram index on lower(title) — every new item issues a
// FindSimilarPages probe and the union of top-K hits across all items is
// the candidate set. This replaces the legacy "ListAllPages + Go-side
// surface-form Jaccard" path that scaled O(P × N) on large KBs.
//
// The KB-id-keyed query relies on idx_wiki_pages_title_trgm (added in
// migration 000041); pg_search environments load pg_trgm in the same
// init step (see migrations/paradedb/00-init-db.sql).
func (s *wikiIngestService) deduplicateExtractedBatch(
ctx context.Context,
chatModel chat.Chat,
kbID string,
entities, concepts []extractedItem,
batchCtx *WikiBatchContext,
) ([]extractedItem, []extractedItem) {
if len(entities) == 0 && len(concepts) == 0 {
return entities, concepts
}
if s.wikiService == nil {
return s.stabilizeExtractedIdentities(ctx, kbID, types.WikiPageTypeEntity, entities, nil, nil, batchCtx),
s.stabilizeExtractedIdentities(ctx, kbID, types.WikiPageTypeConcept, concepts, nil, nil, batchCtx)
}
// Build the candidate set: for each new item, ask the repo for
// the top-K trigram-similar pages and union the results. Dedup by
// slug as we go so the prompt only carries each candidate once.
//
// itemCandidates additionally records, per new item, the slugs that
// surfaced for THAT item specifically. The prompt only ever sees the
// flattened union, so validMerge below uses this per-item scoping to
// reject a merge whose target was pulled in for a *different* item —
// the class of hallucination the union otherwise enables (e.g. weak
// models emitting entity/tencent-open → entity/hiring-agent, which
// share no trigram signal and were never candidates for each other).
candidatePages := make(map[string]*types.WikiPageLite)
itemCandidates := make(map[string]map[string]bool)
probe := func(item extractedItem) {
queries := make([]string, 0, 1+len(item.Aliases))
if item.Name != "" {
queries = append(queries, item.Name)
}
for _, alias := range item.Aliases {
if alias != "" {
queries = append(queries, alias)
}
}
own := itemCandidates[item.Slug]
if own == nil {
own = make(map[string]bool)
itemCandidates[item.Slug] = own
}
for _, q := range queries {
pages, err := s.wikiService.FindSimilarPages(ctx, kbID, q,
[]string{types.WikiPageTypeEntity, types.WikiPageTypeConcept},
dedupCandidateTopK)
if err != nil {
logger.Warnf(ctx, "wiki ingest: dedup FindSimilarPages(%q) failed: %v", q, err)
continue
}
for _, p := range pages {
if p == nil || p.Slug == "" {
continue
}
if _, ok := candidatePages[p.Slug]; !ok {
candidatePages[p.Slug] = p
}
own[p.Slug] = true
}
}
}
for _, e := range entities {
probe(e)
}
for _, c := range concepts {
probe(c)
}
s.attachExactIdentityPages(ctx, kbID, types.WikiPageTypeEntity, entities, candidatePages, itemCandidates, batchCtx)
s.attachExactIdentityPages(ctx, kbID, types.WikiPageTypeConcept, concepts, candidatePages, itemCandidates, batchCtx)
// Resolve exact same-type, same-title candidates deterministically before
// asking the model about semantic/alias variants. Besides avoiding an LLM
// call for the obvious case, this makes an already-materialized page
// authoritative for the identity reservation below.
exactTargets := make(map[string]string)
mergeTargets := make(map[string]string)
collectExactIdentityTargets(entities, types.WikiPageTypeEntity, itemCandidates, candidatePages, exactTargets)
collectExactIdentityTargets(concepts, types.WikiPageTypeConcept, itemCandidates, candidatePages, exactTargets)
stabilize := func() ([]extractedItem, []extractedItem) {
return s.stabilizeExtractedIdentities(
ctx, kbID, types.WikiPageTypeEntity, entities, mergeTargets, exactTargets, batchCtx,
), s.stabilizeExtractedIdentities(
ctx, kbID, types.WikiPageTypeConcept, concepts, mergeTargets, exactTargets, batchCtx,
)
}
if len(candidatePages) == 0 {
// No similar existing pages — identity reservations still make
// concurrent batches converge before they materialize new pages.
logger.Infof(ctx, "wiki ingest: no similar existing pages found for %d new items", len(entities)+len(concepts))
return stabilize()
}
logger.Infof(ctx, "wiki ingest: %d similar existing pages selected for %d new items",
len(candidatePages), len(entities)+len(concepts))
// Group each new item with ONLY the existing pages that surfaced for
// its own similarity probe. Presenting the model two flat lists (all
// new items × all candidates) invites cross-item mispairings — it has
// no way to tell which candidate is relevant to which item, so a weak
// model pairs unrelated slugs that merely coexist in the prompt. A
// per-item shortlist turns dedup into a local yes/no decision against
// a handful of genuinely-similar pages and makes cross-item pairings
// structurally unnatural to express. Items with no candidate are
// omitted entirely (they cannot merge and only add hallucination
// surface + tokens).
var candBuf strings.Builder
groups := 0
renderGroup := func(item extractedItem, itemType string) {
if exactTargets[item.Slug] != "" {
return
}
cset := itemCandidates[item.Slug]
if len(cset) == 0 {
return
}
slugs := make([]string, 0, len(cset))
for slug := range cset {
// Skip the item's own slug: an identically-slugged existing
// page is a re-ingest/update, not a merge target.
if slug == item.Slug {
continue
}
if _, ok := candidatePages[slug]; ok {
slugs = append(slugs, slug)
}
}
if len(slugs) == 0 {
return
}
sort.Strings(slugs)
pages := make([]*types.WikiPageLite, 0, len(slugs))
for _, slug := range slugs {
pages = append(pages, candidatePages[slug])
}
writeDedupCandidateGroup(&candBuf, item, itemType, pages)
groups++
}
for _, item := range entities {
renderGroup(item, "entity")
}
for _, item := range concepts {
renderGroup(item, "concept")
}
if groups == 0 {
// Every item was resolved exactly or has no safe semantic candidate.
return stabilize()
}
dedupeJSON, err := s.generateWithTemplate(ctx, chatModel, agent.WikiDeduplicationPrompt, map[string]string{
"Candidates": candBuf.String(),
})
if err != nil {
logger.Warnf(ctx, "wiki ingest: deduplication LLM call failed: %v", err)
return stabilize()
}
dedupeJSON = cleanLLMJSON(dedupeJSON)
var dedupeResult struct {
Merges map[string]string `json:"merges"`
}
if err := json.Unmarshal([]byte(dedupeJSON), &dedupeResult); err != nil {
logger.Warnf(ctx, "wiki ingest: failed to parse dedup JSON: %v\nRaw: %s", err, dedupeJSON)
return stabilize()
}
validMerge := func(srcSlug, dstSlug string) bool {
if reason := dedupMergeRejectReason(srcSlug, dstSlug, itemCandidates[srcSlug]); reason != "" {
logger.Warnf(ctx, "wiki ingest: dedup rejected %s → %s (%s)", srcSlug, dstSlug, reason)
return false
}
return true
}
for _, item := range entities {
if exactTargets[item.Slug] != "" {
continue
}
if existingSlug, ok := dedupeResult.Merges[item.Slug]; ok && validMerge(item.Slug, existingSlug) {
logger.Infof(ctx, "wiki ingest: dedup merge %s → %s", item.Slug, existingSlug)
mergeTargets[item.Slug] = existingSlug
}
}
for _, item := range concepts {
if exactTargets[item.Slug] != "" {
continue
}
if existingSlug, ok := dedupeResult.Merges[item.Slug]; ok && validMerge(item.Slug, existingSlug) {
logger.Infof(ctx, "wiki ingest: dedup merge %s → %s", item.Slug, existingSlug)
mergeTargets[item.Slug] = existingSlug
}
}
return stabilize()
}
// generateWithTemplate executes a prompt template and calls the LLM with
// bounded exponential-backoff retries for transient infrastructure errors.
//
// Retry policy:
// - Up to wikiLLMMaxAttempts total attempts (initial + retries).
// - Only retry errors classified as transient by isTransientLLMError:
// HTTP 408/429/5xx, context deadline exceeded (when the parent ctx is
// still alive), or generic "timeout"/"connection reset" wording.
// 4xx (except 408/429) is a caller-side fault and fails fast.
// - Backoff is exponential base 2s: 2s, 4s, 8s — roughly wikiLLMBackoffBase
// - 2^(attempt-1). Honors ctx cancellation so the task can abort.
//
// This exists because wiki ingest makes several independent LLM calls per
// document (extraction, summary, dedup, citations, intro) and a single
// transient 504 from the upstream gateway used to drop the document's
// summary page permanently. Retries plus failedOps requeuing (see
// mapOneDocument) turn those events into at-most-a-few-minute hiccups.
func (s *wikiIngestService) generateWithTemplate(ctx context.Context, chatModel chat.Chat, promptTpl string, data map[string]string) (string, error) {
return s.generateWithTemplateRequest(ctx, chatModel, promptTpl, data, "", wikiLLMMaxAttempts)
}
func (s *wikiIngestService) generateWithTemplateRequest(ctx context.Context, chatModel chat.Chat, promptTpl string, data map[string]string, repairInstruction string, maxAttempts int) (string, error) {
if err := requestbudget.Check(ctx); err != nil {
return "", err
}
tmpl, err := template.New("wiki").Parse(promptTpl)
if err != nil {
return "", fmt.Errorf("parse template: %w", err)
}
maskedData, urlMap := maskTemplateDataImageURLs(data)
var buf strings.Builder
if err := tmpl.Execute(&buf, maskedData); err != nil {
return "", fmt.Errorf("execute template: %w", err)
}
prompt := buf.String()
purpose := wikiPromptPurpose(promptTpl)
ctx = requestbudget.WithPhase(requestbudget.WithModel(ctx, chatModel.GetModelID()), purpose)
messages := []chat.Message{{Role: "user", Content: prompt}}
if promptTpl == agent.WikiPageModifyUserPrompt {
systemPrompt := types.AppendCustomPromptInstructions(
agent.WikiPageModifySystemPrompt,
maskedData["CustomInstructions"],
maskedData["InstructionScope"],
)
messages = []chat.Message{
{Role: "system", Content: systemPrompt},
{Role: "user", Content: prompt},
}
} else {
messages[0].Content = types.AppendCustomPromptInstructions(
prompt, maskedData["CustomInstructions"], maskedData["InstructionScope"],
)
}
if repairInstruction != "" {
messages[len(messages)-1].Content += "\n\n" + repairInstruction
}
thinking := false
opts := &chat.ChatOptions{Temperature: 0.3, Thinking: &thinking, MaxTokens: wikiLLMMaxTokens}
prefixFingerprint := chat.PromptPrefixFingerprint(messages, opts)
warmupKey := ""
if promptTpl == agent.WikiPageModifyUserPrompt {
prefixFingerprint = chat.FingerprintPromptPrefix(
messages[0].Content, maskedData["SharedSourceContexts"],
)
if tenantID, ok := types.TenantIDFromContext(ctx); ok {
warmupKey = chat.BuildPromptCacheKey(
tenantID, chatModel.GetModelID(), purpose, prefixFingerprint,
)
}
}
ctx = types.WithLLMCallMetadata(ctx, purpose, prefixFingerprint)
tenantID, tenantScoped := types.TenantIDFromContext(ctx)
requestJSON, _ := json.Marshal(struct {
Messages []chat.Message `json:"messages"`
Options *chat.ChatOptions `json:"options"`
}{Messages: messages, Options: opts})
requestKey := chat.BuildPromptCacheKey(
tenantID, chatModel.GetModelID(), "wiki_exact_request",
chat.FingerprintPromptPrefix(string(requestJSON)),
)
if guard := requestbudget.FromContext(ctx); guard != nil {
requestKey += ":recovery:" + guard.Key()
}
execute := func() (interface{}, error) {
providerKey := wikiProviderCircuitKey(tenantID, chatModel.GetModelID())
if tenantScoped {
if err := s.checkWikiProvider(ctx, providerKey); err != nil {
return "", err
}
}
releaseWarmup := func() {}
if tenantScoped && promptTpl == agent.WikiPageModifyUserPrompt && strings.TrimSpace(maskedData["SharedSourceContexts"]) != "" {
var warmupErr error
releaseWarmup, warmupErr = s.awaitWikiPromptWarmup(ctx, warmupKey)
if warmupErr != nil {
return "", warmupErr
}
}
defer releaseWarmup()
if tenantScoped {
release, acquireErr := acquireWikiProviderSlot(ctx, providerKey)
if acquireErr != nil {
return "", acquireErr
}
defer release()
}
var lastErr error
for attempt := 1; attempt <= maxAttempts; attempt++ {
if err := requestbudget.Check(ctx); err != nil {
return "", err
}
if tenantScoped {
if err := s.checkWikiProvider(ctx, providerKey); err != nil {
return "", err
}
}
callCtx, cancelCall := context.WithTimeout(ctx, wikiLLMCallTimeout)
response, callErr := chatModel.Chat(callCtx, messages, opts)
cancelCall()
if callErr == nil && response != nil {
if outputErr := validateWikiResponse(response); outputErr != nil {
return "", outputErr
}
return response.Content, nil
}
if callErr == nil {
callErr = errors.New("LLM returned nil response")
}
lastErr = callErr
if errors.Is(callErr, ErrWikiRecoveryBudgetStopped) {
if guard := wikiRecoveryFromContext(ctx); guard != nil {
guard.fail(callErr)
}
return "", callErr
}
if tenantScoped {
s.blockWikiProvider(ctx, providerKey, callErr)
if blocked := s.checkWikiProvider(ctx, providerKey); blocked != nil {
return "", blocked
}
}
if !isTransientLLMError(ctx, callErr) {
return "", fmt.Errorf("LLM call failed: %w", callErr)
}
if attempt == maxAttempts {
break
}
backoff := wikiLLMBackoffBase << (attempt - 1)
logger.Warnf(ctx, "wiki ingest: LLM call failed (attempt %d/%d), retrying in %s: %v",
attempt, maxAttempts, backoff, callErr)
timer := time.NewTimer(backoff)
select {
case <-ctx.Done():
timer.Stop()
return "", fmt.Errorf("LLM call aborted during backoff: %w", ctx.Err())
case <-timer.C:
}
}
return "", fmt.Errorf("LLM call failed after %d attempts: %w", maxAttempts, lastErr)
}
// Missing tenant context is unexpected for production Wiki work. Fail safe
// by skipping cross-call coalescing instead of putting unrelated requests
// into a synthetic tenant-0 bucket.
if !tenantScoped {
value, executeErr := execute()
if executeErr != nil {
return "", executeErr
}
content, _ := value.(string)
return unmaskImageURLs(content, urlMap), nil
}
resultCh := s.llmRequests.DoChan(requestKey, execute)
select {
case <-ctx.Done():
return "", ctx.Err()
case result := <-resultCh:
if result.Err != nil {
return "", result.Err
}
content, _ := result.Val.(string)
return unmaskImageURLs(content, urlMap), nil
}
}
func wikiPromptPurpose(promptTpl string) string {
switch promptTpl {
case agent.WikiPageModifyUserPrompt:
return "wiki_page_modify"
case agent.WikiChunkCitationPrompt:
return "wiki_chunk_citation"
case agent.WikiCandidateSlugPrompt:
return "wiki_candidate_slug"
case agent.WikiSummaryPrompt:
return "wiki_summary"
case agent.WikiKnowledgeExtractPrompt:
return "wiki_knowledge_extract"
case agent.WikiTaxonomyPlanPrompt:
return "wiki_taxonomy_plan"
case agent.WikiDeduplicationPrompt:
return "wiki_deduplication"
case agent.WikiIndexIntroPrompt, agent.WikiIndexIntroUpdatePrompt:
return "wiki_index_intro"
default:
return "wiki_generation"
}
}
func (s *wikiIngestService) awaitWikiPromptWarmup(ctx context.Context, key string) (func(), error) {
if key == "" {
return func() {}, nil
}
candidate := &wikiPromptWarmup{done: make(chan struct{})}
actual, loaded := s.promptWarmups.LoadOrStore(key, candidate)
entry := actual.(*wikiPromptWarmup)
if !loaded {
return func() {
entry.once.Do(func() { close(entry.done) })
// Keep the local warmed marker long enough to cover the parallel
// reduce burst without turning it into a persistent application cache.
time.AfterFunc(4*time.Minute, func() {
s.promptWarmups.CompareAndDelete(key, entry)
})
}, nil
}
select {
case <-ctx.Done():
return nil, ctx.Err()
case <-entry.done:
return func() {}, nil
}
}
// isTransientLLMError reports whether an error from the chat provider
// looks like an infrastructure hiccup worth retrying. Classification is
// intentionally conservative: the truthful "could not tell, assume
// permanent" choice keeps retries cheap and avoids masking real bugs.
//
// We treat the following as transient:
// - HTTP 408 (client request timeout — upstream usually didn't process),
// 429 (rate-limited — retry after backoff may succeed), 5xx (any
// server-side fault, including the 504 "Remote error, timeout with
// 60" we see from the gateway in front of several LLM providers).
// - Wrapped context.DeadlineExceeded when the parent ctx is still alive
// (nested per-call timeouts).
// - Substring matches on the error text for common transport failures
// ("timeout", "connection reset", "EOF") that providers surface
// without a structured status code.
func isTransientLLMError(ctx context.Context, err error) bool {
if err == nil {
return false
}
// Never retry after the parent ctx itself expired — the task is
// being cancelled and the next attempt would just fail again.
if ctx.Err() != nil {
return false
}
msg := err.Error()
// Providers that bubble HTTP status up formatted as
// "API request failed with status NNN: ..." — match that first.
for _, s := range []string{
"status 408", "status 429",
"status 500", "status 501", "status 502", "status 503", "status 504",
"status 520", "status 521", "status 522", "status 523", "status 524",
} {
if strings.Contains(msg, s) {
return true
}
}
lower := strings.ToLower(msg)
for _, s := range []string{
"timeout",
"timed out",
"connection reset",
"connection refused",
"broken pipe",
"no such host", // DNS hiccup
"i/o timeout",
"unexpected eof",
"tls handshake",
"context deadline exceeded", // nested per-call deadline
} {
if strings.Contains(lower, s) {
return true
}
}
return false
}
// --- Helpers ---
// isKnowledgeGone returns true if the given knowledge has been deleted or is
// in the middle of being deleted. It first consults the Redis tombstone
// (written by cleanupWikiOnKnowledgeDelete) as a fast path, then falls back
// to the DB. A nil result from GetKnowledgeByIDOnly also counts as gone: the
// repo layer uses GORM First() which filters soft-deleted rows, so a
// soft-deleted knowledge surfaces as "not found" here — exactly what we want.
func (s *wikiIngestService) isKnowledgeGone(ctx context.Context, kbID, knowledgeID string) (bool, error) {
if knowledgeID == "" {
return true, nil
}
if s.redisClient != nil {
if exists, err := s.redisClient.Exists(ctx, WikiDeletedTombstoneKey(kbID, knowledgeID)).Result(); err == nil && exists > 0 {
return true, nil
}
}
kn, err := s.knowledgeSvc.GetKnowledgeByIDOnly(ctx, knowledgeID)
if errors.Is(err, apprepo.ErrKnowledgeNotFound) || errors.Is(err, gorm.ErrRecordNotFound) {
return true, nil
}
if err != nil {
return false, fmt.Errorf("check wiki source knowledge: %w", err)
}
if kn == nil {
return true, nil
}
if kn.KnowledgeBaseID != kbID {
return false, errors.New("wiki source knowledge belongs to a different knowledge base")
}
switch kn.ParseStatus {
case types.ParseStatusDeleting, types.ParseStatusCancelled:
return true, nil
}
return false, nil
}
// filterLiveUpdates drops additions/summaries whose source knowledge has been
// deleted since the Map phase finished. Retract updates are preserved so
// pages still get cleaned up. Caches per-knowledge results to avoid DB
// hammering when a single reduce slug carries many updates for the same doc.
func (s *wikiIngestService) filterLiveUpdates(ctx context.Context, kbID string, updates []SlugUpdate) ([]SlugUpdate, error) {
if len(updates) == 0 {
return updates, nil
}
goneCache := make(map[string]bool)
isGone := func(kid string) (bool, error) {
if kid == "" {
return false, nil
}
if v, ok := goneCache[kid]; ok {
return v, nil
}
v, err := s.isKnowledgeGone(ctx, kbID, kid)
if err != nil {
return false, err
}
goneCache[kid] = v
return v, nil
}
filtered := make([]SlugUpdate, 0, len(updates))
dropped := 0
for _, u := range updates {
switch u.Type {
case "retract", "retractStale":
filtered = append(filtered, u)
default:
gone, err := isGone(u.KnowledgeID)
if err != nil {
return nil, err
}
if gone {
dropped++
continue
}
filtered = append(filtered, u)
}
}
if dropped > 0 {
logger.Infof(ctx, "wiki ingest: reduce dropped %d updates for deleted knowledge(s)", dropped)
}
return filtered, nil
}
// reconstructContent rebuilds document text from chunks.
//
// This only concatenates text-type chunks — image OCR / caption information is
// stored on image_ocr / image_caption child chunks (see image_multimodal.go),
// not on the parent text chunk's ImageInfo field. Callers that need the full
// enriched content (with OCR / captions inlined) should call
// reconstructEnrichedContent instead so image info is fetched from child
// chunks and embedded alongside Markdown image links.
func reconstructContent(chunks []*types.Chunk) string {
var textChunks []*types.Chunk
for _, c := range chunks {
if c.ChunkType == types.ChunkTypeText || c.ChunkType == "" {
textChunks = append(textChunks, c)
}
}
// 重叠去重与排序统一交给公共逻辑(按文本匹配,兼容补写表头 / HTML 实体)。
return searchutil.MergeTextChunks(textChunks, "\n")
}
// reconstructEnrichedContent rebuilds document text and inlines image_info
// (OCR text + caption) pulled from image_ocr / image_caption child chunks.
//
// Without this enrichment, image-heavy documents (e.g. a scanned PDF or a
// standalone .jpg) reach the LLM as bare Markdown image links, causing
// extraction / summarization to produce empty or "no textual content" output.
func reconstructEnrichedContent(
ctx context.Context,
chunkRepo interfaces.ChunkRepository,
tenantID uint64,
chunks []*types.Chunk,
) string {
content := reconstructContent(chunks)
var textChunkIDs []string
for _, c := range chunks {
if c.ChunkType == types.ChunkTypeText || c.ChunkType == "" {
if c.ID != "" {
textChunkIDs = append(textChunkIDs, c.ID)
}
}
}
if len(textChunkIDs) == 0 || chunkRepo == nil {
return content
}
imageInfoMap := searchutil.CollectImageInfoByChunkIDs(ctx, chunkRepo, tenantID, textChunkIDs)
mergedImageInfo := searchutil.MergeImageInfoJSON(imageInfoMap)
if mergedImageInfo == "" {
return content
}
return searchutil.EnrichContentWithImageInfo(content, mergedImageInfo)
}
// slugify creates a URL-friendly slug from a string
func slugify(s string) string {
s = strings.ToLower(strings.TrimSpace(s))
s = strings.Map(func(r rune) rune {
if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '-' || r == '/' {
return r
}
if r == ' ' || r == '_' {
return '-'
}
// Keep CJK characters
if r >= 0x4E00 && r <= 0x9FFF {
return r
}
return -1
}, s)
// Collapse multiple hyphens
for strings.Contains(s, "--") {
s = strings.ReplaceAll(s, "--", "-")
}
s = strings.Trim(s, "-")
if len(s) > 200 {
s = s[:200]
}
return s
}
// truncateString truncates a string to maxLen runes
func truncateString(s string, maxLen int) string {
runes := []rune(s)
if len(runes) <= maxLen {
return s
}
return string(runes[:maxLen]) + "..."
}
// appendUnique appends a string to a StringArray if not already present
func appendUnique(arr types.StringArray, s string) types.StringArray {
for _, v := range arr {
if v == s {
return arr
}
}
return append(arr, s)
}
// minTextContentRunes is the minimum number of non-whitespace, non-image-reference
// runes required for content to be considered substantive enough for LLM
// summarization or wiki extraction. Documents below this threshold (e.g. a
// scanned PDF where OCR yielded nothing AND no caption either) are routed to
// a deterministic empty-content fallback instead of being passed to the LLM,
// which would otherwise hallucinate based on metadata alone.
//
// The threshold is intentionally low: legitimate short documents (brief
// memos, single-line notes) must still pass. The goal is only to catch
// the empty-image-only case.
//
// Declared as a var (not const) so tests can override it and future config
// plumbing can adjust it at runtime without a rebuild.
var minTextContentRunes = 10
var (
// Markdown image references like ![alt](path) — pure visual placeholders
// with no extractable text, so the whole reference is removed.
mdImageRefRE = regexp.MustCompile(`!\[[^\]]*\]\([^)]*\)`)
// <image_original>...</image_original> blocks wrap the verbatim Markdown
// image reference inside an enriched <image> block (see
// searchutil.EnrichContentWithImageInfo). The content is just a redundant
// copy of an already-stripped image link, so the whole block (tags +
// content) is removed.
imageOriginalBlockRE = regexp.MustCompile(`(?is)<image_original\b[^>]*>.*?</image_original>`)
// Self-closing or attribute-only HTML <img> tags.
htmlImgTagRE = regexp.MustCompile(`(?i)<img\b[^>]*/?>`)
// Wrapper-style <image>, <images>, <image_caption>, <image_ocr> tags
// (opening or closing). Matches ONLY the tag; the text content between
// open and close tags is preserved. This is critical: VLM-generated OCR
// and caption text live inside <image_ocr>...</image_ocr> and
// <image_caption>...</image_caption> blocks, and stripping the content
// would silently destroy the very text we want to keep.
imageWrapperTagRE = regexp.MustCompile(`(?i)</?image[a-z_]*\b[^>]*/?>`)
// Markdown image references with the URL captured separately so LLM-bound
// image URLs can be frozen while captions remain editable.
mdImageURLRE = regexp.MustCompile(`!\[[^\]]*\]\(([^)]*)\)`)
// Enriched image blocks store the original object URL as an attribute,
// e.g. <image url="...">. Capture both double- and single-quoted forms.
imageURLAttrRE = regexp.MustCompile(`(?i)<image\b[^>]*\surl\s*=\s*(?:"([^"]*)"|'([^']*)')`)
imagePlaceholderTokenRE = regexp.MustCompile(`wkimg:[A-Za-z0-9_-]+`)
)
func maskTemplateDataImageURLs(data map[string]string) (map[string]string, map[string]string) {
if len(data) == 0 {
return data, nil
}
masked := make(map[string]string, len(data))
urlToToken := make(map[string]string)
tokenToURL := make(map[string]string)
keys := make([]string, 0, len(data))
for key := range data {
keys = append(keys, key)
}
sort.Strings(keys)
for _, key := range keys {
masked[key] = maskImageURLsWithState(data[key], urlToToken, tokenToURL)
}
return masked, tokenToURL
}
// maskImageURLs replaces image URLs with low-entropy placeholders. It only
// freezes URLs; alt/caption text remains in place for the LLM to edit.
func maskImageURLs(s string) (string, map[string]string) {
urlToToken := make(map[string]string)
tokenToURL := make(map[string]string)
return maskImageURLsWithState(s, urlToToken, tokenToURL), tokenToURL
}
func maskImageURLsWithState(s string, urlToToken, tokenToURL map[string]string) string {
urls := collectMaskableImageURLs(s)
if len(urls) == 0 {
return s
}
for _, url := range urls {
if _, ok := urlToToken[url]; ok {
continue
}
token := fmt.Sprintf("wkimg:%04d", len(tokenToURL)+1)
urlToToken[url] = token
tokenToURL[token] = url
}
replaceURLs := append([]string(nil), urls...)
sort.SliceStable(replaceURLs, func(i, j int) bool {
return len(replaceURLs[i]) > len(replaceURLs[j])
})
masked := s
for _, url := range replaceURLs {
masked = strings.ReplaceAll(masked, url, urlToToken[url])
}
return masked
}
func collectMaskableImageURLs(s string) []string {
seen := make(map[string]struct{})
var urls []string
addURL := func(url string) {
url = strings.TrimSpace(url)
if url == "" {
return
}
if _, ok := seen[url]; ok {
return
}
seen[url] = struct{}{}
urls = append(urls, url)
}
for _, match := range mdImageURLRE.FindAllStringSubmatch(s, -1) {
addURL(match[1])
}
for _, match := range imageURLAttrRE.FindAllStringSubmatch(s, -1) {
if match[1] != "" {
addURL(match[1])
continue
}
addURL(match[2])
}
return urls
}
// unmaskImageURLs restores known placeholders and drops any corrupted or
// invented image placeholders so broken image links never reach storage.
func unmaskImageURLs(out string, urlMap map[string]string) string {
out = mdImageURLRE.ReplaceAllStringFunc(out, func(match string) string {
parts := mdImageURLRE.FindStringSubmatch(match)
if len(parts) != 2 {
return match
}
url := strings.TrimSpace(parts[1])
if realURL, ok := urlMap[url]; ok {
idx := strings.LastIndex(match, "(")
if idx < 0 {
return match
}
return match[:idx+1] + realURL + ")"
}
if strings.HasPrefix(url, "wkimg:") {
return ""
}
return match
})
return replaceImagePlaceholderTokensOutsideMarkdown(out, urlMap)
}
func replaceImagePlaceholderTokensOutsideMarkdown(s string, urlMap map[string]string) string {
matches := mdImageURLRE.FindAllStringIndex(s, -1)
if len(matches) == 0 {
return replaceImagePlaceholderTokens(s, urlMap)
}
var b strings.Builder
last := 0
for _, match := range matches {
if match[0] > last {
b.WriteString(replaceImagePlaceholderTokens(s[last:match[0]], urlMap))
}
b.WriteString(s[match[0]:match[1]])
last = match[1]
}
if last < len(s) {
b.WriteString(replaceImagePlaceholderTokens(s[last:], urlMap))
}
return b.String()
}
func replaceImagePlaceholderTokens(s string, urlMap map[string]string) string {
return imagePlaceholderTokenRE.ReplaceAllStringFunc(s, func(token string) string {
if realURL, ok := urlMap[token]; ok {
return realURL
}
return ""
})
}
// stripImageMarkup removes image-only placeholders (Markdown image refs,
// <img> tags, <image_original> redundancy blocks) and unwraps the
// <image>/<image_caption>/<image_ocr> XML wrappers produced by the search
// enrichment layer, leaving any OCR or caption text as plain inline text.
//
// This shape matters: when VLM OCR succeeds on a scanned PDF page, the
// extracted text reaches downstream code wrapped in <image_ocr> tags inside
// an <image> block. A naive "strip the whole <image>...</image> block"
// approach would discard the OCR text — the exact opposite of what we want.
func stripImageMarkup(s string) string {
s = imageOriginalBlockRE.ReplaceAllString(s, "")
s = mdImageRefRE.ReplaceAllString(s, "")
s = htmlImgTagRE.ReplaceAllString(s, "")
s = imageWrapperTagRE.ReplaceAllString(s, "")
return s
}
// extractRealText returns the trimmed content with image markup stripped.
// Cached at the call site for use both in the threshold check and in any
// subsequent log message, avoiding redundant regex passes over large docs.
func extractRealText(content string) string {
return strings.TrimSpace(stripImageMarkup(content))
}
// hasSufficientTextContent reports whether the given content carries enough
// real text (after image markup is stripped, with OCR/caption text retained)
// to warrant an LLM call. It is the primary defence against filename-driven
// hallucinations on scanned PDFs that have NO usable text at all.
func hasSufficientTextContent(content string) bool {
return realTextRuneCount(content) >= minTextContentRunes
}
// realTextRuneCount returns the rune length of the content after image
// markup is stripped. Uses utf8.RuneCountInString to avoid allocating a
// rune slice for the count.
func realTextRuneCount(content string) int {
return utf8.RuneCountInString(extractRealText(content))
}
// cleanLLMJSON strips markdown code-fence wrappers and sanitizes control characters
// from LLM-generated JSON output so it can be safely unmarshalled.
func cleanLLMJSON(s string) string {
s = strings.TrimSpace(s)
s = strings.TrimPrefix(s, "```json")
s = strings.TrimPrefix(s, "```")
s = strings.TrimSuffix(s, "```")
s = strings.TrimSpace(s)
return sanitizeJSONString(s)
}
// sanitizeJSONString sanitizes a string that is intended to be parsed as JSON,
// by properly escaping unescaped control characters (like newlines) inside string literals.
func sanitizeJSONString(s string) string {
var buf strings.Builder
buf.Grow(len(s))
inString := false
escape := false
for _, r := range s {
if escape {
if r == '\n' {
buf.WriteString(`n`)
} else if r == '\r' {
buf.WriteString(`r`)
} else if r == '\t' {
buf.WriteString(`t`)
} else {
buf.WriteRune(r)
}
escape = false
continue
}
if r == '\\' {
escape = true
buf.WriteRune(r)
continue
}
if r == '"' {
inString = !inString
buf.WriteRune(r)
continue
}
if inString {
if r == '\n' {
buf.WriteString(`\n`)
continue
}
if r == '\r' {
buf.WriteString(`\r`)
continue
}
if r == '\t' {
buf.WriteString(`\t`)
continue
}
}
buf.WriteRune(r)
}
return buf.String()
}