Skip to content
This repository has been archived by the owner on Jun 6, 2024. It is now read-only.

Commit

Permalink
Looking at using the ext section.
Browse files Browse the repository at this point in the history
  • Loading branch information
jzonthemtn committed Apr 19, 2024
1 parent b17791c commit ed18c0c
Show file tree
Hide file tree
Showing 6 changed files with 267 additions and 4 deletions.
2 changes: 1 addition & 1 deletion docker-compose.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ services:
http.cors.allow-methods: OPTIONS,TRACE,HEAD,GET,POST,PUT,DELETE
http.cors.allow-credentials: true
http.cors.allow-headers: X-Requested-With,X-Auth-Token,Content-Type,Content-Length,Authorization,X-ubi-store,X-ubi-query-id,X-ubi-user-id,X-ubi-session-id
logger.level: info
logger.level: debug
OPENSEARCH_INITIAL_ADMIN_PASSWORD: SuperSecretPassword_123
ulimits:
memlock:
Expand Down
16 changes: 16 additions & 0 deletions search.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
#!/bin/bash -e

curl -X GET "http://localhost:9200/ecommerce/_search" -H 'Content-Type: application/json' -d'
{
"query": {
"match": {
"text_entry": "I am the night"
}
},
"ext": {
"ubi_parameters": {
"query_id": "1234512345"
}
}
}
}'
48 changes: 45 additions & 3 deletions src/main/java/com/o19s/ubi/UserBehaviorInsightsPlugin.java
Original file line number Diff line number Diff line change
Expand Up @@ -13,20 +13,26 @@
import com.o19s.ubi.model.HeaderConstants;
import com.o19s.ubi.model.SettingsConstants;
import com.o19s.ubi.rest.UserBehaviorInsightsRestHandler;
import com.o19s.ubi.ext.UbiParamExtBuilder;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.opensearch.action.support.ActionFilter;
import org.opensearch.client.Client;
import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
import org.opensearch.cluster.node.DiscoveryNodes;
import org.opensearch.cluster.service.ClusterService;
import org.opensearch.common.settings.*;
import org.opensearch.common.settings.ClusterSettings;
import org.opensearch.common.settings.IndexScopedSettings;
import org.opensearch.common.settings.Setting;
import org.opensearch.common.settings.Settings;
import org.opensearch.common.settings.SettingsFilter;
import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
import org.opensearch.core.xcontent.NamedXContentRegistry;
import org.opensearch.env.Environment;
import org.opensearch.env.NodeEnvironment;
import org.opensearch.plugins.ActionPlugin;
import org.opensearch.plugins.Plugin;
import org.opensearch.plugins.SearchPlugin;
import org.opensearch.repositories.RepositoriesService;
import org.opensearch.rest.RestController;
import org.opensearch.rest.RestHandler;
Expand All @@ -35,7 +41,12 @@
import org.opensearch.threadpool.ThreadPool;
import org.opensearch.watcher.ResourceWatcherService;

import java.util.*;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;

