use of io.openk9.model.SuggestionCategoryField in project openk9 by smclab.
the class SuggestionsHTTPHandler method _getSuggestionsResponse.
private SuggestionsResponse _getSuggestionsResponse(List<Datasource> datasourceList, PluginDriverDTOList pluginDriverDTOList, SearchRequest searchRequest, SearchResponse searchResponse, List<SuggestionCategoryField> fields, Map<String, String[]> entityMap) {
Aggregations aggregations = searchResponse.getAggregations();
if (aggregations == null) {
return SuggestionsResponse.of(List.of(), null);
}
Map<String, Aggregation> aggregationMap = aggregations.asMap();
if (!aggregationMap.containsKey("composite")) {
return SuggestionsResponse.of(List.of(), null);
}
CompositeAggregation compositeAggregation = (CompositeAggregation) aggregationMap.get("composite");
Map<String, Long> fieldNameCategoryIdMap = fields.stream().collect(Collectors.toMap(SuggestionCategoryField::getFieldName, SuggestionCategoryField::getCategoryId, (a1, a2) -> a2));
Long datasourceIdCategoryId = fieldNameCategoryIdMap.getOrDefault("datasourceId", 1L);
Long entityIdCategoryId = fieldNameCategoryIdMap.getOrDefault("entities.id", 2L);
Long entitiesContextCategoryId = fieldNameCategoryIdMap.getOrDefault("entities.context", entityIdCategoryId);
Long documentTypesCategoryId = fieldNameCategoryIdMap.getOrDefault("documentTypes", 4L);
List<? extends CompositeAggregation.Bucket> buckets = compositeAggregation.getBuckets();
LinkedList<Suggestions> suggestions = new LinkedList<>();
String suggestKeyword = searchRequest.getSuggestKeyword();
BiConsumer<String, Suggestions> addSuggestions;
if (suggestKeyword != null) {
addSuggestions = (key, sugg) -> {
if (!suggestions.contains(sugg)) {
if (key.contains(suggestKeyword)) {
suggestions.addFirst(sugg);
}
}
};
} else {
addSuggestions = (key, sugg) -> {
if (!suggestions.contains(sugg)) {
suggestions.add(sugg);
}
};
}
for (CompositeAggregation.Bucket bucket : buckets) {
Map<String, Object> keys = new HashMap<>(bucket.getKey());
for (Map.Entry<String, Object> entry : keys.entrySet()) {
String key = entry.getKey();
String value = (String) entry.getValue();
if (value == null) {
continue;
}
switch(key) {
case "datasourceId":
long datasourceIdL = Long.parseLong(value);
_datasource(datasourceList, pluginDriverDTOList, datasourceIdCategoryId, addSuggestions, datasourceIdL);
break;
case "entities.context":
break;
case "entities.id":
String[] typeName = entityMap.get(value);
if (typeName != null) {
String type = typeName[0];
String name = typeName[1];
String entitiesContext = (String) keys.get("entities.context");
if (entitiesContext != null) {
addSuggestions.accept(name, Suggestions.entity(value, entitiesContextCategoryId, type, name, entitiesContext));
} else {
addSuggestions.accept(name, Suggestions.entity(value, entityIdCategoryId, type, name));
}
}
break;
case "documentTypes":
addSuggestions.accept(value, Suggestions.docType(value, documentTypesCategoryId));
break;
default:
Long textCategoryId = fieldNameCategoryIdMap.getOrDefault(key, 5L);
addSuggestions.accept(value, Suggestions.text(value, textCategoryId, key));
}
}
}
Map<String, Object> map = compositeAggregation.afterKey();
String afterKey = null;
int[] range = searchRequest.getRange();
if (map != null) {
afterKey = _jsonFactory.toJson(map);
afterKey = Base64.getEncoder().encodeToString(afterKey.getBytes(StandardCharsets.UTF_8));
}
return SuggestionsResponse.of(suggestions, afterKey);
}
use of io.openk9.model.SuggestionCategoryField in project openk9 by smclab.
the class SuggestionsHTTPHandler method searchHitToResponseMono.
@Override
protected Mono<Object> searchHitToResponseMono(Tenant tenant, List<Datasource> datasourceList, PluginDriverDTOList pluginDriverDTOList, HttpServerRequest httpServerRequest, SearchRequest searchRequest, SearchResponse searchResponse) {
Long suggestionCategoryId = searchRequest.getSuggestionCategoryId();
Mono<List<SuggestionCategoryField>> suggestionCategoryFields;
if (suggestionCategoryId == null) {
suggestionCategoryFields = _datasourceClient.findSuggestionCategoryFieldsByTenantId(tenant.getTenantId());
} else {
suggestionCategoryFields = _datasourceClient.findSuggestionCategoryFieldsByTenantIdAndCategoryId(tenant.getTenantId(), suggestionCategoryId);
}
if (_enableEntityAggregation) {
return _search.search(factory -> {
org.elasticsearch.action.search.SearchRequest searchRequestEntity = factory.createSearchRequestEntity(tenant.getTenantId());
Aggregations aggregations = searchResponse.getAggregations();
CompositeAggregation compositeAggregation = aggregations.get("composite");
List<? extends CompositeAggregation.Bucket> buckets = compositeAggregation.getBuckets();
BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery();
buckets.stream().map(bucket -> (String) bucket.getKey().get("entities.id")).filter(Objects::nonNull).distinct().forEach(entityId -> boolQueryBuilder.should(QueryBuilders.matchQuery("id", entityId)));
if (_log.isDebugEnabled()) {
_log.debug("entities query: " + boolQueryBuilder);
}
SearchSourceBuilder ssb = new SearchSourceBuilder();
ssb.query(boolQueryBuilder);
ssb.size(1000);
ssb.fetchSource(new String[] { "name", "id", "type" }, null);
return searchRequestEntity.source(ssb);
}).flatMap(entityResponse -> suggestionCategoryFields.map(fields -> {
Map<String, String[]> entityMap = new HashMap<>();
for (SearchHit hit : entityResponse.getHits()) {
Map<String, Object> sourceAsMap = hit.getSourceAsMap();
String name = (String) sourceAsMap.get("name");
String type = (String) sourceAsMap.get("type");
String entityId = (String) sourceAsMap.get("id");
entityMap.put(entityId, new String[] { type, name });
}
return _getSuggestionsResponse(datasourceList, pluginDriverDTOList, searchRequest, searchResponse, fields, entityMap);
}));
} else {
return suggestionCategoryFields.map(fields -> _getSuggestionsResponse(datasourceList, pluginDriverDTOList, searchRequest, searchResponse, fields, Map.of()));
}
}
use of io.openk9.model.SuggestionCategoryField in project openk9 by smclab.
the class SuggestionsHTTPHandler method customizeSearchSourceBuilderMono.
@Override
protected Mono<org.elasticsearch.action.search.SearchRequest> customizeSearchSourceBuilderMono(Tenant tenant, List<Datasource> datasources, SearchRequest searchRequest, List<PluginDriverDTO> documentTypeList, SearchSourceBuilder searchSourceBuilder, org.elasticsearch.action.search.SearchRequest elasticSearchQuery) {
return Mono.defer(() -> {
Long suggestionCategoryId = searchRequest.getSuggestionCategoryId();
if (suggestionCategoryId == null) {
return _datasourceClient.findSuggestionCategoryFieldsByTenantId(tenant.getTenantId());
} else {
return _datasourceClient.findSuggestionCategoryFieldsByTenantIdAndCategoryId(tenant.getTenantId(), suggestionCategoryId);
}
}).map(fields -> {
if (!(fields == null || fields.isEmpty())) {
Function<String, CompositeValuesSourceBuilder<?>> fieldToTerms = nameField -> new TermsValuesSourceBuilder(nameField).field(nameField).missingBucket(true);
CompositeAggregationBuilder compositeAggregation = fields.stream().map(SuggestionCategoryField::getFieldName).map(fieldToTerms).collect(Collectors.collectingAndThen(Collectors.toList(), list -> AggregationBuilders.composite("composite", list)));
String afterKey = searchRequest.getAfterKey();
if (afterKey != null) {
byte[] afterKeyDecoded = Base64.getDecoder().decode(afterKey);
Map<String, Object> map = _jsonFactory.fromJsonMap(new String(afterKeyDecoded), Object.class);
compositeAggregation.aggregateAfter(map);
}
int[] range = searchRequest.getRange();
if (range != null && range.length == 2) {
int size = range[1];
compositeAggregation.size(size);
}
searchSourceBuilder.aggregation(compositeAggregation);
}
searchSourceBuilder.from(0);
searchSourceBuilder.size(0);
searchSourceBuilder.highlighter(null);
return elasticSearchQuery.source(searchSourceBuilder);
});
}
Aggregations