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

Enhance Message and Memory API Validation and storage #3283

Merged
merged 5 commits into from
Dec 24, 2024
Merged

Enhance Message and Memory API Validation and storage #3283

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
c0ac084
Enchance Message and Memory API Validation and storage
rithin-pullela-aws Dec 16, 2024
c8b314c
Update unit test to check for null instead of empty map
rithin-pullela-aws Dec 16, 2024
4e91ea6
Refactored userstr to Camel Case
rithin-pullela-aws Dec 18, 2024
41325c4
Addressing comments
rithin-pullela-aws Dec 18, 2024
e577f61
Undo: throw an error when an unknown field is provided in CreateConve…
rithin-pullela-aws Dec 19, 2024
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
16 changes: 12 additions & 4 deletions common/src/main/java/org/opensearch/ml/common/conversation/Interaction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,10 +184,18 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContentObject.Para
builder.field(ActionConstants.CONVERSATION_ID_FIELD, conversationId);
builder.field(ActionConstants.RESPONSE_INTERACTION_ID_FIELD, id);
builder.field(ConversationalIndexConstants.INTERACTIONS_CREATE_TIME_FIELD, createTime);
builder.field(ConversationalIndexConstants.INTERACTIONS_INPUT_FIELD, input);
builder.field(ConversationalIndexConstants.INTERACTIONS_PROMPT_TEMPLATE_FIELD, promptTemplate);
builder.field(ConversationalIndexConstants.INTERACTIONS_RESPONSE_FIELD, response);
builder.field(ConversationalIndexConstants.INTERACTIONS_ORIGIN_FIELD, origin);
if (input != null && !input.trim().isEmpty()) {
builder.field(ConversationalIndexConstants.INTERACTIONS_INPUT_FIELD, input);
}
if (promptTemplate != null && !promptTemplate.trim().isEmpty()) {
builder.field(ConversationalIndexConstants.INTERACTIONS_PROMPT_TEMPLATE_FIELD, promptTemplate);
}
if (response != null && !response.trim().isEmpty()) {
builder.field(ConversationalIndexConstants.INTERACTIONS_RESPONSE_FIELD, response);
}
if (origin != null && !origin.trim().isEmpty()) {
builder.field(ConversationalIndexConstants.INTERACTIONS_ORIGIN_FIELD, origin);
}
if (additionalInfo != null) {
builder.field(ConversationalIndexConstants.INTERACTIONS_ADDITIONAL_INFO_FIELD, additionalInfo);
}
Expand Down
3 changes: 2 additions & 1 deletion common/src/test/java/org/opensearch/ml/common/conversation/InteractionTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,13 +124,14 @@ public void test_ToXContent() throws IOException {
.origin("amazon bedrock")
.parentInteractionId("parant id")
.additionalInfo(Collections.singletonMap("suggestion", "new suggestion"))
.response("sample response")
.traceNum(1)
.build();
XContentBuilder builder = XContentBuilder.builder(XContentType.JSON.xContent());
interaction.toXContent(builder, EMPTY_PARAMS);
String interactionContent = TestHelper.xContentBuilderToString(builder);
assertEquals(
"{\"memory_id\":\"conversation id\",\"message_id\":null,\"create_time\":null,\"input\":null,\"prompt_template\":null,\"response\":null,\"origin\":\"amazon bedrock\",\"additional_info\":{\"suggestion\":\"new suggestion\"},\"parent_message_id\":\"parant id\",\"trace_number\":1}",
"{\"memory_id\":\"conversation id\",\"message_id\":null,\"create_time\":null,\"response\":\"sample response\",\"origin\":\"amazon bedrock\",\"additional_info\":{\"suggestion\":\"new suggestion\"},\"parent_message_id\":\"parant id\",\"trace_number\":1}",
rithin-pullela-aws marked this conversation as resolved.
Show resolved Hide resolved
interactionContent
);
}
Expand Down
25 changes: 20 additions & 5 deletions ...src/main/java/org/opensearch/ml/memory/action/conversation/CreateConversationRequest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,12 +137,27 @@ public static CreateConversationRequest fromRestRequest(RestRequest restRequest)
}
try (XContentParser parser = restRequest.contentParser()) {
Map<String, Object> body = parser.map();
String name = null;
String applicationType = null;
Map<String, String> additionalInfo = null;

for (String key : body.keySet()) {
rithin-pullela-aws marked this conversation as resolved.
Show resolved Hide resolved
switch (key) {
case ActionConstants.REQUEST_CONVERSATION_NAME_FIELD:
name = (String) body.get(ActionConstants.REQUEST_CONVERSATION_NAME_FIELD);
break;
case APPLICATION_TYPE_FIELD:
applicationType = (String) body.get(APPLICATION_TYPE_FIELD);
break;
case META_ADDITIONAL_INFO_FIELD:
additionalInfo = (Map<String, String>) body.get(META_ADDITIONAL_INFO_FIELD);
break;
default:
throw new IllegalArgumentException("Invalid field [" + key + "] found in request body");
}
}
if (body.get(ActionConstants.REQUEST_CONVERSATION_NAME_FIELD) != null) {
return new CreateConversationRequest(
(String) body.get(ActionConstants.REQUEST_CONVERSATION_NAME_FIELD),
body.get(APPLICATION_TYPE_FIELD) == null ? null : (String) body.get(APPLICATION_TYPE_FIELD),
body.get(META_ADDITIONAL_INFO_FIELD) == null ? null : (Map<String, String>) body.get(META_ADDITIONAL_INFO_FIELD)
);
return new CreateConversationRequest(name, applicationType, additionalInfo);
} else {
return new CreateConversationRequest();
}
Expand Down
13 changes: 11 additions & 2 deletions .../src/main/java/org/opensearch/ml/memory/action/conversation/CreateInteractionRequest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,11 +166,20 @@ public static CreateInteractionRequest fromRestRequest(RestRequest request) thro
tracenum = parser.intValue(false);
break;
default:
parser.skipChildren();
break;
throw new IllegalArgumentException("Invalid field [" + fieldName + "] found in request body");
}
}

