351 lines
13 KiB
Go
351 lines
13 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"fmt"
|
||
|
|
"path/filepath"
|
||
|
|
"sort"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/models/chat"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
// TestMergeCitationsIntoItems_PopulatesSourceChunksOnCandidates verifies that
|
||
|
|
// citations returned by the chunk-classification pass are attached back onto
|
||
|
|
// the matching candidate items while non-cited candidates are left untouched.
|
||
|
|
func TestMergeCitationsIntoItems_PopulatesSourceChunksOnCandidates(t *testing.T) {
|
||
|
|
entities := []extractedItem{
|
||
|
|
{Name: "Acme", Slug: "entity/acme"},
|
||
|
|
{Name: "Beta", Slug: "entity/beta"},
|
||
|
|
}
|
||
|
|
concepts := []extractedItem{
|
||
|
|
{Name: "RAG", Slug: "concept/rag"},
|
||
|
|
}
|
||
|
|
citations := map[string][]string{
|
||
|
|
"entity/acme": {"chunk-1", "chunk-3"},
|
||
|
|
"concept/rag": {"chunk-2"},
|
||
|
|
}
|
||
|
|
|
||
|
|
gotE, gotC, uncited := mergeCitationsIntoItems(entities, concepts, citations, nil)
|
||
|
|
|
||
|
|
if len(gotE) != 2 || len(gotC) != 1 {
|
||
|
|
t.Fatalf("expected 2 entities + 1 concept, got %d + %d", len(gotE), len(gotC))
|
||
|
|
}
|
||
|
|
acme := findBySlug(gotE, "entity/acme")
|
||
|
|
if acme == nil {
|
||
|
|
t.Fatalf("entity/acme missing from result")
|
||
|
|
}
|
||
|
|
if !equalStrings(acme.SourceChunks, []string{"chunk-1", "chunk-3"}) {
|
||
|
|
t.Errorf("entity/acme source_chunks = %v, want [chunk-1 chunk-3]", acme.SourceChunks)
|
||
|
|
}
|
||
|
|
beta := findBySlug(gotE, "entity/beta")
|
||
|
|
if beta == nil {
|
||
|
|
t.Fatalf("entity/beta missing from result")
|
||
|
|
}
|
||
|
|
if len(beta.SourceChunks) != 0 {
|
||
|
|
t.Errorf("entity/beta should have no citations, got %v", beta.SourceChunks)
|
||
|
|
}
|
||
|
|
rag := findBySlug(gotC, "concept/rag")
|
||
|
|
if rag == nil {
|
||
|
|
t.Fatalf("concept/rag missing")
|
||
|
|
}
|
||
|
|
if !equalStrings(rag.SourceChunks, []string{"chunk-2"}) {
|
||
|
|
t.Errorf("concept/rag source_chunks = %v, want [chunk-2]", rag.SourceChunks)
|
||
|
|
}
|
||
|
|
if uncited != 1 {
|
||
|
|
t.Errorf("uncited = %d, want 1", uncited)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// TestMergeCitationsIntoItems_AddsNewSlugsAndUnionsChunksAcrossBatches checks
|
||
|
|
// that genuinely new slugs (ones Pass 0 missed) are appended to the right
|
||
|
|
// type slice, and that a slug surfacing in two batches ends up with the union
|
||
|
|
// of its source chunks.
|
||
|
|
func TestMergeCitationsIntoItems_AddsNewSlugsAndUnionsChunksAcrossBatches(t *testing.T) {
|
||
|
|
entities := []extractedItem{
|
||
|
|
{Name: "Known", Slug: "entity/known"},
|
||
|
|
}
|
||
|
|
concepts := []extractedItem{}
|
||
|
|
|
||
|
|
newSlugs := []newSlugFromCitation{
|
||
|
|
{
|
||
|
|
Type: "entity",
|
||
|
|
Name: "Fresh Entity",
|
||
|
|
Slug: "entity/fresh",
|
||
|
|
Description: "desc",
|
||
|
|
Details: "details",
|
||
|
|
SourceChunks: []string{"c001", "c002"},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
// Same slug as above, appears in another batch — must union.
|
||
|
|
Type: "entity",
|
||
|
|
Name: "Fresh Entity",
|
||
|
|
Slug: "entity/fresh",
|
||
|
|
SourceChunks: []string{"c002", "c003"},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Type: "concept",
|
||
|
|
Name: "New Concept",
|
||
|
|
Slug: "concept/new-concept",
|
||
|
|
SourceChunks: []string{"c010"},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
// Duplicate of an existing candidate — should NOT produce a
|
||
|
|
// duplicate entry (Known already exists in `entities`).
|
||
|
|
Type: "entity",
|
||
|
|
Name: "Known",
|
||
|
|
Slug: "entity/known",
|
||
|
|
SourceChunks: []string{"c020"},
|
||
|
|
},
|
||
|
|
}
|
||
|
|
|
||
|
|
gotE, gotC, _ := mergeCitationsIntoItems(entities, concepts, nil, newSlugs)
|
||
|
|
|
||
|
|
if len(gotE) != 2 {
|
||
|
|
t.Fatalf("expected 2 entities, got %d (%+v)", len(gotE), gotE)
|
||
|
|
}
|
||
|
|
if len(gotC) != 1 {
|
||
|
|
t.Fatalf("expected 1 concept, got %d (%+v)", len(gotC), gotC)
|
||
|
|
}
|
||
|
|
fresh := findBySlug(gotE, "entity/fresh")
|
||
|
|
if fresh == nil {
|
||
|
|
t.Fatalf("entity/fresh missing")
|
||
|
|
}
|
||
|
|
sort.Strings(fresh.SourceChunks)
|
||
|
|
if !equalStrings(fresh.SourceChunks, []string{"c001", "c002", "c003"}) {
|
||
|
|
t.Errorf("entity/fresh source_chunks = %v, want union [c001 c002 c003]", fresh.SourceChunks)
|
||
|
|
}
|
||
|
|
newC := findBySlug(gotC, "concept/new-concept")
|
||
|
|
if newC == nil || !equalStrings(newC.SourceChunks, []string{"c010"}) {
|
||
|
|
t.Errorf("concept/new-concept missing or wrong chunks: %+v", newC)
|
||
|
|
}
|
||
|
|
known := findBySlug(gotE, "entity/known")
|
||
|
|
require.NotNil(t, known)
|
||
|
|
require.Equal(t, []string{"c020"}, known.SourceChunks)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestMergeCitationsIntoItems_PreservesValidatedEvidenceRepresentations(t *testing.T) {
|
||
|
|
batch := splitChunksIntoCitationBatches([]*types.Chunk{
|
||
|
|
{ID: "first", Content: "Alpha first", ChunkIndex: 0},
|
||
|
|
{ID: "second", Content: "Alpha second", ChunkIndex: 1},
|
||
|
|
})[0]
|
||
|
|
for _, kind := range []string{"entity", "concept"} {
|
||
|
|
for _, existing := range []bool{false, true} {
|
||
|
|
for _, test := range []struct {
|
||
|
|
name string
|
||
|
|
citations, declared []string
|
||
|
|
want []string
|
||
|
|
}{
|
||
|
|
{name: "source-chunks-only", declared: []string{"c000"}, want: []string{"first"}},
|
||
|
|
{name: "citations-only", citations: []string{"c001"}, want: []string{"second"}},
|
||
|
|
{name: "both-with-duplicates", citations: []string{"c000", "c000"}, declared: []string{"c000", "c001", "c001"}, want: []string{"first", "second"}},
|
||
|
|
} {
|
||
|
|
t.Run(fmt.Sprintf("%s/existing-%t/%s", kind, existing, test.name), func(t *testing.T) {
|
||
|
|
slug := kind + "/alpha"
|
||
|
|
result := citationBatchResult{
|
||
|
|
Citations: map[string][]string{},
|
||
|
|
NewSlugs: []newSlugFromCitation{{Type: kind, Name: "Alpha", Slug: slug, SourceChunks: test.declared}},
|
||
|
|
}
|
||
|
|
if test.citations != nil {
|
||
|
|
result.Citations[slug] = test.citations
|
||
|
|
}
|
||
|
|
allowed := map[string]bool{}
|
||
|
|
var entities, concepts []extractedItem
|
||
|
|
if existing {
|
||
|
|
allowed[slug] = true
|
||
|
|
candidate := extractedItem{Name: "Alpha", Slug: slug}
|
||
|
|
if kind == "entity" {
|
||
|
|
entities = append(entities, candidate)
|
||
|
|
} else {
|
||
|
|
concepts = append(concepts, candidate)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
raw, err := json.Marshal(result)
|
||
|
|
require.NoError(t, err)
|
||
|
|
parsed, err := parseWikiCitationBatch(string(raw), batch, allowed)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.NoError(t, validateWikiCitationCheckpoint(&parsed, batch, allowed))
|
||
|
|
gotE, gotC, uncited := mergeCitationsIntoItems(entities, concepts, parsed.Citations, parsed.NewSlugs)
|
||
|
|
all := append(gotE, gotC...)
|
||
|
|
require.Len(t, all, 1)
|
||
|
|
require.Equal(t, slug, all[0].Slug)
|
||
|
|
require.Equal(t, test.want, all[0].SourceChunks)
|
||
|
|
require.Zero(t, uncited)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestMergeCitationsIntoItems_MergesDeclarationIdentityMetadata(t *testing.T) {
|
||
|
|
entities := []extractedItem{{Name: "Alpha", Slug: "entity/alpha", Aliases: []string{"original"}, Description: "short", Details: "short"}}
|
||
|
|
newSlugs := []newSlugFromCitation{
|
||
|
|
{Type: "entity", Name: "A lpha", Slug: "entity/alpha", Aliases: []string{"first alias"}, Description: "longer description", SourceChunks: []string{"first"}},
|
||
|
|
{Type: "entity", Name: "Alpha", Slug: "entity/alpha", Aliases: []string{"second alias", "first alias"}, Details: "longer details", SourceChunks: []string{"second", "first"}},
|
||
|
|
{Type: "concept", Name: "B eta", Slug: "concept/beta", Aliases: []string{"first alias"}, Description: "short", Details: "longer details", SourceChunks: []string{"first"}},
|
||
|
|
{Type: "concept", Name: "Beta", Slug: "concept/beta", Aliases: []string{"second alias"}, Description: "longer description", Details: "short", SourceChunks: []string{"second"}},
|
||
|
|
}
|
||
|
|
gotE, gotC, uncited := mergeCitationsIntoItems(entities, nil, nil, newSlugs)
|
||
|
|
require.Zero(t, uncited)
|
||
|
|
require.Len(t, gotE, 1)
|
||
|
|
require.Len(t, gotC, 1)
|
||
|
|
require.Equal(t, "Alpha", gotE[0].Name)
|
||
|
|
require.ElementsMatch(t, []string{"original", "A lpha", "first alias", "second alias"}, gotE[0].Aliases)
|
||
|
|
require.Equal(t, "Beta", gotC[0].Name)
|
||
|
|
require.ElementsMatch(t, []string{"B eta", "first alias", "second alias"}, gotC[0].Aliases)
|
||
|
|
for _, item := range append(gotE, gotC...) {
|
||
|
|
require.Equal(t, "longer description", item.Description)
|
||
|
|
require.Equal(t, "longer details", item.Details)
|
||
|
|
require.Equal(t, []string{"first", "second"}, item.SourceChunks)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestMergeCitationsIntoItems_DeduplicatesExistingCandidateCitations(t *testing.T) {
|
||
|
|
got, _, uncited := mergeCitationsIntoItems(
|
||
|
|
[]extractedItem{{Name: "Alpha", Slug: "entity/alpha", SourceChunks: []string{"unvalidated-extraction-source"}}}, nil,
|
||
|
|
map[string][]string{"entity/alpha": {"first", "first", "second"}}, nil,
|
||
|
|
)
|
||
|
|
require.Zero(t, uncited)
|
||
|
|
require.Equal(t, []string{"first", "second"}, got[0].SourceChunks)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCitationEvidenceUnionStillRejectsInvalidBatchEvidence(t *testing.T) {
|
||
|
|
batch := splitChunksIntoCitationBatches([]*types.Chunk{{ID: "first", Content: "first"}})[0]
|
||
|
|
allowed := map[string]bool{"entity/alpha": true}
|
||
|
|
for _, test := range []struct{ name, output string }{
|
||
|
|
{"undeclared-slug", `{"citations":{"concept/new":["c000"]}}`},
|
||
|
|
{"invalid-citation-source", `{"citations":{"entity/alpha":["c999"]}}`},
|
||
|
|
{"invalid-declaration-source", `{"citations":{},"new_slugs":[{"type":"entity","name":"Alpha","slug":"entity/alpha","source_chunks":["c999"]}]}`},
|
||
|
|
{"invalid-new-slug-type", `{"citations":{},"new_slugs":[{"type":"entity","name":"New","slug":"concept/new","source_chunks":["c000"]}]}`},
|
||
|
|
} {
|
||
|
|
t.Run(test.name, func(t *testing.T) {
|
||
|
|
_, err := parseWikiCitationBatch(test.output, batch, allowed)
|
||
|
|
require.Error(t, err)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
for _, result := range []citationBatchResult{
|
||
|
|
{Citations: map[string][]string{"entity/alpha": {"foreign-source"}}},
|
||
|
|
{Citations: map[string][]string{}, NewSlugs: []newSlugFromCitation{{Type: "entity", Name: "Alpha", Slug: "entity/alpha", SourceChunks: []string{"foreign-source"}}}},
|
||
|
|
} {
|
||
|
|
require.Error(t, validateWikiCitationCheckpoint(&result, batch, allowed))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiCitationEvidenceUnionSurvivesCheckpointReload(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 := []*types.Chunk{{ID: "first", Content: "first Alpha observation", ChunkIndex: 0}, {ID: "second", Content: "second Alpha observation", ChunkIndex: 1}}
|
||
|
|
batch := splitChunksIntoCitationBatches(chunks)[0]
|
||
|
|
entities := []extractedItem{{Name: "Alpha", Slug: "entity/alpha"}}
|
||
|
|
candidates := renderCandidateSlugsXML(entities, nil)
|
||
|
|
key := wikiCheckpointHash([]any{candidates, "English", wikiCheckpointChunks(batch.chunks)})
|
||
|
|
result := citationBatchResult{
|
||
|
|
Citations: map[string][]string{"entity/alpha": {"first"}, "concept/beta": {"second"}},
|
||
|
|
NewSlugs: []newSlugFromCitation{
|
||
|
|
{Type: "entity", Name: "Alpha", Slug: "entity/alpha", SourceChunks: []string{"second"}},
|
||
|
|
{Type: "concept", Name: "Beta", Slug: "concept/beta"},
|
||
|
|
},
|
||
|
|
}
|
||
|
|
require.NoError(t, validateWikiCitationCheckpoint(&result, batch, wikiCandidateSlugSet(candidates)))
|
||
|
|
require.NoError(t, svc.saveWikiCheckpoint(context.Background(), scope, "citation", key, result))
|
||
|
|
sqlDB, err := db.DB()
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.NoError(t, sqlDB.Close())
|
||
|
|
svc = &wikiIngestService{checkpointDB: openWikiCheckpointDB(t, path)}
|
||
|
|
model := &wikiScriptedChat{id: t.Name(), call: func(context.Context, []chat.Message) (*types.ChatResponse, error) {
|
||
|
|
t.Fatal("validated citation checkpoint must not call the model again")
|
||
|
|
return nil, nil
|
||
|
|
}}
|
||
|
|
citations, declarations, completed, err := svc.classifyChunkCitations(context.Background(), model, candidates, chunks, "English", &WikiBatchContext{}, scope)
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, 1, completed)
|
||
|
|
gotE, gotC, uncited := mergeCitationsIntoItems(entities, nil, citations, declarations)
|
||
|
|
require.Zero(t, uncited)
|
||
|
|
require.Len(t, gotE, 1)
|
||
|
|
require.Len(t, gotC, 1)
|
||
|
|
require.Equal(t, []string{"first", "second"}, gotE[0].SourceChunks)
|
||
|
|
require.Equal(t, []string{"second"}, gotC[0].SourceChunks)
|
||
|
|
}
|
||
|
|
|
||
|
|
// TestSplitChunksIntoCitationBatches_RespectsBudgetAndOrder verifies that the
|
||
|
|
// batcher never puts too many runes in one batch, preserves document order,
|
||
|
|
// and that an oversized chunk gets its own batch.
|
||
|
|
func TestSplitChunksIntoCitationBatches_RespectsBudgetAndOrder(t *testing.T) {
|
||
|
|
// Each small chunk is 5k runes → 3 of them should fit in one batch
|
||
|
|
// (15k > 12k limit would spill to a second batch).
|
||
|
|
mk := func(idx int, runes int, id string) *types.Chunk {
|
||
|
|
return &types.Chunk{
|
||
|
|
ID: id,
|
||
|
|
ChunkIndex: idx,
|
||
|
|
Content: repeatRune('a', runes),
|
||
|
|
ChunkType: types.ChunkTypeText,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
chunks := []*types.Chunk{
|
||
|
|
mk(0, 5000, "id-0"),
|
||
|
|
mk(1, 5000, "id-1"),
|
||
|
|
mk(2, 5000, "id-2"), // this should start a new batch (15k > 12k)
|
||
|
|
// An oversized chunk gets a dedicated batch.
|
||
|
|
mk(3, 20000, "id-big"),
|
||
|
|
mk(4, 1000, "id-small"),
|
||
|
|
}
|
||
|
|
batches := splitChunksIntoCitationBatches(chunks)
|
||
|
|
if len(batches) < 3 {
|
||
|
|
t.Fatalf("expected at least 3 batches, got %d", len(batches))
|
||
|
|
}
|
||
|
|
// All input IDs should show up in some batch, exactly once, in order.
|
||
|
|
seen := []string{}
|
||
|
|
for _, b := range batches {
|
||
|
|
for _, c := range b.chunks {
|
||
|
|
seen = append(seen, c.ID)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
wantOrder := []string{"id-0", "id-1", "id-2", "id-big", "id-small"}
|
||
|
|
if !equalStrings(seen, wantOrder) {
|
||
|
|
t.Errorf("batch order = %v, want %v", seen, wantOrder)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Verify the typed handle table is populated per batch.
|
||
|
|
for bi, b := range batches {
|
||
|
|
if b.handles.Len() != len(b.chunks) {
|
||
|
|
t.Errorf("batch %d handle count %d != chunk count %d", bi, b.handles.Len(), len(b.chunks))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// --- helpers ---
|
||
|
|
|
||
|
|
func findBySlug(items []extractedItem, slug string) *extractedItem {
|
||
|
|
for i := range items {
|
||
|
|
if items[i].Slug == slug {
|
||
|
|
return &items[i]
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func equalStrings(a, b []string) bool {
|
||
|
|
if len(a) != len(b) {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
for i := range a {
|
||
|
|
if a[i] != b[i] {
|
||
|
|
return false
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
|
||
|
|
func repeatRune(r rune, n int) string {
|
||
|
|
out := make([]rune, n)
|
||
|
|
for i := range out {
|
||
|
|
out[i] = r
|
||
|
|
}
|
||
|
|
return string(out)
|
||
|
|
}
|