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

package chatpipeline

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

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

func TestGraphPriceAttributesRequireExplicitProductEvidence(t *testing.T) {
	const unheaded = "| 私觅玉丹 | 紧紧丹 | 0.5g/7粒 | 39元/盒 | 34元/一盒 99盒/箱 | 199元/一盒 399元/三盒 |\n\n| 私觅玉丹 | 水光胶原棒 | 3g/6支 | 39元/盒 | 34元/盒 99盒/箱 | 199元/一盒 399元/三盒 |"
	const table = "| 产品名称 | 零售价 | 批发价 |\n| --- | --- | --- |\n| 甲产品 | 39元/盒 | 34元/盒 |\n| 乙产品 | 59元/盒 | 54元/盒 |"
	for _, tc := range []struct {
		name, source, product, attribute string
		valid                            bool
	}{
		{name: "probe cannot infer retail column", source: unheaded, product: "紧紧丹", attribute: "零售价：39元/盒"},
		{name: "probe cannot infer wholesale column", source: unheaded, product: "紧紧丹", attribute: "批发价：34元/一盒，99盒/箱"},
		{name: "probe cannot infer promotion column", source: unheaded, product: "水光胶原棒", attribute: "促销价：199元/一盒，399元/三盒"},
		{name: "raw source amount does not invent price category", source: unheaded, product: "紧紧丹", attribute: "39元/盒", valid: true},
		{name: "non price summaries remain supported", source: "甲产品采用简便包装，每盒7粒。", product: "甲产品", attribute: "包装便于使用；规格：7粒/盒", valid: true},
		{name: "unresolved price without amount is not a price assertion", source: "甲产品每盒7粒，价格待确认。", product: "甲产品", attribute: "规格：7粒/盒；价格待确认", valid: true},
		{name: "explicit sentence", source: "甲产品的零售价为39元/盒。", product: "甲产品", attribute: "零售价：39元/盒", valid: true},
		{name: "sentence punctuation and spaces may differ", source: "甲产品零售价格：39 元 / 盒，批发价：34元/盒。", product: "甲产品", attribute: "零售价39元/盒；批发价格34元/盒", valid: true},
		{name: "label cannot borrow other labelled value", source: "甲产品零售价39元/盒，批发价34元/盒。", product: "甲产品", attribute: "零售价：34元/盒"},
		{name: "cannot borrow another product sentence", source: "甲产品零售价39元/盒。乙产品零售价59元/盒。", product: "甲产品", attribute: "零售价：59元/盒"},
		{name: "cannot borrow another product in same sentence", source: "甲产品零售价39元/盒，乙产品零售价59元/盒。", product: "甲产品", attribute: "零售价：59元/盒"},
		{name: "cannot match product name substring", source: "甲产品加强版零售价39元/盒。", product: "甲产品", attribute: "零售价：39元/盒"},
		{name: "different packaging unit", source: "甲产品零售价39元/箱。", product: "甲产品", attribute: "零售价：39元/盒"},
		{name: "different amount cannot match numeric substring", source: "甲产品零售价399元/盒。", product: "甲产品", attribute: "零售价：39元/盒"},
		{name: "promotion pack amounts preserved", source: "甲产品促销价199元/一盒，399元/三盒。", product: "甲产品", attribute: "促销价：199元/一盒，399元/三盒", valid: true},
		{name: "promotion pack counts cannot change", source: "甲产品促销价199元/一盒，399元/三盒。", product: "甲产品", attribute: "促销价：399元/一盒"},
		{name: "unparsed packaging count cannot be ignored", source: "甲产品批发价34元/盒，99盒/箱。", product: "甲产品", attribute: "批发价34元/盒（199盒/箱）"},
		{name: "table row and header", source: table, product: "甲产品", attribute: "零售价：39元/盒；批发价：34元/盒", valid: true},
		{name: "table cannot swap columns", source: table, product: "甲产品", attribute: "零售价：34元/盒"},
		{name: "table cannot borrow other row", source: table, product: "甲产品", attribute: "零售价：59元/盒"},
		{name: "table cannot use shared brand as product", source: "| 品牌 | 产品名称 | 零售价 |\n| --- | --- | --- |\n| 甲 | 甲精华 | 39元/盒 |", product: "甲", attribute: "零售价：39元/盒"},
		{name: "table header supplies explicit unit", source: "| 产品名称 | 零售价（元/盒） |\n| --- | --- |\n| 甲产品 | 39 |", product: "甲产品", attribute: "零售价：39元/盒", valid: true},
		{name: "table missing unit cannot invent yuan", source: "| 产品名称 | 零售价 |\n| --- | --- |\n| 甲产品 | 39 |", product: "甲产品", attribute: "零售价：39元/盒"},
		{name: "header cannot bleed into another table", source: table + "\n\n| 丙产品 | 69元/盒 | 64元/盒 |", product: "丙产品", attribute: "零售价：69元/盒"},
		{name: "headerless row with explicit cell label", source: "| 甲产品 | 零售价：39元/盒 |", product: "甲产品", attribute: "零售价：39元/盒", valid: true},
		{name: "named section", source: "〔产品名称〕甲产品\n零售价：39元/盒。\n\n〔产品名称〕乙产品\n零售价：59元/盒。", product: "甲产品", attribute: "零售价：39元/盒", valid: true},
		{name: "sections cannot borrow price", source: "〔产品名称〕甲产品\n零售价：39元/盒。\n\n〔产品名称〕乙产品\n零售价：59元/盒。", product: "甲产品", attribute: "零售价：59元/盒"},
		{name: "numeric unsupported phrasing cannot bypass validation", source: "甲产品39元/盒。", product: "甲产品", attribute: "零售价约39元/盒"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			graph := &types.GraphData{Node: []*types.GraphNode{{Name: tc.product, Attributes: []string{tc.attribute}}}}
			err := validateGraphPriceAttributes(graph, tc.source)
			if tc.valid && err != nil {
				t.Fatal(err)
			}
			if !tc.valid && !errors.Is(err, ErrGraphNeedsReview) {
				t.Fatalf("ambiguous price must require review: %v", err)
			}
			if !reflect.DeepEqual(graph.Node[0].Attributes, []string{tc.attribute}) {
				t.Fatal("validation silently rewrote or removed an attribute")
			}
		})
	}
}

