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
|
package ai.vespa.search.llm;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.yahoo.search.Result;
import com.yahoo.search.rendering.EventRenderer;
import com.yahoo.search.searchchain.Execution;
import com.yahoo.text.Utf8;
import org.junit.jupiter.api.Test;
import java.io.ByteArrayOutputStream;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
public class RAGWithEventRendererTest {
@Test
public void testPromptAndHitsAreRendered() throws Exception {
var params = Map.of(
"query", "why are ducks better than cats?",
"llm.stream", "false",
"llm.includePrompt", "true",
"llm.includeHits", "true"
);
var llm = LLMSearcherTest.createLLMClient();
var searcher = RAGSearcherTest.createRAGSearcher(Map.of("mock", llm));
var results = RAGSearcherTest.runMockSearch(searcher, params);
var result = render(results);
var promptEvent = extractEvent(result, "prompt");
assertNotNull(promptEvent);
assertTrue(promptEvent.has("prompt"));
var resultsEvent = extractEvent(result, "hits");
assertNotNull(resultsEvent);
assertTrue(resultsEvent.has("root"));
assertEquals(2, resultsEvent.get("root").get("children").size());
}
private JsonNode extractEvent(String result, String eventName) throws JsonProcessingException {
var lines = result.split("\n");
for (int i = 0; i < lines.length; i++) {
if (lines[i].startsWith("event: " + eventName)) {
var data = lines[i + 1].substring("data: ".length()).trim();
ObjectMapper objectMapper = new ObjectMapper();
return objectMapper.readTree(data);
}
}
return null;
}
private String render(Result r) throws InterruptedException, ExecutionException {
var execution = new Execution(Execution.Context.createContextStub());
return render(execution, r);
}
private String render(Execution execution, Result r) throws ExecutionException, InterruptedException {
var renderer = new EventRenderer();
try {
renderer.init();
ByteArrayOutputStream bs = new ByteArrayOutputStream();
CompletableFuture<Boolean> f = renderer.renderResponse(bs, r, execution, null);
assertTrue(f.get());
return Utf8.toString(bs.toByteArray());
} finally {
renderer.deconstruct();
}
}
}
|
