ragserver/document_test.go at main · RichardKnop/ragserver · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
package ragserver

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestMatchSnippetsToDocuments(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name                      string
		snippets                  []string
		documents                 []Document
		expectedMatchedDocuments  []Document
		expectedUnmatchedSnippets []string
	}{
		{
			"no matches",
			[]string{
				"nonexistent1",
				"nonexistent2",
			},
			[]Document{
				{Content: "snippet1"},
				{Content: "snippet2"},
			},
			nil,
			[]string{
				"nonexistent1",
				"nonexistent2",
			},
		},
		{
			"all matching",
			[]string{
				"snippet1",
				"snippet2",
			},
			[]Document{
				{Content: "snippet1"},
				{Content: "snippet2"},
			},
			[]Document{
				{Content: "snippet1"},
				{Content: "snippet2"},
			},
			nil,
		},
		{
			"some matching",
			[]string{
				"snippet1",
				"nonexistent2",
			},
			[]Document{
				{Content: "snippet1"},
				{Content: "snippet2"},
			},
			[]Document{
				{Content: "snippet1"},
			},
			[]string{
				"nonexistent2",
			},
		},
		{
			"partial matches",
			[]string{
				"(7) We set our target using the International Energy Agency Net-Zero Emissions by 2050 scenario.",
				"nonexistent2",
				"Achieve net-zero GHG emissions by 2050, including operational emissions (Scope 1 and 2) and emissions attributable to our financing (Scope 3, Category 15).",
				"While the Company set sector- specific targets to enable it to track the alignment of its financing activities to its net-zero goal, these targets, even if met, do not guarantee reductions of absolute greenhouse gas emissions in the real economy.",
			},
			[]Document{
				{Content: "snippet1"},
				{Content: "(6) Metric tons of CO2 per metric ton of steel (7) We set our target using the International Energy Agency Net-Zero Emissions by 2050 scenario."},
				{Content: "snippet2"},
				{Content: "• Achieve net-zero GHG emissions by 2050, including operational emissions (Scope 1 and 2) and emissions attributable to our financing (Scope 3, Category 15)."},
				{Content: "Also, while the Company set sector- specific targets to enable it to track the alignment of its financing activities to its net-zero goal, these targets, even if met, do not guarantee reductions of absolute greenhouse gas emissions in the real economy."},
			},
			[]Document{
				{Content: "(6) Metric tons of CO2 per metric ton of steel (7) We set our target using the International Energy Agency Net-Zero Emissions by 2050 scenario."},
				{Content: "• Achieve net-zero GHG emissions by 2050, including operational emissions (Scope 1 and 2) and emissions attributable to our financing (Scope 3, Category 15)."},
				{Content: "Also, while the Company set sector- specific targets to enable it to track the alignment of its financing activities to its net-zero goal, these targets, even if met, do not guarantee reductions of absolute greenhouse gas emissions in the real economy."},
			},
			[]string{
				"nonexistent2",
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			matchedDocuments, unmatchedSnippets := MatchSnippetsToDocuments(tc.snippets, tc.documents)
			assert.Equal(t, tc.expectedMatchedDocuments, matchedDocuments)
			assert.Equal(t, tc.expectedUnmatchedSnippets, unmatchedSnippets)
		})
	}
}