xinmi/internal/application/service/chat_pipeline/extract_entity_coverage_test.go

package chatpipeline

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"os"
	"strings"
	"testing"

	"github.com/Tencent/WeKnora/internal/types"
)

func graphCoverageResponse(t *testing.T, graph *types.GraphData) *types.ChatResponse {
	t.Helper()
	items := make([]map[string]any, 0, len(graph.Node)+len(graph.Relation))
	for _, node := range graph.Node {
		items = append(items, map[string]any{"entity": node.Name, "entity_attributes": node.Attributes})
	}
	for _, relation := range graph.Relation {
		items = append(items, map[string]any{"entity1": relation.Node1, "entity2": relation.Node2, "relation": relation.Type, "evidence": relation.Evidence})
	}
	raw, err := json.Marshal(items)
	if err != nil {
		t.Fatal(err)
	}
	return &types.ChatResponse{Content: string(raw), FinishReason: "stop"}
}

func TestGraphStorageCoverageSavedProbeRequiresReview(t *testing.T) {
	raw, err := os.ReadFile("testdata/graph_storage_coverage_add45b26.json")
	if err != nil {
		t.Fatal(err)
	}
	var fixture struct {
		SourceSHA256 string           `json:"source_sha256"`
		Source       string           `json:"source_content"`
		Graph        *types.GraphData `json:"graph"`
	}
	if err := json.Unmarshal(raw, &fixture); err != nil {
		t.Fatal(err)
	}
	digest := sha256.Sum256([]byte(fixture.Source))
	if hex.EncodeToString(digest[:]) != fixture.SourceSHA256 || len(fixture.Graph.Relation) != 10 {
		t.Fatal("saved probe source or returned relations changed")
	}
	model := &extractionResponseChat{response: graphCoverageResponse(t, fixture.Graph)}
	extractor := NewGraphExtractor(model, &types.PromptTemplateStructured{Description: "Extract graph"})
	graph, err := extractor.Extract(context.Background(), fixture.Source)
	if graph != nil || !errors.Is(err, ErrGraphNeedsReview) || !strings.Contains(err.Error(), "置于丿 L 童不易触及处") || model.calls != 2 {
		t.Fatalf("saved omission must require review after one correction: graph=%+v, err=%v, calls=%d", graph, err, model.calls)
	}
	if !strings.Contains(model.messages[len(model.messages)-1].Content, "置于丿 L 童不易触及处") {
		t.Fatal("correction lost the original OCR clause")
	}
}

func TestGraphStorageCoverageRequiresCompleteConstraint(t *testing.T) {
	const clause = "置于丿 L 童不易触及处"
	const source = "本品不得口服，" + clause + "。"
	for _, tc := range []struct {
		name, target, label, evidence string
		attribute                     bool
	}{
		{name: "omitted relation"},
		{name: "attribute alone", attribute: true},
		{name: "normalized OCR target", target: "置于儿童不易触及处", label: "约束", evidence: clause},
		{name: "restriction removed from target", target: "置于丿 L 童触及处", label: "约束", evidence: clause},
		{name: "evidence alone", target: "储存要求", label: "约束", evidence: clause},
		{name: "unrelated evidence", target: clause, label: "约束", evidence: "本品不得口服"},
		{name: "wrong relationship", target: clause, label: "关联", evidence: clause},
	} {
		t.Run(tc.name, func(t *testing.T) {
			input := &types.GraphData{Relation: []*types.GraphRelation{{Node1: "本品", Node2: "口服", Type: "禁用于", Evidence: "本品不得口服"}}}
			if tc.target != "" {
				input.Relation = append(input.Relation, &types.GraphRelation{Node1: "本品", Node2: tc.target, Type: tc.label, Evidence: tc.evidence})
			}
			if tc.attribute {
				input.Node = []*types.GraphNode{{Name: "本品", Attributes: []string{clause}}}
			}
			model := &extractionResponseChat{response: graphCoverageResponse(t, input)}
			extractor := NewGraphExtractor(model, &types.PromptTemplateStructured{Description: "Extract graph"})
			graph, err := extractor.Extract(context.Background(), source)
			if graph != nil || !errors.Is(err, ErrGraphNeedsReview) || model.calls != 2 {
				t.Fatalf("missing constraint passed: graph=%+v, err=%v, calls=%d", graph, err, model.calls)
			}
		})
	}
}

func TestGraphStorageCoverageCorrectionPreservesOCR(t *testing.T) {
	const clause = "置于丿 L 童不易触及处"
	const source = "本品不得口服，" + clause + "。"
	input := &types.GraphData{Relation: []*types.GraphRelation{{Node1: "本品", Node2: "口服", Type: "禁用于", Evidence: "本品不得口服"}}}
	incomplete := graphCoverageResponse(t, input)
	input.Relation = append(input.Relation, &types.GraphRelation{Node1: "本品", Node2: clause, Type: "约束", Evidence: clause})
	model := &extractionResponseChat{responses: []*types.ChatResponse{incomplete, graphCoverageResponse(t, input)}}
	extractor := NewGraphExtractor(model, &types.PromptTemplateStructured{Description: "Extract graph"})
	graph, err := extractor.Extract(context.Background(), source)
	if err != nil || graph == nil || len(graph.Relation) != 2 || model.calls != 2 {
		t.Fatalf("source-backed correction failed: graph=%+v, err=%v, calls=%d", graph, err, model.calls)
	}
	if graph.Relation[1].Node2 != clause || graph.Relation[1].Evidence != clause {
		t.Fatalf("OCR was changed: %+v", graph.Relation[1])
	}
}

func TestGraphStorageCoverageScope(t *testing.T) {
	for _, tc := range []struct {
		source     string
		wantReview bool
	}{
		{"请放置于儿童不易触及之处。", true},
		{"〔 储存方法 〗 置于丿 L 童不易触及处。", true},
		{"放在儿童不易触及处。", true},
		{"教材介绍储存要求。", false},
		{"作者讨论“置于儿童不易触及处”这句话。", false},
		{"他认为置于儿童不易触及处能够保证安全。", false},
	} {
		err := validateGraphStorageCoverage(&types.GraphData{}, tc.source)
		if errors.Is(err, ErrGraphNeedsReview) != tc.wantReview {
			t.Errorf("source=%q, error=%v, want review=%v", tc.source, err, tc.wantReview)
		}
	}
	model := &extractionResponseChat{response: &types.ChatResponse{Content: `[{"entity":"储存"}]`, FinishReason: "stop"}}
	extractor := NewExtractor(model, &types.PromptTemplateStructured{Description: "Extract query entities"})
	graph, err := extractor.Extract(context.Background(), "置于丿 L 童不易触及处")
	if err != nil || graph == nil || len(graph.Node) != 1 || model.calls != 1 {
		t.Fatalf("query extraction should not require persisted-graph coverage: graph=%+v, err=%v", graph, err)
	}
}