func TestGraphPriceTableFormattingPreservesProductAndPackageScope(t *testing.T) {
	const headers = "| 品牌名称 | 产品名称 | 图片 | 规格 | 终端零售价 | 产品说明 | 其他 |\n| --- | --- | --- | --- | --- | --- | --- |\n"
	const truncatedRow = "| 甲品牌 | 甲精华 | | 30ml/1瓶 | 139元/一瓶 239元/俩瓶 | 原文说明"
	const fullRow = truncatedRow + " | 补充说明 |"
	const otherRow = "| 乙品牌 | 乙精华 | | 30ml/1瓶 | 159元/一瓶 259元/俩瓶 | 原文说明 | 补充说明 |"
	for _, tc := range []struct {
		name, source, product, attribute string
		valid                            bool
	}{
		{name: "complete row with same-row brand and product", source: headers + fullRow, product: "甲品牌 甲精华", attribute: "终端零售价：139元/一瓶 239元/俩瓶", valid: true},
		{name: "trailing pipe is optional", source: headers + strings.TrimSuffix(fullRow, "|"), product: "甲精华", attribute: "终端零售价：239元/俩瓶", valid: true},
		{name: "header trailing pipe is optional", source: strings.ReplaceAll(headers, "|\n", "\n") + fullRow, product: "甲精华", attribute: "终端零售价：239元/俩瓶", valid: true},
		{name: "available columns in truncated row", source: headers + truncatedRow, product: "甲品牌甲精华", attribute: "终端零售价：139元/一瓶 239元/俩瓶", valid: true},
		{name: "available columns with truncated row delimiter", source: headers + truncatedRow + " |", product: "甲精华", attribute: "零售价：239元/俩瓶", valid: true},
		{name: "terminal retail sentence", source: "甲精华终端零售价格139元/一瓶，239元/俩瓶。", product: "甲精华", attribute: "终端零售价：139元/一瓶，239元/俩瓶", valid: true},
		{name: "terminal retail is explicit retail category", source: headers + fullRow, product: "甲精华", attribute: "零售价：239元/俩瓶", valid: true},
		{name: "brand alone cannot own the price", source: headers + fullRow, product: "甲品牌", attribute: "终端零售价：239元/俩瓶"},
		{name: "product substring cannot own the price", source: headers + fullRow, product: "甲精", attribute: "终端零售价：239元/俩瓶"},
		{name: "brand substring cannot own the price", source: headers + fullRow, product: "甲牌甲精华", attribute: "终端零售价：239元/俩瓶"},
		{name: "brand and product must come from same row", source: headers + fullRow + "\n" + otherRow, product: "乙品牌甲精华", attribute: "终端零售价：239元/俩瓶"},
		{name: "same product under another brand cannot supply price", source: headers + fullRow + "\n" + strings.Replace(otherRow, "乙精华", "甲精华", 1), product: "甲品牌甲精华", attribute: "终端零售价：259元/俩瓶"},
		{name: "another row cannot supply price", source: headers + truncatedRow + "\n" + otherRow, product: "甲品牌甲精华", attribute: "终端零售价：259元/俩瓶"},
		{name: "two bottle price cannot become single bottle", source: headers + truncatedRow, product: "甲精华", attribute: "终端零售价：239元/一瓶"},
		{name: "two bottle quantity cannot be dropped", source: headers + truncatedRow, product: "甲精华", attribute: "终端零售价：239元/瓶"},
		{name: "two bottle unit cannot become box", source: headers + truncatedRow, product: "甲精华", attribute: "终端零售价：239元/俩盒"},
		{name: "two bottle token cannot be silently changed", source: headers + truncatedRow, product: "甲精华", attribute: "终端零售价：239元/两瓶"},
		{name: "unclosed price in a short row cannot be accepted", source: headers + "| 甲品牌 | 甲精华 | | 30ml/1瓶 | 239元", product: "甲精华", attribute: "终端零售价：239元"},
		{name: "truncated package cannot be ignored", source: "甲精华终端零售价239元/俩", product: "甲精华", attribute: "终端零售价：239元"},
		{name: "unknown package count cannot be ignored", source: "甲精华终端零售价239元/仨瓶", product: "甲精华", attribute: "终端零售价：239元"},
		{name: "terminal retail cannot become suggested retail", source: headers + truncatedRow, product: "甲精华", attribute: "建议零售价：239元/俩瓶"},
		{name: "terminal retail cannot become wholesale", source: headers + truncatedRow, product: "甲精华", attribute: "批发价：239元/俩瓶"},
		{name: "unheaded amounts still cannot infer retail", source: truncatedRow, product: "甲精华", attribute: "终端零售价：239元/俩瓶"},
		{name: "missing price column cannot borrow other cells", source: headers + "| 甲品牌 | 甲精华 | 239元/俩瓶", product: "甲品牌甲精华", attribute: "终端零售价：239元/俩瓶"},
		{name: "missing product column cannot borrow heading", source: "| 品牌名称 | 终端零售价 | 产品名称 |\n| --- | --- | --- |\n| 甲品牌 | 239元/俩瓶", product: "甲品牌甲精华", attribute: "终端零售价：239元/俩瓶"},
		{name: "blank product cannot match brand", source: headers + strings.Replace(fullRow, "甲精华", "", 1), product: "甲品牌", attribute: "终端零售价：239元/俩瓶"},
		{name: "overlong row is ambiguous", source: headers + fullRow + " 多余单元格 |", product: "甲品牌甲精华", attribute: "终端零售价：239元/俩瓶"},
		{name: "header cannot survive intervening prose", source: headers + "表格结束\n" + truncatedRow, product: "甲精华", attribute: "终端零售价：239元/俩瓶"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			graph := &types.GraphData{Node: []*types.GraphNode{{Name: tc.product, Attributes: []string{tc.attribute}}}}
			err := validateGraphPriceAttributes(graph, tc.source)
			if tc.valid && err != nil {
				t.Fatal(err)
			}
			if !tc.valid && !errors.Is(err, ErrGraphNeedsReview) {
				t.Fatalf("unsupported product, column or package was accepted: %v", err)
			}
			if !reflect.DeepEqual(graph.Node[0].Attributes, []string{tc.attribute}) {
				t.Fatal("validation rewrote the source-supported attribute")
			}
		})
	}
}