boolean allFieldsEmpty = (input == null || input.trim().isEmpty())
&& (prompt == null || prompt.trim().isEmpty())
&& (response == null || response.trim().isEmpty())
&& (origin == null || origin.trim().isEmpty())
&& (addinf == null || addinf.isEmpty());
if (allFieldsEmpty) {
throw new IllegalArgumentException(
"At least one of the following parameters must be non-empty: " + "input, prompt_template, response, origin, additional_info"
rithin-pullela-aws marked this conversation as resolved.
Show resolved Hide resolved
);
}
return new CreateInteractionRequest(cid, input, prompt, response, origin, addinf, parintid, tracenum);
}

Expand Down
33 changes: 17 additions & 16 deletions memory/src/main/java/org/opensearch/ml/memory/index/ConversationMetaIndex.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@

import java.io.IOException;
import java.time.Instant;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
Expand Down Expand Up @@ -141,22 +142,22 @@ public void createConversation(
if (indexExists) {
String userstr = getUserStrFromThreadContext();
Instant now = Instant.now();
IndexRequest request = Requests
.indexRequest(META_INDEX_NAME)
.source(
ConversationalIndexConstants.META_CREATED_TIME_FIELD,
now,
ConversationalIndexConstants.META_UPDATED_TIME_FIELD,
now,
ConversationalIndexConstants.META_NAME_FIELD,
name,
ConversationalIndexConstants.USER_FIELD,
userstr == null ? null : User.parse(userstr).getName(),
ConversationalIndexConstants.APPLICATION_TYPE_FIELD,
applicationType,
ConversationalIndexConstants.META_ADDITIONAL_INFO_FIELD,
additionalInfos == null ? Map.of() : additionalInfos
);
Map<String, Object> sourceMap = new HashMap<>();
sourceMap.put(ConversationalIndexConstants.META_CREATED_TIME_FIELD, now);
sourceMap.put(ConversationalIndexConstants.META_UPDATED_TIME_FIELD, now);
if (name != null && !name.trim().isEmpty()) {
sourceMap.put(ConversationalIndexConstants.META_NAME_FIELD, name);
}
if (userstr != null && !userstr.trim().isEmpty()) {
rithin-pullela-aws marked this conversation as resolved.
Show resolved Hide resolved
sourceMap.put(ConversationalIndexConstants.USER_FIELD, User.parse(userstr).getName());
}
if (applicationType != null && !applicationType.trim().isEmpty()) {
sourceMap.put(ConversationalIndexConstants.APPLICATION_TYPE_FIELD, applicationType);
}
if (additionalInfos != null && !additionalInfos.isEmpty()) {
sourceMap.put(ConversationalIndexConstants.META_ADDITIONAL_INFO_FIELD, additionalInfos);
}
IndexRequest request = Requests.indexRequest(META_INDEX_NAME).source(sourceMap);
try (ThreadContext.StoredContext threadContext = client.threadPool().getThreadContext().stashContext()) {
ActionListener<String> internalListener = ActionListener.runBefore(listener, () -> threadContext.restore());
ActionListener<IndexResponse> al = ActionListener.wrap(resp -> {
Expand Down
45 changes: 23 additions & 22 deletions memory/src/main/java/org/opensearch/ml/memory/index/InteractionsIndex.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@

import java.io.IOException;
import java.time.Instant;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
Expand Down Expand Up @@ -161,28 +162,28 @@ public void createInteraction(
if (indexExists) {
this.conversationMetaIndex.checkAccess(conversationId, ActionListener.wrap(access -> {
if (access) {
IndexRequest request = Requests
.indexRequest(INTERACTIONS_INDEX_NAME)
.source(
ConversationalIndexConstants.INTERACTIONS_ORIGIN_FIELD,
origin,
ConversationalIndexConstants.INTERACTIONS_CONVERSATION_ID_FIELD,
conversationId,
ConversationalIndexConstants.INTERACTIONS_INPUT_FIELD,
input,
ConversationalIndexConstants.INTERACTIONS_PROMPT_TEMPLATE_FIELD,
promptTemplate,
ConversationalIndexConstants.INTERACTIONS_RESPONSE_FIELD,
response,
ConversationalIndexConstants.INTERACTIONS_ADDITIONAL_INFO_FIELD,
additionalInfo,
ConversationalIndexConstants.INTERACTIONS_CREATE_TIME_FIELD,
timestamp,
ConversationalIndexConstants.PARENT_INTERACTIONS_ID_FIELD,
parintid,
ConversationalIndexConstants.INTERACTIONS_TRACE_NUMBER_FIELD,
traceNumber
);
Map<String, Object> sourceMap = new HashMap<>();
sourceMap.put(ConversationalIndexConstants.INTERACTIONS_CONVERSATION_ID_FIELD, conversationId);
sourceMap.put(ConversationalIndexConstants.INTERACTIONS_CREATE_TIME_FIELD, timestamp);
sourceMap.put(ConversationalIndexConstants.PARENT_INTERACTIONS_ID_FIELD, parintid);
rithin-pullela-aws marked this conversation as resolved.
Show resolved Hide resolved
sourceMap.put(ConversationalIndexConstants.INTERACTIONS_TRACE_NUMBER_FIELD, traceNumber);

if (input != null && !input.trim().isEmpty()) {
sourceMap.put(ConversationalIndexConstants.INTERACTIONS_INPUT_FIELD, input);
}
if (promptTemplate != null && !promptTemplate.trim().isEmpty()) {
sourceMap.put(ConversationalIndexConstants.INTERACTIONS_PROMPT_TEMPLATE_FIELD, promptTemplate);
}
if (response != null && !response.trim().isEmpty()) {
sourceMap.put(ConversationalIndexConstants.INTERACTIONS_RESPONSE_FIELD, response);
}
if (origin != null && !origin.trim().isEmpty()) {
sourceMap.put(ConversationalIndexConstants.INTERACTIONS_ORIGIN_FIELD, origin);
}
if (additionalInfo != null && !additionalInfo.isEmpty()) {
sourceMap.put(ConversationalIndexConstants.INTERACTIONS_ADDITIONAL_INFO_FIELD, additionalInfo);
}
IndexRequest request = Requests.indexRequest(INTERACTIONS_INDEX_NAME).source(sourceMap);
try (ThreadContext.StoredContext threadContext = client.threadPool().getThreadContext().stashContext()) {
ActionListener<String> internalListener = ActionListener.runBefore(listener, () -> threadContext.restore());
ActionListener<IndexResponse> al = ActionListener.wrap(resp -> {
Expand Down
21 changes: 21 additions & 0 deletions ...est/java/org/opensearch/ml/memory/action/conversation/CreateConversationRequestTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import org.junit.Before;
import org.junit.Rule;
import org.junit.rules.ExpectedException;
import org.opensearch.OpenSearchParseException;
import org.opensearch.common.io.stream.BytesStreamOutput;
import org.opensearch.core.common.bytes.BytesArray;
import org.opensearch.core.common.bytes.BytesReference;
Expand Down Expand Up @@ -132,4 +133,24 @@ public void testRestRequest_WithAdditionalInfo() throws IOException {
Assert.assertEquals("value1", request.getAdditionalInfos().get("key1"));
Assert.assertEquals(123, request.getAdditionalInfos().get("key2"));
}

public void testRestRequest_WithUnknownFields_Fails() throws IOException {
rithin-pullela-aws marked this conversation as resolved.
Show resolved Hide resolved
String name = "test-name";
RestRequest req = new FakeRestRequest.Builder(NamedXContentRegistry.EMPTY)
.withContent(
new BytesArray(gson.toJson(Map.of(ActionConstants.REQUEST_CONVERSATION_NAME_FIELD, name, "unknown_field", "some value"))),
MediaTypeRegistry.JSON
)
.build();

try {
CreateConversationRequest request = CreateConversationRequest.fromRestRequest(req);
fail("Expected IllegalArgumentException due to unknown field");
} catch (OpenSearchParseException e) {
assertEquals(e.getMessage(), "Invalid field [unknown_field] found in request body");
} catch (Exception e) {
fail("Expected OpenSearchParseException due to unknown field, got " + e.getClass().getName());
}

}
}
60 changes: 60 additions & 0 deletions ...test/java/org/opensearch/ml/memory/action/conversation/CreateInteractionRequestTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

import org.junit.Before;
Expand Down Expand Up @@ -153,4 +154,63 @@ public void testFromRestRequest_Trace() throws IOException {
assert (request.getParentIid().equals("parentId"));
assert (request.getTraceNumber().equals(1));
}

public void testRestRequest_WithUnknownFields_Fails() throws IOException {
Map<String, Object> params = Map
.of(
ActionConstants.INPUT_FIELD,
"input",
ActionConstants.PROMPT_TEMPLATE_FIELD,
"pt",
ActionConstants.AI_RESPONSE_FIELD,
"response",
ActionConstants.RESPONSE_ORIGIN_FIELD,
"origin",
ActionConstants.ADDITIONAL_INFO_FIELD,
Collections.singletonMap("metadata", "some meta"),
"unknown_field",
"some value"
);

RestRequest rrequest = new FakeRestRequest.Builder(NamedXContentRegistry.EMPTY)
.withParams(Map.of(ActionConstants.MEMORY_ID, "cid"))
.withContent(new BytesArray(gson.toJson(params)), MediaTypeRegistry.JSON)
.build();

try {
rithin-pullela-aws marked this conversation as resolved.
Show resolved Hide resolved
CreateInteractionRequest request = CreateInteractionRequest.fromRestRequest(rrequest);
fail("Expected IllegalArgumentException due to unknown field");
} catch (IllegalArgumentException e) {
assertEquals(e.getMessage(), "Invalid field [unknown_field] found in request body");
} catch (Exception e) {
fail("Expected IllegalArgumentException due to unknown field, got " + e.getClass().getName());
}
}

public void testFromRestRequest_WithAllFieldsEmpty_Fails() throws IOException {
Map<String, Object> params = new HashMap<>();

params.put(ActionConstants.INPUT_FIELD, "");
params.put(ActionConstants.PROMPT_TEMPLATE_FIELD, null);
params.put(ActionConstants.AI_RESPONSE_FIELD, " ");
params.put(ActionConstants.RESPONSE_ORIGIN_FIELD, null);
params.put(ActionConstants.ADDITIONAL_INFO_FIELD, Collections.emptyMap());

RestRequest rrequest = new FakeRestRequest.Builder(NamedXContentRegistry.EMPTY)
.withParams(Map.of(ActionConstants.MEMORY_ID, "cid"))
.withContent(new BytesArray(gson.toJson(params)), MediaTypeRegistry.JSON)
.build();

try {
CreateInteractionRequest request = CreateInteractionRequest.fromRestRequest(rrequest);
fail("Expected IllegalArgumentException due to all fields empty");
} catch (IllegalArgumentException e) {
assertEquals(
e.getMessage(),
"At least one of the following parameters must be non-empty: input, prompt_template, response, origin, additional_info"
);
} catch (Exception e) {
fail("Expected IllegalArgumentException due to all fields empty, got " + e.getClass().getName());
}
}
}
4 changes: 2 additions & 2 deletions memory/src/test/java/org/opensearch/ml/memory/index/ConversationMetaIndexITTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -564,8 +564,8 @@ public void testCanGetAConversationById() {
assert (cid2.result().equals(get2.result().getId()));
assert (get1.result().getName().equals("convo1"));
assert (get2.result().getName().equals("convo2"));
Assert.assertTrue(convo2.getAdditionalInfos().isEmpty());
Assert.assertTrue(get1.result().getAdditionalInfos().isEmpty());
Assert.assertTrue(convo2.getAdditionalInfos() == null);
Assert.assertTrue(get1.result().getAdditionalInfos() == null);
cdl.countDown();
}, e -> {
cdl.countDown();
Expand Down
Loading