Skip to content

Commit 439c3bc

Browse files
jamesamcljamesamcl
committed
working model selection
1 parent c6d9ab3 commit 439c3bc

10 files changed

Lines changed: 173 additions & 34 deletions

File tree

‎backend/src/main/java/uk/ac/ebi/spot/ols/controller/api/v1/V1SearchController.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -127,8 +127,8 @@ public void search(
127127
@RequestParam(value = "lang", defaultValue = "en") String lang,
128128
@RequestParam(value = "model", required = false)
129129
@Parameter(name = "model",
130-
description = "Optional: Use vector similarity search with the specified embedding model (e.g., text-embedding-3-small). When provided, the query will be embedded and vector search will be performed.",
131-
example = "text-embedding-3-small") String model,
130+
description = "Optional: Use vector similarity search with the specified embedding model. When provided, the query will be embedded and vector search will be performed.",
131+
example = "llama-embed-megatron-8b") String model,
132132
HttpServletResponse response
133133
) throws IOException, SolrServerException {
134134

‎backend/src/main/java/uk/ac/ebi/spot/ols/controller/api/v2/V2EntityController.java‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,10 @@ public HttpEntity<V2PagedAndFacetedResponse<V2Entity>> getEntities(
7878
description = "Specify any other search field here which are not specified by searchFields or boostFields.",
7979
example = "{}") MultiValueMap<String,String> searchProperties,
8080
@RequestParam(value = "lang", required = false, defaultValue = "en") String lang,
81+
@RequestParam(value = "model", required = false)
82+
@Parameter(name = "model",
83+
description = "Optional: Use vector similarity search with the specified embedding model. When provided, the query will be embedded and vector search will be performed.",
84+
example = "llama-embed-megatron-8b") String model,
8185
JsonTransformOptions outputOpts
8286
) throws ResourceNotFoundException, IOException {
8387

@@ -88,7 +92,7 @@ public HttpEntity<V2PagedAndFacetedResponse<V2Entity>> getEntities(
8892

8993
return new ResponseEntity<>(
9094
new V2PagedAndFacetedResponse<V2Entity>(
91-
entityRepository.find(pageable, lang, search, searchFields, boostFields, facetFields, exactMatch, DynamicQueryHelper.filterProperties(properties), outputOpts) .map(V2Entity::new)
95+
entityRepository.find(pageable, lang, search, searchFields, boostFields, facetFields, exactMatch, DynamicQueryHelper.filterProperties(properties), model, outputOpts) .map(V2Entity::new)
9296
),
9397
HttpStatus.OK);
9498
}
@@ -136,6 +140,10 @@ public HttpEntity<V2PagedAndFacetedResponse<V2Entity>> getTerms(
136140
description = "Specify any other search field here which are not specified by searchFields or boostFields.",
137141
example = "{}") MultiValueMap<String,String> searchProperties,
138142
@RequestParam(value = "lang", required = false, defaultValue = "en") String lang,
143+
@RequestParam(value = "model", required = false)
144+
@Parameter(name = "model",
145+
description = "Optional: Use vector similarity search with the specified embedding model. When provided, the query will be embedded and vector search will be performed.",
146+
example = "llama-embed-megatron-8b") String model,
139147
JsonTransformOptions outputOpts
140148
) throws ResourceNotFoundException, IOException {
141149

@@ -146,7 +154,7 @@ public HttpEntity<V2PagedAndFacetedResponse<V2Entity>> getTerms(
146154

147155
return new ResponseEntity<>(
148156
new V2PagedAndFacetedResponse<V2Entity>(
149-
entityRepository.findByOntologyId(ontologyId, pageable, lang, search, searchFields, boostFields, facetFields, exactMatch, DynamicQueryHelper.filterProperties(properties), outputOpts).map(V2Entity::new)
157+
entityRepository.findByOntologyId(ontologyId, pageable, lang, search, searchFields, boostFields, facetFields, exactMatch, DynamicQueryHelper.filterProperties(properties), model, outputOpts).map(V2Entity::new)
150158
),
151159
HttpStatus.OK);
152160
}

‎backend/src/main/java/uk/ac/ebi/spot/ols/controller/api/v2/V2LLMController.java‎

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -61,13 +61,6 @@ public class V2LLMController {
6161
@Autowired
6262
OlsSolrClient solrClient;
6363

64-
@RequestMapping(path = "/models", produces = {MediaType.APPLICATION_JSON_VALUE }, method = RequestMethod.GET)
65-
public HttpEntity<Map<String, Object>> getAvailableModels() {
66-
Map<String, Object> response = new HashMap<>();
67-
response.put("models", embeddingServiceClient.getAvailableModels());
68-
return new ResponseEntity<>(response, HttpStatus.OK);
69-
}
70-
7164
@RequestMapping(path = "/llm_models", produces = {MediaType.APPLICATION_JSON_VALUE }, method = RequestMethod.GET)
7265
@Parameter(name = "llm_models",
7366
description = "Returns a list of embedding models, indicating which can be used for embedding (via the embedding service) and which only have pre-computed embeddings stored in Solr")

‎backend/src/main/java/uk/ac/ebi/spot/ols/controller/api/v2/helpers/DynamicQueryHelper.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,9 @@ public static Map<String,Collection<String>> filterProperties(Map<String,Collect
2020
|| k.equals("boostFields") || k.equals("page") || k.equals("size") || k.equals("exactMatch")
2121
|| k.equals("includeObsoleteEntities")
2222
|| k.equals("resolveReferences")
23-
|| k.equals("manchesterSyntax"))
23+
|| k.equals("manchesterSyntax")
24+
|| k.equals("model")
25+
)
2426
continue;
2527

2628
newProps.put(k, properties.get(k));

‎backend/src/main/java/uk/ac/ebi/spot/ols/controller/mcp/McpClassService.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ McpPage<McpClass> searchClasses(
6464
null,
6565
false,
6666
properties,
67+
null, // model
6768
outputOpts
6869
);
6970

‎backend/src/main/java/uk/ac/ebi/spot/ols/controller/mcp/McpSearchService.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ String search(
4747
null,
4848
false,
4949
Map.of(),
50+
null, // model
5051
outputOpts
5152
);
5253

‎backend/src/main/java/uk/ac/ebi/spot/ols/repository/EntityRepository.java‎

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import uk.ac.ebi.spot.ols.repository.solr.OlsSolrClient;
1919
import uk.ac.ebi.spot.ols.repository.helpers.DynamicFilterParser;
2020
import uk.ac.ebi.spot.ols.repository.helpers.SearchFieldsParser;
21+
import uk.ac.ebi.spot.ols.service.EmbeddingServiceClient;
2122

2223
import java.io.IOException;
2324
import java.util.Collection;
@@ -32,16 +33,30 @@ public class EntityRepository {
3233

3334
@Autowired
3435
OlsNeo4jClient neo4jClient;
36+
37+
@Autowired
38+
EmbeddingServiceClient embeddingServiceClient;
3539

3640

3741
public OlsFacetedResultsPage<JsonElement> find(
38-
Pageable pageable, String lang, String search, String searchFields, String boostFields, String facetFields, boolean exactMatch, Map<String, Collection<String>> properties, JsonTransformOptions outputOpts) throws IOException {
42+
Pageable pageable, String lang, String search, String searchFields, String boostFields, String facetFields, boolean exactMatch, Map<String, Collection<String>> properties, String model, JsonTransformOptions outputOpts) throws IOException {
3943

4044
Validation.validateLang(lang);
4145

4246
OlsSolrQuery query = new OlsSolrQuery();
43-
query.setSearchText(search);
44-
query.setExactMatch(exactMatch);
47+
48+
// Choose between vector search or text search based on model parameter
49+
if (model != null && !model.isEmpty()) {
50+
// Vector search: embed the query text
51+
float[] embeddings = embeddingServiceClient.embedText(model, search);
52+
query.setEmbeddingVector(embeddings, model);
53+
query.setTopK(pageable.getPageSize());
54+
} else {
55+
// Text search
56+
query.setSearchText(search);
57+
query.setExactMatch(exactMatch);
58+
}
59+
4560
query.addFilter("type", List.of("entity"), SearchType.WHOLE_FIELD);
4661
SearchFieldsParser.addSearchFieldsToQuery(query, searchFields);
4762
SearchFieldsParser.addBoostFieldsToQuery(query, boostFields);
@@ -54,14 +69,24 @@ public OlsFacetedResultsPage<JsonElement> find(
5469
}
5570

5671
public OlsFacetedResultsPage<JsonElement> findByOntologyId(
57-
String ontologyId, Pageable pageable, String lang, String search, String searchFields, String boostFields, String facetFields, boolean exactMatch, Map<String,Collection<String>> properties, JsonTransformOptions outputOpts) throws IOException {
72+
String ontologyId, Pageable pageable, String lang, String search, String searchFields, String boostFields, String facetFields, boolean exactMatch, Map<String,Collection<String>> properties, String model, JsonTransformOptions outputOpts) throws IOException {
5873

5974
Validation.validateOntologyId(ontologyId);
6075
Validation.validateLang(lang);
6176

6277
OlsSolrQuery query = new OlsSolrQuery();
63-
query.setSearchText(search);
64-
query.setExactMatch(exactMatch);
78+
79+
// Choose between vector search or text search based on model parameter
80+
if (model != null && !model.isEmpty()) {
81+
// Vector search: embed the query text
82+
float[] embeddings = embeddingServiceClient.embedText(model, search);
83+
query.setEmbeddingVector(embeddings, model);
84+
query.setTopK(pageable.getPageSize());
85+
} else {
86+
// Text search
87+
query.setSearchText(search);
88+
query.setExactMatch(exactMatch);
89+
}
6590

6691
query.addFilter("type", List.of("entity"), SearchType.WHOLE_FIELD);
6792
query.addFilter("ontologyId", List.of(ontologyId), SearchType.CASE_INSENSITIVE_TOKENS);

‎backend/src/main/java/uk/ac/ebi/spot/ols/repository/solr/OlsSolrClient.java‎

Lines changed: 77 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,83 @@ private JsonElement getOlsEntityFromSolrResult(SolrDocument doc) {
164164
}
165165

166166
public QueryResponse runSolrQuery(OlsSolrQuery query, Pageable pageable) {
167-
return runSolrQuery(query.constructQuery(), pageable);
167+
// Check if this is a vector search
168+
if (query.isVectorSearch()) {
169+
return runVectorQuery(query, pageable);
170+
}
171+
return runSolrQuery(query.constructQuery(), pageable);
172+
}
173+
174+
private QueryResponse runVectorQuery(OlsSolrQuery query, Pageable pageable) {
175+
float[] vector = query.getEmbeddingVector();
176+
String modelName = query.getEmbeddingModel();
177+
178+
SolrQuery solrQuery = new SolrQuery();
179+
180+
// Build vector string for KNN query
181+
StringBuilder vectorStr = new StringBuilder("[");
182+
for (int i = 0; i < vector.length; i++) {
183+
if (i > 0) vectorStr.append(",");
184+
vectorStr.append(vector[i]);
185+
}
186+
vectorStr.append("]");
187+
188+
// Use Solr's KNN query parser
189+
String embeddingField = "embeddings_" + modelName;
190+
int topK = query.getTopK() != null ? query.getTopK() : (pageable != null ? pageable.getPageSize() : 10);
191+
solrQuery.setQuery("{!knn f=" + embeddingField + " topK=" + topK + "}" + vectorStr.toString());
192+
193+
// Request all standard fields plus score
194+
solrQuery.setFields("_json", "score");
195+
196+
// Apply filters from the query
197+
SolrQuery baseQuery = query.constructQuery();
198+
String[] filterQueries = baseQuery.getFilterQueries();
199+
if (filterQueries != null) {
200+
for (String fq : filterQueries) {
201+
solrQuery.addFilterQuery(fq);
202+
}
203+
}
204+
205+
// Apply facets
206+
if (baseQuery.getFacetFields() != null) {
207+
for (String facetField : baseQuery.getFacetFields()) {
208+
solrQuery.addFacetField(facetField);
209+
}
210+
}
211+
212+
// Apply pagination
213+
if (pageable != null) {
214+
solrQuery.setStart((int) pageable.getOffset());
215+
solrQuery.setRows(pageable.getPageSize() > maxRows ? maxRows : pageable.getPageSize());
216+
} else {
217+
solrQuery.setStart(0);
218+
solrQuery.setRows(topK > maxRows ? maxRows : topK);
219+
}
220+
221+
logger.debug("Vector search query (length: {}): {}", solrQuery.toQueryString().length(),
222+
solrQuery.toQueryString().length() > 200 ? solrQuery.toQueryString().substring(0, 200) + "..." : solrQuery.toQueryString());
223+
224+
QueryResponse qr = null;
225+
org.apache.solr.client.solrj.SolrClient mySolrClient = new HttpSolrClient.Builder(host + "/solr/ols4_entities").build();
226+
227+
try {
228+
// Use POST method via QueryRequest to avoid URI too long errors with large vectors
229+
org.apache.solr.client.solrj.request.QueryRequest req = new org.apache.solr.client.solrj.request.QueryRequest(solrQuery);
230+
req.setMethod(org.apache.solr.client.solrj.SolrRequest.METHOD.POST);
231+
qr = req.process(mySolrClient);
232+
logger.debug("Vector search found {} result(s)", qr.getResults().getNumFound());
233+
} catch (SolrServerException | IOException e) {
234+
throw new RuntimeException("Vector search failed", e);
235+
} finally {
236+
try {
237+
mySolrClient.close();
238+
} catch (IOException ioe) {
239+
logger.error("Failed to close Solr client with exception \"{}\"", ioe.getMessage());
240+
}
241+
}
242+
243+
return qr;
168244
}
169245

170246
public QueryResponse runSolrQuery(SolrQuery query, Pageable pageable) {

‎backend/src/main/java/uk/ac/ebi/spot/ols/repository/solr/OlsSolrQuery.java‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,11 @@ public class OlsSolrQuery {
1515
List<BoostField> boostFields = new ArrayList<>();
1616
List<String> facetFields = new ArrayList<>();
1717
List<Filter> filters = new ArrayList<>();
18+
19+
// Vector search fields
20+
float[] embeddingVector = null;
21+
String embeddingModel = null;
22+
Integer topK = null;
1823

1924
public OlsSolrQuery() {
2025
}
@@ -30,6 +35,31 @@ public String getSearchText() {
3035
public void setExactMatch(boolean exactMatch) {
3136
this.exactMatch = exactMatch;
3237
}
38+
39+
public void setEmbeddingVector(float[] vector, String model) {
40+
this.embeddingVector = vector;
41+
this.embeddingModel = model;
42+
}
43+
44+
public void setTopK(Integer topK) {
45+
this.topK = topK;
46+
}
47+
48+
public float[] getEmbeddingVector() {
49+
return this.embeddingVector;
50+
}
51+
52+
public String getEmbeddingModel() {
53+
return this.embeddingModel;
54+
}
55+
56+
public Integer getTopK() {
57+
return this.topK;
58+
}
59+
60+
public boolean isVectorSearch() {
61+
return embeddingVector != null && embeddingModel != null;
62+
}
3363

3464
public void addSearchField(String propertyName, int weight, SearchType searchType) {
3565
this.searchFields.add(new SearchField(propertyName, weight, searchType));

‎frontend/src/components/SearchBox.tsx‎

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import Entity from "../model/Entity";
88
import Ontology from "../model/Ontology";
99
import { Suggest } from "../model/Suggest";
1010
import Thing from "../model/Thing";
11+
import Model from "../model/Model";
1112

1213
let curSearchToken: any = null;
1314

@@ -37,7 +38,7 @@ export default function SearchBox({
3738
const [arrowKeySelectedN, setArrowKeySelectedN] = useState<
3839
number | undefined
3940
>(undefined);
40-
const [availableModels, setAvailableModels] = useState<string[]>([]);
41+
const [availableModels, setAvailableModels] = useState<Model[]>([]);
4142
const [selectedModel, setSelectedModel] = useState<string>("lexical");
4243

4344
let exact = searchParams.get("exactMatch") === "true";
@@ -93,22 +94,22 @@ export default function SearchBox({
9394
useEffect(() => {
9495
mounted.current = true;
9596

96-
// Fetch available models
97-
get<any>("api/v2/models")
98-
.then((response) => {
99-
if (response && response.models) {
100-
setAvailableModels(["lexical", ...response.models]);
101-
}
102-
})
103-
.catch((error) => {
104-
console.error("Error fetching models:", error);
105-
});
106-
10797
return () => {
10898
mounted.current = false;
10999
};
110100
});
111101

102+
useEffect(() => {
103+
104+
async function fetchModels() {
105+
setAvailableModels(await get<Model[]>("api/v2/llm_models"));
106+
}
107+
108+
fetchModels();
109+
110+
}, []);
111+
112+
112113
const cancelPromisesRef = useRef(false);
113114
useEffect(() => {
114115
async function loadSuggestions() {
@@ -128,6 +129,7 @@ export default function SearchBox({
128129
includeObsoleteEntities: obsolete.toString(),
129130
...(ontologyId ? { ontologyId } : {}),
130131
...((canonical ? { isDefiningOntology: true } : {}) as any),
132+
...((selectedModel && selectedModel !== 'lexical') ? { model: selectedModel } : {})
131133
})}`
132134
),
133135
searchForOntologies
@@ -137,7 +139,7 @@ export default function SearchBox({
137139
size: "5",
138140
lang: "en",
139141
exactMatch: exact.toString(),
140-
includeObsoleteEntities: obsolete.toString(),
142+
includeObsoleteEntities: obsolete.toString()
141143
})}`
142144
)
143145
: null,
@@ -463,9 +465,10 @@ export default function SearchBox({
463465
label="Search Model"
464466
onChange={(e) => setSelectedModel(e.target.value)}
465467
>
468+
<MenuItem key="lexical" value="lexical">Lexical</MenuItem>
466469
{availableModels.map((model) => (
467-
<MenuItem key={model} value={model}>
468-
{model === "lexical" ? "Lexical (No embeddings)" : model}
470+
<MenuItem key={model.model} value={model.model}>
471+
{model.model === "lexical" ? "Lexical" : model.model}
469472
</MenuItem>
470473
))}
471474
</Select>

0 commit comments

Comments
 (0)