func graphPriceTestResponse(t *testing.T, attribute string) *types.ChatResponse {
	t.Helper()
	items := []map[string]interface{}{
		{"entity": "甲产品", "entity_attributes": []string{attribute}},
		{"entity1": "甲产品", "entity2": "成分甲", "relation": "包含", "evidence": "甲产品包含成分甲"},
	}
	raw, err := json.Marshal(items)
	if err != nil {
		t.Fatal(err)
	}
	return &types.ChatResponse{Content: string(raw), FinishReason: "stop"}
}

func TestGraphPriceAmbiguityRequiresReviewAfterBoundedCorrection(t *testing.T) {
	const source = "甲产品包含成分甲。\n| 甲产品 | 39元/盒 | 34元/盒 |"
	model := &extractionResponseChat{response: graphPriceTestResponse(t, "零售价：39元/盒")}
	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("ambiguous attribute completed or retried indefinitely: graph=%+v, err=%v, calls=%d", graph, err, model.calls)
	}
	if !strings.Contains(model.messages[len(model.messages)-1].Content, "unsupported price label/value") {
		t.Fatal("bounded correction did not receive the price attribution failure")
	}
	if !strings.Contains(model.messages[0].Content, "无表头的金额列不能推断") {
		t.Fatal("graph instruction must cover headerless sources without product markers")
	}
}

