From 679f59b1420185afd2e2d709babfc41e4a17c0e0 Mon Sep 17 00:00:00 2001 From: Even1115 Date: Tue, 28 Jul 2026 09:19:36 +0800 Subject: [PATCH 01/11] feat: add versioned artifact DAG reconciliation --- internal/application/repository/chunk.go | 14 + internal/application/repository/knowledge.go | 67 + .../knowledge_attempt_fence_test.go | 99 ++ .../repository/knowledge_span_repo.go | 67 +- .../repository/knowledge_span_repo_test.go | 45 + .../repository/processing_artifact.go | 247 ++++ .../repository/processing_artifact_test.go | 106 ++ .../repository/retriever/neo4j/repository.go | 315 ++++- .../application/service/docreader_artifact.go | 266 ++++ .../service/docreader_artifact_test.go | 161 +++ internal/application/service/extract.go | 235 +++- internal/application/service/graph.go | 41 +- .../application/service/image_multimodal.go | 347 ++++- .../service/image_multimodal_prompt_test.go | 63 + internal/application/service/knowledge.go | 6 +- .../application/service/knowledge_create.go | 84 +- .../application/service/knowledge_process.go | 1211 ++++++++++++----- .../service/knowledge_reconcile.go | 252 ++++ .../service/knowledge_reconcile_test.go | 284 ++++ internal/application/service/model.go | 43 +- .../keywords_vector_hybrid_indexer.go | 30 +- internal/application/service/wiki_ingest.go | 43 + .../application/service/wiki_ingest_test.go | 24 + internal/artifact/codec.go | 195 +++ internal/artifact/codec_test.go | 48 + internal/artifact/fence.go | 47 + internal/artifact/fence_test.go | 28 + internal/artifact/identity.go | 139 ++ internal/artifact/identity_test.go | 76 ++ internal/artifact/key.go | 304 +++++ internal/artifact/key_test.go | 133 ++ internal/artifact/lease.go | 132 ++ internal/artifact/reconcile.go | 97 ++ internal/artifact/reconcile_test.go | 48 + internal/artifact/runtime.go | 415 ++++++ internal/artifact/runtime_test.go | 229 ++++ internal/container/container.go | 11 + .../processing_artifact_migration_test.go | 58 + internal/models/chat/artifact_cache.go | 221 +++ internal/models/chat/artifact_cache_test.go | 130 ++ internal/models/embedding/artifact_cache.go | 342 +++++ .../models/embedding/artifact_cache_test.go | 188 +++ internal/models/embedding/batch.go | 8 + internal/models/vlm/artifact_cache.go | 188 +++ internal/models/vlm/artifact_cache_test.go | 95 ++ internal/testutil/artifactrepo/repository.go | 114 ++ internal/types/const.go | 3 + internal/types/interfaces/chunk.go | 3 + internal/types/interfaces/knowledge.go | 19 + .../types/interfaces/processing_artifact.go | 34 + internal/types/interfaces/retriever_graph.go | 20 + internal/types/knowledge_span.go | 13 + internal/types/processing_artifact.go | 121 ++ internal/types/task.go | 1 + migrations/mysql/00-init-db.sql | 37 + migrations/paradedb/00-init-db.sql | 39 + .../000001_processing_artifacts.down.sql | 2 + .../sqlite/000001_processing_artifacts.up.sql | 38 + .../000077_processing_artifacts.down.sql | 3 + .../000077_processing_artifacts.up.sql | 44 + 60 files changed, 7192 insertions(+), 481 deletions(-) create mode 100644 internal/application/repository/knowledge_attempt_fence_test.go create mode 100644 internal/application/repository/processing_artifact.go create mode 100644 internal/application/repository/processing_artifact_test.go create mode 100644 internal/application/service/docreader_artifact.go create mode 100644 internal/application/service/docreader_artifact_test.go create mode 100644 internal/application/service/knowledge_reconcile.go create mode 100644 internal/application/service/knowledge_reconcile_test.go create mode 100644 internal/artifact/codec.go create mode 100644 internal/artifact/codec_test.go create mode 100644 internal/artifact/fence.go create mode 100644 internal/artifact/fence_test.go create mode 100644 internal/artifact/identity.go create mode 100644 internal/artifact/identity_test.go create mode 100644 internal/artifact/key.go create mode 100644 internal/artifact/key_test.go create mode 100644 internal/artifact/lease.go create mode 100644 internal/artifact/reconcile.go create mode 100644 internal/artifact/reconcile_test.go create mode 100644 internal/artifact/runtime.go create mode 100644 internal/artifact/runtime_test.go create mode 100644 internal/database/processing_artifact_migration_test.go create mode 100644 internal/models/chat/artifact_cache.go create mode 100644 internal/models/chat/artifact_cache_test.go create mode 100644 internal/models/embedding/artifact_cache.go create mode 100644 internal/models/embedding/artifact_cache_test.go create mode 100644 internal/models/vlm/artifact_cache.go create mode 100644 internal/models/vlm/artifact_cache_test.go create mode 100644 internal/testutil/artifactrepo/repository.go create mode 100644 internal/types/interfaces/processing_artifact.go create mode 100644 internal/types/processing_artifact.go create mode 100644 migrations/sqlite/000001_processing_artifacts.down.sql create mode 100644 migrations/sqlite/000001_processing_artifacts.up.sql create mode 100644 migrations/versioned/000077_processing_artifacts.down.sql create mode 100644 migrations/versioned/000077_processing_artifacts.up.sql diff --git a/internal/application/repository/chunk.go b/internal/application/repository/chunk.go index 08723f3b64..e50fa8c1a9 100644 --- a/internal/application/repository/chunk.go +++ b/internal/application/repository/chunk.go @@ -159,6 +159,20 @@ func (r *chunkRepository) ListChunksByKnowledgeID( return chunks, nil } +// ListAllChunksByKnowledgeID lists all active chunk types for reconciliation. +func (r *chunkRepository) ListAllChunksByKnowledgeID( + ctx context.Context, tenantID uint64, knowledgeID string, +) ([]*types.Chunk, error) { + var chunks []*types.Chunk + if err := r.db.WithContext(ctx). + Where("tenant_id = ? AND knowledge_id = ?", tenantID, knowledgeID). + Order("chunk_index ASC, created_at ASC"). + Find(&chunks).Error; err != nil { + return nil, err + } + return chunks, nil +} + // ListPagedChunksByKnowledgeID lists chunks for a knowledge ID with pagination func (r *chunkRepository) ListPagedChunksByKnowledgeID( ctx context.Context, diff --git a/internal/application/repository/knowledge.go b/internal/application/repository/knowledge.go index 095dcaadac..a59b90e554 100644 --- a/internal/application/repository/knowledge.go +++ b/internal/application/repository/knowledge.go @@ -198,6 +198,73 @@ func (r *knowledgeRepository) UpdateKnowledge(ctx context.Context, knowledge *ty return err } +func (r *knowledgeRepository) UpdateKnowledgeIfAttemptCurrent( + ctx context.Context, + knowledge *types.Knowledge, + attempt int, +) (bool, error) { + if knowledge == nil { + return false, errors.New("knowledge must not be nil") + } + if attempt <= 0 { + return true, r.UpdateKnowledge(ctx, knowledge) + } + newerAttempt := r.db. + Table(types.KnowledgeProcessingSpan{}.TableName()). + Select("1"). + Where("knowledge_id = ? AND kind = ? AND attempt > ?", + knowledge.ID, + types.SpanKindRoot, + attempt, + ) + result := r.db.WithContext(ctx). + Model(&types.Knowledge{}). + Where("tenant_id = ? AND id = ?", knowledge.TenantID, knowledge.ID). + Where("NOT EXISTS (?)", newerAttempt). + Select("*"). + Omit(omitFieldsOnUpdate...). + Updates(knowledge) + if result.Error != nil { + return false, result.Error + } + return result.RowsAffected == 1, nil +} + +func (r *knowledgeRepository) UpdateKnowledgeColumnsIfAttemptCurrent( + ctx context.Context, + tenantID uint64, + knowledgeID string, + attempt int, + values map[string]interface{}, +) (bool, error) { + if tenantID == 0 || knowledgeID == "" { + return false, errors.New("tenant and knowledge IDs must not be empty") + } + if len(values) == 0 { + return false, errors.New("knowledge update values must not be empty") + } + query := r.db.WithContext(ctx). + Model(&types.Knowledge{}). + Where("tenant_id = ? AND id = ?", tenantID, knowledgeID) + if attempt > 0 { + newerAttempt := r.db. + Table(types.KnowledgeProcessingSpan{}.TableName()). + Select("1"). + Where( + "knowledge_id = ? AND kind = ? AND attempt > ?", + knowledgeID, + types.SpanKindRoot, + attempt, + ) + query = query.Where("NOT EXISTS (?)", newerAttempt) + } + result := query.Updates(values) + if result.Error != nil { + return false, result.Error + } + return result.RowsAffected == 1, nil +} + // UpdateKnowledgeBatch updates knowledge items in batch func (r *knowledgeRepository) UpdateKnowledgeBatch(ctx context.Context, knowledgeList []*types.Knowledge) error { if len(knowledgeList) == 0 { diff --git a/internal/application/repository/knowledge_attempt_fence_test.go b/internal/application/repository/knowledge_attempt_fence_test.go new file mode 100644 index 0000000000..c8fa50e232 --- /dev/null +++ b/internal/application/repository/knowledge_attempt_fence_test.go @@ -0,0 +1,99 @@ +package repository + +import ( + "context" + "testing" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" +) + +func TestUpdateKnowledgeIfAttemptCurrentIsDatabaseFenced(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&types.Knowledge{})) + require.NoError(t, db.Exec(` + CREATE TABLE knowledge_processing_spans ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + knowledge_id TEXT NOT NULL, + attempt INTEGER NOT NULL, + kind TEXT NOT NULL + ) + `).Error) + + knowledge := &types.Knowledge{ + ID: uuid.New().String(), + TenantID: 1, + KnowledgeBaseID: uuid.New().String(), + Type: types.KnowledgeTypeManual, + Title: "before", + Source: "manual", + ParseStatus: types.ParseStatusProcessing, + EnableStatus: "enabled", + } + require.NoError(t, db.Create(knowledge).Error) + repository := &knowledgeRepository{db: db} + + knowledge.Title = "attempt one" + published, err := repository.UpdateKnowledgeIfAttemptCurrent( + context.Background(), + knowledge, + 1, + ) + require.NoError(t, err) + require.True(t, published) + + require.NoError(t, db.Exec(` + INSERT INTO knowledge_processing_spans (knowledge_id, attempt, kind) + VALUES (?, 2, ?) + `, knowledge.ID, types.SpanKindRoot).Error) + knowledge.Title = "stale attempt one" + published, err = repository.UpdateKnowledgeIfAttemptCurrent( + context.Background(), + knowledge, + 1, + ) + require.NoError(t, err) + assert.False(t, published) + + var stored types.Knowledge + require.NoError(t, db.First(&stored, "id = ?", knowledge.ID).Error) + assert.Equal(t, "attempt one", stored.Title) + + knowledge.Title = "attempt two" + published, err = repository.UpdateKnowledgeIfAttemptCurrent( + context.Background(), + knowledge, + 2, + ) + require.NoError(t, err) + assert.True(t, published) + require.NoError(t, db.First(&stored, "id = ?", knowledge.ID).Error) + assert.Equal(t, "attempt two", stored.Title) + + published, err = repository.UpdateKnowledgeColumnsIfAttemptCurrent( + context.Background(), + knowledge.TenantID, + knowledge.ID, + 1, + map[string]interface{}{"pending_subtasks_count": 7}, + ) + require.NoError(t, err) + assert.False(t, published) + + published, err = repository.UpdateKnowledgeColumnsIfAttemptCurrent( + context.Background(), + knowledge.TenantID, + knowledge.ID, + 2, + map[string]interface{}{"pending_subtasks_count": 3}, + ) + require.NoError(t, err) + assert.True(t, published) + require.NoError(t, db.First(&stored, "id = ?", knowledge.ID).Error) + assert.Equal(t, 3, stored.PendingSubtasksCount) +} diff --git a/internal/application/repository/knowledge_span_repo.go b/internal/application/repository/knowledge_span_repo.go index 1defb0f786..5426389654 100644 --- a/internal/application/repository/knowledge_span_repo.go +++ b/internal/application/repository/knowledge_span_repo.go @@ -3,6 +3,7 @@ package repository import ( "context" "errors" + "strings" "time" "github.com/Tencent/WeKnora/internal/types" @@ -105,15 +106,65 @@ func (r *knowledgeSpanRepository) Upsert(ctx context.Context, row *types.Knowled } func (r *knowledgeSpanRepository) NextAttempt(ctx context.Context, knowledgeID string) (int, error) { - var max int - err := r.db.WithContext(ctx).Model(&types.KnowledgeProcessingSpan{}). - Where("knowledge_id = ?", knowledgeID). - Select("COALESCE(MAX(attempt), 0)"). - Row().Scan(&max) - if err != nil { - return 0, err + if knowledgeID == "" { + return 0, errors.New("knowledgeSpanRepository.NextAttempt: knowledge_id required") + } + var attempt int + var err error + for retry := 0; retry < 5; retry++ { + err = r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var max int + if scanErr := tx.Model(&types.KnowledgeProcessingSpan{}). + Where("knowledge_id = ?", knowledgeID). + Select("COALESCE(MAX(attempt), 0)"). + Row().Scan(&max); scanErr != nil { + return scanErr + } + + // The insert is intentionally the first write in the transaction. + // PostgreSQL/MySQL wait on the unique-key conflict; SQLite takes its + // write lock here, before the read-modify-write below. + if createErr := tx.Clauses(clause.OnConflict{DoNothing: true}). + Create(&types.KnowledgeAttemptCounter{ + KnowledgeID: knowledgeID, + LastAttempt: max, + }).Error; createErr != nil { + return createErr + } + + var counter types.KnowledgeAttemptCounter + if readErr := tx.Clauses(clause.Locking{Strength: "UPDATE"}). + Where("knowledge_id = ?", knowledgeID). + Take(&counter).Error; readErr != nil { + return readErr + } + if counter.LastAttempt < max { + counter.LastAttempt = max + } + counter.LastAttempt++ + if updateErr := tx.Model(&types.KnowledgeAttemptCounter{}). + Where("knowledge_id = ?", knowledgeID). + Update("last_attempt", counter.LastAttempt).Error; updateErr != nil { + return updateErr + } + attempt = counter.LastAttempt + return nil + }) + if err == nil || !isSQLiteBusy(err) { + return attempt, err + } + time.Sleep(time.Duration(retry+1) * time.Millisecond) + } + return 0, err +} + +func isSQLiteBusy(err error) bool { + if err == nil { + return false } - return max + 1, nil + message := strings.ToLower(err.Error()) + return strings.Contains(message, "database is locked") || + strings.Contains(message, "database table is locked") } func (r *knowledgeSpanRepository) LatestAttempt(ctx context.Context, knowledgeID string) (int, error) { diff --git a/internal/application/repository/knowledge_span_repo_test.go b/internal/application/repository/knowledge_span_repo_test.go index 341a3d4fd6..e115255281 100644 --- a/internal/application/repository/knowledge_span_repo_test.go +++ b/internal/application/repository/knowledge_span_repo_test.go @@ -2,6 +2,9 @@ package repository import ( "context" + "fmt" + "sort" + "sync" "testing" "time" @@ -40,6 +43,11 @@ CREATE TABLE IF NOT EXISTS knowledge_processing_spans ( updated_at DATETIME DEFAULT CURRENT_TIMESTAMP, UNIQUE (knowledge_id, attempt, span_id) ); +CREATE TABLE IF NOT EXISTS knowledge_attempt_counters ( + knowledge_id VARCHAR(64) PRIMARY KEY, + last_attempt INTEGER NOT NULL, + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP +); ` func setupSpanTestRepo(t *testing.T) (KnowledgeSpanRepository, *gorm.DB) { @@ -113,6 +121,43 @@ func TestKnowledgeSpanRepo_NextAttempt(t *testing.T) { assert.Equal(t, 1, other, "NextAttempt must scope to the knowledge_id") } +func TestKnowledgeSpanRepo_NextAttemptConcurrent(t *testing.T) { + dsn := fmt.Sprintf( + "file:attempt-counter-%d?mode=memory&cache=shared&_busy_timeout=5000", + time.Now().UnixNano(), + ) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.Exec(spansTestDDL).Error) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(16) + + repo := NewKnowledgeSpanRepository(db) + const workers = 12 + attempts := make([]int, workers) + errorsByWorker := make([]error, workers) + var wg sync.WaitGroup + for i := range workers { + wg.Add(1) + go func(index int) { + defer wg.Done() + attempts[index], errorsByWorker[index] = repo.NextAttempt( + context.Background(), + "concurrent-knowledge", + ) + }(i) + } + wg.Wait() + for _, allocateErr := range errorsByWorker { + require.NoError(t, allocateErr) + } + sort.Ints(attempts) + for index, attempt := range attempts { + assert.Equal(t, index+1, attempt) + } +} + // TestKnowledgeSpanRepo_CancelDescendants verifies the cascade walk: // failing a stage cancels every pending/running descendant in its // subtree, while terminal states (done/skipped/failed) are left intact. diff --git a/internal/application/repository/processing_artifact.go b/internal/application/repository/processing_artifact.go new file mode 100644 index 0000000000..f076876812 --- /dev/null +++ b/internal/application/repository/processing_artifact.go @@ -0,0 +1,247 @@ +package repository + +import ( + "context" + "errors" + "fmt" + "sort" + "time" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +const processingArtifactBatchSize = 500 + +type processingArtifactRepository struct { + db *gorm.DB +} + +func NewProcessingArtifactRepository(db *gorm.DB) interfaces.ProcessingArtifactRepository { + return &processingArtifactRepository{db: db} +} + +func (r *processingArtifactRepository) Get( + ctx context.Context, + key types.ProcessingArtifactLookup, +) (*types.ProcessingArtifact, error) { + if err := key.Validate(); err != nil { + return nil, err + } + var artifact types.ProcessingArtifact + err := r.db.WithContext(ctx). + Where( + "tenant_id = ? AND stage = ? AND key_version = ? AND artifact_key = ?", + key.TenantID, + key.Stage, + key.KeyVersion, + key.ArtifactKey, + ). + Take(&artifact).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, types.ErrProcessingArtifactNotFound + } + if err != nil { + return nil, err + } + return &artifact, nil +} + +func (r *processingArtifactRepository) BatchGet( + ctx context.Context, + keys []types.ProcessingArtifactLookup, +) (map[types.ProcessingArtifactLookup]*types.ProcessingArtifact, error) { + result := make(map[types.ProcessingArtifactLookup]*types.ProcessingArtifact, len(keys)) + groups, err := groupProcessingArtifactKeys(keys) + if err != nil { + return nil, err + } + + for group, artifactKeys := range groups { + for start := 0; start < len(artifactKeys); start += processingArtifactBatchSize { + end := min(start+processingArtifactBatchSize, len(artifactKeys)) + var artifacts []*types.ProcessingArtifact + err := r.db.WithContext(ctx). + Where( + "tenant_id = ? AND stage = ? AND key_version = ?", + group.tenantID, + group.stage, + group.keyVersion, + ). + Where("artifact_key IN ?", artifactKeys[start:end]). + Find(&artifacts).Error + if err != nil { + return nil, err + } + for _, artifact := range artifacts { + result[artifact.Lookup()] = artifact + } + } + } + return result, nil +} + +func (r *processingArtifactRepository) PutIfAbsent( + ctx context.Context, + candidate *types.ProcessingArtifact, +) (*types.ProcessingArtifact, bool, error) { + if candidate == nil { + return nil, false, errors.New("processing artifact candidate must not be nil") + } + if err := candidate.Validate(); err != nil { + return nil, false, err + } + + result := r.db.WithContext(ctx).Clauses(clause.OnConflict{ + Columns: []clause.Column{ + {Name: "tenant_id"}, + {Name: "stage"}, + {Name: "key_version"}, + {Name: "artifact_key"}, + }, + DoNothing: true, + }).Create(candidate) + if result.Error != nil { + return nil, false, result.Error + } + if result.RowsAffected == 1 { + return candidate, true, nil + } + + winner, err := r.Get(ctx, candidate.Lookup()) + if err != nil { + return nil, false, fmt.Errorf("load processing artifact conflict winner: %w", err) + } + return winner, false, nil +} + +func (r *processingArtifactRepository) PutManyIfAbsent( + ctx context.Context, + candidates []*types.ProcessingArtifact, +) (map[types.ProcessingArtifactLookup]*types.ProcessingArtifact, error) { + if len(candidates) == 0 { + return map[types.ProcessingArtifactLookup]*types.ProcessingArtifact{}, nil + } + keys := make([]types.ProcessingArtifactLookup, 0, len(candidates)) + unique := make(map[types.ProcessingArtifactLookup]struct{}, len(candidates)) + for _, candidate := range candidates { + if candidate == nil { + return nil, errors.New("processing artifact candidate must not be nil") + } + if err := candidate.Validate(); err != nil { + return nil, err + } + key := candidate.Lookup() + if _, exists := unique[key]; exists { + return nil, fmt.Errorf("duplicate processing artifact candidate %v", key) + } + unique[key] = struct{}{} + keys = append(keys, key) + } + + err := r.db.WithContext(ctx).Clauses(clause.OnConflict{ + Columns: []clause.Column{ + {Name: "tenant_id"}, + {Name: "stage"}, + {Name: "key_version"}, + {Name: "artifact_key"}, + }, + DoNothing: true, + }).CreateInBatches(candidates, processingArtifactBatchSize).Error + if err != nil { + return nil, err + } + return r.BatchGet(ctx, keys) +} + +func (r *processingArtifactRepository) DeleteCorrupt( + ctx context.Context, + key types.ProcessingArtifactLookup, + observedChecksum string, +) error { + if err := key.Validate(); err != nil { + return err + } + return r.db.WithContext(ctx). + Where( + "tenant_id = ? AND stage = ? AND key_version = ? AND artifact_key = ? AND payload_checksum = ?", + key.TenantID, + key.Stage, + key.KeyVersion, + key.ArtifactKey, + observedChecksum, + ). + Delete(&types.ProcessingArtifact{}).Error +} + +func (r *processingArtifactRepository) TouchHits( + ctx context.Context, + keys []types.ProcessingArtifactLookup, +) error { + groups, err := groupProcessingArtifactKeys(keys) + if err != nil { + return err + } + now := time.Now().UTC() + for group, artifactKeys := range groups { + for start := 0; start < len(artifactKeys); start += processingArtifactBatchSize { + end := min(start+processingArtifactBatchSize, len(artifactKeys)) + err := r.db.WithContext(ctx). + Model(&types.ProcessingArtifact{}). + Where( + "tenant_id = ? AND stage = ? AND key_version = ?", + group.tenantID, + group.stage, + group.keyVersion, + ). + Where("artifact_key IN ?", artifactKeys[start:end]). + Updates(map[string]any{ + "hit_count": gorm.Expr("hit_count + 1"), + "last_hit_at": now, + }).Error + if err != nil { + return err + } + } + } + return nil +} + +type processingArtifactKeyGroup struct { + tenantID uint64 + stage string + keyVersion uint16 +} + +func groupProcessingArtifactKeys( + keys []types.ProcessingArtifactLookup, +) (map[processingArtifactKeyGroup][]string, error) { + sets := make(map[processingArtifactKeyGroup]map[string]struct{}) + for _, key := range keys { + if err := key.Validate(); err != nil { + return nil, err + } + group := processingArtifactKeyGroup{ + tenantID: key.TenantID, + stage: key.Stage, + keyVersion: key.KeyVersion, + } + if sets[group] == nil { + sets[group] = make(map[string]struct{}) + } + sets[group][key.ArtifactKey] = struct{}{} + } + + groups := make(map[processingArtifactKeyGroup][]string, len(sets)) + for group, set := range sets { + values := make([]string, 0, len(set)) + for value := range set { + values = append(values, value) + } + sort.Strings(values) + groups[group] = values + } + return groups, nil +} diff --git a/internal/application/repository/processing_artifact_test.go b/internal/application/repository/processing_artifact_test.go new file mode 100644 index 0000000000..d3d575ecaa --- /dev/null +++ b/internal/application/repository/processing_artifact_test.go @@ -0,0 +1,106 @@ +package repository + +import ( + "context" + "testing" + + "github.com/Tencent/WeKnora/internal/artifact" + "github.com/Tencent/WeKnora/internal/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" +) + +func setupProcessingArtifactRepository(t *testing.T) (*gorm.DB, *processingArtifactRepository) { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&types.ProcessingArtifact{})) + return db, &processingArtifactRepository{db: db} +} + +func processingArtifactCandidate(t *testing.T, tenantID uint64, payload string) *types.ProcessingArtifact { + t.Helper() + key, err := artifact.BuildKey(tenantID, artifact.KeyMaterial{ + KeyVersion: 1, + Stage: "embedding", + DirectInputs: []artifact.DirectInput{{ + Role: "text", + Digest: artifact.SHA256Hex([]byte("input")), + }}, + Processor: artifact.ProcessorIdentity{ModelID: "model"}, + RenderedRequest: "exact input", + Options: map[string]any{"dimensions": 3}, + CanonicalizerVersion: artifact.CanonicalJSONVersion, + OutputSchemaVersion: "embedding.float32.v1", + }) + require.NoError(t, err) + candidate, err := artifact.NewInlineArtifact(key, artifact.CodecFloat32BEV1, []byte(payload)) + require.NoError(t, err) + return candidate +} + +func TestProcessingArtifactPutIfAbsentKeepsWinner(t *testing.T) { + _, repository := setupProcessingArtifactRepository(t) + ctx := context.Background() + first := processingArtifactCandidate(t, 1, "first") + second := processingArtifactCandidate(t, 1, "second") + + winner, created, err := repository.PutIfAbsent(ctx, first) + require.NoError(t, err) + assert.True(t, created) + assert.Equal(t, "first", string(winner.Payload)) + + winner, created, err = repository.PutIfAbsent(ctx, second) + require.NoError(t, err) + assert.False(t, created) + assert.Equal(t, "first", string(winner.Payload)) +} + +func TestProcessingArtifactTenantBoundaryAndBatchGet(t *testing.T) { + _, repository := setupProcessingArtifactRepository(t) + ctx := context.Background() + first := processingArtifactCandidate(t, 1, "tenant-one") + second := processingArtifactCandidate(t, 2, "tenant-two") + + winners, err := repository.PutManyIfAbsent(ctx, []*types.ProcessingArtifact{first, second}) + require.NoError(t, err) + require.Len(t, winners, 2) + assert.Equal(t, first.ArtifactKey, second.ArtifactKey) + assert.Equal(t, "tenant-one", string(winners[first.Lookup()].Payload)) + assert.Equal(t, "tenant-two", string(winners[second.Lookup()].Payload)) +} + +func TestProcessingArtifactDeleteCorruptIsConditional(t *testing.T) { + db, repository := setupProcessingArtifactRepository(t) + ctx := context.Background() + candidate := processingArtifactCandidate(t, 1, "value") + _, _, err := repository.PutIfAbsent(ctx, candidate) + require.NoError(t, err) + + require.NoError(t, repository.DeleteCorrupt(ctx, candidate.Lookup(), artifact.SHA256Hex([]byte("other")))) + var count int64 + require.NoError(t, db.Model(&types.ProcessingArtifact{}).Count(&count).Error) + assert.Equal(t, int64(1), count) + + require.NoError(t, repository.DeleteCorrupt(ctx, candidate.Lookup(), candidate.PayloadChecksum)) + require.NoError(t, db.Model(&types.ProcessingArtifact{}).Count(&count).Error) + assert.Zero(t, count) +} + +func TestProcessingArtifactSuccessfulEmptyPayloadRoundTrips(t *testing.T) { + _, repository := setupProcessingArtifactRepository(t) + candidate := processingArtifactCandidate(t, 1, "") + + winner, created, err := repository.PutIfAbsent(context.Background(), candidate) + require.NoError(t, err) + require.True(t, created) + require.NotNil(t, winner.Payload) + assert.Empty(t, winner.Payload) + + reloaded, err := repository.Get(context.Background(), candidate.Lookup()) + require.NoError(t, err) + require.NotNil(t, reloaded.Payload) + assert.Empty(t, reloaded.Payload) +} diff --git a/internal/application/repository/retriever/neo4j/repository.go b/internal/application/repository/retriever/neo4j/repository.go index 0c563ff133..4a759c6a1d 100644 --- a/internal/application/repository/retriever/neo4j/repository.go +++ b/internal/application/repository/retriever/neo4j/repository.go @@ -2,8 +2,11 @@ package neo4j import ( "context" + "crypto/sha256" + "encoding/hex" "fmt" "strings" + "sync" "github.com/Tencent/WeKnora/internal/logger" "github.com/Tencent/WeKnora/internal/types" @@ -13,8 +16,10 @@ import ( // Neo4jRepository is a repository for Neo4j type Neo4jRepository struct { - driver neo4j.Driver - nodePrefix string + driver neo4j.Driver + nodePrefix string + contributionConstraintMu sync.Mutex + contributionConstraintReady bool } // NewNeo4jRepository creates a new Neo4j repository @@ -56,6 +61,312 @@ func (n *Neo4jRepository) AddGraph(ctx context.Context, namespace types.NameSpac return nil } +// ReplaceGraphContribution atomically replaces the graph facts owned by one +// stable chunk. A marker node fences attempts inside Neo4j so an older worker +// cannot publish after a newer attempt has claimed the same contribution. +func (n *Neo4jRepository) ReplaceGraphContribution( + ctx context.Context, + namespace types.NameSpace, + chunkID string, + attempt int, + graph *types.GraphData, +) (bool, error) { + if n.driver == nil { + logger.Warnf(ctx, "NOT SUPPORT RETRIEVE GRAPH") + return false, nil + } + if chunkID == "" || graph == nil { + return false, fmt.Errorf("graph contribution chunk and graph must not be empty") + } + if err := n.ensureContributionConstraint(ctx); err != nil { + return false, err + } + session := n.driver.NewSession(ctx, neo4j.SessionConfig{AccessMode: neo4j.AccessModeWrite}) + defer session.Close(ctx) + + result, err := session.ExecuteWrite(ctx, func(tx neo4j.ManagedTransaction) (interface{}, error) { + applied, err := claimGraphContributionAttempt(ctx, tx, namespace, chunkID, attempt) + if err != nil || !applied { + return applied, err + } + if err := n.deleteGraphContributionsTx(ctx, tx, namespace, []string{chunkID}); err != nil { + return false, err + } + if err := n.importGraphContributionTx(ctx, tx, namespace, chunkID, graph); err != nil { + return false, err + } + return true, nil + }) + if err != nil { + return false, err + } + applied, _ := result.(bool) + return applied, nil +} + +// DeleteGraphContributions removes only facts owned by the supplied chunks. +// Per-chunk attempt markers make stale cleanup a no-op after a newer attempt +// has already replaced that contribution. +func (n *Neo4jRepository) DeleteGraphContributions( + ctx context.Context, + namespace types.NameSpace, + chunkIDs []string, + attempt int, +) error { + if n.driver == nil || len(chunkIDs) == 0 { + return nil + } + if err := n.ensureContributionConstraint(ctx); err != nil { + return err + } + session := n.driver.NewSession(ctx, neo4j.SessionConfig{AccessMode: neo4j.AccessModeWrite}) + defer session.Close(ctx) + + _, err := session.ExecuteWrite(ctx, func(tx neo4j.ManagedTransaction) (interface{}, error) { + for _, chunkID := range chunkIDs { + applied, err := claimGraphContributionAttempt(ctx, tx, namespace, chunkID, attempt) + if err != nil { + return nil, err + } + if !applied { + continue + } + if err := n.deleteGraphContributionsTx( + ctx, + tx, + namespace, + []string{chunkID}, + ); err != nil { + return nil, err + } + if _, err := tx.Run(ctx, ` + MATCH (c:WEKNORA_GRAPH_CONTRIBUTION { + key: $contribution_key + }) + WHERE c.attempt <= $attempt + DELETE c + `, map[string]interface{}{ + "contribution_key": graphContributionKey(namespace, chunkID), + "attempt": attempt, + }); err != nil { + return nil, err + } + } + return nil, nil + }) + return err +} + +func (n *Neo4jRepository) ensureContributionConstraint(ctx context.Context) error { + n.contributionConstraintMu.Lock() + defer n.contributionConstraintMu.Unlock() + if n.contributionConstraintReady { + return nil + } + session := n.driver.NewSession( + ctx, + neo4j.SessionConfig{AccessMode: neo4j.AccessModeWrite}, + ) + defer session.Close(ctx) + result, err := session.Run(ctx, ` + CREATE CONSTRAINT weknora_graph_contribution_key IF NOT EXISTS + FOR (c:WEKNORA_GRAPH_CONTRIBUTION) + REQUIRE c.key IS UNIQUE + `, nil) + if err == nil { + _, err = result.Consume(ctx) + } + if err != nil { + return fmt.Errorf("ensure graph contribution uniqueness: %w", err) + } + n.contributionConstraintReady = true + return nil +} + +func graphContributionKey(namespace types.NameSpace, chunkID string) string { + sum := sha256.Sum256([]byte( + namespace.KnowledgeBase + "\x00" + namespace.Knowledge + "\x00" + chunkID, + )) + return hex.EncodeToString(sum[:]) +} + +func claimGraphContributionAttempt( + ctx context.Context, + tx neo4j.ManagedTransaction, + namespace types.NameSpace, + chunkID string, + attempt int, +) (bool, error) { + result, err := tx.Run(ctx, ` + MERGE (c:WEKNORA_GRAPH_CONTRIBUTION { + key: $contribution_key + }) + ON CREATE SET + c.knowledge_base_id = $knowledge_base_id, + c.knowledge_id = $knowledge_id, + c.chunk_id = $chunk_id, + c.attempt = $attempt + WITH c + WHERE c.attempt <= $attempt + SET c.attempt = $attempt + RETURN count(c) AS applied + `, map[string]interface{}{ + "contribution_key": graphContributionKey(namespace, chunkID), + "knowledge_base_id": namespace.KnowledgeBase, + "knowledge_id": namespace.Knowledge, + "chunk_id": chunkID, + "attempt": attempt, + }) + if err != nil { + return false, err + } + if !result.Next(ctx) { + return false, result.Err() + } + raw, _ := result.Record().Get("applied") + applied, _ := raw.(int64) + return applied > 0, nil +} + +func (n *Neo4jRepository) deleteGraphContributionsTx( + ctx context.Context, + tx neo4j.ManagedTransaction, + namespace types.NameSpace, + chunkIDs []string, +) error { + labelExpr := n.Label(namespace) + const attributeSeparator = "\u001f" + chunkPrefixes := make([]string, len(chunkIDs)) + for index, chunkID := range chunkIDs { + chunkPrefixes[index] = chunkID + attributeSeparator + } + deleteRelationships := ` + MATCH (n:` + labelExpr + ` {kg: $knowledge_id})-[r]-(m:` + labelExpr + ` {kg: $knowledge_id}) + WHERE ANY(chunk_id IN $chunk_ids WHERE chunk_id IN coalesce(r.chunks, [])) + SET r.chunks = [chunk_id IN coalesce(r.chunks, []) WHERE NOT chunk_id IN $chunk_ids] + WITH DISTINCT r + WHERE size(r.chunks) = 0 + DELETE r + ` + if _, err := tx.Run(ctx, deleteRelationships, map[string]interface{}{ + "knowledge_id": namespace.Knowledge, + "chunk_ids": chunkIDs, + }); err != nil { + return fmt.Errorf("delete graph contribution relationships: %w", err) + } + deleteNodes := ` + MATCH (n:` + labelExpr + ` {kg: $knowledge_id}) + WHERE ANY(chunk_id IN $chunk_ids WHERE chunk_id IN coalesce(n.chunks, [])) + SET n.attribute_contributions = + CASE + WHEN n.attribute_contributions IS NULL + THEN reduce( + entries = [], + owner IN coalesce(n.chunks, []) | + entries + [ + attribute IN coalesce(n.attributes, []) | + owner + $attribute_separator + attribute + ] + ) + ELSE n.attribute_contributions + END + SET n.chunks = [chunk_id IN coalesce(n.chunks, []) WHERE NOT chunk_id IN $chunk_ids] + SET n.attribute_contributions = [ + entry IN coalesce(n.attribute_contributions, []) + WHERE NOT ANY(prefix IN $chunk_prefixes WHERE entry STARTS WITH prefix) + ] + SET n.attributes = apoc.coll.toSet([ + entry IN coalesce(n.attribute_contributions, []) | + substring( + entry, + size(split(entry, $attribute_separator)[0]) + 1 + ) + ]) + WITH n + WHERE size(n.chunks) = 0 + DETACH DELETE n + ` + if _, err := tx.Run(ctx, deleteNodes, map[string]interface{}{ + "knowledge_id": namespace.Knowledge, + "chunk_ids": chunkIDs, + "chunk_prefixes": chunkPrefixes, + "attribute_separator": attributeSeparator, + }); err != nil { + return fmt.Errorf("delete graph contribution nodes: %w", err) + } + return nil +} + +func (n *Neo4jRepository) importGraphContributionTx( + ctx context.Context, + tx neo4j.ManagedTransaction, + namespace types.NameSpace, + chunkID string, + graph *types.GraphData, +) error { + const attributeSeparator = "\u001f" + nodeData := make([]map[string]interface{}, 0, len(graph.Node)) + for _, node := range graph.Node { + attributeContributions := make([]string, len(node.Attributes)) + for index, attribute := range node.Attributes { + attributeContributions[index] = chunkID + attributeSeparator + attribute + } + nodeData = append(nodeData, map[string]interface{}{ + "name": node.Name, + "knowledge_id": namespace.Knowledge, + "attribute_contributions": attributeContributions, + "chunks": []string{chunkID}, + "labels": n.Labels(namespace), + }) + } + if _, err := tx.Run(ctx, ` + UNWIND $data AS row + CALL apoc.merge.node(row.labels, {name: row.name, kg: row.knowledge_id}, {}, {}) YIELD node + SET node.chunks = apoc.coll.union(coalesce(node.chunks, []), row.chunks) + SET node.attribute_contributions = apoc.coll.union( + coalesce(node.attribute_contributions, []), + row.attribute_contributions + ) + SET node.attributes = apoc.coll.toSet([ + entry IN node.attribute_contributions | + substring( + entry, + size(split(entry, $attribute_separator)[0]) + 1 + ) + ]) + RETURN distinct 'done' AS result + `, map[string]interface{}{ + "data": nodeData, + "attribute_separator": attributeSeparator, + }); err != nil { + return fmt.Errorf("import graph contribution nodes: %w", err) + } + + relData := make([]map[string]interface{}, 0, len(graph.Relation)) + for _, relation := range graph.Relation { + relData = append(relData, map[string]interface{}{ + "source": relation.Node1, + "target": relation.Node2, + "knowledge_id": namespace.Knowledge, + "type": relation.Type, + "source_labels": n.Labels(namespace), + "target_labels": n.Labels(namespace), + "chunks": []string{chunkID}, + }) + } + if _, err := tx.Run(ctx, ` + UNWIND $data AS row + CALL apoc.merge.node(row.source_labels, {name: row.source, kg: row.knowledge_id}, {}, {}) YIELD node as source + CALL apoc.merge.node(row.target_labels, {name: row.target, kg: row.knowledge_id}, {}, {}) YIELD node as target + CALL apoc.merge.relationship(source, row.type, {}, {chunks: row.chunks}, target) YIELD rel + SET rel.chunks = apoc.coll.union(coalesce(rel.chunks, []), row.chunks) + RETURN distinct 'done' + `, map[string]interface{}{"data": relData}); err != nil { + return fmt.Errorf("import graph contribution relationships: %w", err) + } + return nil +} + // addGraph adds a graph to the Neo4j repository func (n *Neo4jRepository) addGraph(ctx context.Context, namespace types.NameSpace, graph *types.GraphData) error { session := n.driver.NewSession(ctx, neo4j.SessionConfig{AccessMode: neo4j.AccessModeWrite}) diff --git a/internal/application/service/docreader_artifact.go b/internal/application/service/docreader_artifact.go new file mode 100644 index 0000000000..260b58f289 --- /dev/null +++ b/internal/application/service/docreader_artifact.go @@ -0,0 +1,266 @@ +package service + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/url" + "strings" + "unicode/utf8" + + "github.com/Tencent/WeKnora/internal/artifact" + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" +) + +const ( + docReaderArtifactStage = "parse" + docReaderArtifactKeyVersion = uint16(1) + docReaderArtifactSchema = "docreader.read-result.v1" + docReaderArtifactMaxInline = 16 << 20 +) + +var errUncacheableDocReaderResult = errors.New("uncacheable DocReader result") + +type docReaderArtifactRequest struct { + FileContentDigest string `json:"file_content_digest"` + FileName string `json:"file_name"` + FileType string `json:"file_type"` + Title string `json:"title"` + ParserEngine string `json:"parser_engine"` + ParserEngineOverrides map[string]string `json:"parser_engine_overrides"` +} + +type docReaderArtifactImage struct { + Filename string `json:"filename"` + OriginalRef string `json:"original_ref"` + MimeType string `json:"mime_type"` + ImageData []byte `json:"image_data"` + IsOriginal bool `json:"is_original"` +} + +type docReaderArtifactValue struct { + Version uint8 `json:"version"` + MarkdownContent string `json:"markdown_content"` + ImageRefs []docReaderArtifactImage `json:"image_refs"` + Metadata map[string]string `json:"metadata"` + IsAudio bool `json:"is_audio"` + AudioData []byte `json:"audio_data"` +} + +func (s *knowledgeService) callDocReaderWithArtifact( + ctx context.Context, + reader interfaces.DocReader, + request *types.ReadRequest, +) (*types.ReadResult, error) { + if s.artifactRuntime == nil || request == nil || request.URL != "" { + return s.callDocReaderWithTimeout(ctx, reader, request) + } + expected, cacheable := s.docReaderArtifactExpected(ctx, request) + if !cacheable { + return s.callDocReaderWithTimeout(ctx, reader, request) + } + + var computed *types.ReadResult + value, err := s.artifactRuntime.LoadOrCompute(ctx, expected, func(ctx context.Context) ([]byte, error) { + result, callErr := s.callDocReaderWithTimeout(ctx, reader, request) + if callErr != nil { + return nil, callErr + } + if result == nil { + return nil, errors.New("DocReader returned a nil result") + } + computed = result + if result.Error != "" { + return nil, fmt.Errorf("%w: provider reported an error", errUncacheableDocReaderResult) + } + payload, encodeErr := encodeDocReaderArtifact(result) + if encodeErr != nil { + return nil, fmt.Errorf("%w: %v", errUncacheableDocReaderResult, encodeErr) + } + return payload, nil + }) + if errors.Is(err, errUncacheableDocReaderResult) && computed != nil { + return computed, nil + } + if err != nil { + return nil, err + } + return decodeDocReaderArtifact(value.Payload) +} + +func (s *knowledgeService) docReaderArtifactExpected( + ctx context.Context, + request *types.ReadRequest, +) (artifact.Expected, bool) { + overrides, cacheable := sanitizedDocReaderOverrides(request.ParserEngineOverrides) + if !cacheable { + return artifact.Expected{}, false + } + endpoint := "" + transport := "" + if s.config != nil && s.config.DocReader != nil { + endpoint = s.config.DocReader.Addr + transport = s.config.DocReader.Transport + if !safeArtifactEndpoint(endpoint) { + return artifact.Expected{}, false + } + } + inputDigest := artifact.SHA256Hex(request.FileContent) + rendered := docReaderArtifactRequest{ + FileContentDigest: inputDigest, + FileName: request.FileName, + FileType: request.FileType, + Title: request.Title, + ParserEngine: request.ParserEngine, + ParserEngineOverrides: overrides, + } + tenantID, _ := ctx.Value(types.TenantIDContextKey).(uint64) + if tenantID == 0 { + return artifact.Expected{}, false + } + key, err := artifact.BuildKey(tenantID, artifact.KeyMaterial{ + KeyVersion: docReaderArtifactKeyVersion, + Stage: docReaderArtifactStage, + DirectInputs: []artifact.DirectInput{{ + Role: "file_content", + Digest: inputDigest, + }}, + Processor: artifact.ProcessorIdentity{ + ModelName: request.ParserEngine, + Source: "docreader", + Provider: transport, + EndpointIdentity: endpoint, + Revision: "docreader-contract.v1", + Parameters: map[string]any{ + "parser_engine": request.ParserEngine, + "config_overrides": overrides, + }, + }, + RenderedRequest: rendered, + Options: overrides, + CanonicalizerVersion: artifact.CanonicalJSONVersion, + OutputSchemaVersion: docReaderArtifactSchema, + }) + if err != nil { + return artifact.Expected{}, false + } + return artifact.Expected{ + Key: key, + Codec: artifact.CodecJSONV1, + Validate: func(payload []byte) error { + _, err := decodeDocReaderArtifact(payload) + return err + }, + Cacheable: func(payload []byte) bool { + return len(payload) <= docReaderArtifactMaxInline + }, + }, true +} + +func encodeDocReaderArtifact(result *types.ReadResult) ([]byte, error) { + if result == nil { + return nil, errors.New("DocReader result must not be nil") + } + value := docReaderArtifactValue{ + Version: 1, + MarkdownContent: result.MarkdownContent, + ImageRefs: make([]docReaderArtifactImage, 0, len(result.ImageRefs)), + Metadata: result.Metadata, + IsAudio: result.IsAudio, + AudioData: append([]byte(nil), result.AudioData...), + } + if value.Metadata == nil { + value.Metadata = map[string]string{} + } + if !utf8.ValidString(value.MarkdownContent) { + return nil, errors.New("DocReader markdown is not valid UTF-8") + } + for _, image := range result.ImageRefs { + if !utf8.ValidString(image.Filename) || + !utf8.ValidString(image.OriginalRef) || + !utf8.ValidString(image.MimeType) { + return nil, errors.New("DocReader image metadata is not valid UTF-8") + } + value.ImageRefs = append(value.ImageRefs, docReaderArtifactImage{ + Filename: image.Filename, + OriginalRef: image.OriginalRef, + MimeType: image.MimeType, + ImageData: append([]byte(nil), image.ImageData...), + IsOriginal: image.IsOriginal, + }) + } + if err := artifact.ValidateOwnershipFreeJSON(value); err != nil { + return nil, err + } + return artifact.CanonicalJSON(value) +} + +func decodeDocReaderArtifact(payload []byte) (*types.ReadResult, error) { + decoder := json.NewDecoder(bytes.NewReader(payload)) + decoder.DisallowUnknownFields() + var value docReaderArtifactValue + if err := decoder.Decode(&value); err != nil { + return nil, err + } + var trailing any + if err := decoder.Decode(&trailing); err != io.EOF { + return nil, errors.New("DocReader artifact has trailing JSON") + } + if value.Version != 1 || value.ImageRefs == nil || value.Metadata == nil { + return nil, errors.New("DocReader artifact schema is incomplete") + } + result := &types.ReadResult{ + MarkdownContent: value.MarkdownContent, + Metadata: value.Metadata, + IsAudio: value.IsAudio, + AudioData: append([]byte(nil), value.AudioData...), + ImageRefs: make([]types.ImageRef, 0, len(value.ImageRefs)), + } + for _, image := range value.ImageRefs { + result.ImageRefs = append(result.ImageRefs, types.ImageRef{ + Filename: image.Filename, + OriginalRef: image.OriginalRef, + MimeType: image.MimeType, + ImageData: append([]byte(nil), image.ImageData...), + IsOriginal: image.IsOriginal, + }) + } + return result, nil +} + +func sanitizedDocReaderOverrides(overrides map[string]string) (map[string]string, bool) { + result := make(map[string]string) + for key, value := range overrides { + normalized := strings.NewReplacer("_", "", "-", "", ".", "").Replace(strings.ToLower(key)) + if strings.Contains(normalized, "apikey") || + strings.Contains(normalized, "token") || + strings.Contains(normalized, "secret") || + strings.Contains(normalized, "password") || + strings.Contains(normalized, "signature") { + continue + } + if (strings.Contains(normalized, "endpoint") || strings.Contains(normalized, "url")) && + !safeArtifactEndpoint(value) { + return nil, false + } + result[key] = value + } + return result, true +} + +func safeArtifactEndpoint(value string) bool { + if value == "" || !strings.Contains(value, "://") { + return true + } + parsed, err := url.Parse(value) + return err == nil && + parsed.Scheme != "" && + parsed.Host != "" && + parsed.User == nil && + parsed.RawQuery == "" && + parsed.Fragment == "" +} diff --git a/internal/application/service/docreader_artifact_test.go b/internal/application/service/docreader_artifact_test.go new file mode 100644 index 0000000000..646edecd9c --- /dev/null +++ b/internal/application/service/docreader_artifact_test.go @@ -0,0 +1,161 @@ +package service + +import ( + "context" + "sync" + "testing" + + "github.com/Tencent/WeKnora/internal/application/repository" + "github.com/Tencent/WeKnora/internal/artifact" + "github.com/Tencent/WeKnora/internal/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" +) + +type countingDocReader struct { + mu sync.Mutex + calls int + result *types.ReadResult +} + +func (r *countingDocReader) Read(context.Context, *types.ReadRequest) (*types.ReadResult, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.calls++ + copy := *r.result + copy.ImageRefs = append([]types.ImageRef(nil), r.result.ImageRefs...) + copy.Metadata = make(map[string]string, len(r.result.Metadata)) + for key, value := range r.result.Metadata { + copy.Metadata[key] = value + } + return ©, nil +} + +func (r *countingDocReader) callCount() int { + r.mu.Lock() + defer r.mu.Unlock() + return r.calls +} + +func setupDocReaderArtifactService(t *testing.T) *knowledgeService { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&types.ProcessingArtifact{})) + return &knowledgeService{ + artifactRuntime: artifact.NewRuntime(repository.NewProcessingArtifactRepository(db), nil), + } +} + +func docReaderArtifactContext() context.Context { + return context.WithValue(context.Background(), types.TenantIDContextKey, uint64(1)) +} + +func TestDocReaderArtifactReusesExactSuccessfulResult(t *testing.T) { + service := setupDocReaderArtifactService(t) + reader := &countingDocReader{result: &types.ReadResult{ + MarkdownContent: "# Exact\r\ncontent ", + Metadata: map[string]string{"pages": "1"}, + ImageRefs: []types.ImageRef{{ + Filename: "image.png", + OriginalRef: "images/image.png", + MimeType: "image/png", + ImageData: []byte{1, 2, 3}, + }}, + }} + request := &types.ReadRequest{ + FileContent: []byte("document bytes"), + FileName: "document.pdf", + FileType: "pdf", + Title: "Document", + ParserEngine: "builtin", + RequestID: "trace-one", + } + + first, err := service.callDocReaderWithArtifact(docReaderArtifactContext(), reader, request) + require.NoError(t, err) + request.RequestID = "trace-two" + second, err := service.callDocReaderWithArtifact(docReaderArtifactContext(), reader, request) + require.NoError(t, err) + + assert.Equal(t, 1, reader.callCount()) + assert.Equal(t, first.MarkdownContent, second.MarkdownContent) + require.Len(t, second.ImageRefs, 1) + assert.Equal(t, []byte{1, 2, 3}, second.ImageRefs[0].ImageData) +} + +func TestDocReaderArtifactInvalidatesExactContentAndOptions(t *testing.T) { + service := setupDocReaderArtifactService(t) + reader := &countingDocReader{result: &types.ReadResult{ + MarkdownContent: "content", + Metadata: map[string]string{}, + }} + request := &types.ReadRequest{ + FileContent: []byte("first"), + FileName: "document.pdf", + FileType: "pdf", + ParserEngine: "builtin", + ParserEngineOverrides: map[string]string{ + "pdf_force_scanned": "false", + }, + } + _, err := service.callDocReaderWithArtifact(docReaderArtifactContext(), reader, request) + require.NoError(t, err) + + request.FileContent = []byte("second") + _, err = service.callDocReaderWithArtifact(docReaderArtifactContext(), reader, request) + require.NoError(t, err) + request.ParserEngineOverrides["pdf_force_scanned"] = "true" + _, err = service.callDocReaderWithArtifact(docReaderArtifactContext(), reader, request) + require.NoError(t, err) + assert.Equal(t, 3, reader.callCount()) +} + +func TestDocReaderArtifactExcludesCredentialRotation(t *testing.T) { + service := setupDocReaderArtifactService(t) + reader := &countingDocReader{result: &types.ReadResult{ + MarkdownContent: "content", + Metadata: map[string]string{}, + }} + request := &types.ReadRequest{ + FileContent: []byte("same"), + FileName: "document.pdf", + FileType: "pdf", + ParserEngine: "mineru_cloud", + ParserEngineOverrides: map[string]string{ + "mineru_api_key": "secret-one", + "mineru_model": "pipeline-v1", + }, + } + _, err := service.callDocReaderWithArtifact(docReaderArtifactContext(), reader, request) + require.NoError(t, err) + request.ParserEngineOverrides["mineru_api_key"] = "secret-two" + _, err = service.callDocReaderWithArtifact(docReaderArtifactContext(), reader, request) + require.NoError(t, err) + + assert.Equal(t, 1, reader.callCount()) +} + +func TestDocReaderArtifactDoesNotCacheProviderErrors(t *testing.T) { + service := setupDocReaderArtifactService(t) + reader := &countingDocReader{result: &types.ReadResult{ + Error: "parse failed", + Metadata: map[string]string{}, + }} + request := &types.ReadRequest{ + FileContent: []byte("same"), + FileName: "document.pdf", + FileType: "pdf", + ParserEngine: "builtin", + } + + first, err := service.callDocReaderWithArtifact(docReaderArtifactContext(), reader, request) + require.NoError(t, err) + second, err := service.callDocReaderWithArtifact(docReaderArtifactContext(), reader, request) + require.NoError(t, err) + assert.Equal(t, "parse failed", first.Error) + assert.Equal(t, "parse failed", second.Error) + assert.Equal(t, 2, reader.callCount()) +} diff --git a/internal/application/service/extract.go b/internal/application/service/extract.go index f9a9d59c84..156fe84b2e 100644 --- a/internal/application/service/extract.go +++ b/internal/application/service/extract.go @@ -4,6 +4,7 @@ import ( "context" "database/sql" "encoding/json" + "errors" "fmt" "os" "strings" @@ -12,6 +13,7 @@ import ( "github.com/Tencent/WeKnora/internal/agent/tools" chatpipeline "github.com/Tencent/WeKnora/internal/application/service/chat_pipeline" "github.com/Tencent/WeKnora/internal/application/service/retriever" + "github.com/Tencent/WeKnora/internal/artifact" "github.com/Tencent/WeKnora/internal/config" "github.com/Tencent/WeKnora/internal/logger" "github.com/Tencent/WeKnora/internal/models/chat" @@ -80,6 +82,8 @@ Please output in the following format (one paragraph per column): - Write descriptions in the same language as the data content` ) +var errTableSummarySuperseded = errors.New("table summary attempt superseded") + // NewChunkExtractTask creates a new chunk extract task. It returns // (enqueued, err): enqueued is true only when a task was actually placed on // the queue. When NEO4J is disabled the call is a no-op and returns @@ -132,12 +136,14 @@ func NewDataTableSummaryTask( knowledgeID string, summaryModel string, embeddingModel string, + attempt int, ) error { taskPayload := DataTableSummaryPayload{ TenantID: tenantID, KnowledgeID: knowledgeID, SummaryModel: summaryModel, EmbeddingModel: embeddingModel, + Attempt: attempt, } langfuse.InjectTracing(ctx, &taskPayload) payload, err := json.Marshal(taskPayload) @@ -164,6 +170,7 @@ func enqueueDataTableSummaryIfNeeded( tenantID uint64, knowledgeID string, fileName, fileType, summaryModelID, embeddingModelID string, + attempt int, ) { ft := normalizeFileExtension(fileType) if ft == "" && fileName != "" { @@ -172,7 +179,15 @@ func enqueueDataTableSummaryIfNeeded( if !isDataTableFileType(ft) { return } - if err := NewDataTableSummaryTask(ctx, client, tenantID, knowledgeID, summaryModelID, embeddingModelID); err != nil { + if err := NewDataTableSummaryTask( + ctx, + client, + tenantID, + knowledgeID, + summaryModelID, + embeddingModelID, + attempt, + ); err != nil { logger.Warnf(ctx, "Failed to enqueue data table summary task for knowledge %s: %v", knowledgeID, err) } } @@ -259,6 +274,7 @@ func (s *ChunkExtractService) Handle(ctx context.Context, t *asynq.Task) error { } } var handleErr error + superseded := false graphOut := types.JSONMap{} defer func() { // Decrement the parent's enrichment counter on terminal exit so a @@ -267,7 +283,7 @@ func (s *ChunkExtractService) Handle(ctx context.Context, t *asynq.Task) error { // payload field; legacy in-flight tasks without it are skipped. finalizeSubtaskDetached(ctx, s.knowledgeRepo, p.KnowledgeID, fmt.Sprintf("graph_chunk[%d]", p.ChunkIndex), - handleErr, false, isFinalAsynqAttempt(ctx)) + handleErr, superseded, isFinalAsynqAttempt(ctx)) if gSpan == nil { return } @@ -356,6 +372,11 @@ func (s *ChunkExtractService) Handle(ctx context.Context, t *asynq.Task) error { handleErr = err return err } + if attemptSuperseded(ctx, s.tracker(), p.KnowledgeID, p.Attempt) { + superseded = true + graphOut["skipped"] = "superseded_after_provider" + return nil + } chunk, err = s.chunkRepo.GetChunkByID(ctx, p.TenantID, p.ChunkID) if err != nil { @@ -367,14 +388,44 @@ func (s *ChunkExtractService) Handle(ctx context.Context, t *asynq.Task) error { for _, node := range graph.Node { node.Chunks = []string{chunk.ID} } - if err = s.graphEngine.AddGraph(ctx, - types.NameSpace{KnowledgeBase: chunk.KnowledgeBaseID, Knowledge: chunk.KnowledgeID}, - []*types.GraphData{graph}, - ); err != nil { + namespace := types.NameSpace{ + KnowledgeBase: chunk.KnowledgeBaseID, + Knowledge: chunk.KnowledgeID, + } + if publisher, ok := s.graphEngine.(interfaces.GraphContributionRepository); ok { + var applied bool + applied, err = publisher.ReplaceGraphContribution( + ctx, + namespace, + chunk.ID, + p.Attempt, + graph, + ) + if err == nil && !applied { + superseded = true + graphOut["skipped"] = "newer_graph_contribution" + return nil + } + } else { + // Compatibility fallback for graph backends that have not implemented + // contribution replacement yet. Fence immediately before the append. + if attemptSuperseded(ctx, s.tracker(), p.KnowledgeID, p.Attempt) { + superseded = true + graphOut["skipped"] = "superseded_before_graph_publish" + return nil + } + err = s.graphEngine.AddGraph(ctx, namespace, []*types.GraphData{graph}) + } + if err != nil { logger.Errorf(ctx, "failed to add graph: %v", err) handleErr = err return err } + if attemptSuperseded(ctx, s.tracker(), p.KnowledgeID, p.Attempt) { + superseded = true + graphOut["status"] = "superseded_after_graph_publish" + return nil + } graphOut["nodes_added"] = len(graph.Node) graphOut["relations_added"] = len(graph.Relation) // Capture a couple of sample nodes/relations so the trace viewer can @@ -413,6 +464,7 @@ type DataTableSummaryPayload struct { KnowledgeID string `json:"knowledge_id"` SummaryModel string `json:"summary_model"` EmbeddingModel string `json:"embedding_model"` + Attempt int `json:"attempt,omitempty"` } // DataTableSummaryService is a service for extracting tables @@ -427,6 +479,7 @@ type DataTableSummaryService struct { ownership retriever.TenantStoreOwnership sqlDB *sql.DB storageResolver interfaces.StorageBackendResolver + spanTracker SpanTracker } // NewDataTableSummaryService creates a new DataTableSummaryService @@ -441,6 +494,7 @@ func NewDataTableSummaryService( ownership retriever.TenantStoreOwnership, sqlDB *sql.DB, storageResolver interfaces.StorageBackendResolver, + spanTracker SpanTracker, ) interfaces.TaskHandler { return &DataTableSummaryService{ modelService: modelService, @@ -453,6 +507,7 @@ func NewDataTableSummaryService( ownership: ownership, sqlDB: sqlDB, storageResolver: storageResolver, + spanTracker: spanTracker, } } @@ -471,12 +526,30 @@ func (s *DataTableSummaryService) Handle(ctx context.Context, t *asynq.Task) err ctx = context.WithValue(ctx, types.TenantIDContextKey, payload.TenantID) logger.Infof(ctx, "Processing table extraction for knowledge: %s", payload.KnowledgeID) + if attemptSuperseded(ctx, s.spanTracker, payload.KnowledgeID, payload.Attempt) { + logger.Infof( + ctx, + "Table extraction attempt %d superseded for %s", + payload.Attempt, + payload.KnowledgeID, + ) + return nil + } // 2. 准备所有必需的资源（知识、模型、引擎等） resources, err := s.prepareResources(ctx, payload) if err != nil { return err } + if attemptSuperseded(ctx, s.spanTracker, payload.KnowledgeID, payload.Attempt) { + logger.Infof( + ctx, + "Table extraction attempt %d superseded after provider work for %s", + payload.Attempt, + payload.KnowledgeID, + ) + return nil + } // 3. 加载表格数据并生成摘要 chunks, err := s.processTableData(ctx, resources) @@ -485,8 +558,18 @@ func (s *DataTableSummaryService) Handle(ctx context.Context, t *asynq.Task) err } // 4. 索引到向量数据库 - if err := s.indexToVectorDB(ctx, chunks, resources.retrieveEngine, resources.embeddingModel); err != nil { - s.cleanupOnFailure(ctx, resources, chunks, err) + added, err := s.indexToVectorDB( + ctx, + chunks, + resources.retrieveEngine, + resources.embeddingModel, + payload.Attempt, + ) + if err != nil { + if errors.Is(err, errTableSummarySuperseded) { + return nil + } + s.cleanupOnFailure(ctx, resources, added, err) return err } @@ -677,18 +760,35 @@ func (s *DataTableSummaryService) processTableData(ctx context.Context, resource logger.Debugf(ctx, "column describe of knowledge %s: %s", resources.knowledge.ID, columnDescription) // 构建chunks：一个表格摘要chunk + 多个列描述chunks - chunks := s.buildChunks(resources, tableDescription, columnDescription) + chunks, err := s.buildChunks(resources, tableDescription, columnDescription) + if err != nil { + return nil, err + } return chunks, nil } // buildChunks 构建chunk对象 // tableDescription和columnDescriptions分别生成一个chunk -func (s *DataTableSummaryService) buildChunks(resources *extractionResources, tableDescription string, columnDescription string) []*types.Chunk { +func (s *DataTableSummaryService) buildChunks( + resources *extractionResources, + tableDescription string, + columnDescription string, +) ([]*types.Chunk, error) { chunks := make([]*types.Chunk, 0, 2) + summaryIdentity, err := artifact.StableEntityIdentity(artifact.EntityIdentityInput{ + KnowledgeID: resources.knowledge.ID, + IDVersion: artifact.StableEntityIDVersion, + EntityType: string(types.ChunkTypeTableSummary), + SourceAnchor: "table-summary", + Content: tableDescription, + }) + if err != nil { + return nil, fmt.Errorf("build table summary identity: %w", err) + } // 表格摘要chunk summaryChunk := &types.Chunk{ - ID: uuid.New().String(), + ID: summaryIdentity.ID, TenantID: resources.knowledge.TenantID, KnowledgeID: resources.knowledge.ID, KnowledgeBaseID: resources.knowledge.KnowledgeBaseID, @@ -700,9 +800,20 @@ func (s *DataTableSummaryService) buildChunks(resources *extractionResources, ta } chunks = append(chunks, summaryChunk) + columnIdentity, err := artifact.StableEntityIdentity(artifact.EntityIdentityInput{ + KnowledgeID: resources.knowledge.ID, + IDVersion: artifact.StableEntityIDVersion, + EntityType: string(types.ChunkTypeTableColumn), + ParentSemanticID: summaryIdentity.SemanticKey, + SourceAnchor: "table-columns", + Content: columnDescription, + }) + if err != nil { + return nil, fmt.Errorf("build table column identity: %w", err) + } // 列描述chunk（所有列的描述合并为一个chunk） columnChunk := &types.Chunk{ - ID: uuid.New().String(), + ID: columnIdentity.ID, TenantID: resources.knowledge.TenantID, KnowledgeID: resources.knowledge.ID, KnowledgeBaseID: resources.knowledge.KnowledgeBaseID, @@ -718,7 +829,7 @@ func (s *DataTableSummaryService) buildChunks(resources *extractionResources, ta summaryChunk.NextChunkID = columnChunk.ID columnChunk.PreChunkID = summaryChunk.ID - return chunks + return chunks, nil } // indexToVectorDB 将chunks索引到向量数据库 @@ -728,7 +839,14 @@ func (s *DataTableSummaryService) indexToVectorDB( chunks []*types.Chunk, engine *retriever.CompositeRetrieveEngine, embedder embedding.Embedder, -) error { + attempt int, +) ([]*types.Chunk, error) { + if len(chunks) == 0 { + return nil, nil + } + if attemptSuperseded(ctx, s.spanTracker, chunks[0].KnowledgeID, attempt) { + return nil, errTableSummarySuperseded + } // 构建索引信息列表 indexInfoList := make([]*types.IndexInfo, 0, len(chunks)) for _, chunk := range chunks { @@ -743,17 +861,63 @@ func (s *DataTableSummaryService) indexToVectorDB( }) } - // 保存到数据库 - if err := s.chunkService.CreateChunks(ctx, chunks); err != nil { - logger.Errorf(ctx, "failed to create chunks: %v", err) - return err + existing, err := s.chunkService.ListChunksByKnowledgeID( + ctx, + chunks[0].KnowledgeID, + ) + if err != nil { + return nil, fmt.Errorf("list table summary chunks: %w", err) + } + existingByID := make(map[string]*types.Chunk, len(existing)) + desiredIDs := make(map[string]struct{}, len(chunks)) + for _, chunk := range existing { + existingByID[chunk.ID] = chunk + } + added := make([]*types.Chunk, 0, len(chunks)) + updated := make([]*types.Chunk, 0, len(chunks)) + for _, desired := range chunks { + desiredIDs[desired.ID] = struct{}{} + if live := existingByID[desired.ID]; live != nil { + desired.CreatedAt = live.CreatedAt + updated = append(updated, desired) + } else { + added = append(added, desired) + } + } + stale := make([]*types.Chunk, 0) + for _, live := range existing { + if live.ChunkType != types.ChunkTypeTableSummary && + live.ChunkType != types.ChunkTypeTableColumn { + continue + } + if _, keep := desiredIDs[live.ID]; !keep { + stale = append(stale, live) + } + } + + // Materialize desired rows before indexing; retries update the stable IDs + // instead of inserting duplicates. + if len(added) > 0 { + if err := s.chunkService.CreateChunks(ctx, added); err != nil { + logger.Errorf(ctx, "failed to create chunks: %v", err) + return nil, err + } + } + if len(updated) > 0 { + if err := s.chunkService.UpdateChunks(ctx, updated); err != nil { + logger.Errorf(ctx, "failed to update chunks: %v", err) + return added, err + } } logger.Infof(ctx, "Created %d chunks for data table", len(chunks)) // 批量索引 if err := engine.BatchIndex(ctx, embedder, indexInfoList); err != nil { logger.Errorf(ctx, "failed to index chunks: %v", err) - return err + return added, err + } + if attemptSuperseded(ctx, s.spanTracker, chunks[0].KnowledgeID, attempt) { + return added, errTableSummarySuperseded } // 更新chunk状态为已索引 @@ -762,10 +926,25 @@ func (s *DataTableSummaryService) indexToVectorDB( } if err := s.chunkService.UpdateChunks(ctx, chunks); err != nil { logger.Errorf(ctx, "failed to update chunk status: %v", err) - return err + return added, err + } + + // Desired state is live. Exact stale table-summary contributions are + // cleaned last and never trigger rollback of the new generation. + staleIDs := chunkIDList(stale) + if len(staleIDs) > 0 { + if err := engine.DeleteByChunkIDList( + ctx, + staleIDs, + embedder.GetDimensions(), + types.KnowledgeBaseTypeDocument, + ); err != nil { + logger.Warnf(ctx, "failed to delete stale table vectors: %v", err) + } else if err := s.chunkService.DeleteChunks(ctx, staleIDs); err != nil { + logger.Warnf(ctx, "failed to delete stale table chunks: %v", err) + } } - - return nil + return added, nil } // cleanupOnFailure 索引失败时的清理工作 @@ -820,7 +999,11 @@ func (s *DataTableSummaryService) generateTableDescription(ctx context.Context, // logger.Debugf(ctx, "generateTableDescription prompt: %s", prompt) thinking := false - response, err := chatModel.Chat(ctx, []chat.Message{ + modelCtx := chat.WithArtifactStage(ctx, chat.ArtifactStage{ + Stage: "summary.table", + OutputSchema: "table-summary.text.v1", + }) + response, err := chatModel.Chat(modelCtx, []chat.Message{ {Role: "user", Content: prompt}, }, &chat.ChatOptions{ Temperature: 0.3, @@ -845,7 +1028,11 @@ func (s *DataTableSummaryService) generateColumnDescriptions(ctx context.Context // Call LLM once for all columns thinking := false - response, err := chatModel.Chat(ctx, []chat.Message{ + modelCtx := chat.WithArtifactStage(ctx, chat.ArtifactStage{ + Stage: "summary.table_columns", + OutputSchema: "table-columns.text.v1", + }) + response, err := chatModel.Chat(modelCtx, []chat.Message{ {Role: "user", Content: prompt}, }, &chat.ChatOptions{ Temperature: 0.3, diff --git a/internal/application/service/graph.go b/internal/application/service/graph.go index 767d59351f..b16e5745de 100644 --- a/internal/application/service/graph.go +++ b/internal/application/service/graph.go @@ -120,7 +120,15 @@ func (b *graphBuilder) extractEntities(ctx context.Context, chunk *types.Chunk) // Call LLM to extract entities log.Debug("Calling LLM to extract entities") - resp, err := b.chatModel.Chat(ctx, messages, &chat.ChatOptions{ + entityCtx := chat.WithArtifactStage(ctx, chat.ArtifactStage{ + Stage: "graph_extract.entities", + OutputSchema: "graph.entities.v1", + Validate: func(content string) error { + var value []*types.Entity + return common.ParseLLMJsonResponse(content, &value) + }, + }) + resp, err := b.chatModel.Chat(entityCtx, messages, &chat.ChatOptions{ Temperature: DefaultLLMTemperature, Thinking: &thinking, }) @@ -203,7 +211,26 @@ func (b *graphBuilder) extractRelationships(ctx context.Context, } // Serialize entities to build prompt - entitiesJSON, err := json.Marshal(entities) + canonicalEntities := make([]struct { + Title string `json:"title"` + Type string `json:"type"` + Description string `json:"description"` + }, 0, len(entities)) + for _, entity := range entities { + if entity == nil { + continue + } + canonicalEntities = append(canonicalEntities, struct { + Title string `json:"title"` + Type string `json:"type"` + Description string `json:"description"` + }{ + Title: entity.Title, + Type: entity.Type, + Description: entity.Description, + }) + } + entitiesJSON, err := json.Marshal(canonicalEntities) if err != nil { log.WithError(err).Error("Failed to serialize entities to JSON") return fmt.Errorf("failed to serialize entities: %w", err) @@ -231,7 +258,15 @@ func (b *graphBuilder) extractRelationships(ctx context.Context, // Call LLM to extract relationships log.Debug("Calling LLM to extract relationships") - resp, err := b.chatModel.Chat(ctx, messages, &chat.ChatOptions{ + relationCtx := chat.WithArtifactStage(ctx, chat.ArtifactStage{ + Stage: "graph_extract.relationships", + OutputSchema: "graph.relationships.v1", + Validate: func(content string) error { + var value []*types.Relationship + return common.ParseLLMJsonResponse(content, &value) + }, + }) + resp, err := b.chatModel.Chat(relationCtx, messages, &chat.ChatOptions{ Temperature: DefaultLLMTemperature, Thinking: &thinking, }) diff --git a/internal/application/service/image_multimodal.go b/internal/application/service/image_multimodal.go index 64e471f6bc..d2ad264773 100644 --- a/internal/application/service/image_multimodal.go +++ b/internal/application/service/image_multimodal.go @@ -12,14 +12,15 @@ import ( "github.com/Tencent/WeKnora/internal/application/repository" "github.com/Tencent/WeKnora/internal/application/service/retriever" + "github.com/Tencent/WeKnora/internal/artifact" "github.com/Tencent/WeKnora/internal/logger" + "github.com/Tencent/WeKnora/internal/models/embedding" "github.com/Tencent/WeKnora/internal/models/utils/ollama" "github.com/Tencent/WeKnora/internal/models/vlm" "github.com/Tencent/WeKnora/internal/tracing/langfuse" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" secutils "github.com/Tencent/WeKnora/internal/utils" - "github.com/google/uuid" "github.com/hibiken/asynq" "github.com/redis/go-redis/v9" ) @@ -147,6 +148,15 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e if payload.Language != "" { ctx = context.WithValue(ctx, types.LanguageContextKey, payload.Language) } + if attemptSuperseded(ctx, s.tracker(), payload.KnowledgeID, payload.Attempt) { + logger.Infof( + ctx, + "[ImageMultimodal] Dropping superseded attempt %d for knowledge %s", + payload.Attempt, + payload.KnowledgeID, + ) + return nil + } // Drop orphaned or user-aborted work before touching VLM. Missing // knowledge/KB rows are permanent failures — retrying only burns queue @@ -215,7 +225,14 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e handleErr) } } - if handleErr == nil || isFinalAsynqAttempt(ctx) { + if attemptSuperseded(ctx, tracker, payload.KnowledgeID, payload.Attempt) { + logger.Infof( + ctx, + "[ImageMultimodal] Skip finalize for superseded attempt %d on %s", + payload.Attempt, + payload.KnowledgeID, + ) + } else if handleErr == nil || isFinalAsynqAttempt(ctx) { s.checkAndFinalizeAllImages(ctx, payload) } else { logger.Infof(ctx, @@ -256,6 +273,8 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e URL: payload.ImageURL, OriginalURL: payload.ImageURL, } + ocrDesiredKnown := !payload.EnableOCR + captionDesiredKnown := !payload.EnableCaption if payload.EnableOCR { prompt := vlmOCRPrompt @@ -268,11 +287,16 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e } prompt = types.AppendCustomPromptInstructions(prompt, vlmCfg.CustomInstructions, "image_ocr") - ocrText, ocrErr := vlmModel.Predict(ctx, [][]byte{imgBytes}, prompt) + ocrCtx := vlm.WithArtifactStage(ctx, vlm.ArtifactStage{ + Stage: "vlm_ocr", + OutputSchema: "vlm-ocr.text.v1", + }) + ocrText, ocrErr := vlmModel.Predict(ocrCtx, [][]byte{imgBytes}, prompt) if ocrErr != nil { logger.Warnf(ctx, "[ImageMultimodal] OCR failed for %s: %v", payload.ImageURL, ocrErr) imgOut["ocr_error"] = ocrErr.Error() } else { + ocrDesiredKnown = true ocrText = sanitizeOCRText(ocrText) if ocrText != "" { imageInfo.OCRText = ocrText @@ -286,74 +310,199 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e } } - caption, capErr := vlmModel.Predict(ctx, [][]byte{imgBytes}, buildVLMCaptionPrompt(ctx, vlmCfg)) - if capErr != nil { - logger.Warnf(ctx, "[ImageMultimodal] Caption failed for %s: %v", payload.ImageURL, capErr) - imgOut["caption_error"] = capErr.Error() - } else if caption != "" { - imageInfo.Caption = caption - imgOut["caption_chars"] = len([]rune(caption)) - imgOut["caption_preview"] = previewText(caption, 200) + if payload.EnableCaption { + captionCtx := vlm.WithArtifactStage(ctx, vlm.ArtifactStage{ + Stage: "vlm_caption", + OutputSchema: "vlm-caption.text.v1", + }) + caption, capErr := vlmModel.Predict( + captionCtx, + [][]byte{imgBytes}, + buildVLMCaptionPrompt(ctx, vlmCfg), + ) + if capErr != nil { + logger.Warnf(ctx, "[ImageMultimodal] Caption failed for %s: %v", payload.ImageURL, capErr) + imgOut["caption_error"] = capErr.Error() + } else { + captionDesiredKnown = true + if caption != "" { + imageInfo.Caption = caption + imgOut["caption_chars"] = len([]rune(caption)) + imgOut["caption_preview"] = previewText(caption, 200) + } + } + } + // The provider may finish after a newer reparse has started. Its + // ownership-free artifacts remain reusable, but it must not bind chunks, + // write vectors, publish, or trigger post-processing for the old attempt. + if attemptSuperseded(ctx, tracker, payload.KnowledgeID, payload.Attempt) { + logger.Infof( + ctx, + "[ImageMultimodal] Attempt %d superseded after VLM for knowledge %s", + payload.Attempt, + payload.KnowledgeID, + ) + return nil } // Build child chunks for OCR and caption results imageInfoJSON, _ := json.Marshal([]types.ImageInfo{imageInfo}) + existingChildren, err := s.chunkService.ListChunkByParentID( + ctx, + payload.TenantID, + payload.ChunkID, + ) + if err != nil { + handleErr = fmt.Errorf("list existing multimodal chunks: %w", err) + return handleErr + } var newChunks []*types.Chunk if imageInfo.OCRText != "" { - newChunks = append(newChunks, &types.Chunk{ - ID: uuid.New().String(), - TenantID: payload.TenantID, - KnowledgeID: payload.KnowledgeID, - KnowledgeBaseID: payload.KnowledgeBaseID, - Content: imageInfo.OCRText, - ChunkType: types.ChunkTypeImageOCR, - ParentChunkID: payload.ChunkID, - IsEnabled: true, - Flags: types.ChunkFlagRecommended, - ImageInfo: string(imageInfoJSON), - CreatedAt: time.Now(), - UpdatedAt: time.Now(), - }) + chunk, buildErr := buildStableMultimodalChunk( + payload, + types.ChunkTypeImageOCR, + imageInfo.OCRText, + string(imageInfoJSON), + imgBytes, + existingChildren, + ) + if buildErr != nil { + handleErr = buildErr + return handleErr + } + newChunks = append(newChunks, chunk) } if imageInfo.Caption != "" { - newChunks = append(newChunks, &types.Chunk{ - ID: uuid.New().String(), - TenantID: payload.TenantID, - KnowledgeID: payload.KnowledgeID, - KnowledgeBaseID: payload.KnowledgeBaseID, - Content: imageInfo.Caption, - ChunkType: types.ChunkTypeImageCaption, - ParentChunkID: payload.ChunkID, - IsEnabled: true, - Flags: types.ChunkFlagRecommended, - ImageInfo: string(imageInfoJSON), - CreatedAt: time.Now(), - UpdatedAt: time.Now(), - }) + chunk, buildErr := buildStableMultimodalChunk( + payload, + types.ChunkTypeImageCaption, + imageInfo.Caption, + string(imageInfoJSON), + imgBytes, + existingChildren, + ) + if buildErr != nil { + handleErr = buildErr + return handleErr + } + newChunks = append(newChunks, chunk) } imgOut["chunks_created"] = len(newChunks) if len(newChunks) == 0 { - // Deferred finalize will count this image on success. imgOut["skipped"] = "no_extracted_content" - return nil } - // Persist chunks - if err := s.chunkService.CreateChunks(ctx, newChunks); err != nil { - handleErr = fmt.Errorf("create multimodal chunks: %w", err) - return handleErr + // Reuse a stable row on retry. New rows are written only after the + // post-provider attempt fence above; metadata-only refreshes keep seq_id + // and creation time from the published row. + existingByID := make(map[string]*types.Chunk, len(existingChildren)) + for _, existing := range existingChildren { + existingByID[existing.ID] = existing + } + desiredIDs := make(map[string]struct{}, len(newChunks)) + for _, desired := range newChunks { + desiredIDs[desired.ID] = struct{}{} + } + staleChildren := make([]*types.Chunk, 0) + for _, existing := range existingChildren { + switch existing.ChunkType { + case types.ChunkTypeImageOCR: + if !ocrDesiredKnown { + continue + } + case types.ChunkTypeImageCaption: + if !captionDesiredKnown { + continue + } + default: + continue + } + if _, desired := desiredIDs[existing.ID]; !desired { + staleChildren = append(staleChildren, existing) + } + } + toCreate := make([]*types.Chunk, 0, len(newChunks)) + for _, desired := range newChunks { + existing, found := existingByID[desired.ID] + if !found { + toCreate = append(toCreate, desired) + continue + } + existing.Content = desired.Content + existing.ImageInfo = desired.ImageInfo + existing.IsEnabled = desired.IsEnabled + existing.Flags = desired.Flags + existing.UpdatedAt = desired.UpdatedAt + if err := s.chunkService.UpdateChunk(ctx, existing); err != nil { + handleErr = fmt.Errorf("refresh multimodal chunk %s: %w", existing.ID, err) + return handleErr + } + desired.CreatedAt = existing.CreatedAt + } + if len(toCreate) > 0 { + if err := s.chunkService.CreateChunks(ctx, toCreate); err != nil { + handleErr = fmt.Errorf("create multimodal chunks: %w", err) + return handleErr + } } for _, c := range newChunks { logger.Infof(ctx, "[ImageMultimodal] Created %s chunk %s for image %s, len=%d", c.ChunkType, c.ID, payload.ImageURL, len(c.Content)) } - // Index chunks so they can be retrieved - s.indexChunks(ctx, payload, newChunks) - imgOut["indexed"] = true + // Index desired chunks before exact stale cleanup. + indexResult, indexErr := s.indexChunks(ctx, payload, newChunks) + if indexErr != nil { + handleErr = indexErr + return handleErr + } + if attemptSuperseded(ctx, tracker, payload.KnowledgeID, payload.Attempt) { + logger.Infof( + ctx, + "[ImageMultimodal] Attempt %d superseded after indexing knowledge %s", + payload.Attempt, + payload.KnowledgeID, + ) + return nil + } + imgOut["indexed"] = len(newChunks) > 0 + + staleIDs := chunkIDList(staleChildren) + if len(staleIDs) > 0 { + staleVectorsDeleted := indexResult.engine != nil && indexResult.model != nil + if indexResult.engine != nil && indexResult.model != nil { + if err := indexResult.engine.DeleteByChunkIDList( + ctx, + staleIDs, + indexResult.model.GetDimensions(), + indexResult.kb.Type, + ); err != nil { + staleVectorsDeleted = false + logger.Warnf( + ctx, + "[ImageMultimodal] Failed to delete stale child vectors: %v", + err, + ) + } + } else { + logger.Warnf( + ctx, + "[ImageMultimodal] Retaining stale child rows because the previous vector backend cannot be resolved", + ) + } + if staleVectorsDeleted { + if err := s.chunkService.DeleteChunks(ctx, staleIDs); err != nil { + logger.Warnf( + ctx, + "[ImageMultimodal] Failed to delete stale child chunks: %v", + err, + ) + } + } + } // Enqueue question generation for the caption/OCR content if KB has it enabled. // During initial processChunks, question generation is skipped for image-type @@ -365,6 +514,55 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e return nil } +func buildStableMultimodalChunk( + payload types.ImageMultimodalPayload, + chunkType types.ChunkType, + content string, + imageInfo string, + imageBytes []byte, + existing []*types.Chunk, +) (*types.Chunk, error) { + identity, err := artifact.StableEntityIdentity(artifact.EntityIdentityInput{ + KnowledgeID: payload.KnowledgeID, + IDVersion: artifact.StableEntityIDVersion, + EntityType: string(chunkType), + ParentSemanticID: payload.ChunkID, + SourceAnchor: "image:" + artifact.SHA256Hex(imageBytes), + Content: content, + }) + if err != nil { + return nil, fmt.Errorf("build stable multimodal chunk identity: %w", err) + } + + // Preserve one uniquely matching pre-upgrade random UUID. Ambiguous legacy + // duplicates are never guessed; the deterministic ID wins instead. + matchingLegacyIDs := make([]string, 0, 1) + for _, candidate := range existing { + if candidate.ChunkType == chunkType && candidate.Content == content { + matchingLegacyIDs = append(matchingLegacyIDs, candidate.ID) + } + } + if len(matchingLegacyIDs) == 1 { + identity.ID = matchingLegacyIDs[0] + } + + now := time.Now() + return &types.Chunk{ + ID: identity.ID, + TenantID: payload.TenantID, + KnowledgeID: payload.KnowledgeID, + KnowledgeBaseID: payload.KnowledgeBaseID, + Content: content, + ChunkType: chunkType, + ParentChunkID: payload.ChunkID, + IsEnabled: true, + Flags: types.ChunkFlagRecommended, + ImageInfo: imageInfo, + CreatedAt: now, + UpdatedAt: now, + }, nil +} + // shouldDropOrphanedMultimodal reports whether the task should exit without // retrying. True for user-cancelled/deleting knowledge, or when the parent // knowledge / knowledge-base row no longer exists (deleted while queue entries @@ -422,14 +620,28 @@ func isFinalAsynqAttempt(ctx context.Context) bool { return retried >= maxRetry } -// indexChunks indexes the newly created multimodal chunks into the retrieval engine -// so they can participate in semantic search. -func (s *ImageMultimodalService) indexChunks(ctx context.Context, payload types.ImageMultimodalPayload, chunks []*types.Chunk) { +type multimodalIndexResult struct { + engine *retriever.CompositeRetrieveEngine + model embedding.Embedder + kb *types.KnowledgeBase +} + +// indexChunks indexes the desired multimodal chunks and returns the exact +// backend identity needed for stale cleanup. +func (s *ImageMultimodalService) indexChunks( + ctx context.Context, + payload types.ImageMultimodalPayload, + chunks []*types.Chunk, +) (multimodalIndexResult, error) { kb, err := s.kbService.GetKnowledgeBaseByIDOnly(ctx, payload.KnowledgeBaseID) if err != nil || kb == nil { logger.Warnf(ctx, "[ImageMultimodal] Failed to get KB for indexing: %v", err) - return + if err == nil { + err = errors.New("knowledge base not found") + } + return multimodalIndexResult{}, err } + result := multimodalIndexResult{kb: kb} // Skip vector/keyword indexing when the KB has no embedding-based pipeline enabled // (e.g. Wiki-only KBs). Without this check, GetEmbeddingModel would fail because @@ -450,19 +662,20 @@ func (s *ImageMultimodalService) indexChunks(ctx context.Context, payload types. logger.Warnf(ctx, "[ImageMultimodal] Failed to update chunk %s status to indexed: %v", chunk.ID, uerr) } } - return + return result, nil } embeddingModel, err := s.modelService.GetEmbeddingModel(ctx, kb.EmbeddingModelID) if err != nil { logger.Warnf(ctx, "[ImageMultimodal] Failed to get embedding model for indexing: %v", err) - return + return multimodalIndexResult{}, err } + result.model = embeddingModel tenantInfo, err := s.tenantRepo.GetTenantByID(ctx, payload.TenantID) if err != nil { logger.Warnf(ctx, "[ImageMultimodal] Failed to get tenant for indexing: %v", err) - return + return multimodalIndexResult{}, err } // The factory's unbound path reads TenantInfo from ctx; make sure it's there. ctx = context.WithValue(ctx, types.TenantInfoContextKey, tenantInfo) @@ -473,8 +686,9 @@ func (s *ImageMultimodalService) indexChunks(ctx context.Context, payload types. ctx, s.retrieveEngine, s.ownership, payload.TenantID, kb.VectorStoreID) if err != nil { logger.Warnf(ctx, "[ImageMultimodal] Failed to init retrieve engine: %v", err) - return + return multimodalIndexResult{}, err } + result.engine = engine indexInfoList := make([]*types.IndexInfo, 0, len(chunks)) for _, chunk := range chunks { @@ -488,9 +702,11 @@ func (s *ImageMultimodalService) indexChunks(ctx context.Context, payload types. }) } - if err := engine.BatchIndex(ctx, embeddingModel, indexInfoList); err != nil { - logger.Errorf(ctx, "[ImageMultimodal] Failed to index multimodal chunks: %v", err) - return + if len(indexInfoList) > 0 { + if err := engine.BatchIndex(ctx, embeddingModel, indexInfoList); err != nil { + logger.Errorf(ctx, "[ImageMultimodal] Failed to index multimodal chunks: %v", err) + return multimodalIndexResult{}, err + } } // Mark chunks as indexed. @@ -509,6 +725,7 @@ func (s *ImageMultimodalService) indexChunks(ctx context.Context, payload types. } logger.Infof(ctx, "[ImageMultimodal] Indexed %d multimodal chunks for image %s", len(chunks), payload.ImageURL) + return result, nil } // resolveVLM creates a vlm.VLM instance for the given knowledge base, @@ -647,13 +864,22 @@ func downloadImageFromURL(imageURL string) ([]byte, error) { return secutils.DownloadBytes(imageURL) } +func multimodalPendingKey(knowledgeID string, attempt int) string { + if attempt > 0 { + return fmt.Sprintf("multimodal:pending:%s:%d", knowledgeID, attempt) + } + // Preserve compatibility with tasks already queued before attempts were + // included in the pending counter key. + return fmt.Sprintf("multimodal:pending:%s", knowledgeID) +} + func (s *ImageMultimodalService) checkAndFinalizeAllImages(ctx context.Context, payload types.ImageMultimodalPayload) { if s.redisClient == nil { s.enqueueKnowledgePostProcessTask(ctx, payload) return } - redisKey := fmt.Sprintf("multimodal:pending:%s", payload.KnowledgeID) + redisKey := multimodalPendingKey(payload.KnowledgeID, payload.Attempt) pendingCount, err := s.redisClient.Decr(ctx, redisKey).Result() if err != nil && err != redis.Nil { @@ -688,6 +914,7 @@ func (s *ImageMultimodalService) enqueueKnowledgePostProcessTask(ctx context.Con KnowledgeID: payload.KnowledgeID, KnowledgeBaseID: payload.KnowledgeBaseID, Language: payload.Language, + Attempt: payload.Attempt, } langfuse.InjectTracing(ctx, &taskPayload) payloadBytes, err := json.Marshal(taskPayload) diff --git a/internal/application/service/image_multimodal_prompt_test.go b/internal/application/service/image_multimodal_prompt_test.go index 4eaa1eaff7..5007f666a2 100644 --- a/internal/application/service/image_multimodal_prompt_test.go +++ b/internal/application/service/image_multimodal_prompt_test.go @@ -6,6 +6,9 @@ import ( "testing" "github.com/Tencent/WeKnora/internal/types" + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestBuildVLMCaptionPrompt(t *testing.T) { @@ -27,3 +30,63 @@ func TestBuildVLMCaptionPrompt(t *testing.T) { } }) } + +func TestBuildStableMultimodalChunkUsesExactInputsAndUniqueLegacyMatch(t *testing.T) { + payload := types.ImageMultimodalPayload{ + TenantID: 1, + KnowledgeID: uuid.New().String(), + KnowledgeBaseID: uuid.New().String(), + ChunkID: uuid.New().String(), + } + first, err := buildStableMultimodalChunk( + payload, + types.ChunkTypeImageOCR, + "exact OCR", + "[]", + []byte{0, 1, 2}, + nil, + ) + require.NoError(t, err) + second, err := buildStableMultimodalChunk( + payload, + types.ChunkTypeImageOCR, + "exact OCR", + "[]", + []byte{0, 1, 2}, + nil, + ) + require.NoError(t, err) + assert.Equal(t, first.ID, second.ID) + + changed, err := buildStableMultimodalChunk( + payload, + types.ChunkTypeImageOCR, + "exact OCR", + "[]", + []byte{0, 1, 3}, + nil, + ) + require.NoError(t, err) + assert.NotEqual(t, first.ID, changed.ID) + + legacyID := uuid.New().String() + reused, err := buildStableMultimodalChunk( + payload, + types.ChunkTypeImageOCR, + "exact OCR", + "[]", + []byte{0, 1, 2}, + []*types.Chunk{{ + ID: legacyID, + ChunkType: types.ChunkTypeImageOCR, + Content: "exact OCR", + }}, + ) + require.NoError(t, err) + assert.Equal(t, legacyID, reused.ID) +} + +func TestMultimodalPendingKeyIsAttemptScoped(t *testing.T) { + assert.Equal(t, "multimodal:pending:k:7", multimodalPendingKey("k", 7)) + assert.Equal(t, "multimodal:pending:k", multimodalPendingKey("k", 0)) +} diff --git a/internal/application/service/knowledge.go b/internal/application/service/knowledge.go index e633ba5422..aa3ed4c346 100644 --- a/internal/application/service/knowledge.go +++ b/internal/application/service/knowledge.go @@ -13,6 +13,7 @@ import ( "github.com/Tencent/WeKnora/internal/application/repository" "github.com/Tencent/WeKnora/internal/application/service/retriever" + "github.com/Tencent/WeKnora/internal/artifact" "github.com/Tencent/WeKnora/internal/config" werrors "github.com/Tencent/WeKnora/internal/errors" "github.com/Tencent/WeKnora/internal/infrastructure/docparser" @@ -67,6 +68,7 @@ type knowledgeService struct { kbShareService interfaces.KBShareService imageResolver *docparser.ImageResolver taskPendingRepo interfaces.TaskPendingOpsRepository + artifactRuntime *artifact.Runtime // In-memory fallbacks for Lite mode (no Redis) memFAQProgress sync.Map // taskID -> *types.FAQImportProgress @@ -117,6 +119,7 @@ func NewKnowledgeService( taskPendingRepo interfaces.TaskPendingOpsRepository, spanTracker SpanTracker, audit interfaces.AuditLogService, + artifactRuntime *artifact.Runtime, ) (interfaces.KnowledgeService, error) { return &knowledgeService{ config: config, @@ -146,6 +149,7 @@ func NewKnowledgeService( taskPendingRepo: taskPendingRepo, spanTracker: spanTracker, audit: audit, + artifactRuntime: artifactRuntime, }, nil } @@ -192,7 +196,7 @@ func attemptFromCtx(ctx context.Context) int { // of 0 predates attempt tracking (or tracking is disabled) and is never treated // as superseded. func attemptSuperseded(ctx context.Context, tracker SpanTracker, knowledgeID string, attempt int) bool { - if attempt <= 0 || knowledgeID == "" { + if tracker == nil || attempt <= 0 || knowledgeID == "" { return false } return tracker.LatestAttempt(ctx, knowledgeID) > attempt diff --git a/internal/application/service/knowledge_create.go b/internal/application/service/knowledge_create.go index b7473c01a3..f86dbba8f0 100644 --- a/internal/application/service/knowledge_create.go +++ b/internal/application/service/knowledge_create.go @@ -194,6 +194,21 @@ func (s *knowledgeService) CreateKnowledgeFromFile(ctx context.Context, } lang, _ := types.LanguageFromContext(ctx) + attempt := 0 + if root, allocated, openErr := s.tracker().OpenAttempt( + ctx, + knowledge.ID, + "", + ); openErr == nil && root != nil { + attempt = allocated + } else if openErr != nil { + logger.Warnf( + ctx, + "Failed to allocate processing attempt for %s: %v", + knowledge.ID, + openErr, + ) + } taskPayload := types.DocumentProcessPayload{ TenantID: tenantID, KnowledgeID: knowledge.ID, @@ -205,6 +220,7 @@ func (s *knowledgeService) CreateKnowledgeFromFile(ctx context.Context, EnableQuestionGeneration: enableQuestionGeneration, QuestionCount: questionCount, Language: lang, + Attempt: attempt, } langfuse.InjectTracing(ctx, &taskPayload) @@ -251,7 +267,17 @@ func (s *knowledgeService) CreateKnowledgeFromFile(ctx context.Context, knowledge.ID, ) - enqueueDataTableSummaryIfNeeded(ctx, s.task, tenantID, knowledge.ID, safeFilename, getFileType(safeFilename), kb.SummaryModelID, kb.EmbeddingModelID) + enqueueDataTableSummaryIfNeeded( + ctx, + s.task, + tenantID, + knowledge.ID, + safeFilename, + getFileType(safeFilename), + kb.SummaryModelID, + kb.EmbeddingModelID, + attempt, + ) logger.Infof(ctx, "Knowledge from file created successfully, ID: %s", knowledge.ID) return knowledge, nil @@ -683,7 +709,17 @@ func (s *knowledgeService) createKnowledgeFromFileURL( }) logger.Infof(ctx, "Enqueued file URL process task: id=%s queue=%s knowledge_id=%s", info.ID, info.Queue, knowledge.ID) - enqueueDataTableSummaryIfNeeded(ctx, s.task, tenantID, knowledge.ID, fileName, fileType, kb.SummaryModelID, kb.EmbeddingModelID) + enqueueDataTableSummaryIfNeeded( + ctx, + s.task, + tenantID, + knowledge.ID, + fileName, + fileType, + kb.SummaryModelID, + kb.EmbeddingModelID, + 0, + ) logger.Infof(ctx, "Knowledge from file URL created successfully, ID: %s", knowledge.ID) return knowledge, nil @@ -798,13 +834,13 @@ func (s *knowledgeService) CreateKnowledgeFromManual(ctx context.Context, if status == types.ManualKnowledgeStatusPublish { logger.Infof(ctx, "Manual knowledge created, enqueuing async processing task, ID: %s", knowledge.ID) - taskID, err := s.enqueueManualProcessing(ctx, knowledge, cleanContent, false) + taskID, attempt, err := s.enqueueManualProcessing(ctx, knowledge, cleanContent, false) if err != nil { logger.Errorf(ctx, "Failed to enqueue manual processing task for new knowledge: %v", err) // Non-fatal: mark as failed so user can retry knowledge.ParseStatus = "failed" knowledge.ErrorMessage = "Failed to enqueue processing task" - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, attempt) recordKBActivity(ctx, s.audit, tenantID, kbID, types.AuditActionKnowledgeCreated, "knowledge", knowledge.ID, types.AuditOutcomeFailed, map[string]any{ "title": knowledge.Title, "source_type": "manual", "status": status, @@ -1030,12 +1066,11 @@ func (s *knowledgeService) UpdateManualKnowledge(ctx context.Context, existing.FileType = types.KnowledgeTypeManual existing.Type = types.KnowledgeTypeManual existing.Source = types.KnowledgeTypeManual - existing.EnableStatus = "disabled" existing.UpdatedAt = time.Now() - existing.EmbeddingModelID = kb.EmbeddingModelID if status == types.ManualKnowledgeStatusDraft { existing.ParseStatus = types.ManualKnowledgeStatusDraft + existing.EnableStatus = "disabled" existing.Description = "" existing.ProcessedAt = nil @@ -1050,10 +1085,9 @@ func (s *knowledgeService) UpdateManualKnowledge(ctx context.Context, return existing, nil } - // Publish: persist pending status and enqueue async task for cleanup + re-indexing + // Publish: keep the last live description/enable/model binding while the + // desired generation is built, then switch them in processChunks. existing.ParseStatus = "pending" - existing.Description = "" - existing.ProcessedAt = nil if _, err := ApplyKnowledgeProcessOverrides(ctx, kb, existing, payload.ProcessConfig, nil, nil); err != nil { return nil, err @@ -1065,13 +1099,13 @@ func (s *knowledgeService) UpdateManualKnowledge(ctx context.Context, } logger.Infof(ctx, "Manual knowledge updated, enqueuing async processing task, ID: %s", existing.ID) - taskID, err := s.enqueueManualProcessing(ctx, existing, cleanContent, true) + taskID, attempt, err := s.enqueueManualProcessing(ctx, existing, cleanContent, true) if err != nil { logger.Errorf(ctx, "Failed to enqueue manual processing task: %v", err) // Non-fatal: mark as failed so user can retry existing.ParseStatus = "failed" existing.ErrorMessage = "Failed to enqueue processing task" - s.repo.UpdateKnowledge(ctx, existing) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, existing, attempt) recordKBActivity(ctx, s.audit, tenantID, existing.KnowledgeBaseID, types.AuditActionKnowledgeUpdated, "knowledge", existing.ID, types.AuditOutcomeFailed, map[string]any{ "title": existing.Title, "status": status, @@ -1087,10 +1121,27 @@ func (s *knowledgeService) UpdateManualKnowledge(ctx context.Context, return existing, nil } -// enqueueManualProcessing enqueues a manual:process Asynq task for async cleanup + re-indexing. +// enqueueManualProcessing enqueues desired-state manual processing. func (s *knowledgeService) enqueueManualProcessing(ctx context.Context, knowledge *types.Knowledge, content string, needCleanup bool, -) (string, error) { +) (string, int, error) { + attempt := attemptFromCtx(ctx) + if attempt <= 0 { + if root, allocated, err := s.tracker().OpenAttempt( + ctx, + knowledge.ID, + "", + ); err == nil && root != nil { + attempt = allocated + } else if err != nil { + logger.Warnf( + ctx, + "enqueueManualProcessing: OpenAttempt failed for %s: %v", + knowledge.ID, + err, + ) + } + } requestID, _ := types.RequestIDFromContext(ctx) payload := types.ManualProcessPayload{ RequestId: requestID, @@ -1099,21 +1150,22 @@ func (s *knowledgeService) enqueueManualProcessing(ctx context.Context, KnowledgeBaseID: knowledge.KnowledgeBaseID, Content: content, NeedCleanup: needCleanup, + Attempt: attempt, } langfuse.InjectTracing(ctx, &payload) payloadBytes, err := json.Marshal(payload) if err != nil { - return "", fmt.Errorf("failed to marshal manual process payload: %w", err) + return "", attempt, fmt.Errorf("failed to marshal manual process payload: %w", err) } task := asynq.NewTask(types.TypeManualProcess, payloadBytes, asynq.Queue(types.QueueDefault), asynq.MaxRetry(3), asynq.Timeout(30*time.Minute)) info, err := s.task.Enqueue(task) if err != nil { - return "", fmt.Errorf("failed to enqueue manual process task: %w", err) + return "", attempt, fmt.Errorf("failed to enqueue manual process task: %w", err) } logger.Infof(ctx, "Enqueued manual process task: knowledge_id=%s, asynq_id=%s", knowledge.ID, info.ID) - return info.ID, nil + return info.ID, attempt, nil } // markKnowledgeEnqueueFailed prevents a durable knowledge row from remaining diff --git a/internal/application/service/knowledge_process.go b/internal/application/service/knowledge_process.go index 8fbdbb5dfd..955132c19e 100644 --- a/internal/application/service/knowledge_process.go +++ b/internal/application/service/knowledge_process.go @@ -11,6 +11,7 @@ import ( "time" "github.com/Tencent/WeKnora/internal/application/service/retriever" + "github.com/Tencent/WeKnora/internal/artifact" werrors "github.com/Tencent/WeKnora/internal/errors" "github.com/Tencent/WeKnora/internal/infrastructure/chunker" "github.com/Tencent/WeKnora/internal/infrastructure/docparser" @@ -272,6 +273,10 @@ func (s *knowledgeService) processChunks(ctx context.Context, logger.Infof(ctx, "Knowledge aborted (%s), skipping chunk processing: %s", status, knowledge.ID) return } + if attemptSuperseded(ctx, s.tracker(), knowledge.ID, attemptFromCtx(ctx)) { + logger.Infof(ctx, "Knowledge attempt superseded, skipping chunk processing: %s", knowledge.ID) + return + } // Get embedding model for vectorization — only needed when vector/keyword indexing is enabled var embeddingModel embedding.Embedder @@ -286,37 +291,22 @@ func (s *knowledgeService) processChunks(ctx context.Context, logger.Infof(ctx, "Vector/keyword indexing disabled for KB %s, skipping embedding model", kb.ID) } - // 幂等性处理：清理旧的chunks和索引数据，避免重复数据 - logger.Infof(ctx, "Cleaning up existing chunks and index data for knowledge: %s", knowledge.ID) - - // 删除旧的chunks - if err := s.chunkService.DeleteChunksByKnowledgeID(ctx, knowledge.ID); err != nil { - logger.Warnf(ctx, "Failed to delete existing chunks (may not exist): %v", err) - // 不返回错误，继续处理（可能没有旧数据） + tenantInfo := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + existingChunks, err := s.chunkRepo.ListAllChunksByKnowledgeID(ctx, knowledge.TenantID, knowledge.ID) + if err != nil { + logger.Errorf(ctx, "Failed to load current chunks for reconciliation: %v", err) + return } - // 删除旧的索引数据 — only when vector/keyword indexing is enabled - tenantInfo := ctx.Value(types.TenantInfoContextKey).(*types.Tenant) + // Resolve the index once, but keep the currently published rows and vectors + // intact until desired state has been materialized successfully. retrieveEngine, err := retriever.CreateRetrieveEngineForKB( ctx, s.retrieveEngine, s.ownership, tenantInfo.ID, kb.VectorStoreID) - if err == nil && embeddingModel != nil { - if err := retrieveEngine.DeleteByKnowledgeIDList(ctx, []string{knowledge.ID}, embeddingModel.GetDimensions(), knowledge.Type); err != nil { - logger.Warnf(ctx, "Failed to delete existing index data (may not exist): %v", err) - // 不返回错误，继续处理（可能没有旧数据） - } else { - logger.Infof(ctx, "Successfully deleted existing index data for knowledge: %s", knowledge.ID) - } - } - - // 删除知识图谱数据（如果存在） - namespace := types.NameSpace{KnowledgeBase: knowledge.KnowledgeBaseID, Knowledge: knowledge.ID} - if err := s.graphEngine.DelGraph(ctx, []types.NameSpace{namespace}); err != nil { - logger.Warnf(ctx, "Failed to delete existing graph data (may not exist): %v", err) - // 不返回错误，继续处理 + if err != nil && embeddingModel != nil { + logger.Errorf(ctx, "Failed to resolve retrieve engine for reconciliation: %v", err) + return } - logger.Infof(ctx, "Cleanup completed, starting to process new chunks") - // ========== DocReader 解析结果日志 ========== logger.Infof(ctx, "[DocReader] ========== 解析结果概览 ==========") logger.Infof(ctx, "[DocReader] 知识ID: %s, 知识库ID: %s", knowledge.ID, knowledge.KnowledgeBaseID) @@ -366,119 +356,19 @@ func (s *knowledgeService) processChunks(ctx context.Context, } logger.Infof(ctx, "[DocReader] ========== 解析结果概览结束 ==========") - // Create chunk objects from proto chunks - maxSeq := 0 - - // 统计图片相关的子Chunk数量，用于扩展insertChunks的容量 - imageChunkCount := 0 - for _, chunkData := range chunks { - if len(chunkData.Images) > 0 { - // 为每个图片的OCR和Caption分别创建一个Chunk - imageChunkCount += len(chunkData.Images) * 2 - } - if int(chunkData.Seq) > maxSeq { - maxSeq = int(chunkData.Seq) - } - } - - // === Parent-Child Chunking: create parent chunks first === - hasParentChild := len(options.ParentChunks) > 0 - var parentDBChunks []*types.Chunk // indexed by ParsedParentChunk position - if hasParentChild { - parentDBChunks = make([]*types.Chunk, len(options.ParentChunks)) - for i, pc := range options.ParentChunks { - parentDBChunks[i] = &types.Chunk{ - ID: uuid.New().String(), - TenantID: knowledge.TenantID, - KnowledgeID: knowledge.ID, - KnowledgeBaseID: knowledge.KnowledgeBaseID, - Content: pc.Content, - ChunkIndex: pc.Seq, - IsEnabled: true, - CreatedAt: time.Now(), - UpdatedAt: time.Now(), - StartAt: pc.Start, - EndAt: pc.End, - ChunkType: types.ChunkTypeParentText, - } - } - // Set prev/next links for parent chunks - for i := range parentDBChunks { - if i > 0 { - parentDBChunks[i-1].NextChunkID = parentDBChunks[i].ID - parentDBChunks[i].PreChunkID = parentDBChunks[i-1].ID - } - } - logger.Infof(ctx, "Created %d parent chunks for parent-child strategy", len(parentDBChunks)) - } - - // 重新分配容量，考虑图片相关的Chunk + parent chunks - parentCount := len(options.ParentChunks) - insertChunks := make([]*types.Chunk, 0, len(chunks)+imageChunkCount+parentCount) - // Add parent chunks first (they go into DB but NOT into the vector index) - if hasParentChild { - insertChunks = append(insertChunks, parentDBChunks...) - } - - for idx, chunkData := range chunks { - if strings.TrimSpace(chunkData.Content) == "" { - continue - } - - // 创建主文本Chunk - textChunk := &types.Chunk{ - ID: uuid.New().String(), - TenantID: knowledge.TenantID, - KnowledgeID: knowledge.ID, - KnowledgeBaseID: knowledge.KnowledgeBaseID, - Content: chunkData.Content, - ContextHeader: chunkData.ContextHeader, - ChunkIndex: int(chunkData.Seq), - IsEnabled: true, - CreatedAt: time.Now(), - UpdatedAt: time.Now(), - StartAt: int(chunkData.Start), - EndAt: int(chunkData.End), - ChunkType: types.ChunkTypeText, - } - - // Wire up ParentChunkID for child chunks - if hasParentChild && chunkData.ParentIndex >= 0 && chunkData.ParentIndex < len(parentDBChunks) { - textChunk.ParentChunkID = parentDBChunks[chunkData.ParentIndex].ID - } - - chunks[idx].ChunkID = textChunk.ID - insertChunks = append(insertChunks, textChunk) - } - - // Sort chunks by index for proper ordering - sort.Slice(insertChunks, func(i, j int) bool { - return insertChunks[i].ChunkIndex < insertChunks[j].ChunkIndex - }) - - // 仅为文本类型的Chunk设置前后关系（child chunks only, parents already linked above） - textChunks := make([]*types.Chunk, 0, len(chunks)) - for _, chunk := range insertChunks { - if chunk.ChunkType == types.ChunkTypeText && chunk.ParentChunkID != "" { - // This is a child chunk in parent-child mode - textChunks = append(textChunks, chunk) - } else if chunk.ChunkType == types.ChunkTypeText && !hasParentChild { - // Normal flat chunk (no parent-child mode) - textChunks = append(textChunks, chunk) - } - } - - // 设置文本Chunk之间的前后关系 (skip if parent-child, children don't need prev/next links) - if !hasParentChild { - for i, chunk := range textChunks { - if i > 0 { - textChunks[i-1].NextChunkID = chunk.ID - } - if i < len(textChunks)-1 { - textChunks[i+1].PreChunkID = chunk.ID - } - } + desired, err := buildDesiredDocumentChunks( + knowledge, + chunks, + options.ParentChunks, + existingChunks, + ) + if err != nil { + logger.Errorf(ctx, "Failed to build desired chunk state: %v", err) + return } + insertChunks := desired.All + textChunks := desired.Text + previousStorageSize := knowledge.StorageSize // Check if knowledge is being deleted/cancelled before writing chunks. // Nothing has been persisted yet, so both branches just bail. @@ -486,6 +376,10 @@ func (s *knowledgeService) processChunks(ctx context.Context, logger.Infof(ctx, "Knowledge aborted (%s), skipping chunk write: %s", status, knowledge.ID) return } + if attemptSuperseded(ctx, s.tracker(), knowledge.ID, attemptFromCtx(ctx)) { + logger.Infof(ctx, "Knowledge attempt superseded before desired chunk write: %s", knowledge.ID) + return + } // Save chunks to database — ALWAYS, regardless of indexing strategy. // Chunks are needed for wiki generation, graph extraction, and summary generation @@ -493,21 +387,48 @@ func (s *knowledgeService) processChunks(ctx context.Context, s.beginStage(ctx, knowledge.ID, types.StageChunking, types.JSONMap{ "chunks_planned": len(insertChunks), }) - if err := s.chunkService.CreateChunks(ctx, insertChunks); err != nil { + if len(desired.Added) > 0 { + err = s.chunkRepo.CreateChunks(ctx, desired.Added) + } + if err != nil { knowledge.ParseStatus = types.ParseStatusFailed knowledge.ErrorMessage = err.Error() knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent( + ctx, + knowledge, + attemptFromCtx(ctx), + ) s.failStage(ctx, knowledge.ID, types.StageChunking, werrors.ErrCodeChunkingFailed, "create chunks failed", err) return } + if len(desired.Updated) > 0 { + err = s.chunkRepo.SaveChunks(ctx, desired.Updated) + } + if err != nil { + _ = s.chunkRepo.DeleteChunks(ctx, knowledge.TenantID, chunkIDList(desired.Added)) + knowledge.ParseStatus = types.ParseStatusFailed + knowledge.ErrorMessage = err.Error() + knowledge.UpdatedAt = time.Now() + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent( + ctx, + knowledge, + attemptFromCtx(ctx), + ) + s.failStage(ctx, knowledge.ID, types.StageChunking, + werrors.ErrCodeChunkingFailed, "update desired chunks failed", err) + return + } totalChunkChars := 0 for _, c := range insertChunks { totalChunkChars += len(c.Content) } s.endStage(ctx, knowledge.ID, types.StageChunking, types.JSONMap{ - "chunks_written": len(insertChunks), + "chunks_written": len(desired.Added) + len(desired.Updated), + "chunks_added": len(desired.Added), + "chunks_reused": len(desired.Updated), + "chunks_stale": len(desired.Stale), "total_text_chars": totalChunkChars, }) @@ -554,29 +475,36 @@ func (s *knowledgeService) processChunks(ctx context.Context, knowledge.ParseStatus = types.ParseStatusFailed knowledge.ErrorMessage = err.Error() knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent( + ctx, + knowledge, + attemptFromCtx(ctx), + ) return } // Check if there's enough storage quota available - if tenantInfo.StorageUsed+totalStorageSize > tenantInfo.StorageQuota { + projectedStorage := tenantInfo.StorageUsed - previousStorageSize + totalStorageSize + if projectedStorage > tenantInfo.StorageQuota { knowledge.ParseStatus = types.ParseStatusFailed knowledge.ErrorMessage = "存储空间不足" knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent( + ctx, + knowledge, + attemptFromCtx(ctx), + ) return } } - // Check again before batch indexing (heavy operation). - // deleting → row is going away anyway, drop the chunks we just wrote. - // cancelled → user wants to keep what was already persisted, just stop. + // Check again before the heavy provider/index operation. A superseded + // attempt may leave ownership-free artifacts, but must not publish. if aborted, status := s.isKnowledgeAborted(ctx, knowledge.TenantID, knowledge.ID); aborted { logger.Infof(ctx, "Knowledge aborted (%s) before indexing: %s", status, knowledge.ID) - if status == types.ParseStatusDeleting { - if err := s.chunkService.DeleteChunksByKnowledgeID(ctx, knowledge.ID); err != nil { - logger.Warnf(ctx, "Failed to cleanup chunks after deletion detected: %v", err) - } - } + return + } + if attemptSuperseded(ctx, s.tracker(), knowledge.ID, attemptFromCtx(ctx)) { + logger.Infof(ctx, "Knowledge attempt superseded before indexing: %s", knowledge.ID) return } @@ -585,18 +513,27 @@ func (s *knowledgeService) processChunks(ctx context.Context, knowledge.ParseStatus = types.ParseStatusFailed knowledge.ErrorMessage = err.Error() knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) - - // delete failed chunks - if err := s.chunkService.DeleteChunksByKnowledgeID(ctx, knowledge.ID); err != nil { - logger.Errorf(ctx, "Delete chunks failed: %v", err) - } - - // delete index - if err := retrieveEngine.DeleteByKnowledgeIDList( - ctx, []string{knowledge.ID}, embeddingModel.GetDimensions(), kb.Type, - ); err != nil { - logger.Errorf(ctx, "Delete index failed: %v", err) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent( + ctx, + knowledge, + attemptFromCtx(ctx), + ) + + // Roll back only IDs introduced by this failed attempt. Existing + // published chunks and their index entries remain available. + addedIDs := chunkIDList(desired.Added) + if len(addedIDs) > 0 { + if err := retrieveEngine.DeleteByChunkIDList( + ctx, + addedIDs, + embeddingModel.GetDimensions(), + kb.Type, + ); err != nil { + logger.Warnf(ctx, "Failed to cleanup newly-added vectors: %v", err) + } + if err := s.chunkRepo.DeleteChunks(ctx, knowledge.TenantID, addedIDs); err != nil { + logger.Warnf(ctx, "Failed to cleanup newly-added chunks: %v", err) + } } // Map vector store / embedding rate-limit errors to a // stable code so the UI can offer "retry later" hints. @@ -614,25 +551,21 @@ func (s *knowledgeService) processChunks(ctx context.Context, "storage_bytes": totalStorageSize, }) - // Final check before marking as completed. - // deleting → drop chunks+index we just wrote. - // cancelled → keep persisted data; the row stays in cancelled status - // and downstream stages skip via the entry guards. + // Fence again after provider work and before publishing live state. if aborted, status := s.isKnowledgeAborted(ctx, knowledge.TenantID, knowledge.ID); aborted { logger.Infof(ctx, "Knowledge aborted (%s) after indexing: %s", status, knowledge.ID) - if status == types.ParseStatusDeleting { - if err := s.chunkService.DeleteChunksByKnowledgeID(ctx, knowledge.ID); err != nil { - logger.Warnf(ctx, "Failed to cleanup chunks after deletion detected: %v", err) - } - if err := retrieveEngine.DeleteByKnowledgeIDList(ctx, []string{knowledge.ID}, embeddingModel.GetDimensions(), kb.Type); err != nil { - logger.Warnf(ctx, "Failed to cleanup index after deletion detected: %v", err) - } - } + return + } + if attemptSuperseded(ctx, s.tracker(), knowledge.ID, attemptFromCtx(ctx)) { + logger.Infof(ctx, "Knowledge attempt superseded after indexing: %s", knowledge.ID) return } } else { logger.Infof(ctx, "Vector/keyword indexing disabled for KB %s, skipping BatchIndex", kb.ID) s.skipStage(ctx, knowledge.ID, types.StageEmbedding, "skipped") + // Until backend-specific deletion can use the previous model dimension, + // retain the accounted size rather than falsely subtracting live data. + totalStorageSize = previousStorageSize } // Check if this document has extracted images that will be processed asynchronously @@ -641,6 +574,10 @@ func (s *knowledgeService) processChunks(ctx context.Context, pendingMultimodal := isImage && options.EnableMultimodel && len(options.StoredImages) > 0 pendingPDFMultimodal := !isImage && !isVideo && options.EnableMultimodel && len(options.StoredImages) > 0 + if attemptSuperseded(ctx, s.tracker(), knowledge.ID, attemptFromCtx(ctx)) { + logger.Infof(ctx, "Knowledge attempt superseded before publish: %s", knowledge.ID) + return + } now := time.Now() finalizeIndexedKnowledgeState( knowledge, @@ -649,9 +586,66 @@ func (s *knowledgeService) processChunks(ctx context.Context, pendingMultimodal || pendingPDFMultimodal, now, ) + knowledge.EmbeddingModelID = kb.EmbeddingModelID - if err := s.repo.UpdateKnowledge(ctx, knowledge); err != nil { + published, err := s.repo.UpdateKnowledgeIfAttemptCurrent( + ctx, + knowledge, + attemptFromCtx(ctx), + ) + if err != nil { logger.GetLogger(ctx).WithField("error", err).Errorf("processChunks update knowledge failed") + return + } + if !published { + logger.Infof(ctx, "Knowledge attempt lost conditional publish fence: %s", knowledge.ID) + return + } + + // Cleanup is deliberately last and exact. If a newer attempt starts after + // publish, it owns cleanup; this attempt must leave both generations intact. + if attemptSuperseded(ctx, s.tracker(), knowledge.ID, attemptFromCtx(ctx)) { + logger.Infof(ctx, "Knowledge attempt superseded before stale cleanup: %s", knowledge.ID) + return + } + staleIDs := chunkIDList(desired.Stale) + staleVectorsDeleted := true + if len(staleIDs) > 0 && embeddingModel != nil { + if err := retrieveEngine.DeleteByChunkIDList( + ctx, + staleIDs, + embeddingModel.GetDimensions(), + kb.Type, + ); err != nil { + staleVectorsDeleted = false + logger.Warnf(ctx, "Failed to cleanup stale vectors; retaining stale chunks for retry: %v", err) + } + } + if len(staleIDs) > 0 { + if graphContributions, ok := s.graphEngine.(interfaces.GraphContributionRepository); ok { + err := graphContributions.DeleteGraphContributions( + ctx, + types.NameSpace{ + KnowledgeBase: knowledge.KnowledgeBaseID, + Knowledge: knowledge.ID, + }, + staleIDs, + attemptFromCtx(ctx), + ) + if err != nil { + staleVectorsDeleted = false + logger.Warnf( + ctx, + "Failed to cleanup stale graph contributions; retaining stale chunks for retry: %v", + err, + ) + } + } + } + if len(staleIDs) > 0 && staleVectorsDeleted { + if err := s.chunkRepo.DeleteChunks(ctx, knowledge.TenantID, staleIDs); err != nil { + logger.Warnf(ctx, "Failed to cleanup stale chunks: %v", err) + } } // Enqueue multimodal tasks for images (async, non-blocking) @@ -688,9 +682,11 @@ func (s *knowledgeService) processChunks(ctx context.Context, } } - // Update tenant's storage usage - tenantInfo.StorageUsed += totalStorageSize - if err := s.tenantRepo.AdjustStorageUsed(ctx, tenantInfo.ID, totalStorageSize); err != nil { + // Update tenant storage by the desired-state delta, not by re-adding the + // full size on every reparse. + storageDelta := totalStorageSize - previousStorageSize + tenantInfo.StorageUsed += storageDelta + if err := s.tenantRepo.AdjustStorageUsed(ctx, tenantInfo.ID, storageDelta); err != nil { logger.GetLogger(ctx).WithField("error", err).Errorf("processChunks update tenant storage used failed") } logger.GetLogger(ctx).Infof("processChunks successfully") @@ -873,6 +869,16 @@ func (s *knowledgeService) getSummary(ctx context.Context, }) thinking := false modelCtx := types.WithLLMCallMetadata(ctx, "document_summary", "") + modelCtx = chat.WithArtifactStage(modelCtx, chat.ArtifactStage{ + Stage: "summary", + OutputSchema: "summary.text.v1", + Validate: func(content string) error { + if strings.TrimSpace(content) == "" { + return errors.New("summary output must not be empty") + } + return nil + }, + }) summary, err := summaryModel.Chat(modelCtx, []chat.Message{ { Role: "system", @@ -990,6 +996,7 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn "language": payload.Language, }) var summaryErr error + superseded := false summaryOut := types.JSONMap{} defer func() { // Decrement the parent's enrichment counter on terminal exit. @@ -1002,7 +1009,7 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn // "finalizing". When we DO return an error asynq will retry, so // we only drain on the final attempt. finalizeSubtaskDetached(ctx, s.repo, payload.KnowledgeID, "summary", - retErr, false, isFinalAsynqAttempt(ctx)) + retErr, superseded, isFinalAsynqAttempt(ctx)) if span == nil { return } @@ -1054,15 +1061,27 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn // Update summary status to processing knowledge.SummaryStatus = types.SummaryStatusProcessing knowledge.UpdatedAt = time.Now() - if err := s.repo.UpdateKnowledge(ctx, knowledge); err != nil { + if current, err := s.repo.UpdateKnowledgeIfAttemptCurrent( + ctx, + knowledge, + payload.Attempt, + ); err != nil { logger.Warnf(ctx, "Failed to update summary status to processing: %v", err) + } else if !current { + superseded = true + summaryOut["status"] = "superseded_before_summary" + return nil } // Helper function to mark summary as failed markSummaryFailed := func() { knowledge.SummaryStatus = types.SummaryStatusFailed knowledge.UpdatedAt = time.Now() - if err := s.repo.UpdateKnowledge(ctx, knowledge); err != nil { + if _, err := s.repo.UpdateKnowledgeIfAttemptCurrent( + ctx, + knowledge, + payload.Attempt, + ); err != nil { logger.Warnf(ctx, "Failed to update summary status to failed: %v", err) } } @@ -1078,19 +1097,43 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn // Filter text chunks only textChunks := make([]*types.Chunk, 0) + staleSummaryIDs := make([]string, 0) for _, chunk := range chunks { if chunk.ChunkType == types.ChunkTypeText { textChunks = append(textChunks, chunk) + } else if chunk.ChunkType == types.ChunkTypeSummary { + staleSummaryIDs = append(staleSummaryIDs, chunk.ID) } } summaryOut["text_chunks"] = len(textChunks) if len(textChunks) == 0 { logger.Infof(ctx, "No text chunks found for knowledge: %s", payload.KnowledgeID) - // Mark as completed since there's nothing to summarize + // Successful empty desired state: publish an empty description and then + // remove the previous exact summary contribution. + knowledge.Description = "" knowledge.SummaryStatus = types.SummaryStatusCompleted knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + published, updateErr := s.repo.UpdateKnowledgeIfAttemptCurrent( + ctx, + knowledge, + payload.Attempt, + ) + if updateErr != nil { + summaryErr = updateErr + return updateErr + } + if published { + s.cleanupStaleSummaryContributions( + ctx, + payload, + kb, + knowledge, + staleSummaryIDs, + nil, + nil, + ) + } summaryOut["skipped"] = "no_text_chunks" return nil } @@ -1121,34 +1164,17 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn // (deadline exceeded vs unexpected EOF vs 5xx, etc.). summaryOut["error"] = previewText(err.Error(), 500) summaryOut["error_type"] = fmt.Sprintf("%T", err) - // For the insufficient-content case (scanned PDF without OCR, etc.) - // we deliberately do NOT fall back to the first chunk's raw content, - // since that chunk is typically just a bare markdown image reference - // and surfacing it in the description is misleading. + // A failed provider call must not replace the last good description or + // summary contribution with a fallback. Mark this attempt failed and + // leave the published generation intact; retryable errors are returned + // to Asynq. + markSummaryFailed() + summaryErr = err if errors.Is(err, errInsufficientSummaryContent) { - knowledge.Description = "" - knowledge.SummaryStatus = types.SummaryStatusFailed - knowledge.UpdatedAt = time.Now() - if updateErr := s.repo.UpdateKnowledge(ctx, knowledge); updateErr != nil { - logger.Errorf(ctx, "Failed to mark summary as failed: %v", updateErr) - summaryErr = updateErr - return fmt.Errorf("failed to update knowledge: %w", updateErr) - } summaryOut["fallback"] = "insufficient_content" - summaryErr = err return nil } - // For other errors (LLM API issues etc.), fall back to the first chunk. - if len(textChunks) > 0 { - summary = textChunks[0].Content - if len(summary) > 500 { - runes := []rune(summary) - if len(runes) > 500 { - summary = string(runes[:500]) - } - } - summaryOut["fallback"] = "first_chunk" - } + return fmt.Errorf("failed to generate summary: %w", err) } // Do not publish an answer derived from a superseded chunk or metadata // version. A user can explicitly refresh again from the latest revision. @@ -1167,21 +1193,23 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn return nil } - // Update knowledge description - knowledge.Description = summary - knowledge.SummaryStatus = types.SummaryStatusCompleted - knowledge.UpdatedAt = time.Now() + // Provider work may finish after a newer reparse starts. The summary + // artifact is ownership-free and remains reusable, but this attempt must + // not bind it into live knowledge state. + if attemptSuperseded(ctx, s.tracker(), payload.KnowledgeID, payload.Attempt) { + superseded = true + summaryOut["status"] = "superseded_after_provider" + return nil + } summaryOut["summary_chars"] = len([]rune(summary)) // Preview the generated summary on the span output so the trace // viewer can show "this is what the LLM produced" at a glance, // without hopping to the knowledge-detail page. Capped to keep // span rows compact. summaryOut["summary_preview"] = previewText(summary, 240) - if err := s.repo.UpdateKnowledge(ctx, knowledge); err != nil { - logger.Errorf(ctx, "Failed to update knowledge description: %v", err) - summaryErr = err - return fmt.Errorf("failed to update knowledge: %w", err) - } + + var summaryRetrieveEngine *retriever.CompositeRetrieveEngine + var summaryEmbeddingModel embedding.Embedder // Create summary chunk and index it — only when RAG indexing is enabled. // Wiki-only KBs don't need summary chunks in the vector index. @@ -1199,12 +1227,51 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn // unreliable signal (e.g. "MX5280.pdf" for a scanned legal letter) // and surfacing them in retrieved RAG context can re-introduce the // hallucination vector this branch is meant to close. + summaryContent := fmt.Sprintf("# Summary\n%s", summary) + identity, err := artifact.StableEntityIdentity(artifact.EntityIdentityInput{ + KnowledgeID: knowledge.ID, + IDVersion: artifact.StableEntityIDVersion, + EntityType: string(types.ChunkTypeSummary), + ParentSemanticID: textChunks[0].ID, + SourceAnchor: "summary", + Content: summaryContent, + }) + if err != nil { + summaryErr = err + return fmt.Errorf("build summary chunk identity: %w", err) + } + allChunks, err := s.chunkRepo.ListAllChunksByKnowledgeID( + ctx, + knowledge.TenantID, + knowledge.ID, + ) + if err != nil { + summaryErr = err + return fmt.Errorf("list summary chunks: %w", err) + } + existingByID := make(map[string]*types.Chunk, len(allChunks)) + var exactLegacy []*types.Chunk + for _, existing := range allChunks { + existingByID[existing.ID] = existing + if existing.ChunkType == types.ChunkTypeSummary && existing.Content == summaryContent { + exactLegacy = append(exactLegacy, existing) + } + } + if len(exactLegacy) == 1 { + identity.ID = exactLegacy[0].ID + } + staleSummaryIDs = staleSummaryIDs[:0] + for _, existing := range allChunks { + if existing.ChunkType == types.ChunkTypeSummary && existing.ID != identity.ID { + staleSummaryIDs = append(staleSummaryIDs, existing.ID) + } + } summaryChunk := &types.Chunk{ - ID: uuid.New().String(), + ID: identity.ID, TenantID: knowledge.TenantID, KnowledgeID: knowledge.ID, KnowledgeBaseID: knowledge.KnowledgeBaseID, - Content: fmt.Sprintf("# Summary\n%s", summary), + Content: summaryContent, ChunkIndex: maxChunkIndex + 1, IsEnabled: true, CreatedAt: time.Now(), @@ -1215,8 +1282,19 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn ParentChunkID: textChunks[0].ID, } - // Save summary chunk - if err := s.chunkService.CreateChunks(ctx, []*types.Chunk{summaryChunk}); err != nil { + // Save desired summary state without duplicating the row on a retry. + if existing, found := existingByID[summaryChunk.ID]; found { + existing.Content = summaryChunk.Content + existing.ParentChunkID = summaryChunk.ParentChunkID + existing.IsEnabled = true + existing.UpdatedAt = summaryChunk.UpdatedAt + summaryChunk.CreatedAt = existing.CreatedAt + if err := s.chunkService.UpdateChunk(ctx, existing); err != nil { + logger.Errorf(ctx, "Failed to refresh summary chunk: %v", err) + summaryErr = err + return fmt.Errorf("failed to refresh summary chunk: %w", err) + } + } else if err := s.chunkService.CreateChunks(ctx, []*types.Chunk{summaryChunk}); err != nil { logger.Errorf(ctx, "Failed to create summary chunk: %v", err) summaryErr = err return fmt.Errorf("failed to create summary chunk: %w", err) @@ -1231,7 +1309,7 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn } ctx = context.WithValue(ctx, types.TenantInfoContextKey, tenantInfo) - retrieveEngine, err := retriever.CreateRetrieveEngineForKB( + summaryRetrieveEngine, err = retriever.CreateRetrieveEngineForKB( ctx, s.retrieveEngine, s.ownership, tenantInfo.ID, kb.VectorStoreID) if err != nil { logger.Errorf(ctx, "Failed to init retrieve engine: %v", err) @@ -1239,7 +1317,7 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn return fmt.Errorf("failed to init retrieve engine: %w", err) } - embeddingModel, err := s.modelService.GetEmbeddingModel(ctx, kb.EmbeddingModelID) + summaryEmbeddingModel, err = s.modelService.GetEmbeddingModel(ctx, kb.EmbeddingModelID) if err != nil { logger.Errorf(ctx, "Failed to get embedding model: %v", err) summaryErr = err @@ -1256,21 +1334,124 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn IsEnabled: true, }} - if err := retrieveEngine.BatchIndex(ctx, embeddingModel, indexInfo); err != nil { + if err := summaryRetrieveEngine.BatchIndex(ctx, summaryEmbeddingModel, indexInfo); err != nil { logger.Errorf(ctx, "Failed to index summary chunk: %v", err) summaryErr = err return fmt.Errorf("failed to index summary chunk: %w", err) } + if attemptSuperseded(ctx, s.tracker(), payload.KnowledgeID, payload.Attempt) { + superseded = true + summaryOut["status"] = "superseded_after_index" + return nil + } logger.Infof(ctx, "Successfully created and indexed summary chunk for knowledge: %s", payload.KnowledgeID) summaryOut["summary_chunk_indexed"] = true } + if attemptSuperseded(ctx, s.tracker(), payload.KnowledgeID, payload.Attempt) { + superseded = true + summaryOut["status"] = "superseded_before_publish" + return nil + } + knowledge.Description = summary + knowledge.SummaryStatus = types.SummaryStatusCompleted + knowledge.UpdatedAt = time.Now() + published, err := s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, payload.Attempt) + if err != nil { + logger.Errorf(ctx, "Failed to update knowledge description: %v", err) + summaryErr = err + return fmt.Errorf("failed to update knowledge: %w", err) + } + if !published { + superseded = true + summaryOut["status"] = "superseded_at_publish" + return nil + } + s.cleanupStaleSummaryContributions( + ctx, + payload, + kb, + knowledge, + staleSummaryIDs, + summaryRetrieveEngine, + summaryEmbeddingModel, + ) + logger.Infof(ctx, "Successfully generated summary for knowledge: %s", payload.KnowledgeID) summaryOut["status"] = "completed" return nil } +func (s *knowledgeService) cleanupStaleSummaryContributions( + ctx context.Context, + payload types.SummaryGenerationPayload, + kb *types.KnowledgeBase, + knowledge *types.Knowledge, + staleIDs []string, + engine *retriever.CompositeRetrieveEngine, + model embedding.Embedder, +) { + if len(staleIDs) == 0 || + kb == nil || + knowledge == nil || + attemptSuperseded(ctx, s.tracker(), payload.KnowledgeID, payload.Attempt) { + return + } + + if model == nil { + modelID := knowledge.EmbeddingModelID + if modelID == "" { + modelID = kb.EmbeddingModelID + } + if modelID != "" { + var err error + model, err = s.modelService.GetEmbeddingModel(ctx, modelID) + if err != nil { + logger.Warnf(ctx, "Retaining stale summary chunks; embedding model lookup failed: %v", err) + return + } + } + } + if model != nil { + if engine == nil { + tenantInfo, err := s.tenantRepo.GetTenantByID(ctx, payload.TenantID) + if err != nil { + logger.Warnf(ctx, "Retaining stale summary chunks; tenant lookup failed: %v", err) + return + } + ctx = context.WithValue(ctx, types.TenantInfoContextKey, tenantInfo) + engine, err = retriever.CreateRetrieveEngineForKB( + ctx, + s.retrieveEngine, + s.ownership, + tenantInfo.ID, + kb.VectorStoreID, + ) + if err != nil { + logger.Warnf(ctx, "Retaining stale summary chunks; vector backend lookup failed: %v", err) + return + } + } + if err := engine.DeleteByChunkIDList( + ctx, + staleIDs, + model.GetDimensions(), + kb.Type, + ); err != nil { + logger.Warnf(ctx, "Failed to delete stale summary vectors: %v", err) + return + } + } + if err := s.chunkRepo.DeleteChunks( + ctx, + knowledge.TenantID, + staleIDs, + ); err != nil { + logger.Warnf(ctx, "Failed to delete stale summary chunks: %v", err) + } +} + // ProcessQuestionGeneration handles async question generation task. It // dispatches between the batched fan-out path (current: one task per window of // text chunks, payload.ChunkIDs set) and the legacy whole-knowledge path (kept @@ -1561,6 +1742,8 @@ func (s *knowledgeService) processQuestionGenerationForKnowledge(ctx context.Con // Generate questions for each chunk with context var indexInfoList []*types.IndexInfo + var questionChunksToPublish []*types.Chunk + staleQuestionSourceIDsByChunk := make(map[string][]string) for i, chunk := range textChunks { if strings.TrimSpace(chunk.Content) == "" { emptyContentChunks++ @@ -1588,7 +1771,11 @@ func (s *knowledgeService) processQuestionGenerationForKnowledge(ctx context.Con if len(questions) == 0 { llmCallEmpty++ - continue + } + if attemptSuperseded(ctx, s.tracker(), payload.KnowledgeID, payload.Attempt) { + superseded = true + exitStatus = "superseded_after_provider" + return nil } latestChunk, latestErr := s.chunkRepo.GetChunkByID(ctx, payload.TenantID, chunk.ID) if latestErr != nil || latestChunk.ContentRevision != generationRevision { @@ -1602,17 +1789,18 @@ func (s *knowledgeService) processQuestionGenerationForKnowledge(ctx context.Con sampleQuestion = previewText(questions[0], 200) } - // Update chunk metadata with unique IDs for each question - generatedQuestions := make([]types.GeneratedQuestion, len(questions)) - questionRevision := chunk.ContentRevision - for j, question := range questions { - questionID := fmt.Sprintf("q%d", time.Now().UnixNano()+int64(j)) - generatedQuestions[j] = types.GeneratedQuestion{ - ID: questionID, - Question: question, - ContentRevision: &questionRevision, - } + // Question IDs are semantic UUIDv5 values, stable across retries and + // unrelated chunk reordering. A unique exact legacy match is retained. + generatedQuestions, err := stableGeneratedQuestions(knowledge.ID, chunk, questions) + if err != nil { + chunkMetadataSetFailed++ + logger.Warnf(ctx, "Failed to build stable question IDs for chunk %s: %v", chunk.ID, err) + continue } + staleQuestionSourceIDsByChunk[chunk.ID] = staleGeneratedQuestionSourceIDs( + chunk, + generatedQuestions, + ) meta := &types.DocumentChunkMetadata{ GeneratedQuestions: generatedQuestions, GeneratedQuestionsRevision: chunk.ContentRevision, } @@ -1621,13 +1809,7 @@ func (s *knowledgeService) processQuestionGenerationForKnowledge(ctx context.Con logger.Warnf(ctx, "Failed to set document metadata for chunk %s: %v", chunk.ID, err) continue } - - // Update chunk in database - if err := s.chunkService.UpdateChunk(ctx, chunk); err != nil { - chunkUpdateFailed++ - logger.Warnf(ctx, "Failed to update chunk %s: %v", chunk.ID, err) - continue - } + questionChunksToPublish = append(questionChunksToPublish, chunk) // Create index entries for generated questions for _, gq := range generatedQuestions { @@ -1654,9 +1836,51 @@ func (s *knowledgeService) processQuestionGenerationForKnowledge(ctx context.Con logger.Errorf(ctx, "Failed to index generated questions: %v", err) return fmt.Errorf("failed to index questions: %w", err) } + if attemptSuperseded(ctx, s.tracker(), payload.KnowledgeID, payload.Attempt) { + superseded = true + exitStatus = "superseded_after_index" + return nil + } indexBatchSucceeded = true logger.Infof(ctx, "Successfully indexed %d generated questions for knowledge: %s", len(indexInfoList), payload.KnowledgeID) } + if attemptSuperseded(ctx, s.tracker(), payload.KnowledgeID, payload.Attempt) { + superseded = true + exitStatus = "superseded_before_question_publish" + return nil + } + var publishedStaleQuestionSourceIDs []string + for _, chunk := range questionChunksToPublish { + if err := s.chunkService.UpdateChunk(ctx, chunk); err != nil { + chunkUpdateFailed++ + logger.Warnf(ctx, "Failed to publish question metadata for chunk %s: %v", chunk.ID, err) + continue + } + publishedStaleQuestionSourceIDs = append( + publishedStaleQuestionSourceIDs, + staleQuestionSourceIDsByChunk[chunk.ID]..., + ) + } + if chunkUpdateFailed > 0 { + exitStatus = "publish_question_metadata_failed" + return fmt.Errorf("failed to publish question metadata for %d chunks", chunkUpdateFailed) + } + if attemptSuperseded(ctx, s.tracker(), payload.KnowledgeID, payload.Attempt) { + superseded = true + exitStatus = "superseded_before_question_cleanup" + return nil + } + if len(publishedStaleQuestionSourceIDs) > 0 { + if err := retrieveEngine.DeleteBySourceIDList( + ctx, + publishedStaleQuestionSourceIDs, + embeddingModel.GetDimensions(), + kb.Type, + ); err != nil { + exitStatus = "cleanup_stale_questions_failed" + return fmt.Errorf("delete stale question vectors: %w", err) + } + } return nil } @@ -1910,6 +2134,8 @@ func (s *knowledgeService) processQuestionGenerationForChunks(ctx context.Contex } var indexInfoList []*types.IndexInfo + var questionChunksToPublish []*types.Chunk + staleQuestionSourceIDsByChunk := make(map[string][]string) for i, chunk := range batchChunks { if chunk == nil || strings.TrimSpace(chunk.Content) == "" { emptyChunks++ @@ -1925,8 +2151,10 @@ func (s *knowledgeService) processQuestionGenerationForChunks(ctx context.Contex logger.Warnf(ctx, "Failed to generate questions for chunk %s: %v", chunk.ID, gerr) continue } - if len(questions) == 0 { - continue + if attemptSuperseded(ctx, s.tracker(), payload.KnowledgeID, payload.Attempt) { + superseded = true + exitStatus = "superseded_after_provider" + return nil } latestChunk, latestErr := s.chunkRepo.GetChunkByID(ctx, payload.TenantID, chunk.ID) if latestErr != nil || latestChunk.ContentRevision != generationRevision { @@ -1936,30 +2164,28 @@ func (s *knowledgeService) processQuestionGenerationForChunks(ctx context.Contex chunk = latestChunk chunksProcessed++ generatedQuestionsTotal += len(questions) - if sampleQuestion == "" { + if sampleQuestion == "" && len(questions) > 0 { sampleQuestion = previewText(questions[0], 200) } - generatedQuestions := make([]types.GeneratedQuestion, len(questions)) - questionRevision := chunk.ContentRevision - for j, question := range questions { - generatedQuestions[j] = types.GeneratedQuestion{ - ID: fmt.Sprintf("q%d", time.Now().UnixNano()+int64(j)), - Question: question, - ContentRevision: &questionRevision, - } + generatedQuestions, err := stableGeneratedQuestions(knowledge.ID, chunk, questions) + if err != nil { + logger.Warnf(ctx, "Failed to build stable question IDs for chunk %s: %v", chunk.ID, err) + continue } + staleQuestionSourceIDsByChunk[chunk.ID] = staleGeneratedQuestionSourceIDs( + chunk, + generatedQuestions, + ) meta := &types.DocumentChunkMetadata{ - GeneratedQuestions: generatedQuestions, GeneratedQuestionsRevision: chunk.ContentRevision, + GeneratedQuestions: generatedQuestions, + GeneratedQuestionsRevision: chunk.ContentRevision, } if err := chunk.SetDocumentMetadata(meta); err != nil { logger.Warnf(ctx, "Failed to set document metadata for chunk %s: %v", chunk.ID, err) continue } - if err := s.chunkService.UpdateChunk(ctx, chunk); err != nil { - logger.Warnf(ctx, "Failed to update chunk %s: %v", chunk.ID, err) - continue - } + questionChunksToPublish = append(questionChunksToPublish, chunk) for _, gq := range generatedQuestions { indexInfoList = append(indexInfoList, &types.IndexInfo{ Content: buildKnowledgeIndexContent(knowledge, gq.Question), @@ -1981,13 +2207,126 @@ func (s *knowledgeService) processQuestionGenerationForChunks(ctx context.Contex logger.Errorf(ctx, "Failed to index generated questions for batch %d: %v", payload.BatchIndex, err) return fmt.Errorf("failed to index questions: %w", err) } + if attemptSuperseded(ctx, s.tracker(), payload.KnowledgeID, payload.Attempt) { + superseded = true + exitStatus = "superseded_after_index" + return nil + } indexBatchSucceeded = true logger.Infof(ctx, "Indexed %d generated questions for knowledge=%s batch=%d", len(indexInfoList), payload.KnowledgeID, payload.BatchIndex) } + if attemptSuperseded(ctx, s.tracker(), payload.KnowledgeID, payload.Attempt) { + superseded = true + exitStatus = "superseded_before_question_publish" + return nil + } + var publishedStaleQuestionSourceIDs []string + for _, chunk := range questionChunksToPublish { + if err := s.chunkService.UpdateChunk(ctx, chunk); err != nil { + logger.Warnf(ctx, "Failed to publish question metadata for chunk %s: %v", chunk.ID, err) + exitStatus = "publish_question_metadata_failed" + qErr = err + return fmt.Errorf("publish question metadata for chunk %s: %w", chunk.ID, err) + } + publishedStaleQuestionSourceIDs = append( + publishedStaleQuestionSourceIDs, + staleQuestionSourceIDsByChunk[chunk.ID]..., + ) + } + if attemptSuperseded(ctx, s.tracker(), payload.KnowledgeID, payload.Attempt) { + superseded = true + exitStatus = "superseded_before_question_cleanup" + return nil + } + if len(publishedStaleQuestionSourceIDs) > 0 { + if err := retrieveEngine.DeleteBySourceIDList( + ctx, + publishedStaleQuestionSourceIDs, + embeddingModel.GetDimensions(), + kb.Type, + ); err != nil { + exitStatus = "cleanup_stale_questions_failed" + qErr = err + return fmt.Errorf("delete stale question vectors: %w", err) + } + } return nil } +func stableGeneratedQuestions( + knowledgeID string, + parent *types.Chunk, + questions []string, +) ([]types.GeneratedQuestion, error) { + if parent == nil { + return nil, errors.New("question parent chunk must not be nil") + } + allocator := artifact.NewIdentityAllocator(knowledgeID, artifact.StableEntityIDVersion) + result := make([]types.GeneratedQuestion, len(questions)) + + existingByContent := make(map[string][]string) + if metadata, err := parent.DocumentMetadata(); err == nil && metadata != nil { + for _, question := range metadata.GeneratedQuestions { + existingByContent[question.Question] = append( + existingByContent[question.Question], + question.ID, + ) + } + } + desiredCount := make(map[string]int, len(questions)) + for _, question := range questions { + desiredCount[question]++ + } + + for index, question := range questions { + identity, err := allocator.Next( + "generated_question", + parent.ID, + "generated-question", + question, + ) + if err != nil { + return nil, err + } + id := identity.ID + if existing := existingByContent[question]; len(existing) == 1 && desiredCount[question] == 1 { + id = existing[0] + } + questionRevision := parent.ContentRevision + result[index] = types.GeneratedQuestion{ + ID: id, + Question: question, + ContentRevision: &questionRevision, + } + } + return result, nil +} + +func staleGeneratedQuestionSourceIDs( + parent *types.Chunk, + desired []types.GeneratedQuestion, +) []string { + if parent == nil { + return nil + } + metadata, err := parent.DocumentMetadata() + if err != nil || metadata == nil { + return nil + } + desiredIDs := make(map[string]struct{}, len(desired)) + for _, question := range desired { + desiredIDs[question.ID] = struct{}{} + } + result := make([]string, 0) + for _, question := range metadata.GeneratedQuestions { + if _, keep := desiredIDs[question.ID]; !keep { + result = append(result, fmt.Sprintf("%s-%s", parent.ID, question.ID)) + } + } + return result +} + // generateQuestionsWithContext generates questions for a chunk with surrounding context func (s *knowledgeService) generateQuestionsWithContext(ctx context.Context, chatModel chat.Chat, content, prevContent, nextContent, docName string, questionCount int, @@ -2027,6 +2366,10 @@ func (s *knowledgeService) generateQuestionsWithContext(ctx context.Context, thinking := false modelCtx := types.WithLLMCallMetadata(ctx, "question_generation", "") + modelCtx = chat.WithArtifactStage(modelCtx, chat.ArtifactStage{ + Stage: "question", + OutputSchema: "questions.text.v1", + }) response, err := chatModel.Chat(modelCtx, []chat.Message{ { Role: "user", @@ -2126,12 +2469,9 @@ func (s *knowledgeService) RegenerateChunkQuestions( return nil, ErrChunkRevisionConflict } chunk = latestChunk - generated := make([]types.GeneratedQuestion, 0, len(questions)) - questionRevision := chunk.ContentRevision - for _, question := range questions { - generated = append(generated, types.GeneratedQuestion{ - ID: uuid.NewString(), Question: question, ContentRevision: &questionRevision, - }) + generated, err := stableGeneratedQuestions(knowledge.ID, chunk, questions) + if err != nil { + return nil, err } meta := &types.DocumentChunkMetadata{ GeneratedQuestions: generated, GeneratedQuestionsRevision: chunk.ContentRevision, @@ -2231,9 +2571,21 @@ func (s *knowledgeService) RegenerateKnowledgeSummary( } } if !found { + summaryContent := "# Summary\n" + summary + identity, identityErr := artifact.StableEntityIdentity(artifact.EntityIdentityInput{ + KnowledgeID: knowledge.ID, + IDVersion: artifact.StableEntityIDVersion, + EntityType: string(types.ChunkTypeSummary), + ParentSemanticID: textChunks[0].ID, + SourceAnchor: "summary", + Content: summaryContent, + }) + if identityErr != nil { + return nil, fmt.Errorf("build summary chunk identity: %w", identityErr) + } summaryChunk := &types.Chunk{ - ID: uuid.NewString(), TenantID: tenantID, KnowledgeID: knowledge.ID, - KnowledgeBaseID: knowledge.KnowledgeBaseID, Content: "# Summary\n" + summary, + ID: identity.ID, TenantID: tenantID, KnowledgeID: knowledge.ID, + KnowledgeBaseID: knowledge.KnowledgeBaseID, Content: summaryContent, ChunkIndex: maxIndex + 1, IsEnabled: true, ChunkType: types.ChunkTypeSummary, ParentChunkID: textChunks[0].ID, CreatedAt: time.Now(), UpdatedAt: time.Now(), } @@ -2249,8 +2601,8 @@ func (s *knowledgeService) RegenerateKnowledgeSummary( return knowledge, nil } -// ReparseKnowledge deletes existing document content and re-parses the knowledge asynchronously. -// This method reuses the logic from UpdateManualKnowledge for resource cleanup and async parsing. +// ReparseKnowledge schedules a desired-state rebuild. Existing document +// resources remain published until the worker has materialized replacements. func (s *knowledgeService) ReparseKnowledge( ctx context.Context, knowledgeID string, @@ -2297,10 +2649,23 @@ func (s *knowledgeService) ReparseKnowledge( logger.Errorf(ctx, "Failed to set process overrides on reparse: %v", err) return nil, err } - if err := s.repo.UpdateKnowledgeColumn(ctx, existing.ID, "metadata", existing.Metadata); err != nil { + published, err := s.repo.UpdateKnowledgeColumnsIfAttemptCurrent( + ctx, + existing.TenantID, + existing.ID, + reparseAttempt, + map[string]interface{}{ + "metadata": existing.Metadata, + "updated_at": time.Now(), + }, + ) + if err != nil { logger.Errorf(ctx, "Failed to persist process overrides on reparse: %v", err) return nil, err } + if !published { + return s.repo.GetKnowledgeByID(ctx, tenantID, knowledgeID) + } } processOverrides, _ = existing.ProcessOverrides() @@ -2321,7 +2686,7 @@ func (s *knowledgeService) ReparseKnowledge( map[string]any{"title": existing.Title, "type": existing.Type, "attempt": reparseAttempt}) } - // For manual knowledge, use async manual processing (cleanup + re-indexing in worker) + // Manual knowledge uses the same desired-state worker as document imports. if existing.IsManual() { meta, metaErr := existing.ManualMetadata() if metaErr != nil || meta == nil { @@ -2329,20 +2694,50 @@ func (s *knowledgeService) ReparseKnowledge( return nil, werrors.NewBadRequestError("无法获取手工知识内容") } - resetKnowledgeForReparse(existing, kb) + existing.ParseStatus = types.ParseStatusPending + existing.ErrorMessage = "" + // Reset the enrichment counter so a leftover value from a + // previous attempt (e.g. cancelled before all subtasks decremented) + // cannot block the new finalizing transition later. This must be + // an explicit column write: UpdateKnowledge (full-row Save) omits + // pending_subtasks_count, so the struct assignment alone would not + // persist. + existing.PendingSubtasksCount = 0 - if err := s.repo.UpdateKnowledge(ctx, existing); err != nil { + published, err := s.repo.UpdateKnowledgeColumnsIfAttemptCurrent( + ctx, + existing.TenantID, + existing.ID, + reparseAttempt, + map[string]interface{}{ + "parse_status": types.ParseStatusPending, + "error_message": "", + "pending_subtasks_count": 0, + "updated_at": time.Now(), + }, + ) + if err != nil { logger.Errorf(ctx, "Failed to update knowledge status before reparse: %v", err) return nil, err } - if err := s.repo.UpdateKnowledgeColumn(ctx, existing.ID, "pending_subtasks_count", 0); err != nil { - logger.Errorf(ctx, "Failed to reset pending_subtasks_count before reparse: %v", err) - return nil, err + if !published { + return s.repo.GetKnowledgeByID(ctx, tenantID, knowledgeID) } - if _, err := s.enqueueManualProcessing(ctx, existing, meta.Content, true); err != nil { + if _, _, err := s.enqueueManualProcessing( + withAttempt(ctx, reparseAttempt), + existing, + meta.Content, + true, + ); err != nil { logger.Errorf(ctx, "Failed to enqueue manual reparse task: %v", err) - s.markKnowledgeEnqueueFailed(ctx, existing) + existing.ParseStatus = "failed" + existing.ErrorMessage = "Failed to enqueue processing task" + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent( + ctx, + existing, + reparseAttempt, + ) return existing, werrors.NewInternalServerError("Failed to submit processing task") } else { recordReparseStarted() @@ -2350,25 +2745,35 @@ func (s *knowledgeService) ReparseKnowledge( return existing, nil } - // For non-manual knowledge, cleanup synchronously then enqueue document processing - logger.Infof(ctx, "Cleaning up existing resources for knowledge: %s", knowledgeID) - if err := s.cleanupKnowledgeResources(ctx, existing); err != nil { - logger.ErrorWithFields(ctx, err, map[string]interface{}{ - "knowledge_id": knowledgeID, - }) - return nil, err - } - - // Step 2: Update knowledge status and metadata - resetKnowledgeForReparse(existing, kb) + // For non-manual knowledge, retain the last published state while the new + // attempt computes desired chunks and artifacts. + existing.ParseStatus = types.ParseStatusPending + existing.ErrorMessage = "" + // Reset the enrichment counter so a leftover value from a previous + // attempt cannot block the new finalizing transition later. This must + // be an explicit column write: UpdateKnowledge (full-row Save) omits + // pending_subtasks_count, so the struct assignment alone would not + // persist. + existing.PendingSubtasksCount = 0 - if err := s.repo.UpdateKnowledge(ctx, existing); err != nil { + published, err := s.repo.UpdateKnowledgeColumnsIfAttemptCurrent( + ctx, + existing.TenantID, + existing.ID, + reparseAttempt, + map[string]interface{}{ + "parse_status": types.ParseStatusPending, + "error_message": "", + "pending_subtasks_count": 0, + "updated_at": time.Now(), + }, + ) + if err != nil { logger.Errorf(ctx, "Failed to update knowledge status before reparse: %v", err) return nil, err } - if err := s.repo.UpdateKnowledgeColumn(ctx, existing.ID, "pending_subtasks_count", 0); err != nil { - logger.Errorf(ctx, "Failed to reset pending_subtasks_count before reparse: %v", err) - return nil, err + if !published { + return s.repo.GetKnowledgeByID(ctx, tenantID, knowledgeID) } // Step 3: Trigger async re-parsing based on knowledge type @@ -2423,7 +2828,17 @@ func (s *knowledgeService) ReparseKnowledge( recordReparseStarted() // For data tables (csv, xlsx, xls), also enqueue summary task - enqueueDataTableSummaryIfNeeded(ctx, s.task, tenantID, existing.ID, existing.FileName, existing.FileType, kb.SummaryModelID, kb.EmbeddingModelID) + enqueueDataTableSummaryIfNeeded( + ctx, + s.task, + tenantID, + existing.ID, + existing.FileName, + existing.FileType, + kb.SummaryModelID, + kb.EmbeddingModelID, + reparseAttempt, + ) return existing, nil } @@ -2476,7 +2891,17 @@ func (s *knowledgeService) ReparseKnowledge( logger.Infof(ctx, "Enqueued file URL reparse task: id=%s queue=%s knowledge_id=%s", info.ID, info.Queue, existing.ID) recordReparseStarted() - enqueueDataTableSummaryIfNeeded(ctx, s.task, tenantID, existing.ID, existing.FileName, existing.FileType, kb.SummaryModelID, kb.EmbeddingModelID) + enqueueDataTableSummaryIfNeeded( + ctx, + s.task, + tenantID, + existing.ID, + existing.FileName, + existing.FileType, + kb.SummaryModelID, + kb.EmbeddingModelID, + reparseAttempt, + ) return existing, nil } @@ -2835,8 +3260,17 @@ func (s *knowledgeService) UpdateImageInfo( // Create a new caption chunk if it doesn't exist and we have caption data if !hasCaptionChunk && image.Caption != "" { + id, err := stableImageDerivedChunkID( + chunk, + types.ChunkTypeImageCaption, + image.OriginalURL, + image.Caption, + ) + if err != nil { + return err + } captionChunk := &types.Chunk{ - ID: uuid.New().String(), + ID: id, TenantID: tenantID, KnowledgeID: chunk.KnowledgeID, KnowledgeBaseID: chunk.KnowledgeBaseID, @@ -2851,8 +3285,17 @@ func (s *knowledgeService) UpdateImageInfo( // Create a new OCR chunk if it doesn't exist and we have OCR data if !hasOCRChunk && image.OCRText != "" { + id, err := stableImageDerivedChunkID( + chunk, + types.ChunkTypeImageOCR, + image.OriginalURL, + image.OCRText, + ) + if err != nil { + return err + } ocrChunk := &types.Chunk{ - ID: uuid.New().String(), + ID: id, TenantID: tenantID, KnowledgeID: chunk.KnowledgeID, KnowledgeBaseID: chunk.KnowledgeBaseID, @@ -2915,6 +3358,31 @@ func (s *knowledgeService) UpdateImageInfo( return nil } +func stableImageDerivedChunkID( + parent *types.Chunk, + chunkType string, + originalURL string, + content string, +) (string, error) { + if parent == nil { + return "", errors.New("image-derived parent chunk must not be nil") + } + identity, err := artifact.StableEntityIdentity(artifact.EntityIdentityInput{ + KnowledgeID: parent.KnowledgeID, + IDVersion: artifact.StableEntityIDVersion, + EntityType: chunkType, + ParentSemanticID: parent.ID, + SourceAnchor: "image-ref:" + artifact.SHA256Hex( + []byte(originalURL), + ), + Content: content, + }) + if err != nil { + return "", fmt.Errorf("build image-derived chunk identity: %w", err) + } + return identity.ID, nil +} + // ProcessManualUpdate handles Asynq manual knowledge update tasks. // It performs cleanup of old indexes/chunks (when NeedCleanup is true) and re-indexes the content. func (s *knowledgeService) ProcessManualUpdate(ctx context.Context, t *asynq.Task) error { @@ -2975,38 +3443,44 @@ func (s *knowledgeService) ProcessManualUpdate(ctx context.Context, t *asynq.Tas logger.Infof(ctx, "ProcessManualUpdate: knowledge aborted (%s), skipping: %s", status, knowledge.ID) return nil } - // Update status to processing - knowledge.ParseStatus = "processing" - knowledge.UpdatedAt = time.Now() - if err := s.repo.UpdateKnowledge(ctx, knowledge); err != nil { - logger.Errorf(ctx, "ProcessManualUpdate: failed to update status to processing: %v", err) - return nil - } - // Allocate a fresh span-tracking attempt for this manual (re)index. // Without it attemptFromCtx stays 0, so processChunks drops all stage // spans and KnowledgePostProcess falls back to LatestAttempt — piling // this run's summary/wiki subspans onto the previous attempt's trace. - attempt := 0 - if root, n, err := s.tracker().OpenAttempt(ctx, knowledge.ID, payload.LangfuseTraceID); err == nil && root != nil { - attempt = n - } else if err != nil { - logger.Warnf(ctx, "ProcessManualUpdate: OpenAttempt failed for %s: %v", knowledge.ID, err) + attempt := payload.Attempt + if attempt <= 0 { + if root, n, err := s.tracker().OpenAttempt(ctx, knowledge.ID, payload.LangfuseTraceID); err == nil && root != nil { + attempt = n + } else if err != nil { + logger.Warnf(ctx, "ProcessManualUpdate: OpenAttempt failed for %s: %v", knowledge.ID, err) + } } ctx = withAttempt(ctx, attempt) - // Cleanup old resources (indexes, chunks, graph) for update operations + // Mark processing only after the attempt exists. The conditional update + // prevents a delayed manual task from overwriting a newer attempt. + knowledge.ParseStatus = "processing" + knowledge.UpdatedAt = time.Now() + published, err := s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, attempt) + if err != nil { + logger.Errorf(ctx, "ProcessManualUpdate: failed to update status to processing: %v", err) + return nil + } + if !published { + logger.Infof(ctx, "ProcessManualUpdate: superseded before processing: %s", knowledge.ID) + return nil + } + + // NeedCleanup historically triggered a delete-first pass. Desired-state + // processChunks now writes/indexes the new generation and removes only its + // exact stale IDs after conditional publish, so pre-cleanup would recreate + // the empty-index failure window this path is meant to avoid. if payload.NeedCleanup { - if err := s.cleanupKnowledgeResources(ctx, knowledge); err != nil { - logger.ErrorWithFields(ctx, err, map[string]interface{}{ - "knowledge_id": payload.KnowledgeID, - }) - knowledge.ParseStatus = "failed" - knowledge.ErrorMessage = fmt.Sprintf("failed to cleanup old resources: %v", err) - knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) - return nil - } + logger.Infof( + ctx, + "ProcessManualUpdate: deferring exact stale cleanup until desired publish: %s", + knowledge.ID, + ) } // Run manual processing (image resolution + chunking + embedding) synchronously within the worker @@ -3090,6 +3564,22 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e logger.Warnf(ctx, "Unexpected parse status: %s for knowledge: %s", knowledge.ParseStatus, payload.KnowledgeID) } + // Resolve the attempt before any mutable state. The enqueue site sets + // payload.Attempt to a fresh number for the initial parse and each + // user-initiated reparse; Asynq retries retain it. Legacy payloads allocate + // one here. + attempt := payload.Attempt + if attempt <= 0 { + if root, n, openErr := s.tracker().OpenAttempt( + ctx, + knowledge.ID, + payload.LangfuseTraceID, + ); openErr == nil && root != nil { + attempt = n + } + } + ctx = withAttempt(ctx, attempt) + // 获取知识库信息 kb, err := s.kbService.GetKnowledgeBaseByID(ctx, payload.KnowledgeBaseID) if err != nil { @@ -3097,7 +3587,7 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e knowledge.ParseStatus = "failed" knowledge.ErrorMessage = fmt.Sprintf("failed to get knowledge base: %v", err) knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, attempt) return nil } @@ -3112,26 +3602,18 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e logger.Infof(ctx, "Knowledge aborted (%s) before marking processing: %s", status, knowledge.ID) return nil } + knowledge.ParseStatus = "processing" knowledge.UpdatedAt = time.Now() - if err := s.repo.UpdateKnowledge(ctx, knowledge); err != nil { + published, err := s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, attempt) + if err != nil { logger.Errorf(ctx, "failed to update knowledge status to processing: %v", err) return nil } - - // Resolve the attempt for span tracking. The enqueue site sets - // payload.Attempt to a fresh number for the initial parse and to - // max+1 for each user-initiated reparse. Asynq retries within a - // single user action keep the same payload (so retries record - // onto the same attempt). For payloads predating this code we - // fall back to OpenAttempt. - attempt := payload.Attempt - if attempt <= 0 { - if root, n, err := s.tracker().OpenAttempt(ctx, knowledge.ID, payload.LangfuseTraceID); err == nil && root != nil { - attempt = n - } + if !published { + logger.Infof(ctx, "Document attempt superseded before processing: %s", knowledge.ID) + return nil } - ctx = withAttempt(ctx, attempt) // 检查多模态配置（仅对文件导入） if payload.FilePath != "" && !payload.EnableMultimodel && IsImageType(payload.FileType) { @@ -3140,7 +3622,7 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e knowledge.ParseStatus = "failed" knowledge.ErrorMessage = ErrImageNotParse.Error() knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, attempt) return nil } @@ -3151,7 +3633,7 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e knowledge.ParseStatus = "failed" knowledge.ErrorMessage = "上传音频文件需要设置ASR语音识别模型" knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, attempt) return nil } @@ -3162,7 +3644,7 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e knowledge.ParseStatus = "failed" knowledge.ErrorMessage = "暂不支持视频文件" knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, attempt) return nil } @@ -3177,7 +3659,7 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e knowledge.ParseStatus = "failed" knowledge.ErrorMessage = "File URL is not allowed for security reasons" knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, attempt) return nil } @@ -3190,7 +3672,7 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e knowledge.ParseStatus = "failed" knowledge.ErrorMessage = err.Error() knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, attempt) } return fmt.Errorf("failed to download file from URL: %w", err) } @@ -3200,7 +3682,7 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e knowledge.ParseStatus = "failed" knowledge.ErrorMessage = fmt.Sprintf("unsupported file type: %s", resolvedFileType) knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, attempt) return nil } @@ -3209,7 +3691,7 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e } if resolvedFileType != "" && knowledge.FileType == "" { knowledge.FileType = resolvedFileType - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, attempt) } fileSvc := s.resolveFileService(ctx, kb) @@ -3219,7 +3701,7 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e knowledge.ParseStatus = "failed" knowledge.ErrorMessage = err.Error() knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, attempt) } return fmt.Errorf("failed to save downloaded file: %w", err) } @@ -3248,9 +3730,14 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e if extractedTitle := convertResult.Metadata["title"]; extractedTitle != "" { knowledge.Title = extractedTitle knowledge.UpdatedAt = time.Now() - if err := s.repo.UpdateKnowledge(ctx, knowledge); err != nil { + titlePublished, err := s.repo.UpdateKnowledgeIfAttemptCurrent( + ctx, + knowledge, + attempt, + ) + if err != nil { logger.Warnf(ctx, "Failed to update knowledge title from extracted page title: %v", err) - } else { + } else if titlePublished { logger.Infof(ctx, "Updated knowledge title to extracted page title: %s", extractedTitle) } } @@ -3296,7 +3783,7 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e knowledge.ParseStatus = "failed" knowledge.ErrorMessage = "ASR model is not configured for audio transcription" knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, attempt) return nil } @@ -3309,7 +3796,7 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e knowledge.ParseStatus = "failed" knowledge.ErrorMessage = fmt.Sprintf("failed to get ASR model: %v", err) knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, attempt) return nil } @@ -3320,7 +3807,7 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e knowledge.ParseStatus = "failed" knowledge.ErrorMessage = fmt.Sprintf("audio transcription failed: %v", err) knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, attempt) } return fmt.Errorf("audio transcription failed: %w", err) } @@ -3467,7 +3954,11 @@ func (s *knowledgeService) convert( knowledge.ParseStatus = "failed" knowledge.ErrorMessage = "URL is not allowed for security reasons" knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent( + ctx, + knowledge, + attemptFromCtx(ctx), + ) s.failStage(ctx, knowledge.ID, types.StageDocReader, werrors.ErrCodeDocReaderParseFailed, "URL rejected for security reasons", err) return nil, nil @@ -3489,7 +3980,11 @@ func (s *knowledgeService) convert( knowledge.ParseStatus = "failed" knowledge.ErrorMessage = "Document parsing service is not configured. Please use text/paragraph import or set DOCREADER_ADDR." knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent( + ctx, + knowledge, + attemptFromCtx(ctx), + ) s.failStage(ctx, knowledge.ID, types.StageDocReader, werrors.ErrCodeDocReaderUnavailable, knowledge.ErrorMessage, nil) return nil, nil @@ -3522,7 +4017,7 @@ func (s *knowledgeService) convert( req.FileType = fileType } - result, err := s.callDocReaderWithTimeout(ctx, reader, req) + result, err := s.callDocReaderWithArtifact(ctx, reader, req) if err != nil { // Distinguish DocReader timeout (a knowable user-facing // failure) from generic read errors so the UI can suggest @@ -3541,7 +4036,11 @@ func (s *knowledgeService) convert( knowledge.ParseStatus = "failed" knowledge.ErrorMessage = result.Error knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent( + ctx, + knowledge, + attemptFromCtx(ctx), + ) s.failStage(ctx, knowledge.ID, types.StageDocReader, werrors.ErrCodeDocReaderParseFailed, result.Error, nil) return nil, nil @@ -3662,7 +4161,11 @@ func (s *knowledgeService) failKnowledge( knowledge.ParseStatus = "failed" knowledge.ErrorMessage = errMsg knowledge.UpdatedAt = time.Now() - s.repo.UpdateKnowledge(ctx, knowledge) + _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent( + ctx, + knowledge, + attemptFromCtx(ctx), + ) } return nil, fmt.Errorf(format, args...) } @@ -3681,7 +4184,7 @@ func (s *knowledgeService) enqueueImageMultimodalTasks( } attempt := attemptFromCtx(ctx) - redisKey := fmt.Sprintf("multimodal:pending:%s", knowledge.ID) + redisKey := multimodalPendingKey(knowledge.ID, attempt) if s.redisClient != nil { if err := s.redisClient.Set(ctx, redisKey, len(images), 24*time.Hour).Err(); err != nil { logger.Warnf(ctx, "Failed to set multimodal pending count for %s: %v", knowledge.ID, err) diff --git a/internal/application/service/knowledge_reconcile.go b/internal/application/service/knowledge_reconcile.go new file mode 100644 index 0000000000..9b352855d1 --- /dev/null +++ b/internal/application/service/knowledge_reconcile.go @@ -0,0 +1,252 @@ +package service + +import ( + "fmt" + "strings" + "time" + + "github.com/Tencent/WeKnora/internal/artifact" + "github.com/Tencent/WeKnora/internal/types" +) + +type desiredChunkSet struct { + All []*types.Chunk + Text []*types.Chunk + Added []*types.Chunk + Updated []*types.Chunk + Stale []*types.Chunk +} + +type desiredParentIdentity struct { + chunk *types.Chunk + semanticKey string + contentHash string +} + +// buildDesiredDocumentChunks assigns order-independent UUIDv5 identities. The +// current parser does not expose structural block IDs, so the safest available +// source anchor is the exact persisted content digest; unrelated insertion or +// reordering therefore leaves IDs stable without risking a false content reuse. +func buildDesiredDocumentChunks( + knowledge *types.Knowledge, + parsed []types.ParsedChunk, + parsedParents []types.ParsedParentChunk, + existing []*types.Chunk, +) (desiredChunkSet, error) { + if knowledge == nil { + return desiredChunkSet{}, fmt.Errorf("knowledge must not be nil") + } + now := time.Now() + existingByID := make(map[string]*types.Chunk, len(existing)) + existingParentContent := make(map[string]string) + legacyCandidates := make(map[string][]*types.Chunk) + for _, chunk := range existing { + if chunk == nil { + continue + } + existingByID[chunk.ID] = chunk + if chunk.ChunkType == types.ChunkTypeParentText { + existingParentContent[chunk.ID] = artifact.SHA256Hex([]byte(chunk.Content)) + } + } + for _, chunk := range existing { + if chunk == nil || !isDocumentBaseChunk(chunk.ChunkType) { + continue + } + parentDigest := existingParentContent[chunk.ParentChunkID] + key := legacyChunkSemanticKey(chunk.ChunkType, parentDigest, chunk.Content) + legacyCandidates[key] = append(legacyCandidates[key], chunk) + } + + desiredLegacyCounts := make(map[string]int) + for _, parent := range parsedParents { + desiredLegacyCounts[legacyChunkSemanticKey(types.ChunkTypeParentText, "", parent.Content)]++ + } + for _, chunk := range parsed { + if strings.TrimSpace(chunk.Content) == "" { + continue + } + parentDigest := "" + if chunk.ParentIndex >= 0 && chunk.ParentIndex < len(parsedParents) { + parentDigest = artifact.SHA256Hex([]byte(parsedParents[chunk.ParentIndex].Content)) + } + desiredLegacyCounts[legacyChunkSemanticKey(types.ChunkTypeText, parentDigest, chunk.Content)]++ + } + + allocator := artifact.NewIdentityAllocator(knowledge.ID, artifact.StableEntityIDVersion) + parents := make([]desiredParentIdentity, len(parsedParents)) + result := desiredChunkSet{} + for index, parsedParent := range parsedParents { + contentHash := artifact.SHA256Hex([]byte(parsedParent.Content)) + identity, err := allocator.Next( + types.ChunkTypeParentText, + "", + "content:"+contentHash, + parsedParent.Content, + ) + if err != nil { + return desiredChunkSet{}, err + } + legacyKey := legacyChunkSemanticKey(types.ChunkTypeParentText, "", parsedParent.Content) + id := uniqueLegacyChunkID(identity.ID, legacyKey, desiredLegacyCounts, legacyCandidates) + chunk := &types.Chunk{ + ID: id, + TenantID: knowledge.TenantID, + KnowledgeID: knowledge.ID, + KnowledgeBaseID: knowledge.KnowledgeBaseID, + Content: parsedParent.Content, + ContentHash: contentHash, + ChunkIndex: parsedParent.Seq, + IsEnabled: true, + CreatedAt: now, + UpdatedAt: now, + StartAt: parsedParent.Start, + EndAt: parsedParent.End, + ChunkType: types.ChunkTypeParentText, + } + preserveLiveChunkFields(chunk, existingByID[id]) + parents[index] = desiredParentIdentity{ + chunk: chunk, + semanticKey: identity.SemanticKey, + contentHash: contentHash, + } + result.All = append(result.All, chunk) + } + linkDesiredChunks(parentChunks(parents)) + + for index, parsedChunk := range parsed { + if strings.TrimSpace(parsedChunk.Content) == "" { + continue + } + parentSemantic := "" + parentDigest := "" + parentID := "" + if parsedChunk.ParentIndex >= 0 && parsedChunk.ParentIndex < len(parents) { + parent := parents[parsedChunk.ParentIndex] + parentSemantic = parent.semanticKey + parentDigest = parent.contentHash + parentID = parent.chunk.ID + } + contentHash := artifact.SHA256Hex([]byte(parsedChunk.Content)) + identity, err := allocator.Next( + types.ChunkTypeText, + parentSemantic, + "content:"+contentHash, + parsedChunk.Content, + ) + if err != nil { + return desiredChunkSet{}, err + } + legacyKey := legacyChunkSemanticKey(types.ChunkTypeText, parentDigest, parsedChunk.Content) + id := uniqueLegacyChunkID(identity.ID, legacyKey, desiredLegacyCounts, legacyCandidates) + chunk := &types.Chunk{ + ID: id, + TenantID: knowledge.TenantID, + KnowledgeID: knowledge.ID, + KnowledgeBaseID: knowledge.KnowledgeBaseID, + Content: parsedChunk.Content, + ContentHash: contentHash, + ContextHeader: parsedChunk.ContextHeader, + ChunkIndex: int(parsedChunk.Seq), + IsEnabled: true, + CreatedAt: now, + UpdatedAt: now, + StartAt: int(parsedChunk.Start), + EndAt: int(parsedChunk.End), + ChunkType: types.ChunkTypeText, + ParentChunkID: parentID, + } + preserveLiveChunkFields(chunk, existingByID[id]) + parsed[index].ChunkID = id + result.All = append(result.All, chunk) + result.Text = append(result.Text, chunk) + } + if len(parents) == 0 { + linkDesiredChunks(result.Text) + } + + desiredIDs := make(map[string]struct{}, len(result.All)) + for _, chunk := range result.All { + desiredIDs[chunk.ID] = struct{}{} + if existingByID[chunk.ID] == nil { + result.Added = append(result.Added, chunk) + } else { + result.Updated = append(result.Updated, chunk) + } + } + for _, chunk := range existing { + if chunk == nil || !isDocumentBaseChunk(chunk.ChunkType) { + continue + } + if _, desired := desiredIDs[chunk.ID]; !desired { + result.Stale = append(result.Stale, chunk) + } + } + return result, nil +} + +func parentChunks(parents []desiredParentIdentity) []*types.Chunk { + result := make([]*types.Chunk, 0, len(parents)) + for _, parent := range parents { + result = append(result, parent.chunk) + } + return result +} + +func linkDesiredChunks(chunks []*types.Chunk) { + for index, chunk := range chunks { + chunk.PreChunkID = "" + chunk.NextChunkID = "" + if index > 0 { + chunk.PreChunkID = chunks[index-1].ID + } + if index+1 < len(chunks) { + chunk.NextChunkID = chunks[index+1].ID + } + } +} + +func uniqueLegacyChunkID( + fallback string, + semanticKey string, + desiredCounts map[string]int, + candidates map[string][]*types.Chunk, +) string { + if desiredCounts[semanticKey] == 1 && len(candidates[semanticKey]) == 1 { + return candidates[semanticKey][0].ID + } + return fallback +} + +func legacyChunkSemanticKey(chunkType, parentDigest, content string) string { + return chunkType + "\x00" + parentDigest + "\x00" + artifact.SHA256Hex([]byte(content)) +} + +func preserveLiveChunkFields(desired, live *types.Chunk) { + if live == nil { + return + } + desired.SeqID = live.SeqID + desired.CreatedAt = live.CreatedAt + desired.IsEnabled = live.IsEnabled + desired.Flags = live.Flags + desired.Status = live.Status + desired.Metadata = append(types.JSON(nil), live.Metadata...) + desired.RelationChunks = append(types.JSON(nil), live.RelationChunks...) + desired.IndirectRelationChunks = append(types.JSON(nil), live.IndirectRelationChunks...) + desired.ImageInfo = live.ImageInfo +} + +func isDocumentBaseChunk(chunkType string) bool { + return chunkType == types.ChunkTypeText || chunkType == types.ChunkTypeParentText +} + +func chunkIDList(chunks []*types.Chunk) []string { + result := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + if chunk != nil { + result = append(result, chunk.ID) + } + } + return result +} diff --git a/internal/application/service/knowledge_reconcile_test.go b/internal/application/service/knowledge_reconcile_test.go new file mode 100644 index 0000000000..b4d26ca34f --- /dev/null +++ b/internal/application/service/knowledge_reconcile_test.go @@ -0,0 +1,284 @@ +package service + +import ( + "testing" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func reconcileKnowledge() *types.Knowledge { + return &types.Knowledge{ + ID: uuid.New().String(), + TenantID: 1, + KnowledgeBaseID: uuid.New().String(), + } +} + +func parsedText(content string, seq int) types.ParsedChunk { + return types.ParsedChunk{ + Content: content, + Seq: seq, + Start: seq * 10, + End: seq*10 + len(content), + } +} + +func TestBuildDesiredDocumentChunksStableAcrossInsertionAndReorder(t *testing.T) { + knowledge := reconcileKnowledge() + first, err := buildDesiredDocumentChunks( + knowledge, + []types.ParsedChunk{parsedText("alpha", 0), parsedText("beta", 1)}, + nil, + nil, + ) + require.NoError(t, err) + firstIDs := map[string]string{} + for _, chunk := range first.Text { + firstIDs[chunk.Content] = chunk.ID + } + + reordered, err := buildDesiredDocumentChunks( + knowledge, + []types.ParsedChunk{ + parsedText("inserted", 0), + parsedText("beta", 1), + parsedText("alpha", 2), + }, + nil, + first.All, + ) + require.NoError(t, err) + reorderedIDs := map[string]string{} + for _, chunk := range reordered.Text { + reorderedIDs[chunk.Content] = chunk.ID + } + + assert.Equal(t, firstIDs["alpha"], reorderedIDs["alpha"]) + assert.Equal(t, firstIDs["beta"], reorderedIDs["beta"]) + assert.NotEmpty(t, reorderedIDs["inserted"]) + require.Len(t, reordered.Added, 1) + require.Empty(t, reordered.Stale) +} + +func TestBuildDesiredDocumentChunksReusesUniqueLegacyRandomUUID(t *testing.T) { + knowledge := reconcileKnowledge() + legacy := &types.Chunk{ + ID: uuid.New().String(), + TenantID: knowledge.TenantID, + KnowledgeID: knowledge.ID, + KnowledgeBaseID: knowledge.KnowledgeBaseID, + Content: "unchanged", + ChunkType: types.ChunkTypeText, + } + + desired, err := buildDesiredDocumentChunks( + knowledge, + []types.ParsedChunk{parsedText("unchanged", 0)}, + nil, + []*types.Chunk{legacy}, + ) + require.NoError(t, err) + require.Len(t, desired.Text, 1) + assert.Equal(t, legacy.ID, desired.Text[0].ID) + require.Len(t, desired.Updated, 1) + require.Empty(t, desired.Added) + require.Empty(t, desired.Stale) +} + +func TestBuildDesiredDocumentChunksPreservesLiveEnableState(t *testing.T) { + knowledge := reconcileKnowledge() + first, err := buildDesiredDocumentChunks( + knowledge, + []types.ParsedChunk{parsedText("keep disabled", 0)}, + nil, + nil, + ) + require.NoError(t, err) + require.Len(t, first.Text, 1) + live := *first.Text[0] + live.IsEnabled = false + + next, err := buildDesiredDocumentChunks( + knowledge, + []types.ParsedChunk{parsedText("keep disabled", 0)}, + nil, + []*types.Chunk{&live}, + ) + require.NoError(t, err) + require.Len(t, next.Text, 1) + assert.False(t, next.Text[0].IsEnabled) +} + +func TestBuildDesiredDocumentChunksDoesNotGuessAmbiguousLegacyDuplicates(t *testing.T) { + knowledge := reconcileKnowledge() + existing := []*types.Chunk{ + { + ID: uuid.New().String(), + TenantID: knowledge.TenantID, + KnowledgeID: knowledge.ID, + KnowledgeBaseID: knowledge.KnowledgeBaseID, + Content: "repeat", + ChunkType: types.ChunkTypeText, + }, + { + ID: uuid.New().String(), + TenantID: knowledge.TenantID, + KnowledgeID: knowledge.ID, + KnowledgeBaseID: knowledge.KnowledgeBaseID, + Content: "repeat", + ChunkType: types.ChunkTypeText, + }, + } + + desired, err := buildDesiredDocumentChunks( + knowledge, + []types.ParsedChunk{parsedText("repeat", 0), parsedText("repeat", 1)}, + nil, + existing, + ) + require.NoError(t, err) + require.Len(t, desired.Text, 2) + assert.NotEqual(t, desired.Text[0].ID, desired.Text[1].ID) + assert.NotEqual(t, existing[0].ID, desired.Text[0].ID) + assert.NotEqual(t, existing[1].ID, desired.Text[1].ID) + require.Len(t, desired.Added, 2) + require.Len(t, desired.Stale, 2) +} + +func TestBuildDesiredDocumentChunksParentIdentityAffectsChildren(t *testing.T) { + knowledge := reconcileKnowledge() + parents := []types.ParsedParentChunk{ + {Content: "parent A"}, + {Content: "parent B"}, + } + desired, err := buildDesiredDocumentChunks( + knowledge, + []types.ParsedChunk{ + {Content: "same child", ParentIndex: 0}, + {Content: "same child", ParentIndex: 1}, + }, + parents, + nil, + ) + require.NoError(t, err) + require.Len(t, desired.Text, 2) + assert.NotEqual(t, desired.Text[0].ID, desired.Text[1].ID) + assert.NotEqual(t, desired.Text[0].ParentChunkID, desired.Text[1].ParentChunkID) +} + +func TestStableGeneratedQuestionsSurviveReordering(t *testing.T) { + knowledge := reconcileKnowledge() + parent := &types.Chunk{ID: uuid.New().String()} + + first, err := stableGeneratedQuestions( + knowledge.ID, + parent, + []string{"What is alpha?", "What is beta?"}, + ) + require.NoError(t, err) + second, err := stableGeneratedQuestions( + knowledge.ID, + parent, + []string{"What is beta?", "What is alpha?"}, + ) + require.NoError(t, err) + + firstIDs := map[string]string{} + for _, question := range first { + firstIDs[question.Question] = question.ID + } + for _, question := range second { + assert.Equal(t, firstIDs[question.Question], question.ID) + } +} + +func TestStableGeneratedQuestionsReusesUniqueLegacyID(t *testing.T) { + knowledge := reconcileKnowledge() + parent := &types.Chunk{ID: uuid.New().String()} + require.NoError(t, parent.SetDocumentMetadata(&types.DocumentChunkMetadata{ + GeneratedQuestions: []types.GeneratedQuestion{{ + ID: "q-legacy", + Question: "What is alpha?", + }}, + })) + + questions, err := stableGeneratedQuestions( + knowledge.ID, + parent, + []string{"What is alpha?"}, + ) + require.NoError(t, err) + require.Len(t, questions, 1) + assert.Equal(t, "q-legacy", questions[0].ID) +} + +func TestStaleGeneratedQuestionSourceIDsAreExact(t *testing.T) { + parent := &types.Chunk{ID: uuid.New().String()} + require.NoError(t, parent.SetDocumentMetadata(&types.DocumentChunkMetadata{ + GeneratedQuestions: []types.GeneratedQuestion{ + {ID: "keep", Question: "kept"}, + {ID: "stale", Question: "removed"}, + }, + })) + + stale := staleGeneratedQuestionSourceIDs(parent, []types.GeneratedQuestion{{ + ID: "keep", + Question: "kept", + }}) + assert.Equal(t, []string{parent.ID + "-stale"}, stale) +} + +func TestStableImageDerivedChunkID(t *testing.T) { + knowledge := reconcileKnowledge() + parent := &types.Chunk{ + ID: uuid.New().String(), + KnowledgeID: knowledge.ID, + } + first, err := stableImageDerivedChunkID( + parent, + types.ChunkTypeImageOCR, + "provider://images/one", + "recognized text", + ) + require.NoError(t, err) + retry, err := stableImageDerivedChunkID( + parent, + types.ChunkTypeImageOCR, + "provider://images/one", + "recognized text", + ) + require.NoError(t, err) + changed, err := stableImageDerivedChunkID( + parent, + types.ChunkTypeImageOCR, + "provider://images/one", + "changed text", + ) + require.NoError(t, err) + + assert.Equal(t, first, retry) + assert.NotEqual(t, first, changed) +} + +func TestDataTableDerivedChunkIDsAreStable(t *testing.T) { + knowledge := reconcileKnowledge() + service := &DataTableSummaryService{} + resources := &extractionResources{knowledge: knowledge} + + first, err := service.buildChunks(resources, "summary", "columns") + require.NoError(t, err) + retry, err := service.buildChunks(resources, "summary", "columns") + require.NoError(t, err) + changed, err := service.buildChunks(resources, "new summary", "columns") + require.NoError(t, err) + + require.Len(t, first, 2) + assert.Equal(t, first[0].ID, retry[0].ID) + assert.Equal(t, first[1].ID, retry[1].ID) + assert.NotEqual(t, first[0].ID, changed[0].ID) + assert.NotEqual(t, first[1].ID, changed[1].ID) + assert.Equal(t, first[0].ID, first[1].ParentChunkID) +} diff --git a/internal/application/service/model.go b/internal/application/service/model.go index d25da56c56..68baddc001 100644 --- a/internal/application/service/model.go +++ b/internal/application/service/model.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" + "github.com/Tencent/WeKnora/internal/artifact" apperrors "github.com/Tencent/WeKnora/internal/errors" "github.com/Tencent/WeKnora/internal/logger" "github.com/Tencent/WeKnora/internal/models/asr" @@ -24,12 +25,27 @@ var ErrModelNotFound = errors.New("model not found") // modelService implements the model service interface type modelService struct { - repo interfaces.ModelRepository - kbRepo interfaces.KnowledgeBaseRepository - agentRepo interfaces.CustomAgentRepository - ollamaService *ollama.OllamaService - pooler embedding.EmbedderPooler - tenantService interfaces.TenantService + repo interfaces.ModelRepository + kbRepo interfaces.KnowledgeBaseRepository + agentRepo interfaces.CustomAgentRepository + ollamaService *ollama.OllamaService + pooler embedding.EmbedderPooler + tenantService interfaces.TenantService + artifactRuntime *artifact.Runtime +} + +// ConfigureModelArtifactCache attaches the shared fail-open artifact runtime +// after dependency injection constructs the model service. Keeping the public +// constructor stable avoids forcing artifact storage into lightweight tests. +func ConfigureModelArtifactCache( + service interfaces.ModelService, + runtime *artifact.Runtime, +) { + models, ok := service.(*modelService) + if !ok || runtime == nil { + return + } + models.artifactRuntime = runtime } // NewModelService creates a new model service instance @@ -431,6 +447,12 @@ func (s *modelService) GetEmbeddingModel(ctx context.Context, modelId string) (e }) return nil, err } + if config, cacheable := embedding.ArtifactCacheConfigFromModel( + model, + types.MustTenantIDFromContext(ctx), + ); cacheable { + embedder = embedding.NewArtifactCachedEmbedder(embedder, s.artifactRuntime, config) + } logger.Info(ctx, "Embedding model initialized successfully") return embedder, nil @@ -479,6 +501,9 @@ func (s *modelService) GetEmbeddingModelForTenant(ctx context.Context, modelId s }) return nil, err } + if config, cacheable := embedding.ArtifactCacheConfigFromModel(model, tenantID); cacheable { + embedder = embedding.NewArtifactCachedEmbedder(embedder, s.artifactRuntime, config) + } logger.Info(ctx, "Cross-tenant embedding model initialized successfully") return embedder, nil @@ -551,6 +576,9 @@ func (s *modelService) GetChatModel(ctx context.Context, modelId string) (chat.C }) return nil, err } + if config, cacheable := chat.ArtifactCacheConfigFromModel(model, tenantID); cacheable { + chatModel = chat.NewArtifactCachedChat(chatModel, s.artifactRuntime, config) + } return chatModel, nil } @@ -588,6 +616,9 @@ func (s *modelService) GetVLMModel(ctx context.Context, modelId string) (vlm.VLM }) return nil, err } + if config, cacheable := vlm.ArtifactCacheConfigFromModel(model, tenantID); cacheable { + vlmModel = vlm.NewArtifactCachedVLM(vlmModel, s.artifactRuntime, config) + } return vlmModel, nil } diff --git a/internal/application/service/retriever/keywords_vector_hybrid_indexer.go b/internal/application/service/retriever/keywords_vector_hybrid_indexer.go index 3e62ceb86e..71289a408f 100644 --- a/internal/application/service/retriever/keywords_vector_hybrid_indexer.go +++ b/internal/application/service/retriever/keywords_vector_hybrid_indexer.go @@ -9,7 +9,7 @@ import ( "unicode/utf8" "github.com/Tencent/WeKnora/internal/logger" - "github.com/Tencent/WeKnora/internal/models/embedding" + embeddingmodel "github.com/Tencent/WeKnora/internal/models/embedding" "github.com/Tencent/WeKnora/internal/models/utils" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" @@ -67,15 +67,23 @@ func (v *KeywordsVectorHybridRetrieveEngineService) Retrieve(ctx context.Context // Index creates embeddings for the content and saves it to the repository // if vector retrieval is enabled in the retriever types func (v *KeywordsVectorHybridRetrieveEngineService) Index(ctx context.Context, - embedder embedding.Embedder, indexInfo *types.IndexInfo, retrieverTypes []types.RetrieverType, + embedder embeddingmodel.Embedder, indexInfo *types.IndexInfo, retrieverTypes []types.RetrieverType, ) error { params := make(map[string]any) embeddingMap := make(map[string][]float32) if slices.Contains(retrieverTypes, types.VectorRetrieverType) { - embedding, err := embedder.Embed(ctx, sanitizeForEmbedding(ctx, indexInfo.Content)) + embedCtx := context.WithValue(ctx, types.EmbedDocumentContextKey, true) + embedding, err := embedder.Embed(embedCtx, sanitizeForEmbedding(ctx, indexInfo.Content)) if err != nil { return err } + if err := embeddingmodel.ValidateEmbeddingBatch( + [][]float32{embedding}, + 1, + embedder.GetDimensions(), + ); err != nil { + return err + } embeddingMap[indexInfo.SourceID] = embedding } params["embedding"] = embeddingMap @@ -85,7 +93,7 @@ func (v *KeywordsVectorHybridRetrieveEngineService) Index(ctx context.Context, // BatchIndex creates embeddings for multiple content items and saves them to the repository // in batches for efficiency. Uses concurrent batch saving to improve performance. func (v *KeywordsVectorHybridRetrieveEngineService) BatchIndex(ctx context.Context, - embedder embedding.Embedder, indexInfoList []*types.IndexInfo, retrieverTypes []types.RetrieverType, + embedder embeddingmodel.Embedder, indexInfoList []*types.IndexInfo, retrieverTypes []types.RetrieverType, ) error { if len(indexInfoList) == 0 { return nil @@ -96,10 +104,18 @@ func (v *KeywordsVectorHybridRetrieveEngineService) BatchIndex(ctx context.Conte for _, indexInfo := range indexInfoList { contentList = append(contentList, sanitizeForEmbedding(ctx, indexInfo.Content)) } - embeddings, err := batchEmbedWithBackoff(ctx, embedder, contentList) + embedCtx := context.WithValue(ctx, types.EmbedDocumentContextKey, true) + embeddings, err := batchEmbedWithBackoff(embedCtx, embedder, contentList) if err != nil { return err } + if err := embeddingmodel.ValidateEmbeddingBatch( + embeddings, + len(contentList), + embedder.GetDimensions(), + ); err != nil { + return err + } batchSize := 40 chunks := utils.ChunkSlice(indexInfoList, batchSize) @@ -128,7 +144,7 @@ func (v *KeywordsVectorHybridRetrieveEngineService) BatchIndex(ctx context.Conte // batchEmbedWithBackoff calls BatchEmbedWithPool with exponential backoff on // transient failures (200 / 400 / 800 / 1600 / 3200 ms). It returns the last // embedding result on success or the last error if every attempt failed. -func batchEmbedWithBackoff(ctx context.Context, embedder embedding.Embedder, contentList []string) ([][]float32, error) { +func batchEmbedWithBackoff(ctx context.Context, embedder embeddingmodel.Embedder, contentList []string) ([][]float32, error) { delay := embedRetryBaseDelay var ( embeddings [][]float32 @@ -297,7 +313,7 @@ func (v *KeywordsVectorHybridRetrieveEngineService) Support() []types.RetrieverT // EstimateStorageSize estimates the storage space needed for the provided index information func (v *KeywordsVectorHybridRetrieveEngineService) EstimateStorageSize( ctx context.Context, - embedder embedding.Embedder, + embedder embeddingmodel.Embedder, indexInfoList []*types.IndexInfo, retrieverTypes []types.RetrieverType, ) int64 { diff --git a/internal/application/service/wiki_ingest.go b/internal/application/service/wiki_ingest.go index be05f50e7e..775342b9d5 100644 --- a/internal/application/service/wiki_ingest.go +++ b/internal/application/service/wiki_ingest.go @@ -2477,6 +2477,9 @@ func (s *wikiIngestService) generateWithTemplate(ctx context.Context, chatModel } } ctx = types.WithLLMCallMetadata(ctx, purpose, prefixFingerprint) + if artifactStage, cacheable := wikiMapArtifactStage(purpose); cacheable { + ctx = chat.WithArtifactStage(ctx, artifactStage) + } tenantID, tenantScoped := types.TenantIDFromContext(ctx) requestJSON, _ := json.Marshal(struct { @@ -2556,6 +2559,46 @@ func (s *wikiIngestService) generateWithTemplate(ctx context.Context, chatModel } } +func wikiMapArtifactStage(purpose string) (chat.ArtifactStage, bool) { + jsonOutput := func(content string) error { + cleaned := strings.TrimSpace(cleanLLMJSON(content)) + if cleaned == "" { + return errors.New("wiki map output is empty") + } + if !json.Valid([]byte(cleaned)) { + return errors.New("wiki map output is not valid JSON") + } + return nil + } + nonEmpty := func(content string) error { + if strings.TrimSpace(content) == "" { + return errors.New("wiki map output is empty") + } + return nil + } + + stage := chat.ArtifactStage{ + Stage: "wiki_map", + OutputSchema: "wiki-map." + purpose + ".v1", + } + switch purpose { + case "wiki_knowledge_extract", + "wiki_candidate_slug", + "wiki_chunk_citation", + "wiki_taxonomy_plan": + stage.Validate = jsonOutput + return stage, true + case "wiki_summary": + stage.Validate = nonEmpty + return stage, true + default: + // Page modification, deduplication and index introduction are reduce + // operations over mutable live state. They deliberately bypass + // immutable map-artifact reuse. + return chat.ArtifactStage{}, false + } +} + func wikiPromptPurpose(promptTpl string) string { switch promptTpl { case agent.WikiPageModifyUserPrompt: diff --git a/internal/application/service/wiki_ingest_test.go b/internal/application/service/wiki_ingest_test.go index 71712e7ad1..ef2dd432a3 100644 --- a/internal/application/service/wiki_ingest_test.go +++ b/internal/application/service/wiki_ingest_test.go @@ -611,3 +611,27 @@ func (r *wikiPendingRepoForCleanupTest) DeleteByDedupKey( ) error { return nil } + +func TestWikiMapArtifactStageExcludesLiveReduceOperations(t *testing.T) { + for _, purpose := range []string{ + "wiki_knowledge_extract", + "wiki_candidate_slug", + "wiki_chunk_citation", + "wiki_taxonomy_plan", + "wiki_summary", + } { + stage, ok := wikiMapArtifactStage(purpose) + if !ok || stage.Stage != "wiki_map" || stage.OutputSchema == "" { + t.Fatalf("purpose %q was not classified as wiki map: %#v, %v", purpose, stage, ok) + } + } + for _, purpose := range []string{ + "wiki_page_modify", + "wiki_deduplication", + "wiki_index_intro", + } { + if stage, ok := wikiMapArtifactStage(purpose); ok { + t.Fatalf("live reduce purpose %q must not be cached: %#v", purpose, stage) + } + } +} diff --git a/internal/artifact/codec.go b/internal/artifact/codec.go new file mode 100644 index 0000000000..abe12c410a --- /dev/null +++ b/internal/artifact/codec.go @@ -0,0 +1,195 @@ +package artifact + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "io" + + "github.com/Tencent/WeKnora/internal/types" +) + +const ( + CodecJSONV1 = "json.v1" + CodecFloat32BEV1 = "float32be.v1" + CodecTextUTF8V1 = "text.utf8.v1" +) + +var ErrCorruptArtifact = errors.New("corrupt processing artifact") + +// NewInlineArtifact freezes a successful ownership-free output. A nil payload +// is represented as a non-nil empty byte slice so successful empty output is +// distinct from an external-object manifest. +func NewInlineArtifact(key Key, codec string, payload []byte) (*types.ProcessingArtifact, error) { + if err := key.Lookup.Validate(); err != nil { + return nil, err + } + if !isSHA256(key.ProcessorDigest) { + return nil, errors.New("artifact processor digest must be 64 lowercase hex characters") + } + if key.OutputSchema == "" { + return nil, errors.New("artifact output schema must not be empty") + } + if codec == "" { + return nil, errors.New("artifact codec must not be empty") + } + frozen := make([]byte, len(payload)) + copy(frozen, payload) + digest := SHA256Hex(frozen) + return &types.ProcessingArtifact{ + TenantID: key.Lookup.TenantID, + Stage: key.Lookup.Stage, + KeyVersion: key.Lookup.KeyVersion, + ArtifactKey: key.Lookup.ArtifactKey, + ProcessorDigest: key.ProcessorDigest, + OutputDigest: digest, + OutputSchema: key.OutputSchema, + Codec: codec, + InlinePayload: true, + Payload: frozen, + PayloadChecksum: digest, + SizeBytes: int64(len(frozen)), + }, nil +} + +// EncodeJSON canonicalizes and checks ownership before constructing a manifest. +func EncodeJSON(key Key, value any) (*types.ProcessingArtifact, error) { + if err := ValidateOwnershipFreeJSON(value); err != nil { + return nil, err + } + payload, err := CanonicalJSON(value) + if err != nil { + return nil, err + } + return NewInlineArtifact(key, CodecJSONV1, payload) +} + +// DecodeInline validates manifest identity, schema, size, storage checksum and +// semantic output digest before returning a detached payload. +func DecodeInline( + artifact *types.ProcessingArtifact, + key types.ProcessingArtifactLookup, + outputSchema string, + codec string, +) ([]byte, error) { + if artifact == nil { + return nil, fmt.Errorf("%w: manifest is nil", ErrCorruptArtifact) + } + if err := artifact.Validate(); err != nil { + return nil, fmt.Errorf("%w: %v", ErrCorruptArtifact, err) + } + if artifact.Lookup() != key { + return nil, fmt.Errorf("%w: manifest key mismatch", ErrCorruptArtifact) + } + if artifact.OutputSchema != outputSchema { + return nil, fmt.Errorf("%w: output schema mismatch", ErrCorruptArtifact) + } + if artifact.Codec != codec { + return nil, fmt.Errorf("%w: codec mismatch", ErrCorruptArtifact) + } + if !artifact.InlinePayload { + return nil, fmt.Errorf("%w: object payload reader is not configured", ErrCorruptArtifact) + } + if int64(len(artifact.Payload)) != artifact.SizeBytes { + return nil, fmt.Errorf("%w: payload size mismatch", ErrCorruptArtifact) + } + checksum := SHA256Hex(artifact.Payload) + if checksum != artifact.PayloadChecksum { + return nil, fmt.Errorf("%w: payload checksum mismatch", ErrCorruptArtifact) + } + if checksum != artifact.OutputDigest { + return nil, fmt.Errorf("%w: output digest mismatch", ErrCorruptArtifact) + } + result := make([]byte, len(artifact.Payload)) + copy(result, artifact.Payload) + return result, nil +} + +// DecodeJSON additionally rejects trailing data and validates that cached +// payloads remain ownership-free. +func DecodeJSON( + artifact *types.ProcessingArtifact, + key types.ProcessingArtifactLookup, + outputSchema string, + target any, +) error { + payload, err := DecodeInline(artifact, key, outputSchema, CodecJSONV1) + if err != nil { + return err + } + decoder := json.NewDecoder(bytes.NewReader(payload)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(target); err != nil { + return fmt.Errorf("%w: decode JSON: %v", ErrCorruptArtifact, err) + } + var trailing any + if err := decoder.Decode(&trailing); err != io.EOF { + return fmt.Errorf("%w: trailing JSON value", ErrCorruptArtifact) + } + if err := ValidateOwnershipFreeJSON(target); err != nil { + return fmt.Errorf("%w: %v", ErrCorruptArtifact, err) + } + return nil +} + +// ValidateOwnershipFreeJSON rejects identifiers that bind a reusable artifact +// to one live knowledge entity or processing attempt. +func ValidateOwnershipFreeJSON(value any) error { + raw, err := json.Marshal(value) + if err != nil { + return err + } + var decoded any + if err := json.Unmarshal(raw, &decoded); err != nil { + return err + } + return walkOwnershipFields(decoded, "") +} + +func walkOwnershipFields(value any, path string) error { + switch typed := value.(type) { + case []any: + for index, item := range typed { + if err := walkOwnershipFields(item, fmt.Sprintf("%s[%d]", path, index)); err != nil { + return err + } + } + case map[string]any: + for key, item := range typed { + normalized := normalizeFieldName(key) + if _, forbidden := ownershipFieldNames[normalized]; forbidden { + return fmt.Errorf("ownership field %q must not enter processing artifact payload", joinJSONPath(path, key)) + } + if err := walkOwnershipFields(item, joinJSONPath(path, key)); err != nil { + return err + } + } + } + return nil +} + +func normalizeFieldName(value string) string { + result := make([]byte, 0, len(value)) + for index := 0; index < len(value); index++ { + switch value[index] { + case '_', '-', '.': + continue + default: + if value[index] >= 'A' && value[index] <= 'Z' { + result = append(result, value[index]+'a'-'A') + } else { + result = append(result, value[index]) + } + } + } + return string(result) +} + +var ownershipFieldNames = map[string]struct{}{ + "attemptid": {}, + "attempt": {}, + "chunkid": {}, + "knowledgeid": {}, + "tenantid": {}, +} diff --git a/internal/artifact/codec_test.go b/internal/artifact/codec_test.go new file mode 100644 index 0000000000..ba91d6c944 --- /dev/null +++ b/internal/artifact/codec_test.go @@ -0,0 +1,48 @@ +package artifact + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestInlineArtifactAllowsSuccessfulEmptyOutput(t *testing.T) { + key, err := BuildKey(1, testKeyMaterial()) + require.NoError(t, err) + manifest, err := NewInlineArtifact(key, "empty.v1", nil) + require.NoError(t, err) + + assert.NotNil(t, manifest.Payload) + assert.Zero(t, manifest.SizeBytes) + payload, err := DecodeInline(manifest, key.Lookup, key.OutputSchema, "empty.v1") + require.NoError(t, err) + assert.Empty(t, payload) +} + +func TestDecodeInlineRejectsCorruptionAndSchemaMismatch(t *testing.T) { + key, err := BuildKey(1, testKeyMaterial()) + require.NoError(t, err) + manifest, err := NewInlineArtifact(key, CodecJSONV1, []byte(`{"ok":true}`)) + require.NoError(t, err) + + manifest.Payload[1] = 'X' + _, err = DecodeInline(manifest, key.Lookup, key.OutputSchema, CodecJSONV1) + require.ErrorIs(t, err, ErrCorruptArtifact) + + manifest, err = NewInlineArtifact(key, CodecJSONV1, []byte(`{"ok":true}`)) + require.NoError(t, err) + _, err = DecodeInline(manifest, key.Lookup, "other.v1", CodecJSONV1) + require.ErrorIs(t, err, ErrCorruptArtifact) +} + +func TestEncodeJSONRejectsOwnershipFields(t *testing.T) { + key, err := BuildKey(1, testKeyMaterial()) + require.NoError(t, err) + + _, err = EncodeJSON(key, map[string]any{ + "value": "valid provider output", + "knowledge_id": "live-owner", + }) + require.ErrorContains(t, err, "ownership field") +} diff --git a/internal/artifact/fence.go b/internal/artifact/fence.go new file mode 100644 index 0000000000..4d48ecac11 --- /dev/null +++ b/internal/artifact/fence.go @@ -0,0 +1,47 @@ +package artifact + +import ( + "context" + "errors" + "fmt" +) + +var ErrAttemptSuperseded = errors.New("processing attempt superseded") + +type LatestAttemptReader interface { + LatestAttempt(ctx context.Context, knowledgeID string) (int, error) +} + +// AttemptFence uses persisted attempt state. Artifact computation may continue +// after this returns ErrAttemptSuperseded, but live bind/publish/cleanup must not. +type AttemptFence struct { + reader LatestAttemptReader +} + +func NewAttemptFence(reader LatestAttemptReader) AttemptFence { + return AttemptFence{reader: reader} +} + +func (f AttemptFence) EnsureCurrent( + ctx context.Context, + knowledgeID string, + attempt int, +) error { + if f.reader == nil { + return errors.New("attempt fence reader is not configured") + } + if knowledgeID == "" { + return errors.New("attempt fence knowledge ID must not be empty") + } + if attempt <= 0 { + return errors.New("attempt fence attempt must be positive") + } + latest, err := f.reader.LatestAttempt(ctx, knowledgeID) + if err != nil { + return fmt.Errorf("load latest processing attempt: %w", err) + } + if latest != attempt { + return fmt.Errorf("%w: expected %d, latest %d", ErrAttemptSuperseded, attempt, latest) + } + return nil +} diff --git a/internal/artifact/fence_test.go b/internal/artifact/fence_test.go new file mode 100644 index 0000000000..aaad062143 --- /dev/null +++ b/internal/artifact/fence_test.go @@ -0,0 +1,28 @@ +package artifact + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" +) + +type staticAttemptReader struct { + attempt int + err error +} + +func (r staticAttemptReader) LatestAttempt(context.Context, string) (int, error) { + return r.attempt, r.err +} + +func TestAttemptFenceRequiresExactPersistedAttempt(t *testing.T) { + fence := NewAttemptFence(staticAttemptReader{attempt: 3}) + require.NoError(t, fence.EnsureCurrent(context.Background(), "knowledge", 3)) + require.ErrorIs(t, fence.EnsureCurrent(context.Background(), "knowledge", 2), ErrAttemptSuperseded) + + readErr := errors.New("database unavailable") + fence = NewAttemptFence(staticAttemptReader{err: readErr}) + require.ErrorIs(t, fence.EnsureCurrent(context.Background(), "knowledge", 3), readErr) +} diff --git a/internal/artifact/identity.go b/internal/artifact/identity.go new file mode 100644 index 0000000000..1628c1f79e --- /dev/null +++ b/internal/artifact/identity.go @@ -0,0 +1,139 @@ +package artifact + +import ( + "errors" + "fmt" + "strconv" + "strings" + + "github.com/google/uuid" +) + +const StableEntityIDVersion uint16 = 1 + +var weknoraEntityNamespace = uuid.MustParse("8298b3ec-ccef-5ae4-a326-7e29c2d07b41") + +// EntityIdentityInput contains only semantic identity fields. Position or a +// global sequence number must never be used as SourceAnchor. +type EntityIdentityInput struct { + KnowledgeID string + IDVersion uint16 + EntityType string + ParentSemanticID string + SourceAnchor string + Content string + DuplicateOrdinal uint32 +} + +type EntityIdentity struct { + ID string + MatchKey string + SemanticKey string + ContentDigest string +} + +// StableEntityIdentity returns a UUIDv5 scoped to the knowledge namespace. +func StableEntityIdentity(input EntityIdentityInput) (EntityIdentity, error) { + if strings.TrimSpace(input.KnowledgeID) == "" { + return EntityIdentity{}, errors.New("entity knowledge ID must not be empty") + } + if input.IDVersion == 0 { + return EntityIdentity{}, errors.New("entity ID version must not be zero") + } + if strings.TrimSpace(input.EntityType) == "" { + return EntityIdentity{}, errors.New("entity type must not be empty") + } + if strings.TrimSpace(input.SourceAnchor) == "" { + return EntityIdentity{}, errors.New("entity source anchor must not be empty") + } + + contentDigest := SHA256Hex([]byte(input.Content)) + matchMaterial := strings.Join([]string{ + "weknora.entity-match.v1", + strconv.FormatUint(uint64(input.IDVersion), 10), + input.EntityType, + input.ParentSemanticID, + input.SourceAnchor, + strconv.FormatUint(uint64(input.DuplicateOrdinal), 10), + }, "\x00") + matchKey := SHA256Hex([]byte(matchMaterial)) + semanticMaterial := strings.Join([]string{ + "weknora.entity-id.v1", + matchKey, + contentDigest, + }, "\x00") + semanticKey := SHA256Hex([]byte(semanticMaterial)) + + namespace := uuid.NewSHA1(weknoraEntityNamespace, []byte(input.KnowledgeID)) + if parsed, err := uuid.Parse(input.KnowledgeID); err == nil { + namespace = parsed + } + id := uuid.NewSHA1(namespace, []byte(semanticMaterial)) + return EntityIdentity{ + ID: id.String(), + MatchKey: matchKey, + SemanticKey: semanticKey, + ContentDigest: contentDigest, + }, nil +} + +// IdentityAllocator supplies a local duplicate ordinal only when all other +// identity fields are identical. +type IdentityAllocator struct { + knowledgeID string + idVersion uint16 + occurrences map[string]uint32 +} + +func NewIdentityAllocator(knowledgeID string, idVersion uint16) *IdentityAllocator { + return &IdentityAllocator{ + knowledgeID: knowledgeID, + idVersion: idVersion, + occurrences: make(map[string]uint32), + } +} + +func (a *IdentityAllocator) Next( + entityType string, + parentSemanticID string, + sourceAnchor string, + content string, +) (EntityIdentity, error) { + if a == nil { + return EntityIdentity{}, errors.New("entity identity allocator must not be nil") + } + occurrenceKey := strings.Join([]string{ + entityType, + parentSemanticID, + sourceAnchor, + SHA256Hex([]byte(content)), + }, "\x00") + ordinal := a.occurrences[occurrenceKey] + a.occurrences[occurrenceKey] = ordinal + 1 + return StableEntityIdentity(EntityIdentityInput{ + KnowledgeID: a.knowledgeID, + IDVersion: a.idVersion, + EntityType: entityType, + ParentSemanticID: parentSemanticID, + SourceAnchor: sourceAnchor, + Content: content, + DuplicateOrdinal: ordinal, + }) +} + +// ReuseUniqueLegacyID preserves an existing random UUID when exactly one live +// entity has the same semantic key. Ambiguous matches deliberately fall back to +// the deterministic ID. +func ReuseUniqueLegacyID( + desired EntityIdentity, + existing map[string][]string, +) (string, bool, error) { + ids := existing[desired.SemanticKey] + if len(ids) != 1 { + return desired.ID, false, nil + } + if _, err := uuid.Parse(ids[0]); err != nil { + return "", false, fmt.Errorf("legacy entity ID %q is not a UUID: %w", ids[0], err) + } + return ids[0], true, nil +} diff --git a/internal/artifact/identity_test.go b/internal/artifact/identity_test.go new file mode 100644 index 0000000000..2e0c479a03 --- /dev/null +++ b/internal/artifact/identity_test.go @@ -0,0 +1,76 @@ +package artifact + +import ( + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestStableEntityIdentityIgnoresGlobalOrder(t *testing.T) { + knowledgeID := uuid.New().String() + input := EntityIdentityInput{ + KnowledgeID: knowledgeID, + IDVersion: StableEntityIDVersion, + EntityType: "text", + SourceAnchor: "page=2;block=heading-intro", + Content: "Exact persisted content\r\n", + } + before, err := StableEntityIdentity(input) + require.NoError(t, err) + + // A new unrelated block inserted before this block does not alter any + // semantic field or introduce a sequence number. + after, err := StableEntityIdentity(input) + require.NoError(t, err) + assert.Equal(t, before, after) + + input.Content = "Exact persisted content\n" + changed, err := StableEntityIdentity(input) + require.NoError(t, err) + assert.NotEqual(t, before.ID, changed.ID) + assert.Equal(t, before.MatchKey, changed.MatchKey) +} + +func TestStableEntityIdentityIncludesParentAndLocalDuplicateOrdinal(t *testing.T) { + knowledgeID := uuid.New().String() + allocator := NewIdentityAllocator(knowledgeID, StableEntityIDVersion) + + first, err := allocator.Next("text", "parent-a", "block=7", "same") + require.NoError(t, err) + second, err := allocator.Next("text", "parent-a", "block=7", "same") + require.NoError(t, err) + otherParent, err := allocator.Next("text", "parent-b", "block=7", "same") + require.NoError(t, err) + + assert.NotEqual(t, first.ID, second.ID) + assert.NotEqual(t, first.MatchKey, second.MatchKey) + assert.NotEqual(t, first.ID, otherParent.ID) +} + +func TestReuseUniqueLegacyIDOnlyWhenUnambiguous(t *testing.T) { + desired, err := StableEntityIdentity(EntityIdentityInput{ + KnowledgeID: uuid.New().String(), + IDVersion: StableEntityIDVersion, + EntityType: "text", + SourceAnchor: "block=1", + Content: "content", + }) + require.NoError(t, err) + legacyID := uuid.New().String() + + reused, ok, err := ReuseUniqueLegacyID(desired, map[string][]string{ + desired.SemanticKey: {legacyID}, + }) + require.NoError(t, err) + assert.True(t, ok) + assert.Equal(t, legacyID, reused) + + reused, ok, err = ReuseUniqueLegacyID(desired, map[string][]string{ + desired.SemanticKey: {legacyID, uuid.New().String()}, + }) + require.NoError(t, err) + assert.False(t, ok) + assert.Equal(t, desired.ID, reused) +} diff --git a/internal/artifact/key.go b/internal/artifact/key.go new file mode 100644 index 0000000000..09c0e815c7 --- /dev/null +++ b/internal/artifact/key.go @@ -0,0 +1,304 @@ +// Package artifact implements content-addressed processing artifacts and the +// stable identities used to reconcile them into live knowledge state. +package artifact + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "math" + "net/url" + "sort" + "strconv" + "strings" + + "github.com/Tencent/WeKnora/internal/types" +) + +const CanonicalJSONVersion = "weknora.canonical-json.v1" + +// DirectInput identifies one ordered direct dependency. Downstream stages use +// an upstream OutputDigest here, not the upstream stage's entire configuration. +type DirectInput struct { + Role string `json:"role"` + Digest string `json:"digest"` +} + +// ProcessorIdentity contains only effective, non-secret provider identity. +// Callers must render the exact provider request separately in KeyMaterial. +type ProcessorIdentity struct { + ModelID string `json:"model_id,omitempty"` + ModelName string `json:"model_name,omitempty"` + Source string `json:"source,omitempty"` + Provider string `json:"provider,omitempty"` + EndpointIdentity string `json:"endpoint_identity,omitempty"` + Revision string `json:"revision,omitempty"` + Parameters map[string]any `json:"parameters,omitempty"` +} + +// KeyMaterial is the versioned canonical envelope hashed into an artifact key. +// DirectInputs is deliberately ordered. +type KeyMaterial struct { + KeyVersion uint16 `json:"key_version"` + Stage string `json:"stage"` + DirectInputs []DirectInput `json:"direct_inputs"` + Processor ProcessorIdentity `json:"processor"` + RenderedRequest any `json:"rendered_request"` + Options any `json:"options"` + CanonicalizerVersion string `json:"canonicalizer_version"` + OutputSchemaVersion string `json:"output_schema_version"` +} + +// Key is the complete result of canonical key construction. +type Key struct { + Lookup types.ProcessingArtifactLookup + ProcessorDigest string + OutputSchema string + Canonical []byte +} + +// BuildKey hashes the exact canonical envelope without altering provider +// inputs. Tenant identity is enforced by the database key, not mixed into the +// reusable content digest. +func BuildKey(tenantID uint64, material KeyMaterial) (Key, error) { + lookup := types.ProcessingArtifactLookup{ + TenantID: tenantID, + Stage: material.Stage, + KeyVersion: material.KeyVersion, + ArtifactKey: strings.Repeat("0", sha256.Size*2), + } + if err := lookup.Validate(); err != nil { + return Key{}, err + } + if len(material.DirectInputs) == 0 { + return Key{}, errors.New("artifact direct inputs must not be empty") + } + for index, input := range material.DirectInputs { + if strings.TrimSpace(input.Role) == "" { + return Key{}, fmt.Errorf("artifact direct input %d role must not be empty", index) + } + if !isSHA256(input.Digest) { + return Key{}, fmt.Errorf("artifact direct input %d digest must be 64 lowercase hex characters", index) + } + } + if material.CanonicalizerVersion == "" { + return Key{}, errors.New("artifact canonicalizer version must not be empty") + } + if material.OutputSchemaVersion == "" { + return Key{}, errors.New("artifact output schema version must not be empty") + } + if err := validateEndpointIdentity(material.Processor.EndpointIdentity); err != nil { + return Key{}, err + } + if err := RejectSecretFields(material); err != nil { + return Key{}, err + } + + processorCanonical, err := CanonicalJSON(material.Processor) + if err != nil { + return Key{}, fmt.Errorf("canonicalize artifact processor identity: %w", err) + } + canonical, err := CanonicalJSON(material) + if err != nil { + return Key{}, fmt.Errorf("canonicalize artifact key material: %w", err) + } + processorDigest := SHA256Hex(processorCanonical) + lookup.ArtifactKey = SHA256Hex(canonical) + return Key{ + Lookup: lookup, + ProcessorDigest: processorDigest, + OutputSchema: material.OutputSchemaVersion, + Canonical: canonical, + }, nil +} + +// CanonicalJSON returns stable JSON with sorted object keys, preserved array +// order, no insignificant whitespace and normalized finite numbers. +func CanonicalJSON(value any) ([]byte, error) { + raw, err := json.Marshal(value) + if err != nil { + return nil, err + } + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + var decoded any + if err := decoder.Decode(&decoded); err != nil { + return nil, err + } + var trailing any + if err := decoder.Decode(&trailing); err == nil { + return nil, errors.New("canonical JSON contains a trailing value") + } + + var result bytes.Buffer + if err := writeCanonicalJSON(&result, decoded); err != nil { + return nil, err + } + return result.Bytes(), nil +} + +func writeCanonicalJSON(result *bytes.Buffer, value any) error { + switch typed := value.(type) { + case nil: + result.WriteString("null") + case bool: + if typed { + result.WriteString("true") + } else { + result.WriteString("false") + } + case string: + encoded, _ := json.Marshal(typed) + result.Write(encoded) + case json.Number: + number, err := canonicalNumber(typed.String()) + if err != nil { + return err + } + result.WriteString(number) + case []any: + result.WriteByte('[') + for index, item := range typed { + if index > 0 { + result.WriteByte(',') + } + if err := writeCanonicalJSON(result, item); err != nil { + return err + } + } + result.WriteByte(']') + case map[string]any: + keys := make([]string, 0, len(typed)) + for key := range typed { + keys = append(keys, key) + } + sort.Strings(keys) + result.WriteByte('{') + for index, key := range keys { + if index > 0 { + result.WriteByte(',') + } + encoded, _ := json.Marshal(key) + result.Write(encoded) + result.WriteByte(':') + if err := writeCanonicalJSON(result, typed[key]); err != nil { + return err + } + } + result.WriteByte('}') + default: + return fmt.Errorf("unsupported canonical JSON type %T", value) + } + return nil +} + +func canonicalNumber(raw string) (string, error) { + if !strings.ContainsAny(raw, ".eE") { + if raw == "-0" { + return "0", nil + } + return raw, nil + } + value, err := strconv.ParseFloat(raw, 64) + if err != nil || math.IsNaN(value) || math.IsInf(value, 0) { + return "", fmt.Errorf("invalid canonical JSON number %q", raw) + } + encoded, err := json.Marshal(value) + if err != nil { + return "", err + } + return string(encoded), nil +} + +// RejectSecretFields prevents credentials from influencing cache identity. +func RejectSecretFields(value any) error { + raw, err := json.Marshal(value) + if err != nil { + return err + } + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.UseNumber() + var decoded any + if err := decoder.Decode(&decoded); err != nil { + return err + } + return walkSecretFields(decoded, "") +} + +func walkSecretFields(value any, path string) error { + switch typed := value.(type) { + case []any: + for index, item := range typed { + if err := walkSecretFields(item, fmt.Sprintf("%s[%d]", path, index)); err != nil { + return err + } + } + case map[string]any: + for key, item := range typed { + normalized := strings.NewReplacer("_", "", "-", "", ".", "").Replace(strings.ToLower(key)) + if _, secret := secretFieldNames[normalized]; secret { + return fmt.Errorf("secret field %q must not enter artifact key material", joinJSONPath(path, key)) + } + if err := walkSecretFields(item, joinJSONPath(path, key)); err != nil { + return err + } + } + } + return nil +} + +var secretFieldNames = map[string]struct{}{ + "apikey": {}, + "authorization": {}, + "accesstoken": {}, + "refreshtoken": {}, + "clientsecret": {}, + "cookie": {}, + "password": {}, + "secret": {}, + "signature": {}, + "xamzsignature": {}, +} + +func joinJSONPath(parent, child string) string { + if parent == "" { + return child + } + return parent + "." + child +} + +func validateEndpointIdentity(raw string) error { + if raw == "" { + return nil + } + if strings.TrimSpace(raw) != raw { + return errors.New("artifact endpoint identity must not have surrounding whitespace") + } + if !strings.Contains(raw, "://") { + return nil + } + endpoint, err := url.Parse(raw) + if err != nil || endpoint.Scheme == "" || endpoint.Host == "" { + return errors.New("artifact endpoint identity is not a valid absolute URL") + } + if endpoint.User != nil || endpoint.RawQuery != "" || endpoint.Fragment != "" { + return errors.New("artifact endpoint identity must not contain credentials, query parameters or fragments") + } + return nil +} + +func isSHA256(value string) bool { + if len(value) != sha256.Size*2 { + return false + } + decoded, err := hex.DecodeString(value) + return err == nil && len(decoded) == sha256.Size && strings.ToLower(value) == value +} + +func SHA256Hex(value []byte) string { + digest := sha256.Sum256(value) + return hex.EncodeToString(digest[:]) +} diff --git a/internal/artifact/key_test.go b/internal/artifact/key_test.go new file mode 100644 index 0000000000..cd17fee2ab --- /dev/null +++ b/internal/artifact/key_test.go @@ -0,0 +1,133 @@ +package artifact + +import ( + "encoding/json" + "math" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func testKeyMaterial() KeyMaterial { + return KeyMaterial{ + KeyVersion: 1, + Stage: "summary", + DirectInputs: []DirectInput{ + {Role: "chunk", Digest: SHA256Hex([]byte("one"))}, + {Role: "chunk", Digest: SHA256Hex([]byte("two"))}, + }, + Processor: ProcessorIdentity{ + ModelID: "model-1", + ModelName: "chat-model", + Source: "openai", + Provider: "compatible", + EndpointIdentity: "https://example.com/v1", + Revision: "weights-v1", + Parameters: map[string]any{ + "temperature": 0.3, + }, + }, + RenderedRequest: []map[string]string{ + {"role": "system", "content": "Summarize exactly."}, + {"role": "user", "content": " alpha \r\n beta "}, + }, + Options: map[string]any{ + "thinking": false, + }, + CanonicalizerVersion: CanonicalJSONVersion, + OutputSchemaVersion: "summary.v1", + } +} + +func TestCanonicalJSONStableMapOrderAndNumbers(t *testing.T) { + left := map[string]any{"z": 1.0, "a": map[string]any{"b": true, "a": -0.0}} + right := json.RawMessage(`{"a":{"a":0,"b":true},"z":1}`) + + leftJSON, err := CanonicalJSON(left) + require.NoError(t, err) + rightJSON, err := CanonicalJSON(right) + require.NoError(t, err) + + assert.Equal(t, `{"a":{"a":0,"b":true},"z":1}`, string(leftJSON)) + assert.Equal(t, leftJSON, rightJSON) +} + +func TestCanonicalJSONRejectsNonFiniteNumbers(t *testing.T) { + _, err := CanonicalJSON(map[string]any{"value": math.NaN()}) + require.Error(t, err) +} + +func TestBuildKeyExactInvalidationAndOrderedInputs(t *testing.T) { + base := testKeyMaterial() + original, err := BuildKey(7, base) + require.NoError(t, err) + + unchanged, err := BuildKey(7, testKeyMaterial()) + require.NoError(t, err) + assert.Equal(t, original.Lookup.ArtifactKey, unchanged.Lookup.ArtifactKey) + assert.Equal(t, original.ProcessorDigest, unchanged.ProcessorDigest) + + reordered := testKeyMaterial() + reordered.DirectInputs[0], reordered.DirectInputs[1] = reordered.DirectInputs[1], reordered.DirectInputs[0] + reorderedKey, err := BuildKey(7, reordered) + require.NoError(t, err) + assert.NotEqual(t, original.Lookup.ArtifactKey, reorderedKey.Lookup.ArtifactKey) + + requestChanged := testKeyMaterial() + requestChanged.RenderedRequest.([]map[string]string)[1]["content"] = "alpha\nbeta" + requestKey, err := BuildKey(7, requestChanged) + require.NoError(t, err) + assert.NotEqual(t, original.Lookup.ArtifactKey, requestKey.Lookup.ArtifactKey) + + optionsChanged := testKeyMaterial() + optionsChanged.Options.(map[string]any)["thinking"] = true + optionsKey, err := BuildKey(7, optionsChanged) + require.NoError(t, err) + assert.NotEqual(t, original.Lookup.ArtifactKey, optionsKey.Lookup.ArtifactKey) + + schemaChanged := testKeyMaterial() + schemaChanged.OutputSchemaVersion = "summary.v2" + schemaKey, err := BuildKey(7, schemaChanged) + require.NoError(t, err) + assert.NotEqual(t, original.Lookup.ArtifactKey, schemaKey.Lookup.ArtifactKey) +} + +func TestBuildKeyTenantBoundaryIsExplicit(t *testing.T) { + material := testKeyMaterial() + first, err := BuildKey(7, material) + require.NoError(t, err) + second, err := BuildKey(8, material) + require.NoError(t, err) + + assert.Equal(t, first.Lookup.ArtifactKey, second.Lookup.ArtifactKey) + assert.NotEqual(t, first.Lookup, second.Lookup) +} + +func TestBuildKeyRejectsSecretsAndSignedEndpoints(t *testing.T) { + withSecret := testKeyMaterial() + withSecret.Options = map[string]any{"api_key": "do-not-hash"} + _, err := BuildKey(7, withSecret) + require.ErrorContains(t, err, "secret field") + + withSignedEndpoint := testKeyMaterial() + withSignedEndpoint.Processor.EndpointIdentity = "https://example.com/v1?signature=secret" + _, err = BuildKey(7, withSignedEndpoint) + require.ErrorContains(t, err, "query parameters") +} + +func TestDownstreamKeyDependsOnUpstreamOutputDigestOnly(t *testing.T) { + first := testKeyMaterial() + first.DirectInputs = []DirectInput{{Role: "parse", Digest: SHA256Hex([]byte("canonical-output"))}} + first.RenderedRequest = "downstream prompt" + firstKey, err := BuildKey(7, first) + require.NoError(t, err) + + second := testKeyMaterial() + second.DirectInputs = []DirectInput{{Role: "parse", Digest: SHA256Hex([]byte("canonical-output"))}} + second.RenderedRequest = "downstream prompt" + secondKey, err := BuildKey(7, second) + require.NoError(t, err) + + assert.Equal(t, firstKey.Lookup.ArtifactKey, secondKey.Lookup.ArtifactKey) +} diff --git a/internal/artifact/lease.go b/internal/artifact/lease.go new file mode 100644 index 0000000000..80d5671d91 --- /dev/null +++ b/internal/artifact/lease.go @@ -0,0 +1,132 @@ +package artifact + +import ( + "context" + "crypto/rand" + "encoding/hex" + "errors" + "time" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/redis/go-redis/v9" +) + +const ( + artifactLeaseTTL = 5 * time.Minute + artifactLeaseKeyPrefix = "weknora:artifact:lease:" +) + +// Lease suppresses duplicate provider work across processes. It is only an +// optimization: immutable database put-if-absent remains the correctness +// boundary, and lease failures make Runtime compute normally. +type Lease interface { + TryAcquire( + ctx context.Context, + key types.ProcessingArtifactLookup, + ) (handle LeaseHandle, acquired bool, err error) +} + +type LeaseHandle interface { + Release() +} + +type redisLease struct { + client redis.UniversalClient + ttl time.Duration +} + +func NewRedisLease(client *redis.Client) Lease { + if client == nil { + return nil + } + return &redisLease{client: client, ttl: artifactLeaseTTL} +} + +func (l *redisLease) TryAcquire( + ctx context.Context, + key types.ProcessingArtifactLookup, +) (LeaseHandle, bool, error) { + if l == nil || l.client == nil { + return nil, false, errors.New("artifact Redis lease is not configured") + } + tokenBytes := make([]byte, 16) + if _, err := rand.Read(tokenBytes); err != nil { + return nil, false, err + } + token := hex.EncodeToString(tokenBytes) + redisKey := artifactLeaseKeyPrefix + SHA256Hex([]byte(singleflightKey(key))) + acquired, err := l.client.SetNX(ctx, redisKey, token, l.ttl).Result() + if err != nil || !acquired { + return nil, acquired, err + } + handle := &redisLeaseHandle{ + client: l.client, + key: redisKey, + token: token, + ttl: l.ttl, + stop: make(chan struct{}), + } + go handle.renew() + return handle, true, nil +} + +type redisLeaseHandle struct { + client redis.UniversalClient + key string + token string + ttl time.Duration + stop chan struct{} +} + +var releaseLeaseScript = redis.NewScript(` +if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("del", KEYS[1]) +end +return 0 +`) + +var renewLeaseScript = redis.NewScript(` +if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("pexpire", KEYS[1], ARGV[2]) +end +return 0 +`) + +func (h *redisLeaseHandle) Release() { + if h == nil { + return + } + select { + case <-h.stop: + return + default: + close(h.stop) + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + _ = releaseLeaseScript.Run(ctx, h.client, []string{h.key}, h.token).Err() +} + +func (h *redisLeaseHandle) renew() { + ticker := time.NewTicker(h.ttl / 3) + defer ticker.Stop() + for { + select { + case <-h.stop: + return + case <-ticker.C: + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + result, err := renewLeaseScript.Run( + ctx, + h.client, + []string{h.key}, + h.token, + h.ttl.Milliseconds(), + ).Int() + cancel() + if err != nil || result == 0 { + return + } + } + } +} diff --git a/internal/artifact/reconcile.go b/internal/artifact/reconcile.go new file mode 100644 index 0000000000..67ef62efd0 --- /dev/null +++ b/internal/artifact/reconcile.go @@ -0,0 +1,97 @@ +package artifact + +import ( + "errors" + "fmt" +) + +type EntityState struct { + ID string + MatchKey string + ContentDigest string + ArtifactDigest string + MetadataDigest string +} + +type EntityPair struct { + Desired EntityState + Live EntityState +} + +type ReconcilePlan struct { + Kept []EntityPair + MetadataOnly []EntityPair + Added []EntityState + Changed []EntityPair + Stale []EntityState +} + +// PlanDesiredState computes a non-destructive diff. Callers materialize Added +// and Changed first, publish only while fenced, then clean exactly Stale last. +func PlanDesiredState(desired, live []EntityState) (ReconcilePlan, error) { + plan := ReconcilePlan{} + liveByMatch := make(map[string]EntityState, len(live)) + for _, entity := range live { + if err := validateEntityState(entity); err != nil { + return ReconcilePlan{}, fmt.Errorf("invalid live entity: %w", err) + } + if _, exists := liveByMatch[entity.MatchKey]; exists { + return ReconcilePlan{}, fmt.Errorf("duplicate live entity match key %q", entity.MatchKey) + } + liveByMatch[entity.MatchKey] = entity + } + + desiredMatches := make(map[string]struct{}, len(desired)) + for _, entity := range desired { + if err := validateEntityState(entity); err != nil { + return ReconcilePlan{}, fmt.Errorf("invalid desired entity: %w", err) + } + if _, exists := desiredMatches[entity.MatchKey]; exists { + return ReconcilePlan{}, fmt.Errorf("duplicate desired entity match key %q", entity.MatchKey) + } + desiredMatches[entity.MatchKey] = struct{}{} + + current, found := liveByMatch[entity.MatchKey] + if !found { + plan.Added = append(plan.Added, entity) + continue + } + pair := EntityPair{Desired: entity, Live: current} + switch { + case entity.ID != current.ID || + entity.ContentDigest != current.ContentDigest || + entity.ArtifactDigest != current.ArtifactDigest: + plan.Changed = append(plan.Changed, pair) + case entity.MetadataDigest != current.MetadataDigest: + plan.MetadataOnly = append(plan.MetadataOnly, pair) + default: + plan.Kept = append(plan.Kept, pair) + } + } + + for _, entity := range live { + if _, desired := desiredMatches[entity.MatchKey]; !desired { + plan.Stale = append(plan.Stale, entity) + } + } + return plan, nil +} + +func validateEntityState(entity EntityState) error { + if entity.ID == "" { + return errors.New("entity ID must not be empty") + } + if entity.MatchKey == "" { + return errors.New("entity match key must not be empty") + } + if !isSHA256(entity.ContentDigest) { + return errors.New("entity content digest must be a SHA-256") + } + if entity.ArtifactDigest != "" && !isSHA256(entity.ArtifactDigest) { + return errors.New("entity artifact digest must be empty or a SHA-256") + } + if entity.MetadataDigest != "" && !isSHA256(entity.MetadataDigest) { + return errors.New("entity metadata digest must be empty or a SHA-256") + } + return nil +} diff --git a/internal/artifact/reconcile_test.go b/internal/artifact/reconcile_test.go new file mode 100644 index 0000000000..df1ba25a71 --- /dev/null +++ b/internal/artifact/reconcile_test.go @@ -0,0 +1,48 @@ +package artifact + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func entityState(id, match, content, artifactValue, metadata string) EntityState { + return EntityState{ + ID: id, + MatchKey: match, + ContentDigest: SHA256Hex([]byte(content)), + ArtifactDigest: SHA256Hex([]byte(artifactValue)), + MetadataDigest: SHA256Hex([]byte(metadata)), + } +} + +func TestPlanDesiredStateClassifiesExactChanges(t *testing.T) { + kept := entityState("kept", "m1", "a", "p1", "meta") + metadataLive := entityState("metadata", "m2", "b", "p2", "old") + metadataDesired := entityState("metadata", "m2", "b", "p2", "new") + changedLive := entityState("old-id", "m3", "old", "p3", "meta") + changedDesired := entityState("new-id", "m3", "new", "p4", "meta") + stale := entityState("stale", "m4", "d", "p5", "meta") + added := entityState("added", "m5", "e", "p6", "meta") + + plan, err := PlanDesiredState( + []EntityState{kept, metadataDesired, changedDesired, added}, + []EntityState{kept, metadataLive, changedLive, stale}, + ) + require.NoError(t, err) + + require.Len(t, plan.Kept, 1) + require.Len(t, plan.MetadataOnly, 1) + require.Len(t, plan.Changed, 1) + require.Len(t, plan.Added, 1) + require.Len(t, plan.Stale, 1) + assert.Equal(t, "stale", plan.Stale[0].ID) +} + +func TestPlanDesiredStateRejectsAmbiguousMatches(t *testing.T) { + first := entityState("one", "duplicate", "a", "p", "m") + second := entityState("two", "duplicate", "b", "p", "m") + _, err := PlanDesiredState(nil, []EntityState{first, second}) + require.ErrorContains(t, err, "duplicate live") +} diff --git a/internal/artifact/runtime.go b/internal/artifact/runtime.go new file mode 100644 index 0000000000..5bfb8ff61e --- /dev/null +++ b/internal/artifact/runtime.go @@ -0,0 +1,415 @@ +package artifact + +import ( + "context" + "errors" + "fmt" + "strconv" + "time" + + "github.com/Tencent/WeKnora/internal/types" + "golang.org/x/sync/singleflight" +) + +type EventKind string + +const ( + EventHit EventKind = "hit" + EventMiss EventKind = "miss" + EventCorrupt EventKind = "corrupt" + EventStoreFailure EventKind = "store_failure" + EventStored EventKind = "stored" + EventLostRace EventKind = "lost_race" +) + +type Event struct { + Kind EventKind + Lookup types.ProcessingArtifactLookup + Err error +} + +type Observer func(Event) + +// Repository is declared in this leaf package to keep artifact mechanics +// independent from the broad application interfaces package (which itself +// references model packages that consume artifacts). +type Repository interface { + Get( + ctx context.Context, + key types.ProcessingArtifactLookup, + ) (*types.ProcessingArtifact, error) + BatchGet( + ctx context.Context, + keys []types.ProcessingArtifactLookup, + ) (map[types.ProcessingArtifactLookup]*types.ProcessingArtifact, error) + PutIfAbsent( + ctx context.Context, + candidate *types.ProcessingArtifact, + ) (winner *types.ProcessingArtifact, created bool, err error) + PutManyIfAbsent( + ctx context.Context, + candidates []*types.ProcessingArtifact, + ) (map[types.ProcessingArtifactLookup]*types.ProcessingArtifact, error) + DeleteCorrupt( + ctx context.Context, + key types.ProcessingArtifactLookup, + observedChecksum string, + ) error + TouchHits(ctx context.Context, keys []types.ProcessingArtifactLookup) error +} + +// Runtime turns repository failures into misses while preserving provider +// failures. Database uniqueness remains the correctness boundary across +// processes; singleflight only suppresses duplicate work within one process. +type Runtime struct { + repository Repository + observer Observer + lease Lease + group singleflight.Group +} + +func NewRuntime(repository Repository, observer Observer) *Runtime { + return &Runtime{repository: repository, observer: observer} +} + +func (r *Runtime) ConfigureLease(lease Lease) { + if r != nil { + r.lease = lease + } +} + +type Expected struct { + Key Key + Codec string + Validate func([]byte) error + Cacheable func([]byte) bool +} + +type Value struct { + Payload []byte + OutputDigest string + CacheHit bool +} + +type Candidate struct { + Expected Expected + Payload []byte +} + +func (r *Runtime) Load(ctx context.Context, expected Expected) (Value, bool) { + values := r.BatchLoad(ctx, []Expected{expected}) + value, ok := values[expected.Key.Lookup] + return value, ok +} + +// BatchLoad performs one manifest query per database-sized batch, validates +// every result independently, and evicts only the corrupt row observed. +func (r *Runtime) BatchLoad( + ctx context.Context, + expected []Expected, +) map[types.ProcessingArtifactLookup]Value { + result := make(map[types.ProcessingArtifactLookup]Value, len(expected)) + if r == nil || r.repository == nil || len(expected) == 0 { + return result + } + + keys := make([]types.ProcessingArtifactLookup, 0, len(expected)) + byKey := make(map[types.ProcessingArtifactLookup]Expected, len(expected)) + for _, item := range expected { + keys = append(keys, item.Key.Lookup) + byKey[item.Key.Lookup] = item + } + artifacts, err := r.repository.BatchGet(ctx, keys) + if err != nil { + for _, key := range keys { + r.emit(Event{Kind: EventStoreFailure, Lookup: key, Err: err}) + } + return result + } + + hits := make([]types.ProcessingArtifactLookup, 0, len(artifacts)) + for _, key := range keys { + item := byKey[key] + manifest, found := artifacts[key] + if !found { + r.emit(Event{Kind: EventMiss, Lookup: key}) + continue + } + payload, decodeErr := DecodeInline(manifest, key, item.Key.OutputSchema, item.Codec) + if decodeErr == nil && item.Validate != nil { + decodeErr = item.Validate(payload) + } + if decodeErr != nil { + r.emit(Event{Kind: EventCorrupt, Lookup: key, Err: decodeErr}) + if deleteErr := r.repository.DeleteCorrupt(ctx, key, manifest.PayloadChecksum); deleteErr != nil { + r.emit(Event{Kind: EventStoreFailure, Lookup: key, Err: deleteErr}) + } + continue + } + result[key] = Value{ + Payload: payload, + OutputDigest: manifest.OutputDigest, + CacheHit: true, + } + hits = append(hits, key) + r.emit(Event{Kind: EventHit, Lookup: key}) + } + if len(hits) > 0 { + if err := r.repository.TouchHits(ctx, hits); err != nil { + for _, key := range hits { + r.emit(Event{Kind: EventStoreFailure, Lookup: key, Err: err}) + } + } + } + return result +} + +// LoadOrCompute caches only validated successful output. Artifact read/write +// failures are fail-open; compute errors are returned unchanged. +func (r *Runtime) LoadOrCompute( + ctx context.Context, + expected Expected, + compute func(context.Context) ([]byte, error), +) (Value, error) { + if compute == nil { + return Value{}, errors.New("artifact compute function must not be nil") + } + if value, hit := r.Load(ctx, expected); hit { + return value, nil + } + if r == nil { + payload, err := compute(ctx) + return uncachedValue(payload, expected.Validate, err) + } + + result := r.group.DoChan(singleflightKey(expected.Key.Lookup), func() (any, error) { + if value, hit := r.Load(ctx, expected); hit { + return value, nil + } + if r.lease != nil { + for { + handle, acquired, leaseErr := r.lease.TryAcquire(ctx, expected.Key.Lookup) + if leaseErr != nil { + r.emit(Event{ + Kind: EventStoreFailure, + Lookup: expected.Key.Lookup, + Err: leaseErr, + }) + break + } + if acquired { + defer handle.Release() + // The winner may have committed between our last read and + // lease acquisition. + if value, hit := r.Load(ctx, expected); hit { + return value, nil + } + break + } + timer := time.NewTimer(100 * time.Millisecond) + select { + case <-ctx.Done(): + if !timer.Stop() { + <-timer.C + } + return Value{}, ctx.Err() + case <-timer.C: + } + if value, hit := r.Load(ctx, expected); hit { + return value, nil + } + } + } + payload, err := compute(ctx) + if err != nil { + return Value{}, err + } + if expected.Validate != nil { + if err := expected.Validate(payload); err != nil { + return Value{}, fmt.Errorf("validate processing artifact output: %w", err) + } + } + if expected.Cacheable != nil && !expected.Cacheable(payload) { + return uncachedValue(payload, expected.Validate, nil) + } + candidate, err := NewInlineArtifact(expected.Key, expected.Codec, payload) + if err != nil { + return Value{}, err + } + return r.freeze(ctx, expected, candidate), nil + }) + + select { + case <-ctx.Done(): + return Value{}, ctx.Err() + case completed := <-result: + if completed.Err != nil { + return Value{}, completed.Err + } + return completed.Val.(Value), nil + } +} + +// BatchFreeze performs one immutable batch insert and then returns the database +// winners. Invalid winners are conditionally evicted and the caller's validated +// candidate remains usable, preserving fail-open processing. +func (r *Runtime) BatchFreeze( + ctx context.Context, + candidates []Candidate, +) map[types.ProcessingArtifactLookup]Value { + result := make(map[types.ProcessingArtifactLookup]Value, len(candidates)) + manifests := make([]*types.ProcessingArtifact, 0, len(candidates)) + byKey := make(map[types.ProcessingArtifactLookup]Candidate, len(candidates)) + for _, candidate := range candidates { + if candidate.Expected.Cacheable != nil && !candidate.Expected.Cacheable(candidate.Payload) { + continue + } + if candidate.Expected.Validate != nil { + if err := candidate.Expected.Validate(candidate.Payload); err != nil { + continue + } + } + manifest, err := NewInlineArtifact( + candidate.Expected.Key, + candidate.Expected.Codec, + candidate.Payload, + ) + if err != nil { + continue + } + key := manifest.Lookup() + manifests = append(manifests, manifest) + byKey[key] = candidate + result[key] = Value{ + Payload: append([]byte(nil), manifest.Payload...), + OutputDigest: manifest.OutputDigest, + } + } + if r == nil || r.repository == nil || len(manifests) == 0 { + return result + } + + winners, err := r.repository.PutManyIfAbsent(ctx, manifests) + if err != nil { + for _, manifest := range manifests { + r.emit(Event{Kind: EventStoreFailure, Lookup: manifest.Lookup(), Err: err}) + } + return result + } + for _, manifest := range manifests { + key := manifest.Lookup() + winner := winners[key] + candidate := byKey[key] + if winner == nil { + r.emit(Event{ + Kind: EventStoreFailure, + Lookup: key, + Err: errors.New("artifact repository omitted an inserted winner"), + }) + continue + } + payload, decodeErr := DecodeInline( + winner, + key, + candidate.Expected.Key.OutputSchema, + candidate.Expected.Codec, + ) + if decodeErr == nil && candidate.Expected.Validate != nil { + decodeErr = candidate.Expected.Validate(payload) + } + if decodeErr != nil { + r.emit(Event{Kind: EventCorrupt, Lookup: key, Err: decodeErr}) + if winner != nil { + if deleteErr := r.repository.DeleteCorrupt(ctx, key, winner.PayloadChecksum); deleteErr != nil { + r.emit(Event{Kind: EventStoreFailure, Lookup: key, Err: deleteErr}) + } + } + continue + } + result[key] = Value{ + Payload: payload, + OutputDigest: winner.OutputDigest, + CacheHit: winner.ID != manifest.ID, + } + } + return result +} + +func (r *Runtime) freeze( + ctx context.Context, + expected Expected, + candidate *types.ProcessingArtifact, +) Value { + fallback := Value{ + Payload: append([]byte(nil), candidate.Payload...), + OutputDigest: candidate.OutputDigest, + } + if r.repository == nil { + return fallback + } + winner, created, err := r.repository.PutIfAbsent(ctx, candidate) + if err != nil { + r.emit(Event{Kind: EventStoreFailure, Lookup: candidate.Lookup(), Err: err}) + return fallback + } + if winner == nil { + r.emit(Event{ + Kind: EventStoreFailure, + Lookup: candidate.Lookup(), + Err: errors.New("artifact repository returned a nil winner"), + }) + return fallback + } + payload, err := DecodeInline( + winner, + expected.Key.Lookup, + expected.Key.OutputSchema, + expected.Codec, + ) + if err == nil && expected.Validate != nil { + err = expected.Validate(payload) + } + if err != nil { + r.emit(Event{Kind: EventCorrupt, Lookup: candidate.Lookup(), Err: err}) + if deleteErr := r.repository.DeleteCorrupt(ctx, winner.Lookup(), winner.PayloadChecksum); deleteErr != nil { + r.emit(Event{Kind: EventStoreFailure, Lookup: winner.Lookup(), Err: deleteErr}) + } + return fallback + } + if created { + r.emit(Event{Kind: EventStored, Lookup: candidate.Lookup()}) + } else { + r.emit(Event{Kind: EventLostRace, Lookup: candidate.Lookup()}) + } + return Value{ + Payload: payload, + OutputDigest: winner.OutputDigest, + CacheHit: !created, + } +} + +func uncachedValue(payload []byte, validate func([]byte) error, err error) (Value, error) { + if err != nil { + return Value{}, err + } + if validate != nil { + if err := validate(payload); err != nil { + return Value{}, fmt.Errorf("validate processing output: %w", err) + } + } + frozen := append([]byte(nil), payload...) + return Value{Payload: frozen, OutputDigest: SHA256Hex(frozen)}, nil +} + +func singleflightKey(key types.ProcessingArtifactLookup) string { + return strconv.FormatUint(key.TenantID, 10) + "\x00" + + key.Stage + "\x00" + + strconv.FormatUint(uint64(key.KeyVersion), 10) + "\x00" + + key.ArtifactKey +} + +func (r *Runtime) emit(event Event) { + if r != nil && r.observer != nil { + r.observer(event) + } +} diff --git a/internal/artifact/runtime_test.go b/internal/artifact/runtime_test.go new file mode 100644 index 0000000000..1208d065df --- /dev/null +++ b/internal/artifact/runtime_test.go @@ -0,0 +1,229 @@ +package artifact + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type memoryArtifactRepository struct { + mu sync.Mutex + values map[types.ProcessingArtifactLookup]*types.ProcessingArtifact + getErr error + putErr error + deleted int +} + +func newMemoryArtifactRepository() *memoryArtifactRepository { + return &memoryArtifactRepository{ + values: make(map[types.ProcessingArtifactLookup]*types.ProcessingArtifact), + } +} + +func (r *memoryArtifactRepository) Get( + _ context.Context, + key types.ProcessingArtifactLookup, +) (*types.ProcessingArtifact, error) { + values, err := r.BatchGet(context.Background(), []types.ProcessingArtifactLookup{key}) + if err != nil { + return nil, err + } + value, ok := values[key] + if !ok { + return nil, types.ErrProcessingArtifactNotFound + } + return value, nil +} + +func (r *memoryArtifactRepository) BatchGet( + _ context.Context, + keys []types.ProcessingArtifactLookup, +) (map[types.ProcessingArtifactLookup]*types.ProcessingArtifact, error) { + r.mu.Lock() + defer r.mu.Unlock() + if r.getErr != nil { + return nil, r.getErr + } + result := make(map[types.ProcessingArtifactLookup]*types.ProcessingArtifact) + for _, key := range keys { + if value := r.values[key]; value != nil { + copy := *value + copy.Payload = append([]byte(nil), value.Payload...) + result[key] = © + } + } + return result, nil +} + +func (r *memoryArtifactRepository) PutIfAbsent( + _ context.Context, + candidate *types.ProcessingArtifact, +) (*types.ProcessingArtifact, bool, error) { + r.mu.Lock() + defer r.mu.Unlock() + if r.putErr != nil { + return nil, false, r.putErr + } + if winner := r.values[candidate.Lookup()]; winner != nil { + copy := *winner + copy.Payload = append([]byte(nil), winner.Payload...) + return ©, false, nil + } + copy := *candidate + copy.Payload = append([]byte(nil), candidate.Payload...) + r.values[candidate.Lookup()] = © + return candidate, true, nil +} + +func (r *memoryArtifactRepository) PutManyIfAbsent( + ctx context.Context, + candidates []*types.ProcessingArtifact, +) (map[types.ProcessingArtifactLookup]*types.ProcessingArtifact, error) { + for _, candidate := range candidates { + if _, _, err := r.PutIfAbsent(ctx, candidate); err != nil { + return nil, err + } + } + keys := make([]types.ProcessingArtifactLookup, 0, len(candidates)) + for _, candidate := range candidates { + keys = append(keys, candidate.Lookup()) + } + return r.BatchGet(ctx, keys) +} + +func (r *memoryArtifactRepository) DeleteCorrupt( + _ context.Context, + key types.ProcessingArtifactLookup, + observedChecksum string, +) error { + r.mu.Lock() + defer r.mu.Unlock() + if value := r.values[key]; value != nil && value.PayloadChecksum == observedChecksum { + delete(r.values, key) + r.deleted++ + } + return nil +} + +func (r *memoryArtifactRepository) TouchHits(context.Context, []types.ProcessingArtifactLookup) error { + return nil +} + +func testExpected(t *testing.T) Expected { + t.Helper() + key, err := BuildKey(1, testKeyMaterial()) + require.NoError(t, err) + return Expected{ + Key: key, + Codec: CodecJSONV1, + Validate: func(payload []byte) error { + if string(payload) != `{"value":"computed"}` { + return errors.New("unexpected payload") + } + return nil + }, + } +} + +func TestRuntimeStoreFailureIsFailOpen(t *testing.T) { + repository := newMemoryArtifactRepository() + repository.getErr = errors.New("database down") + runtime := NewRuntime(repository, nil) + + value, err := runtime.LoadOrCompute(context.Background(), testExpected(t), func(context.Context) ([]byte, error) { + return []byte(`{"value":"computed"}`), nil + }) + require.NoError(t, err) + assert.Equal(t, `{"value":"computed"}`, string(value.Payload)) + assert.False(t, value.CacheHit) +} + +func TestRuntimeCorruptEntrySelfHeals(t *testing.T) { + repository := newMemoryArtifactRepository() + expected := testExpected(t) + corrupt, err := NewInlineArtifact(expected.Key, CodecJSONV1, []byte(`{"value":"computed"}`)) + require.NoError(t, err) + corrupt.PayloadChecksum = SHA256Hex([]byte("wrong")) + repository.values[expected.Key.Lookup] = corrupt + runtime := NewRuntime(repository, nil) + + value, err := runtime.LoadOrCompute(context.Background(), expected, func(context.Context) ([]byte, error) { + return []byte(`{"value":"computed"}`), nil + }) + require.NoError(t, err) + assert.Equal(t, `{"value":"computed"}`, string(value.Payload)) + assert.Equal(t, 1, repository.deleted) +} + +func TestRuntimeSingleflightComputesOnce(t *testing.T) { + repository := newMemoryArtifactRepository() + runtime := NewRuntime(repository, nil) + expected := testExpected(t) + var calls atomic.Int32 + + const workers = 16 + start := make(chan struct{}) + results := make(chan error, workers) + for index := 0; index < workers; index++ { + go func() { + <-start + _, err := runtime.LoadOrCompute(context.Background(), expected, func(context.Context) ([]byte, error) { + calls.Add(1) + time.Sleep(10 * time.Millisecond) + return []byte(`{"value":"computed"}`), nil + }) + results <- err + }() + } + close(start) + for index := 0; index < workers; index++ { + require.NoError(t, <-results) + } + assert.Equal(t, int32(1), calls.Load()) +} + +func TestRuntimeRedisLeaseSuppressesCrossProcessDuplicateCompute(t *testing.T) { + server := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: server.Addr()}) + t.Cleanup(func() { _ = client.Close() }) + + repository := newMemoryArtifactRepository() + firstRuntime := NewRuntime(repository, nil) + secondRuntime := NewRuntime(repository, nil) + firstRuntime.ConfigureLease(NewRedisLease(client)) + secondRuntime.ConfigureLease(NewRedisLease(client)) + expected := testExpected(t) + var calls atomic.Int32 + + start := make(chan struct{}) + results := make(chan error, 2) + run := func(runtime *Runtime) { + <-start + _, err := runtime.LoadOrCompute( + context.Background(), + expected, + func(context.Context) ([]byte, error) { + calls.Add(1) + time.Sleep(75 * time.Millisecond) + return []byte(`{"value":"computed"}`), nil + }, + ) + results <- err + } + go run(firstRuntime) + go run(secondRuntime) + close(start) + + require.NoError(t, <-results) + require.NoError(t, <-results) + assert.Equal(t, int32(1), calls.Load()) +} diff --git a/internal/container/container.go b/internal/container/container.go index 8a9265f5ee..6245106fb4 100644 --- a/internal/container/container.go +++ b/internal/container/container.go @@ -50,6 +50,7 @@ import ( chatpipeline "github.com/Tencent/WeKnora/internal/application/service/chat_pipeline" "github.com/Tencent/WeKnora/internal/application/service/file" "github.com/Tencent/WeKnora/internal/application/service/retriever" + "github.com/Tencent/WeKnora/internal/artifact" "github.com/Tencent/WeKnora/internal/common" "github.com/Tencent/WeKnora/internal/config" "github.com/Tencent/WeKnora/internal/database" @@ -146,6 +147,15 @@ func BuildContainer(container *dig.Container) *dig.Container { must(container.Provide(repository.NewKnowledgeBaseRepository)) must(container.Provide(repository.NewKnowledgeRepository)) must(container.Provide(repository.NewKnowledgeSpanRepository)) + must(container.Provide(repository.NewProcessingArtifactRepository)) + must(container.Provide(func( + repo interfaces.ProcessingArtifactRepository, + redisClient *redis.Client, + ) *artifact.Runtime { + runtime := artifact.NewRuntime(repo, nil) + runtime.ConfigureLease(artifact.NewRedisLease(redisClient)) + return runtime + })) must(container.Provide(repository.NewChunkRepository)) must(container.Provide(repository.NewKnowledgeTagRepository)) must(container.Provide(repository.NewSessionRepository)) @@ -196,6 +206,7 @@ func BuildContainer(container *dig.Container) *dig.Container { must(container.Provide(service.NewKnowledgeTagService)) must(container.Provide(embedding.NewBatchEmbedder)) must(container.Provide(service.NewModelService)) + must(container.Invoke(service.ConfigureModelArtifactCache)) must(container.Provide(service.NewDatasetService)) must(container.Provide(service.NewEvaluationService)) must(container.Provide(service.NewUserService)) diff --git a/internal/database/processing_artifact_migration_test.go b/internal/database/processing_artifact_migration_test.go new file mode 100644 index 0000000000..65417432da --- /dev/null +++ b/internal/database/processing_artifact_migration_test.go @@ -0,0 +1,58 @@ +package database + +import ( + "database/sql" + "os" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestSQLiteProcessingArtifactMigrationAndTenantBoundary(t *testing.T) { + db, err := sql.Open("sqlite3", ":memory:") + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + + up, err := os.ReadFile("../../migrations/sqlite/000001_processing_artifacts.up.sql") + require.NoError(t, err) + _, err = db.Exec(string(up)) + require.NoError(t, err) + + insert := ` + INSERT INTO processing_artifacts ( + tenant_id, stage, key_version, artifact_key, processor_digest, + output_digest, output_schema, codec, inline_payload, payload, + object_ref, payload_checksum, size_bytes + ) VALUES (?, 'embedding', 1, ?, ?, ?, 'embedding.v1', 'float32be.v1', 1, ?, '', ?, ?) + ` + hash := "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + payload := []byte{} + _, err = db.Exec(insert, 1, hash, hash, hash, payload, hash, 0) + require.NoError(t, err) + _, err = db.Exec(insert, 2, hash, hash, hash, payload, hash, 0) + require.NoError(t, err, "the same content key must be isolated by tenant") + _, err = db.Exec(insert, 1, hash, hash, hash, payload, hash, 0) + require.Error(t, err, "the tenant-scoped immutable key must be unique") + + _, err = db.Exec( + "INSERT INTO knowledge_attempt_counters (knowledge_id, last_attempt) VALUES (?, ?)", + "knowledge-1", + 1, + ) + require.NoError(t, err) + _, err = db.Exec( + "INSERT INTO knowledge_attempt_counters (knowledge_id, last_attempt) VALUES (?, ?)", + "knowledge-1", + 2, + ) + require.Error(t, err, "each knowledge must have exactly one attempt allocator row") + + down, err := os.ReadFile("../../migrations/sqlite/000001_processing_artifacts.down.sql") + require.NoError(t, err) + _, err = db.Exec(string(down)) + require.NoError(t, err) + _, err = db.Exec("SELECT 1 FROM processing_artifacts LIMIT 1") + require.Error(t, err) + _, err = db.Exec("SELECT 1 FROM knowledge_attempt_counters LIMIT 1") + require.Error(t, err) +} diff --git a/internal/models/chat/artifact_cache.go b/internal/models/chat/artifact_cache.go new file mode 100644 index 0000000000..5a2be753b0 --- /dev/null +++ b/internal/models/chat/artifact_cache.go @@ -0,0 +1,221 @@ +package chat + +import ( + "context" + "errors" + "fmt" + "net/url" + "strings" + "unicode/utf8" + + "github.com/Tencent/WeKnora/internal/artifact" + "github.com/Tencent/WeKnora/internal/types" +) + +type artifactStageContextKey struct{} + +type ArtifactStage struct { + Stage string + OutputSchema string + Validate func(string) error +} + +func WithArtifactStage(ctx context.Context, stage ArtifactStage) context.Context { + return context.WithValue(ctx, artifactStageContextKey{}, stage) +} + +type ArtifactCacheConfig struct { + TenantID uint64 + Processor artifact.ProcessorIdentity +} + +type artifactCachedChat struct { + inner Chat + runtime *artifact.Runtime + config ArtifactCacheConfig +} + +func ArtifactCacheConfigFromModel(model *types.Model, tenantID uint64) (ArtifactCacheConfig, bool) { + if model == nil || tenantID == 0 || len(model.Parameters.CustomHeaders) > 0 { + return ArtifactCacheConfig{}, false + } + endpoint := model.Parameters.BaseURL + if endpoint != "" { + parsed, err := url.Parse(endpoint) + if err != nil || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" { + return ArtifactCacheConfig{}, false + } + } + extra := make(map[string]any, len(model.Parameters.ExtraConfig)) + for key, value := range model.Parameters.ExtraConfig { + normalized := strings.NewReplacer("_", "", "-", "", ".", "").Replace(strings.ToLower(key)) + if strings.Contains(normalized, "key") || + strings.Contains(normalized, "token") || + strings.Contains(normalized, "secret") || + strings.Contains(normalized, "password") || + strings.Contains(normalized, "signature") { + return ArtifactCacheConfig{}, false + } + extra[key] = value + } + parameters := map[string]any{ + "interface_type": model.Parameters.InterfaceType, + "parameter_size": model.Parameters.ParameterSize, + "supports_vision": model.Parameters.SupportsVision, + } + if len(extra) > 0 { + parameters["extra_config"] = extra + } + return ArtifactCacheConfig{ + TenantID: tenantID, + Processor: artifact.ProcessorIdentity{ + ModelID: model.ID, + ModelName: model.Name, + Source: string(model.Source), + Provider: model.Parameters.Provider, + EndpointIdentity: endpoint, + Parameters: parameters, + }, + }, true +} + +func NewArtifactCachedChat(inner Chat, runtime *artifact.Runtime, config ArtifactCacheConfig) Chat { + if inner == nil || runtime == nil || config.TenantID == 0 { + return inner + } + return &artifactCachedChat{inner: inner, runtime: runtime, config: config} +} + +func (c *artifactCachedChat) Chat( + ctx context.Context, + messages []Message, + options *ChatOptions, +) (*types.ChatResponse, error) { + stage, ok := ctx.Value(artifactStageContextKey{}).(ArtifactStage) + if !ok || stage.Stage == "" || stage.OutputSchema == "" || !cacheableChatRequest(messages, options) { + return c.inner.Chat(ctx, messages, options) + } + expected, err := c.expected(stage, messages, options) + if err != nil { + return c.inner.Chat(ctx, messages, options) + } + + var providerResponse *types.ChatResponse + value, err := c.runtime.LoadOrCompute(ctx, expected, func(ctx context.Context) ([]byte, error) { + response, callErr := c.inner.Chat(ctx, messages, options) + if callErr != nil { + return nil, callErr + } + if response == nil { + return nil, errors.New("chat provider returned a nil response") + } + if stage.Validate != nil { + if err := stage.Validate(response.Content); err != nil { + return nil, fmt.Errorf("validate chat stage %s output: %w", stage.Stage, err) + } + } + if !utf8.ValidString(response.Content) { + return nil, errors.New("chat provider response is not valid UTF-8") + } + providerResponse = response + return []byte(response.Content), nil + }) + if err != nil { + return nil, err + } + if providerResponse == nil { + return &types.ChatResponse{Content: string(value.Payload)}, nil + } + providerResponse.Content = string(value.Payload) + return providerResponse, nil +} + +func (c *artifactCachedChat) ChatStream( + ctx context.Context, + messages []Message, + options *ChatOptions, +) (<-chan types.StreamResponse, error) { + return c.inner.ChatStream(ctx, messages, options) +} + +func (c *artifactCachedChat) GetModelName() string { + return c.inner.GetModelName() +} + +func (c *artifactCachedChat) GetModelID() string { + return c.inner.GetModelID() +} + +func (c *artifactCachedChat) expected( + stage ArtifactStage, + messages []Message, + options *ChatOptions, +) (artifact.Expected, error) { + request := struct { + Messages []Message `json:"messages"` + Options *ChatOptions `json:"options"` + }{ + Messages: messages, + Options: options, + } + directInputs := make([]artifact.DirectInput, 0, len(messages)+1) + for index, message := range messages { + canonical, err := artifact.CanonicalJSON(message) + if err != nil { + return artifact.Expected{}, err + } + directInputs = append(directInputs, artifact.DirectInput{ + Role: fmt.Sprintf("message.%d", index), + Digest: artifact.SHA256Hex(canonical), + }) + } + optionsCanonical, err := artifact.CanonicalJSON(options) + if err != nil { + return artifact.Expected{}, err + } + directInputs = append(directInputs, artifact.DirectInput{ + Role: "options", + Digest: artifact.SHA256Hex(optionsCanonical), + }) + key, err := artifact.BuildKey(c.config.TenantID, artifact.KeyMaterial{ + KeyVersion: 1, + Stage: stage.Stage, + DirectInputs: directInputs, + Processor: c.config.Processor, + RenderedRequest: request, + Options: options, + CanonicalizerVersion: artifact.CanonicalJSONVersion, + OutputSchemaVersion: stage.OutputSchema, + }) + if err != nil { + return artifact.Expected{}, err + } + return artifact.Expected{ + Key: key, + Codec: artifact.CodecTextUTF8V1, + Validate: func(payload []byte) error { + if !utf8.Valid(payload) { + return errors.New("cached chat output is not valid UTF-8") + } + if stage.Validate != nil { + return stage.Validate(string(payload)) + } + return nil + }, + }, nil +} + +func cacheableChatRequest(messages []Message, options *ChatOptions) bool { + if options == nil || len(options.Tools) > 0 || options.ToolChoice != "" { + return false + } + for _, message := range messages { + if len(message.MultiContent) > 0 || + len(message.Images) > 0 || + len(message.ToolCalls) > 0 || + message.ToolCallID != "" { + return false + } + } + return true +} diff --git a/internal/models/chat/artifact_cache_test.go b/internal/models/chat/artifact_cache_test.go new file mode 100644 index 0000000000..929e82ce8f --- /dev/null +++ b/internal/models/chat/artifact_cache_test.go @@ -0,0 +1,130 @@ +package chat + +import ( + "context" + "errors" + "sync" + "testing" + + "github.com/Tencent/WeKnora/internal/artifact" + "github.com/Tencent/WeKnora/internal/testutil/artifactrepo" + "github.com/Tencent/WeKnora/internal/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type countingChat struct { + mu sync.Mutex + calls int + response string +} + +func (c *countingChat) Chat( + context.Context, + []Message, + *ChatOptions, +) (*types.ChatResponse, error) { + c.mu.Lock() + defer c.mu.Unlock() + c.calls++ + return &types.ChatResponse{Content: c.response}, nil +} + +func (c *countingChat) ChatStream( + context.Context, + []Message, + *ChatOptions, +) (<-chan types.StreamResponse, error) { + result := make(chan types.StreamResponse) + close(result) + return result, nil +} + +func (c *countingChat) GetModelName() string { return "chat-model" } +func (c *countingChat) GetModelID() string { return "chat-id" } + +func (c *countingChat) callCount() int { + c.mu.Lock() + defer c.mu.Unlock() + return c.calls +} + +func setupChatArtifactRuntime(t *testing.T) *artifact.Runtime { + t.Helper() + return artifact.NewRuntime(artifactrepo.New(), nil) +} + +func chatArtifactConfig() ArtifactCacheConfig { + return ArtifactCacheConfig{ + TenantID: 1, + Processor: artifact.ProcessorIdentity{ + ModelID: "chat-id", + ModelName: "chat-model", + Source: "remote", + Provider: "openai", + }, + } +} + +func TestArtifactCachedChatUsesExactRenderedRequest(t *testing.T) { + provider := &countingChat{response: "summary"} + cached := NewArtifactCachedChat(provider, setupChatArtifactRuntime(t), chatArtifactConfig()) + ctx := WithArtifactStage(context.Background(), ArtifactStage{ + Stage: "summary", + OutputSchema: "summary.text.v1", + }) + thinking := false + options := &ChatOptions{Temperature: 0.3, Thinking: &thinking} + + first, err := cached.Chat(ctx, []Message{{Role: "user", Content: " alpha\r\n "}}, options) + require.NoError(t, err) + second, err := cached.Chat(ctx, []Message{{Role: "user", Content: " alpha\r\n "}}, options) + require.NoError(t, err) + assert.Equal(t, first.Content, second.Content) + assert.Equal(t, 1, provider.callCount()) + + _, err = cached.Chat(ctx, []Message{{Role: "user", Content: "alpha\n"}}, options) + require.NoError(t, err) + assert.Equal(t, 2, provider.callCount()) +} + +func TestArtifactCachedChatDoesNotCacheInvalidOutput(t *testing.T) { + provider := &countingChat{response: ""} + cached := NewArtifactCachedChat(provider, setupChatArtifactRuntime(t), chatArtifactConfig()) + ctx := WithArtifactStage(context.Background(), ArtifactStage{ + Stage: "summary", + OutputSchema: "summary.text.v1", + Validate: func(content string) error { + if content == "" { + return errors.New("empty") + } + return nil + }, + }) + + _, err := cached.Chat(ctx, []Message{{Role: "user", Content: "content"}}, &ChatOptions{}) + require.Error(t, err) + _, err = cached.Chat(ctx, []Message{{Role: "user", Content: "content"}}, &ChatOptions{}) + require.Error(t, err) + assert.Equal(t, 2, provider.callCount()) +} + +func TestChatArtifactConfigExcludesCredentialRotation(t *testing.T) { + model := &types.Model{ + ID: "chat-id", + Name: "chat-model", + Source: types.ModelSourceRemote, + Parameters: types.ModelParameters{ + BaseURL: "https://example.com/v1", + APIKey: "secret-one", + Provider: "openai", + }, + } + first, ok := ArtifactCacheConfigFromModel(model, 1) + require.True(t, ok) + model.Parameters.APIKey = "secret-two" + model.Description = "unrelated" + second, ok := ArtifactCacheConfigFromModel(model, 1) + require.True(t, ok) + assert.Equal(t, first, second) +} diff --git a/internal/models/embedding/artifact_cache.go b/internal/models/embedding/artifact_cache.go new file mode 100644 index 0000000000..9fe3a0da5d --- /dev/null +++ b/internal/models/embedding/artifact_cache.go @@ -0,0 +1,342 @@ +package embedding + +import ( + "context" + "encoding/binary" + "errors" + "fmt" + "math" + "net/url" + + "github.com/Tencent/WeKnora/internal/artifact" + "github.com/Tencent/WeKnora/internal/types" +) + +const ( + embeddingArtifactStage = "embedding" + embeddingArtifactKeyVersion = uint16(1) + embeddingArtifactSchemaPrefix = "embedding.float32" +) + +type ArtifactCacheConfig struct { + TenantID uint64 + Processor artifact.ProcessorIdentity + Dimensions int + TruncatePromptTokens int +} + +type artifactCachedEmbedder struct { + inner Embedder + runtime *artifact.Runtime + config ArtifactCacheConfig +} + +// ArtifactCacheConfigFromModel builds an effective identity without credentials. +// Unknown extra configuration and custom headers disable caching because their +// effect and secrecy cannot be classified safely. +func ArtifactCacheConfigFromModel(model *types.Model, tenantID uint64) (ArtifactCacheConfig, bool) { + if model == nil || tenantID == 0 { + return ArtifactCacheConfig{}, false + } + if len(model.Parameters.CustomHeaders) > 0 { + return ArtifactCacheConfig{}, false + } + endpoint := model.Parameters.BaseURL + if endpoint != "" { + parsed, err := url.Parse(endpoint) + if err != nil || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" { + return ArtifactCacheConfig{}, false + } + } + extra := make(map[string]any) + for key, value := range model.Parameters.ExtraConfig { + switch key { + case "api_version", "remote_model_name": + extra[key] = value + default: + return ArtifactCacheConfig{}, false + } + } + parameters := map[string]any{ + "dimensions": model.Parameters.EmbeddingParameters.Dimension, + "truncate_prompt_tokens": model.Parameters.EmbeddingParameters.TruncatePromptTokens, + "supports_dimension_override": model.Parameters.EmbeddingParameters.SupportsDimensionOverride, + } + if len(extra) > 0 { + parameters["extra_config"] = extra + } + return ArtifactCacheConfig{ + TenantID: tenantID, + Processor: artifact.ProcessorIdentity{ + ModelID: model.ID, + ModelName: model.Name, + Source: string(model.Source), + Provider: model.Parameters.Provider, + EndpointIdentity: endpoint, + Parameters: parameters, + }, + Dimensions: model.Parameters.EmbeddingParameters.Dimension, + TruncatePromptTokens: model.Parameters.EmbeddingParameters.TruncatePromptTokens, + }, model.Parameters.EmbeddingParameters.Dimension > 0 +} + +func NewArtifactCachedEmbedder( + inner Embedder, + runtime *artifact.Runtime, + config ArtifactCacheConfig, +) Embedder { + if inner == nil || runtime == nil || config.TenantID == 0 || config.Dimensions <= 0 { + return inner + } + return &artifactCachedEmbedder{inner: inner, runtime: runtime, config: config} +} + +func (e *artifactCachedEmbedder) Embed(ctx context.Context, text string) ([]float32, error) { + if !isDocumentEmbedding(ctx) { + return e.inner.Embed(ctx, text) + } + expected, err := e.expected(text) + if err != nil { + return nil, err + } + value, err := e.runtime.LoadOrCompute(ctx, expected, func(ctx context.Context) ([]byte, error) { + vector, err := e.inner.Embed(ctx, text) + if err != nil { + return nil, err + } + return encodeEmbeddingArtifact(vector, e.config.Dimensions) + }) + if err != nil { + return nil, err + } + return decodeEmbeddingArtifact(value.Payload, e.config.Dimensions) +} + +func (e *artifactCachedEmbedder) BatchEmbed(ctx context.Context, texts []string) ([][]float32, error) { + return e.batchEmbed(ctx, texts, e.inner.BatchEmbed) +} + +func (e *artifactCachedEmbedder) BatchEmbedWithPool( + ctx context.Context, + _ Embedder, + texts []string, +) ([][]float32, error) { + return e.batchEmbed(ctx, texts, func(ctx context.Context, missing []string) ([][]float32, error) { + return e.inner.BatchEmbedWithPool(ctx, e.inner, missing) + }) +} + +func (e *artifactCachedEmbedder) batchEmbed( + ctx context.Context, + texts []string, + provider func(context.Context, []string) ([][]float32, error), +) ([][]float32, error) { + if !isDocumentEmbedding(ctx) { + return provider(ctx, texts) + } + if len(texts) == 0 { + return [][]float32{}, nil + } + + keys := make([]types.ProcessingArtifactLookup, len(texts)) + unique := make([]artifact.Expected, 0, len(texts)) + uniqueIndex := make(map[types.ProcessingArtifactLookup]int, len(texts)) + uniqueTexts := make([]string, 0, len(texts)) + for index, text := range texts { + expected, err := e.expected(text) + if err != nil { + return nil, err + } + key := expected.Key.Lookup + keys[index] = key + if _, found := uniqueIndex[key]; found { + continue + } + uniqueIndex[key] = len(unique) + unique = append(unique, expected) + uniqueTexts = append(uniqueTexts, text) + } + + cached := e.runtime.BatchLoad(ctx, unique) + uniqueVectors := make([][]float32, len(unique)) + missingIndexes := make([]int, 0, len(unique)) + for index, expected := range unique { + value, found := cached[expected.Key.Lookup] + if !found { + missingIndexes = append(missingIndexes, index) + continue + } + vector, err := decodeEmbeddingArtifact(value.Payload, e.config.Dimensions) + if err != nil { + missingIndexes = append(missingIndexes, index) + continue + } + uniqueVectors[index] = vector + } + + if len(missingIndexes) > 0 { + missingTexts := make([]string, len(missingIndexes)) + for index, uniquePosition := range missingIndexes { + missingTexts[index] = uniqueTexts[uniquePosition] + } + vectors, err := provider(ctx, missingTexts) + if err != nil { + return nil, err + } + if err := ValidateEmbeddingBatch(vectors, len(missingTexts), e.config.Dimensions); err != nil { + return nil, err + } + + candidates := make([]artifact.Candidate, 0, len(missingIndexes)) + for index, uniquePosition := range missingIndexes { + payload, err := encodeEmbeddingArtifact(vectors[index], e.config.Dimensions) + if err != nil { + return nil, err + } + candidates = append(candidates, artifact.Candidate{ + Expected: unique[uniquePosition], + Payload: payload, + }) + } + frozen := e.runtime.BatchFreeze(ctx, candidates) + for index, uniquePosition := range missingIndexes { + value, found := frozen[unique[uniquePosition].Key.Lookup] + if !found { + uniqueVectors[uniquePosition] = append([]float32(nil), vectors[index]...) + continue + } + vector, err := decodeEmbeddingArtifact(value.Payload, e.config.Dimensions) + if err != nil { + uniqueVectors[uniquePosition] = append([]float32(nil), vectors[index]...) + continue + } + uniqueVectors[uniquePosition] = vector + } + } + + result := make([][]float32, len(texts)) + for index, key := range keys { + result[index] = append([]float32(nil), uniqueVectors[uniqueIndex[key]]...) + } + return result, nil +} + +func (e *artifactCachedEmbedder) expected(text string) (artifact.Expected, error) { + key, err := artifact.BuildKey(e.config.TenantID, artifact.KeyMaterial{ + KeyVersion: embeddingArtifactKeyVersion, + Stage: embeddingArtifactStage, + DirectInputs: []artifact.DirectInput{{ + Role: "provider_input", + Digest: artifact.SHA256Hex([]byte(text)), + }}, + Processor: e.config.Processor, + RenderedRequest: map[string]any{ + "input": text, + }, + Options: map[string]any{ + "dimensions": e.config.Dimensions, + "truncate_prompt_tokens": e.config.TruncatePromptTokens, + }, + CanonicalizerVersion: artifact.CanonicalJSONVersion, + OutputSchemaVersion: fmt.Sprintf( + "%s.%d.v1", + embeddingArtifactSchemaPrefix, + e.config.Dimensions, + ), + }) + if err != nil { + return artifact.Expected{}, err + } + return artifact.Expected{ + Key: key, + Codec: artifact.CodecFloat32BEV1, + Validate: func(payload []byte) error { + _, err := decodeEmbeddingArtifact(payload, e.config.Dimensions) + return err + }, + }, nil +} + +func (e *artifactCachedEmbedder) GetModelName() string { + return e.inner.GetModelName() +} + +func (e *artifactCachedEmbedder) GetDimensions() int { + return e.inner.GetDimensions() +} + +func (e *artifactCachedEmbedder) GetModelID() string { + return e.inner.GetModelID() +} + +func isDocumentEmbedding(ctx context.Context) bool { + document, _ := ctx.Value(types.EmbedDocumentContextKey).(bool) + query, _ := ctx.Value(types.EmbedQueryContextKey).(bool) + return document && !query +} + +func encodeEmbeddingArtifact(vector []float32, expectedDimensions int) ([]byte, error) { + if err := ValidateEmbeddingBatch([][]float32{vector}, 1, expectedDimensions); err != nil { + return nil, err + } + payload := make([]byte, 4+len(vector)*4) + binary.BigEndian.PutUint32(payload[:4], uint32(len(vector))) + for index, value := range vector { + binary.BigEndian.PutUint32(payload[4+index*4:], math.Float32bits(value)) + } + return payload, nil +} + +func decodeEmbeddingArtifact(payload []byte, expectedDimensions int) ([]float32, error) { + if len(payload) < 4 || (len(payload)-4)%4 != 0 { + return nil, errors.New("invalid embedding artifact payload length") + } + count := int(binary.BigEndian.Uint32(payload[:4])) + if count != expectedDimensions || len(payload) != 4+count*4 { + return nil, fmt.Errorf( + "embedding artifact has %d dimensions, expected %d", + count, + expectedDimensions, + ) + } + vector := make([]float32, count) + for index := range vector { + vector[index] = math.Float32frombits(binary.BigEndian.Uint32(payload[4+index*4:])) + if math.IsNaN(float64(vector[index])) || math.IsInf(float64(vector[index]), 0) { + return nil, fmt.Errorf("embedding artifact dimension %d is not finite", index) + } + } + return vector, nil +} + +// ValidateEmbeddingBatch prevents short responses, wrong dimensions and +// non-finite values from reaching vector indexing. +func ValidateEmbeddingBatch(vectors [][]float32, expectedCount, expectedDimensions int) error { + if len(vectors) != expectedCount { + return fmt.Errorf( + "embedding provider returned %d vectors for %d inputs", + len(vectors), + expectedCount, + ) + } + for vectorIndex, vector := range vectors { + if expectedDimensions > 0 && len(vector) != expectedDimensions { + return fmt.Errorf( + "embedding provider vector %d has %d dimensions, expected %d", + vectorIndex, + len(vector), + expectedDimensions, + ) + } + for dimension, value := range vector { + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { + return fmt.Errorf( + "embedding provider vector %d dimension %d is not finite", + vectorIndex, + dimension, + ) + } + } + } + return nil +} diff --git a/internal/models/embedding/artifact_cache_test.go b/internal/models/embedding/artifact_cache_test.go new file mode 100644 index 0000000000..c3931a0875 --- /dev/null +++ b/internal/models/embedding/artifact_cache_test.go @@ -0,0 +1,188 @@ +package embedding + +import ( + "context" + "math" + "sync" + "testing" + + "github.com/Tencent/WeKnora/internal/artifact" + "github.com/Tencent/WeKnora/internal/testutil/artifactrepo" + "github.com/Tencent/WeKnora/internal/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type countingEmbedder struct { + mu sync.Mutex + batches [][]string + dimensions int + shortResult bool + invalid float32 +} + +func (e *countingEmbedder) Embed(ctx context.Context, text string) ([]float32, error) { + values, err := e.BatchEmbed(ctx, []string{text}) + if err != nil { + return nil, err + } + return values[0], nil +} + +func (e *countingEmbedder) BatchEmbed(_ context.Context, texts []string) ([][]float32, error) { + e.mu.Lock() + e.batches = append(e.batches, append([]string(nil), texts...)) + e.mu.Unlock() + count := len(texts) + if e.shortResult && count > 0 { + count-- + } + result := make([][]float32, count) + for index := range result { + result[index] = make([]float32, e.dimensions) + for dimension := range result[index] { + result[index][dimension] = float32(len(texts[index]) + dimension) + } + if e.invalid != 0 { + result[index][0] = e.invalid + } + } + return result, nil +} + +func (e *countingEmbedder) BatchEmbedWithPool( + ctx context.Context, + _ Embedder, + texts []string, +) ([][]float32, error) { + return e.BatchEmbed(ctx, texts) +} + +func (e *countingEmbedder) GetModelName() string { return "model" } +func (e *countingEmbedder) GetDimensions() int { return e.dimensions } +func (e *countingEmbedder) GetModelID() string { return "model-id" } + +func (e *countingEmbedder) calls() [][]string { + e.mu.Lock() + defer e.mu.Unlock() + result := make([][]string, len(e.batches)) + for index := range e.batches { + result[index] = append([]string(nil), e.batches[index]...) + } + return result +} + +func setupEmbeddingArtifactRuntime(t *testing.T) (*artifact.Runtime, *artifactrepo.Repository) { + t.Helper() + repository := artifactrepo.New() + return artifact.NewRuntime(repository, nil), repository +} + +func embeddingArtifactConfig() ArtifactCacheConfig { + return ArtifactCacheConfig{ + TenantID: 1, + Processor: artifact.ProcessorIdentity{ + ModelID: "model-id", + ModelName: "model", + Source: "remote", + Provider: "openai", + Parameters: map[string]any{ + "dimensions": 3, + }, + }, + Dimensions: 3, + } +} + +func documentEmbeddingContext() context.Context { + return context.WithValue(context.Background(), types.EmbedDocumentContextKey, true) +} + +func TestArtifactCachedEmbeddingDeduplicatesAndPreservesExactInputOrder(t *testing.T) { + runtime, _ := setupEmbeddingArtifactRuntime(t) + provider := &countingEmbedder{dimensions: 3} + cached := NewArtifactCachedEmbedder(provider, runtime, embeddingArtifactConfig()) + inputs := []string{" alpha \r\n", "beta", " alpha \r\n"} + + first, err := cached.BatchEmbed(documentEmbeddingContext(), inputs) + require.NoError(t, err) + require.Len(t, first, 3) + assert.Equal(t, first[0], first[2]) + assert.Equal(t, [][]string{{" alpha \r\n", "beta"}}, provider.calls()) + + second, err := cached.BatchEmbed(documentEmbeddingContext(), inputs) + require.NoError(t, err) + assert.Equal(t, first, second) + assert.Len(t, provider.calls(), 1) +} + +func TestArtifactCachedEmbeddingUsesPartialHits(t *testing.T) { + runtime, _ := setupEmbeddingArtifactRuntime(t) + provider := &countingEmbedder{dimensions: 3} + cached := NewArtifactCachedEmbedder(provider, runtime, embeddingArtifactConfig()) + ctx := documentEmbeddingContext() + + _, err := cached.Embed(ctx, "cached") + require.NoError(t, err) + provider.mu.Lock() + provider.batches = nil + provider.mu.Unlock() + + vectors, err := cached.BatchEmbed(ctx, []string{"cached", "missing"}) + require.NoError(t, err) + require.Len(t, vectors, 2) + assert.Equal(t, [][]string{{"missing"}}, provider.calls()) +} + +func TestArtifactCachedEmbeddingBypassesInteractiveQuery(t *testing.T) { + runtime, repository := setupEmbeddingArtifactRuntime(t) + provider := &countingEmbedder{dimensions: 3} + cached := NewArtifactCachedEmbedder(provider, runtime, embeddingArtifactConfig()) + ctx := context.WithValue(documentEmbeddingContext(), types.EmbedQueryContextKey, true) + + _, err := cached.Embed(ctx, "query") + require.NoError(t, err) + _, err = cached.Embed(ctx, "query") + require.NoError(t, err) + assert.Len(t, provider.calls(), 2) + + assert.Zero(t, repository.Count()) +} + +func TestArtifactCachedEmbeddingRejectsShortAndNonFiniteResponses(t *testing.T) { + runtime, _ := setupEmbeddingArtifactRuntime(t) + short := &countingEmbedder{dimensions: 3, shortResult: true} + cached := NewArtifactCachedEmbedder(short, runtime, embeddingArtifactConfig()) + _, err := cached.BatchEmbed(documentEmbeddingContext(), []string{"one", "two"}) + require.ErrorContains(t, err, "returned 1 vectors for 2 inputs") + + nonFinite := &countingEmbedder{dimensions: 3, invalid: float32(math.Inf(1))} + cached = NewArtifactCachedEmbedder(nonFinite, runtime, embeddingArtifactConfig()) + _, err = cached.BatchEmbed(documentEmbeddingContext(), []string{"three"}) + require.ErrorContains(t, err, "not finite") +} + +func TestArtifactCacheConfigExcludesCredentialsAndUnrelatedMetadata(t *testing.T) { + model := &types.Model{ + ID: "model-id", + Name: "model", + Source: types.ModelSourceRemote, + TenantID: 1, + Parameters: types.ModelParameters{ + BaseURL: "https://example.com/v1", + APIKey: "first-secret", + Provider: "openai", + EmbeddingParameters: types.EmbeddingParameters{ + Dimension: 3, + }, + }, + } + first, ok := ArtifactCacheConfigFromModel(model, 1) + require.True(t, ok) + model.Parameters.APIKey = "rotated-secret" + model.Description = "unrelated metadata" + second, ok := ArtifactCacheConfigFromModel(model, 1) + require.True(t, ok) + + assert.Equal(t, first, second) +} diff --git a/internal/models/embedding/batch.go b/internal/models/embedding/batch.go index 506948f088..53248ac2d9 100644 --- a/internal/models/embedding/batch.go +++ b/internal/models/embedding/batch.go @@ -60,6 +60,14 @@ func (e *batchEmbedder) BatchEmbedWithPool(ctx context.Context, model Embedder, mu.Unlock() return } + if err := ValidateEmbeddingBatch(embedding, len(texts), model.GetDimensions()); err != nil { + mu.Lock() + if firstErr == nil { + firstErr = err + } + mu.Unlock() + return + } mu.Lock() for i, text := range texts { if text == nil { diff --git a/internal/models/vlm/artifact_cache.go b/internal/models/vlm/artifact_cache.go new file mode 100644 index 0000000000..e3ad8d8b7c --- /dev/null +++ b/internal/models/vlm/artifact_cache.go @@ -0,0 +1,188 @@ +package vlm + +import ( + "context" + "errors" + "fmt" + "net/url" + "strings" + "unicode/utf8" + + "github.com/Tencent/WeKnora/internal/artifact" + "github.com/Tencent/WeKnora/internal/types" +) + +type artifactStageContextKey struct{} + +// ArtifactStage explicitly opts an ingestion VLM call into artifact reuse. +// Interactive VLM calls do not carry this marker and always reach the provider. +type ArtifactStage struct { + Stage string + OutputSchema string + Validate func(string) error +} + +func WithArtifactStage(ctx context.Context, stage ArtifactStage) context.Context { + return context.WithValue(ctx, artifactStageContextKey{}, stage) +} + +type ArtifactCacheConfig struct { + TenantID uint64 + Processor artifact.ProcessorIdentity +} + +type artifactCachedVLM struct { + inner VLM + runtime *artifact.Runtime + config ArtifactCacheConfig +} + +// ArtifactCacheConfigFromModel keeps only effective, non-secret provider +// identity. Unknown provider options and custom headers disable reuse because +// they cannot safely be classified as semantic or credential-only. +func ArtifactCacheConfigFromModel(model *types.Model, tenantID uint64) (ArtifactCacheConfig, bool) { + if model == nil || tenantID == 0 || len(model.Parameters.CustomHeaders) > 0 { + return ArtifactCacheConfig{}, false + } + endpoint := model.Parameters.BaseURL + if endpoint != "" { + parsed, err := url.Parse(endpoint) + if err != nil || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" { + return ArtifactCacheConfig{}, false + } + } + extra := make(map[string]any, len(model.Parameters.ExtraConfig)) + for key, value := range model.Parameters.ExtraConfig { + normalized := strings.NewReplacer("_", "", "-", "", ".", "").Replace(strings.ToLower(key)) + if strings.Contains(normalized, "key") || + strings.Contains(normalized, "token") || + strings.Contains(normalized, "secret") || + strings.Contains(normalized, "password") || + strings.Contains(normalized, "signature") { + return ArtifactCacheConfig{}, false + } + extra[key] = value + } + parameters := map[string]any{ + "interface_type": model.Parameters.InterfaceType, + "supports_vision": model.Parameters.SupportsVision, + } + if len(extra) > 0 { + parameters["extra_config"] = extra + } + return ArtifactCacheConfig{ + TenantID: tenantID, + Processor: artifact.ProcessorIdentity{ + ModelID: model.ID, + ModelName: model.Name, + Source: string(model.Source), + Provider: model.Parameters.Provider, + EndpointIdentity: endpoint, + Parameters: parameters, + }, + }, true +} + +func NewArtifactCachedVLM(inner VLM, runtime *artifact.Runtime, config ArtifactCacheConfig) VLM { + if inner == nil || runtime == nil || config.TenantID == 0 { + return inner + } + return &artifactCachedVLM{inner: inner, runtime: runtime, config: config} +} + +func (v *artifactCachedVLM) Predict( + ctx context.Context, + images [][]byte, + prompt string, +) (string, error) { + stage, ok := ctx.Value(artifactStageContextKey{}).(ArtifactStage) + if !ok || stage.Stage == "" || stage.OutputSchema == "" || len(images) == 0 { + return v.inner.Predict(ctx, images, prompt) + } + expected, err := v.expected(stage, images, prompt) + if err != nil { + return v.inner.Predict(ctx, images, prompt) + } + + value, err := v.runtime.LoadOrCompute(ctx, expected, func(ctx context.Context) ([]byte, error) { + output, callErr := v.inner.Predict(ctx, images, prompt) + if callErr != nil { + return nil, callErr + } + if !utf8.ValidString(output) { + return nil, errors.New("VLM provider response is not valid UTF-8") + } + if stage.Validate != nil { + if err := stage.Validate(output); err != nil { + return nil, fmt.Errorf("validate VLM stage %s output: %w", stage.Stage, err) + } + } + return []byte(output), nil + }) + if err != nil { + return "", err + } + return string(value.Payload), nil +} + +func (v *artifactCachedVLM) GetModelName() string { + return v.inner.GetModelName() +} + +func (v *artifactCachedVLM) GetModelID() string { + return v.inner.GetModelID() +} + +func (v *artifactCachedVLM) expected( + stage ArtifactStage, + images [][]byte, + prompt string, +) (artifact.Expected, error) { + type imageDescriptor struct { + Digest string `json:"digest"` + Bytes int `json:"bytes"` + } + descriptors := make([]imageDescriptor, len(images)) + directInputs := make([]artifact.DirectInput, 0, len(images)+1) + for index, image := range images { + digest := artifact.SHA256Hex(image) + descriptors[index] = imageDescriptor{Digest: digest, Bytes: len(image)} + directInputs = append(directInputs, artifact.DirectInput{ + Role: fmt.Sprintf("image.%d", index), + Digest: digest, + }) + } + directInputs = append(directInputs, artifact.DirectInput{ + Role: "prompt", + Digest: artifact.SHA256Hex([]byte(prompt)), + }) + key, err := artifact.BuildKey(v.config.TenantID, artifact.KeyMaterial{ + KeyVersion: 1, + Stage: stage.Stage, + DirectInputs: directInputs, + Processor: v.config.Processor, + RenderedRequest: map[string]any{ + "images": descriptors, + "prompt": prompt, + }, + Options: map[string]any{}, + CanonicalizerVersion: artifact.CanonicalJSONVersion, + OutputSchemaVersion: stage.OutputSchema, + }) + if err != nil { + return artifact.Expected{}, err + } + return artifact.Expected{ + Key: key, + Codec: artifact.CodecTextUTF8V1, + Validate: func(payload []byte) error { + if !utf8.Valid(payload) { + return errors.New("cached VLM output is not valid UTF-8") + } + if stage.Validate != nil { + return stage.Validate(string(payload)) + } + return nil + }, + }, nil +} diff --git a/internal/models/vlm/artifact_cache_test.go b/internal/models/vlm/artifact_cache_test.go new file mode 100644 index 0000000000..03d8fca3c4 --- /dev/null +++ b/internal/models/vlm/artifact_cache_test.go @@ -0,0 +1,95 @@ +package vlm + +import ( + "context" + "sync" + "testing" + + "github.com/Tencent/WeKnora/internal/artifact" + "github.com/Tencent/WeKnora/internal/testutil/artifactrepo" + "github.com/Tencent/WeKnora/internal/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type countingVLM struct { + mu sync.Mutex + calls int + response string +} + +func (v *countingVLM) Predict(context.Context, [][]byte, string) (string, error) { + v.mu.Lock() + defer v.mu.Unlock() + v.calls++ + return v.response, nil +} + +func (v *countingVLM) GetModelName() string { return "vision-model" } +func (v *countingVLM) GetModelID() string { return "vision-id" } + +func (v *countingVLM) callCount() int { + v.mu.Lock() + defer v.mu.Unlock() + return v.calls +} + +func setupVLMArtifactRuntime(t *testing.T) *artifact.Runtime { + t.Helper() + return artifact.NewRuntime(artifactrepo.New(), nil) +} + +func vlmArtifactConfig() ArtifactCacheConfig { + return ArtifactCacheConfig{ + TenantID: 1, + Processor: artifact.ProcessorIdentity{ + ModelID: "vision-id", + ModelName: "vision-model", + Source: "remote", + Provider: "openai", + }, + } +} + +func TestArtifactCachedVLMUsesExactImagesAndPrompt(t *testing.T) { + provider := &countingVLM{response: "caption"} + cached := NewArtifactCachedVLM(provider, setupVLMArtifactRuntime(t), vlmArtifactConfig()) + ctx := WithArtifactStage(context.Background(), ArtifactStage{ + Stage: "vlm_caption", + OutputSchema: "vlm-caption.text.v1", + }) + + first, err := cached.Predict(ctx, [][]byte{[]byte{0, 1, 2}}, " describe ") + require.NoError(t, err) + second, err := cached.Predict(ctx, [][]byte{[]byte{0, 1, 2}}, " describe ") + require.NoError(t, err) + assert.Equal(t, first, second) + assert.Equal(t, 1, provider.callCount()) + + _, err = cached.Predict(ctx, [][]byte{[]byte{0, 1, 3}}, " describe ") + require.NoError(t, err) + _, err = cached.Predict(ctx, [][]byte{[]byte{0, 1, 2}}, "describe") + require.NoError(t, err) + assert.Equal(t, 3, provider.callCount()) +} + +func TestVLMArtifactConfigExcludesCredentialRotation(t *testing.T) { + model := &types.Model{ + ID: "vision-id", + Name: "vision-model", + Source: types.ModelSourceRemote, + Parameters: types.ModelParameters{ + BaseURL: "https://example.com/v1", + APIKey: "secret-one", + Provider: "openai", + SupportsVision: true, + }, + } + first, ok := ArtifactCacheConfigFromModel(model, 1) + require.True(t, ok) + model.Parameters.APIKey = "secret-two" + model.Description = "unrelated" + second, ok := ArtifactCacheConfigFromModel(model, 1) + require.True(t, ok) + assert.Equal(t, first, second) +} diff --git a/internal/testutil/artifactrepo/repository.go b/internal/testutil/artifactrepo/repository.go new file mode 100644 index 0000000000..dc5e97fef9 --- /dev/null +++ b/internal/testutil/artifactrepo/repository.go @@ -0,0 +1,114 @@ +// Package artifactrepo provides a concurrency-safe in-memory artifact +// repository for model-package tests without importing the application +// repository layer back into those model packages. +package artifactrepo + +import ( + "context" + "sync" + + "github.com/Tencent/WeKnora/internal/types" +) + +type Repository struct { + mu sync.Mutex + values map[types.ProcessingArtifactLookup]*types.ProcessingArtifact +} + +func New() *Repository { + return &Repository{ + values: make(map[types.ProcessingArtifactLookup]*types.ProcessingArtifact), + } +} + +func (r *Repository) Get( + ctx context.Context, + key types.ProcessingArtifactLookup, +) (*types.ProcessingArtifact, error) { + values, err := r.BatchGet(ctx, []types.ProcessingArtifactLookup{key}) + if err != nil { + return nil, err + } + value, found := values[key] + if !found { + return nil, types.ErrProcessingArtifactNotFound + } + return value, nil +} + +func (r *Repository) BatchGet( + _ context.Context, + keys []types.ProcessingArtifactLookup, +) (map[types.ProcessingArtifactLookup]*types.ProcessingArtifact, error) { + r.mu.Lock() + defer r.mu.Unlock() + result := make(map[types.ProcessingArtifactLookup]*types.ProcessingArtifact, len(keys)) + for _, key := range keys { + if value := r.values[key]; value != nil { + result[key] = clone(value) + } + } + return result, nil +} + +func (r *Repository) PutIfAbsent( + _ context.Context, + candidate *types.ProcessingArtifact, +) (*types.ProcessingArtifact, bool, error) { + r.mu.Lock() + defer r.mu.Unlock() + key := candidate.Lookup() + if winner := r.values[key]; winner != nil { + return clone(winner), false, nil + } + r.values[key] = clone(candidate) + return clone(candidate), true, nil +} + +func (r *Repository) PutManyIfAbsent( + ctx context.Context, + candidates []*types.ProcessingArtifact, +) (map[types.ProcessingArtifactLookup]*types.ProcessingArtifact, error) { + for _, candidate := range candidates { + if _, _, err := r.PutIfAbsent(ctx, candidate); err != nil { + return nil, err + } + } + keys := make([]types.ProcessingArtifactLookup, len(candidates)) + for index, candidate := range candidates { + keys[index] = candidate.Lookup() + } + return r.BatchGet(ctx, keys) +} + +func (r *Repository) DeleteCorrupt( + _ context.Context, + key types.ProcessingArtifactLookup, + observedChecksum string, +) error { + r.mu.Lock() + defer r.mu.Unlock() + if value := r.values[key]; value != nil && value.PayloadChecksum == observedChecksum { + delete(r.values, key) + } + return nil +} + +func (r *Repository) TouchHits(context.Context, []types.ProcessingArtifactLookup) error { + return nil +} + +func (r *Repository) Count() int { + r.mu.Lock() + defer r.mu.Unlock() + return len(r.values) +} + +func clone(value *types.ProcessingArtifact) *types.ProcessingArtifact { + if value == nil { + return nil + } + copy := *value + copy.Payload = append([]byte(nil), value.Payload...) + return © +} diff --git a/internal/types/const.go b/internal/types/const.go index 82d6113458..b24dba2dc2 100644 --- a/internal/types/const.go +++ b/internal/types/const.go @@ -33,6 +33,9 @@ const ( // page write (user / agent / revert). Absent means the wiki ingest // pipeline. See types.WithWikiEditSource. WikiEditSourceContextKey ContextKey = "WikiEditSource" + // EmbedDocumentContextKey marks exact, already-sanitized document inputs + // that are eligible for durable processing artifact reuse. + EmbedDocumentContextKey ContextKey = "EmbedDocument" // LanguageContextKey is the context key for user language preference (e.g. "zh-CN", "en-US") LanguageContextKey ContextKey = "Language" // EmbedVisitorContextKey is the anonymous visitor id for embed OAuth isolation. diff --git a/internal/types/interfaces/chunk.go b/internal/types/interfaces/chunk.go index 9f6434d61d..5b2d30a152 100644 --- a/internal/types/interfaces/chunk.go +++ b/internal/types/interfaces/chunk.go @@ -30,6 +30,9 @@ type ChunkRepository interface { ListChunksBySeqID(ctx context.Context, tenantID uint64, seqIDs []int64) ([]*types.Chunk, error) // ListChunksByKnowledgeID lists chunks by knowledge id ListChunksByKnowledgeID(ctx context.Context, tenantID uint64, knowledgeID string) ([]*types.Chunk, error) + // ListAllChunksByKnowledgeID lists every active base and generated chunk. + // Reconciliation uses this to avoid broad delete-by-knowledge operations. + ListAllChunksByKnowledgeID(ctx context.Context, tenantID uint64, knowledgeID string) ([]*types.Chunk, error) // ListPagedChunksByKnowledgeID lists paged chunks by knowledge id. // When tagIDs is non-empty, results are filtered by tag_id (OR semantics). // knowledgeType: "faq" or "manual" - determines sort order and search behavior diff --git a/internal/types/interfaces/knowledge.go b/internal/types/interfaces/knowledge.go index daf86eca7a..71ca8646c6 100644 --- a/internal/types/interfaces/knowledge.go +++ b/internal/types/interfaces/knowledge.go @@ -226,6 +226,25 @@ type KnowledgeRepository interface { tenantID uint64, kbID string, page *types.Pagination, filter types.KnowledgeListFilter, ) ([]*types.Knowledge, int64, error) UpdateKnowledge(ctx context.Context, knowledge *types.Knowledge) error + // UpdateKnowledgeIfAttemptCurrent performs the final publish as one + // conditional statement. It returns false when a newer root attempt already + // exists, closing the race between an application-level fence check and the + // database write. + UpdateKnowledgeIfAttemptCurrent( + ctx context.Context, + knowledge *types.Knowledge, + attempt int, + ) (bool, error) + // UpdateKnowledgeColumnsIfAttemptCurrent conditionally updates orchestration + // columns that full-row updates intentionally omit (for example the + // pending-subtask counter). + UpdateKnowledgeColumnsIfAttemptCurrent( + ctx context.Context, + tenantID uint64, + knowledgeID string, + attempt int, + values map[string]interface{}, + ) (bool, error) // UpdateKnowledgeBatch updates knowledge items in batch UpdateKnowledgeBatch(ctx context.Context, knowledgeList []*types.Knowledge) error DeleteKnowledge(ctx context.Context, tenantID uint64, id string) error diff --git a/internal/types/interfaces/processing_artifact.go b/internal/types/interfaces/processing_artifact.go new file mode 100644 index 0000000000..fb306e7312 --- /dev/null +++ b/internal/types/interfaces/processing_artifact.go @@ -0,0 +1,34 @@ +package interfaces + +import ( + "context" + + "github.com/Tencent/WeKnora/internal/types" +) + +// ProcessingArtifactRepository provides immutable, tenant-scoped artifact +// manifests. Put methods never overwrite an existing successful value. +type ProcessingArtifactRepository interface { + Get( + ctx context.Context, + key types.ProcessingArtifactLookup, + ) (*types.ProcessingArtifact, error) + BatchGet( + ctx context.Context, + keys []types.ProcessingArtifactLookup, + ) (map[types.ProcessingArtifactLookup]*types.ProcessingArtifact, error) + PutIfAbsent( + ctx context.Context, + candidate *types.ProcessingArtifact, + ) (winner *types.ProcessingArtifact, created bool, err error) + PutManyIfAbsent( + ctx context.Context, + candidates []*types.ProcessingArtifact, + ) (map[types.ProcessingArtifactLookup]*types.ProcessingArtifact, error) + DeleteCorrupt( + ctx context.Context, + key types.ProcessingArtifactLookup, + observedChecksum string, + ) error + TouchHits(ctx context.Context, keys []types.ProcessingArtifactLookup) error +} diff --git a/internal/types/interfaces/retriever_graph.go b/internal/types/interfaces/retriever_graph.go index 9b6333ce0f..34d7728de8 100644 --- a/internal/types/interfaces/retriever_graph.go +++ b/internal/types/interfaces/retriever_graph.go @@ -15,3 +15,23 @@ type RetrieveGraphRepository interface { // SearchNode searches for nodes in the repository SearchNode(ctx context.Context, namespace types.NameSpace, nodes []string) (*types.GraphData, error) } + +// GraphContributionRepository is the desired-state graph publishing extension. +// A contribution belongs to one stable chunk and one attempt. Implementations +// must ignore a replace/delete from an attempt older than the latest marker for +// that contribution. +type GraphContributionRepository interface { + ReplaceGraphContribution( + ctx context.Context, + namespace types.NameSpace, + chunkID string, + attempt int, + graph *types.GraphData, + ) (applied bool, err error) + DeleteGraphContributions( + ctx context.Context, + namespace types.NameSpace, + chunkIDs []string, + attempt int, + ) error +} diff --git a/internal/types/knowledge_span.go b/internal/types/knowledge_span.go index 9d121c00c0..6071868217 100644 --- a/internal/types/knowledge_span.go +++ b/internal/types/knowledge_span.go @@ -109,6 +109,19 @@ func (KnowledgeProcessingSpan) TableName() string { return "knowledge_processing_spans" } +// KnowledgeAttemptCounter allocates monotonically increasing processing +// attempt numbers independently from span insertion. Keeping the allocation in +// one locked row closes the MAX(attempt)+1 race between concurrent reparses. +type KnowledgeAttemptCounter struct { + KnowledgeID string `gorm:"column:knowledge_id;primaryKey;size:64"` + LastAttempt int `gorm:"column:last_attempt;not null"` + UpdatedAt time.Time `gorm:"column:updated_at;autoUpdateTime"` +} + +func (KnowledgeAttemptCounter) TableName() string { + return "knowledge_attempt_counters" +} + // SpanTreeNode is the API-only tree projection. The repo returns flat // rows; the handler/tracker assembles SpanTreeNode for the response. type SpanTreeNode struct { diff --git a/internal/types/processing_artifact.go b/internal/types/processing_artifact.go new file mode 100644 index 0000000000..0f337816c6 --- /dev/null +++ b/internal/types/processing_artifact.go @@ -0,0 +1,121 @@ +package types + +import ( + "errors" + "fmt" + "regexp" + "time" +) + +var ( + processingArtifactStagePattern = regexp.MustCompile(`^[a-z0-9][a-z0-9._-]{0,63}$`) + processingArtifactHashPattern = regexp.MustCompile(`^[0-9a-f]{64}$`) + + ErrProcessingArtifactNotFound = errors.New("processing artifact not found") +) + +// ProcessingArtifactLookup is the complete tenant-scoped identity of an +// immutable processing artifact. +type ProcessingArtifactLookup struct { + TenantID uint64 + Stage string + KeyVersion uint16 + ArtifactKey string +} + +// Validate rejects incomplete or ambiguous artifact identities before they +// reach a database query. +func (k ProcessingArtifactLookup) Validate() error { + if k.TenantID == 0 { + return errors.New("processing artifact tenant ID must not be zero") + } + if !processingArtifactStagePattern.MatchString(k.Stage) { + return fmt.Errorf("invalid processing artifact stage %q", k.Stage) + } + if k.KeyVersion == 0 { + return errors.New("processing artifact key version must not be zero") + } + if !processingArtifactHashPattern.MatchString(k.ArtifactKey) { + return errors.New("processing artifact key must be 64 lowercase hex characters") + } + return nil +} + +// ProcessingArtifact is an ownership-free, immutable value produced by a +// processing stage. Knowledge IDs, chunk IDs and attempt IDs belong in live +// reconciliation state, never in this record or its payload. +type ProcessingArtifact struct { + ID uint64 `gorm:"primaryKey"` + TenantID uint64 `gorm:"not null;uniqueIndex:uq_processing_artifacts_key,priority:1"` + Stage string `gorm:"size:64;not null;uniqueIndex:uq_processing_artifacts_key,priority:2"` + KeyVersion uint16 `gorm:"not null;uniqueIndex:uq_processing_artifacts_key,priority:3"` + ArtifactKey string `gorm:"size:64;not null;uniqueIndex:uq_processing_artifacts_key,priority:4"` + ProcessorDigest string `gorm:"size:64;not null"` + OutputDigest string `gorm:"size:64;not null"` + OutputSchema string `gorm:"size:64;not null"` + Codec string `gorm:"size:32;not null"` + InlinePayload bool `gorm:"not null;default:true"` + Payload []byte `gorm:"type:bytea"` + ObjectRef string `gorm:"type:text;not null;default:''"` + PayloadChecksum string `gorm:"size:64;not null"` + SizeBytes int64 `gorm:"not null"` + HitCount uint64 `gorm:"not null;default:0"` + LastHitAt *time.Time + CreatedAt time.Time `gorm:"not null"` +} + +func (ProcessingArtifact) TableName() string { + return "processing_artifacts" +} + +// Lookup returns the complete immutable key represented by the row. +func (a ProcessingArtifact) Lookup() ProcessingArtifactLookup { + return ProcessingArtifactLookup{ + TenantID: a.TenantID, + Stage: a.Stage, + KeyVersion: a.KeyVersion, + ArtifactKey: a.ArtifactKey, + } +} + +// Validate checks manifest invariants. It intentionally does not decode the +// payload; codec- and stage-specific validation happens when the value is read. +func (a ProcessingArtifact) Validate() error { + if err := a.Lookup().Validate(); err != nil { + return err + } + if !processingArtifactHashPattern.MatchString(a.ProcessorDigest) { + return errors.New("processing artifact processor digest must be 64 lowercase hex characters") + } + if !processingArtifactHashPattern.MatchString(a.OutputDigest) { + return errors.New("processing artifact output digest must be 64 lowercase hex characters") + } + if a.OutputSchema == "" { + return errors.New("processing artifact output schema must not be empty") + } + if a.Codec == "" { + return errors.New("processing artifact codec must not be empty") + } + if !processingArtifactHashPattern.MatchString(a.PayloadChecksum) { + return errors.New("processing artifact payload checksum must be 64 lowercase hex characters") + } + if a.SizeBytes < 0 { + return errors.New("processing artifact size must not be negative") + } + if a.InlinePayload { + if a.Payload == nil { + return errors.New("inline processing artifact payload must not be nil") + } + if a.ObjectRef != "" { + return errors.New("inline processing artifact must not have an object reference") + } + } else { + if a.ObjectRef == "" { + return errors.New("object processing artifact reference must not be empty") + } + if a.Payload != nil { + return errors.New("object processing artifact payload must be nil") + } + } + return nil +} diff --git a/internal/types/task.go b/internal/types/task.go index 9f16b2e727..ca649e4ca4 100644 --- a/internal/types/task.go +++ b/internal/types/task.go @@ -463,6 +463,7 @@ type ManualProcessPayload struct { KnowledgeBaseID string `json:"knowledge_base_id"` Content string `json:"content"` // cleaned markdown content NeedCleanup bool `json:"need_cleanup"` // true for update, false for create + Attempt int `json:"attempt,omitempty"` } // ImageMultimodalPayload represents the image multimodal processing task payload. diff --git a/migrations/mysql/00-init-db.sql b/migrations/mysql/00-init-db.sql index a4c7aec608..721ab21fbb 100644 --- a/migrations/mysql/00-init-db.sql +++ b/migrations/mysql/00-init-db.sql @@ -227,3 +227,40 @@ CREATE TABLE chunk_revisions ( UNIQUE KEY idx_chunk_revisions_chunk_revision (chunk_id, revision), KEY idx_chunk_revisions_tenant_chunk (tenant_id, chunk_id) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; + +CREATE TABLE processing_artifacts ( + id BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY, + tenant_id BIGINT UNSIGNED NOT NULL, + stage VARCHAR(64) NOT NULL, + key_version INTEGER UNSIGNED NOT NULL, + artifact_key CHAR(64) NOT NULL, + processor_digest CHAR(64) NOT NULL, + output_digest CHAR(64) NOT NULL, + output_schema VARCHAR(64) NOT NULL, + codec VARCHAR(32) NOT NULL, + inline_payload BOOLEAN NOT NULL DEFAULT TRUE, + payload LONGBLOB, + object_ref TEXT NOT NULL, + payload_checksum CHAR(64) NOT NULL, + size_bytes BIGINT NOT NULL, + hit_count BIGINT UNSIGNED NOT NULL DEFAULT 0, + last_hit_at TIMESTAMP NULL DEFAULT NULL, + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + UNIQUE KEY uq_processing_artifacts_key + (tenant_id, stage, key_version, artifact_key), + KEY idx_processing_artifacts_tenant_created (tenant_id, created_at), + CONSTRAINT ck_processing_artifacts_payload + CHECK ( + (inline_payload = TRUE AND payload IS NOT NULL AND object_ref = '') + OR + (inline_payload = FALSE AND payload IS NULL AND object_ref <> '') + ), + CONSTRAINT ck_processing_artifacts_size CHECK (size_bytes >= 0) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; + +CREATE TABLE knowledge_attempt_counters ( + knowledge_id VARCHAR(64) PRIMARY KEY, + last_attempt INTEGER UNSIGNED NOT NULL, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP + ON UPDATE CURRENT_TIMESTAMP +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; diff --git a/migrations/paradedb/00-init-db.sql b/migrations/paradedb/00-init-db.sql index 7380224960..dda0cc9158 100644 --- a/migrations/paradedb/00-init-db.sql +++ b/migrations/paradedb/00-init-db.sql @@ -212,3 +212,42 @@ WITH ( ); CREATE INDEX ON embeddings USING hnsw ((embedding::halfvec(3584)) halfvec_cosine_ops) WITH (m = 16, ef_construction = 64) WHERE (dimension = 3584); CREATE INDEX ON embeddings USING hnsw ((embedding::halfvec(798)) halfvec_cosine_ops) WITH (m = 16, ef_construction = 64) WHERE (dimension = 798); + +CREATE TABLE IF NOT EXISTS processing_artifacts ( + id BIGSERIAL PRIMARY KEY, + tenant_id BIGINT NOT NULL, + stage VARCHAR(64) NOT NULL, + key_version INTEGER NOT NULL, + artifact_key CHAR(64) NOT NULL, + processor_digest CHAR(64) NOT NULL, + output_digest CHAR(64) NOT NULL, + output_schema VARCHAR(64) NOT NULL, + codec VARCHAR(32) NOT NULL, + inline_payload BOOLEAN NOT NULL DEFAULT TRUE, + payload BYTEA, + object_ref TEXT NOT NULL DEFAULT '', + payload_checksum CHAR(64) NOT NULL, + size_bytes BIGINT NOT NULL, + hit_count BIGINT NOT NULL DEFAULT 0, + last_hit_at TIMESTAMP WITH TIME ZONE, + created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + CONSTRAINT uq_processing_artifacts_key + UNIQUE (tenant_id, stage, key_version, artifact_key), + CONSTRAINT ck_processing_artifacts_payload + CHECK ( + (inline_payload AND payload IS NOT NULL AND object_ref = '') + OR + (NOT inline_payload AND payload IS NULL AND object_ref <> '') + ), + CONSTRAINT ck_processing_artifacts_size CHECK (size_bytes >= 0), + CONSTRAINT ck_processing_artifacts_hit_count CHECK (hit_count >= 0) +); + +CREATE INDEX IF NOT EXISTS idx_processing_artifacts_tenant_created + ON processing_artifacts (tenant_id, created_at); + +CREATE TABLE IF NOT EXISTS knowledge_attempt_counters ( + knowledge_id VARCHAR(64) PRIMARY KEY, + last_attempt INTEGER NOT NULL CHECK (last_attempt >= 0), + updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP +); diff --git a/migrations/sqlite/000001_processing_artifacts.down.sql b/migrations/sqlite/000001_processing_artifacts.down.sql new file mode 100644 index 0000000000..217f81cdca --- /dev/null +++ b/migrations/sqlite/000001_processing_artifacts.down.sql @@ -0,0 +1,2 @@ +DROP TABLE IF EXISTS knowledge_attempt_counters; +DROP TABLE IF EXISTS processing_artifacts; diff --git a/migrations/sqlite/000001_processing_artifacts.up.sql b/migrations/sqlite/000001_processing_artifacts.up.sql new file mode 100644 index 0000000000..1f2fa3c80c --- /dev/null +++ b/migrations/sqlite/000001_processing_artifacts.up.sql @@ -0,0 +1,38 @@ +CREATE TABLE IF NOT EXISTS processing_artifacts ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + tenant_id INTEGER NOT NULL, + stage VARCHAR(64) NOT NULL, + key_version INTEGER NOT NULL, + artifact_key CHAR(64) NOT NULL, + processor_digest CHAR(64) NOT NULL, + output_digest CHAR(64) NOT NULL, + output_schema VARCHAR(64) NOT NULL, + codec VARCHAR(32) NOT NULL, + inline_payload BOOLEAN NOT NULL DEFAULT 1, + payload BLOB, + object_ref TEXT NOT NULL DEFAULT '', + payload_checksum CHAR(64) NOT NULL, + size_bytes INTEGER NOT NULL, + hit_count INTEGER NOT NULL DEFAULT 0, + last_hit_at DATETIME, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + CONSTRAINT uq_processing_artifacts_key + UNIQUE (tenant_id, stage, key_version, artifact_key), + CONSTRAINT ck_processing_artifacts_payload + CHECK ( + (inline_payload = 1 AND payload IS NOT NULL AND object_ref = '') + OR + (inline_payload = 0 AND payload IS NULL AND object_ref <> '') + ), + CONSTRAINT ck_processing_artifacts_size CHECK (size_bytes >= 0), + CONSTRAINT ck_processing_artifacts_hit_count CHECK (hit_count >= 0) +); + +CREATE INDEX IF NOT EXISTS idx_processing_artifacts_tenant_created + ON processing_artifacts (tenant_id, created_at); + +CREATE TABLE IF NOT EXISTS knowledge_attempt_counters ( + knowledge_id VARCHAR(64) PRIMARY KEY, + last_attempt INTEGER NOT NULL CHECK (last_attempt >= 0), + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP +); diff --git a/migrations/versioned/000077_processing_artifacts.down.sql b/migrations/versioned/000077_processing_artifacts.down.sql new file mode 100644 index 0000000000..e6969c0e4e --- /dev/null +++ b/migrations/versioned/000077_processing_artifacts.down.sql @@ -0,0 +1,3 @@ +-- Roll back processing-artifact persistence and attempt fencing. +DROP TABLE IF EXISTS knowledge_attempt_counters; +DROP TABLE IF EXISTS processing_artifacts; diff --git a/migrations/versioned/000077_processing_artifacts.up.sql b/migrations/versioned/000077_processing_artifacts.up.sql new file mode 100644 index 0000000000..4af80224c5 --- /dev/null +++ b/migrations/versioned/000077_processing_artifacts.up.sql @@ -0,0 +1,44 @@ +DO $$ BEGIN RAISE NOTICE '[Migration 000077] Creating processing artifacts...'; END $$; + +CREATE TABLE IF NOT EXISTS processing_artifacts ( + id BIGSERIAL PRIMARY KEY, + tenant_id BIGINT NOT NULL, + stage VARCHAR(64) NOT NULL, + key_version INTEGER NOT NULL, + artifact_key CHAR(64) NOT NULL, + processor_digest CHAR(64) NOT NULL, + output_digest CHAR(64) NOT NULL, + output_schema VARCHAR(64) NOT NULL, + codec VARCHAR(32) NOT NULL, + inline_payload BOOLEAN NOT NULL DEFAULT TRUE, + payload BYTEA, + object_ref TEXT NOT NULL DEFAULT '', + payload_checksum CHAR(64) NOT NULL, + size_bytes BIGINT NOT NULL, + hit_count BIGINT NOT NULL DEFAULT 0, + last_hit_at TIMESTAMP WITH TIME ZONE, + created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + CONSTRAINT uq_processing_artifacts_key + UNIQUE (tenant_id, stage, key_version, artifact_key), + CONSTRAINT ck_processing_artifacts_payload + CHECK ( + (inline_payload AND payload IS NOT NULL AND object_ref = '') + OR + (NOT inline_payload AND payload IS NULL AND object_ref <> '') + ), + CONSTRAINT ck_processing_artifacts_size CHECK (size_bytes >= 0), + CONSTRAINT ck_processing_artifacts_hit_count CHECK (hit_count >= 0) +); + +CREATE INDEX IF NOT EXISTS idx_processing_artifacts_tenant_created + ON processing_artifacts (tenant_id, created_at); + +CREATE TABLE IF NOT EXISTS knowledge_attempt_counters ( + knowledge_id VARCHAR(64) PRIMARY KEY, + last_attempt INTEGER NOT NULL, + updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP, + CONSTRAINT ck_knowledge_attempt_counters_last + CHECK (last_attempt >= 0) +); + +DO $$ BEGIN RAISE NOTICE '[Migration 000075] Processing artifacts ready'; END $$; From ed6c2d565af85dab0152b880e9fef9ff535b2e06 Mon Sep 17 00:00:00 2001 From: Even1115 Date: Wed, 29 Jul 2026 17:37:14 +0800 Subject: [PATCH 02/11] fix: harden artifact cache concurrency and rollout --- config/config.yaml | 16 + docs/dev/artifact-dag-reconciliation.md | 164 ++++++++++ .../service/knowledge_span_tracker_test.go | 6 + internal/artifact/runtime.go | 301 ++++++++++++++++-- internal/artifact/runtime_test.go | 109 +++++++ internal/config/config.go | 12 + internal/container/artifact_runtime.go | 67 ++++ internal/container/artifact_runtime_test.go | 47 +++ internal/container/container.go | 12 +- internal/models/embedding/artifact_cache.go | 181 +++++++++-- .../models/embedding/artifact_cache_test.go | 57 ++++ .../000077_processing_artifacts.up.sql | 2 +- 12 files changed, 919 insertions(+), 55 deletions(-) create mode 100644 docs/dev/artifact-dag-reconciliation.md create mode 100644 internal/container/artifact_runtime.go create mode 100644 internal/container/artifact_runtime_test.go diff --git a/config/config.yaml b/config/config.yaml index c6b69f3c1d..b5c62b4601 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -3,6 +3,22 @@ server: port: 8080 host: "0.0.0.0" +# Reusable ingestion artifacts. Environment overrides use +# ARTIFACT_CACHE_READ_ENABLED and ARTIFACT_CACHE_WRITE_ENABLED. +artifact_cache: + read_enabled: true + write_enabled: true + stages: + parse: true + embedding: true + summary: true + question: true + vlm_ocr: true + vlm_caption: true + wiki_map: true + graph_extract.entities: true + graph_extract.relationships: true + # Conversation service configuration # NOTE: Prompt content is resolved from prompt_templates/ YAML files via xxx_id fields. # Set the _id to the template ID you want; the system will load its content at startup. diff --git a/docs/dev/artifact-dag-reconciliation.md b/docs/dev/artifact-dag-reconciliation.md new file mode 100644 index 0000000000..9b09dbddf1 --- /dev/null +++ b/docs/dev/artifact-dag-reconciliation.md @@ -0,0 +1,164 @@ +# Versioned Artifact DAG and Knowledge Reconciliation + +This document describes the implementation added for +[issue #1679](https://github.com/Tencent/WeKnora/issues/1679). The goal is to +reuse expensive deterministic processing outputs without allowing cached data +to own mutable knowledge, chunk, or attempt state. + +## Invariants + +- Artifact identity is tenant scoped and immutable. +- An artifact key contains the exact direct-input digests, processor identity, + rendered request, effective options, canonicalizer version, and output schema + version. +- Downstream keys depend on upstream `output_digest` values, not only on the + upstream input key. +- Stored payloads contain reusable provider output only. Tenant, knowledge, + chunk, and attempt ownership is rebound when the desired state is built. +- Cache failures are fail-open. Provider failures and invalid provider output + remain visible to the caller. +- Publication is add/update first and stale deletion last. A stale attempt must + not publish final knowledge state. + +## Data flow + +```text +source bytes + -> DocReader artifact + -> normalized source chunks + -> chat/VLM/wiki artifacts + -> stable desired chunk IDs + -> embedding and graph artifacts + -> desired-state diff + -> add/update/index + -> attempt fence + -> delete stale vector, graph, and chunk state + -> conditional knowledge publication +``` + +`internal/artifact` owns canonical keys, codecs, payload validation, immutable +freeze semantics, cache lookup, corruption eviction, singleflight, Redis +leases, stable UUIDv5 identities, desired-state diffs, and attempt fences. +Model and service adapters keep provider-specific request and response details +outside the reconciliation layer. + +## Persistence + +`processing_artifacts` is keyed by: + +```text +(tenant_id, stage, key_version, artifact_key) +``` + +The row records processor and output digests, schema, codec, checksum, size, and +hit metadata. `object_ref` is reserved for a future object-store implementation; +the current implementation stores bounded inline payloads and bypasses +DocReader caching above 16 MiB. + +`knowledge_attempt_counters` allocates monotonically increasing attempts +independently of span history. This prevents attempt reuse after spans are +cleaned up. + +Migration locations: + +- PostgreSQL: `migrations/versioned/000077_processing_artifacts.{up,down}.sql` +- SQLite: `migrations/sqlite/000001_processing_artifacts.{up,down}.sql` +- MySQL bootstrap: `migrations/mysql/00-init-db.sql` +- ParadeDB bootstrap: `migrations/paradedb/00-init-db.sql` + +## Concurrency and publication + +The runtime suppresses duplicate work in three layers: + +1. in-process `singleflight` for an artifact key; +2. a Redis lease for workers in different processes; +3. database `put-if-absent` uniqueness as the final correctness boundary. + +Batch embedding selects a deterministic missing artifact as the batch leader. +The leader lease covers the provider batch, while the remaining results are +frozen together. This keeps one provider call for concurrent identical batches +without changing caller output order. + +Knowledge processing allocates an attempt before work begins. Final publication +and destructive stale cleanup recheck that attempt. Graph storage uses +per-chunk contributions so stale chunks can be removed exactly without deleting +unchanged contributions. + +## Compatibility and rollback + +The schema addition is non-destructive to existing knowledge and chunk tables. +Artifact misses rebuild data through the existing providers, so an empty +artifact table is valid. + +`artifact_cache` supports independent read/write controls and exact stage +overrides. Missing configuration defaults to read/write enabled. + +| Mode | `read_enabled` | `write_enabled` | +| --- | --- | --- | +| Disabled | `false` | `false` | +| Shadow write | `false` | `true` | +| Read fallback | `true` | `true` | +| Read only | `true` | `false` | + +Environment overrides use `ARTIFACT_CACHE_READ_ENABLED` and +`ARTIFACT_CACHE_WRITE_ENABLED`. Individual stages can be disabled in +`artifact_cache.stages`; omitted stages remain enabled. + +To roll back: + +1. disable artifact reads, then writes, and restart or drain workers; +2. deploy the previous application version; +3. verify no process reads or writes artifact/attempt tables; +4. optionally run migration `000077` down. + +Dropping the new tables discards only reusable artifacts and attempt counters; +it does not delete knowledge, chunks, vectors, or graph data. Do not run the +down migration while the new worker version is active. + +## Observability and log safety + +The runtime observer emits structured fields for stage, outcome, reason, key +version, output schema, provider calls, singleflight wait time, and embedding +batch totals/hits/misses/deduplication. Outcomes use `hit`, `miss`, `computed`, +`wait`, `bypass`, `corrupt`, and `error_fallback`. + +Logs intentionally omit tenant IDs, complete artifact keys, payloads, prompts, +URLs, and raw database error strings. Database errors are represented only by +their concrete error class because driver messages can contain SQL arguments. + +## Validation + +The implementation has focused tests for: + +- canonical keys, tenant isolation, credential exclusion, and schema versions; +- checksum validation, corrupt-row eviction, fail-open storage, and immutable + first-writer-wins behavior; +- local and Redis-backed concurrent provider suppression; +- exact input bytes, duplicate input ordering, partial batch hits, and invalid + provider response rejection; +- stable chunk/generated IDs, desired-state diffs, and stale-attempt fencing; +- SQLite migration up/down and uniqueness behavior; +- DocReader, chat, embedding, VLM, wiki, multimodal, and graph stage adapters. + +Run the focused suite with: + +```bash +go test -count=1 \ + ./internal/artifact \ + ./internal/application/repository \ + ./internal/application/service \ + ./internal/database \ + ./internal/models/chat \ + ./internal/models/embedding \ + ./internal/models/vlm + +go test -race -count=1 \ + ./internal/artifact \ + ./internal/application/repository \ + ./internal/models/embedding +``` + +Live PostgreSQL/MySQL/Redis/Neo4j/vector integration should also be exercised in +the deployment environment before a broad rollout. URL-based DocReader inputs +currently bypass artifacts because their remote content cannot be proven stable +from the URL alone. diff --git a/internal/application/service/knowledge_span_tracker_test.go b/internal/application/service/knowledge_span_tracker_test.go index 219f944690..cf958ce9f8 100644 --- a/internal/application/service/knowledge_span_tracker_test.go +++ b/internal/application/service/knowledge_span_tracker_test.go @@ -49,6 +49,12 @@ CREATE TABLE IF NOT EXISTS knowledge_processing_spans ( updated_at DATETIME DEFAULT CURRENT_TIMESTAMP, UNIQUE (knowledge_id, attempt, span_id) ); + +CREATE TABLE IF NOT EXISTS knowledge_attempt_counters ( + knowledge_id VARCHAR(64) PRIMARY KEY, + last_attempt INTEGER NOT NULL CHECK (last_attempt >= 0), + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP +); ` func setupSpanTrackerTest(t *testing.T) (SpanTracker, *gorm.DB) { diff --git a/internal/artifact/runtime.go b/internal/artifact/runtime.go index 5bfb8ff61e..da81215638 100644 --- a/internal/artifact/runtime.go +++ b/internal/artifact/runtime.go @@ -5,6 +5,9 @@ import ( "errors" "fmt" "strconv" + "strings" + "sync" + "sync/atomic" "time" "github.com/Tencent/WeKnora/internal/types" @@ -20,12 +23,23 @@ const ( EventStoreFailure EventKind = "store_failure" EventStored EventKind = "stored" EventLostRace EventKind = "lost_race" + EventComputed EventKind = "computed" + EventWait EventKind = "wait" + EventBypass EventKind = "bypass" ) type Event struct { - Kind EventKind - Lookup types.ProcessingArtifactLookup - Err error + Kind EventKind + Lookup types.ProcessingArtifactLookup + OutputSchema string + Reason string + ProviderCall bool + BatchTotal int + BatchHits int + BatchMisses int + BatchDeduplicated int + SingleflightWaitMS int64 + Err error } type Observer func(Event) @@ -66,10 +80,19 @@ type Runtime struct { observer Observer lease Lease group singleflight.Group + configMu sync.RWMutex + read bool + write bool + stages map[string]bool } func NewRuntime(repository Repository, observer Observer) *Runtime { - return &Runtime{repository: repository, observer: observer} + return &Runtime{ + repository: repository, + observer: observer, + read: true, + write: true, + } } func (r *Runtime) ConfigureLease(lease Lease) { @@ -78,6 +101,25 @@ func (r *Runtime) ConfigureLease(lease Lease) { } } +// ConfigureCacheMode controls artifact reads and writes without changing the +// provider path. Stages omitted from the map remain enabled. +func (r *Runtime) ConfigureCacheMode(read, write bool, stages map[string]bool) { + if r == nil { + return + } + r.configMu.Lock() + defer r.configMu.Unlock() + r.read = read + r.write = write + r.stages = make(map[string]bool, len(stages)) + for stage, enabled := range stages { + normalized := strings.TrimSpace(stage) + if normalized != "" { + r.stages[normalized] = enabled + } + } +} + type Expected struct { Key Key Codec string @@ -116,13 +158,31 @@ func (r *Runtime) BatchLoad( keys := make([]types.ProcessingArtifactLookup, 0, len(expected)) byKey := make(map[types.ProcessingArtifactLookup]Expected, len(expected)) for _, item := range expected { + if !r.readEnabled(item.Key.Lookup.Stage) { + r.emit(Event{ + Kind: EventBypass, + Lookup: item.Key.Lookup, + OutputSchema: item.Key.OutputSchema, + Reason: "read_disabled", + }) + continue + } keys = append(keys, item.Key.Lookup) byKey[item.Key.Lookup] = item } + if len(keys) == 0 { + return result + } artifacts, err := r.repository.BatchGet(ctx, keys) if err != nil { for _, key := range keys { - r.emit(Event{Kind: EventStoreFailure, Lookup: key, Err: err}) + r.emit(Event{ + Kind: EventStoreFailure, + Lookup: key, + OutputSchema: byKey[key].Key.OutputSchema, + Reason: "read_error", + Err: err, + }) } return result } @@ -132,7 +192,12 @@ func (r *Runtime) BatchLoad( item := byKey[key] manifest, found := artifacts[key] if !found { - r.emit(Event{Kind: EventMiss, Lookup: key}) + r.emit(Event{ + Kind: EventMiss, + Lookup: key, + OutputSchema: item.Key.OutputSchema, + Reason: "not_found", + }) continue } payload, decodeErr := DecodeInline(manifest, key, item.Key.OutputSchema, item.Codec) @@ -140,9 +205,21 @@ func (r *Runtime) BatchLoad( decodeErr = item.Validate(payload) } if decodeErr != nil { - r.emit(Event{Kind: EventCorrupt, Lookup: key, Err: decodeErr}) + r.emit(Event{ + Kind: EventCorrupt, + Lookup: key, + OutputSchema: item.Key.OutputSchema, + Reason: corruptReason(decodeErr), + Err: decodeErr, + }) if deleteErr := r.repository.DeleteCorrupt(ctx, key, manifest.PayloadChecksum); deleteErr != nil { - r.emit(Event{Kind: EventStoreFailure, Lookup: key, Err: deleteErr}) + r.emit(Event{ + Kind: EventStoreFailure, + Lookup: key, + OutputSchema: item.Key.OutputSchema, + Reason: "corrupt_delete_error", + Err: deleteErr, + }) } continue } @@ -152,12 +229,23 @@ func (r *Runtime) BatchLoad( CacheHit: true, } hits = append(hits, key) - r.emit(Event{Kind: EventHit, Lookup: key}) + r.emit(Event{ + Kind: EventHit, + Lookup: key, + OutputSchema: item.Key.OutputSchema, + Reason: "found_valid", + }) } if len(hits) > 0 { if err := r.repository.TouchHits(ctx, hits); err != nil { for _, key := range hits { - r.emit(Event{Kind: EventStoreFailure, Lookup: key, Err: err}) + r.emit(Event{ + Kind: EventStoreFailure, + Lookup: key, + OutputSchema: byKey[key].Key.OutputSchema, + Reason: "touch_error", + Err: err, + }) } } } @@ -182,6 +270,8 @@ func (r *Runtime) LoadOrCompute( return uncachedValue(payload, expected.Validate, err) } + waitStarted := time.Now() + var computedByCaller atomic.Bool result := r.group.DoChan(singleflightKey(expected.Key.Lookup), func() (any, error) { if value, hit := r.Load(ctx, expected); hit { return value, nil @@ -191,9 +281,11 @@ func (r *Runtime) LoadOrCompute( handle, acquired, leaseErr := r.lease.TryAcquire(ctx, expected.Key.Lookup) if leaseErr != nil { r.emit(Event{ - Kind: EventStoreFailure, - Lookup: expected.Key.Lookup, - Err: leaseErr, + Kind: EventStoreFailure, + Lookup: expected.Key.Lookup, + OutputSchema: expected.Key.OutputSchema, + Reason: "lease_error", + Err: leaseErr, }) break } @@ -220,6 +312,7 @@ func (r *Runtime) LoadOrCompute( } } } + computedByCaller.Store(true) payload, err := compute(ctx) if err != nil { return Value{}, err @@ -230,8 +323,22 @@ func (r *Runtime) LoadOrCompute( } } if expected.Cacheable != nil && !expected.Cacheable(payload) { + r.emit(Event{ + Kind: EventBypass, + Lookup: expected.Key.Lookup, + OutputSchema: expected.Key.OutputSchema, + Reason: "output_not_cacheable", + ProviderCall: true, + }) return uncachedValue(payload, expected.Validate, nil) } + r.emit(Event{ + Kind: EventComputed, + Lookup: expected.Key.Lookup, + OutputSchema: expected.Key.OutputSchema, + Reason: "provider_success", + ProviderCall: true, + }) candidate, err := NewInlineArtifact(expected.Key, expected.Codec, payload) if err != nil { return Value{}, err @@ -246,6 +353,15 @@ func (r *Runtime) LoadOrCompute( if completed.Err != nil { return Value{}, completed.Err } + if completed.Shared && !computedByCaller.Load() { + r.emit(Event{ + Kind: EventWait, + Lookup: expected.Key.Lookup, + OutputSchema: expected.Key.OutputSchema, + Reason: "singleflight", + SingleflightWaitMS: time.Since(waitStarted).Milliseconds(), + }) + } return completed.Val.(Value), nil } } @@ -261,7 +377,14 @@ func (r *Runtime) BatchFreeze( manifests := make([]*types.ProcessingArtifact, 0, len(candidates)) byKey := make(map[types.ProcessingArtifactLookup]Candidate, len(candidates)) for _, candidate := range candidates { + key := candidate.Expected.Key.Lookup if candidate.Expected.Cacheable != nil && !candidate.Expected.Cacheable(candidate.Payload) { + r.emit(Event{ + Kind: EventBypass, + Lookup: key, + OutputSchema: candidate.Expected.Key.OutputSchema, + Reason: "output_not_cacheable", + }) continue } if candidate.Expected.Validate != nil { @@ -277,7 +400,7 @@ func (r *Runtime) BatchFreeze( if err != nil { continue } - key := manifest.Lookup() + key = manifest.Lookup() manifests = append(manifests, manifest) byKey[key] = candidate result[key] = Value{ @@ -289,10 +412,36 @@ func (r *Runtime) BatchFreeze( return result } + filtered := manifests[:0] + for _, manifest := range manifests { + key := manifest.Lookup() + if !r.writeEnabled(key.Stage) { + r.emit(Event{ + Kind: EventBypass, + Lookup: key, + OutputSchema: byKey[key].Expected.Key.OutputSchema, + Reason: "write_disabled", + }) + continue + } + filtered = append(filtered, manifest) + } + manifests = filtered + if len(manifests) == 0 { + return result + } + winners, err := r.repository.PutManyIfAbsent(ctx, manifests) if err != nil { for _, manifest := range manifests { - r.emit(Event{Kind: EventStoreFailure, Lookup: manifest.Lookup(), Err: err}) + key := manifest.Lookup() + r.emit(Event{ + Kind: EventStoreFailure, + Lookup: key, + OutputSchema: byKey[key].Expected.Key.OutputSchema, + Reason: "write_error", + Err: err, + }) } return result } @@ -302,9 +451,11 @@ func (r *Runtime) BatchFreeze( candidate := byKey[key] if winner == nil { r.emit(Event{ - Kind: EventStoreFailure, - Lookup: key, - Err: errors.New("artifact repository omitted an inserted winner"), + Kind: EventStoreFailure, + Lookup: key, + OutputSchema: candidate.Expected.Key.OutputSchema, + Reason: "winner_missing", + Err: errors.New("artifact repository omitted an inserted winner"), }) continue } @@ -318,10 +469,22 @@ func (r *Runtime) BatchFreeze( decodeErr = candidate.Expected.Validate(payload) } if decodeErr != nil { - r.emit(Event{Kind: EventCorrupt, Lookup: key, Err: decodeErr}) + r.emit(Event{ + Kind: EventCorrupt, + Lookup: key, + OutputSchema: candidate.Expected.Key.OutputSchema, + Reason: corruptReason(decodeErr), + Err: decodeErr, + }) if winner != nil { if deleteErr := r.repository.DeleteCorrupt(ctx, key, winner.PayloadChecksum); deleteErr != nil { - r.emit(Event{Kind: EventStoreFailure, Lookup: key, Err: deleteErr}) + r.emit(Event{ + Kind: EventStoreFailure, + Lookup: key, + OutputSchema: candidate.Expected.Key.OutputSchema, + Reason: "corrupt_delete_error", + Err: deleteErr, + }) } } continue @@ -344,19 +507,35 @@ func (r *Runtime) freeze( Payload: append([]byte(nil), candidate.Payload...), OutputDigest: candidate.OutputDigest, } - if r.repository == nil { + if r.repository == nil || !r.writeEnabled(candidate.Stage) { + if r != nil && r.repository != nil { + r.emit(Event{ + Kind: EventBypass, + Lookup: candidate.Lookup(), + OutputSchema: expected.Key.OutputSchema, + Reason: "write_disabled", + }) + } return fallback } winner, created, err := r.repository.PutIfAbsent(ctx, candidate) if err != nil { - r.emit(Event{Kind: EventStoreFailure, Lookup: candidate.Lookup(), Err: err}) + r.emit(Event{ + Kind: EventStoreFailure, + Lookup: candidate.Lookup(), + OutputSchema: expected.Key.OutputSchema, + Reason: "write_error", + Err: err, + }) return fallback } if winner == nil { r.emit(Event{ - Kind: EventStoreFailure, - Lookup: candidate.Lookup(), - Err: errors.New("artifact repository returned a nil winner"), + Kind: EventStoreFailure, + Lookup: candidate.Lookup(), + OutputSchema: expected.Key.OutputSchema, + Reason: "winner_missing", + Err: errors.New("artifact repository returned a nil winner"), }) return fallback } @@ -370,16 +549,38 @@ func (r *Runtime) freeze( err = expected.Validate(payload) } if err != nil { - r.emit(Event{Kind: EventCorrupt, Lookup: candidate.Lookup(), Err: err}) + r.emit(Event{ + Kind: EventCorrupt, + Lookup: candidate.Lookup(), + OutputSchema: expected.Key.OutputSchema, + Reason: corruptReason(err), + Err: err, + }) if deleteErr := r.repository.DeleteCorrupt(ctx, winner.Lookup(), winner.PayloadChecksum); deleteErr != nil { - r.emit(Event{Kind: EventStoreFailure, Lookup: winner.Lookup(), Err: deleteErr}) + r.emit(Event{ + Kind: EventStoreFailure, + Lookup: winner.Lookup(), + OutputSchema: expected.Key.OutputSchema, + Reason: "corrupt_delete_error", + Err: deleteErr, + }) } return fallback } if created { - r.emit(Event{Kind: EventStored, Lookup: candidate.Lookup()}) + r.emit(Event{ + Kind: EventStored, + Lookup: candidate.Lookup(), + OutputSchema: expected.Key.OutputSchema, + Reason: "stored", + }) } else { - r.emit(Event{Kind: EventLostRace, Lookup: candidate.Lookup()}) + r.emit(Event{ + Kind: EventLostRace, + Lookup: candidate.Lookup(), + OutputSchema: expected.Key.OutputSchema, + Reason: "immutable_winner", + }) } return Value{ Payload: payload, @@ -413,3 +614,45 @@ func (r *Runtime) emit(event Event) { r.observer(event) } } + +// Observe records adapter-level batch metrics through the same safe observer +// used by the runtime. +func (r *Runtime) Observe(event Event) { + r.emit(event) +} + +func (r *Runtime) readEnabled(stage string) bool { + if r == nil { + return false + } + r.configMu.RLock() + defer r.configMu.RUnlock() + return r.read && r.stageEnabledLocked(stage) +} + +func (r *Runtime) writeEnabled(stage string) bool { + if r == nil { + return false + } + r.configMu.RLock() + defer r.configMu.RUnlock() + return r.write && r.stageEnabledLocked(stage) +} + +func (r *Runtime) stageEnabledLocked(stage string) bool { + enabled, configured := r.stages[stage] + return !configured || enabled +} + +func corruptReason(err error) string { + switch { + case err == nil: + return "decode_failed" + case strings.Contains(err.Error(), "checksum"): + return "checksum_mismatch" + case strings.Contains(err.Error(), "schema"): + return "schema_mismatch" + default: + return "decode_failed" + } +} diff --git a/internal/artifact/runtime_test.go b/internal/artifact/runtime_test.go index 1208d065df..d69162f8bb 100644 --- a/internal/artifact/runtime_test.go +++ b/internal/artifact/runtime_test.go @@ -227,3 +227,112 @@ func TestRuntimeRedisLeaseSuppressesCrossProcessDuplicateCompute(t *testing.T) { require.NoError(t, <-results) assert.Equal(t, int32(1), calls.Load()) } + +func TestRuntimeCacheModes(t *testing.T) { + t.Run("shadow write does not read", func(t *testing.T) { + repository := newMemoryArtifactRepository() + runtime := NewRuntime(repository, nil) + runtime.ConfigureCacheMode(false, true, nil) + expected := testExpected(t) + var calls int + + for index := 0; index < 2; index++ { + _, err := runtime.LoadOrCompute( + context.Background(), + expected, + func(context.Context) ([]byte, error) { + calls++ + return []byte(`{"value":"computed"}`), nil + }, + ) + require.NoError(t, err) + } + + assert.Equal(t, 2, calls) + assert.Len(t, repository.values, 1) + }) + + t.Run("read only does not write misses", func(t *testing.T) { + repository := newMemoryArtifactRepository() + runtime := NewRuntime(repository, nil) + runtime.ConfigureCacheMode(true, false, nil) + + _, err := runtime.LoadOrCompute( + context.Background(), + testExpected(t), + func(context.Context) ([]byte, error) { + return []byte(`{"value":"computed"}`), nil + }, + ) + require.NoError(t, err) + assert.Empty(t, repository.values) + }) + + t.Run("disabled stage bypasses reads and writes", func(t *testing.T) { + repository := newMemoryArtifactRepository() + runtime := NewRuntime(repository, nil) + runtime.ConfigureCacheMode(true, true, map[string]bool{"summary": false}) + expected := testExpected(t) + var calls int + + for index := 0; index < 2; index++ { + _, err := runtime.LoadOrCompute( + context.Background(), + expected, + func(context.Context) ([]byte, error) { + calls++ + return []byte(`{"value":"computed"}`), nil + }, + ) + require.NoError(t, err) + } + + assert.Equal(t, 2, calls) + assert.Empty(t, repository.values) + }) +} + +func TestRuntimeObserverRecordsSingleflightWaitWithoutExposingPayload(t *testing.T) { + repository := newMemoryArtifactRepository() + var eventsMu sync.Mutex + var events []Event + runtime := NewRuntime(repository, func(event Event) { + eventsMu.Lock() + defer eventsMu.Unlock() + events = append(events, event) + }) + expected := testExpected(t) + + start := make(chan struct{}) + results := make(chan error, 2) + for index := 0; index < 2; index++ { + go func() { + <-start + _, err := runtime.LoadOrCompute( + context.Background(), + expected, + func(context.Context) ([]byte, error) { + time.Sleep(20 * time.Millisecond) + return []byte(`{"value":"computed"}`), nil + }, + ) + results <- err + }() + } + close(start) + require.NoError(t, <-results) + require.NoError(t, <-results) + + eventsMu.Lock() + defer eventsMu.Unlock() + assert.Condition(t, func() bool { + for _, event := range events { + if event.Kind == EventWait && + event.Reason == "singleflight" && + event.SingleflightWaitMS >= 0 { + return true + } + } + return false + }) +} diff --git a/internal/config/config.go b/internal/config/config.go index 769d9fd03b..d1557e88ec 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -33,6 +33,7 @@ type Config struct { PromptTemplates *PromptTemplatesConfig `yaml:"prompt_templates" json:"prompt_templates"` IM *IMConfig `yaml:"im" json:"im"` Agent *AgentConfig `yaml:"agent" json:"agent"` + ArtifactCache *ArtifactCacheConfig `yaml:"artifact_cache" json:"artifact_cache"` // FrontendBaseURL is the externally-visible origin of the SPA, used // to compose absolute share-link URLs. Empty falls back to a host- // relative URL ("/register?token=…") which the SPA then resolves @@ -41,6 +42,14 @@ type Config struct { FrontendBaseURL string `yaml:"frontend_base_url" json:"frontend_base_url"` } +// ArtifactCacheConfig controls reusable ingestion artifacts. A nil config +// keeps reads and writes enabled for backward-compatible defaults. +type ArtifactCacheConfig struct { + ReadEnabled bool `yaml:"read_enabled" json:"read_enabled"` + WriteEnabled bool `yaml:"write_enabled" json:"write_enabled"` + Stages map[string]bool `yaml:"stages" json:"stages"` +} + // AgentConfig represents the global agent settings. type AgentConfig struct { // LLMCallTimeout is the default timeout for a single LLM call in seconds. @@ -489,6 +498,9 @@ func ConfigDir() string { // LoadConfig 从配置文件加载配置 func LoadConfig() (*Config, error) { + viper.SetDefault("artifact_cache.read_enabled", true) + viper.SetDefault("artifact_cache.write_enabled", true) + // 设置配置文件名和路径 viper.SetConfigName("config") // 配置文件名称(不带扩展名) viper.SetConfigType("yaml") // 配置文件类型 diff --git a/internal/container/artifact_runtime.go b/internal/container/artifact_runtime.go new file mode 100644 index 0000000000..f62cbdbc5b --- /dev/null +++ b/internal/container/artifact_runtime.go @@ -0,0 +1,67 @@ +package container + +import ( + "context" + "fmt" + + "github.com/Tencent/WeKnora/internal/artifact" + "github.com/Tencent/WeKnora/internal/logger" +) + +func newArtifactObserver() artifact.Observer { + return func(event artifact.Event) { + fields := artifactEventFields(event) + entry := logger.GetLogger(context.Background()).WithFields(fields) + switch event.Kind { + case artifact.EventCorrupt, artifact.EventStoreFailure: + entry.Warn("processing artifact cache fallback") + default: + entry.Debug("processing artifact cache") + } + } +} + +func artifactEventFields(event artifact.Event) logger.Fields { + fields := logger.Fields{ + "artifact_stage": event.Lookup.Stage, + "cache_outcome": artifactCacheOutcome(event.Kind), + "cache_reason": event.Reason, + "key_version": event.Lookup.KeyVersion, + "output_schema_version": event.OutputSchema, + "provider_call": event.ProviderCall, + "singleflight_wait_ms": event.SingleflightWaitMS, + } + if event.BatchTotal > 0 { + fields["batch_total"] = event.BatchTotal + fields["batch_hits"] = event.BatchHits + fields["batch_misses"] = event.BatchMisses + fields["batch_deduplicated"] = event.BatchDeduplicated + } + if event.Err != nil { + // Error values from database drivers may include SQL arguments. Log + // only the concrete class so keys and payloads cannot leak. + fields["cache_error_class"] = fmt.Sprintf("%T", event.Err) + } + return fields +} + +func artifactCacheOutcome(kind artifact.EventKind) string { + switch kind { + case artifact.EventHit: + return "hit" + case artifact.EventMiss: + return "miss" + case artifact.EventComputed, artifact.EventStored: + return "computed" + case artifact.EventWait, artifact.EventLostRace: + return "wait" + case artifact.EventBypass: + return "bypass" + case artifact.EventCorrupt: + return "corrupt" + case artifact.EventStoreFailure: + return "error_fallback" + default: + return "error_fallback" + } +} diff --git a/internal/container/artifact_runtime_test.go b/internal/container/artifact_runtime_test.go new file mode 100644 index 0000000000..d0e5d23c66 --- /dev/null +++ b/internal/container/artifact_runtime_test.go @@ -0,0 +1,47 @@ +package container + +import ( + "errors" + "fmt" + "testing" + + "github.com/Tencent/WeKnora/internal/artifact" + "github.com/Tencent/WeKnora/internal/types" + "github.com/stretchr/testify/assert" +) + +func TestArtifactEventFieldsExcludeKeysPayloadsAndErrorMessages(t *testing.T) { + fields := artifactEventFields(artifact.Event{ + Kind: artifact.EventStoreFailure, + Lookup: types.ProcessingArtifactLookup{ + TenantID: 7, + Stage: "embedding", + KeyVersion: 1, + ArtifactKey: "full-secret-equivalence-key", + }, + OutputSchema: "embedding.float32.3.v1", + Reason: "write_error", + Err: errors.New("driver exposed full-secret-equivalence-key and prompt body"), + }) + + assert.Equal(t, "embedding", fields["artifact_stage"]) + assert.Equal(t, "error_fallback", fields["cache_outcome"]) + assert.Equal(t, "*errors.errorString", fields["cache_error_class"]) + assert.NotContains(t, fields, "artifact_key") + assert.NotContains(t, fields, "tenant_id") + for _, value := range fields { + rendered := fmt.Sprint(value) + assert.NotContains(t, rendered, "full-secret-equivalence-key") + assert.NotContains(t, rendered, "prompt body") + } +} + +func TestArtifactCacheOutcomeUsesRequiredVocabulary(t *testing.T) { + assert.Equal(t, "hit", artifactCacheOutcome(artifact.EventHit)) + assert.Equal(t, "miss", artifactCacheOutcome(artifact.EventMiss)) + assert.Equal(t, "computed", artifactCacheOutcome(artifact.EventComputed)) + assert.Equal(t, "wait", artifactCacheOutcome(artifact.EventWait)) + assert.Equal(t, "bypass", artifactCacheOutcome(artifact.EventBypass)) + assert.Equal(t, "corrupt", artifactCacheOutcome(artifact.EventCorrupt)) + assert.Equal(t, "error_fallback", artifactCacheOutcome(artifact.EventStoreFailure)) +} diff --git a/internal/container/container.go b/internal/container/container.go index 6245106fb4..595f693ce5 100644 --- a/internal/container/container.go +++ b/internal/container/container.go @@ -151,8 +151,18 @@ func BuildContainer(container *dig.Container) *dig.Container { must(container.Provide(func( repo interfaces.ProcessingArtifactRepository, redisClient *redis.Client, + cfg *config.Config, ) *artifact.Runtime { - runtime := artifact.NewRuntime(repo, nil) + readEnabled := true + writeEnabled := true + var stages map[string]bool + if cfg != nil && cfg.ArtifactCache != nil { + readEnabled = cfg.ArtifactCache.ReadEnabled + writeEnabled = cfg.ArtifactCache.WriteEnabled + stages = cfg.ArtifactCache.Stages + } + runtime := artifact.NewRuntime(repo, newArtifactObserver()) + runtime.ConfigureCacheMode(readEnabled, writeEnabled, stages) runtime.ConfigureLease(artifact.NewRedisLease(redisClient)) return runtime })) diff --git a/internal/models/embedding/artifact_cache.go b/internal/models/embedding/artifact_cache.go index 9fe3a0da5d..ad9c21e71a 100644 --- a/internal/models/embedding/artifact_cache.go +++ b/internal/models/embedding/artifact_cache.go @@ -173,51 +173,184 @@ func (e *artifactCachedEmbedder) batchEmbed( } uniqueVectors[index] = vector } + initialMisses := len(missingIndexes) + providerCalled := false - if len(missingIndexes) > 0 { - missingTexts := make([]string, len(missingIndexes)) - for index, uniquePosition := range missingIndexes { - missingTexts[index] = uniqueTexts[uniquePosition] + for len(missingIndexes) > 0 { + // Every worker picks the same deterministic missing key as the batch + // leader. LoadOrCompute then applies both the in-process singleflight + // and the cross-process Redis lease to the entire provider batch. + leader := missingIndexes[0] + for _, uniquePosition := range missingIndexes[1:] { + if unique[uniquePosition].Key.Lookup.ArtifactKey < + unique[leader].Key.Lookup.ArtifactKey { + leader = uniquePosition + } } - vectors, err := provider(ctx, missingTexts) + computed := make(map[int][]float32, len(missingIndexes)) + + leaderValue, err := e.runtime.LoadOrCompute( + ctx, + unique[leader], + func(computeContext context.Context) ([]byte, error) { + pendingExpected := make([]artifact.Expected, len(missingIndexes)) + for index, uniquePosition := range missingIndexes { + pendingExpected[index] = unique[uniquePosition] + } + latest := e.runtime.BatchLoad(computeContext, pendingExpected) + + providerPositions := make([]int, 0, len(missingIndexes)) + var leaderPayload []byte + for _, uniquePosition := range missingIndexes { + value, found := latest[unique[uniquePosition].Key.Lookup] + if !found { + providerPositions = append(providerPositions, uniquePosition) + continue + } + vector, decodeErr := decodeEmbeddingArtifact( + value.Payload, + e.config.Dimensions, + ) + if decodeErr != nil { + providerPositions = append(providerPositions, uniquePosition) + continue + } + computed[uniquePosition] = vector + if uniquePosition == leader { + leaderPayload = value.Payload + } + } + + if len(providerPositions) > 0 { + missingTexts := make([]string, len(providerPositions)) + for index, uniquePosition := range providerPositions { + missingTexts[index] = uniqueTexts[uniquePosition] + } + vectors, providerErr := provider(computeContext, missingTexts) + if providerErr != nil { + return nil, providerErr + } + providerCalled = true + if validateErr := ValidateEmbeddingBatch( + vectors, + len(missingTexts), + e.config.Dimensions, + ); validateErr != nil { + return nil, validateErr + } + + candidates := make([]artifact.Candidate, 0, len(providerPositions)-1) + candidatePositions := make([]int, 0, len(providerPositions)-1) + for index, uniquePosition := range providerPositions { + payload, encodeErr := encodeEmbeddingArtifact( + vectors[index], + e.config.Dimensions, + ) + if encodeErr != nil { + return nil, encodeErr + } + if uniquePosition == leader { + leaderPayload = payload + continue + } + computed[uniquePosition] = append([]float32(nil), vectors[index]...) + candidates = append(candidates, artifact.Candidate{ + Expected: unique[uniquePosition], + Payload: payload, + }) + candidatePositions = append(candidatePositions, uniquePosition) + } + + frozen := e.runtime.BatchFreeze(computeContext, candidates) + for _, uniquePosition := range candidatePositions { + value, found := frozen[unique[uniquePosition].Key.Lookup] + if !found { + continue + } + vector, decodeErr := decodeEmbeddingArtifact( + value.Payload, + e.config.Dimensions, + ) + if decodeErr == nil { + computed[uniquePosition] = vector + } + } + } + + if len(leaderPayload) == 0 { + return nil, errors.New("embedding batch leader did not produce a payload") + } + return leaderPayload, nil + }, + ) if err != nil { return nil, err } - if err := ValidateEmbeddingBatch(vectors, len(missingTexts), e.config.Dimensions); err != nil { + leaderVector, err := decodeEmbeddingArtifact( + leaderValue.Payload, + e.config.Dimensions, + ) + if err != nil { return nil, err } + uniqueVectors[leader] = leaderVector - candidates := make([]artifact.Candidate, 0, len(missingIndexes)) - for index, uniquePosition := range missingIndexes { - payload, err := encodeEmbeddingArtifact(vectors[index], e.config.Dimensions) - if err != nil { - return nil, err + remainingExpected := make([]artifact.Expected, 0, len(missingIndexes)-1) + for _, uniquePosition := range missingIndexes { + if uniquePosition != leader { + remainingExpected = append(remainingExpected, unique[uniquePosition]) } - candidates = append(candidates, artifact.Candidate{ - Expected: unique[uniquePosition], - Payload: payload, - }) } - frozen := e.runtime.BatchFreeze(ctx, candidates) - for index, uniquePosition := range missingIndexes { - value, found := frozen[unique[uniquePosition].Key.Lookup] - if !found { - uniqueVectors[uniquePosition] = append([]float32(nil), vectors[index]...) + latest := e.runtime.BatchLoad(ctx, remainingExpected) + nextMissing := make([]int, 0, len(remainingExpected)) + for _, uniquePosition := range missingIndexes { + if uniquePosition == leader { continue } - vector, err := decodeEmbeddingArtifact(value.Payload, e.config.Dimensions) - if err != nil { - uniqueVectors[uniquePosition] = append([]float32(nil), vectors[index]...) + value, found := latest[unique[uniquePosition].Key.Lookup] + if found { + vector, decodeErr := decodeEmbeddingArtifact( + value.Payload, + e.config.Dimensions, + ) + if decodeErr == nil { + uniqueVectors[uniquePosition] = vector + continue + } + } + if vector, found := computed[uniquePosition]; found { + uniqueVectors[uniquePosition] = append([]float32(nil), vector...) continue } - uniqueVectors[uniquePosition] = vector + nextMissing = append(nextMissing, uniquePosition) } + missingIndexes = nextMissing } result := make([][]float32, len(texts)) for index, key := range keys { result[index] = append([]float32(nil), uniqueVectors[uniqueIndex[key]]...) } + outcome := artifact.EventHit + reason := "batch_hit" + if providerCalled { + outcome = artifact.EventComputed + reason = "batch_computed" + } else if initialMisses > 0 { + outcome = artifact.EventWait + reason = "batch_filled_by_concurrent_worker" + } + e.runtime.Observe(artifact.Event{ + Kind: outcome, + Lookup: unique[0].Key.Lookup, + OutputSchema: unique[0].Key.OutputSchema, + Reason: reason, + ProviderCall: providerCalled, + BatchTotal: len(texts), + BatchHits: len(unique) - initialMisses, + BatchMisses: initialMisses, + BatchDeduplicated: len(texts) - len(unique), + }) return result, nil } diff --git a/internal/models/embedding/artifact_cache_test.go b/internal/models/embedding/artifact_cache_test.go index c3931a0875..74304d4789 100644 --- a/internal/models/embedding/artifact_cache_test.go +++ b/internal/models/embedding/artifact_cache_test.go @@ -5,10 +5,13 @@ import ( "math" "sync" "testing" + "time" "github.com/Tencent/WeKnora/internal/artifact" "github.com/Tencent/WeKnora/internal/testutil/artifactrepo" "github.com/Tencent/WeKnora/internal/types" + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -19,6 +22,7 @@ type countingEmbedder struct { dimensions int shortResult bool invalid float32 + delay time.Duration } func (e *countingEmbedder) Embed(ctx context.Context, text string) ([]float32, error) { @@ -33,6 +37,9 @@ func (e *countingEmbedder) BatchEmbed(_ context.Context, texts []string) ([][]fl e.mu.Lock() e.batches = append(e.batches, append([]string(nil), texts...)) e.mu.Unlock() + if e.delay > 0 { + time.Sleep(e.delay) + } count := len(texts) if e.shortResult && count > 0 { count-- @@ -134,6 +141,56 @@ func TestArtifactCachedEmbeddingUsesPartialHits(t *testing.T) { assert.Equal(t, [][]string{{"missing"}}, provider.calls()) } +func TestArtifactCachedEmbeddingBatchSuppressesConcurrentProviderCalls(t *testing.T) { + server := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: server.Addr()}) + t.Cleanup(func() { require.NoError(t, client.Close()) }) + + repository := artifactrepo.New() + firstRuntime := artifact.NewRuntime(repository, nil) + secondRuntime := artifact.NewRuntime(repository, nil) + firstRuntime.ConfigureLease(artifact.NewRedisLease(client)) + secondRuntime.ConfigureLease(artifact.NewRedisLease(client)) + + provider := &countingEmbedder{dimensions: 3, delay: 75 * time.Millisecond} + cached := []Embedder{ + NewArtifactCachedEmbedder(provider, firstRuntime, embeddingArtifactConfig()), + NewArtifactCachedEmbedder(provider, secondRuntime, embeddingArtifactConfig()), + } + inputs := []string{"alpha", "beta", "alpha", "gamma"} + + const workers = 10 + start := make(chan struct{}) + results := make(chan [][]float32, workers) + errors := make(chan error, workers) + for index := 0; index < workers; index++ { + go func(worker int) { + <-start + vectors, err := cached[worker%len(cached)].BatchEmbed( + documentEmbeddingContext(), + inputs, + ) + results <- vectors + errors <- err + }(index) + } + close(start) + + var first [][]float32 + for index := 0; index < workers; index++ { + require.NoError(t, <-errors) + vectors := <-results + require.Len(t, vectors, len(inputs)) + assert.Equal(t, vectors[0], vectors[2]) + if index == 0 { + first = vectors + continue + } + assert.Equal(t, first, vectors) + } + assert.Equal(t, [][]string{{"alpha", "beta", "gamma"}}, provider.calls()) +} + func TestArtifactCachedEmbeddingBypassesInteractiveQuery(t *testing.T) { runtime, repository := setupEmbeddingArtifactRuntime(t) provider := &countingEmbedder{dimensions: 3} diff --git a/migrations/versioned/000077_processing_artifacts.up.sql b/migrations/versioned/000077_processing_artifacts.up.sql index 4af80224c5..34d71299db 100644 --- a/migrations/versioned/000077_processing_artifacts.up.sql +++ b/migrations/versioned/000077_processing_artifacts.up.sql @@ -41,4 +41,4 @@ CREATE TABLE IF NOT EXISTS knowledge_attempt_counters ( CHECK (last_attempt >= 0) ); -DO $$ BEGIN RAISE NOTICE '[Migration 000075] Processing artifacts ready'; END $$; +DO $$ BEGIN RAISE NOTICE '[Migration 000077] Processing artifacts ready'; END $$; From 5bd8be887f30ff05ad134270bb074fbe7224b0ca Mon Sep 17 00:00:00 2001 From: Even1115 Date: Sun, 2 Aug 2026 23:07:33 +0800 Subject: [PATCH 03/11] fix: complete artifact DAG recovery validation --- docs/dev/artifact-dag-reconciliation.md | 32 +- internal/application/repository/knowledge.go | 92 ++++ .../knowledge_attempt_fence_test.go | 72 +++ .../service/artifact_fault_recovery_test.go | 287 ++++++++++++ .../service/artifact_pipeline_e2e_test.go | 410 ++++++++++++++++++ internal/application/service/extract.go | 50 +-- .../application/service/image_multimodal.go | 83 ++-- .../service/image_multimodal_prompt_test.go | 24 + internal/application/service/knowledge.go | 5 +- .../service/knowledge_clone_move.go | 2 +- .../application/service/knowledge_create.go | 14 +- .../application/service/knowledge_process.go | 259 ++++++----- .../service/knowledge_reconcile_lock.go | 165 +++++++ .../service/knowledge_reconcile_lock_test.go | 96 ++++ .../service/knowledge_span_tracker_test.go | 31 ++ .../application/service/wiki_ingest_batch.go | 85 ++-- internal/artifact/fault.go | 46 ++ internal/artifact/runtime.go | 20 +- internal/artifact/runtime_benchmark_test.go | 70 +++ internal/artifact/runtime_test.go | 55 +++ internal/database/migration_version_test.go | 52 +++ .../processing_artifact_migration_test.go | 4 +- internal/models/chat/artifact_cache_test.go | 23 + internal/models/embedding/artifact_cache.go | 6 + .../models/embedding/artifact_cache_test.go | 18 + internal/models/vlm/artifact_cache_test.go | 6 +- ...l => 000002_processing_artifacts.down.sql} | 1 + ...sql => 000002_processing_artifacts.up.sql} | 1 + ...l => 000079_processing_artifacts.down.sql} | 2 +- ...sql => 000079_processing_artifacts.up.sql} | 4 +- 30 files changed, 1749 insertions(+), 266 deletions(-) create mode 100644 internal/application/service/artifact_fault_recovery_test.go create mode 100644 internal/application/service/artifact_pipeline_e2e_test.go create mode 100644 internal/application/service/knowledge_reconcile_lock.go create mode 100644 internal/application/service/knowledge_reconcile_lock_test.go create mode 100644 internal/artifact/fault.go create mode 100644 internal/artifact/runtime_benchmark_test.go create mode 100644 internal/database/migration_version_test.go rename migrations/sqlite/{000001_processing_artifacts.down.sql => 000002_processing_artifacts.down.sql} (52%) rename migrations/sqlite/{000001_processing_artifacts.up.sql => 000002_processing_artifacts.up.sql} (95%) rename migrations/versioned/{000077_processing_artifacts.down.sql => 000079_processing_artifacts.down.sql} (52%) rename migrations/versioned/{000077_processing_artifacts.up.sql => 000079_processing_artifacts.up.sql} (92%) diff --git a/docs/dev/artifact-dag-reconciliation.md b/docs/dev/artifact-dag-reconciliation.md index 9b09dbddf1..d4a9b6ae52 100644 --- a/docs/dev/artifact-dag-reconciliation.md +++ b/docs/dev/artifact-dag-reconciliation.md @@ -32,8 +32,8 @@ source bytes -> desired-state diff -> add/update/index -> attempt fence + -> conditional knowledge publication + storage accounting -> delete stale vector, graph, and chunk state - -> conditional knowledge publication ``` `internal/artifact` owns canonical keys, codecs, payload validation, immutable @@ -61,8 +61,8 @@ cleaned up. Migration locations: -- PostgreSQL: `migrations/versioned/000077_processing_artifacts.{up,down}.sql` -- SQLite: `migrations/sqlite/000001_processing_artifacts.{up,down}.sql` +- PostgreSQL: `migrations/versioned/000079_processing_artifacts.{up,down}.sql` +- SQLite: `migrations/sqlite/000002_processing_artifacts.{up,down}.sql` - MySQL bootstrap: `migrations/mysql/00-init-db.sql` - ParadeDB bootstrap: `migrations/paradedb/00-init-db.sql` @@ -80,9 +80,14 @@ frozen together. This keeps one provider call for concurrent identical batches without changing caller output order. Knowledge processing allocates an attempt before work begins. Final publication -and destructive stale cleanup recheck that attempt. Graph storage uses -per-chunk contributions so stale chunks can be removed exactly without deleting -unchanged contributions. +and destructive stale cleanup recheck that attempt. A per-knowledge mutation +lock spans chunk/vector/graph binding through exact stale cleanup (local gate in +Lite mode, ownership-token Redis lease in standard mode), preventing a newer +generation from being deleted between a fence check and an external-store +delete. Knowledge publication and tenant storage accounting share one database +transaction; retries calculate the delta from the already-published row and +therefore cannot double-charge. Graph storage uses per-chunk contributions so +stale chunks can be removed exactly without deleting unchanged contributions. ## Compatibility and rollback @@ -109,7 +114,7 @@ To roll back: 1. disable artifact reads, then writes, and restart or drain workers; 2. deploy the previous application version; 3. verify no process reads or writes artifact/attempt tables; -4. optionally run migration `000077` down. +4. optionally run migration `000079` down. Dropping the new tables discards only reusable artifacts and attempt counters; it does not delete knowledge, chunks, vectors, or graph data. Do not run the @@ -122,9 +127,10 @@ version, output schema, provider calls, singleflight wait time, and embedding batch totals/hits/misses/deduplication. Outcomes use `hit`, `miss`, `computed`, `wait`, `bypass`, `corrupt`, and `error_fallback`. -Logs intentionally omit tenant IDs, complete artifact keys, payloads, prompts, -URLs, and raw database error strings. Database errors are represented only by -their concrete error class because driver messages can contain SQL arguments. +The artifact observer never receives request or payload bytes. New DAG logs and +span fields omit complete artifact keys, bodies, prompts, file/image references, +and provider error details; they retain IDs, booleans, lengths, counts, and +concrete error classes. ## Validation @@ -137,7 +143,12 @@ The implementation has focused tests for: - exact input bytes, duplicate input ordering, partial batch hits, and invalid provider response rejection; - stable chunk/generated IDs, desired-state diffs, and stale-attempt fencing; +- all eight crash boundaries, with final DB/vector/Wiki/Graph snapshots equal + to a clean run after retry; +- local and Redis per-knowledge mutation-lock ownership and wait behavior; +- atomic, idempotent knowledge publication plus tenant storage accounting; - SQLite migration up/down and uniqueness behavior; +- duplicate migration-version rejection for PostgreSQL and SQLite directories; - DocReader, chat, embedding, VLM, wiki, multimodal, and graph stage adapters. Run the focused suite with: @@ -155,6 +166,7 @@ go test -count=1 \ go test -race -count=1 \ ./internal/artifact \ ./internal/application/repository \ + ./internal/application/service \ ./internal/models/embedding ``` diff --git a/internal/application/repository/knowledge.go b/internal/application/repository/knowledge.go index a59b90e554..7c42e6ac92 100644 --- a/internal/application/repository/knowledge.go +++ b/internal/application/repository/knowledge.go @@ -9,6 +9,7 @@ import ( "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" "gorm.io/gorm" + "gorm.io/gorm/clause" ) var ErrKnowledgeNotFound = errors.New("knowledge not found") @@ -230,6 +231,97 @@ func (r *knowledgeRepository) UpdateKnowledgeIfAttemptCurrent( return result.RowsAffected == 1, nil } +// PublishKnowledgeIfAttemptCurrent atomically publishes the desired knowledge +// generation and applies its tenant-storage delta. Reading the currently +// published storage size inside the same transaction makes a retry after a +// worker crash idempotent: a repeated publish computes a zero delta. +func (r *knowledgeRepository) PublishKnowledgeIfAttemptCurrent( + ctx context.Context, + knowledge *types.Knowledge, + attempt int, +) (published bool, err error) { + if knowledge == nil { + return false, errors.New("knowledge must not be nil") + } + + err = r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var current types.Knowledge + if err := tx. + Clauses(clause.Locking{Strength: "UPDATE"}). + Select("id", "tenant_id", "storage_size"). + Where("tenant_id = ? AND id = ?", knowledge.TenantID, knowledge.ID). + Take(¤t).Error; err != nil { + return err + } + + query := tx. + Model(&types.Knowledge{}). + Where("tenant_id = ? AND id = ?", knowledge.TenantID, knowledge.ID) + if attempt > 0 { + newerAttempt := tx. + Table(types.KnowledgeProcessingSpan{}.TableName()). + Select("1"). + Where("knowledge_id = ? AND kind = ? AND attempt > ?", + knowledge.ID, + types.SpanKindRoot, + attempt, + ) + query = query.Where("NOT EXISTS (?)", newerAttempt) + } + result := query. + Select("*"). + Omit(omitFieldsOnUpdate...). + Updates(knowledge) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != 1 { + if attempt > 0 { + var newerCount int64 + if err := tx. + Table(types.KnowledgeProcessingSpan{}.TableName()). + Where("knowledge_id = ? AND kind = ? AND attempt > ?", + knowledge.ID, + types.SpanKindRoot, + attempt, + ). + Count(&newerCount).Error; err != nil { + return err + } + if newerCount > 0 { + return nil + } + } + // MySQL reports zero affected rows when the target values are + // already identical (notably second-resolution TIMESTAMP fields). + // With no newer attempt, that is an idempotent successful publish. + } + + delta := knowledge.StorageSize - current.StorageSize + if delta != 0 { + var tenant types.Tenant + if err := tx. + Clauses(clause.Locking{Strength: "UPDATE"}). + Where("id = ?", knowledge.TenantID). + Take(&tenant).Error; err != nil { + return err + } + tenant.StorageUsed += delta + if tenant.StorageUsed < 0 { + tenant.StorageUsed = 0 + } + if err := tx.Model(&tenant). + Select("storage_used", "updated_at"). + Updates(&tenant).Error; err != nil { + return err + } + } + published = true + return nil + }) + return published, err +} + func (r *knowledgeRepository) UpdateKnowledgeColumnsIfAttemptCurrent( ctx context.Context, tenantID uint64, diff --git a/internal/application/repository/knowledge_attempt_fence_test.go b/internal/application/repository/knowledge_attempt_fence_test.go index c8fa50e232..a82f444986 100644 --- a/internal/application/repository/knowledge_attempt_fence_test.go +++ b/internal/application/repository/knowledge_attempt_fence_test.go @@ -3,6 +3,7 @@ package repository import ( "context" "testing" + "time" "github.com/Tencent/WeKnora/internal/types" "github.com/google/uuid" @@ -97,3 +98,74 @@ func TestUpdateKnowledgeIfAttemptCurrentIsDatabaseFenced(t *testing.T) { require.NoError(t, db.First(&stored, "id = ?", knowledge.ID).Error) assert.Equal(t, 3, stored.PendingSubtasksCount) } + +func TestPublishKnowledgeIfAttemptCurrentAtomicallyAdjustsStorageOnce(t *testing.T) { + db, err := gorm.Open(sqlite.Open("file:"+uuid.NewString()+"?mode=memory&cache=shared"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&types.Knowledge{}, &types.Tenant{})) + require.NoError(t, db.Exec(` + CREATE TABLE knowledge_processing_spans ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + knowledge_id TEXT NOT NULL, + attempt INTEGER NOT NULL, + kind TEXT NOT NULL + ) + `).Error) + + tenant := &types.Tenant{ID: 7, Name: "tenant", StorageUsed: 1_000} + require.NoError(t, db.Create(tenant).Error) + knowledge := &types.Knowledge{ + ID: uuid.NewString(), + TenantID: tenant.ID, + KnowledgeBaseID: uuid.NewString(), + Type: types.KnowledgeTypeManual, + Source: "manual", + ParseStatus: types.ParseStatusProcessing, + EnableStatus: "enabled", + StorageSize: 20, + } + require.NoError(t, db.Create(knowledge).Error) + repository := &knowledgeRepository{db: db} + + knowledge.StorageSize = 50 + knowledge.UpdatedAt = time.Now() + published, err := repository.PublishKnowledgeIfAttemptCurrent( + context.Background(), + knowledge, + 1, + ) + require.NoError(t, err) + require.True(t, published) + + var storedTenant types.Tenant + require.NoError(t, db.First(&storedTenant, tenant.ID).Error) + assert.Equal(t, int64(1_030), storedTenant.StorageUsed) + + // A worker retry after crashing immediately after publish must not charge + // the same 30-byte generation a second time. + knowledge.UpdatedAt = time.Now().Add(time.Second) + published, err = repository.PublishKnowledgeIfAttemptCurrent( + context.Background(), + knowledge, + 1, + ) + require.NoError(t, err) + require.True(t, published) + require.NoError(t, db.First(&storedTenant, tenant.ID).Error) + assert.Equal(t, int64(1_030), storedTenant.StorageUsed) + + require.NoError(t, db.Exec(` + INSERT INTO knowledge_processing_spans (knowledge_id, attempt, kind) + VALUES (?, 2, ?) + `, knowledge.ID, types.SpanKindRoot).Error) + knowledge.StorageSize = 80 + published, err = repository.PublishKnowledgeIfAttemptCurrent( + context.Background(), + knowledge, + 1, + ) + require.NoError(t, err) + assert.False(t, published) + require.NoError(t, db.First(&storedTenant, tenant.ID).Error) + assert.Equal(t, int64(1_030), storedTenant.StorageUsed) +} diff --git a/internal/application/service/artifact_fault_recovery_test.go b/internal/application/service/artifact_fault_recovery_test.go new file mode 100644 index 0000000000..640cd68ab3 --- /dev/null +++ b/internal/application/service/artifact_fault_recovery_test.go @@ -0,0 +1,287 @@ +package service + +import ( + "context" + "errors" + "sort" + "sync" + "testing" + + "github.com/Tencent/WeKnora/internal/artifact" + "github.com/Tencent/WeKnora/internal/models/embedding" + "github.com/Tencent/WeKnora/internal/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type recoveryDocReader struct { + mu sync.Mutex + calls int +} + +func (r *recoveryDocReader) Read( + _ context.Context, + request *types.ReadRequest, +) (*types.ReadResult, error) { + r.mu.Lock() + r.calls++ + r.mu.Unlock() + return &types.ReadResult{ + MarkdownContent: string(request.FileContent), + Metadata: map[string]string{"pages": "1"}, + }, nil +} + +type faultRecoveryState struct { + chunks map[string]*types.Chunk + vectors map[string][]float32 + wiki map[string]string + graph map[string]string + publishedIDs []string + storageBytes int64 +} + +func newFaultRecoveryState() *faultRecoveryState { + return &faultRecoveryState{ + chunks: make(map[string]*types.Chunk), + vectors: make(map[string][]float32), + wiki: make(map[string]string), + graph: make(map[string]string), + } +} + +func (s *faultRecoveryState) existingChunks() []*types.Chunk { + result := make([]*types.Chunk, 0, len(s.chunks)) + for _, chunk := range s.chunks { + copy := *chunk + result = append(result, ©) + } + sort.Slice(result, func(i, j int) bool { return result[i].ID < result[j].ID }) + return result +} + +type faultRecoverySnapshot struct { + chunks map[string]string + vectors map[string][]float32 + wiki map[string]string + graph map[string]string + publishedIDs []string + storageBytes int64 +} + +func (s *faultRecoveryState) snapshot() faultRecoverySnapshot { + chunks := make(map[string]string, len(s.chunks)) + for id, chunk := range s.chunks { + chunks[id] = chunk.Content + } + vectors := make(map[string][]float32, len(s.vectors)) + for id, vector := range s.vectors { + vectors[id] = append([]float32(nil), vector...) + } + wiki := make(map[string]string, len(s.wiki)) + for id, digest := range s.wiki { + wiki[id] = digest + } + graph := make(map[string]string, len(s.graph)) + for id, digest := range s.graph { + graph[id] = digest + } + return faultRecoverySnapshot{ + chunks: chunks, + vectors: vectors, + wiki: wiki, + graph: graph, + publishedIDs: append([]string(nil), s.publishedIDs...), + storageBytes: s.storageBytes, + } +} + +type faultRecoveryHarness struct { + service *knowledgeService + reader *recoveryDocReader + embedder embedding.Embedder +} + +func newFaultRecoveryHarness(t *testing.T) *faultRecoveryHarness { + t.Helper() + service := setupDocReaderArtifactService(t) + provider := &pipelineCountingEmbedder{} + return &faultRecoveryHarness{ + service: service, + reader: &recoveryDocReader{}, + embedder: embedding.NewArtifactCachedEmbedder( + provider, + service.artifactRuntime, + embedding.ArtifactCacheConfig{ + TenantID: 1, + Processor: artifact.ProcessorIdentity{ + ModelID: "fault-recovery-embedding-v1", + ModelName: "fault-recovery-embedding", + Provider: "counting", + }, + Dimensions: 2, + }, + ), + } +} + +// reconcileFaultRecoveryState uses the production artifact adapters, stable +// chunk identity builder and the same non-destructive ordering as processChunks. +// The in-memory bindings make the final DB/vector/Wiki/Graph state directly +// comparable after each injected crash boundary. +func (h *faultRecoveryHarness) reconcileFaultRecoveryState( + ctx context.Context, + state *faultRecoveryState, + content string, +) error { + ctx = context.WithValue(ctx, types.TenantIDContextKey, uint64(1)) + read, err := h.service.callDocReaderWithArtifact(ctx, h.reader, &types.ReadRequest{ + FileContent: []byte(content), + FileName: "fault-recovery.md", + FileType: "md", + ParserEngine: "counting", + }) + if err != nil { + return err + } + + knowledge := &types.Knowledge{ + ID: "fault-recovery-knowledge", + TenantID: 1, + KnowledgeBaseID: "fault-recovery-kb", + } + desired, err := buildDesiredDocumentChunks( + knowledge, + []types.ParsedChunk{{ + Content: read.MarkdownContent, + Seq: 0, + Start: 0, + End: len(read.MarkdownContent), + ParentIndex: -1, + }}, + nil, + state.existingChunks(), + ) + if err != nil { + return err + } + + for _, chunk := range desired.All { + copy := *chunk + state.chunks[chunk.ID] = © + } + if err := artifact.InjectFault(ctx, artifact.FaultAfterChunkUpsert); err != nil { + return err + } + + inputs := make([]string, len(desired.Text)) + for index, chunk := range desired.Text { + inputs[index] = chunk.EmbeddingContent() + } + embedCtx := context.WithValue(ctx, types.EmbedDocumentContextKey, true) + vectors, err := h.embedder.BatchEmbed(embedCtx, inputs) + if err != nil { + return err + } + for index, chunk := range desired.Text { + state.vectors[chunk.ID] = append([]float32(nil), vectors[index]...) + } + if err := artifact.InjectFault(ctx, artifact.FaultAfterVectorUpsert); err != nil { + return err + } + for _, chunk := range desired.Text { + digest := artifact.SHA256Hex([]byte(chunk.Content)) + state.wiki[chunk.ID] = digest + state.graph[chunk.ID] = digest + } + if err := artifact.InjectFault(ctx, artifact.FaultAfterGraphBinding); err != nil { + return err + } + if err := artifact.InjectFault(ctx, artifact.FaultBeforeFence); err != nil { + return err + } + + state.publishedIDs = chunkIDList(desired.All) + sort.Strings(state.publishedIDs) + state.storageBytes = int64(len(desired.Text) * 2 * 4) + if err := artifact.InjectFault(ctx, artifact.FaultAfterPublish); err != nil { + return err + } + + staleIDs := chunkIDList(desired.Stale) + for _, id := range staleIDs { + delete(state.vectors, id) + } + if len(staleIDs) > 0 { + if err := artifact.InjectFault(ctx, artifact.FaultDuringStaleCleanup); err != nil { + return err + } + } + for _, id := range staleIDs { + delete(state.wiki, id) + delete(state.graph, id) + delete(state.chunks, id) + } + return nil +} + +func TestArtifactPipelineCrashBoundariesConvergeToCleanRun(t *testing.T) { + points := []artifact.FaultPoint{ + artifact.FaultAfterProviderCall, + artifact.FaultAfterArtifactPut, + artifact.FaultAfterChunkUpsert, + artifact.FaultAfterVectorUpsert, + artifact.FaultAfterGraphBinding, + artifact.FaultBeforeFence, + artifact.FaultAfterPublish, + artifact.FaultDuringStaleCleanup, + } + + for _, point := range points { + t.Run(string(point), func(t *testing.T) { + cleanHarness := newFaultRecoveryHarness(t) + clean := newFaultRecoveryState() + require.NoError(t, cleanHarness.reconcileFaultRecoveryState( + context.Background(), + clean, + "old generation", + )) + require.NoError(t, cleanHarness.reconcileFaultRecoveryState( + context.Background(), + clean, + "new generation", + )) + + recoveryHarness := newFaultRecoveryHarness(t) + recovered := newFaultRecoveryState() + require.NoError(t, recoveryHarness.reconcileFaultRecoveryState( + context.Background(), + recovered, + "old generation", + )) + injected := errors.New("injected crash boundary") + faultCtx := artifact.WithFaultInjector( + context.Background(), + func(current artifact.FaultPoint) error { + if current == point { + return injected + } + return nil + }, + ) + err := recoveryHarness.reconcileFaultRecoveryState( + faultCtx, + recovered, + "new generation", + ) + require.ErrorIs(t, err, injected) + + require.NoError(t, recoveryHarness.reconcileFaultRecoveryState( + context.Background(), + recovered, + "new generation", + )) + assert.Equal(t, clean.snapshot(), recovered.snapshot()) + }) + } +} diff --git a/internal/application/service/artifact_pipeline_e2e_test.go b/internal/application/service/artifact_pipeline_e2e_test.go new file mode 100644 index 0000000000..f41841587c --- /dev/null +++ b/internal/application/service/artifact_pipeline_e2e_test.go @@ -0,0 +1,410 @@ +package service + +import ( + "context" + "fmt" + "strings" + "sync" + "testing" + + "github.com/Tencent/WeKnora/internal/artifact" + "github.com/Tencent/WeKnora/internal/models/chat" + "github.com/Tencent/WeKnora/internal/models/embedding" + "github.com/Tencent/WeKnora/internal/models/vlm" + "github.com/Tencent/WeKnora/internal/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type pipelineCountingEmbedder struct { + mu sync.Mutex + batches [][]string +} + +func (e *pipelineCountingEmbedder) Embed(ctx context.Context, text string) ([]float32, error) { + vectors, err := e.BatchEmbed(ctx, []string{text}) + if err != nil { + return nil, err + } + return vectors[0], nil +} + +func (e *pipelineCountingEmbedder) BatchEmbed( + _ context.Context, + texts []string, +) ([][]float32, error) { + e.mu.Lock() + e.batches = append(e.batches, append([]string(nil), texts...)) + e.mu.Unlock() + result := make([][]float32, len(texts)) + for index, text := range texts { + result[index] = []float32{float32(len(text)), float32(len([]rune(text)))} + } + return result, nil +} + +func (e *pipelineCountingEmbedder) BatchEmbedWithPool( + ctx context.Context, + _ embedding.Embedder, + texts []string, +) ([][]float32, error) { + return e.BatchEmbed(ctx, texts) +} + +func (e *pipelineCountingEmbedder) GetModelName() string { return "counting-embedding" } +func (e *pipelineCountingEmbedder) GetDimensions() int { return 2 } +func (e *pipelineCountingEmbedder) GetModelID() string { return "counting-embedding-v1" } + +func (e *pipelineCountingEmbedder) callCount() int { + e.mu.Lock() + defer e.mu.Unlock() + return len(e.batches) +} + +type pipelineCountingChat struct { + mu sync.Mutex + calls map[string]int +} + +func newPipelineCountingChat() *pipelineCountingChat { + return &pipelineCountingChat{calls: make(map[string]int)} +} + +func (c *pipelineCountingChat) Chat( + _ context.Context, + messages []chat.Message, + _ *chat.ChatOptions, +) (*types.ChatResponse, error) { + if len(messages) == 0 { + return nil, fmt.Errorf("counting chat requires a message") + } + stage, _, _ := strings.Cut(messages[0].Content, "|") + c.mu.Lock() + c.calls[stage]++ + c.mu.Unlock() + return &types.ChatResponse{Content: "canonical-" + stage}, nil +} + +func (c *pipelineCountingChat) ChatStream( + context.Context, + []chat.Message, + *chat.ChatOptions, +) (<-chan types.StreamResponse, error) { + result := make(chan types.StreamResponse) + close(result) + return result, nil +} + +func (c *pipelineCountingChat) GetModelName() string { return "counting-chat" } +func (c *pipelineCountingChat) GetModelID() string { return "counting-chat-v1" } + +func (c *pipelineCountingChat) callCount(stage string) int { + c.mu.Lock() + defer c.mu.Unlock() + return c.calls[stage] +} + +type pipelineCountingVLM struct { + mu sync.Mutex + prompts []string +} + +func (v *pipelineCountingVLM) Predict( + _ context.Context, + _ [][]byte, + prompt string, +) (string, error) { + v.mu.Lock() + v.prompts = append(v.prompts, prompt) + v.mu.Unlock() + // Deliberately canonical across prompt revisions. This proves that a + // changed VLM request invalidates that stage without invalidating consumers + // when the verified output is unchanged. + return "canonical-image-output", nil +} + +func (v *pipelineCountingVLM) GetModelName() string { return "counting-vlm" } +func (v *pipelineCountingVLM) GetModelID() string { return "counting-vlm-v1" } + +func (v *pipelineCountingVLM) callCount() int { + v.mu.Lock() + defer v.mu.Unlock() + return len(v.prompts) +} + +type artifactPipelineHarness struct { + service *knowledgeService + reader *countingDocReader + embeddingProvider *pipelineCountingEmbedder + chatProvider *pipelineCountingChat + vlmProvider *pipelineCountingVLM + embedder embedding.Embedder + chat chat.Chat + vlm vlm.VLM +} + +type artifactPipelineResult struct { + chunkIDs []string + vectors [][]float32 + parseOutput string + ocrOutput string + caption string + summary string + questions string + wikiMap string + graphEntity string + graphEdge string + wikiReduce string + desired desiredChunkSet +} + +func newArtifactPipelineHarness( + t *testing.T, + runtime *artifact.Runtime, +) *artifactPipelineHarness { + t.Helper() + service := setupDocReaderArtifactService(t) + if runtime != nil { + service.artifactRuntime = runtime + } + reader := &countingDocReader{result: &types.ReadResult{ + MarkdownContent: "# Report\nstable body", + Metadata: map[string]string{"pages": "1"}, + }} + embeddingProvider := &pipelineCountingEmbedder{} + chatProvider := newPipelineCountingChat() + vlmProvider := &pipelineCountingVLM{} + return &artifactPipelineHarness{ + service: service, + reader: reader, + embeddingProvider: embeddingProvider, + chatProvider: chatProvider, + vlmProvider: vlmProvider, + embedder: embedding.NewArtifactCachedEmbedder( + embeddingProvider, + service.artifactRuntime, + embedding.ArtifactCacheConfig{ + TenantID: 1, + Processor: artifact.ProcessorIdentity{ + ModelID: "counting-embedding-v1", + ModelName: "counting-embedding", + Provider: "counting", + }, + Dimensions: 2, + }, + ), + chat: chat.NewArtifactCachedChat( + chatProvider, + service.artifactRuntime, + chat.ArtifactCacheConfig{ + TenantID: 1, + Processor: artifact.ProcessorIdentity{ + ModelID: "counting-chat-v1", + ModelName: "counting-chat", + Provider: "counting", + }, + }, + ), + vlm: vlm.NewArtifactCachedVLM( + vlmProvider, + service.artifactRuntime, + vlm.ArtifactCacheConfig{ + TenantID: 1, + Processor: artifact.ProcessorIdentity{ + ModelID: "counting-vlm-v1", + ModelName: "counting-vlm", + Provider: "counting", + }, + }, + ), + } +} + +func (h *artifactPipelineHarness) run( + t *testing.T, + parserOverride string, + vlmPromptRevision string, + existing []*types.Chunk, +) artifactPipelineResult { + t.Helper() + ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(1)) + parsed, err := h.service.callDocReaderWithArtifact(ctx, h.reader, &types.ReadRequest{ + FileContent: []byte("stable file bytes"), + FileName: "report.pdf", + FileType: "pdf", + Title: "Report", + ParserEngine: "counting-parser", + ParserEngineOverrides: map[string]string{ + "pdf_force_scanned": parserOverride, + }, + }) + require.NoError(t, err) + + knowledge := &types.Knowledge{ + ID: "knowledge-e2e", + TenantID: 1, + KnowledgeBaseID: "kb-e2e", + } + sourceChunks := []types.ParsedChunk{ + {Content: parsed.MarkdownContent, Seq: 0, Start: 0, End: len(parsed.MarkdownContent), ParentIndex: -1}, + {Content: "second stable block", Seq: 1, Start: len(parsed.MarkdownContent), End: len(parsed.MarkdownContent) + 19, ParentIndex: -1}, + } + desired, err := buildDesiredDocumentChunks(knowledge, sourceChunks, nil, existing) + require.NoError(t, err) + + imageBytes := [][]byte{[]byte("stable-image-bytes")} + ocr, err := h.vlm.Predict( + vlm.WithArtifactStage(ctx, vlm.ArtifactStage{ + Stage: "vlm_ocr", + OutputSchema: "vlm.ocr.text.v1", + }), + imageBytes, + "ocr|"+vlmPromptRevision, + ) + require.NoError(t, err) + caption, err := h.vlm.Predict( + vlm.WithArtifactStage(ctx, vlm.ArtifactStage{ + Stage: "vlm_caption", + OutputSchema: "vlm.caption.text.v1", + }), + imageBytes, + "caption|"+vlmPromptRevision, + ) + require.NoError(t, err) + + embeddingInputs := make([]string, len(desired.Text)) + for index, chunk := range desired.Text { + embeddingInputs[index] = chunk.EmbeddingContent() + } + embeddingContext := context.WithValue(ctx, types.EmbedDocumentContextKey, true) + vectors, err := h.embedder.BatchEmbed(embeddingContext, embeddingInputs) + require.NoError(t, err) + + callStage := func(stage, schema, input string) string { + response, callErr := h.chat.Chat( + chat.WithArtifactStage(ctx, chat.ArtifactStage{ + Stage: stage, + OutputSchema: schema, + }), + []chat.Message{{Role: "user", Content: stage + "|" + input}}, + &chat.ChatOptions{Temperature: 0, Seed: 7, MaxTokens: 128}, + ) + require.NoError(t, callErr) + return response.Content + } + summary := callStage("summary", "summary.text.v1", parsed.MarkdownContent+"|"+caption) + questions := callStage("question", "questions.text.v1", parsed.MarkdownContent) + wikiMap := callStage("wiki_map", "wiki.map.text.v1", parsed.MarkdownContent+"|"+caption) + graphEntity := callStage("graph_extract.entities", "graph.entities.text.v1", parsed.MarkdownContent) + graphEdge := callStage( + "graph_extract.relationships", + "graph.relationships.text.v1", + parsed.MarkdownContent+"|"+graphEntity, + ) + + // Wiki Reduce intentionally has no ArtifactStage marker. It must reflect + // the current live contributor set on every run. + reduce, err := h.chat.Chat( + ctx, + []chat.Message{{Role: "user", Content: "wiki_reduce|" + wikiMap}}, + &chat.ChatOptions{Temperature: 0, Seed: 7, MaxTokens: 128}, + ) + require.NoError(t, err) + + chunkIDs := make([]string, len(desired.All)) + for index, chunk := range desired.All { + chunkIDs[index] = chunk.ID + } + return artifactPipelineResult{ + chunkIDs: chunkIDs, + vectors: vectors, + parseOutput: parsed.MarkdownContent, + ocrOutput: ocr, + caption: caption, + summary: summary, + questions: questions, + wikiMap: wikiMap, + graphEntity: graphEntity, + graphEdge: graphEdge, + wikiReduce: reduce.Content, + desired: desired, + } +} + +func TestArtifactPipelineColdWarmCountingProviders(t *testing.T) { + harness := newArtifactPipelineHarness(t, nil) + cold := harness.run(t, "false", "prompt-v1", nil) + warm := harness.run(t, "false", "prompt-v1", cold.desired.All) + + assert.Equal(t, 1, harness.reader.callCount(), "Parse increment on warm run") + assert.Equal(t, 2, harness.vlmProvider.callCount(), "OCR/Caption increment on warm run") + assert.Equal(t, 1, harness.embeddingProvider.callCount(), "Embedding increment on warm run") + for _, stage := range []string{ + "summary", + "question", + "wiki_map", + "graph_extract.entities", + "graph_extract.relationships", + } { + assert.Equal(t, 1, harness.chatProvider.callCount(stage), stage+" increment on warm run") + } + assert.Equal(t, 2, harness.chatProvider.callCount("wiki_reduce"), "Wiki Reduce must remain live") + + assert.Equal(t, cold.chunkIDs, warm.chunkIDs) + assert.Equal(t, cold.vectors, warm.vectors) + assert.Equal(t, cold.parseOutput, warm.parseOutput) + assert.Equal(t, cold.ocrOutput, warm.ocrOutput) + assert.Equal(t, cold.caption, warm.caption) + assert.Equal(t, cold.summary, warm.summary) + assert.Equal(t, cold.questions, warm.questions) + assert.Equal(t, cold.wikiMap, warm.wikiMap) + assert.Equal(t, cold.graphEntity, warm.graphEntity) + assert.Equal(t, cold.graphEdge, warm.graphEdge) + assert.Equal(t, cold.wikiReduce, warm.wikiReduce) + assert.Empty(t, warm.desired.Added) + assert.Empty(t, warm.desired.Stale) + assert.Len(t, warm.desired.Updated, len(cold.desired.All)) + for _, vector := range warm.vectors { + assert.Len(t, vector, 2) + } +} + +func TestArtifactPipelineLayeredInvalidationStopsAtCanonicalOutput(t *testing.T) { + harness := newArtifactPipelineHarness(t, nil) + cold := harness.run(t, "false", "prompt-v1", nil) + + // Parser options and both VLM prompts change, but their deterministic + // canonical outputs remain identical. Direct stages recompute; consumers + // continue to hit. + changed := harness.run(t, "true", "prompt-v2", cold.desired.All) + assert.Equal(t, 2, harness.reader.callCount()) + assert.Equal(t, 4, harness.vlmProvider.callCount()) + assert.Equal(t, 1, harness.embeddingProvider.callCount()) + for _, stage := range []string{ + "summary", + "question", + "wiki_map", + "graph_extract.entities", + "graph_extract.relationships", + } { + assert.Equal(t, 1, harness.chatProvider.callCount(stage)) + } + assert.Equal(t, cold.chunkIDs, changed.chunkIDs) + assert.Equal(t, cold.vectors, changed.vectors) +} + +func TestArtifactPipelineStoreUnavailableCompletesFailOpen(t *testing.T) { + harness := newArtifactPipelineHarness(t, artifact.NewRuntime(nil, nil)) + result := harness.run(t, "false", "prompt-v1", nil) + + assert.Len(t, result.chunkIDs, 2) + assert.Len(t, result.vectors, 2) + assert.Equal(t, 1, harness.reader.callCount()) + assert.Equal(t, 2, harness.vlmProvider.callCount()) + assert.Equal(t, 1, harness.embeddingProvider.callCount()) + assert.Equal(t, 1, harness.chatProvider.callCount("summary")) + assert.Equal(t, 1, harness.chatProvider.callCount("question")) + assert.Equal(t, 1, harness.chatProvider.callCount("wiki_map")) + assert.Equal(t, 1, harness.chatProvider.callCount("graph_extract.entities")) + assert.Equal(t, 1, harness.chatProvider.callCount("graph_extract.relationships")) +} diff --git a/internal/application/service/extract.go b/internal/application/service/extract.go index 156fe84b2e..b6637d3276 100644 --- a/internal/application/service/extract.go +++ b/internal/application/service/extract.go @@ -288,7 +288,13 @@ func (s *ChunkExtractService) Handle(ctx context.Context, t *asynq.Task) error { return } if handleErr != nil { - s.tracker().FailSpan(ctx, gSpan, "GRAPH_EXTRACT_FAILED", handleErr.Error(), handleErr) + s.tracker().FailSpan( + ctx, + gSpan, + "GRAPH_EXTRACT_FAILED", + fmt.Sprintf("%T", handleErr), + nil, + ) } else { s.tracker().EndSpan(ctx, gSpan, graphOut) } @@ -316,12 +322,9 @@ func (s *ChunkExtractService) Handle(ctx context.Context, t *asynq.Task) error { handleErr = err return err } - // Capture chunk content shape on output — lets traces answer "WHAT - // did the LLM call see?" without joining back to the chunk store. - // Preview is truncated to keep span rows reasonable. + // Capture only chunk shape; body content must not enter traces. if gSpan != nil { graphOut["chunk_chars"] = len([]rune(chunk.Content)) - graphOut["chunk_preview"] = previewText(chunk.Content, 200) } kb, err := s.knowledgeBaseRepo.GetKnowledgeBaseByID(ctx, chunk.KnowledgeBaseID) if err != nil { @@ -369,6 +372,7 @@ func (s *ChunkExtractService) Handle(ctx context.Context, t *asynq.Task) error { extractor := chatpipeline.NewExtractor(chatModel, template) graph, err := extractor.Extract(ctx, chunk.Content) if err != nil { + logger.Warnf(ctx, "graph extraction provider failed: error_class=%T", err) handleErr = err return err } @@ -421,6 +425,10 @@ func (s *ChunkExtractService) Handle(ctx context.Context, t *asynq.Task) error { handleErr = err return err } + if err := artifact.InjectFault(ctx, artifact.FaultAfterGraphBinding); err != nil { + handleErr = err + return err + } if attemptSuperseded(ctx, s.tracker(), p.KnowledgeID, p.Attempt) { superseded = true graphOut["status"] = "superseded_after_graph_publish" @@ -428,32 +436,6 @@ func (s *ChunkExtractService) Handle(ctx context.Context, t *asynq.Task) error { } graphOut["nodes_added"] = len(graph.Node) graphOut["relations_added"] = len(graph.Relation) - // Capture a couple of sample nodes/relations so the trace viewer can - // answer "what did the LLM actually extract?" without round-tripping - // to the graph store. Cap to two each — anything more bloats span - // rows and the full graph is queryable elsewhere. - if len(graph.Node) > 0 { - samples := graph.Node - if len(samples) > 2 { - samples = samples[:2] - } - names := make([]string, 0, len(samples)) - for _, n := range samples { - names = append(names, n.Name) - } - graphOut["sample_nodes"] = names - } - if len(graph.Relation) > 0 { - samples := graph.Relation - if len(samples) > 2 { - samples = samples[:2] - } - out := make([]string, 0, len(samples)) - for _, r := range samples { - out = append(out, fmt.Sprintf("%s --[%s]--> %s", r.Node1, r.Type, r.Node2)) - } - graphOut["sample_relations"] = out - } return nil } @@ -749,7 +731,8 @@ func (s *DataTableSummaryService) processTableData(ctx context.Context, resource logger.Errorf(ctx, "failed to generate table description: %v", err) return nil, err } - logger.Debugf(ctx, "table describe of knowledge %s: %s", resources.knowledge.ID, tableDescription) + logger.Debugf(ctx, "generated table description for knowledge %s: chars=%d", + resources.knowledge.ID, len([]rune(tableDescription))) columnDescription, err := s.generateColumnDescriptions(ctx, resources.chatModel, tableSchema.TableName, schemaDesc, sampleDesc, customInstructions) @@ -757,7 +740,8 @@ func (s *DataTableSummaryService) processTableData(ctx context.Context, resource logger.Errorf(ctx, "failed to generate column descriptions: %v", err) return nil, err } - logger.Debugf(ctx, "column describe of knowledge %s: %s", resources.knowledge.ID, columnDescription) + logger.Debugf(ctx, "generated column descriptions for knowledge %s: chars=%d", + resources.knowledge.ID, len([]rune(columnDescription))) // 构建chunks：一个表格摘要chunk + 多个列描述chunks chunks, err := s.buildChunks(resources, tableDescription, columnDescription) diff --git a/internal/application/service/image_multimodal.go b/internal/application/service/image_multimodal.go index d2ad264773..bd811d7f63 100644 --- a/internal/application/service/image_multimodal.go +++ b/internal/application/service/image_multimodal.go @@ -141,8 +141,8 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e return fmt.Errorf("unmarshal image multimodal payload: %w", err) } - logger.Infof(ctx, "[ImageMultimodal] Processing image: chunk=%s, url=%s, ocr=%v, caption=%v", - payload.ChunkID, payload.ImageURL, payload.EnableOCR, payload.EnableCaption) + logger.Infof(ctx, "[ImageMultimodal] Processing image: chunk=%s, has_url=%v, ocr=%v, caption=%v", + payload.ChunkID, payload.ImageURL != "", payload.EnableOCR, payload.EnableCaption) ctx = context.WithValue(ctx, types.TenantIDContextKey, payload.TenantID) if payload.Language != "" { @@ -167,8 +167,8 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e } if drop { logger.Infof(ctx, - "[ImageMultimodal] Dropping task chunk=%s knowledge=%s kb=%s image=%s", - payload.ChunkID, payload.KnowledgeID, payload.KnowledgeBaseID, payload.ImageURL) + "[ImageMultimodal] Dropping task chunk=%s knowledge=%s kb=%s", + payload.ChunkID, payload.KnowledgeID, payload.KnowledgeBaseID) // Still count this image toward the parent finalize gate so a batch // of dropped orphans cannot strand multimodal:pending forever. s.checkAndFinalizeAllImages(ctx, payload) @@ -187,7 +187,7 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e if parent != nil { name := fmt.Sprintf("multimodal.image[%d]", payload.ImageIndex) imgSpan = tracker.BeginSubSpan(ctx, parent, name, types.SpanKindGeneration, types.JSONMap{ - "image_url": payload.ImageURL, + "has_image_url": payload.ImageURL != "", "image_source_type": payload.ImageSourceType, "enable_ocr": payload.EnableOCR, "enable_caption": payload.EnableCaption, @@ -197,10 +197,9 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e } // Output map populated as we go — the deferred close picks it up. - // Captures real VLM results (model id, byte count, OCR/caption - // previews, downstream chunk counts) so the trace viewer can answer - // "what did this image actually produce?" without joining back to - // the chunks table. + // Captures safe VLM result metadata (model id, byte count, output + // lengths, downstream chunk counts) without placing image references + // or generated content in the trace. imgOut := types.JSONMap{} // finalize-once semantics: on success we always decrement the parent's @@ -221,8 +220,8 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e } else if isFinalAsynqAttempt(ctx) { tracker.FailSpan(ctx, imgSpan, "MULTIMODAL_VLM_FAILED", - handleErr.Error(), - handleErr) + fmt.Sprintf("%T", handleErr), + nil) } } if attemptSuperseded(ctx, tracker, payload.KnowledgeID, payload.Attempt) { @@ -236,8 +235,7 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e s.checkAndFinalizeAllImages(ctx, payload) } else { logger.Infof(ctx, - "[ImageMultimodal] Skip finalize on retryable error for %s (will count on last attempt)", - payload.ImageURL) + "[ImageMultimodal] Skip finalize on retryable error (will count on last attempt)") } }() @@ -262,9 +260,9 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e // image, skip it (deferred finalize will count it). imgBytes, readErr := s.readImageBytes(ctx, payload) if readErr != nil { - logger.Errorf(ctx, "[ImageMultimodal] Skip unreadable image %s: %v", payload.ImageURL, readErr) + logger.Errorf(ctx, "[ImageMultimodal] Skip unreadable image: error_class=%T", readErr) imgOut["skipped"] = "unreadable_image" - imgOut["read_error"] = readErr.Error() + imgOut["read_error_class"] = fmt.Sprintf("%T", readErr) return nil } imgOut["image_bytes"] = len(imgBytes) @@ -280,7 +278,7 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e prompt := vlmOCRPrompt if payload.ImageSourceType == "scanned_pdf" { prompt = vlmOCRScannedPDFPrompt - logger.Infof(ctx, "[ImageMultimodal] Using scanned PDF prompt for OCR: %s", payload.ImageURL) + logger.Infof(ctx, "[ImageMultimodal] Using scanned PDF prompt for OCR") imgOut["ocr_prompt"] = "scanned_pdf" } else { imgOut["ocr_prompt"] = "default" @@ -293,17 +291,16 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e }) ocrText, ocrErr := vlmModel.Predict(ocrCtx, [][]byte{imgBytes}, prompt) if ocrErr != nil { - logger.Warnf(ctx, "[ImageMultimodal] OCR failed for %s: %v", payload.ImageURL, ocrErr) - imgOut["ocr_error"] = ocrErr.Error() + logger.Warnf(ctx, "[ImageMultimodal] OCR failed: error_class=%T", ocrErr) + imgOut["ocr_error_class"] = fmt.Sprintf("%T", ocrErr) } else { ocrDesiredKnown = true ocrText = sanitizeOCRText(ocrText) if ocrText != "" { imageInfo.OCRText = ocrText imgOut["ocr_chars"] = len([]rune(ocrText)) - imgOut["ocr_preview"] = previewText(ocrText, 200) } else { - logger.Warnf(ctx, "[ImageMultimodal] OCR returned empty/invalid content for %s, discarded", payload.ImageURL) + logger.Warnf(ctx, "[ImageMultimodal] OCR returned empty/invalid content, discarded") imgOut["ocr_chars"] = 0 imgOut["ocr_skipped"] = "empty_or_invalid" } @@ -321,14 +318,13 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e buildVLMCaptionPrompt(ctx, vlmCfg), ) if capErr != nil { - logger.Warnf(ctx, "[ImageMultimodal] Caption failed for %s: %v", payload.ImageURL, capErr) - imgOut["caption_error"] = capErr.Error() + logger.Warnf(ctx, "[ImageMultimodal] Caption failed: error_class=%T", capErr) + imgOut["caption_error_class"] = fmt.Sprintf("%T", capErr) } else { captionDesiredKnown = true if caption != "" { imageInfo.Caption = caption imgOut["caption_chars"] = len([]rune(caption)) - imgOut["caption_preview"] = previewText(caption, 200) } } } @@ -344,6 +340,25 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e ) return nil } + ctx, releaseReconcile, lockErr := acquireKnowledgeReconcileLock( + ctx, + s.redisClient, + payload.KnowledgeID, + ) + if lockErr != nil { + handleErr = fmt.Errorf("acquire multimodal reconciliation lock: %w", lockErr) + return handleErr + } + defer releaseReconcile() + if attemptSuperseded(ctx, tracker, payload.KnowledgeID, payload.Attempt) { + logger.Infof( + ctx, + "[ImageMultimodal] Attempt %d superseded while waiting for reconciliation lock on %s", + payload.Attempt, + payload.KnowledgeID, + ) + return nil + } // Build child chunks for OCR and caption results imageInfoJSON, _ := json.Marshal([]types.ImageInfo{imageInfo}) @@ -449,8 +464,8 @@ func (s *ImageMultimodalService) Handle(ctx context.Context, task *asynq.Task) e } } for _, c := range newChunks { - logger.Infof(ctx, "[ImageMultimodal] Created %s chunk %s for image %s, len=%d", - c.ChunkType, c.ID, payload.ImageURL, len(c.Content)) + logger.Infof(ctx, "[ImageMultimodal] Created %s chunk %s, len=%d", + c.ChunkType, c.ID, len(c.Content)) } // Index desired chunks before exact stale cleanup. @@ -724,7 +739,7 @@ func (s *ImageMultimodalService) indexChunks( } } - logger.Infof(ctx, "[ImageMultimodal] Indexed %d multimodal chunks for image %s", len(chunks), payload.ImageURL) + logger.Infof(ctx, "[ImageMultimodal] Indexed %d multimodal chunks", len(chunks)) return result, nil } @@ -783,7 +798,7 @@ func (s *ImageMultimodalService) resolveFileServiceForPayload(ctx context.Contex // stored on a different backend (multi-backend / post-migration). if _, isResourceRef := types.ParseResourcePath(payload.ImageURL); isResourceRef && s.resourceCatalog != nil { if resource, resErr := s.resourceCatalog.Resolve(ctx, payload.ImageURL); resErr != nil { - logger.Warnf(ctx, "[ImageMultimodal] resolve resource reference failed: url=%s err=%v", payload.ImageURL, resErr) + logger.Warnf(ctx, "[ImageMultimodal] resolve resource reference failed: error_class=%T", resErr) } else if resource != nil { backendID = resource.StorageBackendID provider = strings.ToLower(strings.TrimSpace(resource.Provider)) @@ -806,8 +821,8 @@ func (s *ImageMultimodalService) resolveFileServiceForPayload(ctx context.Contex } baseDir := strings.TrimSpace(os.Getenv("LOCAL_STORAGE_BASE_DIR")) - logger.Infof(ctx, "[ImageMultimodal] resolving file service: tenant=%d provider=%q LOCAL_STORAGE_BASE_DIR=%q imageURL=%s", - payload.TenantID, provider, baseDir, payload.ImageURL) + logger.Infof(ctx, "[ImageMultimodal] resolving file service: tenant=%d provider=%q has_backend=%v", + payload.TenantID, provider, backendID != "") fileSvc, _, svcErr := s.storageResolver.ResolveFileService(ctx, tenant, backendID, provider, baseDir) if svcErr != nil { logger.Warnf(ctx, "[ImageMultimodal] resolve file service failed (falling back to default): tenant=%d provider=%s err=%v", @@ -829,16 +844,16 @@ func (s *ImageMultimodalService) readImageBytes(ctx context.Context, payload typ if isResourceRef || types.ParseProviderScheme(payload.ImageURL) != "" { fileSvc := s.resolveFileServiceForPayload(ctx, payload) if fileSvc == nil { - return nil, fmt.Errorf("no file service available for %s", payload.ImageURL) + return nil, fmt.Errorf("no file service available for image reference") } reader, err := fileSvc.GetFile(ctx, payload.ImageURL) if err != nil { - return nil, fmt.Errorf("file service get %s: %w", payload.ImageURL, err) + return nil, fmt.Errorf("file service get image: %T", err) } defer reader.Close() data, err := io.ReadAll(reader) if err != nil { - return nil, fmt.Errorf("read %s: %w", payload.ImageURL, err) + return nil, fmt.Errorf("read image: %T", err) } return data, nil } @@ -847,13 +862,13 @@ func (s *ImageMultimodalService) readImageBytes(ctx context.Context, payload typ if data, err := os.ReadFile(payload.ImageLocalPath); err == nil { return data, nil } else { - logger.Warnf(ctx, "[ImageMultimodal] Local file %s not available (%v), falling back to URL", payload.ImageLocalPath, err) + logger.Warnf(ctx, "[ImageMultimodal] Local image unavailable: error_class=%T; falling back to URL", err) } } data, err := downloadImageFromURL(payload.ImageURL) if err != nil { - return nil, fmt.Errorf("download %s: %w", payload.ImageURL, err) + return nil, fmt.Errorf("download image: %T", err) } logger.Infof(ctx, "[ImageMultimodal] Image downloaded from URL, len=%d", len(data)) return data, nil diff --git a/internal/application/service/image_multimodal_prompt_test.go b/internal/application/service/image_multimodal_prompt_test.go index 5007f666a2..b99f8a8628 100644 --- a/internal/application/service/image_multimodal_prompt_test.go +++ b/internal/application/service/image_multimodal_prompt_test.go @@ -58,6 +58,19 @@ func TestBuildStableMultimodalChunkUsesExactInputsAndUniqueLegacyMatch(t *testin require.NoError(t, err) assert.Equal(t, first.ID, second.ID) + signedURLChanged := payload + signedURLChanged.ImageURL = "https://objects.example/image.png?signature=rotated" + sameBytes, err := buildStableMultimodalChunk( + signedURLChanged, + types.ChunkTypeImageOCR, + "exact OCR", + "[]", + []byte{0, 1, 2}, + nil, + ) + require.NoError(t, err) + assert.Equal(t, first.ID, sameBytes.ID, "signed URL rotation must not change byte-addressed identity") + changed, err := buildStableMultimodalChunk( payload, types.ChunkTypeImageOCR, @@ -69,6 +82,17 @@ func TestBuildStableMultimodalChunkUsesExactInputsAndUniqueLegacyMatch(t *testin require.NoError(t, err) assert.NotEqual(t, first.ID, changed.ID) + changedOutput, err := buildStableMultimodalChunk( + payload, + types.ChunkTypeImageOCR, + "changed OCR", + "[]", + []byte{0, 1, 2}, + nil, + ) + require.NoError(t, err) + assert.NotEqual(t, first.ID, changedOutput.ID) + legacyID := uuid.New().String() reused, err := buildStableMultimodalChunk( payload, diff --git a/internal/application/service/knowledge.go b/internal/application/service/knowledge.go index aa3ed4c346..b4de7cd653 100644 --- a/internal/application/service/knowledge.go +++ b/internal/application/service/knowledge.go @@ -282,7 +282,10 @@ func (s *knowledgeService) failStage(ctx context.Context, kid, name, code, msg s if span == nil { return } - s.tracker().FailSpan(ctx, span, code, msg, err) + if err != nil { + msg = fmt.Sprintf("%s (error_class=%T)", msg, err) + } + s.tracker().FailSpan(ctx, span, code, msg, nil) } func (s *knowledgeService) skipStage(ctx context.Context, kid, name, reason string) { diff --git a/internal/application/service/knowledge_clone_move.go b/internal/application/service/knowledge_clone_move.go index 6f2b0e2ab0..b0a6b2a90e 100644 --- a/internal/application/service/knowledge_clone_move.go +++ b/internal/application/service/knowledge_clone_move.go @@ -1352,7 +1352,7 @@ func (s *knowledgeService) moveKnowledgeReparse( if err != nil || meta == nil { return fmt.Errorf("failed to get manual metadata for reparse: %w", err) } - s.triggerManualProcessing(ctx, targetKB, knowledge, meta.Content, false) + _ = s.triggerManualProcessing(ctx, targetKB, knowledge, meta.Content, false) return nil } diff --git a/internal/application/service/knowledge_create.go b/internal/application/service/knowledge_create.go index f86dbba8f0..ffe460be6e 100644 --- a/internal/application/service/knowledge_create.go +++ b/internal/application/service/knowledge_create.go @@ -1212,10 +1212,10 @@ func sanitizeManualDownloadFilename(title string) string { func (s *knowledgeService) triggerManualProcessing(ctx context.Context, kb *types.KnowledgeBase, knowledge *types.Knowledge, content string, doSync bool, -) { +) error { clean := strings.TrimSpace(content) if clean == "" { - return + return nil } // Resolve embedded data:base64 images and remote http(s) images → storage, replace URLs. @@ -1293,10 +1293,14 @@ func (s *knowledgeService) triggerManualProcessing(ctx context.Context, } if doSync { - s.processChunks(ctx, kb, knowledge, parsed, opts) - return + return s.processChunks(ctx, kb, knowledge, parsed, opts) } newCtx := logger.CloneContext(ctx) - go s.processChunks(newCtx, kb, knowledge, parsed, opts) + go func() { + if err := s.processChunks(newCtx, kb, knowledge, parsed, opts); err != nil { + logger.Warnf(newCtx, "manual chunk processing failed: error_class=%T", err) + } + }() + return nil } diff --git a/internal/application/service/knowledge_process.go b/internal/application/service/knowledge_process.go index 955132c19e..466a9360a2 100644 --- a/internal/application/service/knowledge_process.go +++ b/internal/application/service/knowledge_process.go @@ -148,7 +148,9 @@ func (s *knowledgeService) processDocumentFromPassage(ctx context.Context, opts.QuestionCount = 3 } } - s.processChunks(ctx, kb, knowledge, chunks, opts) + if err := s.processChunks(ctx, kb, knowledge, chunks, opts); err != nil { + logger.Warnf(ctx, "passage chunk processing failed: error_class=%T", err) + } } // ProcessChunksOptions contains options for processing chunks @@ -259,7 +261,7 @@ func buildParentChildConfigs(cc types.ChunkingConfig, base chunker.SplitterConfi func (s *knowledgeService) processChunks(ctx context.Context, kb *types.KnowledgeBase, knowledge *types.Knowledge, chunks []types.ParsedChunk, opts ...ProcessChunksOptions, -) { +) error { // Get options var options ProcessChunksOptions if len(opts) > 0 { @@ -271,21 +273,29 @@ func (s *knowledgeService) processChunks(ctx context.Context, // up yet so the branch is purely "stop early". if aborted, status := s.isKnowledgeAborted(ctx, knowledge.TenantID, knowledge.ID); aborted { logger.Infof(ctx, "Knowledge aborted (%s), skipping chunk processing: %s", status, knowledge.ID) - return + return nil } if attemptSuperseded(ctx, s.tracker(), knowledge.ID, attemptFromCtx(ctx)) { logger.Infof(ctx, "Knowledge attempt superseded, skipping chunk processing: %s", knowledge.ID) - return + return nil + } + ctx, releaseReconcile, err := acquireKnowledgeReconcileLock(ctx, s.redisClient, knowledge.ID) + if err != nil { + return fmt.Errorf("acquire knowledge reconciliation lock: %w", err) + } + defer releaseReconcile() + if attemptSuperseded(ctx, s.tracker(), knowledge.ID, attemptFromCtx(ctx)) { + logger.Infof(ctx, "Knowledge attempt superseded while waiting for reconciliation lock: %s", knowledge.ID) + return nil } // Get embedding model for vectorization — only needed when vector/keyword indexing is enabled var embeddingModel embedding.Embedder if kb.NeedsEmbeddingModel() { - var err error embeddingModel, err = s.modelService.GetEmbeddingModel(ctx, kb.EmbeddingModelID) if err != nil { logger.GetLogger(ctx).WithField("error", err).Errorf("processChunks get embedding model failed") - return + return err } } else { logger.Infof(ctx, "Vector/keyword indexing disabled for KB %s, skipping embedding model", kb.ID) @@ -295,7 +305,7 @@ func (s *knowledgeService) processChunks(ctx context.Context, existingChunks, err := s.chunkRepo.ListAllChunksByKnowledgeID(ctx, knowledge.TenantID, knowledge.ID) if err != nil { logger.Errorf(ctx, "Failed to load current chunks for reconciliation: %v", err) - return + return err } // Resolve the index once, but keep the currently published rows and vectors @@ -304,7 +314,7 @@ func (s *knowledgeService) processChunks(ctx context.Context, ctx, s.retrieveEngine, s.ownership, tenantInfo.ID, kb.VectorStoreID) if err != nil && embeddingModel != nil { logger.Errorf(ctx, "Failed to resolve retrieve engine for reconciliation: %v", err) - return + return err } // ========== DocReader 解析结果日志 ========== @@ -323,35 +333,25 @@ func (s *knowledgeService) processChunks(ctx context.Context, } logger.Infof(ctx, "[DocReader] 包含图片的Chunk数: %d, 总图片数: %d", chunksWithImages, totalImages) - // 打印每个Chunk的详细信息 + // Record only structural diagnostics. Document text, prompts, captions, + // OCR output, and image references may contain customer data or signed + // URLs and must never enter logs. for idx, chunkData := range chunks { - contentPreview := chunkData.Content - if len(contentPreview) > 200 { - contentPreview = contentPreview[:200] + "..." - } logger.Infof(ctx, "[DocReader] Chunk #%d (seq=%d): 内容长度=%d, 图片数=%d, 范围=[%d-%d]", idx, chunkData.Seq, len(chunkData.Content), len(chunkData.Images), chunkData.Start, chunkData.End) - logger.Debugf(ctx, "[DocReader] Chunk #%d 内容预览: %s", idx, contentPreview) - // 打印图片详细信息 for imgIdx, img := range chunkData.Images { - logger.Infof(ctx, "[DocReader] 图片 #%d: URL=%s", imgIdx, img.URL) - logger.Infof(ctx, "[DocReader] 图片 #%d: OriginalURL=%s", imgIdx, img.OriginalURL) - if img.Caption != "" { - captionPreview := img.Caption - if len(captionPreview) > 100 { - captionPreview = captionPreview[:100] + "..." - } - logger.Infof(ctx, "[DocReader] 图片 #%d: Caption=%s", imgIdx, captionPreview) - } - if img.OCRText != "" { - ocrPreview := img.OCRText - if len(ocrPreview) > 100 { - ocrPreview = ocrPreview[:100] + "..." - } - logger.Infof(ctx, "[DocReader] 图片 #%d: OCRText=%s", imgIdx, ocrPreview) - } - logger.Infof(ctx, "[DocReader] 图片 #%d: 位置=[%d-%d]", imgIdx, img.Start, img.End) + logger.Infof( + ctx, + "[DocReader] 图片 #%d: has_url=%v, has_original_url=%v, caption_chars=%d, ocr_chars=%d, 位置=[%d-%d]", + imgIdx, + img.URL != "", + img.OriginalURL != "", + len([]rune(img.Caption)), + len([]rune(img.OCRText)), + img.Start, + img.End, + ) } } logger.Infof(ctx, "[DocReader] ========== 解析结果概览结束 ==========") @@ -364,7 +364,7 @@ func (s *knowledgeService) processChunks(ctx context.Context, ) if err != nil { logger.Errorf(ctx, "Failed to build desired chunk state: %v", err) - return + return err } insertChunks := desired.All textChunks := desired.Text @@ -374,11 +374,11 @@ func (s *knowledgeService) processChunks(ctx context.Context, // Nothing has been persisted yet, so both branches just bail. if aborted, status := s.isKnowledgeAborted(ctx, knowledge.TenantID, knowledge.ID); aborted { logger.Infof(ctx, "Knowledge aborted (%s), skipping chunk write: %s", status, knowledge.ID) - return + return nil } if attemptSuperseded(ctx, s.tracker(), knowledge.ID, attemptFromCtx(ctx)) { logger.Infof(ctx, "Knowledge attempt superseded before desired chunk write: %s", knowledge.ID) - return + return nil } // Save chunks to database — ALWAYS, regardless of indexing strategy. @@ -401,7 +401,7 @@ func (s *knowledgeService) processChunks(ctx context.Context, ) s.failStage(ctx, knowledge.ID, types.StageChunking, werrors.ErrCodeChunkingFailed, "create chunks failed", err) - return + return err } if len(desired.Updated) > 0 { err = s.chunkRepo.SaveChunks(ctx, desired.Updated) @@ -418,7 +418,10 @@ func (s *knowledgeService) processChunks(ctx context.Context, ) s.failStage(ctx, knowledge.ID, types.StageChunking, werrors.ErrCodeChunkingFailed, "update desired chunks failed", err) - return + return err + } + if err := artifact.InjectFault(ctx, artifact.FaultAfterChunkUpsert); err != nil { + return err } totalChunkChars := 0 for _, c := range insertChunks { @@ -480,7 +483,7 @@ func (s *knowledgeService) processChunks(ctx context.Context, knowledge, attemptFromCtx(ctx), ) - return + return err } // Check if there's enough storage quota available projectedStorage := tenantInfo.StorageUsed - previousStorageSize + totalStorageSize @@ -493,7 +496,7 @@ func (s *knowledgeService) processChunks(ctx context.Context, knowledge, attemptFromCtx(ctx), ) - return + return nil } } @@ -501,11 +504,11 @@ func (s *knowledgeService) processChunks(ctx context.Context, // attempt may leave ownership-free artifacts, but must not publish. if aborted, status := s.isKnowledgeAborted(ctx, knowledge.TenantID, knowledge.ID); aborted { logger.Infof(ctx, "Knowledge aborted (%s) before indexing: %s", status, knowledge.ID) - return + return nil } if attemptSuperseded(ctx, s.tracker(), knowledge.ID, attemptFromCtx(ctx)) { logger.Infof(ctx, "Knowledge attempt superseded before indexing: %s", knowledge.ID) - return + return nil } err = retrieveEngine.BatchIndex(ctx, embeddingModel, indexInfoList) @@ -543,7 +546,10 @@ func (s *knowledgeService) processChunks(ctx context.Context, } s.failStage(ctx, knowledge.ID, types.StageEmbedding, code, "batch index failed", err) - return + return err + } + if err := artifact.InjectFault(ctx, artifact.FaultAfterVectorUpsert); err != nil { + return err } logger.GetLogger(ctx).Infof("processChunks batch index successfully, with %d index", len(indexInfoList)) s.endStage(ctx, knowledge.ID, types.StageEmbedding, types.JSONMap{ @@ -554,11 +560,11 @@ func (s *knowledgeService) processChunks(ctx context.Context, // Fence again after provider work and before publishing live state. if aborted, status := s.isKnowledgeAborted(ctx, knowledge.TenantID, knowledge.ID); aborted { logger.Infof(ctx, "Knowledge aborted (%s) after indexing: %s", status, knowledge.ID) - return + return nil } if attemptSuperseded(ctx, s.tracker(), knowledge.ID, attemptFromCtx(ctx)) { logger.Infof(ctx, "Knowledge attempt superseded after indexing: %s", knowledge.ID) - return + return nil } } else { logger.Infof(ctx, "Vector/keyword indexing disabled for KB %s, skipping BatchIndex", kb.ID) @@ -574,9 +580,12 @@ func (s *knowledgeService) processChunks(ctx context.Context, pendingMultimodal := isImage && options.EnableMultimodel && len(options.StoredImages) > 0 pendingPDFMultimodal := !isImage && !isVideo && options.EnableMultimodel && len(options.StoredImages) > 0 + if err := artifact.InjectFault(ctx, artifact.FaultBeforeFence); err != nil { + return err + } if attemptSuperseded(ctx, s.tracker(), knowledge.ID, attemptFromCtx(ctx)) { logger.Infof(ctx, "Knowledge attempt superseded before publish: %s", knowledge.ID) - return + return nil } now := time.Now() finalizeIndexedKnowledgeState( @@ -588,28 +597,45 @@ func (s *knowledgeService) processChunks(ctx context.Context, ) knowledge.EmbeddingModelID = kb.EmbeddingModelID - published, err := s.repo.UpdateKnowledgeIfAttemptCurrent( - ctx, - knowledge, - attemptFromCtx(ctx), - ) + storagePublished := false + var published bool + if publisher, ok := s.repo.(interface { + PublishKnowledgeIfAttemptCurrent(context.Context, *types.Knowledge, int) (bool, error) + }); ok { + published, err = publisher.PublishKnowledgeIfAttemptCurrent( + ctx, + knowledge, + attemptFromCtx(ctx), + ) + storagePublished = published && err == nil + } else { + published, err = s.repo.UpdateKnowledgeIfAttemptCurrent( + ctx, + knowledge, + attemptFromCtx(ctx), + ) + } if err != nil { logger.GetLogger(ctx).WithField("error", err).Errorf("processChunks update knowledge failed") - return + return err } if !published { logger.Infof(ctx, "Knowledge attempt lost conditional publish fence: %s", knowledge.ID) - return + return nil + } + if err := artifact.InjectFault(ctx, artifact.FaultAfterPublish); err != nil { + return err } // Cleanup is deliberately last and exact. If a newer attempt starts after // publish, it owns cleanup; this attempt must leave both generations intact. if attemptSuperseded(ctx, s.tracker(), knowledge.ID, attemptFromCtx(ctx)) { logger.Infof(ctx, "Knowledge attempt superseded before stale cleanup: %s", knowledge.ID) - return + return nil } staleIDs := chunkIDList(desired.Stale) staleVectorsDeleted := true + var staleCleanupErr error if len(staleIDs) > 0 && embeddingModel != nil { if err := retrieveEngine.DeleteByChunkIDList( ctx, @@ -618,9 +644,15 @@ func (s *knowledgeService) processChunks(ctx context.Context, kb.Type, ); err != nil { staleVectorsDeleted = false + staleCleanupErr = errors.Join(staleCleanupErr, err) logger.Warnf(ctx, "Failed to cleanup stale vectors; retaining stale chunks for retry: %v", err) } } + if len(staleIDs) > 0 { + if err := artifact.InjectFault(ctx, artifact.FaultDuringStaleCleanup); err != nil { + return err + } + } if len(staleIDs) > 0 { if graphContributions, ok := s.graphEngine.(interfaces.GraphContributionRepository); ok { err := graphContributions.DeleteGraphContributions( @@ -634,6 +666,7 @@ func (s *knowledgeService) processChunks(ctx context.Context, ) if err != nil { staleVectorsDeleted = false + staleCleanupErr = errors.Join(staleCleanupErr, err) logger.Warnf( ctx, "Failed to cleanup stale graph contributions; retaining stale chunks for retry: %v", @@ -644,9 +677,13 @@ func (s *knowledgeService) processChunks(ctx context.Context, } if len(staleIDs) > 0 && staleVectorsDeleted { if err := s.chunkRepo.DeleteChunks(ctx, knowledge.TenantID, staleIDs); err != nil { + staleCleanupErr = errors.Join(staleCleanupErr, err) logger.Warnf(ctx, "Failed to cleanup stale chunks: %v", err) } } + if staleCleanupErr != nil { + return fmt.Errorf("cleanup stale desired-state bindings: %w", staleCleanupErr) + } // Enqueue multimodal tasks for images (async, non-blocking) if options.EnableMultimodel && len(options.StoredImages) > 0 { @@ -674,22 +711,30 @@ func (s *knowledgeService) processChunks(ctx context.Context, knowledgePostProcessTaskOptions()...) if _, err := s.task.Enqueue(task); err != nil { logger.Errorf(ctx, "Failed to enqueue knowledge post process task: %v", err) + return err } else { logger.Infof(ctx, "Enqueued knowledge post process task for %s", knowledge.ID) } } else { logger.Errorf(ctx, "Failed to marshal knowledge post process payload: %v", err) + return err } } // Update tenant storage by the desired-state delta, not by re-adding the - // full size on every reparse. - storageDelta := totalStorageSize - previousStorageSize - tenantInfo.StorageUsed += storageDelta - if err := s.tenantRepo.AdjustStorageUsed(ctx, tenantInfo.ID, storageDelta); err != nil { - logger.GetLogger(ctx).WithField("error", err).Errorf("processChunks update tenant storage used failed") + // full size on every reparse. Production repositories publish knowledge and + // this delta atomically; the fallback preserves compatibility with fakes and + // third-party repository implementations. + if !storagePublished { + storageDelta := totalStorageSize - previousStorageSize + tenantInfo.StorageUsed += storageDelta + if err := s.tenantRepo.AdjustStorageUsed(ctx, tenantInfo.ID, storageDelta); err != nil { + logger.GetLogger(ctx).WithField("error", err).Errorf("processChunks update tenant storage used failed") + return err + } } logger.GetLogger(ctx).Infof("processChunks successfully") + return nil } // defaultMaxInputChars is the default maximum characters used as input for summary generation. @@ -1156,13 +1201,14 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn summaryMetadataVersion := string(knowledge.CustomMetadata) summary, err := s.getSummary(ctx, chatModel, knowledge, textChunks) if err != nil { - logger.Errorf(ctx, "Failed to generate summary for knowledge %s: %v", payload.KnowledgeID, err) + logger.Errorf(ctx, "Failed to generate summary for knowledge %s: error_class=%T", + payload.KnowledgeID, err) // Surface the underlying LLM/IO error on the span so the trace UI // can explain "why did this stage take 60s and then fall back?" // without forcing the operator to grep worker logs. We also capture // the error type to disambiguate timeouts from upstream HTTP errors // (deadline exceeded vs unexpected EOF vs 5xx, etc.). - summaryOut["error"] = previewText(err.Error(), 500) + summaryOut["error_class"] = fmt.Sprintf("%T", err) summaryOut["error_type"] = fmt.Sprintf("%T", err) // A failed provider call must not replace the last good description or // summary contribution with a fallback. Mark this attempt failed and @@ -1202,11 +1248,6 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn return nil } summaryOut["summary_chars"] = len([]rune(summary)) - // Preview the generated summary on the span output so the trace - // viewer can show "this is what the LLM produced" at a glance, - // without hopping to the knowledge-detail page. Capped to keep - // span rows compact. - summaryOut["summary_preview"] = previewText(summary, 240) var summaryRetrieveEngine *retriever.CompositeRetrieveEngine var summaryEmbeddingModel embedding.Embedder @@ -1492,11 +1533,7 @@ func (s *knowledgeService) processQuestionGenerationForKnowledge(ctx context.Con indexEntriesPrepared := 0 indexBatchAttempted := false indexBatchSucceeded := false - // Sample question + model id surfaced on the span output so the - // trace viewer can answer "what did the LLM actually produce?" and - // "which model did it run on?" without joining back to the chunk - // store. Captured the first time we see a non-empty question batch. - var sampleQuestion string + var sampleQuestionChars int var resolvedModelID string // Postprocess subspan for the trace viewer. Opened lazily after we // unmarshal the payload (so we have payload.Attempt) and closed in @@ -1563,26 +1600,21 @@ func (s *knowledgeService) processQuestionGenerationForKnowledge(ctx context.Con "retry": retryCount, "max_retry": maxRetry, } - // Surface the resolved model id and a sample question on the - // span output. These help debugging "why is question generation - // slow" — both questions ("which model was hit?") and ("what - // did it produce?") are hard to answer from logs alone. if resolvedModelID != "" { out["model_id"] = resolvedModelID } - if sampleQuestion != "" { - out["sample_question"] = sampleQuestion + if sampleQuestionChars > 0 { + out["sample_question_chars"] = sampleQuestionChars } // Treat any non-success exitStatus as a failed run; the // existing stats-string already enumerates them. qErr stays // optional for callers that want to surface a Go error. if exitStatus != "success" || qErr != nil { msg := exitStatus - var detailErr error = qErr if qErr != nil { - msg = qErr.Error() + msg += " (" + fmt.Sprintf("%T", qErr) + ")" } - s.failPostprocessSubspan(ctx, qSpan, "QUESTION_FAILED", msg, detailErr) + s.failPostprocessSubspan(ctx, qSpan, "QUESTION_FAILED", msg, nil) } else { s.endPostprocessSubspan(ctx, qSpan, out) } @@ -1765,7 +1797,7 @@ func (s *knowledgeService) processQuestionGenerationForKnowledge(ctx context.Con knowledge.Title, questionCount, customInstructions) if err != nil { llmCallFailed++ - logger.Warnf(ctx, "Failed to generate questions for chunk %s: %v", chunk.ID, err) + logger.Warnf(ctx, "Failed to generate questions for chunk %s: error_class=%T", chunk.ID, err) continue } @@ -1785,8 +1817,8 @@ func (s *knowledgeService) processQuestionGenerationForKnowledge(ctx context.Con chunk = latestChunk llmCallSuccess++ generatedQuestionsTotal += len(questions) - if sampleQuestion == "" && len(questions) > 0 { - sampleQuestion = previewText(questions[0], 200) + if sampleQuestionChars == 0 && len(questions) > 0 { + sampleQuestionChars = len([]rune(questions[0])) } // Question IDs are semantic UUIDv5 values, stable across retries and @@ -1912,7 +1944,7 @@ func (s *knowledgeService) processQuestionGenerationForChunks(ctx context.Contex generatedQuestionsTotal := 0 indexEntriesPrepared := 0 indexBatchSucceeded := false - var sampleQuestion string + var sampleQuestionChars int var resolvedModelID string var qSpan *Span var qErr error @@ -1958,15 +1990,15 @@ func (s *knowledgeService) processQuestionGenerationForChunks(ctx context.Contex if resolvedModelID != "" { out["model_id"] = resolvedModelID } - if sampleQuestion != "" { - out["sample_question"] = sampleQuestion + if sampleQuestionChars > 0 { + out["sample_question_chars"] = sampleQuestionChars } if exitStatus != "success" || qErr != nil { msg := exitStatus if qErr != nil { - msg = qErr.Error() + msg += " (" + fmt.Sprintf("%T", qErr) + ")" } - s.failPostprocessSubspan(ctx, qSpan, "QUESTION_FAILED", msg, qErr) + s.failPostprocessSubspan(ctx, qSpan, "QUESTION_FAILED", msg, nil) } else { s.endPostprocessSubspan(ctx, qSpan, out) } @@ -2148,7 +2180,7 @@ func (s *knowledgeService) processQuestionGenerationForChunks(ctx context.Contex customInstructions) if gerr != nil { llmCallFailed++ - logger.Warnf(ctx, "Failed to generate questions for chunk %s: %v", chunk.ID, gerr) + logger.Warnf(ctx, "Failed to generate questions for chunk %s: error_class=%T", chunk.ID, gerr) continue } if attemptSuperseded(ctx, s.tracker(), payload.KnowledgeID, payload.Attempt) { @@ -2164,8 +2196,8 @@ func (s *knowledgeService) processQuestionGenerationForChunks(ctx context.Contex chunk = latestChunk chunksProcessed++ generatedQuestionsTotal += len(questions) - if sampleQuestion == "" && len(questions) > 0 { - sampleQuestion = previewText(questions[0], 200) + if sampleQuestionChars == 0 && len(questions) > 0 { + sampleQuestionChars = len([]rune(questions[0])) } generatedQuestions, err := stableGeneratedQuestions(knowledge.ID, chunk, questions) @@ -3232,8 +3264,7 @@ func (s *knowledgeService) UpdateImageInfo( continue } if cImageInfo[0].OriginalURL != image.OriginalURL { - logger.Warnf(ctx, "Skipping chunk ID: %s, image URL mismatch: %s != %s", - child.ID, cImageInfo[0].OriginalURL, image.OriginalURL) + logger.Warnf(ctx, "Skipping chunk ID %s because image references differ", child.ID) continue } @@ -3280,7 +3311,7 @@ func (s *knowledgeService) UpdateImageInfo( ImageInfo: imageInfo, } addChunk = append(addChunk, captionChunk) - logger.Infof(ctx, "Created new caption chunk ID: %s for image URL: %s", captionChunk.ID, image.OriginalURL) + logger.Infof(ctx, "Created new caption chunk ID: %s", captionChunk.ID) } // Create a new OCR chunk if it doesn't exist and we have OCR data @@ -3305,7 +3336,7 @@ func (s *knowledgeService) UpdateImageInfo( ImageInfo: imageInfo, } addChunk = append(addChunk, ocrChunk) - logger.Infof(ctx, "Created new OCR chunk ID: %s for image URL: %s", ocrChunk.ID, image.OriginalURL) + logger.Infof(ctx, "Created new OCR chunk ID: %s", ocrChunk.ID) } logger.Infof(ctx, "Updated %d chunks out of %d total chunks", len(updateChunk), len(chunkChildren)+1) @@ -3484,8 +3515,7 @@ func (s *knowledgeService) ProcessManualUpdate(ctx context.Context, t *asynq.Tas } // Run manual processing (image resolution + chunking + embedding) synchronously within the worker - s.triggerManualProcessing(ctx, kb, knowledge, payload.Content, true) - return nil + return s.triggerManualProcessing(ctx, kb, knowledge, payload.Content, true) } // ProcessDocument handles Asynq document processing tasks @@ -3515,8 +3545,8 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e } ctx = context.WithValue(ctx, types.TenantInfoContextKey, tenantInfo) - logger.Infof(ctx, "Processing document task: knowledge_id=%s, file_path=%s, retry=%d/%d", - payload.KnowledgeID, payload.FilePath, retryCount, maxRetry) + logger.Infof(ctx, "Processing document task: knowledge_id=%s, has_file_path=%v, retry=%d/%d", + payload.KnowledgeID, payload.FilePath != "", retryCount, maxRetry) // 幂等性检查：获取knowledge记录 knowledge, err := s.repo.GetKnowledgeByID(ctx, payload.TenantID, payload.KnowledgeID) @@ -3655,7 +3685,7 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e if payload.FileURL != "" { // file_url import: SSRF re-check (防 DNS 重绑定), download, persist, then delegate to convert() if err := secutils.ValidateURLForSSRF(payload.FileURL); err != nil { - logger.Errorf(ctx, "File URL rejected for SSRF protection in ProcessDocument: %s, err: %v", payload.FileURL, err) + logger.Errorf(ctx, "File URL rejected for SSRF protection in ProcessDocument: error_class=%T", err) knowledge.ParseStatus = "failed" knowledge.ErrorMessage = "File URL is not allowed for security reasons" knowledge.UpdatedAt = time.Now() @@ -3667,14 +3697,14 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e resolvedFileType := payload.FileType contentBytes, err := downloadFileFromURL(ctx, payload.FileURL, &resolvedFileName, &resolvedFileType) if err != nil { - logger.Errorf(ctx, "Failed to download file from URL: %s, error: %v", payload.FileURL, err) + logger.Errorf(ctx, "Failed to download file from URL: error_class=%T", err) if isLastRetry { knowledge.ParseStatus = "failed" - knowledge.ErrorMessage = err.Error() + knowledge.ErrorMessage = "failed to download file from URL" knowledge.UpdatedAt = time.Now() _, _ = s.repo.UpdateKnowledgeIfAttemptCurrent(ctx, knowledge, attempt) } - return fmt.Errorf("failed to download file from URL: %w", err) + return fmt.Errorf("failed to download file from URL: %T", err) } if resolvedFileType != "" && !isSupportedImportExtension(resolvedFileType) { @@ -3738,7 +3768,8 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e if err != nil { logger.Warnf(ctx, "Failed to update knowledge title from extracted page title: %v", err) } else if titlePublished { - logger.Infof(ctx, "Updated knowledge title to extracted page title: %s", extractedTitle) + logger.Infof(ctx, "Updated knowledge title from extracted page title: chars=%d", + len([]rune(extractedTitle))) } } } @@ -3763,8 +3794,7 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e EnableQuestionGeneration: payload.EnableQuestionGeneration, QuestionCount: payload.QuestionCount, } - s.processChunks(ctx, kb, knowledge, passageChunks, passageOpts) - return nil + return s.processChunks(ctx, kb, knowledge, passageChunks, passageOpts) } else { // File import convertResult, err = s.convert(ctx, payload, kb, knowledge, eff, isLastRetry) @@ -3910,9 +3940,7 @@ func (s *knowledgeService) ProcessDocument(ctx context.Context, t *asynq.Task) e } // Step 4: Process chunks (vectorize + index + enqueue async tasks) - s.processChunks(ctx, kb, knowledge, chunks, processOpts) - - return nil + return s.processChunks(ctx, kb, knowledge, chunks, processOpts) } // convert handles both file and URL reading using a unified ReadRequest. @@ -3934,9 +3962,6 @@ func (s *knowledgeService) convert( "file_type": payload.FileType, "is_url": payload.URL != "", } - if payload.URL != "" { - docInput["url"] = payload.URL - } s.beginStage(ctx, knowledge.ID, types.StageDocReader, docInput) isURL := payload.URL != "" fileType := payload.FileType @@ -3950,7 +3975,7 @@ func (s *knowledgeService) convert( if isURL { if err := secutils.ValidateURLForSSRF(payload.URL); err != nil { - logger.Errorf(ctx, "URL rejected for SSRF protection: %s, err: %v", payload.URL, err) + logger.Errorf(ctx, "URL rejected for SSRF protection: error_class=%T", err) knowledge.ParseStatus = "failed" knowledge.ErrorMessage = "URL is not allowed for security reasons" knowledge.UpdatedAt = time.Now() @@ -4031,8 +4056,8 @@ func (s *knowledgeService) convert( return s.failKnowledge(ctx, knowledge, isLastRetry, "document read failed: %v", err) } if result.Error != "" { - logger.Errorf(ctx, "[convert] parser returned error kb=%s knowledge=%s file=%q type=%s engine=%q: %s", - kb.ID, knowledge.ID, req.FileName, fileType, parserEngine, result.Error) + logger.Errorf(ctx, "[convert] parser returned error kb=%s knowledge=%s type=%s engine=%q error_chars=%d", + kb.ID, knowledge.ID, fileType, parserEngine, len([]rune(result.Error))) knowledge.ParseStatus = "failed" knowledge.ErrorMessage = result.Error knowledge.UpdatedAt = time.Now() @@ -4042,7 +4067,7 @@ func (s *knowledgeService) convert( attemptFromCtx(ctx), ) s.failStage(ctx, knowledge.ID, types.StageDocReader, - werrors.ErrCodeDocReaderParseFailed, result.Error, nil) + werrors.ErrCodeDocReaderParseFailed, "document parser returned an error", nil) return nil, nil } docOutput := types.JSONMap{ @@ -4230,9 +4255,9 @@ func (s *knowledgeService) enqueueImageMultimodalTasks( task := asynq.NewTask(types.TypeImageMultimodal, payloadBytes, asynq.Queue(types.QueueMultimodal), asynq.MaxRetry(3), asynq.Timeout(30*time.Minute)) if _, err := s.task.Enqueue(task); err != nil { - logger.Warnf(ctx, "Failed to enqueue image multimodal task for %s: %v", img.ServingURL, err) + logger.Warnf(ctx, "Failed to enqueue image multimodal task: error_class=%T", err) } else { - logger.Infof(ctx, "Enqueued image:multimodal task for %s", img.ServingURL) + logger.Infof(ctx, "Enqueued image:multimodal task") } } } diff --git a/internal/application/service/knowledge_reconcile_lock.go b/internal/application/service/knowledge_reconcile_lock.go new file mode 100644 index 0000000000..436c06e040 --- /dev/null +++ b/internal/application/service/knowledge_reconcile_lock.go @@ -0,0 +1,165 @@ +package service + +import ( + "context" + "errors" + "sync" + "time" + + "github.com/google/uuid" + "github.com/redis/go-redis/v9" +) + +const knowledgeReconcileLockTTL = 2 * time.Hour + +const knowledgeReconcileRenewInterval = 30 * time.Second + +type knowledgeReconcileGate struct { + token chan struct{} + refs int +} + +var localKnowledgeReconcileGates = struct { + sync.Mutex + gates map[string]*knowledgeReconcileGate +}{ + gates: make(map[string]*knowledgeReconcileGate), +} + +var ( + renewKnowledgeReconcileLock = redis.NewScript(` + if redis.call("GET", KEYS[1]) == ARGV[1] then + return redis.call("PEXPIRE", KEYS[1], ARGV[2]) + end + return 0 + `) + releaseKnowledgeReconcileLock = redis.NewScript(` + if redis.call("GET", KEYS[1]) == ARGV[1] then + return redis.call("DEL", KEYS[1]) + end + return 0 + `) +) + +// acquireKnowledgeReconcileLock serializes mutable binding operations for one +// knowledge generation. Lite mode uses a process-local gate; standard mode +// uses an ownership-token Redis lease shared by document and image workers. +func acquireKnowledgeReconcileLock( + ctx context.Context, + client *redis.Client, + knowledgeID string, +) (context.Context, func(), error) { + return acquireKnowledgeReconcileLockWithRenewInterval( + ctx, + client, + knowledgeID, + knowledgeReconcileRenewInterval, + ) +} + +func acquireKnowledgeReconcileLockWithRenewInterval( + ctx context.Context, + client *redis.Client, + knowledgeID string, + renewInterval time.Duration, +) (context.Context, func(), error) { + if knowledgeID == "" { + return nil, nil, errors.New("knowledge reconciliation requires an ID") + } + if client == nil { + localKnowledgeReconcileGates.Lock() + gate := localKnowledgeReconcileGates.gates[knowledgeID] + if gate == nil { + gate = &knowledgeReconcileGate{token: make(chan struct{}, 1)} + localKnowledgeReconcileGates.gates[knowledgeID] = gate + } + gate.refs++ + localKnowledgeReconcileGates.Unlock() + detach := func() { + localKnowledgeReconcileGates.Lock() + gate.refs-- + if gate.refs == 0 && localKnowledgeReconcileGates.gates[knowledgeID] == gate { + delete(localKnowledgeReconcileGates.gates, knowledgeID) + } + localKnowledgeReconcileGates.Unlock() + } + select { + case gate.token <- struct{}{}: + var once sync.Once + return ctx, func() { + once.Do(func() { + <-gate.token + detach() + }) + }, nil + case <-ctx.Done(): + detach() + return nil, nil, ctx.Err() + } + } + + key := "weknora:knowledge:reconcile:" + knowledgeID + owner := uuid.NewString() + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + for { + acquired, err := client.SetNX(ctx, key, owner, knowledgeReconcileLockTTL).Result() + if err != nil { + return nil, nil, err + } + if acquired { + break + } + select { + case <-ctx.Done(): + return nil, nil, ctx.Err() + case <-ticker.C: + } + } + + lockCtx, cancelLock := context.WithCancel(ctx) + stopRenewal := make(chan struct{}) + renewalDone := make(chan struct{}) + go func() { + defer close(renewalDone) + ticker := time.NewTicker(renewInterval) + defer ticker.Stop() + for { + select { + case <-stopRenewal: + return + case <-ticker.C: + renewCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + renewed, err := renewKnowledgeReconcileLock.Run( + renewCtx, + client, + []string{key}, + owner, + knowledgeReconcileLockTTL.Milliseconds(), + ).Int64() + cancel() + if err != nil || renewed != 1 { + cancelLock() + return + } + } + } + }() + + var once sync.Once + return lockCtx, func() { + once.Do(func() { + close(stopRenewal) + <-renewalDone + cancelLock() + releaseCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _, _ = releaseKnowledgeReconcileLock.Run( + releaseCtx, + client, + []string{key}, + owner, + ).Result() + }) + }, nil +} diff --git a/internal/application/service/knowledge_reconcile_lock_test.go b/internal/application/service/knowledge_reconcile_lock_test.go new file mode 100644 index 0000000000..89ef09fe32 --- /dev/null +++ b/internal/application/service/knowledge_reconcile_lock_test.go @@ -0,0 +1,96 @@ +package service + +import ( + "context" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/require" +) + +func TestKnowledgeReconcileLocalLockSerializesSameKnowledge(t *testing.T) { + _, releaseFirst, err := acquireKnowledgeReconcileLock( + context.Background(), + nil, + "local-serialized", + ) + require.NoError(t, err) + + acquiredSecond := make(chan func(), 1) + go func() { + _, release, acquireErr := acquireKnowledgeReconcileLock( + context.Background(), + nil, + "local-serialized", + ) + if acquireErr == nil { + acquiredSecond <- release + } + }() + + select { + case <-acquiredSecond: + t.Fatal("second reconciler acquired the same knowledge lock early") + case <-time.After(25 * time.Millisecond): + } + releaseFirst() + + select { + case releaseSecond := <-acquiredSecond: + releaseSecond() + case <-time.After(time.Second): + t.Fatal("second reconciler did not acquire after release") + } +} + +func TestKnowledgeReconcileRedisLockIsOwnedAndContextAware(t *testing.T) { + server := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: server.Addr()}) + t.Cleanup(func() { _ = client.Close() }) + + _, releaseFirst, err := acquireKnowledgeReconcileLock( + context.Background(), + client, + "redis-serialized", + ) + require.NoError(t, err) + + waitCtx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond) + defer cancel() + _, _, err = acquireKnowledgeReconcileLock(waitCtx, client, "redis-serialized") + require.ErrorIs(t, err, context.DeadlineExceeded) + + releaseFirst() + _, releaseSecond, err := acquireKnowledgeReconcileLock( + context.Background(), + client, + "redis-serialized", + ) + require.NoError(t, err) + releaseSecond() +} + +func TestKnowledgeReconcileRedisLockCancelsWorkWhenLeaseIsLost(t *testing.T) { + server := miniredis.RunT(t) + client := redis.NewClient(&redis.Options{Addr: server.Addr()}) + t.Cleanup(func() { _ = client.Close() }) + + lockCtx, release, err := acquireKnowledgeReconcileLockWithRenewInterval( + context.Background(), + client, + "redis-lost", + 10*time.Millisecond, + ) + require.NoError(t, err) + defer release() + require.NoError(t, client.Del(context.Background(), "weknora:knowledge:reconcile:redis-lost").Err()) + + select { + case <-lockCtx.Done(): + require.ErrorIs(t, lockCtx.Err(), context.Canceled) + case <-time.After(time.Second): + t.Fatal("lost Redis ownership did not cancel reconciliation work") + } +} diff --git a/internal/application/service/knowledge_span_tracker_test.go b/internal/application/service/knowledge_span_tracker_test.go index cf958ce9f8..e4f2fa5908 100644 --- a/internal/application/service/knowledge_span_tracker_test.go +++ b/internal/application/service/knowledge_span_tracker_test.go @@ -149,6 +149,37 @@ func TestSpanTracker_FailSpan_CascadesDownstream(t *testing.T) { _ = postprocess } +func TestKnowledgeServiceFailStageRedactsProviderErrorDetail(t *testing.T) { + tracker, db := setupSpanTrackerTest(t) + ctx := context.Background() + _, attempt, err := tracker.OpenAttempt(ctx, "kid-redacted", "") + require.NoError(t, err) + tracker.BeginStage(ctx, "kid-redacted", attempt, types.StageDocReader, nil) + + service := &knowledgeService{spanTracker: tracker} + sentinel := "prompt=customer-body api_key=secret signed_url=https://example.invalid/?sig=secret" + service.failStage( + withAttempt(ctx, attempt), + "kid-redacted", + types.StageDocReader, + "DOCREADER_FAILED", + "document read failed", + errors.New(sentinel), + ) + + var row types.KnowledgeProcessingSpan + require.NoError(t, db. + Where("knowledge_id = ? AND attempt = ? AND name = ?", + "kid-redacted", + attempt, + types.StageDocReader, + ). + Take(&row).Error) + assert.Empty(t, row.ErrorDetail) + assert.NotContains(t, row.ErrorMessage, sentinel) + assert.Contains(t, row.ErrorMessage, "error_class=") +} + // TestSpanTracker_LookupStage_FindsAcrossProcesses simulates the // cross-process bridge an asynq worker uses: the upstream pipeline // creates the multimodal stage span, then a separate worker process diff --git a/internal/application/service/wiki_ingest_batch.go b/internal/application/service/wiki_ingest_batch.go index f049000961..d642419637 100644 --- a/internal/application/service/wiki_ingest_batch.go +++ b/internal/application/service/wiki_ingest_batch.go @@ -205,7 +205,6 @@ func (s *wikiIngestService) ProcessWikiIngest(ctx context.Context, t *asynq.Task retractHandled := 0 followUpScheduled := false totalPagesAffected := 0 - docPreview := make([]string, 0, 6) // Tunables resolved from KB.WikiConfig once we've loaded the KB. // Captured up here so the deferred stats log can observe them // regardless of which exit path we took. (Index-intro rebuild moved to @@ -219,7 +218,7 @@ func (s *wikiIngestService) ProcessWikiIngest(ctx context.Context, t *asynq.Task defer func() { logger.Infof( ctx, - "wiki ingest stats: kb=%s tenant=%d retry=%d/%d status=%s elapsed=%s mode=%s pending_ops=%d ops(ingest=%d,retract=%d) ingest(success=%d,failed=%d) retract_handled=%d pages(total=%d) followup=%v tunables(batch=%d,map_par=%d,reduce_par=%d,max_inflight=%d) preview=%s", + "wiki ingest stats: kb=%s tenant=%d retry=%d/%d status=%s elapsed=%s mode=%s pending_ops=%d ops(ingest=%d,retract=%d) ingest(success=%d,failed=%d) retract_handled=%d pages(total=%d) followup=%v tunables(batch=%d,map_par=%d,reduce_par=%d,max_inflight=%d)", payload.KnowledgeBaseID, payload.TenantID, retryCount, @@ -239,7 +238,6 @@ func (s *wikiIngestService) ProcessWikiIngest(ctx context.Context, t *asynq.Task loggedMapPar, loggedReducePar, loggedMaxInflight, - previewStringSlice(docPreview, 6), ) }() @@ -487,7 +485,6 @@ func (s *wikiIngestService) ProcessWikiIngest(ctx context.Context, t *asynq.Task mapMu.Lock() retractOps++ retractHandled++ - docPreview = append(docPreview, fmt.Sprintf("retract[%s]: %s (%d slugs)", previewText(op.KnowledgeID, 24), previewText(op.DocTitle, 48), len(slugSet))) for slug := range slugSet { slugUpdates[slug] = append(slugUpdates[slug], SlugUpdate{ @@ -511,7 +508,7 @@ func (s *wikiIngestService) ProcessWikiIngest(ctx context.Context, t *asynq.Task ingestOps++ mapMu.Unlock() - logger.Infof(mapCtx, "wiki ingest: processing document '%s' (%s)", op.DocTitle, op.KnowledgeID) + logger.Infof(mapCtx, "wiki ingest: processing knowledge %s", op.KnowledgeID) result, updates, err := s.mapOneDocument(mapCtx, chatModel, payload, op, batchCtx) if err != nil { mapMu.Lock() @@ -529,7 +526,6 @@ func (s *wikiIngestService) ProcessWikiIngest(ctx context.Context, t *asynq.Task mapMu.Lock() ingestSucceeded++ docResults = append(docResults, result) - docPreview = append(docPreview, fmt.Sprintf("ingest[%s]: title=%s summary=%s", previewText(result.KnowledgeID, 24), previewText(result.DocTitle, 40), previewText(result.Summary, 64))) for _, u := range updates { slugUpdates[u.Slug] = append(slugUpdates[u.Slug], u) } @@ -793,28 +789,20 @@ func (s *wikiIngestService) ProcessWikiIngest(ctx context.Context, t *asynq.Task if r.WikiSpan == nil { continue } - writtenPages := make([]map[string]string, 0, len(r.Pages)) - droppedPages := make([]map[string]string, 0) + writtenPages := 0 + droppedPages := 0 for _, p := range r.Pages { - entry := map[string]string{ - "slug": p.Slug, - "title": previewText(p.Title, 80), - } if _, bad := failedAdditionSlugs[p.Slug]; bad { - droppedPages = append(droppedPages, entry) + droppedPages++ continue } - writtenPages = append(writtenPages, entry) + writtenPages++ } output := types.JSONMap{ - "pages_written": len(writtenPages), - "pages_dropped": len(droppedPages), - "pages_total": len(r.Pages), - "failed_slug_writes": failedAdditionSlugCount, - "pages_written_preview": writtenPages, - } - if len(droppedPages) > 0 { - output["pages_dropped_preview"] = droppedPages + "pages_written": writtenPages, + "pages_dropped": droppedPages, + "pages_total": len(r.Pages), + "failed_slug_writes": failedAdditionSlugCount, } for k, v := range r.MapStats { output[k] = v @@ -1182,7 +1170,7 @@ func (s *wikiIngestService) mapOneDocument( chunks, err := s.chunkRepo.ListChunksByKnowledgeID(ctx, payload.TenantID, knowledgeID) if err != nil { - s.tracker().FailSpan(ctx, wikiSpan, "LIST_CHUNKS_FAILED", err.Error(), err) + s.tracker().FailSpan(ctx, wikiSpan, "LIST_CHUNKS_FAILED", fmt.Sprintf("%T", err), nil) return nil, nil, fmt.Errorf("get chunks: %w", err) } if len(chunks) == 0 { @@ -1254,17 +1242,15 @@ func (s *wikiIngestService) mapOneDocument( extractedEntities, extractedConcepts, slugItems, err = s.extractEntitiesAndConceptsNoUpsert(ctx, chatModel, payload.KnowledgeBaseID, content, lang, oldPageSlugs, batchCtx) if err != nil { logger.Warnf(ctx, "wiki ingest: legacy fallback also failed for %s: %v", knowledgeID, err) - s.tracker().FailSpan(ctx, extractSpan, "EXTRACT_FAILED", err.Error(), err) - s.tracker().FailSpan(ctx, wikiSpan, "EXTRACT_FAILED", err.Error(), err) + s.tracker().FailSpan(ctx, extractSpan, "EXTRACT_FAILED", fmt.Sprintf("%T", err), nil) + s.tracker().FailSpan(ctx, wikiSpan, "EXTRACT_FAILED", fmt.Sprintf("%T", err), nil) return nil, nil, err } } s.tracker().EndSpan(ctx, extractSpan, types.JSONMap{ - "entities": len(extractedEntities), - "concepts": len(extractedConcepts), - "pass0_fallback": pass0Failed, - "entities_preview": previewExtractedItems(extractedEntities, 8), - "concepts_preview": previewExtractedItems(extractedConcepts, 8), + "entities": len(extractedEntities), + "concepts": len(extractedConcepts), + "pass0_fallback": pass0Failed, }) // Build slug listing for Summary's wiki-link input. @@ -1329,13 +1315,10 @@ func (s *wikiIngestService) mapOneDocument( "InstructionScope": "wiki_content", }) if summaryErr != nil { - s.tracker().FailSpan(ctx, summarySpan, "SUMMARY_FAILED", summaryErr.Error(), summaryErr) + s.tracker().FailSpan(ctx, summarySpan, "SUMMARY_FAILED", fmt.Sprintf("%T", summaryErr), nil) } else { - sumLine, sumBody := splitSummaryLine(summaryContent) s.tracker().EndSpan(ctx, summarySpan, types.JSONMap{ - "chars": utf8.RuneCountInString(summaryContent), - "summary_line": previewText(sumLine, 160), - "body_preview": previewText(sumBody, 320), + "chars": utf8.RuneCountInString(summaryContent), }) } }() @@ -1351,11 +1334,9 @@ func (s *wikiIngestService) mapOneDocument( candidatesXML := renderCandidateSlugsXML(extractedEntities, extractedConcepts) citations, newSlugs, batchCount = s.classifyChunkCitations(ctx, chatModel, candidatesXML, chunks, lang, batchCtx) s.tracker().EndSpan(ctx, classifySpan, types.JSONMap{ - "cited_slugs": len(citations), - "new_slugs": len(newSlugs), - "batches": batchCount, - "top_cited": topCitedSlugs(citations, 8), - "new_slugs_sample": previewNewSlugs(newSlugs, 8), + "cited_slugs": len(citations), + "new_slugs": len(newSlugs), + "batches": batchCount, }) }() wg.Wait() @@ -1422,8 +1403,9 @@ func (s *wikiIngestService) mapOneDocument( // appends back onto the pending list so the next batch retries. // The internal retries in generateWithTemplate already exhaust // the LLM's own transient-error budget before we give up here. - logger.Errorf(ctx, "wiki ingest: generate summary failed for %s, will requeue: %v", knowledgeID, summaryErr) - s.tracker().FailSpan(ctx, wikiSpan, "SUMMARY_FAILED", summaryErr.Error(), summaryErr) + logger.Errorf(ctx, "wiki ingest: generate summary failed for %s, will requeue: error_class=%T", + knowledgeID, summaryErr) + s.tracker().FailSpan(ctx, wikiSpan, "SUMMARY_FAILED", fmt.Sprintf("%T", summaryErr), nil) return nil, nil, fmt.Errorf("generate summary: %w", summaryErr) } sumLine, sumBody := splitSummaryLine(summaryContent) @@ -1551,8 +1533,8 @@ func (s *wikiIngestService) mapOneDocument( } logger.Infof(ctx, - "wiki ingest: mapped knowledge %s title=%q candidates=%d chunks=%d batches=%d cited_chunks=%d uncited_slugs=%d new_slugs=%d updates=%d reparse_slugs=%d stale_slugs=%d pass0_fallback=%v elapsed=%s", - knowledgeID, previewText(docTitle, 80), + "wiki ingest: mapped knowledge %s candidates=%d chunks=%d batches=%d cited_chunks=%d uncited_slugs=%d new_slugs=%d updates=%d reparse_slugs=%d stale_slugs=%d pass0_fallback=%v elapsed=%s", + knowledgeID, len(slugItems), len(chunks), batchCount, len(citedChunkSet), uncited, len(newSlugs), len(updates), reparseOverlap, staleCount, pass0Failed, time.Since(docStartedAt).Round(time.Millisecond), @@ -1566,7 +1548,6 @@ func (s *wikiIngestService) mapOneDocument( // "wiki processing for this knowledge" time the user sees in the // trace viewer, not just the LLM extraction slice. mapStats := types.JSONMap{ - "doc_title": previewText(docTitle, 120), "chunks": len(chunks), "candidate_slugs": len(slugItems), "cited_chunks": len(citedChunkSet), @@ -1579,7 +1560,6 @@ func (s *wikiIngestService) mapOneDocument( "summary_chars": utf8.RuneCountInString(docSummary), "pass0_fallback": pass0Failed, "classify_batches": batchCount, - "summary_preview": previewText(docSummaryLine, 160), } return &docIngestResult{ @@ -1634,7 +1614,8 @@ func (s *wikiIngestService) extractEntitiesAndConceptsNoUpsert( var result combinedExtraction if err := json.Unmarshal([]byte(extractionJSON), &result); err != nil { - logger.Warnf(ctx, "wiki ingest: failed to parse combined extraction JSON: %v\nRaw: %s", err, extractionJSON) + logger.Warnf(ctx, "wiki ingest: failed to parse combined extraction JSON: error_class=%T chars=%d", + err, utf8.RuneCountInString(extractionJSON)) return nil, nil, nil, fmt.Errorf("parse combined extraction JSON: %w", err) } @@ -1737,7 +1718,7 @@ func (s *wikiIngestService) reduceSlugUpdates( return } if err != nil { - s.tracker().FailSpan(ctx, pageSpan, "REDUCE_FAILED", err.Error(), err) + s.tracker().FailSpan(ctx, pageSpan, "REDUCE_FAILED", fmt.Sprintf("%T", err), nil) return } if !changed { @@ -1750,13 +1731,13 @@ func (s *wikiIngestService) reduceSlugUpdates( "contributors": contributors, } if page != nil { - out["page_title"] = previewText(page.Title, 160) out["page_type"] = string(page.PageType) - out["page_summary"] = previewText(page.Summary, 200) - out["content_preview"] = previewText(page.Content, 320) + out["page_title_chars"] = utf8.RuneCountInString(page.Title) + out["page_summary_chars"] = utf8.RuneCountInString(page.Summary) + out["content_chars"] = utf8.RuneCountInString(page.Content) out["source_refs"] = len(page.SourceRefs) out["chunk_refs"] = len(page.ChunkRefs) - out["aliases"] = []string(page.Aliases) + out["aliases"] = len(page.Aliases) } s.tracker().EndSpan(ctx, pageSpan, out) }() diff --git a/internal/artifact/fault.go b/internal/artifact/fault.go new file mode 100644 index 0000000000..9a851bde80 --- /dev/null +++ b/internal/artifact/fault.go @@ -0,0 +1,46 @@ +package artifact + +import "context" + +// FaultPoint names a durability boundary that recovery tests can interrupt. +// Production code never installs an injector, so the check is a no-op outside +// explicitly instrumented tests. +type FaultPoint string + +const ( + FaultAfterProviderCall FaultPoint = "after_provider_call" + FaultAfterArtifactPut FaultPoint = "after_artifact_put" + FaultAfterChunkUpsert FaultPoint = "after_chunk_upsert" + FaultAfterVectorUpsert FaultPoint = "after_vector_upsert" + FaultAfterGraphBinding FaultPoint = "after_graph_binding" + FaultBeforeFence FaultPoint = "before_fence" + FaultAfterPublish FaultPoint = "after_publish" + FaultDuringStaleCleanup FaultPoint = "during_stale_cleanup" +) + +type faultInjectorKey struct{} + +// FaultInjector is invoked synchronously at a named durability boundary. It +// may return an error for graceful fault tests or panic to emulate a process +// crash before the worker acknowledges its task. +type FaultInjector func(FaultPoint) error + +// WithFaultInjector installs a request-scoped test injector. +func WithFaultInjector(ctx context.Context, injector FaultInjector) context.Context { + if injector == nil { + return ctx + } + return context.WithValue(ctx, faultInjectorKey{}, injector) +} + +// InjectFault executes the request-scoped injector, if any. +func InjectFault(ctx context.Context, point FaultPoint) error { + if ctx == nil { + return nil + } + injector, _ := ctx.Value(faultInjectorKey{}).(FaultInjector) + if injector == nil { + return nil + } + return injector(point) +} diff --git a/internal/artifact/runtime.go b/internal/artifact/runtime.go index da81215638..10c0750f6d 100644 --- a/internal/artifact/runtime.go +++ b/internal/artifact/runtime.go @@ -317,6 +317,9 @@ func (r *Runtime) LoadOrCompute( if err != nil { return Value{}, err } + if err := InjectFault(ctx, FaultAfterProviderCall); err != nil { + return Value{}, err + } if expected.Validate != nil { if err := expected.Validate(payload); err != nil { return Value{}, fmt.Errorf("validate processing artifact output: %w", err) @@ -343,7 +346,7 @@ func (r *Runtime) LoadOrCompute( if err != nil { return Value{}, err } - return r.freeze(ctx, expected, candidate), nil + return r.freeze(ctx, expected, candidate) }) select { @@ -502,7 +505,7 @@ func (r *Runtime) freeze( ctx context.Context, expected Expected, candidate *types.ProcessingArtifact, -) Value { +) (Value, error) { fallback := Value{ Payload: append([]byte(nil), candidate.Payload...), OutputDigest: candidate.OutputDigest, @@ -516,7 +519,7 @@ func (r *Runtime) freeze( Reason: "write_disabled", }) } - return fallback + return fallback, nil } winner, created, err := r.repository.PutIfAbsent(ctx, candidate) if err != nil { @@ -527,7 +530,7 @@ func (r *Runtime) freeze( Reason: "write_error", Err: err, }) - return fallback + return fallback, nil } if winner == nil { r.emit(Event{ @@ -537,7 +540,10 @@ func (r *Runtime) freeze( Reason: "winner_missing", Err: errors.New("artifact repository returned a nil winner"), }) - return fallback + return fallback, nil + } + if err := InjectFault(ctx, FaultAfterArtifactPut); err != nil { + return Value{}, err } payload, err := DecodeInline( winner, @@ -565,7 +571,7 @@ func (r *Runtime) freeze( Err: deleteErr, }) } - return fallback + return fallback, nil } if created { r.emit(Event{ @@ -586,7 +592,7 @@ func (r *Runtime) freeze( Payload: payload, OutputDigest: winner.OutputDigest, CacheHit: !created, - } + }, nil } func uncachedValue(payload []byte, validate func([]byte) error, err error) (Value, error) { diff --git a/internal/artifact/runtime_benchmark_test.go b/internal/artifact/runtime_benchmark_test.go new file mode 100644 index 0000000000..09da075b51 --- /dev/null +++ b/internal/artifact/runtime_benchmark_test.go @@ -0,0 +1,70 @@ +package artifact + +import ( + "context" + "testing" +) + +func benchmarkExpected(b *testing.B) Expected { + b.Helper() + key, err := BuildKey(1, testKeyMaterial()) + if err != nil { + b.Fatal(err) + } + return Expected{ + Key: key, + Codec: CodecJSONV1, + Validate: func(payload []byte) error { + return nil + }, + } +} + +func BenchmarkRuntimeColdCompute(b *testing.B) { + ctx := context.Background() + expected := benchmarkExpected(b) + payload := []byte(`{"value":"computed"}`) + b.ReportAllocs() + b.ResetTimer() + for range b.N { + runtime := NewRuntime(newMemoryArtifactRepository(), nil) + if _, err := runtime.LoadOrCompute( + ctx, + expected, + func(context.Context) ([]byte, error) { return payload, nil }, + ); err != nil { + b.Fatal(err) + } + } + b.ReportMetric(1, "provider_calls/op") +} + +func BenchmarkRuntimeWarmHit(b *testing.B) { + ctx := context.Background() + expected := benchmarkExpected(b) + payload := []byte(`{"value":"computed"}`) + runtime := NewRuntime(newMemoryArtifactRepository(), nil) + if _, err := runtime.LoadOrCompute( + ctx, + expected, + func(context.Context) ([]byte, error) { return payload, nil }, + ); err != nil { + b.Fatal(err) + } + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if _, err := runtime.LoadOrCompute( + ctx, + expected, + func(context.Context) ([]byte, error) { + b.Fatal("warm artifact unexpectedly called the provider") + return nil, nil + }, + ); err != nil { + b.Fatal(err) + } + } + b.ReportMetric(0, "provider_calls/op") +} diff --git a/internal/artifact/runtime_test.go b/internal/artifact/runtime_test.go index d69162f8bb..adb6290886 100644 --- a/internal/artifact/runtime_test.go +++ b/internal/artifact/runtime_test.go @@ -164,6 +164,61 @@ func TestRuntimeCorruptEntrySelfHeals(t *testing.T) { assert.Equal(t, 1, repository.deleted) } +func TestRuntimeRecoveryAfterProviderCallFaultRecomputes(t *testing.T) { + repository := newMemoryArtifactRepository() + runtime := NewRuntime(repository, nil) + expected := testExpected(t) + var calls atomic.Int32 + injected := errors.New("injected after provider call") + ctx := WithFaultInjector(context.Background(), func(point FaultPoint) error { + if point == FaultAfterProviderCall { + return injected + } + return nil + }) + compute := func(context.Context) ([]byte, error) { + calls.Add(1) + return []byte(`{"value":"computed"}`), nil + } + + _, err := runtime.LoadOrCompute(ctx, expected, compute) + require.ErrorIs(t, err, injected) + assert.Empty(t, repository.values) + + value, err := runtime.LoadOrCompute(context.Background(), expected, compute) + require.NoError(t, err) + assert.Equal(t, `{"value":"computed"}`, string(value.Payload)) + assert.Equal(t, int32(2), calls.Load()) +} + +func TestRuntimeRecoveryAfterArtifactPutFaultUsesCommittedArtifact(t *testing.T) { + repository := newMemoryArtifactRepository() + runtime := NewRuntime(repository, nil) + expected := testExpected(t) + var calls atomic.Int32 + injected := errors.New("injected after artifact put") + ctx := WithFaultInjector(context.Background(), func(point FaultPoint) error { + if point == FaultAfterArtifactPut { + return injected + } + return nil + }) + compute := func(context.Context) ([]byte, error) { + calls.Add(1) + return []byte(`{"value":"computed"}`), nil + } + + _, err := runtime.LoadOrCompute(ctx, expected, compute) + require.ErrorIs(t, err, injected) + require.Len(t, repository.values, 1) + + value, err := runtime.LoadOrCompute(context.Background(), expected, compute) + require.NoError(t, err) + assert.True(t, value.CacheHit) + assert.Equal(t, `{"value":"computed"}`, string(value.Payload)) + assert.Equal(t, int32(1), calls.Load()) +} + func TestRuntimeSingleflightComputesOnce(t *testing.T) { repository := newMemoryArtifactRepository() runtime := NewRuntime(repository, nil) diff --git a/internal/database/migration_version_test.go b/internal/database/migration_version_test.go new file mode 100644 index 0000000000..997bf607e0 --- /dev/null +++ b/internal/database/migration_version_test.go @@ -0,0 +1,52 @@ +package database + +import ( + "os" + "regexp" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestMigrationsHaveUniqueVersionNumbers(t *testing.T) { + for _, directory := range []string{"versioned", "sqlite"} { + t.Run(directory, func(t *testing.T) { + assertUniqueMigrationVersions(t, "../../migrations/"+directory) + }) + } +} + +func assertUniqueMigrationVersions(t *testing.T, directory string) { + t.Helper() + entries, err := os.ReadDir(directory) + require.NoError(t, err) + pattern := regexp.MustCompile(`^([0-9]{6})_.+\.(up|down)\.sql$`) + seen := make(map[string]map[string]string) + for _, entry := range entries { + if entry.IsDir() { + continue + } + match := pattern.FindStringSubmatch(entry.Name()) + if match == nil { + continue + } + version, direction := match[1], match[2] + if seen[version] == nil { + seen[version] = make(map[string]string) + } + if previous := seen[version][direction]; previous != "" { + t.Fatalf( + "migration version %s has duplicate %s files: %s and %s", + version, + direction, + previous, + entry.Name(), + ) + } + seen[version][direction] = entry.Name() + } + for version, directions := range seen { + require.NotEmpty(t, directions["up"], "migration %s is missing its up file", version) + require.NotEmpty(t, directions["down"], "migration %s is missing its down file", version) + } +} diff --git a/internal/database/processing_artifact_migration_test.go b/internal/database/processing_artifact_migration_test.go index 65417432da..4a3b523e95 100644 --- a/internal/database/processing_artifact_migration_test.go +++ b/internal/database/processing_artifact_migration_test.go @@ -13,7 +13,7 @@ func TestSQLiteProcessingArtifactMigrationAndTenantBoundary(t *testing.T) { require.NoError(t, err) t.Cleanup(func() { _ = db.Close() }) - up, err := os.ReadFile("../../migrations/sqlite/000001_processing_artifacts.up.sql") + up, err := os.ReadFile("../../migrations/sqlite/000002_processing_artifacts.up.sql") require.NoError(t, err) _, err = db.Exec(string(up)) require.NoError(t, err) @@ -47,7 +47,7 @@ func TestSQLiteProcessingArtifactMigrationAndTenantBoundary(t *testing.T) { ) require.Error(t, err, "each knowledge must have exactly one attempt allocator row") - down, err := os.ReadFile("../../migrations/sqlite/000001_processing_artifacts.down.sql") + down, err := os.ReadFile("../../migrations/sqlite/000002_processing_artifacts.down.sql") require.NoError(t, err) _, err = db.Exec(string(down)) require.NoError(t, err) diff --git a/internal/models/chat/artifact_cache_test.go b/internal/models/chat/artifact_cache_test.go index 929e82ce8f..2527c962d2 100644 --- a/internal/models/chat/artifact_cache_test.go +++ b/internal/models/chat/artifact_cache_test.go @@ -109,6 +109,29 @@ func TestArtifactCachedChatDoesNotCacheInvalidOutput(t *testing.T) { assert.Equal(t, 2, provider.callCount()) } +func TestArtifactCachedChatStageSchemaAndPromptInvalidateExactly(t *testing.T) { + provider := &countingChat{response: "canonical"} + cached := NewArtifactCachedChat(provider, setupChatArtifactRuntime(t), chatArtifactConfig()) + call := func(stage, schema, prompt string) { + _, err := cached.Chat( + WithArtifactStage(context.Background(), ArtifactStage{ + Stage: stage, + OutputSchema: schema, + }), + []Message{{Role: "user", Content: prompt}}, + &ChatOptions{Temperature: 0}, + ) + require.NoError(t, err) + } + + call("wiki_map", "wiki.map.v1", "prompt-v1") + call("wiki_map", "wiki.map.v1", "prompt-v1") + call("wiki_map", "wiki.map.v1", "prompt-v2") + call("wiki_map", "wiki.map.v2", "prompt-v2") + call("graph_extract", "graph.v1", "prompt-v2") + assert.Equal(t, 4, provider.callCount()) +} + func TestChatArtifactConfigExcludesCredentialRotation(t *testing.T) { model := &types.Model{ ID: "chat-id", diff --git a/internal/models/embedding/artifact_cache.go b/internal/models/embedding/artifact_cache.go index ad9c21e71a..e2dad2ac17 100644 --- a/internal/models/embedding/artifact_cache.go +++ b/internal/models/embedding/artifact_cache.go @@ -230,6 +230,12 @@ func (e *artifactCachedEmbedder) batchEmbed( if providerErr != nil { return nil, providerErr } + if faultErr := artifact.InjectFault( + computeContext, + artifact.FaultAfterProviderCall, + ); faultErr != nil { + return nil, faultErr + } providerCalled = true if validateErr := ValidateEmbeddingBatch( vectors, diff --git a/internal/models/embedding/artifact_cache_test.go b/internal/models/embedding/artifact_cache_test.go index 74304d4789..a8b6799592 100644 --- a/internal/models/embedding/artifact_cache_test.go +++ b/internal/models/embedding/artifact_cache_test.go @@ -141,6 +141,24 @@ func TestArtifactCachedEmbeddingUsesPartialHits(t *testing.T) { assert.Equal(t, [][]string{{"missing"}}, provider.calls()) } +func TestArtifactCachedEmbeddingModelIdentityInvalidatesOnlyEmbedding(t *testing.T) { + runtime, _ := setupEmbeddingArtifactRuntime(t) + provider := &countingEmbedder{dimensions: 3} + firstConfig := embeddingArtifactConfig() + first := NewArtifactCachedEmbedder(provider, runtime, firstConfig) + _, err := first.Embed(documentEmbeddingContext(), "same canonical chunk") + require.NoError(t, err) + _, err = first.Embed(documentEmbeddingContext(), "same canonical chunk") + require.NoError(t, err) + + secondConfig := firstConfig + secondConfig.Processor.ModelID = "model-id-v2" + second := NewArtifactCachedEmbedder(provider, runtime, secondConfig) + _, err = second.Embed(documentEmbeddingContext(), "same canonical chunk") + require.NoError(t, err) + assert.Len(t, provider.calls(), 2) +} + func TestArtifactCachedEmbeddingBatchSuppressesConcurrentProviderCalls(t *testing.T) { server := miniredis.RunT(t) client := redis.NewClient(&redis.Options{Addr: server.Addr()}) diff --git a/internal/models/vlm/artifact_cache_test.go b/internal/models/vlm/artifact_cache_test.go index 03d8fca3c4..be2839203b 100644 --- a/internal/models/vlm/artifact_cache_test.go +++ b/internal/models/vlm/artifact_cache_test.go @@ -68,8 +68,12 @@ func TestArtifactCachedVLMUsesExactImagesAndPrompt(t *testing.T) { _, err = cached.Predict(ctx, [][]byte{[]byte{0, 1, 3}}, " describe ") require.NoError(t, err) - _, err = cached.Predict(ctx, [][]byte{[]byte{0, 1, 2}}, "describe") + provider.mu.Lock() + provider.response = "changed caption" + provider.mu.Unlock() + changed, err := cached.Predict(ctx, [][]byte{[]byte{0, 1, 2}}, "describe") require.NoError(t, err) + assert.Equal(t, "changed caption", changed) assert.Equal(t, 3, provider.callCount()) } diff --git a/migrations/sqlite/000001_processing_artifacts.down.sql b/migrations/sqlite/000002_processing_artifacts.down.sql similarity index 52% rename from migrations/sqlite/000001_processing_artifacts.down.sql rename to migrations/sqlite/000002_processing_artifacts.down.sql index 217f81cdca..c99e6da03c 100644 --- a/migrations/sqlite/000001_processing_artifacts.down.sql +++ b/migrations/sqlite/000002_processing_artifacts.down.sql @@ -1,2 +1,3 @@ +-- SQLite rollback for migration 000002 processing artifacts and attempt fencing. DROP TABLE IF EXISTS knowledge_attempt_counters; DROP TABLE IF EXISTS processing_artifacts; diff --git a/migrations/sqlite/000001_processing_artifacts.up.sql b/migrations/sqlite/000002_processing_artifacts.up.sql similarity index 95% rename from migrations/sqlite/000001_processing_artifacts.up.sql rename to migrations/sqlite/000002_processing_artifacts.up.sql index 1f2fa3c80c..41faf08b5c 100644 --- a/migrations/sqlite/000001_processing_artifacts.up.sql +++ b/migrations/sqlite/000002_processing_artifacts.up.sql @@ -1,3 +1,4 @@ +-- Migration 000002: versioned processing artifacts and attempt fencing. CREATE TABLE IF NOT EXISTS processing_artifacts ( id INTEGER PRIMARY KEY AUTOINCREMENT, tenant_id INTEGER NOT NULL, diff --git a/migrations/versioned/000077_processing_artifacts.down.sql b/migrations/versioned/000079_processing_artifacts.down.sql similarity index 52% rename from migrations/versioned/000077_processing_artifacts.down.sql rename to migrations/versioned/000079_processing_artifacts.down.sql index e6969c0e4e..3a72071096 100644 --- a/migrations/versioned/000077_processing_artifacts.down.sql +++ b/migrations/versioned/000079_processing_artifacts.down.sql @@ -1,3 +1,3 @@ --- Roll back processing-artifact persistence and attempt fencing. +-- Roll back migration 000079 processing-artifact persistence and attempt fencing. DROP TABLE IF EXISTS knowledge_attempt_counters; DROP TABLE IF EXISTS processing_artifacts; diff --git a/migrations/versioned/000077_processing_artifacts.up.sql b/migrations/versioned/000079_processing_artifacts.up.sql similarity index 92% rename from migrations/versioned/000077_processing_artifacts.up.sql rename to migrations/versioned/000079_processing_artifacts.up.sql index 34d71299db..01d7bc96cc 100644 --- a/migrations/versioned/000077_processing_artifacts.up.sql +++ b/migrations/versioned/000079_processing_artifacts.up.sql @@ -1,4 +1,4 @@ -DO $$ BEGIN RAISE NOTICE '[Migration 000077] Creating processing artifacts...'; END $$; +DO $$ BEGIN RAISE NOTICE '[Migration 000079] Creating processing artifacts...'; END $$; CREATE TABLE IF NOT EXISTS processing_artifacts ( id BIGSERIAL PRIMARY KEY, @@ -41,4 +41,4 @@ CREATE TABLE IF NOT EXISTS knowledge_attempt_counters ( CHECK (last_attempt >= 0) ); -DO $$ BEGIN RAISE NOTICE '[Migration 000077] Processing artifacts ready'; END $$; +DO $$ BEGIN RAISE NOTICE '[Migration 000079] Processing artifacts ready'; END $$; From d29e0d606c15e2f881162d5d8ee1ed973467c605 Mon Sep 17 00:00:00 2001 From: Even1115 Date: Mon, 3 Aug 2026 09:32:10 +0800 Subject: [PATCH 04/11] test: align reparse failure with reconciliation --- .../service/knowledge_batch_reparse_test.go | 34 ++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/internal/application/service/knowledge_batch_reparse_test.go b/internal/application/service/knowledge_batch_reparse_test.go index 5babdab620..2d9153f6aa 100644 --- a/internal/application/service/knowledge_batch_reparse_test.go +++ b/internal/application/service/knowledge_batch_reparse_test.go @@ -33,6 +33,36 @@ func (r *reparseFailureKnowledgeRepo) UpdateKnowledge( return nil } +func (r *reparseFailureKnowledgeRepo) UpdateKnowledgeIfAttemptCurrent( + _ context.Context, + knowledge *types.Knowledge, + _ int, +) (bool, error) { + r.knowledge = knowledge + r.updateCalls++ + return true, nil +} + +func (r *reparseFailureKnowledgeRepo) UpdateKnowledgeColumnsIfAttemptCurrent( + _ context.Context, + _ uint64, + _ string, + _ int, + values map[string]interface{}, +) (bool, error) { + if status, ok := values["parse_status"].(string); ok { + r.knowledge.ParseStatus = status + } + if message, ok := values["error_message"].(string); ok { + r.knowledge.ErrorMessage = message + } + if pending, ok := values["pending_subtasks_count"].(int); ok { + r.knowledge.PendingSubtasksCount = pending + } + r.updateCalls++ + return true, nil +} + func (r *reparseFailureKnowledgeRepo) UpdateKnowledgeColumn( _ context.Context, _ string, @@ -91,7 +121,9 @@ func TestReparseKnowledgeManualEnqueueFailureIsVisible(t *testing.T) { require.Error(t, err) require.Same(t, knowledge, got) require.Equal(t, types.ParseStatusFailed, knowledge.ParseStatus) - require.Equal(t, "disabled", knowledge.EnableStatus) + // Reparse is non-destructive: an enqueue failure must leave the last + // published generation available even though the new attempt failed. + require.Equal(t, "enabled", knowledge.EnableStatus) require.Equal(t, "Failed to enqueue processing task", knowledge.ErrorMessage) require.GreaterOrEqual(t, repo.updateCalls, 2, "pending and failed states must both be persisted") } From c352499131f165599dfa7d6ef95e9c863c0ec96b Mon Sep 17 00:00:00 2001 From: Even1115 Date: Mon, 3 Aug 2026 11:31:59 +0800 Subject: [PATCH 05/11] test: define fake embedding dimensions --- .../service/retriever/keywords_vector_hybrid_indexer_test.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/internal/application/service/retriever/keywords_vector_hybrid_indexer_test.go b/internal/application/service/retriever/keywords_vector_hybrid_indexer_test.go index 3d37490986..5a4ae93d89 100644 --- a/internal/application/service/retriever/keywords_vector_hybrid_indexer_test.go +++ b/internal/application/service/retriever/keywords_vector_hybrid_indexer_test.go @@ -16,6 +16,10 @@ type capturingEmbedder struct { batchTexts []string } +func (e *capturingEmbedder) GetDimensions() int { + return 1 +} + func (e *capturingEmbedder) Embed(ctx context.Context, text string) ([]float32, error) { e.text = text return []float32{1}, nil From 91b2d31cd4d95b65daaa7c6c773ef17470587e33 Mon Sep 17 00:00:00 2001 From: ttommybot Date: Mon, 3 Aug 2026 22:24:28 +0800 Subject: [PATCH 06/11] fix: unify organization settings modal scrolling --- .../OrganizationSettingsModal.test.mjs | 19 +++++++ .../OrganizationSettingsModal.vue | 55 ++++++++++++++++++- 2 files changed, 71 insertions(+), 3 deletions(-) diff --git a/frontend/src/views/organization/OrganizationSettingsModal.test.mjs b/frontend/src/views/organization/OrganizationSettingsModal.test.mjs index 299bd85546..9c761b639d 100644 --- a/frontend/src/views/organization/OrganizationSettingsModal.test.mjs +++ b/frontend/src/views/organization/OrganizationSettingsModal.test.mjs @@ -14,3 +14,22 @@ test('reviewing join requests goes through the store and refreshes modal data', const refreshCalls = source.match(/await refreshOrganizationAfterReview\(\)/g) ?? [] assert.equal(refreshCalls.length, 2) }) + +test('organization settings use one outer content scroller and reset it on navigation', () => { + assert.match(source, /ref="contentWrapperRef" class="content-wrapper"/) + assert.match(source, /class="data-table-shell members-table-shell"/) + assert.match(source, /contentWrapperRef\.value\?\.scrollTo\(\{ top: 0, behavior: 'auto' \}\)/) + assert.match(source, /watch\(\(\) => props\.visible,[\s\S]*?void scrollContentToTop\(\)/) + assert.match(source, /\}, \{ immediate: true \}\)\s*\n\s*watch\(\(\) => props\.orgId/) + assert.match(source, /watch\(\(\) => props\.orgId,[\s\S]*?void scrollContentToTop\(\)/) + assert.match(source, /watch\(currentSection,[\s\S]*?void scrollContentToTop\(\)/) + assert.match(source, /\.members-table-shell\s*\{[\s\S]*?overflow: visible;[\s\S]*?\.t-table__content/) +}) + +test('organization settings lock and restore background scrolling', () => { + assert.match(source, /document\.body\.style\.overflow = 'hidden'/) + assert.match(source, /document\.body\.style\.overflow = previousBodyOverflow/) + assert.match(source, /onBeforeUnmount\(\(\) => \{\s*unlockBackgroundScroll\(\)/) + assert.match(source, /\.settings-overlay\s*\{[\s\S]*?overscroll-behavior: none;/) + assert.match(source, /\.content-wrapper\s*\{[\s\S]*?overscroll-behavior: contain;/) +}) diff --git a/frontend/src/views/organization/OrganizationSettingsModal.vue b/frontend/src/views/organization/OrganizationSettingsModal.vue index e9ebd96c7a..f5b79c742c 100644 --- a/frontend/src/views/organization/OrganizationSettingsModal.vue +++ b/frontend/src/views/organization/OrganizationSettingsModal.vue @@ -38,7 +38,7 @@
-
+
@@ -456,7 +456,7 @@ ? $t('organization.members.emptySearch', { q: memberSearchQuery }) : $t('organization.noMembers')" />
-
+