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
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
|
// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package ai.vespa.search.llm;
import ai.vespa.llm.InferenceParameters;
import ai.vespa.llm.LanguageModel;
import ai.vespa.llm.LanguageModelException;
import ai.vespa.llm.completion.Completion;
import ai.vespa.llm.completion.Prompt;
import ai.vespa.llm.completion.StringPrompt;
import com.yahoo.api.annotations.Beta;
import com.yahoo.component.ComponentId;
import com.yahoo.component.annotation.Inject;
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.ErrorMessage;
import com.yahoo.search.result.EventStream;
import com.yahoo.search.result.HitGroup;
import com.yahoo.search.searchchain.Execution;
import java.util.List;
import java.util.concurrent.RejectedExecutionException;
import java.util.function.Function;
import java.util.logging.Logger;
import java.util.stream.Collectors;
/**
* Base class for LLM searchers. Provides utilities for calling LLMs and handling properties.
*
* @author lesters
*/
@Beta
public class LLMSearcher extends Searcher {
private static Logger log = Logger.getLogger(LLMSearcher.class.getName());
private static final String API_KEY_HEADER = "X-LLM-API-KEY";
private static final String STREAM_PROPERTY = "stream";
private static final String PROMPT_PROPERTY = "prompt";
private final String propertyPrefix;
private final boolean stream;
private final LanguageModel languageModel;
private final String languageModelId;
@Inject
public LLMSearcher(LlmSearcherConfig config, ComponentRegistry<LanguageModel> languageModels) {
this.stream = config.stream();
this.languageModelId = config.providerId();
this.languageModel = findLanguageModel(languageModelId, languageModels);
this.propertyPrefix = config.propertyPrefix();
}
@Override
public Result search(Query query, Execution execution) {
return complete(query, StringPrompt.from(getPrompt(query)));
}
private LanguageModel findLanguageModel(String providerId, ComponentRegistry<LanguageModel> languageModels)
throws IllegalArgumentException
{
if (languageModels.allComponents().isEmpty()) {
throw new IllegalArgumentException("No language models were found");
}
if (providerId == null || providerId.isEmpty()) {
var entry = languageModels.allComponentsById().entrySet().stream().findFirst();
if (entry.isEmpty()) {
throw new IllegalArgumentException("No language models were found"); // shouldn't happen given check above
}
log.info("Language model provider was not found in config. " +
"Fallback to using first available language model: " + entry.get().getKey());
return entry.get().getValue();
}
final LanguageModel languageModel = languageModels.getComponent(providerId);
if (languageModel == null) {
throw new IllegalArgumentException("No component with id '" + providerId + "' was found. " +
"Available LLM components are: " + languageModels.allComponentsById().keySet().stream()
.map(ComponentId::toString).collect(Collectors.joining(",")));
}
return languageModel;
}
protected Result complete(Query query, Prompt prompt) {
var options = new InferenceParameters(getApiKeyHeader(query), s -> lookupProperty(s, query));
var stream = lookupPropertyBool(STREAM_PROPERTY, query, this.stream); // query value overwrites config
try {
return stream ? completeAsync(query, prompt, options) : completeSync(query, prompt, options);
} catch (RejectedExecutionException e) {
return new Result(query, new ErrorMessage(429, e.getMessage()));
}
}
private boolean shouldAddPrompt(Query query) {
return query.getTrace().getLevel() >= 1;
}
private boolean shouldAddTokenStats(Query query) {
return query.getTrace().getLevel() >= 1;
}
private Result completeAsync(Query query, Prompt prompt, InferenceParameters options) {
final EventStream eventStream = new EventStream();
if (shouldAddPrompt(query)) {
eventStream.add(prompt.asString(), "prompt");
}
final TokenStats tokenStats = new TokenStats();
languageModel.completeAsync(prompt, options, completion -> {
tokenStats.onToken();
handleCompletion(eventStream, completion);
}).exceptionally(exception -> {
handleException(eventStream, exception);
eventStream.markComplete();
return Completion.FinishReason.error;
}).thenAccept(finishReason -> {
tokenStats.onCompletion();
if (shouldAddTokenStats(query)) {
eventStream.add(tokenStats.report(), "stats");
}
eventStream.markComplete();
});
HitGroup hitGroup = new HitGroup("token_stream");
hitGroup.add(eventStream);
return new Result(query, hitGroup);
}
private void handleCompletion(EventStream eventStream, Completion completion) {
if (completion.finishReason() == Completion.FinishReason.error) {
eventStream.add(completion.text(), "error");
} else {
eventStream.add(completion.text());
}
}
private void handleException(EventStream eventStream, Throwable exception) {
int errorCode = 400;
if (exception instanceof LanguageModelException languageModelException) {
errorCode = languageModelException.code();
}
eventStream.error(languageModelId, new ErrorMessage(errorCode, exception.getMessage()));
}
private Result completeSync(Query query, Prompt prompt, InferenceParameters options) {
EventStream eventStream = new EventStream();
if (shouldAddPrompt(query)) {
eventStream.add(prompt.asString(), "prompt");
}
List<Completion> completions = languageModel.complete(prompt, options);
eventStream.add(completions.get(0).text(), "completion");
eventStream.markComplete();
HitGroup hitGroup = new HitGroup("completion");
hitGroup.add(eventStream);
return new Result(query, hitGroup);
}
public String getPrompt(Query query) {
// Look for prompt with or without prefix
String prompt = lookupPropertyWithOrWithoutPrefix(PROMPT_PROPERTY, p -> query.properties().getString(p));
if (prompt != null)
return prompt;
// If not found, use query directly
prompt = query.getModel().getQueryString();
if (prompt != null)
return prompt;
// If not, throw exception
throw new IllegalArgumentException("Could not find prompt found for query. Tried looking for " +
"'" + propertyPrefix + "." + PROMPT_PROPERTY + "', '" + PROMPT_PROPERTY + "' or '@query'.");
}
public String getPropertyPrefix() {
return this.propertyPrefix;
}
public String lookupProperty(String property, Query query) {
String propertyWithPrefix = this.propertyPrefix + "." + property;
return query.properties().getString(propertyWithPrefix, null);
}
public Boolean lookupPropertyBool(String property, Query query, boolean defaultValue) {
String propertyWithPrefix = this.propertyPrefix + "." + property;
return query.properties().getBoolean(propertyWithPrefix, defaultValue);
}
public String lookupPropertyWithOrWithoutPrefix(String property, Function<String, String> lookup) {
String value = lookup.apply(getPropertyPrefix() + "." + property);
if (value != null)
return value;
return lookup.apply(property);
}
public String getApiKeyHeader(Query query) {
return lookupPropertyWithOrWithoutPrefix(API_KEY_HEADER, p -> query.getHttpRequest().getHeader(p));
}
private static class TokenStats {
private long start;
private long timeToFirstToken;
private long timeToLastToken;
private long tokens = 0;
TokenStats() {
start = System.currentTimeMillis();
}
void onToken() {
if (tokens == 0) {
timeToFirstToken = System.currentTimeMillis() - start;
}
tokens++;
}
void onCompletion() {
timeToLastToken = System.currentTimeMillis() - start;
}
String report() {
return "Time to first token: " + timeToFirstToken + " ms, " +
"Generation time: " + timeToLastToken + " ms, " +
"Generated tokens: " + tokens + " " +
String.format("(%.2f tokens/sec)", tokens / (timeToLastToken / 1000.0));
}
}
}
|