func TestGraphPriceCorrectionCanRestoreExplicitPrice(t *testing.T) {
	const source = "甲产品包含成分甲。甲产品零售价39元/盒，批发价34元/盒。"
	model := &extractionResponseChat{responses: []*types.ChatResponse{
		graphPriceTestResponse(t, "零售价：34元/盒"),
		graphPriceTestResponse(t, "零售价：39元/盒"),
	}}
	extractor := NewGraphExtractor(model, &types.PromptTemplateStructured{Description: "Extract graph"})
	graph, err := extractor.Extract(context.Background(), source)
	if err != nil || graph == nil || len(graph.Relation) != 1 || model.calls != 2 {
		t.Fatalf("supported price correction failed: graph=%+v, err=%v, calls=%d", graph, err, model.calls)
	}
	if !reflect.DeepEqual(graph.Node[0].Attributes, []string{"零售价：39元/盒"}) {
		t.Fatal("corrected supported price was not preserved")
	}
}

func TestQueryExtractionDoesNotRequirePriceProvenance(t *testing.T) {
	model := &extractionResponseChat{response: graphPriceTestResponse(t, "零售价：39元/盒")}
	extractor := NewExtractor(model, &types.PromptTemplateStructured{Description: "Extract query entities"})
	graph, err := extractor.Extract(context.Background(), "查询甲产品零售价")
	if err != nil || graph == nil || model.calls != 1 {
		t.Fatalf("graph-only price validation changed query extraction: %+v, %v, calls=%d", graph, err, model.calls)
	}
}