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

package chatpipeline

import (
	"context"
	"errors"
	"strings"
	"testing"

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

const graphRestrictionTestBlock = "〔 不适宜人群 〕 对甲类慎用 、 乙类 、 丙类"
const graphRestrictionTestSource = "产品名称：甲品\n" + graphRestrictionTestBlock + "。"

func TestGraphRestrictionStrengthBelongsToItsOwnItem(t *testing.T) {
	for _, tc := range []struct {
		name, target, label, evidence, wantErr string
	}{
		{"local caution", "甲类", "慎用于", "对甲类慎用", ""},
		{"local caution preserves source preposition", "对甲类", "慎用于", "对甲类慎用", ""},
		{"local caution with full list evidence", "甲类", "慎用于", graphRestrictionTestBlock, ""},
		{"first inherited non suitability", "乙类", "不适用于", graphRestrictionTestBlock, ""},
		{"last inherited non suitability", "丙类", "不适用于", graphRestrictionTestBlock, ""},
		{"neighbor caution cannot propagate", "乙类", "慎用于", graphRestrictionTestBlock, "restriction strength mismatch"},
		{"later neighbor caution cannot propagate", "丙类", "慎用于", graphRestrictionTestBlock, "restriction strength mismatch"},
		{"caution cannot strengthen to non suitability", "甲类", "不适用于", graphRestrictionTestBlock, "restriction strength mismatch"},
		{"caution cannot strengthen to prohibition", "甲类", "禁用于", graphRestrictionTestBlock, "restriction strength mismatch"},
		{"non suitability cannot strengthen to prohibition", "乙类", "禁用于", graphRestrictionTestBlock, "restriction strength mismatch"},
		{"permission cannot replace restriction", "乙类", "适用于", graphRestrictionTestBlock, "restriction strength mismatch"},
		{"association cannot replace restriction", "乙类", "关联", graphRestrictionTestBlock, "restriction strength mismatch"},
		{"heading predicate cannot be cropped", "乙类", "不适用于", "乙类", "governing predicate"},
		{"local predicate cannot be cropped", "甲类", "慎用于", "甲类", "governing predicate"},
		{"whole source constraint", graphRestrictionTestBlock, "约束", graphRestrictionTestBlock, ""},
		{"truncated constraint", "乙类", "约束", graphRestrictionTestBlock, "restriction strength mismatch"},
		{"changed target spelling", "乙组", "不适用于", graphRestrictionTestBlock, "ambiguous restriction list"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			graph := &types.GraphData{Relation: []*types.GraphRelation{{Node1: "甲品", Node2: tc.target, Type: tc.label, Evidence: tc.evidence}}}
			err := validateGraphRelationEvidence(graph, graphRestrictionTestSource)
			if tc.wantErr == "" {
				if err != nil {
					t.Fatal(err)
				}
			} else if !errors.Is(err, ErrGraphNeedsReview) || !strings.Contains(err.Error(), tc.wantErr) {
				t.Fatalf("error = %v, want review containing %q", err, tc.wantErr)
			}
		})
	}
}

func TestGraphRestrictionScopesRetainSourceSpans(t *testing.T) {
	for _, tc := range []struct {
		block, target, label string
	}{
		{"【慎用人群】甲类，乙类", "乙类", "慎用于"},
		{"〖禁用对象〕甲类;乙类", "乙类", "禁用于"},
		{"[不适用人群] 丙类、乙类", "丙类", "不适用于"},
		{"不适宜人群：丙类；乙类", "丙类", "不适用于"},
		{"〔 不适宜人群 〕 对甲类慎甩，乙丿 IJ 类、丙类", "乙丿 IJ 类", "不适用于"},
		{"〔不适合对象〕甲类\n 、\t乙类", "乙类", "不适用于"},
		{"〔慎用人群〕甲类禁用、乙类", "甲类", "禁用于"},
		{"〔慎用人群〕甲类禁用、乙类", "乙类", "慎用于"},
	} {
		t.Run(tc.block, func(t *testing.T) {
			source := "产品名称：甲品\n" + tc.block + "\n〔贮存条件〕阴凉处保存。"
			scopes := graphRestrictionScopes(source)
			if len(scopes) != 1 || source[scopes[0].start:scopes[0].end] != tc.block {
				t.Fatal("restriction list boundaries changed source text")
			}
			relation := &types.GraphRelation{Node1: "甲品", Node2: tc.target, Type: tc.label, Evidence: tc.block}
			if err := validateGraphRelationEvidence(&types.GraphData{Relation: []*types.GraphRelation{relation}}, source); err != nil {
				t.Fatal(err)
			}
			if relation.Evidence != tc.block || relation.Node2 != tc.target {
				t.Fatal("validation mutated source text or target")
			}
		})
	}
}

