yawan
f7fc992f57
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s
716 lines
33 KiB
Go
716 lines
33 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
apprepo "github.com/Tencent/WeKnora/internal/application/repository"
|
|
"github.com/Tencent/WeKnora/internal/models/chat"
|
|
"github.com/Tencent/WeKnora/internal/types"
|
|
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
|
"github.com/hibiken/asynq"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/sqlite"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/logger"
|
|
)
|
|
|
|
func openWikiCheckpointDB(t *testing.T, path string) *gorm.DB {
|
|
t.Helper()
|
|
db, err := gorm.Open(sqlite.Open(path), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
|
|
require.NoError(t, err)
|
|
sqlDB, err := db.DB()
|
|
require.NoError(t, err)
|
|
sqlDB.SetMaxOpenConns(1)
|
|
t.Cleanup(func() { _ = sqlDB.Close() })
|
|
require.NoError(t, db.AutoMigrate(&wikiIngestCheckpoint{}))
|
|
return db
|
|
}
|
|
|
|
func checkpointChunks(count int) []*types.Chunk {
|
|
chunks := make([]*types.Chunk, count)
|
|
for i := range chunks {
|
|
chunks[i] = &types.Chunk{ID: fmt.Sprintf("chunk-%03d", i), TenantID: 7,
|
|
KnowledgeBaseID: "kb", KnowledgeID: "doc", ChunkType: types.ChunkTypeText, IsEnabled: true,
|
|
ChunkIndex: i, Content: fmt.Sprintf("batch-%03d Alpha ", i) + strings.Repeat("text ", 2398)}
|
|
}
|
|
return chunks
|
|
}
|
|
|
|
func TestWikiCheckpointCitationResumesAcrossDatabaseReopen(t *testing.T) {
|
|
path := filepath.Join(t.TempDir(), "checkpoints.sqlite")
|
|
db := openWikiCheckpointDB(t, path)
|
|
svc := &wikiIngestService{checkpointDB: db}
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "fingerprint"}
|
|
chunks := checkpointChunks(wikiCitationSlice + 1)
|
|
var calls atomic.Int32
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
|
calls.Add(1)
|
|
return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]}}`}, nil
|
|
}}
|
|
got, _, completed, err := svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
|
|
require.ErrorIs(t, err, errWikiCheckpointYield)
|
|
require.Nil(t, got)
|
|
require.Equal(t, wikiCitationSlice, completed)
|
|
require.Equal(t, int32(wikiCitationSlice), calls.Load())
|
|
sqlDB, _ := db.DB()
|
|
require.NoError(t, sqlDB.Close())
|
|
|
|
svc = &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, path)}
|
|
got, _, completed, err = svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
|
|
require.NoError(t, err)
|
|
require.Len(t, got["entity/alpha"], len(chunks))
|
|
require.Equal(t, len(chunks), completed)
|
|
require.Equal(t, int32(len(chunks)), calls.Load())
|
|
_, _, _, err = svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
|
|
require.NoError(t, err)
|
|
require.Equal(t, int32(len(chunks)), calls.Load(), "restart must not regenerate validated batches")
|
|
}
|
|
|
|
func TestWikiCheckpointCitationFailureKeepsSuccessfulSibling(t *testing.T) {
|
|
db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
|
|
svc := &wikiIngestService{checkpointDB: db}
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "fingerprint"}
|
|
chunks := checkpointChunks(2)
|
|
var calls [2]atomic.Int32
|
|
var failed atomic.Bool
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(ctx context.Context, messages []chat.Message) (*types.ChatResponse, error) {
|
|
if strings.Contains(messages[0].Content, "batch-000") {
|
|
calls[0].Add(1)
|
|
return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]}}`}, nil
|
|
}
|
|
calls[1].Add(1)
|
|
if failed.CompareAndSwap(false, true) {
|
|
// Fail only after the sibling is durably persisted.
|
|
deadline := time.NewTimer(time.Second)
|
|
defer deadline.Stop()
|
|
for {
|
|
var count int64
|
|
db.Model(&wikiIngestCheckpoint{}).Count(&count)
|
|
if count > 0 {
|
|
break
|
|
}
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
case <-deadline.C:
|
|
return nil, errors.New("checkpoint not persisted")
|
|
case <-time.After(time.Millisecond):
|
|
}
|
|
}
|
|
return nil, errors.New("invalid request")
|
|
}
|
|
return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]}}`}, nil
|
|
}}
|
|
got, _, completed, err := svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
|
|
require.Error(t, err)
|
|
require.Nil(t, got)
|
|
require.Equal(t, 1, completed)
|
|
svc = &wikiIngestService{checkpointDB: db}
|
|
got, _, completed, err = svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
|
|
require.NoError(t, err)
|
|
require.Equal(t, 2, completed)
|
|
require.Len(t, got["entity/alpha"], 2)
|
|
require.Equal(t, int32(1), calls[0].Load())
|
|
require.Equal(t, int32(2), calls[1].Load())
|
|
}
|
|
|
|
func TestWikiCheckpointFreshCitationCannotReplaceInvalidType(t *testing.T) {
|
|
for _, test := range []struct{ name, output string }{
|
|
{"invalid-new-slug", `{"citations":{},"new_slugs":[{"type":"entity","name":"New","slug":"concept/new","source_chunks":["c000"]}]}`},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
|
|
svc := &wikiIngestService{checkpointDB: db}
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "fingerprint"}
|
|
chunks := []*types.Chunk{{ID: "first", Content: "first", ChunkIndex: 0}, {ID: "second", Content: "second", ChunkIndex: 1}}
|
|
var calls atomic.Int32
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
|
if calls.Add(1) == 1 {
|
|
return &types.ChatResponse{Content: test.output}, nil
|
|
}
|
|
return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]}}`}, nil
|
|
}}
|
|
got, _, completed, err := svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
|
|
require.Error(t, err)
|
|
require.Zero(t, completed)
|
|
require.Nil(t, got)
|
|
require.Equal(t, int32(1), calls.Load())
|
|
var saved int64
|
|
require.NoError(t, db.Model(&wikiIngestCheckpoint{}).Where("stage = ?", "citation").Count(&saved).Error)
|
|
require.Zero(t, saved, "invalid types cannot become checkpoints")
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestWikiCheckpointUnknownSlugCannotUseSiblingDeclarationOrSavePartialBatch(t *testing.T) {
|
|
db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
|
|
svc := &wikiIngestService{checkpointDB: db}
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "fingerprint"}
|
|
chunks := []*types.Chunk{{ID: "first", Content: "first", ChunkIndex: 0}, {ID: "second", Content: "second", ChunkIndex: 1}}
|
|
var calls atomic.Int32
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
|
call := calls.Add(1)
|
|
if call == 1 {
|
|
return &types.ChatResponse{Content: `{"citations":`, FinishReason: "length"}, nil
|
|
}
|
|
if call == 2 {
|
|
return &types.ChatResponse{Content: `{"citations":{},"new_slugs":[{"type":"concept","name":"New","slug":"concept/new","source_chunks":["c000"]}]}`}, nil
|
|
}
|
|
return &types.ChatResponse{Content: `{"citations":{"concept/new":["c000"]}}`}, nil
|
|
}}
|
|
got, _, completed, err := svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
|
|
require.ErrorContains(t, err, "citation partition 2/2")
|
|
require.ErrorContains(t, err, "undeclared slug")
|
|
require.NotContains(t, err.Error(), "persisted")
|
|
require.Nil(t, got)
|
|
require.Zero(t, completed)
|
|
require.Equal(t, int32(3), calls.Load(), "semantic errors have the same finite retry budget")
|
|
var saved int64
|
|
require.NoError(t, db.Model(&wikiIngestCheckpoint{}).Count(&saved).Error)
|
|
require.Zero(t, saved)
|
|
}
|
|
|
|
func TestWikiCheckpointCorruptCitationDoesNotTriggerModelRepair(t *testing.T) {
|
|
db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
|
|
svc := &wikiIngestService{checkpointDB: db}
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "fingerprint"}
|
|
chunks := checkpointChunks(1)
|
|
batch := splitChunksIntoCitationBatches(chunks)[0]
|
|
key := wikiCheckpointHash([]any{"entity/alpha", "English", wikiCheckpointChunks(batch.chunks)})
|
|
require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "citation", key, citationBatchResult{Citations: map[string][]string{"concept/undeclared": {chunks[0].ID}}}))
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
|
t.Fatal("corrupt persisted work must be investigated without invoking the model")
|
|
return nil, nil
|
|
}}
|
|
got, _, completed, err := svc.classifyChunkCitations(context.Background(), model, "entity/alpha", chunks, "English", &WikiBatchContext{}, scope)
|
|
require.ErrorContains(t, err, "invalid persisted citation")
|
|
require.Nil(t, got)
|
|
require.Zero(t, completed)
|
|
}
|
|
|
|
type checkpointModelService struct {
|
|
interfaces.ModelService
|
|
model *types.Model
|
|
chat chat.Chat
|
|
}
|
|
|
|
func (m *checkpointModelService) GetModelByID(context.Context, string) (*types.Model, error) {
|
|
return m.model, nil
|
|
}
|
|
func (m *checkpointModelService) GetChatModel(context.Context, string) (chat.Chat, error) {
|
|
return m.chat, nil
|
|
}
|
|
|
|
func TestWikiCheckpointSignatureStabilityAndInvalidation(t *testing.T) {
|
|
db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
|
|
m := &checkpointModelService{model: &types.Model{ID: "model", Name: "test", TenantID: 7}}
|
|
svc := &wikiIngestService{checkpointDB: db, modelService: m}
|
|
model := &wikiScriptedChat{id: "model"}
|
|
payload := WikiIngestPayload{TenantID: 7, KnowledgeBaseID: "kb"}
|
|
op := WikiPendingOp{KnowledgeID: "doc", dbID: 42}
|
|
chunks := checkpointChunks(2)
|
|
batch := &WikiBatchContext{ExtractionGranularity: types.WikiExtractionStandard}
|
|
signature := func() string {
|
|
scope, err := svc.newWikiCheckpointScope(context.Background(), model, payload, op, chunks, "content", "title", "English", batch)
|
|
require.NoError(t, err)
|
|
return scope.Fingerprint
|
|
}
|
|
original := signature()
|
|
m.model.UpdatedAt = time.Now()
|
|
m.model.DisplayName = "new display label"
|
|
chunks[0].UpdatedAt = time.Now()
|
|
chunks[0].Status = int(types.ChunkStatusIndexed)
|
|
chunks[0].IndexStatus = "processing"
|
|
chunks[0], chunks[1] = chunks[1], chunks[0]
|
|
require.Equal(t, original, signature(), "runtime timestamps/order must not invalidate resumable work")
|
|
m.model.Parameters.BaseURL = "https://new-provider.invalid"
|
|
require.NotEqual(t, original, signature())
|
|
m.model.Parameters.BaseURL = ""
|
|
chunks[0].Content += "updated source"
|
|
require.NotEqual(t, original, signature())
|
|
chunks[0].Content = strings.TrimSuffix(chunks[0].Content, "updated source")
|
|
batch.ContentInstructions = "changed rules"
|
|
require.NotEqual(t, original, signature())
|
|
batch.ContentInstructions = ""
|
|
op.dbID++
|
|
require.NotEqual(t, original, signature(), "an explicitly new ingest must reconsider existing contributions")
|
|
chunks[0].TenantID = 8
|
|
_, err := svc.newWikiCheckpointScope(context.Background(), model, payload, op, chunks, "content", "title", "English", batch)
|
|
require.ErrorContains(t, err, "scope mismatch")
|
|
}
|
|
|
|
func TestWikiCheckpointIsolationAndCorruptCitationRejection(t *testing.T) {
|
|
svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))}
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "v1"}
|
|
require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "summary", "key", "complete"))
|
|
for _, other := range []wikiCheckpointScope{{8, "kb", "doc", "v1"}, {7, "other", "doc", "v1"}, {7, "kb", "other", "v1"}, {7, "kb", "doc", "v2"}} {
|
|
var value string
|
|
hit, err := svc.loadWikiCheckpoint(context.Background(), &other, "summary", "key", &value)
|
|
require.NoError(t, err)
|
|
require.False(t, hit)
|
|
}
|
|
batch := splitChunksIntoCitationBatches(checkpointChunks(1))[0]
|
|
allowed := map[string]bool{"entity/alpha": true}
|
|
for _, result := range []citationBatchResult{
|
|
{Citations: map[string][]string{"entity/alpha": {"another-document-chunk"}}},
|
|
{Citations: map[string][]string{"concept/unknown": {"chunk-000"}}},
|
|
{Citations: map[string][]string{}, NewSlugs: []newSlugFromCitation{{Name: "wrong type", Type: "entity", Slug: "concept/wrong"}}},
|
|
} {
|
|
require.Error(t, validateWikiCitationCheckpoint(&result, batch, allowed))
|
|
}
|
|
}
|
|
|
|
type checkpointKnowledge struct{ interfaces.KnowledgeService }
|
|
|
|
func (*checkpointKnowledge) GetKnowledgeByIDOnly(_ context.Context, id string) (*types.Knowledge, error) {
|
|
return &types.Knowledge{ID: id, TenantID: 7, KnowledgeBaseID: "kb", Title: "title"}, nil
|
|
}
|
|
|
|
type checkpointWikiPages struct {
|
|
interfaces.WikiPageService
|
|
mu sync.Mutex
|
|
pages map[string]*types.WikiPage
|
|
writes int
|
|
publishErr error
|
|
}
|
|
|
|
func (s *checkpointWikiPages) GetPageBySlug(_ context.Context, _, slug string) (*types.WikiPage, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if p := s.pages[slug]; p != nil {
|
|
out := *p
|
|
return &out, nil
|
|
}
|
|
return nil, apprepo.ErrWikiPageNotFound
|
|
}
|
|
func (s *checkpointWikiPages) ListBySlugs(_ context.Context, _ string, slugs []string) (map[string]*types.WikiPageLite, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
result := make(map[string]*types.WikiPageLite)
|
|
for _, slug := range slugs {
|
|
if page := s.pages[slug]; page != nil {
|
|
result[slug] = &types.WikiPageLite{Slug: page.Slug, Title: page.Title, PageType: page.PageType}
|
|
}
|
|
}
|
|
return result, nil
|
|
}
|
|
func (s *checkpointWikiPages) FindPagesByNormalizedTitles(_ context.Context, _ string, kind string, titles []string) ([]*types.WikiPageLite, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
var result []*types.WikiPageLite
|
|
for _, page := range s.pages {
|
|
for _, title := range titles {
|
|
if page.PageType == kind && normalizeWikiIdentityTitle(page.Title) == title {
|
|
result = append(result, &types.WikiPageLite{Slug: page.Slug, Title: page.Title, PageType: page.PageType})
|
|
break
|
|
}
|
|
}
|
|
}
|
|
return result, nil
|
|
}
|
|
func (s *checkpointWikiPages) CreatePage(_ context.Context, page *types.WikiPage) (*types.WikiPage, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
out := *page
|
|
s.pages[page.Slug] = &out
|
|
s.writes++
|
|
return page, nil
|
|
}
|
|
func (s *checkpointWikiPages) UpdatePage(ctx context.Context, page *types.WikiPage) (*types.WikiPage, error) {
|
|
return s.CreatePage(ctx, page)
|
|
}
|
|
func (*checkpointWikiPages) ListSlugsBySourceRef(context.Context, string, string) ([]string, error) {
|
|
return nil, nil
|
|
}
|
|
func (*checkpointWikiPages) ListDistinctCategoryPaths(context.Context, string, int) ([][]string, error) {
|
|
return nil, nil
|
|
}
|
|
func (s *checkpointWikiPages) UpdatePageMeta(ctx context.Context, page *types.WikiPage) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if s.publishErr != nil {
|
|
return s.publishErr
|
|
}
|
|
out := *page
|
|
s.pages[page.Slug] = &out
|
|
return nil
|
|
}
|
|
|
|
func TestWikiCheckpointPageSliceResumesAndRejectsChangedPage(t *testing.T) {
|
|
db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
|
|
chunks := checkpointChunks(1)
|
|
pages := &checkpointWikiPages{pages: make(map[string]*types.WikiPage)}
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "v1"}
|
|
updates := make([]SlugUpdate, wikiPageSlice+1)
|
|
for i := range updates {
|
|
updates[i] = SlugUpdate{KnowledgeID: "doc", SourceRef: "doc", Slug: fmt.Sprintf("summary/%d", i), Type: types.WikiPageTypeSummary, DocTitle: "title", SummaryBody: "complete summary"}
|
|
}
|
|
svc := &wikiIngestService{checkpointDB: db, wikiService: pages, knowledgeSvc: &checkpointKnowledge{}, chunkRepo: &checkpointChunksRepo{chunks: chunks}}
|
|
run := func() int {
|
|
batch := &WikiBatchContext{}
|
|
batch.checkpointScopes.Store("doc", scope)
|
|
batch.checkpointSourceHashes.Store("doc", wikiEvidenceSourceHash(chunks))
|
|
yielded := 0
|
|
for _, update := range updates {
|
|
_, _, _, err := svc.reduceSlugUpdates(context.Background(), nil, "kb", update.Slug, []SlugUpdate{update}, 7, batch, nil)
|
|
if errors.Is(err, errWikiCheckpointYield) {
|
|
yielded++
|
|
} else {
|
|
require.NoError(t, err)
|
|
}
|
|
}
|
|
return yielded
|
|
}
|
|
require.Equal(t, 1, run())
|
|
require.Equal(t, wikiPageSlice, pages.writes)
|
|
svc = &wikiIngestService{checkpointDB: db, wikiService: pages, knowledgeSvc: &checkpointKnowledge{}, chunkRepo: &checkpointChunksRepo{chunks: chunks}}
|
|
require.Zero(t, run())
|
|
require.Equal(t, wikiPageSlice+1, pages.writes)
|
|
pages.pages[updates[0].Slug].Content = "changed source contribution"
|
|
require.Zero(t, run())
|
|
require.Equal(t, wikiPageSlice+2, pages.writes)
|
|
require.Equal(t, "complete summary", pages.pages[updates[0].Slug].Content)
|
|
pages.pages[updates[0].Slug].SourceRefs = nil
|
|
require.Zero(t, run())
|
|
require.Equal(t, wikiPageSlice+3, pages.writes)
|
|
require.Contains(t, pages.pages[updates[0].Slug].SourceRefs, "doc")
|
|
}
|
|
|
|
func TestWikiCheckpointYieldBudgetDoesNotHideRealFailures(t *testing.T) {
|
|
ctx, cancel := context.WithCancelCause(context.Background())
|
|
cancel(errWikiCheckpointYield)
|
|
require.ErrorIs(t, wikiCheckpointError(ctx, fmt.Errorf("provider: %w", context.Canceled)), errWikiCheckpointYield)
|
|
failure := errors.New("invalid API key")
|
|
require.ErrorIs(t, wikiCheckpointError(ctx, failure), failure)
|
|
require.ErrorIs(t, wikiCheckpointMapError(ctx, context.Canceled, failure), failure)
|
|
require.ErrorIs(t, wikiCheckpointMapError(ctx, failure, errWikiCheckpointYield), failure)
|
|
require.ErrorIs(t, wikiCheckpointMapError(ctx, nil, errWikiCheckpointYield), errWikiCheckpointYield)
|
|
ordinary, cancelOrdinary := context.WithCancel(context.Background())
|
|
cancelOrdinary()
|
|
require.ErrorIs(t, wikiCheckpointError(ordinary, context.Canceled), context.Canceled)
|
|
}
|
|
|
|
type checkpointChunksRepo struct {
|
|
interfaces.ChunkRepository
|
|
chunks []*types.Chunk
|
|
}
|
|
|
|
func (r *checkpointChunksRepo) ListChunksByKnowledgeID(context.Context, uint64, string) ([]*types.Chunk, error) {
|
|
return r.chunks, nil
|
|
}
|
|
func (*checkpointChunksRepo) ListChunksByParentIDs(context.Context, uint64, []string) ([]*types.Chunk, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func (r *checkpointChunksRepo) ListChunksByID(_ context.Context, tenantID uint64, ids []string) ([]*types.Chunk, error) {
|
|
wanted := make(map[string]bool, len(ids))
|
|
for _, id := range ids {
|
|
wanted[id] = true
|
|
}
|
|
var chunks []*types.Chunk
|
|
for _, chunk := range r.chunks {
|
|
if chunk != nil && chunk.TenantID == tenantID && wanted[chunk.ID] {
|
|
chunks = append(chunks, chunk)
|
|
}
|
|
}
|
|
return chunks, nil
|
|
}
|
|
|
|
type checkpointKBService struct {
|
|
interfaces.KnowledgeBaseService
|
|
kb *types.KnowledgeBase
|
|
}
|
|
|
|
func (s *checkpointKBService) GetKnowledgeBaseByIDOnly(context.Context, string) (*types.KnowledgeBase, error) {
|
|
return s.kb, nil
|
|
}
|
|
func (s *checkpointKBService) GetKnowledgeBaseByID(context.Context, string) (*types.KnowledgeBase, error) {
|
|
return s.kb, nil
|
|
}
|
|
|
|
type checkpointPending struct {
|
|
interfaces.TaskPendingOpsRepository
|
|
rows []*types.TaskPendingOp
|
|
released []int64
|
|
deleted []int64
|
|
failures int
|
|
}
|
|
|
|
func (r *checkpointPending) Enqueue(context.Context, *types.TaskPendingOp) error { return nil }
|
|
|
|
func (r *checkpointPending) PeekBatch(context.Context, string, string, string, int) ([]*types.TaskPendingOp, error) {
|
|
return r.rows, nil
|
|
}
|
|
func (r *checkpointPending) DeleteByIDs(_ context.Context, ids []int64) error {
|
|
r.deleted = append(r.deleted, ids...)
|
|
remaining := r.rows[:0]
|
|
for _, row := range r.rows {
|
|
removed := false
|
|
for _, id := range ids {
|
|
removed = removed || row.ID == id
|
|
}
|
|
if !removed {
|
|
remaining = append(remaining, row)
|
|
}
|
|
}
|
|
r.rows = remaining
|
|
return nil
|
|
}
|
|
func (r *checkpointPending) ReleaseByIDs(_ context.Context, ids []int64) error {
|
|
r.released = append(r.released, ids...)
|
|
return nil
|
|
}
|
|
func (r *checkpointPending) IncrFailCount(context.Context, int64) (int, error) {
|
|
r.failures++
|
|
return r.failures, nil
|
|
}
|
|
func (r *checkpointPending) PendingCount(context.Context, string, string, string) (int64, error) {
|
|
return int64(len(r.rows)), nil
|
|
}
|
|
|
|
type checkpointFinalization struct {
|
|
interfaces.KnowledgeRepository
|
|
calls int
|
|
}
|
|
|
|
func (r *checkpointFinalization) FinalizeSubtask(context.Context, string) (int, bool, error) {
|
|
r.calls++
|
|
return 0, true, nil
|
|
}
|
|
|
|
func TestWikiCheckpointYieldReleasesClaimWithoutFailureOrCompletion(t *testing.T) {
|
|
db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
|
|
chunks := checkpointChunks(wikiCitationSlice + 1)
|
|
op := WikiPendingOp{Op: WikiOpIngest, KnowledgeID: "doc", DocTitle: "title", Language: "en-US", dbID: 42}
|
|
opBytes, err := json.Marshal(op)
|
|
require.NoError(t, err)
|
|
pending := &checkpointPending{rows: []*types.TaskPendingOp{{ID: 42, Op: WikiOpIngest, TenantID: 7, Payload: opBytes}}}
|
|
queue := &wikiGuardTaskQueue{}
|
|
finalization := &checkpointFinalization{}
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
|
return &types.ChatResponse{Content: `{"citations":{"entity/alpha":["c000"]}}`}, nil
|
|
}}
|
|
modelSvc := &checkpointModelService{model: &types.Model{ID: t.Name(), Name: "test", TenantID: 7}, chat: model}
|
|
kb := &types.KnowledgeBase{ID: "kb", TenantID: 7, SummaryModelID: t.Name(), WikiConfig: &types.WikiConfig{}}
|
|
kb.IndexingStrategy.WikiEnabled = true
|
|
svc := &wikiIngestService{checkpointDB: db, modelService: modelSvc, kbService: &checkpointKBService{kb: kb}, pendingRepo: pending,
|
|
task: queue, knowledgeSvc: &checkpointKnowledge{}, knowledgeRepo: finalization, chunkRepo: &checkpointChunksRepo{chunks: chunks}, wikiService: &checkpointWikiPages{pages: make(map[string]*types.WikiPage)}}
|
|
payload := WikiIngestPayload{TenantID: 7, KnowledgeBaseID: "kb", Language: "en-US"}
|
|
content := reconstructContent(chunks)
|
|
scope, err := svc.newWikiCheckpointScope(context.Background(), model, payload, op, chunks, content, "title", types.ResolveLanguageName(context.Background(), op.Language), svc.newWikiBatchContext("kb", kb.WikiConfig))
|
|
require.NoError(t, err)
|
|
require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "candidates", "complete", combinedExtraction{Entities: []extractedItem{{Name: "Alpha", Slug: "entity/alpha"}}}))
|
|
require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "summary", wikiCheckpointHash("- [[entity/alpha]] = Alpha\n"), "Summary\n\nComplete document summary"))
|
|
body, err := json.Marshal(payload)
|
|
require.NoError(t, err)
|
|
require.NoError(t, svc.ProcessWikiIngest(context.Background(), asynq.NewTask(types.TypeWikiIngest, body)))
|
|
require.Equal(t, []int64{42}, pending.released)
|
|
require.Empty(t, pending.deleted)
|
|
require.Zero(t, pending.failures)
|
|
require.Zero(t, finalization.calls)
|
|
require.Len(t, queue.tasks, 1)
|
|
require.Equal(t, types.TypeWikiIngest, queue.tasks[0].Type())
|
|
}
|
|
|
|
func checkpointMap(pageCount int) wikiMapCheckpoint {
|
|
saved := wikiMapCheckpoint{Result: docIngestResult{KnowledgeID: "doc", DocTitle: "title", Summary: "Complete summary"}}
|
|
for i := 0; i < pageCount; i++ {
|
|
slug := fmt.Sprintf("concept/item-%d", i)
|
|
update := SlugUpdate{KnowledgeID: "doc", SourceRef: "doc", Slug: slug, Type: types.WikiPageTypeConcept,
|
|
DocTitle: "title", Item: extractedItem{Slug: slug, Name: fmt.Sprintf("Item %d", i), Details: "Grounded source details"}}
|
|
if i == 0 {
|
|
update.Slug = "summary/doc"
|
|
update.Type = types.WikiPageTypeSummary
|
|
update.SummaryBody = "Complete document summary"
|
|
} else {
|
|
update.SourceChunks = []string{"chunk-000"}
|
|
update.Item.SourceChunks = []string{"chunk-000"}
|
|
}
|
|
saved.Updates = append(saved.Updates, update)
|
|
saved.Result.Pages = append(saved.Result.Pages, wikiIngestPageRef{Slug: update.Slug, Title: "title"})
|
|
}
|
|
return saved
|
|
}
|
|
|
|
func TestWikiCheckpointMapRejectsIncompleteOrForeignArtifacts(t *testing.T) {
|
|
scope := &wikiCheckpointScope{7, "kb", "doc", "v1"}
|
|
chunks := checkpointChunks(1)
|
|
require.NoError(t, validateWikiMapCheckpoint(checkpointMap(2), scope, chunks))
|
|
for _, mutate := range []func(*wikiMapCheckpoint){
|
|
func(m *wikiMapCheckpoint) { m.Result.KnowledgeID = "other" },
|
|
func(m *wikiMapCheckpoint) { m.Updates[1].KnowledgeID = "other" },
|
|
func(m *wikiMapCheckpoint) { m.Updates[1].SourceRef = "other" },
|
|
func(m *wikiMapCheckpoint) { m.Updates[1].SourceChunks = []string{"foreign-chunk"} },
|
|
func(m *wikiMapCheckpoint) { m.Updates[1].Item.SourceChunks = []string{"foreign-chunk"} },
|
|
func(m *wikiMapCheckpoint) { m.Updates[1].Item.Slug = "entity/other" },
|
|
func(m *wikiMapCheckpoint) { m.Result.Pages = m.Result.Pages[:1] },
|
|
func(m *wikiMapCheckpoint) { m.Result.Pages = append(m.Result.Pages, m.Result.Pages[0]) },
|
|
func(m *wikiMapCheckpoint) { m.Updates = m.Updates[1:]; m.Result.Pages = m.Result.Pages[1:] },
|
|
func(m *wikiMapCheckpoint) { m.Updates[0].SummaryBody = "" },
|
|
} {
|
|
saved := checkpointMap(2)
|
|
mutate(&saved)
|
|
require.Error(t, validateWikiMapCheckpoint(saved, scope, chunks))
|
|
}
|
|
}
|
|
|
|
func TestWikiCheckpointWholeMapResumesAfterDatabaseReopen(t *testing.T) {
|
|
path := filepath.Join(t.TempDir(), "checkpoints.sqlite")
|
|
db := openWikiCheckpointDB(t, path)
|
|
chunks := checkpointChunks(1)
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
|
t.Error("completed map must not call the model again")
|
|
return nil, errors.New("unexpected model call")
|
|
}}
|
|
modelSvc := &checkpointModelService{model: &types.Model{ID: t.Name(), Name: "test", TenantID: 7}, chat: model}
|
|
svc := &wikiIngestService{checkpointDB: db, modelService: modelSvc, knowledgeSvc: &checkpointKnowledge{}, chunkRepo: &checkpointChunksRepo{chunks: chunks}}
|
|
op := WikiPendingOp{Op: WikiOpIngest, KnowledgeID: "doc", Language: "en-US", dbID: 42}
|
|
payload := WikiIngestPayload{TenantID: 7, KnowledgeBaseID: "kb"}
|
|
batch := svc.newWikiBatchContext("kb", &types.WikiConfig{})
|
|
scope, err := svc.newWikiCheckpointScope(context.Background(), model, payload, op, chunks, reconstructContent(chunks), "title", types.ResolveLanguageName(context.Background(), op.Language), batch)
|
|
require.NoError(t, err)
|
|
saved := checkpointMap(2)
|
|
require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "map", "complete", saved))
|
|
sqlDB, _ := db.DB()
|
|
require.NoError(t, sqlDB.Close())
|
|
svc.checkpointDB = openWikiCheckpointDB(t, path)
|
|
result, updates, err := svc.mapOneDocument(context.Background(), model, payload, op, batch)
|
|
require.NoError(t, err)
|
|
require.Equal(t, saved.Result.Pages, result.Pages)
|
|
require.Equal(t, saved.Updates, updates)
|
|
}
|
|
|
|
func TestWikiCheckpointPageYieldKeepsClaimAndFinalizationSlot(t *testing.T) {
|
|
path := filepath.Join(t.TempDir(), "checkpoints.sqlite")
|
|
db := openWikiCheckpointDB(t, path)
|
|
chunks := checkpointChunks(1)
|
|
op := WikiPendingOp{Op: WikiOpIngest, KnowledgeID: "doc", DocTitle: "title", Language: "en-US", dbID: 42}
|
|
opBytes, err := json.Marshal(op)
|
|
require.NoError(t, err)
|
|
pending := &checkpointPending{rows: []*types.TaskPendingOp{{ID: 42, Op: WikiOpIngest, TenantID: 7, Payload: opBytes}}}
|
|
queue := &wikiGuardTaskQueue{}
|
|
finalization := &checkpointFinalization{}
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
|
return &types.ChatResponse{Content: "SUMMARY: Complete page\n\nGrounded page content"}, nil
|
|
}}
|
|
modelSvc := &checkpointModelService{model: &types.Model{ID: t.Name(), Name: "test", TenantID: 7}, chat: model}
|
|
kb := &types.KnowledgeBase{ID: "kb", TenantID: 7, SummaryModelID: t.Name(), WikiConfig: &types.WikiConfig{}}
|
|
kb.IndexingStrategy.WikiEnabled = true
|
|
pages := &checkpointWikiPages{pages: make(map[string]*types.WikiPage)}
|
|
svc := &wikiIngestService{checkpointDB: db, modelService: modelSvc, kbService: &checkpointKBService{kb: kb}, pendingRepo: pending,
|
|
task: queue, knowledgeSvc: &checkpointKnowledge{}, knowledgeRepo: finalization, chunkRepo: &checkpointChunksRepo{chunks: chunks}, wikiService: pages}
|
|
payload := WikiIngestPayload{TenantID: 7, KnowledgeBaseID: "kb", Language: "en-US"}
|
|
scope, err := svc.newWikiCheckpointScope(context.Background(), model, payload, op, chunks, reconstructContent(chunks), "title", types.ResolveLanguageName(context.Background(), op.Language), svc.newWikiBatchContext("kb", kb.WikiConfig))
|
|
require.NoError(t, err)
|
|
require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "map", "complete", checkpointMap(wikiPageSlice+1)))
|
|
body, err := json.Marshal(payload)
|
|
require.NoError(t, err)
|
|
require.NoError(t, svc.ProcessWikiIngest(context.Background(), asynq.NewTask(types.TypeWikiIngest, body)))
|
|
require.Equal(t, wikiPageSlice, pages.writes)
|
|
require.Equal(t, []int64{42}, pending.released)
|
|
require.Empty(t, pending.deleted)
|
|
require.Zero(t, pending.failures)
|
|
require.Zero(t, finalization.calls)
|
|
require.Len(t, queue.tasks, 1)
|
|
require.Equal(t, types.TypeWikiIngest, queue.tasks[0].Type())
|
|
// Simulate a worker restart after the first bounded page slice. The second
|
|
// run must finish only the missing page and release the finalizing slot once.
|
|
sqlDB, _ := db.DB()
|
|
require.NoError(t, sqlDB.Close())
|
|
svc.checkpointDB = openWikiCheckpointDB(t, path)
|
|
require.NoError(t, svc.ProcessWikiIngest(context.Background(), asynq.NewTask(types.TypeWikiIngest, body)))
|
|
require.Equal(t, wikiPageSlice+1, pages.writes)
|
|
require.Equal(t, []int64{42}, pending.deleted)
|
|
require.Zero(t, pending.failures)
|
|
require.Equal(t, 1, finalization.calls)
|
|
for _, page := range pages.pages {
|
|
require.Equal(t, types.WikiPageStatusPublished, page.Status)
|
|
}
|
|
}
|
|
|
|
type checkpointFailingKnowledge struct {
|
|
interfaces.KnowledgeService
|
|
err error
|
|
}
|
|
|
|
func (s *checkpointFailingKnowledge) GetKnowledgeByIDOnly(context.Context, string) (*types.Knowledge, error) {
|
|
return nil, s.err
|
|
}
|
|
|
|
func TestWikiCheckpointTransientSourceFailureCannotCompleteMapOrPage(t *testing.T) {
|
|
want := errors.New("database temporarily unavailable")
|
|
svc := &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite")), knowledgeSvc: &checkpointFailingKnowledge{err: want}}
|
|
_, _, err := svc.mapOneDocument(context.Background(), nil, WikiIngestPayload{TenantID: 7, KnowledgeBaseID: "kb"}, WikiPendingOp{KnowledgeID: "doc"}, &WikiBatchContext{})
|
|
require.ErrorIs(t, err, want)
|
|
batch := &WikiBatchContext{}
|
|
batch.checkpointScopes.Store("doc", &wikiCheckpointScope{7, "kb", "doc", "v1"})
|
|
_, _, _, err = svc.reduceSlugUpdates(context.Background(), nil, "kb", "summary/doc", checkpointMap(1).Updates, 7, batch, nil)
|
|
require.ErrorIs(t, err, want)
|
|
var count int64
|
|
require.NoError(t, svc.checkpointDB.Model(&wikiIngestCheckpoint{}).Count(&count).Error)
|
|
require.Zero(t, count)
|
|
}
|
|
|
|
func TestWikiCheckpointPublishFailureKeepsPendingOperation(t *testing.T) {
|
|
db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
|
|
chunks := checkpointChunks(1)
|
|
op := WikiPendingOp{Op: WikiOpIngest, KnowledgeID: "doc", DocTitle: "title", Language: "en-US", dbID: 42}
|
|
opBytes, err := json.Marshal(op)
|
|
require.NoError(t, err)
|
|
pending := &checkpointPending{rows: []*types.TaskPendingOp{{ID: 42, Op: WikiOpIngest, TenantID: 7, Payload: opBytes}}}
|
|
queue := &wikiGuardTaskQueue{}
|
|
finalization := &checkpointFinalization{}
|
|
model := &wikiScriptedChat{id: t.Name()}
|
|
modelSvc := &checkpointModelService{model: &types.Model{ID: t.Name(), Name: "test", TenantID: 7}, chat: model}
|
|
kb := &types.KnowledgeBase{ID: "kb", TenantID: 7, SummaryModelID: t.Name(), WikiConfig: &types.WikiConfig{}}
|
|
kb.IndexingStrategy.WikiEnabled = true
|
|
pages := &checkpointWikiPages{pages: make(map[string]*types.WikiPage), publishErr: errors.New("publish transaction failed")}
|
|
svc := &wikiIngestService{checkpointDB: db, modelService: modelSvc, kbService: &checkpointKBService{kb: kb}, pendingRepo: pending,
|
|
task: queue, knowledgeSvc: &checkpointKnowledge{}, knowledgeRepo: finalization, chunkRepo: &checkpointChunksRepo{chunks: chunks}, wikiService: pages}
|
|
payload := WikiIngestPayload{TenantID: 7, KnowledgeBaseID: "kb", Language: "en-US"}
|
|
scope, err := svc.newWikiCheckpointScope(context.Background(), model, payload, op, chunks, reconstructContent(chunks), "title", types.ResolveLanguageName(context.Background(), op.Language), svc.newWikiBatchContext("kb", kb.WikiConfig))
|
|
require.NoError(t, err)
|
|
require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "map", "complete", checkpointMap(1)))
|
|
body, err := json.Marshal(payload)
|
|
require.NoError(t, err)
|
|
require.NoError(t, svc.ProcessWikiIngest(context.Background(), asynq.NewTask(types.TypeWikiIngest, body)))
|
|
require.Empty(t, pending.deleted)
|
|
require.Equal(t, []int64{42}, pending.released)
|
|
require.Equal(t, 1, pending.failures)
|
|
require.Zero(t, finalization.calls)
|
|
require.Equal(t, types.WikiPageStatusDraft, pages.pages["summary/doc"].Status)
|
|
pages.publishErr = nil
|
|
require.NoError(t, svc.ProcessWikiIngest(context.Background(), asynq.NewTask(types.TypeWikiIngest, body)))
|
|
require.Equal(t, 1, pages.writes, "publication retry must reuse the already written page")
|
|
require.Equal(t, []int64{42}, pending.deleted)
|
|
require.Equal(t, 1, finalization.calls)
|
|
require.Equal(t, types.WikiPageStatusPublished, pages.pages["summary/doc"].Status)
|
|
}
|
|
|
|
type checkpointDeadLetter struct {
|
|
interfaces.TaskDeadLetterRepository
|
|
err error
|
|
}
|
|
|
|
func (r *checkpointDeadLetter) Insert(context.Context, *types.TaskDeadLetter) error { return r.err }
|
|
|
|
func TestWikiCheckpointDeadLetterFailurePreservesPendingAndFinalization(t *testing.T) {
|
|
want := errors.New("archive transaction failed")
|
|
pending := &checkpointPending{failures: wikiMaxFailRetries}
|
|
finalization := &checkpointFinalization{}
|
|
svc := &wikiIngestService{pendingRepo: pending, deadLetterRepo: &checkpointDeadLetter{err: want}, knowledgeRepo: finalization}
|
|
err := svc.requeueFailedOps(context.Background(), WikiIngestPayload{TenantID: 7, KnowledgeBaseID: "kb"}, []WikiPendingOp{{Op: WikiOpIngest, KnowledgeID: "doc", dbID: 42}})
|
|
require.ErrorIs(t, err, want)
|
|
require.Empty(t, pending.deleted)
|
|
require.Zero(t, finalization.calls)
|
|
}
|