286 lines
12 KiB
Go
286 lines
12 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"crypto/sha256"
|
||
|
|
"encoding/json"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
chatpipeline "github.com/Tencent/WeKnora/internal/application/service/chat_pipeline"
|
||
|
|
"github.com/Tencent/WeKnora/internal/models/chat"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||
|
|
"github.com/hibiken/asynq"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestNativeGraphReviewedSourceStopsBeforeModelOnChange(t *testing.T) {
|
||
|
|
const source = "Trust supports learning."
|
||
|
|
for _, changed := range []string{"unchanged", "content", "knowledge", "knowledge base", "chunk", "tenant", "empty identity", "uppercase hash", "disabled", "deleted", "kind", "document deleted", "document processing", "document tenant", "document base", "document identity", "document unavailable", "document repository unavailable", "base deleted", "base tenant", "base identity"} {
|
||
|
|
t.Run(changed, func(t *testing.T) {
|
||
|
|
chunk := &types.Chunk{ID: "chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: source, IsEnabled: true, ChunkType: types.ChunkTypeText}
|
||
|
|
expected := &types.GraphSourceIdentity{KnowledgeBaseID: "kb", KnowledgeID: "doc", ContentSHA256: fmt.Sprintf("%x", sha256.Sum256([]byte(source)))}
|
||
|
|
document := &types.Knowledge{ID: "doc", TenantID: 7, KnowledgeBaseID: "kb", ParseStatus: types.ParseStatusCompleted}
|
||
|
|
knowledge := &extractionKnowledge{document: document}
|
||
|
|
base := &types.KnowledgeBase{ID: "kb", TenantID: 7, ExtractConfig: &types.ExtractConfig{Enabled: true}}
|
||
|
|
switch changed {
|
||
|
|
case "content":
|
||
|
|
chunk.Content = "New source."
|
||
|
|
case "knowledge":
|
||
|
|
chunk.KnowledgeID = "other"
|
||
|
|
case "knowledge base":
|
||
|
|
chunk.KnowledgeBaseID = "other"
|
||
|
|
case "chunk":
|
||
|
|
chunk.ID = "other"
|
||
|
|
case "tenant":
|
||
|
|
chunk.TenantID = 8
|
||
|
|
case "empty identity":
|
||
|
|
expected.KnowledgeID = ""
|
||
|
|
case "uppercase hash":
|
||
|
|
expected.ContentSHA256 = strings.ToUpper(expected.ContentSHA256)
|
||
|
|
case "disabled":
|
||
|
|
chunk.IsEnabled = false
|
||
|
|
case "deleted":
|
||
|
|
chunk.DeletedAt.Valid = true
|
||
|
|
case "kind":
|
||
|
|
chunk.ChunkType = types.ChunkTypeImageOCR
|
||
|
|
case "document deleted":
|
||
|
|
document.DeletedAt.Valid = true
|
||
|
|
case "document processing":
|
||
|
|
document.ParseStatus = types.ParseStatusProcessing
|
||
|
|
case "document tenant":
|
||
|
|
document.TenantID = 8
|
||
|
|
case "document base":
|
||
|
|
document.KnowledgeBaseID = "other"
|
||
|
|
case "document identity":
|
||
|
|
document.ID = "other"
|
||
|
|
case "document unavailable":
|
||
|
|
knowledge.lookupErr = errors.New("unavailable")
|
||
|
|
case "base deleted":
|
||
|
|
base.DeletedAt.Valid = true
|
||
|
|
case "base tenant":
|
||
|
|
base.TenantID = 8
|
||
|
|
case "base identity":
|
||
|
|
base.ID = "other"
|
||
|
|
}
|
||
|
|
model := &extractionChat{response: `[{"entity1":"Trust","entity2":"learning","relation":"supports","evidence":"Trust supports learning."}]`}
|
||
|
|
writer := &extractionWriter{}
|
||
|
|
service := &ChunkExtractService{template: &types.PromptTemplateStructured{Description: "Extract graph"}, modelService: extractionModels{model: model}, knowledgeBaseRepo: extractionKB{base: base}, knowledgeRepo: knowledge, chunkRepo: extractionChunks{chunk: chunk}, graphEngine: writer}
|
||
|
|
if changed == "document repository unavailable" {
|
||
|
|
service.knowledgeRepo = nil
|
||
|
|
}
|
||
|
|
service.executionDB = graphBudgetDB(t)
|
||
|
|
payload, _ := json.Marshal(types.ExtractChunkPayload{TenantID: 7, ChunkID: "chunk", ModelID: "model", ExpectedSource: expected})
|
||
|
|
err := service.handle(context.Background(), asynq.NewTask(types.TypeChunkExtract, payload), "task")
|
||
|
|
if changed == "unchanged" {
|
||
|
|
if err != nil || model.calls != 1 || !writer.complete {
|
||
|
|
t.Fatalf("reviewed source rejected: %v", err)
|
||
|
|
}
|
||
|
|
} else if !errors.Is(err, asynq.SkipRetry) || model.calls != 0 || writer.saved != nil {
|
||
|
|
t.Fatalf("changed source reached model or persistence: %v calls=%d", err, model.calls)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestNativeGraphExtractionDoesNotRequireNeo4j(t *testing.T) {
|
||
|
|
t.Setenv("NEO4J_ENABLE", "false")
|
||
|
|
t.Setenv("LEGACY_GRAPH_ENABLE", "true")
|
||
|
|
if !graphExtractionEnabled() {
|
||
|
|
t.Fatal("native extraction was silently skipped")
|
||
|
|
}
|
||
|
|
t.Setenv("LEGACY_GRAPH_ENABLE", "false")
|
||
|
|
if graphExtractionEnabled() {
|
||
|
|
t.Fatal("unconfigured graph extraction enabled")
|
||
|
|
}
|
||
|
|
t.Setenv("NEO4J_ENABLE", "true")
|
||
|
|
if !graphExtractionEnabled() {
|
||
|
|
t.Fatal("Neo4j compatibility lost")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
type extractionChatBase interface{ chat.Chat }
|
||
|
|
type extractionChat struct {
|
||
|
|
extractionChatBase
|
||
|
|
calls int
|
||
|
|
response string
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *extractionChat) Chat(ctx context.Context, messages []chat.Message, opts *chat.ChatOptions) (*types.ChatResponse, error) {
|
||
|
|
m.calls++
|
||
|
|
return &types.ChatResponse{Content: m.response}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
type extractionModels struct {
|
||
|
|
interfaces.ModelService
|
||
|
|
model chat.Chat
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m extractionModels) GetChatModel(context.Context, string) (chat.Chat, error) {
|
||
|
|
return m.model, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
type extractionKB struct {
|
||
|
|
interfaces.KnowledgeBaseRepository
|
||
|
|
base *types.KnowledgeBase
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r extractionKB) GetKnowledgeBaseByID(context.Context, string) (*types.KnowledgeBase, error) {
|
||
|
|
if r.base != nil {
|
||
|
|
return r.base, nil
|
||
|
|
}
|
||
|
|
return &types.KnowledgeBase{ID: "kb", TenantID: 7, ExtractConfig: &types.ExtractConfig{Enabled: true}}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
type extractionChunks struct {
|
||
|
|
interfaces.ChunkRepository
|
||
|
|
chunk *types.Chunk
|
||
|
|
}
|
||
|
|
|
||
|
|
type extractionKnowledge struct {
|
||
|
|
interfaces.KnowledgeRepository
|
||
|
|
drained int
|
||
|
|
document *types.Knowledge
|
||
|
|
lookupErr error
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *extractionKnowledge) GetKnowledgeByIDOnly(context.Context, string) (*types.Knowledge, error) {
|
||
|
|
if r.lookupErr != nil {
|
||
|
|
return nil, r.lookupErr
|
||
|
|
}
|
||
|
|
if r.document != nil {
|
||
|
|
return r.document, nil
|
||
|
|
}
|
||
|
|
return &types.Knowledge{ID: "doc", TenantID: 7, KnowledgeBaseID: "kb", ParseStatus: types.ParseStatusCompleted}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *extractionKnowledge) FinalizeSubtask(context.Context, string) (int, bool, error) {
|
||
|
|
r.drained++
|
||
|
|
return 0, true, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r extractionChunks) GetChunkByID(context.Context, uint64, string) (*types.Chunk, error) {
|
||
|
|
copy := *r.chunk
|
||
|
|
return ©, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
type extractionWriter struct {
|
||
|
|
interfaces.RetrieveGraphRepository
|
||
|
|
saved *types.GraphData
|
||
|
|
source *types.Chunk
|
||
|
|
complete bool
|
||
|
|
fail bool
|
||
|
|
}
|
||
|
|
|
||
|
|
func (w *extractionWriter) HasChunkGraph(context.Context, *types.Chunk) (bool, error) {
|
||
|
|
return w.complete, nil
|
||
|
|
}
|
||
|
|
func (w *extractionWriter) AddChunkGraph(ctx context.Context, ns types.NameSpace, c *types.Chunk, g *types.GraphData) error {
|
||
|
|
if w.fail {
|
||
|
|
return fmt.Errorf("persistence failure")
|
||
|
|
}
|
||
|
|
if ns.Knowledge != "doc" || types.MustTenantIDFromContext(ctx) != 7 {
|
||
|
|
return fmt.Errorf("missing ownership scope")
|
||
|
|
}
|
||
|
|
w.saved, w.source, w.complete = g, c, true
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestNativeGraphTaskExtractsBeforeCheckpointing(t *testing.T) {
|
||
|
|
for _, tc := range []struct {
|
||
|
|
name, response string
|
||
|
|
fail, wantErr bool
|
||
|
|
}{
|
||
|
|
{"actual entities and relationship", `[{"entity":"Trust","entity_attributes":["concept"]},{"entity":"Learning"},{"entity1":"Trust","entity2":"Learning","relation":"supports","evidence":"Trust supports learning."}]`, false, false},
|
||
|
|
{"missing relation evidence", `[{"entity1":"Trust","entity2":"Learning","relation":"supports"}]`, false, true},
|
||
|
|
{"fabricated relation evidence", `[{"entity1":"Trust","entity2":"Learning","relation":"supports","evidence":"Invented source."}]`, false, true},
|
||
|
|
{"explicit empty needs review", `[]`, false, true},
|
||
|
|
{"entities without relations need review", `[{"entity":"Trust"}]`, false, true},
|
||
|
|
{"unrecognized response", `{"error":"provider rejected extraction"}`, false, true},
|
||
|
|
{"database error", `[{"entity1":"Trust","entity2":"Learning","relation":"supports","evidence":"Trust supports learning."}]`, true, true},
|
||
|
|
} {
|
||
|
|
t.Run(tc.name, func(t *testing.T) {
|
||
|
|
model := &extractionChat{response: tc.response}
|
||
|
|
writer := &extractionWriter{fail: tc.fail}
|
||
|
|
service := &ChunkExtractService{template: &types.PromptTemplateStructured{Description: "Extract grounded entities and relationships."}, modelService: extractionModels{model: model}, knowledgeBaseRepo: extractionKB{}, chunkRepo: extractionChunks{chunk: &types.Chunk{ID: "chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: "Trust supports learning."}}, graphEngine: writer}
|
||
|
|
payload, _ := json.Marshal(types.ExtractChunkPayload{TenantID: 7, ChunkID: "chunk", ModelID: "model"})
|
||
|
|
task := asynq.NewTask(types.TypeChunkExtract, payload)
|
||
|
|
err := service.Handle(context.Background(), task)
|
||
|
|
if (err != nil) != tc.wantErr {
|
||
|
|
t.Fatalf("unexpected task result: %v", err)
|
||
|
|
}
|
||
|
|
if tc.wantErr {
|
||
|
|
if writer.complete {
|
||
|
|
t.Fatal("failed task wrote a completion checkpoint")
|
||
|
|
}
|
||
|
|
if tc.name == "explicit empty needs review" || tc.name == "entities without relations need review" {
|
||
|
|
if !errors.Is(err, chatpipeline.ErrGraphNeedsReview) || !errors.Is(err, asynq.SkipRetry) {
|
||
|
|
t.Fatalf("review result should stop automatic retries: %v", err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return
|
||
|
|
}
|
||
|
|
if writer.saved == nil || writer.source == nil || writer.source.Content != "Trust supports learning." {
|
||
|
|
t.Fatal("model output/source not passed to durable writer")
|
||
|
|
}
|
||
|
|
if len(writer.saved.Node) > 0 && writer.saved.Node[0].Chunks[0] != "chunk" {
|
||
|
|
t.Fatal("missing evidence chunk")
|
||
|
|
}
|
||
|
|
if len(writer.saved.Relation) > 0 && writer.saved.Relation[0].Evidence != writer.source.Content {
|
||
|
|
t.Fatal("relation evidence was not passed to durable writer")
|
||
|
|
}
|
||
|
|
if err := service.Handle(context.Background(), task); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if model.calls != 1 {
|
||
|
|
t.Fatal("completed retry billed the model again")
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestNativeGraphTaskDoesNotCheckpointContradictedProductClaim(t *testing.T) {
|
||
|
|
const source = "私蜜妍+蜜桃初醒套盒系列\n过敏者慎甩；粘膜破损处禁用；不得用于性生活中对性病的预防。"
|
||
|
|
model := &extractionChat{response: `[{"entity1":"私蜜妍+蜜桃初醒套盒系列","entity2":"粘膜破损处","relation":"适用于","evidence":"粘膜破损处禁用"}]`}
|
||
|
|
writer := &extractionWriter{}
|
||
|
|
service := &ChunkExtractService{
|
||
|
|
template: &types.PromptTemplateStructured{Description: "Extract grounded entities and relationships."},
|
||
|
|
modelService: extractionModels{model: model}, knowledgeBaseRepo: extractionKB{},
|
||
|
|
chunkRepo: extractionChunks{chunk: &types.Chunk{ID: "chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: source}},
|
||
|
|
graphEngine: writer,
|
||
|
|
}
|
||
|
|
payload, _ := json.Marshal(types.ExtractChunkPayload{TenantID: 7, ChunkID: "chunk", ModelID: "model"})
|
||
|
|
if err := service.Handle(context.Background(), asynq.NewTask(types.TypeChunkExtract, payload)); err == nil {
|
||
|
|
t.Fatal("contradicted claim was accepted")
|
||
|
|
}
|
||
|
|
if writer.saved != nil || writer.complete {
|
||
|
|
t.Fatal("contradicted claim wrote a graph or completion checkpoint")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestNativeGraphReviewDrainsParentWithoutRetryOrCheckpoint(t *testing.T) {
|
||
|
|
model := &extractionChat{response: `[]`}
|
||
|
|
writer := &extractionWriter{}
|
||
|
|
knowledge := &extractionKnowledge{}
|
||
|
|
service := &ChunkExtractService{
|
||
|
|
template: &types.PromptTemplateStructured{Description: "Extract grounded relations."},
|
||
|
|
modelService: extractionModels{model: model}, knowledgeBaseRepo: extractionKB{},
|
||
|
|
knowledgeRepo: knowledge,
|
||
|
|
chunkRepo: extractionChunks{chunk: &types.Chunk{ID: "chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: "Trust supports learning."}},
|
||
|
|
graphEngine: writer,
|
||
|
|
}
|
||
|
|
payload, _ := json.Marshal(types.ExtractChunkPayload{TenantID: 7, KnowledgeID: "doc", ChunkID: "chunk", ModelID: "model"})
|
||
|
|
err := service.Handle(context.Background(), asynq.NewTask(types.TypeChunkExtract, payload))
|
||
|
|
if !errors.Is(err, asynq.SkipRetry) || !errors.Is(err, chatpipeline.ErrGraphNeedsReview) {
|
||
|
|
t.Fatalf("unexpected review outcome: %v", err)
|
||
|
|
}
|
||
|
|
if knowledge.drained != 1 || writer.complete || writer.saved != nil {
|
||
|
|
t.Fatalf("review left the parent pending or wrote a checkpoint: drained=%d complete=%t", knowledge.drained, writer.complete)
|
||
|
|
}
|
||
|
|
if model.calls != 2 {
|
||
|
|
t.Fatalf("expected one bounded correction, got %d calls", model.calls)
|
||
|
|
}
|
||
|
|
}
|