Expand All @@ -44,7 +55,7 @@
/**
* OpenSearch User Behavior Insights Plugin
*/
public class UserBehaviorInsightsPlugin extends Plugin implements ActionPlugin {
public class UserBehaviorInsightsPlugin extends Plugin implements ActionPlugin, SearchPlugin {

private static final Logger LOGGER = LogManager.getLogger(UserBehaviorInsightsPlugin.class);

Expand Down Expand Up @@ -140,4 +151,35 @@ public Collection<Object> createComponents(

}

@Override
public List<SearchExtSpec<?>> getSearchExts() {

final List<SearchExtSpec<?>> searchExts = new ArrayList<>();

searchExts.add(new SearchExtSpec<>(
UbiParamExtBuilder.PARAMETER_NAME,
UbiParamExtBuilder::new,
UbiParamExtBuilder::parse));

return searchExts;
}

// @Override
// public Map<String, Processor.Factory<SearchRequestProcessor>> getRequestProcessors(SearchPipelinePlugin.Parameters parameters) {
//
// final Map<String, Processor.Factory<SearchRequestProcessor>> requestProcessors = new HashMap<>();
// requestProcessors.put(GenerativeQAProcessorConstants.REQUEST_PROCESSOR_TYPE, new GenerativeQARequestProcessor.Factory(() -> this.ragSearchPipelineEnabled));
// return requestProcessors;
//
// }
//
// @Override
// public Map<String, Processor.Factory<SearchResponseProcessor>> getResponseProcessors(Parameters parameters) {
//
// final Map<String, Processor.Factory<SearchResponseProcessor>> responseProcessors = new HashMap<>();
// responseProcessors.put(GenerativeQAProcessorConstants.RESPONSE_PROCESSOR_TYPE, new GenerativeQAResponseProcessor.Factory(this.client, () -> this.ragSearchPipelineEnabled));
//
// return responseProcessors;
// }

}
83 changes: 83 additions & 0 deletions src/main/java/com/o19s/ubi/ext/UbiParamExtBuilder.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
/*
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*/

package com.o19s.ubi.ext;

import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.search.SearchExtBuilder;

import java.io.IOException;
import java.util.Objects;

public class UbiParamExtBuilder extends SearchExtBuilder {

// The name of the "ext" section containing UBI parameters.
public static final String PARAMETER_NAME = "ubi_parameters";

private UbiParameters params;

public UbiParamExtBuilder() {

}

public UbiParamExtBuilder(StreamInput input) throws IOException {
this.params = new UbiParameters(input);
}

@Override
public int hashCode() {
return Objects.hash(this.getClass(), this.params);
}

@Override
public boolean equals(Object obj) {
if (this == obj) {
return true;
}

if (!(obj instanceof UbiParamExtBuilder)) {
return false;
}

return this.params.equals(((UbiParamExtBuilder) obj).getParams());
}

@Override
public String getWriteableName() {
return PARAMETER_NAME;
}

@Override
public void writeTo(StreamOutput out) throws IOException {
this.params.writeTo(out);
}

@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
return builder.value(this.params);
}

public static UbiParamExtBuilder parse(XContentParser parser) throws IOException {
final UbiParamExtBuilder builder = new UbiParamExtBuilder();
UbiParameters params = UbiParameters.parse(parser);
builder.setParams(params);
return builder;
}

public UbiParameters getParams() {
return params;
}

public void setParams(UbiParameters params) {
this.params = params;
}

}
39 changes: 39 additions & 0 deletions src/main/java/com/o19s/ubi/ext/UbiParamUtil.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
/*
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*/

package com.o19s.ubi.ext;

import org.opensearch.action.search.SearchRequest;
import org.opensearch.search.SearchExtBuilder;

import java.util.Optional;

public class UbiParamUtil {

public static UbiParameters getUbiParameters(SearchRequest request) {
UbiParamExtBuilder builder = null;
if (request.source() != null && request.source().ext() != null && !request.source().ext().isEmpty()) {
Optional<SearchExtBuilder> b = request
.source()
.ext()
.stream()
.filter(bldr -> UbiParamExtBuilder.PARAMETER_NAME.equals(bldr.getWriteableName()))
.findFirst();
if (b.isPresent()) {
builder = (UbiParamExtBuilder) b.get();
}
}

UbiParameters params = null;
if (builder != null) {
params = builder.getParams();
}

return params;
}
}
83 changes: 83 additions & 0 deletions src/main/java/com/o19s/ubi/ext/UbiParameters.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
/*
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*/

package com.o19s.ubi.ext;

import org.opensearch.core.ParseField;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.common.io.stream.Writeable;
import org.opensearch.core.xcontent.ObjectParser;
import org.opensearch.core.xcontent.ToXContentObject;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.core.xcontent.XContentParser;

import java.io.IOException;
import java.util.Objects;

public class UbiParameters implements Writeable, ToXContentObject {

private static final ObjectParser<UbiParameters, Void> PARSER;
private static final ParseField QUERY_ID = new ParseField("query_id");

static {
PARSER = new ObjectParser<>(UbiParamExtBuilder.PARAMETER_NAME, UbiParameters::new);
PARSER.declareString(UbiParameters::setQueryId, QUERY_ID);
}

private String queryId;

public UbiParameters(String queryId) {
this.queryId = queryId;
}

public UbiParameters(StreamInput input) throws IOException {
this.queryId = input.readString();
}

public UbiParameters() {

}

@Override
public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params) throws IOException {
return xContentBuilder
.field(QUERY_ID.getPreferredName(), this.queryId);
}

@Override
public void writeTo(StreamOutput out) throws IOException {
out.writeOptionalString(queryId);
}

public static UbiParameters parse(XContentParser parser) throws IOException {
return PARSER.parse(parser, null);
}

@Override
public boolean equals(Object o) {
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
return false;
}

UbiParameters other = (UbiParameters) o;
return Objects.equals(this.queryId, other.getQueryId());
}

public String getQueryId() {
return queryId;
}

public void setQueryId(String conversationId) {
this.queryId = conversationId;
}

}

0 comments on commit ed18c0c

Please sign in to comment.