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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
|
// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package ai.vespa.llm.search;
import ai.vespa.llm.LanguageModel;
import ai.vespa.llm.LlmSearcherConfig;
import com.yahoo.component.ComponentId;
import com.yahoo.component.chain.Chain;
import com.yahoo.component.provider.ComponentRegistry;
import com.yahoo.search.Query;
import com.yahoo.search.Result;
import com.yahoo.search.Searcher;
import com.yahoo.search.result.EventStream;
import com.yahoo.search.result.Hit;
import com.yahoo.search.searchchain.Execution;
import org.junit.jupiter.api.Test;
import java.util.Map;
import static org.junit.jupiter.api.Assertions.assertEquals;
public class RAGSearcherTest {
private static final String DOC1_TITLE = "Exploring the Delightful Qualities of Ducks";
private static final String DOC1_CONTENT = "Ducks, with their gentle quacks and adorable waddling walks, possess a unique " +
"charm that sets them apart as extraordinary pets.";
private static final String DOC2_TITLE = "Why Cats Reign Supreme";
private static final String DOC2_CONTENT = "Cats bring an enchanting allure to households with their independent " +
"companionship, playful nature, natural hunting abilities, low-maintenance grooming, and the " +
"emotional support they offer.";
@Test
public void testRAGGeneration() {
var eventStream = runRAGQuery(Map.of(
"prompt", "why are ducks better than cats?",
"traceLevel", "1"));
var events = eventStream.incoming().drain();
assertEquals(2, events.size());
// Generated prompt
var promptEvent = (EventStream.Event) events.get(0);
assertEquals("prompt", promptEvent.type());
assertEquals("title: " + DOC1_TITLE + "\n" +
"content: " + DOC1_CONTENT + "\n\n" +
"title: " + DOC2_TITLE + "\n" +
"content: " + DOC2_CONTENT + "\n\n\n" +
"why are ducks better than cats?", promptEvent.toString());
// Generated completion
var completionEvent = (EventStream.Event) events.get(1);
assertEquals("completion", completionEvent.type());
assertEquals("Ducks have adorable waddling walks.", completionEvent.toString());
}
@Test
public void testPromptGeneration() {
var eventStream = runRAGQuery(Map.of(
"query", "why are ducks better than cats?",
"prompt", "{context}\nGiven these documents, answer this query as concisely as possible: @query",
"traceLevel", "1"));
var events = eventStream.incoming().drain();
var promptEvent = (EventStream.Event) events.get(0);
assertEquals("prompt", promptEvent.type());
assertEquals("title: " + DOC1_TITLE + "\n" +
"content: " + DOC1_CONTENT + "\n\n" +
"title: " + DOC2_TITLE + "\n" +
"content: " + DOC2_CONTENT + "\n\n\n" +
"Given these documents, answer this query as concisely as possible: " +
"why are ducks better than cats?", promptEvent.toString());
}
@Test
public void testSkipContextInPrompt() {
var eventStream = runRAGQuery(Map.of(
"query", "why are ducks better than cats?",
"llm.context", "skip",
"traceLevel", "1"));
var events = eventStream.incoming().drain();
var promptEvent = (EventStream.Event) events.get(0);
assertEquals("prompt", promptEvent.type());
assertEquals("why are ducks better than cats?", promptEvent.toString());
}
public static class MockSearchResults extends Searcher {
@Override
public Result search(Query query, Execution execution) {
Hit hit1 = new Hit("1");
hit1.setField("title", DOC1_TITLE);
hit1.setField("content", DOC1_CONTENT);
Hit hit2 = new Hit("2");
hit2.setField("title", DOC2_TITLE);
hit2.setField("content", DOC2_CONTENT);
Result r = new Result(query);
r.hits().add(hit1);
r.hits().add(hit2);
return r;
}
}
private EventStream runRAGQuery(Map<String, String> params) {
var llm = LLMSearcherTest.createLLMClient();
var searcher = createRAGSearcher(Map.of("mock", llm));
var result = runMockSearch(searcher, params);
return (EventStream) result.hits().get(0);
}
static Result runMockSearch(Searcher searcher, Map<String, String> parameters) {
Chain<Searcher> chain = new Chain<>(searcher, new MockSearchResults());
Execution execution = new Execution(chain, Execution.Context.createContextStub());
Query query = new Query("?" + LLMSearcherTest.toUrlParams(parameters));
return execution.search(query);
}
private static Searcher createRAGSearcher(Map<String, LanguageModel> llms) {
var config = new LlmSearcherConfig.Builder().stream(false).build();
ComponentRegistry<LanguageModel> models = new ComponentRegistry<>();
llms.forEach((key, value) -> models.register(ComponentId.fromString(key), value));
models.freeze();
return new RAGSearcher(config, models);
}
}
|