func TestGraphRestrictionAmbiguityRequiresCompleteLiteralBlock(t *testing.T) {
	for _, block := range []string{
		"〔不适宜人群〕甲类、乙类慎用",
		"〔不适宜人群〕甲类、甲类",
		"〔不适宜人群〕甲类和乙类慎用",
		"〔不适宜人群〕甲类慎甪、乙类",
		"〔不适冝人群〕甲类、乙类",
		"〔不适宜人群〕甲类（特殊条件除外）、乙类",
	} {
		t.Run(block, func(t *testing.T) {
			source := "本品\n" + block + "。"
			graph := &types.GraphData{Relation: []*types.GraphRelation{{Node1: "本品", Node2: "甲类", Type: "不适用于", Evidence: block}}}
			if err := validateGraphRelationEvidence(graph, source); !errors.Is(err, ErrGraphNeedsReview) {
				t.Fatalf("ambiguous list accepted: %v", err)
			}
			graph.Relation[0].Node2, graph.Relation[0].Type = block, "约束"
			if err := validateGraphRelationEvidence(graph, source); err != nil {
				t.Fatal("literal constraint rejected", err)
			}
			if err := validateGraphRestrictionCoverage(graph, source); err != nil {
				t.Fatal("complete literal constraint lost coverage", err)
			}
			graph.Relation[0].Node2 = strings.TrimPrefix(block, "〔")
			if err := validateGraphRestrictionCoverage(graph, source); !errors.Is(err, ErrGraphNeedsReview) {
				t.Fatal("rewritten target satisfied full-block coverage")
			}
		})
	}
}

func TestGraphRestrictionCoverageCannotDropOrHideMembers(t *testing.T) {
	valid := []*types.GraphRelation{
		{Node1: "甲品", Node2: "甲类", Type: "慎用于", Evidence: "对甲类慎用"},
		{Node1: "甲品", Node2: "乙类", Type: "不适用于", Evidence: graphRestrictionTestBlock},
		{Node1: "甲品", Node2: "丙类", Type: "不适用于", Evidence: graphRestrictionTestBlock},
	}
	if err := validateGraphRestrictionCoverage(&types.GraphData{Relation: valid}, graphRestrictionTestSource); err != nil {
		t.Fatal(err)
	}
	valid[0].Node2 = "对甲类"
	if err := validateGraphRestrictionCoverage(&types.GraphData{Relation: valid}, graphRestrictionTestSource); err != nil {
		t.Fatal("literal source preposition lost coverage", err)
	}
	valid[0].Node2 = "甲类"
	for missing := range valid {
		relations := append([]*types.GraphRelation{}, valid[:missing]...)
		relations = append(relations, valid[missing+1:]...)
		graph := &types.GraphData{Node: []*types.GraphNode{{Name: "甲品", Attributes: []string{graphRestrictionTestBlock}}}, Relation: relations}
		if err := validateGraphRestrictionCoverage(graph, graphRestrictionTestSource); !errors.Is(err, ErrGraphNeedsReview) {
			t.Fatal("attributes or neighboring edges concealed a missing list member")
		}
	}
	const storage = "阴凉处保存"
	graph := &types.GraphData{Relation: []*types.GraphRelation{{Node1: "甲品", Node2: storage, Type: "约束", Evidence: storage}}}
	response := graphCoverageResponse(t, graph)
	source := graphRestrictionTestSource + "\n" + storage + "。"
	if accepted, err := NewFormater().parseGraph(context.Background(), response.Content, &source); accepted != nil || !errors.Is(err, ErrGraphNeedsReview) {
		t.Fatal("unrelated valid relation allowed restriction omission")
	}
}

func TestGraphRestrictionEvidenceAndStrengthErrorsShareOneCorrection(t *testing.T) {
	badQuote := strings.Replace(graphRestrictionTestBlock, "〔", "（", 1)
	bad := &types.GraphData{Relation: []*types.GraphRelation{{Node1: "甲品", Node2: "乙类", Type: "慎用于", Evidence: badQuote}}}
	model := &extractionResponseChat{response: graphCoverageResponse(t, bad)}
	extractor := NewGraphExtractor(model, &types.PromptTemplateStructured{Description: "Extract graph"})
	graph, err := extractor.Extract(context.Background(), graphRestrictionTestSource)
	if graph != nil || err == nil || model.calls != 2 || model.options.MaxTokens != 4096 || model.options.Temperature != 0 {
		t.Fatal("invalid response escaped bounded correction")
	}
	for _, message := range []string{err.Error(), model.messages[3].Content} {
		if !strings.Contains(message, "requires exact evidence") || !strings.Contains(message, "restriction strength mismatch") {
			t.Fatal("quote error hid independently anchored restriction strength")
		}
	}
	if bad.Relation[0].Evidence != badQuote || bad.Relation[0].Type != "慎用于" {
		t.Fatal("invalid model output was silently repaired")
	}
	if !strings.Contains(model.messages[0].Content, "A heading that governs a list supplies its predicate") ||
		strings.Contains(model.messages[0].Content, "section titles need not be included") {
		t.Fatal("persisted prompt still allows predicate headings to disappear")
	}
	if !strings.Contains(model.messages[0].Content, "literal OCR spelling in restriction targets, entity names and attributes") ||
		!strings.Contains(model.messages[0].Content, "Preserve supported usage procedures") ||
		!strings.Contains(model.messages[3].Content, "限制实体名称和属性也必须保留原文 OCR 字形") ||
		!strings.Contains(model.messages[3].Content, "保留有据的用量及使用步骤约束关系") {
		t.Fatal("correction omitted attribute grounding or usage preservation")
	}
}

