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

package chatpipeline

import (
	"context"
	"encoding/json"
	"strings"
	"testing"

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

const graphScopeSource = "不得用于预防疾病。\n〔 贮存方法 〗 避光 、 密封保存 。\n\n## 工作表1\n\n| 品牌名称 | 产品名称 | 终端零售价 |\n| --- | --- | --- |\n| 甲品牌 | 乙产品 | 139元/一瓶 |\n"

func TestGraphProductSourceSectionsReconstructExactInput(t *testing.T) {
	for _, source := range []string{graphScopeSource, "前段。\n\n产品名称：甲\n正文\n\n产品名称：乙\n另段", "1丨甲\n内容\n2丨乙\n内容"} {
		var plan struct {
			Sections []struct {
				Number int    `json:"section"`
				Text   string `json:"text"`
			} `json:"source_sections"`
		}
		if err := json.Unmarshal([]byte(graphProductSourcePrompt(source)), &plan); err != nil {
			t.Fatal(err)
		}
		var reconstructed strings.Builder
		for i, section := range plan.Sections {
			if section.Number != i+1 {
				t.Fatal("nonsequential scope")
			}
			reconstructed.WriteString(section.Text)
		}
		if len(plan.Sections) < 2 || reconstructed.String() != source {
			t.Fatal("source changed or omitted")
		}
	}
	const ordinary = "Alice knows Bob."
	if graphProductSourcePrompt(ordinary) != ordinary {
		t.Fatal("ordinary source was rewritten")
	}
}

func TestGraphScopedCorrectionKeepsWarningsStorageAndPrices(t *testing.T) {
	response := []map[string]any{
		{"entity1": "乙产品", "entity2": "预防疾病", "relation": "不适用于", "evidence": "不得用于预防疾病"},
		{"entity1": "乙产品", "entity2": "避光密封保存", "relation": "约束", "evidence": "避光 、 密封保存"},
		{"entity": "甲品牌乙产品", "entity_attributes": []string{"终端零售价：139元/一瓶"}},
	}
	bad, _ := json.Marshal(response)
	response[0]["entity1"], response[1]["entity1"] = "本品", "本品"
	good, _ := json.Marshal(response)
	model := &extractionResponseChat{responses: []*types.ChatResponse{
		{Content: string(bad), FinishReason: "stop"}, {Content: string(good), FinishReason: "stop"},
	}}
	extractor := NewGraphExtractor(model, &types.PromptTemplateStructured{Description: "Extract graph"})
	graph, err := extractor.Extract(context.Background(), graphScopeSource)
	if err != nil || graph == nil || len(graph.Relation) != 2 || model.calls != 2 {
		t.Fatalf("bounded correction failed: %v", err)
	}
	if !strings.Contains(model.messages[1].Content, `"source_sections"`) ||
		!strings.Contains(model.messages[3].Content, "evidence is in source sections [1]") ||
		!strings.Contains(model.messages[3].Content, "proposed subject occurs in sections [3]") {
		t.Fatal("correction did not identify the conflicting source scopes")
	}
	for _, relation := range graph.Relation {
		if relation.Node1 != "本品" || !strings.Contains(graphScopeSource, relation.Evidence) {
			t.Fatal("scope/evidence changed")
		}
	}
	if graph.Node[len(graph.Node)-1] == nil {
		t.Fatal("lost graph node")
	}
	priceFound := false
	for _, node := range graph.Node {
		for _, attribute := range node.Attributes {
			priceFound = priceFound || attribute == "终端零售价：139元/一瓶"
		}
	}
	if !priceFound {
		t.Fatal("correction dropped the supported table price")
	}
}

func TestGraphSourceEnvelopeCannotBecomeEvidence(t *testing.T) {
	model := &extractionResponseChat{response: &types.ChatResponse{Content: `[{"entity1":"本品","entity2":"预防疾病","relation":"不适用于","evidence":"source_sections"}]`, FinishReason: "stop"}}
	extractor := NewGraphExtractor(model, &types.PromptTemplateStructured{Description: "Extract graph"})
	if graph, err := extractor.Extract(context.Background(), graphScopeSource); err == nil || graph != nil || model.calls != 2 {
		t.Fatal("metadata was accepted as evidence")
	}
}

func TestGraphSeriesHeadingCannotPromoteLocalWarnings(t *testing.T) {
	for _, title := range []string{"## 甲乙套盒系列介绍", "甲乙套盒系列"} {
		source := title + "\n\n过敏者慎甩；粘膜破损处禁用。\n本品不是药品。"
		bad := &types.GraphData{Relation: []*types.GraphRelation{{Node1: "甲乙套盒系列", Node2: "过敏者", Type: "慎用于", Evidence: "过敏者慎甩"}}}
		if err := validateGraphRelationEvidence(bad, source); err == nil || !strings.Contains(err.Error(), "series heading") {
			t.Fatal("series title used as local product")
		}
		bad.Relation[0].Node1 = "本品"
		if err := validateGraphRelationEvidence(bad, source); err != nil {
			t.Fatal(err)
		}
	}
	const explicit = "## 甲乙套盒系列\n\n甲乙套盒系列过敏者慎用。"
	graph := &types.GraphData{Relation: []*types.GraphRelation{{Node1: "甲乙套盒系列", Node2: "过敏者", Type: "慎用于", Evidence: "甲乙套盒系列过敏者慎用"}}}
	if err := validateGraphRelationEvidence(graph, explicit); err != nil {
		t.Fatal("explicit series rule rejected", err)
	}
}