311 lines
15 KiB
Go
311 lines
15 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"path/filepath"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/models/chat"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
"gorm.io/gorm"
|
||
|
|
)
|
||
|
|
|
||
|
|
type wikiEvidenceChunksRepo struct {
|
||
|
|
interfaces.ChunkRepository
|
||
|
|
chunks []*types.Chunk
|
||
|
|
err error
|
||
|
|
queries [][]string
|
||
|
|
tenantIDs []uint64
|
||
|
|
returnAll bool
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *wikiEvidenceChunksRepo) ListChunksByID(_ context.Context, tenantID uint64, ids []string) ([]*types.Chunk, error) {
|
||
|
|
r.queries = append(r.queries, append([]string(nil), ids...))
|
||
|
|
r.tenantIDs = append(r.tenantIDs, tenantID)
|
||
|
|
if r.err != nil {
|
||
|
|
return nil, r.err
|
||
|
|
}
|
||
|
|
if r.returnAll {
|
||
|
|
return r.chunks, nil
|
||
|
|
}
|
||
|
|
wanted := make(map[string]bool)
|
||
|
|
for _, id := range ids {
|
||
|
|
wanted[id] = true
|
||
|
|
}
|
||
|
|
var out []*types.Chunk
|
||
|
|
for _, chunk := range r.chunks {
|
||
|
|
if chunk != nil && wanted[chunk.ID] {
|
||
|
|
out = append(out, chunk)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return out, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func wikiEvidenceChunk() *types.Chunk {
|
||
|
|
return &types.Chunk{ID: "source-chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", IsEnabled: true,
|
||
|
|
Content: "Verbatim source: 42 applies only when condition A holds.", ChunkIndex: 3, StartAt: 80}
|
||
|
|
}
|
||
|
|
|
||
|
|
func wikiEvidenceAddition() SlugUpdate {
|
||
|
|
return SlugUpdate{KnowledgeID: "doc", SourceRef: "doc", Slug: "concept/alpha", Type: types.WikiPageTypeConcept,
|
||
|
|
DocTitle: "Source", SourceChunks: []string{"source-chunk"}, Item: extractedItem{Name: "Alpha", Slug: "concept/alpha", Description: "UNSUPPORTED_DESCRIPTION_MUST_NOT_REPLACE_SOURCE", Details: "UNSUPPORTED_DETAILS_MUST_NOT_REPLACE_SOURCE"}}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiAdditionEvidenceRejectsIncompleteOrForeignSources(t *testing.T) {
|
||
|
|
tests := []struct {
|
||
|
|
name string
|
||
|
|
change func(*SlugUpdate, *wikiEvidenceChunksRepo)
|
||
|
|
}{
|
||
|
|
{"no citations", func(u *SlugUpdate, r *wikiEvidenceChunksRepo) { u.SourceChunks = nil }},
|
||
|
|
{"blank citation", func(u *SlugUpdate, r *wikiEvidenceChunksRepo) { u.SourceChunks = []string{"source-chunk", " "} }},
|
||
|
|
{"missing document", func(u *SlugUpdate, r *wikiEvidenceChunksRepo) { u.KnowledgeID = "" }},
|
||
|
|
{"wrong source ref", func(u *SlugUpdate, r *wikiEvidenceChunksRepo) { u.SourceRef = "other-document" }},
|
||
|
|
{"partially missing", func(u *SlugUpdate, r *wikiEvidenceChunksRepo) {
|
||
|
|
u.SourceChunks = append(u.SourceChunks, "missing-chunk")
|
||
|
|
}},
|
||
|
|
{"wrong tenant", func(u *SlugUpdate, r *wikiEvidenceChunksRepo) { r.chunks[0].TenantID = 8 }},
|
||
|
|
{"wrong knowledge base", func(u *SlugUpdate, r *wikiEvidenceChunksRepo) { r.chunks[0].KnowledgeBaseID = "other-kb" }},
|
||
|
|
{"wrong document", func(u *SlugUpdate, r *wikiEvidenceChunksRepo) { r.chunks[0].KnowledgeID = "other-document" }},
|
||
|
|
{"disabled", func(u *SlugUpdate, r *wikiEvidenceChunksRepo) { r.chunks[0].IsEnabled = false }},
|
||
|
|
{"deleted", func(u *SlugUpdate, r *wikiEvidenceChunksRepo) { r.chunks[0].DeletedAt = gorm.DeletedAt{Valid: true} }},
|
||
|
|
{"empty source", func(u *SlugUpdate, r *wikiEvidenceChunksRepo) { r.chunks[0].Content = " \n" }},
|
||
|
|
{"lookup failure", func(u *SlugUpdate, r *wikiEvidenceChunksRepo) { r.err = errors.New("database unavailable") }},
|
||
|
|
{"unexpected source", func(u *SlugUpdate, r *wikiEvidenceChunksRepo) {
|
||
|
|
r.returnAll = true
|
||
|
|
c := wikiEvidenceChunk()
|
||
|
|
c.ID = "unrequested"
|
||
|
|
r.chunks = append(r.chunks, c)
|
||
|
|
}},
|
||
|
|
{"nil source", func(u *SlugUpdate, r *wikiEvidenceChunksRepo) { r.returnAll = true; r.chunks = append(r.chunks, nil) }},
|
||
|
|
{"duplicate source", func(u *SlugUpdate, r *wikiEvidenceChunksRepo) { r.chunks = append(r.chunks, wikiEvidenceChunk()) }},
|
||
|
|
}
|
||
|
|
for _, test := range tests {
|
||
|
|
t.Run(test.name, func(t *testing.T) {
|
||
|
|
update := wikiEvidenceAddition()
|
||
|
|
repo := &wikiEvidenceChunksRepo{chunks: []*types.Chunk{wikiEvidenceChunk()}}
|
||
|
|
test.change(&update, repo)
|
||
|
|
got, err := (&wikiIngestService{chunkRepo: repo}).resolveWikiAdditionEvidence(context.Background(), 7, "kb", []SlugUpdate{update})
|
||
|
|
require.ErrorIs(t, err, errWikiCitationEvidence)
|
||
|
|
require.Nil(t, got, "never expose partially verified evidence")
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiAdditionEvidenceChecksEveryUniqueSourceInBoundedBatches(t *testing.T) {
|
||
|
|
repo := &wikiEvidenceChunksRepo{}
|
||
|
|
update := wikiEvidenceAddition()
|
||
|
|
update.SourceRef = "doc|Legacy title"
|
||
|
|
update.SourceChunks = nil
|
||
|
|
for i := 0; i < 260; i++ {
|
||
|
|
c := wikiEvidenceChunk()
|
||
|
|
c.ID = fmt.Sprintf("source-%03d", i)
|
||
|
|
repo.chunks = append(repo.chunks, c)
|
||
|
|
update.SourceChunks = append(update.SourceChunks, c.ID)
|
||
|
|
}
|
||
|
|
update.SourceChunks = append(update.SourceChunks, update.SourceChunks[0])
|
||
|
|
got, err := (&wikiIngestService{chunkRepo: repo}).resolveWikiAdditionEvidence(context.Background(), 7, "kb", []SlugUpdate{update, update})
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Len(t, got, 260)
|
||
|
|
require.Len(t, repo.queries, 3)
|
||
|
|
require.Equal(t, []uint64{7, 7, 7}, repo.tenantIDs)
|
||
|
|
for _, query := range repo.queries {
|
||
|
|
require.LessOrEqual(t, len(query), 128)
|
||
|
|
}
|
||
|
|
require.Equal(t, repo.chunks[0], got[repo.chunks[0].ID], "source metadata and exact text must be preserved for ordered page synthesis")
|
||
|
|
require.NotSame(t, repo.chunks[0], got[repo.chunks[0].ID], "keep a snapshot while the model runs")
|
||
|
|
other := update
|
||
|
|
other.KnowledgeID, other.SourceRef = "other-document", "other-document"
|
||
|
|
_, err = (&wikiIngestService{chunkRepo: repo}).resolveWikiAdditionEvidence(context.Background(), 7, "kb", []SlugUpdate{update, other})
|
||
|
|
require.ErrorIs(t, err, errWikiCitationEvidence)
|
||
|
|
require.Len(t, repo.queries, 3, "reject contradictory ownership before any lookup")
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiEvidenceSaveGuardRejectsConcurrentSourceChanges(t *testing.T) {
|
||
|
|
for _, test := range []struct {
|
||
|
|
name string
|
||
|
|
change func(*wikiEvidenceChunksRepo)
|
||
|
|
}{
|
||
|
|
{"content", func(r *wikiEvidenceChunksRepo) { r.chunks[0].Content = "Changed source" }},
|
||
|
|
{"revision", func(r *wikiEvidenceChunksRepo) { r.chunks[0].ContentRevision++ }},
|
||
|
|
{"enabled", func(r *wikiEvidenceChunksRepo) { r.chunks[0].IsEnabled = false }},
|
||
|
|
{"deleted", func(r *wikiEvidenceChunksRepo) { r.chunks = nil }},
|
||
|
|
{"scope", func(r *wikiEvidenceChunksRepo) { r.chunks[0].TenantID = 8 }},
|
||
|
|
{"ordering", func(r *wikiEvidenceChunksRepo) { r.chunks[0].ChunkIndex++ }},
|
||
|
|
{"read failure", func(r *wikiEvidenceChunksRepo) { r.err = errors.New("database unavailable") }},
|
||
|
|
} {
|
||
|
|
t.Run(test.name, func(t *testing.T) {
|
||
|
|
repo := &wikiEvidenceChunksRepo{chunks: []*types.Chunk{wikiEvidenceChunk()}}
|
||
|
|
svc := &wikiIngestService{chunkRepo: repo, knowledgeSvc: &checkpointKnowledge{}}
|
||
|
|
updates := []SlugUpdate{wikiEvidenceAddition()}
|
||
|
|
before, err := svc.resolveWikiAdditionEvidence(context.Background(), 7, "kb", updates)
|
||
|
|
require.NoError(t, err)
|
||
|
|
test.change(repo)
|
||
|
|
err = svc.verifyWikiAdditionEvidenceUnchanged(context.Background(), 7, "kb", updates, before)
|
||
|
|
require.ErrorIs(t, err, errWikiCitationEvidence)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiEvidenceSaveGuardAllowsOnlyUnchangedCompleteSources(t *testing.T) {
|
||
|
|
repo := &wikiEvidenceChunksRepo{chunks: []*types.Chunk{wikiEvidenceChunk()}}
|
||
|
|
svc := &wikiIngestService{chunkRepo: repo, knowledgeSvc: &checkpointKnowledge{}}
|
||
|
|
updates := []SlugUpdate{wikiEvidenceAddition()}
|
||
|
|
before, err := svc.resolveWikiAdditionEvidence(context.Background(), 7, "kb", updates)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.NoError(t, svc.verifyWikiAdditionEvidenceUnchanged(context.Background(), 7, "kb", updates, before))
|
||
|
|
require.ErrorIs(t, svc.verifyWikiAdditionEvidenceUnchanged(context.Background(), 7, "kb", updates, nil), errWikiCitationEvidence)
|
||
|
|
svc.knowledgeSvc = &checkpointFailingKnowledge{err: gorm.ErrRecordNotFound}
|
||
|
|
require.ErrorIs(t, svc.verifyWikiAdditionEvidenceUnchanged(context.Background(), 7, "kb", updates, before), errWikiCitationEvidence)
|
||
|
|
require.NoError(t, svc.verifyWikiAdditionEvidenceUnchanged(context.Background(), 7, "kb", nil, nil))
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiSummaryEvidenceSaveGuardChecksWholeScopedSourceSnapshot(t *testing.T) {
|
||
|
|
for _, test := range []struct {
|
||
|
|
name string
|
||
|
|
change func(*checkpointChunksRepo, *WikiBatchContext)
|
||
|
|
}{
|
||
|
|
{"content", func(r *checkpointChunksRepo, b *WikiBatchContext) { r.chunks[0].Content = "Changed source" }},
|
||
|
|
{"new chunk", func(r *checkpointChunksRepo, b *WikiBatchContext) {
|
||
|
|
c := wikiEvidenceChunk()
|
||
|
|
c.ID = "new"
|
||
|
|
r.chunks = append(r.chunks, c)
|
||
|
|
}},
|
||
|
|
{"deleted chunk", func(r *checkpointChunksRepo, b *WikiBatchContext) { r.chunks = nil }},
|
||
|
|
{"foreign chunk", func(r *checkpointChunksRepo, b *WikiBatchContext) { r.chunks[0].KnowledgeID = "foreign" }},
|
||
|
|
{"missing baseline", func(r *checkpointChunksRepo, b *WikiBatchContext) { b.checkpointSourceHashes.Delete("doc") }},
|
||
|
|
{"wrong scope", func(r *checkpointChunksRepo, b *WikiBatchContext) {
|
||
|
|
b.checkpointScopes.Store("doc", &wikiCheckpointScope{8, "kb", "doc", "fingerprint"})
|
||
|
|
}},
|
||
|
|
} {
|
||
|
|
t.Run(test.name, func(t *testing.T) {
|
||
|
|
repo := &checkpointChunksRepo{chunks: []*types.Chunk{wikiEvidenceChunk()}}
|
||
|
|
svc := &wikiIngestService{chunkRepo: repo, knowledgeSvc: &checkpointKnowledge{}}
|
||
|
|
batch := &WikiBatchContext{}
|
||
|
|
batch.checkpointScopes.Store("doc", &wikiCheckpointScope{7, "kb", "doc", "fingerprintprint"})
|
||
|
|
batch.checkpointSourceHashes.Store("doc", wikiEvidenceSourceHash(repo.chunks))
|
||
|
|
update := SlugUpdate{KnowledgeID: "doc", Type: types.WikiPageTypeSummary, SourceRef: "doc"}
|
||
|
|
require.NoError(t, svc.verifyWikiSummaryEvidenceUnchanged(context.Background(), 7, "kb", update, batch))
|
||
|
|
test.change(repo, batch)
|
||
|
|
require.ErrorIs(t, svc.verifyWikiSummaryEvidenceUnchanged(context.Background(), 7, "kb", update, batch), errWikiCitationEvidence)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiSourceEvidenceHashIgnoresOnlyResultOrder(t *testing.T) {
|
||
|
|
a, b := wikiEvidenceChunk(), wikiEvidenceChunk()
|
||
|
|
b.ID = "other"
|
||
|
|
require.Equal(t, wikiEvidenceSourceHash([]*types.Chunk{a, b}), wikiEvidenceSourceHash([]*types.Chunk{b, a}))
|
||
|
|
b.ContentRevision++
|
||
|
|
require.NotEqual(t, wikiEvidenceSourceHash([]*types.Chunk{a}), wikiEvidenceSourceHash([]*types.Chunk{b}))
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiAdditionEvidenceHonorsCancellationBeforeLookup(t *testing.T) {
|
||
|
|
ctx, cancel := context.WithCancel(context.Background())
|
||
|
|
cancel()
|
||
|
|
repo := &wikiEvidenceChunksRepo{chunks: []*types.Chunk{wikiEvidenceChunk()}}
|
||
|
|
_, err := (&wikiIngestService{chunkRepo: repo}).resolveWikiAdditionEvidence(ctx, 7, "kb", []SlugUpdate{wikiEvidenceAddition()})
|
||
|
|
require.ErrorIs(t, err, context.Canceled)
|
||
|
|
require.Empty(t, repo.queries)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiEvidenceGatePreservesPublishedPageAndNeverCheckpointsUncitedCandidate(t *testing.T) {
|
||
|
|
for _, existing := range []bool{false, true} {
|
||
|
|
for _, partial := range []bool{false, true} {
|
||
|
|
t.Run(fmt.Sprintf("existing-%t-partial-%t", existing, partial), func(t *testing.T) {
|
||
|
|
db := openWikiCheckpointDB(t, filepath.Join(t.TempDir(), "checkpoints.sqlite"))
|
||
|
|
original := &types.WikiPage{ID: "page", TenantID: 7, KnowledgeBaseID: "kb", Slug: "concept/alpha", Title: "Alpha", PageType: types.WikiPageTypeConcept,
|
||
|
|
Status: types.WikiPageStatusPublished, Content: "Original supported content", Summary: "Original summary", SourceRefs: types.StringArray{"old-doc"}, ChunkRefs: types.StringArray{"old-chunk"}, Aliases: types.StringArray{"Old alias"}}
|
||
|
|
before := *original
|
||
|
|
pages := &checkpointWikiPages{pages: make(map[string]*types.WikiPage)}
|
||
|
|
if existing {
|
||
|
|
pages.pages[original.Slug] = original
|
||
|
|
}
|
||
|
|
svc := &wikiIngestService{checkpointDB: db, wikiService: pages, knowledgeSvc: &checkpointKnowledge{}, chunkRepo: &wikiEvidenceChunksRepo{chunks: []*types.Chunk{wikiEvidenceChunk()}}}
|
||
|
|
batch := &WikiBatchContext{SlugTitleMany: func(context.Context, []string) map[string]string { return nil }}
|
||
|
|
batch.checkpointScopes.Store("doc", &wikiCheckpointScope{7, "kb", "doc", "fingerprint"})
|
||
|
|
calls := 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
return &types.ChatResponse{Content: "SUMMARY: Unsafe page\nUNSUPPORTED_DETAILS_MUST_NOT_REPLACE_SOURCE", FinishReason: "stop"}, nil
|
||
|
|
}}
|
||
|
|
update := wikiEvidenceAddition()
|
||
|
|
if partial {
|
||
|
|
update.SourceChunks = append(update.SourceChunks, "missing")
|
||
|
|
} else {
|
||
|
|
update.SourceChunks = nil
|
||
|
|
}
|
||
|
|
updates := []SlugUpdate{{KnowledgeID: "old-doc", SourceRef: "old-doc", Slug: original.Slug, Type: "retract"}, update}
|
||
|
|
for retry := 0; retry < 2; retry++ {
|
||
|
|
changed, _, failed, err := svc.reduceSlugUpdates(context.Background(), model, "kb", original.Slug, updates, 7, batch, nil)
|
||
|
|
require.ErrorIs(t, err, errWikiCitationEvidence)
|
||
|
|
require.False(t, changed)
|
||
|
|
require.True(t, failed)
|
||
|
|
}
|
||
|
|
require.Zero(t, calls)
|
||
|
|
require.Zero(t, pages.writes)
|
||
|
|
if existing {
|
||
|
|
require.Equal(t, before, *pages.pages[original.Slug])
|
||
|
|
} else {
|
||
|
|
require.Empty(t, pages.pages)
|
||
|
|
}
|
||
|
|
var checkpoints int64
|
||
|
|
require.NoError(t, db.Model(&wikiIngestCheckpoint{}).Where("stage = ?", "page").Count(&checkpoints).Error)
|
||
|
|
require.Zero(t, checkpoints)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiEvidenceGateUsesVerbatimSourceWithoutDetailsFallback(t *testing.T) {
|
||
|
|
for _, origin := range []string{"pass-0", "citation-discovery"} {
|
||
|
|
t.Run(origin, func(t *testing.T) {
|
||
|
|
pages := &checkpointWikiPages{pages: make(map[string]*types.WikiPage)}
|
||
|
|
source := wikiEvidenceChunk()
|
||
|
|
source.Content += " Alpha is also called First."
|
||
|
|
update := wikiEvidenceAddition()
|
||
|
|
update.Item.Aliases = nil
|
||
|
|
if origin == "citation-discovery" {
|
||
|
|
batch := splitChunksIntoCitationBatches([]*types.Chunk{source})[0]
|
||
|
|
parsed, err := parseWikiCitationBatch(`{"citations":{},"new_slugs":[{"type":"concept","name":"Alpha","slug":"concept/alpha","aliases":[],"description":"UNSUPPORTED_DESCRIPTION_MUST_NOT_REPLACE_SOURCE","details":"UNSUPPORTED_DETAILS_MUST_NOT_REPLACE_SOURCE","source_chunks":["c000"]}]}`, batch, map[string]bool{})
|
||
|
|
require.NoError(t, err)
|
||
|
|
_, concepts, uncited := mergeCitationsIntoItems(nil, nil, parsed.Citations, parsed.NewSlugs)
|
||
|
|
require.Zero(t, uncited)
|
||
|
|
require.Len(t, concepts, 1)
|
||
|
|
update.Item = concepts[0]
|
||
|
|
update.SourceChunks = concepts[0].SourceChunks
|
||
|
|
}
|
||
|
|
svc := &wikiIngestService{wikiService: pages, knowledgeSvc: &checkpointKnowledge{}, chunkRepo: &wikiEvidenceChunksRepo{chunks: []*types.Chunk{source}}}
|
||
|
|
calls := 0
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(_ context.Context, messages []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
calls++
|
||
|
|
var prompt strings.Builder
|
||
|
|
for _, message := range messages {
|
||
|
|
prompt.WriteString(message.Content)
|
||
|
|
}
|
||
|
|
require.Contains(t, prompt.String(), "<content>\n"+source.Content+"\n</content>")
|
||
|
|
require.Contains(t, prompt.String(), "<slug>concept/alpha</slug>")
|
||
|
|
require.Contains(t, prompt.String(), "<title>Alpha</title>")
|
||
|
|
require.NotContains(t, prompt.String(), "<aliases>")
|
||
|
|
require.NotContains(t, prompt.String(), update.Item.Description)
|
||
|
|
require.NotContains(t, prompt.String(), update.Item.Details)
|
||
|
|
return &types.ChatResponse{Content: "SUMMARY: Supported page\n42 applies only when condition A holds.", FinishReason: "stop"}, nil
|
||
|
|
}}
|
||
|
|
batch := &WikiBatchContext{SlugTitleMany: func(context.Context, []string) map[string]string { return nil }}
|
||
|
|
changed, _, failed, err := svc.reduceSlugUpdatesUncheckpointed(context.Background(), model, "kb", update.Slug, []SlugUpdate{update}, 7, batch, nil)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.True(t, changed)
|
||
|
|
require.False(t, failed)
|
||
|
|
require.Equal(t, 1, calls)
|
||
|
|
require.Equal(t, 1, pages.writes)
|
||
|
|
require.Equal(t, types.StringArray{"source-chunk"}, pages.pages[update.Slug].ChunkRefs)
|
||
|
|
require.Equal(t, types.StringArray{"doc"}, pages.pages[update.Slug].SourceRefs)
|
||
|
|
require.Empty(t, pages.pages[update.Slug].Aliases)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|