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.

Sign up for GitHub

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

Jump to bottom

Backport #57 to 2.x #63

Closed
wants to merge 1 commit into from
Closed

Backport #57 to 2.x #63

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,7 @@ dependencies {
compileOnly group: 'org.json', name: 'json', version: '20231013'
compileOnly("com.google.guava:guava:32.1.3-jre")
compileOnly group: 'org.apache.commons', name: 'commons-lang3', version: '3.10'
compileOnly group: 'org.apache.commons', name: 'commons-text', version: '1.10.0'

// Plugin dependencies
compileOnly group: 'org.opensearch', name:'opensearch-ml-client', version: "${version}"
Expand Down
87 changes: 63 additions & 24 deletions src/main/java/org/opensearch/agent/tools/PPLTool.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,19 +5,20 @@

package org.opensearch.agent.tools;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.security.AccessController;
import java.security.PrivilegedActionException;
import java.security.PrivilegedExceptionAction;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Locale;
import java.util.List;
import java.util.Map;
import java.util.StringJoiner;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import org.apache.commons.lang3.StringUtils;
import org.apache.commons.text.StringSubstitutor;
import org.json.JSONObject;
import org.opensearch.action.ActionRequest;
import org.opensearch.action.admin.indices.mapping.get.GetMappingsRequest;
Expand All @@ -28,9 +29,6 @@
import org.opensearch.cluster.metadata.MappingMetadata;
import org.opensearch.core.action.ActionListener;
import org.opensearch.core.action.ActionResponse;
import org.opensearch.core.common.io.stream.InputStreamStreamInput;
import org.opensearch.core.common.io.stream.OutputStreamStreamOutput;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.index.query.MatchAllQueryBuilder;
import org.opensearch.ml.common.FunctionName;
import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet;
Expand Down Expand Up @@ -58,6 +56,8 @@
import lombok.extern.log4j.Log4j2;

@Log4j2
@Setter
@Getter
@ToolAnnotation(PPLTool.TYPE)
public class PPLTool implements Tool {

Expand Down Expand Up @@ -93,6 +93,9 @@ public PPLTool(Client client, String modelId, String contextPrompt) {
public <T> void run(Map<String, String> parameters, ActionListener<T> listener) {
String indexName = parameters.get("index");
String question = parameters.get("question");
if (StringUtils.isBlank(indexName) || StringUtils.isBlank(question)) {
throw new IllegalArgumentException("Parameter index and question can not be null or empty.");
}
SearchRequest searchRequest = buildSearchRequest(indexName);
GetMappingsRequest getMappingsRequest = buildGetMappingRequest(indexName);
client.admin().indices().getMappings(getMappingsRequest, ActionListener.<GetMappingsResponse>wrap(getMappingsResponse -> {
Expand All @@ -114,7 +117,7 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
ModelTensors modelTensors = modelTensorOutput.getMlModelOutputs().get(0);
ModelTensor modelTensor = modelTensors.getMlModelTensors().get(0);
Map<String, String> dataAsMap = (Map<String, String>) modelTensor.getDataAsMap();
String ppl = dataAsMap.get("output");
String ppl = parseOutput(dataAsMap.get("response"), indexName);
JSONObject jsonContent = new JSONObject(ImmutableMap.of("query", ppl));
PPLQueryRequest pplQueryRequest = new PPLQueryRequest(ppl, jsonContent, null, "jdbc");
TransportPPLQueryRequest transportPPLQueryRequest = new TransportPPLQueryRequest(pplQueryRequest);
Expand Down Expand Up @@ -226,6 +229,8 @@ private String constructTableInfo(SearchHit[] searchHits, Map<String, MappingMet
Map<String, String> fieldsToType = new HashMap<>();
extractNamesTypes(mappingSource, fieldsToType, "");
StringJoiner tableInfoJoiner = new StringJoiner("\n");
List<String> sortedKeys = new ArrayList<>(fieldsToType.keySet());
Collections.sort(sortedKeys);

if (searchHits.length > 0) {
SearchHit hit = searchHits[0];
Expand All @@ -236,12 +241,12 @@ private String constructTableInfo(SearchHit[] searchHits, Map<String, MappingMet
}
extractSamples(sampleSource, fieldsToSample, "");

for (String key : fieldsToType.keySet()) {
for (String key : sortedKeys) {
String line = "- " + key + ": " + fieldsToType.get(key) + " (" + fieldsToSample.get(key) + ")";
tableInfoJoiner.add(line);
}
} else {
for (String key : fieldsToType.keySet()) {
for (String key : sortedKeys) {
String line = "- " + key + ": " + fieldsToType.get(key);
tableInfoJoiner.add(line);
}
Expand All @@ -252,7 +257,10 @@ private String constructTableInfo(SearchHit[] searchHits, Map<String, MappingMet
}

private String constructPrompt(String tableInfo, String question, String indexName) {
return String.format(Locale.getDefault(), contextPrompt, question.strip(), indexName, tableInfo.strip());
Map<String, String> indexInfo = ImmutableMap.of("mappingInfo", tableInfo, "question", question, "indexName", indexName);
StringSubstitutor substitutor = new StringSubstitutor(indexInfo, "${indexInfo.", "}");
String finalPrompt = substitutor.replace(contextPrompt);
return finalPrompt;
}

private void extractNamesTypes(Map<String, Object> mappingSource, Map<String, String> fieldsToType, String prefix) {
Expand Down Expand Up @@ -297,22 +305,53 @@ private static void extractSamples(Map<String, Object> sampleSource, Map<String,
}

private <T extends ActionResponse> ActionListener<T> getPPLTransportActionListener(ActionListener<TransportPPLQueryResponse> listener) {
return ActionListener.wrap(r -> { listener.onResponse(fromActionResponse(r)); }, listener::onFailure);
return ActionListener.wrap(r -> { listener.onResponse(TransportPPLQueryResponse.fromActionResponse(r)); }, listener::onFailure);
}

private static TransportPPLQueryResponse fromActionResponse(ActionResponse actionResponse) {
if (actionResponse instanceof TransportPPLQueryResponse) {
return (TransportPPLQueryResponse) actionResponse;
<<<<<<< HEAD
=======

private Map<String, String> extractFromChatParameters(Map<String, String> parameters) {
if (parameters.containsKey("input")) {
try {
Map<String, String> chatParameters = gson.fromJson(parameters.get("input"), Map.class);
parameters.putAll(chatParameters);
} finally {
return parameters;
}
}
return parameters;
}

private String parseOutput(String llmOutput, String indexName) {
String ppl;
Pattern pattern = Pattern.compile("<ppl>((.|[\\r\\n])+?)</ppl>"); // For ppl like <ppl> source=a \n | fields b </ppl>
Matcher matcher = pattern.matcher(llmOutput);

try (ByteArrayOutputStream baos = new ByteArrayOutputStream(); OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) {
actionResponse.writeTo(osso);
try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) {
return new TransportPPLQueryResponse(input);
if (matcher.find()) {
ppl = matcher.group(1).replaceAll("[\\r\\n]", "").replaceAll("ISNOTNULL", "isnotnull").trim();
} else { // logic for only ppl returned
int sourceIndex = llmOutput.indexOf("source=");
if (sourceIndex != -1) {
llmOutput = llmOutput.substring(sourceIndex);

// Splitting the string at "|"
String[] lists = llmOutput.split("\\|");

// Modifying the first element
if (lists.length > 0) {
lists[0] = "source=" + indexName;
}

// Joining the string back together
ppl = String.join("|", lists);
} else {
throw new IllegalArgumentException("The returned PPL: " + llmOutput + " has wrong format");
}
} catch (IOException e) {
throw new UncheckedIOException("failed to parse ActionResponse into TransportPPLQueryResponse", e);
}

ppl = ppl.replace("`", "");
ppl = ppl.replaceAll("\\bSPAN\\(", "span(");
return ppl;
}

>>>>>>> 30e5255 (Add UT and parser for ppl tool (#57))
}
Loading
Loading