Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adding default use cases #583

Merged
merged 5 commits into from
Mar 17, 2024
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,7 @@ dependencies {

// ZipArchive dependencies used for integration tests
zipArchive group: 'org.opensearch.plugin', name:'opensearch-ml-plugin', version: "${opensearch_build}"

owaiskazi19 marked this conversation as resolved.
Show resolved Hide resolved
secureIntegTestPluginArchive group: 'org.opensearch.plugin', name:'opensearch-security', version: "${opensearch_build}"

configurations.all {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,8 @@ private CommonValue() {}
public static final String PROVISION_WORKFLOW = "provision";
/** The field name for workflow steps. This field represents the name of the workflow steps to be fetched. */
public static final String WORKFLOW_STEP = "workflow_step";
/** The param name for default use case, used by the create workflow API */
public static final String USE_CASE = "use_case";

/*
* Constants associated with plugin configuration
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
/*
* Copyright OpenSearch Contributors
* 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 org.opensearch.flowframework.common;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.opensearch.core.rest.RestStatus;
import org.opensearch.flowframework.exception.FlowFrameworkException;

import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
* Enum encapsulating the different default use cases and templates we have stored
*/
public enum DefaultUseCases {

/** defaults file and substitution ready template for OpenAI embedding model */
OPEN_AI_EMBEDDING_MODEL_DEPLOY(
"open_ai_embedding_model_deploy",
"defaults/open-ai-embedding-defaults.json",
"substitutionTemplates/deploy-remote-model-template.json"
),
/** defaults file and substitution ready template for cohere embedding model */
COHERE_EMBEDDING_MODEL_DEPLOY(
"cohere-embedding_model_deploy",
"defaults/cohere-embedding-defaults.json",
"substitutionTemplates/deploy-remote-model-template-extra-params.json"
),
LOCAL_NEURAL_SPARSE_SEARCH(
"local_neural_sparse_search",
"defaults/local-sparse-search-defaults.json",
"substitutionTemplates/neural-sparse-local-template.json"
);
owaiskazi19 marked this conversation as resolved.
Show resolved Hide resolved

private final String useCaseName;
private final String defaultsFile;
private final String substitutionReadyFile;
private static final Logger logger = LogManager.getLogger(DefaultUseCases.class);
private static final Set<String> allResources = Stream.of(values()).map(DefaultUseCases::getDefaultsFile).collect(Collectors.toSet());
owaiskazi19 marked this conversation as resolved.
Show resolved Hide resolved

DefaultUseCases(String useCaseName, String defaultsFile, String substitutionReadyFile) {
this.useCaseName = useCaseName;
this.defaultsFile = defaultsFile;
this.substitutionReadyFile = substitutionReadyFile;
}

/**
* Returns the useCaseName for the given enum Constant
* @return the useCaseName of this use case.
*/
public String getUseCaseName() {
return useCaseName;
}

/**
* Returns the defaultsFile for the given enum Constant
* @return the defaultsFile of this for the given useCase.
*/
public String getDefaultsFile() {
return defaultsFile;
}

/**
* Returns the substitutionReadyFile for the given enum Constant
* @return the substitutionReadyFile of the given useCase
*/
public String getSubstitutionReadyFile() {
return substitutionReadyFile;
}

/**
* Gets the defaultsFile based on the given use case.
* @param useCaseName name of the given use case
* @return the deafultsFile for that usecase
amitgalitz marked this conversation as resolved.
Show resolved Hide resolved
* @throws FlowFrameworkException if the use case doesn't exist in enum
*/
public static String getDefaultsFileByUseCaseName(String useCaseName) throws FlowFrameworkException {
if (useCaseName != null && !useCaseName.isEmpty()) {
for (DefaultUseCases mapping : values()) {
amitgalitz marked this conversation as resolved.
Show resolved Hide resolved
if (useCaseName.equals(mapping.getUseCaseName())) {
return mapping.getDefaultsFile();
}
}
}
logger.error("Unable to find defaults file for use case: {}", useCaseName);
throw new FlowFrameworkException("Unable to find defaults file for use case: " + useCaseName, RestStatus.BAD_REQUEST);
}

/**
* Gets the substitutionReadyFile based on the given use case
* @param useCaseName name of the given use case
* @return the substitutionReadyFile which has the template
* @throws FlowFrameworkException if the use case doesn't exist in enum
*/
public static String getSubstitutionReadyFileByUseCaseName(String useCaseName) throws FlowFrameworkException {
if (useCaseName != null && !useCaseName.isEmpty()) {
for (DefaultUseCases mapping : values()) {
owaiskazi19 marked this conversation as resolved.
Show resolved Hide resolved
if (mapping.getUseCaseName().equals(useCaseName)) {
return mapping.getSubstitutionReadyFile();
}
}
}
logger.error("Unable to find substitution ready file for use case: {}", useCaseName);
throw new FlowFrameworkException("Unable to find substitution ready file for use case: " + useCaseName, RestStatus.BAD_REQUEST);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -17,16 +17,19 @@
import org.opensearch.core.xcontent.ToXContent;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.flowframework.common.DefaultUseCases;
import org.opensearch.flowframework.common.FlowFrameworkSettings;
import org.opensearch.flowframework.exception.FlowFrameworkException;
import org.opensearch.flowframework.model.Template;
import org.opensearch.flowframework.transport.CreateWorkflowAction;
import org.opensearch.flowframework.transport.WorkflowRequest;
import org.opensearch.flowframework.util.ParseUtils;
import org.opensearch.rest.BaseRestHandler;
import org.opensearch.rest.BytesRestResponse;
import org.opensearch.rest.RestRequest;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.Locale;
import java.util.Map;
Expand All @@ -35,6 +38,7 @@

import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
import static org.opensearch.flowframework.common.CommonValue.PROVISION_WORKFLOW;
import static org.opensearch.flowframework.common.CommonValue.USE_CASE;
import static org.opensearch.flowframework.common.CommonValue.VALIDATION;
import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID;
import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_URI;
Expand Down Expand Up @@ -78,6 +82,7 @@
String workflowId = request.param(WORKFLOW_ID);
String[] validation = request.paramAsStringArray(VALIDATION, new String[] { "all" });
boolean provision = request.paramAsBoolean(PROVISION_WORKFLOW, false);
String useCase = request.param(USE_CASE);
// If provisioning, consume all other params and pass to provision transport action
Map<String, String> params = provision
? request.params()
Expand Down Expand Up @@ -112,11 +117,54 @@
);
}
try {
XContentParser parser = request.contentParser();
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
Template template = Template.parse(parser);

WorkflowRequest workflowRequest = new WorkflowRequest(workflowId, template, validation, provision, params);
Template template;
Map<String, String> useCaseDefaultsMap = Collections.emptyMap();
if (useCase != null) {
String json = ParseUtils.resourceToString("/" + DefaultUseCases.getSubstitutionReadyFileByUseCaseName(useCase));
owaiskazi19 marked this conversation as resolved.
Show resolved Hide resolved
String defaultsFilePath = DefaultUseCases.getDefaultsFileByUseCaseName(useCase);
useCaseDefaultsMap = ParseUtils.parseJsonFileToStringToStringMap("/" + defaultsFilePath);

Check warning on line 126 in src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java#L124-L126

Added lines #L124 - L126 were not covered by tests

if (request.hasContent()) {
try {
XContentParser parser = request.contentParser();
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
amitgalitz marked this conversation as resolved.
Show resolved Hide resolved
Map<String, String> userDefaults = ParseUtils.parseStringToStringMap(parser);

Check warning on line 132 in src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java#L130-L132

Added lines #L130 - L132 were not covered by tests
// updates the default params with anything user has given that matches
for (Map.Entry<String, String> userDefaultsEntry : userDefaults.entrySet()) {
if (useCaseDefaultsMap.containsKey(userDefaultsEntry.getKey())) {
amitgalitz marked this conversation as resolved.
Show resolved Hide resolved
useCaseDefaultsMap.put(userDefaultsEntry.getKey(), userDefaultsEntry.getValue());

Check warning on line 136 in src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java#L136

Added line #L136 was not covered by tests
}
}
} catch (Exception ex) {
String errorMessage = "failure parsing request body when a use case is given";
logger.error(errorMessage, ex);
throw new FlowFrameworkException(errorMessage, ExceptionsHelper.status(ex));
amitgalitz marked this conversation as resolved.
Show resolved Hide resolved
}

Check warning on line 143 in src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java#L138-L143

Added lines #L138 - L143 were not covered by tests

}

json = (String) ParseUtils.conditionallySubstitute(json, null, useCaseDefaultsMap);

Check warning on line 147 in src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java#L147

Added line #L147 was not covered by tests

XContentParser parserTestJson = ParseUtils.jsonToParser(json);
ensureExpectedToken(XContentParser.Token.START_OBJECT, parserTestJson.currentToken(), parserTestJson);
template = Template.parse(parserTestJson);

Check warning on line 151 in src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java#L149-L151

Added lines #L149 - L151 were not covered by tests

} else {

Check warning on line 153 in src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java#L153

Added line #L153 was not covered by tests
XContentParser parser = request.contentParser();
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
template = Template.parse(parser);
}

WorkflowRequest workflowRequest = new WorkflowRequest(
workflowId,
template,
validation,
provision,
params,
useCase,
useCaseDefaultsMap
);

return channel -> client.execute(CreateWorkflowAction.INSTANCE, workflowRequest, ActionListener.wrap(response -> {
XContentBuilder builder = response.toXContent(channel.newBuilder(), ToXContent.EMPTY_PARAMS);
Expand All @@ -134,11 +182,14 @@
channel.sendResponse(new BytesRestResponse(ExceptionsHelper.status(e), errorMessage));
}
}));

} catch (FlowFrameworkException e) {
logger.error("failed to prepare rest request", e);
return channel -> channel.sendResponse(
new BytesRestResponse(e.getRestStatus(), e.toXContent(channel.newErrorBuilder(), ToXContent.EMPTY_PARAMS))
);
} catch (IOException e) {
} catch (Exception e) {
logger.error("failed to prepare rest request", e);

Check warning on line 192 in src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/org/opensearch/flowframework/rest/RestCreateWorkflowAction.java#L191-L192

Added lines #L191 - L192 were not covered by tests
FlowFrameworkException ex = new FlowFrameworkException(
"IOException: template content invalid for specified Content-Type.",
RestStatus.BAD_REQUEST
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -49,13 +49,23 @@
*/
private Map<String, String> params;

/**
* use case flag
*/
private String useCase;

/**
* Deafult params map from use case
*/
private Map<String, String> defaultParams;

/**
* Instantiates a new WorkflowRequest, set validation to all, no provisioning
* @param workflowId the documentId of the workflow
* @param template the use case template which describes the workflow
*/
public WorkflowRequest(@Nullable String workflowId, @Nullable Template template) {
this(workflowId, template, new String[] { "all" }, false, Collections.emptyMap());
this(workflowId, template, new String[] { "all" }, false, Collections.emptyMap(), null, Collections.emptyMap());
}

/**
Expand All @@ -65,7 +75,7 @@
* @param params The parameters from the REST path
*/
public WorkflowRequest(@Nullable String workflowId, @Nullable Template template, Map<String, String> params) {
this(workflowId, template, new String[] { "all" }, true, params);
this(workflowId, template, new String[] { "all" }, true, params, null, Collections.emptyMap());
}

/**
Expand All @@ -75,13 +85,17 @@
* @param validation flag to indicate if validation is necessary
* @param provision flag to indicate if provision is necessary
* @param params map of REST path params. If provision is false, must be an empty map.
* @param useCase default use case given
* @param defaultParams the params to be used in the substitution based on the default use case.
*/
public WorkflowRequest(
@Nullable String workflowId,
@Nullable Template template,
String[] validation,
boolean provision,
Map<String, String> params
Map<String, String> params,
String useCase,
Map<String, String> defaultParams
) {
this.workflowId = workflowId;
this.template = template;
Expand All @@ -91,6 +105,8 @@
throw new IllegalArgumentException("Params may only be included when provisioning.");
}
this.params = params;
this.useCase = useCase;
this.defaultParams = defaultParams;
}

/**
Expand Down Expand Up @@ -150,6 +166,22 @@
return Map.copyOf(this.params);
}

/**
* Gets the params map
owaiskazi19 marked this conversation as resolved.
Show resolved Hide resolved
* @return the params map
*/
public String getUseCase() {
return this.useCase;

Check warning on line 174 in src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java#L174

Added line #L174 was not covered by tests
}

/**
* Gets the params map
* @return the params map
*/
public Map<String, String> getDefaultParams() {
return Map.copyOf(this.defaultParams);

Check warning on line 182 in src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/org/opensearch/flowframework/transport/WorkflowRequest.java#L182

Added line #L182 was not covered by tests
}
amitgalitz marked this conversation as resolved.
Show resolved Hide resolved

@Override
public void writeTo(StreamOutput out) throws IOException {
super.writeTo(out);
Expand Down
39 changes: 34 additions & 5 deletions src/main/java/org/opensearch/flowframework/util/ParseUtils.java
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,8 @@
public static void buildStringToStringMap(XContentBuilder xContentBuilder, Map<?, ?> map) throws IOException {
xContentBuilder.startObject();
for (Entry<?, ?> e : map.entrySet()) {
String key = (String) e.getKey();
amitgalitz marked this conversation as resolved.
Show resolved Hide resolved
String value = (String) e.getValue();
xContentBuilder.field((String) e.getKey(), (String) e.getValue());
}
xContentBuilder.endObject();
Expand Down Expand Up @@ -342,11 +344,18 @@
return inputs;
}

private static Object conditionallySubstitute(Object value, Map<String, WorkflowData> outputs, Map<String, String> params) {
/**
* Executes substitution on the given value by looking at any matching values in either the ouputs or params map
* @param value the Object that will have the substitution done on
* @param outputs potential location of values to be substituted in
* @param params potential location of values to be subsituted in
* @return the substituted object back
*/
public static Object conditionallySubstitute(Object value, Map<String, WorkflowData> outputs, Map<String, String> params) {
if (value instanceof String) {
Matcher m = SUBSTITUTION_PATTERN.matcher((String) value);
StringBuilder result = new StringBuilder();
while (m.find()) {
while (m.find() && outputs != null) {
// outputs content map contains values for previous node input (e.g: deploy_openai_model.model_id)
// Check first if the substitution is looking for the same key, value pair and if yes
// then replace it with the key value pair in the inputs map
Expand All @@ -364,10 +373,17 @@
m.appendTail(result);
value = result.toString();

// Replace all params if present
for (Entry<String, String> e : params.entrySet()) {
for (Map.Entry<String, String> e : params.entrySet()) {
String regex = "\\$\\{\\{\\s*" + Pattern.quote(e.getKey()) + "\\s*\\}\\}";
value = ((String) value).replaceAll(regex, e.getValue());
String replacement = e.getValue();

// Special handling for JSON strings that contain placeholders (connectors action)
replacement = Matcher.quoteReplacement(replacement.replace("\"", "\\\""));

// Use Pattern.compile().matcher() to avoid issues with replaceAll's direct pattern compilation
Pattern pattern = Pattern.compile(regex);
Matcher matcher = pattern.matcher((String) value);
value = matcher.replaceAll(replacement);
}
}
return value;
Expand All @@ -385,4 +401,17 @@
String mappedString = mapper.writeValueAsString(map);
return mappedString;
}

/**
* Generates a String to String map based on a Json File
* @param path file path
* @return instance of the string
* @throws JsonProcessingException JsonProcessingException from Jackson for issues processing map
*/
public static Map<String, String> parseJsonFileToStringToStringMap(String path) throws IOException {
ObjectMapper mapper = new ObjectMapper();
amitgalitz marked this conversation as resolved.
Show resolved Hide resolved
String jsonContent = resourceToString(path);
Map<String, String> mappedJsonFile = mapper.readValue(jsonContent, Map.class);
amitgalitz marked this conversation as resolved.
Show resolved Hide resolved
return mappedJsonFile;

Check warning on line 415 in src/main/java/org/opensearch/flowframework/util/ParseUtils.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/org/opensearch/flowframework/util/ParseUtils.java#L412-L415

Added lines #L412 - L415 were not covered by tests
}
}
Loading
Loading