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
|
// 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.InferenceParameters;
import ai.vespa.llm.LanguageModel;
import ai.vespa.llm.LanguageModelException;
import ai.vespa.llm.LlmSearcherConfig;
import ai.vespa.llm.completion.Completion;
import ai.vespa.llm.completion.Prompt;
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 java.util.List;
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 abstract 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
LLMSearcher(LlmSearcherConfig config, ComponentRegistry<LanguageModel> languageModels) {
this.stream = config.stream();
this.languageModelId = config.providerId();
this.languageModel = findLanguageModel(languageModelId, languageModels);
this.propertyPrefix = config.propertyPrefix();
}
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;
}
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
return stream ? completeAsync(query, prompt, options) : completeSync(query, prompt, options);
}
private Result completeAsync(Query query, Prompt prompt, InferenceParameters options) {
EventStream eventStream = new EventStream();
if (query.getTrace().getLevel() >= 1) {
eventStream.add(prompt.asString(), "prompt");
}
languageModel.completeAsync(prompt, options, token -> {
eventStream.add(token.text());
}).exceptionally(exception -> {
int errorCode = 400;
if (exception instanceof LanguageModelException languageModelException) {
errorCode = languageModelException.code();
}
eventStream.error(languageModelId, new ErrorMessage(errorCode, exception.getMessage()));
eventStream.markComplete();
return Completion.FinishReason.error;
}).thenAccept(finishReason -> {
eventStream.markComplete();
});
HitGroup hitGroup = new HitGroup("token_stream");
hitGroup.add(eventStream);
return new Result(query, hitGroup);
}
private Result completeSync(Query query, Prompt prompt, InferenceParameters options) {
EventStream eventStream = new EventStream();
if (query.getTrace().getLevel() >= 1) {
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);
}
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'.");
}
String getPropertyPrefix() {
return this.propertyPrefix;
}
String lookupProperty(String property, Query query) {
String propertyWithPrefix = this.propertyPrefix + "." + property;
return query.properties().getString(propertyWithPrefix, null);
}
Boolean lookupPropertyBool(String property, Query query, boolean defaultValue) {
String propertyWithPrefix = this.propertyPrefix + "." + property;
return query.properties().getBoolean(propertyWithPrefix, defaultValue);
}
String lookupPropertyWithOrWithoutPrefix(String property, Function<String, String> lookup) {
String value = lookup.apply(getPropertyPrefix() + "." + property);
if (value != null)
return value;
return lookup.apply(property);
}
String getApiKeyHeader(Query query) {
return lookupPropertyWithOrWithoutPrefix(API_KEY_HEADER, p -> query.getHttpRequest().getHeader(p));
}
}
|