func TestGraphRestrictionCorrectionMustExplicitlyPreserveCompleteSource(t *testing.T) {
	bad := &types.GraphData{Relation: []*types.GraphRelation{{Node1: "甲品", Node2: "乙类", Type: "慎用于", Evidence: strings.Replace(graphRestrictionTestBlock, "〔", "（", 1)}}}
	good := &types.GraphData{Relation: []*types.GraphRelation{{Node1: "甲品", Node2: graphRestrictionTestBlock, Type: "约束", Evidence: graphRestrictionTestBlock}}}
	model := &extractionResponseChat{responses: []*types.ChatResponse{graphCoverageResponse(t, bad), graphCoverageResponse(t, good)}}
	extractor := NewGraphExtractor(model, &types.PromptTemplateStructured{Description: "Extract graph"})
	graph, err := extractor.Extract(context.Background(), graphRestrictionTestSource)
	if err != nil || graph == nil || model.calls != 2 || len(graph.Relation) != 1 ||
		graph.Relation[0].Type != "约束" || graph.Relation[0].Node2 != graphRestrictionTestBlock || graph.Relation[0].Evidence != graphRestrictionTestBlock {
		t.Fatal("explicit complete source constraint was not preserved by bounded correction")
	}
	good.Relation = append(good.Relation, &types.GraphRelation{Node1: "甲品", Node2: "乙类", Type: "慎用于", Evidence: graphRestrictionTestBlock})
	response := graphCoverageResponse(t, good)
	source := graphRestrictionTestSource
	if accepted, err := NewFormater().parseGraph(context.Background(), response.Content, &source); accepted != nil ||
		!errors.Is(err, ErrGraphNeedsReview) || !strings.Contains(err.Error(), "restriction strength mismatch") {
		t.Fatal("full constraint concealed a contradictory typed edge")
	}
}

func TestGraphRestrictionQuotationMutationsStayRejected(t *testing.T) {
	const block = "〖 不适宜人群 〕 甲丿 IJ 类\n 、\t乙类"
	source := "本品\n" + block + "。"
	for _, quote := range []string{
		strings.Replace(block, "〖", "〔", 1),
		strings.Replace(block, "甲丿 IJ 类", "甲儿童类", 1),
		strings.ReplaceAll(block, " ", ""),
		strings.ReplaceAll(block, "\n", `\n`),
		strings.ReplaceAll(block, "\t", " "),
	} {
		graph := &types.GraphData{Relation: []*types.GraphRelation{{Node1: "本品", Node2: "乙类", Type: "不适用于", Evidence: quote}}}
		if err := validateGraphRelationEvidence(graph, source); err == nil || !strings.Contains(err.Error(), "requires exact evidence") {
			t.Fatal("mutated source quotation accepted")
		}
	}
}

func TestGraphRestrictionScopeCannotCrossProducts(t *testing.T) {
	const first = "〔不适宜人群〕甲类、乙类"
	const second = "〔慎用人群〕甲类、乙类"
	source := "产品名称：甲品\n" + first + "。\n\n产品名称：乙品\n" + second + "。"
	for _, tc := range []struct {
		subject, quote, label string
		wantErr               bool
	}{
		{"甲品", first, "不适用于", false},
		{"乙品", second, "慎用于", false},
		{"甲品", second, "慎用于", true},
		{"乙品", first, "不适用于", true},
		{"本品", "甲类", "慎用于", true},
		{"未出现产品", first, "不适用于", true},
		{"乙类", first, "不适用于", true},
	} {
		graph := &types.GraphData{Relation: []*types.GraphRelation{{Node1: tc.subject, Node2: "甲类", Type: tc.label, Evidence: tc.quote}}}
		err := validateGraphRelationEvidence(graph, source)
		if (err != nil) != tc.wantErr {
			t.Fatalf("subject=%s, label=%s, error=%v", tc.subject, tc.label, err)
		}
	}
}

func TestGraphRestrictionListsDoNotChangeQueryOrNarrativeExtraction(t *testing.T) {
	const source = "小说章节〔不适宜人群〕甲类、乙类。"
	if len(graphRestrictionScopes(source)) != 0 {
		t.Fatal("non-product narrative became a product restriction list")
	}
	model := &extractionResponseChat{response: &types.ChatResponse{Content: `[{"entity":"甲品"}]`, FinishReason: "stop"}}
	extractor := NewExtractor(model, &types.PromptTemplateStructured{Description: "Extract query entities"})
	graph, err := extractor.Extract(context.Background(), graphRestrictionTestSource)
	if err != nil || graph == nil || model.calls != 1 {
		t.Fatal("query extraction acquired persisted restriction coverage")
	}
}