Skip to content

Commit

Permalink
Add application_type to ConversationMeta; update tests
Browse files Browse the repository at this point in the history
Modify getMemory(Conversation) to return the application_type parameter.
Include application_type in the ConversationMeta data model.
Update existing tests to validate the new parameter.
  • Loading branch information
rithin-pullela-aws committed Dec 16, 2024
1 parent c2a40c1 commit a841543
Show file tree
Hide file tree
Showing 7 changed files with 40 additions and 20 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,8 @@ public class ConversationMeta implements Writeable, ToXContentObject {
@Getter
private String user;
@Getter
private String applicationType;
@Getter
private Map<String, String> additionalInfos;

/**
Expand All @@ -74,8 +76,9 @@ public static ConversationMeta fromMap(String id, Map<String, Object> docFields)
Instant updated = Instant.parse((String) docFields.get(ConversationalIndexConstants.META_UPDATED_TIME_FIELD));
String name = (String) docFields.get(ConversationalIndexConstants.META_NAME_FIELD);
String user = (String) docFields.get(ConversationalIndexConstants.USER_FIELD);
String applicationType = (String) docFields.get(ConversationalIndexConstants.APPLICATION_TYPE_FIELD);
Map<String, String> additionalInfos = (Map<String, String>) docFields.get(ConversationalIndexConstants.META_ADDITIONAL_INFO_FIELD);
return new ConversationMeta(id, created, updated, name, user, additionalInfos);
return new ConversationMeta(id, created, updated, name, user, applicationType, additionalInfos);
}

/**
Expand All @@ -91,13 +94,14 @@ public static ConversationMeta fromStream(StreamInput in) throws IOException {
Instant updated = in.readInstant();
String name = in.readString();
String user = in.readOptionalString();
String applicationType = in.readOptionalString();
Map<String, String> additionalInfos = null;
if (in.getVersion().onOrAfter(MINIMAL_SUPPORTED_VERSION_FOR_ADDITIONAL_INFO)) {
if (in.readBoolean()) {
additionalInfos = in.readMap(StreamInput::readString, StreamInput::readString);
}
}
return new ConversationMeta(id, created, updated, name, user, additionalInfos);
return new ConversationMeta(id, created, updated, name, user, applicationType, additionalInfos);
}

@Override
Expand All @@ -107,6 +111,7 @@ public void writeTo(StreamOutput out) throws IOException {
out.writeInstant(updatedTime);
out.writeString(name);
out.writeOptionalString(user);
out.writeOptionalString(applicationType);
if (out.getVersion().onOrAfter(MINIMAL_SUPPORTED_VERSION_FOR_ADDITIONAL_INFO)) {
if (additionalInfos == null) {
out.writeBoolean(false);
Expand All @@ -129,6 +134,10 @@ public String toString() {
+ updatedTime.toString()
+ ", user="
+ user
+ ", applicationType="
+ applicationType
+ ", additionalInfos="
+ additionalInfos
+ "}";
}

Expand All @@ -142,7 +151,10 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContentObject.Para
if (this.user != null) {
builder.field(ConversationalIndexConstants.USER_FIELD, this.user);
}
if (this.additionalInfos != null) {
if (this.applicationType != null && !this.applicationType.trim().isEmpty()) {
builder.field(ConversationalIndexConstants.APPLICATION_TYPE_FIELD, this.applicationType);
}
if (this.additionalInfos != null && !additionalInfos.isEmpty()) {
builder.field(ConversationalIndexConstants.META_ADDITIONAL_INFO_FIELD, this.additionalInfos);
}
builder.endObject();
Expand All @@ -159,7 +171,9 @@ public boolean equals(Object other) {
&& Objects.equals(this.user, otherConversation.user)
&& Objects.equals(this.createdTime, otherConversation.createdTime)
&& Objects.equals(this.updatedTime, otherConversation.updatedTime)
&& Objects.equals(this.name, otherConversation.name);
&& Objects.equals(this.name, otherConversation.name)
&& Objects.equals(this.applicationType, otherConversation.applicationType)
&& Objects.equals(this.additionalInfos, otherConversation.additionalInfos);
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ public class ConversationMetaTests {
@Before
public void setUp() {
time = Instant.now();
conversationMeta = new ConversationMeta("test_id", time, time, "test_name", "admin", null);
conversationMeta = new ConversationMeta("test_id", time, time, "test_name", "admin", "conversational-search", null);
}

@Test
Expand All @@ -41,6 +41,7 @@ public void test_fromSearchHit() throws IOException {
content.field(ConversationalIndexConstants.META_UPDATED_TIME_FIELD, time);
content.field(ConversationalIndexConstants.META_NAME_FIELD, "meta name");
content.field(ConversationalIndexConstants.USER_FIELD, "admin");
content.field(ConversationalIndexConstants.APPLICATION_TYPE_FIELD, "conversational-search");
content.field(ConversationalIndexConstants.META_ADDITIONAL_INFO_FIELD, Map.of("test_key", "test_value"));
content.endObject();

Expand All @@ -51,6 +52,7 @@ public void test_fromSearchHit() throws IOException {
assertEquals(conversationMeta.getId(), "cId");
assertEquals(conversationMeta.getName(), "meta name");
assertEquals(conversationMeta.getUser(), "admin");
assertEquals(conversationMeta.getApplicationType(), "conversational-search");
assertEquals(conversationMeta.getAdditionalInfos().get("test_key"), "test_value");
}

Expand Down Expand Up @@ -83,6 +85,7 @@ public void test_fromStream() throws IOException {
assertEquals(meta.getId(), conversationMeta.getId());
assertEquals(meta.getName(), conversationMeta.getName());
assertEquals(meta.getUser(), conversationMeta.getUser());
assertEquals(meta.getApplicationType(), conversationMeta.getApplicationType());
}

@Test
Expand All @@ -93,14 +96,15 @@ public void test_ToXContent() throws IOException {
Instant.ofEpochMilli(123),
"test meta",
"admin",
"neural-search",
null
);
XContentBuilder builder = XContentBuilder.builder(XContentType.JSON.xContent());
conversationMeta.toXContent(builder, EMPTY_PARAMS);
String content = TestHelper.xContentBuilderToString(builder);
assertEquals(
content,
"{\"memory_id\":\"test_id\",\"create_time\":\"1970-01-01T00:00:00.123Z\",\"updated_time\":\"1970-01-01T00:00:00.123Z\",\"name\":\"test meta\",\"user\":\"admin\"}"
"{\"memory_id\":\"test_id\",\"create_time\":\"1970-01-01T00:00:00.123Z\",\"updated_time\":\"1970-01-01T00:00:00.123Z\",\"name\":\"test meta\",\"user\":\"admin\",\"application_type\":\"neural-search\"}"
);
}

Expand All @@ -112,10 +116,11 @@ public void test_toString() {
Instant.ofEpochMilli(123),
"test meta",
"admin",
"conversational-search",
null
);
assertEquals(
"{id=test_id, name=test meta, created=1970-01-01T00:00:00.123Z, updated=1970-01-01T00:00:00.123Z, user=admin}",
"{id=test_id, name=test meta, created=1970-01-01T00:00:00.123Z, updated=1970-01-01T00:00:00.123Z, user=admin, applicationType=conversational-search, additionalInfos=null}",
conversationMeta.toString()
);
}
Expand All @@ -128,6 +133,7 @@ public void test_equal() {
Instant.ofEpochMilli(123),
"test meta",
"admin",
"conversational-search",
null
);
assertEquals(meta.equals(conversationMeta), false);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@
public class GetConversationResponseTests extends OpenSearchTestCase {

public void testGetConversationResponseStreaming() throws IOException {
ConversationMeta convo = new ConversationMeta("cid", Instant.now(), Instant.now(), "name", null, null);
ConversationMeta convo = new ConversationMeta("cid", Instant.now(), Instant.now(), "name", null, null, null);
GetConversationResponse response = new GetConversationResponse(convo);
assert (response.getConversation().equals(convo));

Expand All @@ -51,7 +51,7 @@ public void testGetConversationResponseStreaming() throws IOException {
}

public void testToXContent() throws IOException {
ConversationMeta convo = new ConversationMeta("cid", Instant.now(), Instant.now(), "name", null, null);
ConversationMeta convo = new ConversationMeta("cid", Instant.now(), Instant.now(), "name", null, null, null);
GetConversationResponse response = new GetConversationResponse(convo);
XContentBuilder builder = XContentBuilder.builder(XContentType.JSON.xContent());
response.toXContent(builder, ToXContent.EMPTY_PARAMS);
Expand All @@ -68,7 +68,7 @@ public void testToXContent() throws IOException {

public void testToXContent_withAdditionalInfo() throws IOException {
Map<String, String> additionalInfos = Map.of("key1", "value1");
ConversationMeta convo = new ConversationMeta("cid", Instant.now(), Instant.now(), "name", null, additionalInfos);
ConversationMeta convo = new ConversationMeta("cid", Instant.now(), Instant.now(), "name", null, null, additionalInfos);
GetConversationResponse response = new GetConversationResponse(convo);
XContentBuilder builder = XContentBuilder.builder(XContentType.JSON.xContent());
response.toXContent(builder, ToXContent.EMPTY_PARAMS);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ public void setup() throws IOException {
}

public void testGetConversation() {
ConversationMeta result = new ConversationMeta("test-cid", Instant.now(), Instant.now(), "name", null, null);
ConversationMeta result = new ConversationMeta("test-cid", Instant.now(), Instant.now(), "name", null, null, null);
doAnswer(invocation -> {
ActionListener<ConversationMeta> listener = invocation.getArgument(1);
listener.onResponse(result);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -46,9 +46,9 @@ public class GetConversationsResponseTests extends OpenSearchTestCase {
public void setup() {
conversations = List
.of(
new ConversationMeta("0", Instant.now(), Instant.now(), "name0", "user0", null),
new ConversationMeta("1", Instant.now(), Instant.now(), "name1", "user0", null),
new ConversationMeta("2", Instant.now(), Instant.now(), "name2", "user2", null)
new ConversationMeta("0", Instant.now(), Instant.now(), "name0", "user0", null, null),
new ConversationMeta("1", Instant.now(), Instant.now(), "name1", "user0", null, null),
new ConversationMeta("2", Instant.now(), Instant.now(), "name2", "user2", null, null)
);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -114,8 +114,8 @@ public void testGetConversations() {
log.info("testing get conversations transport");
List<ConversationMeta> testResult = List
.of(
new ConversationMeta("testcid1", Instant.now(), Instant.now(), "", null, null),
new ConversationMeta("testcid2", Instant.now(), Instant.now(), "testname", null, null)
new ConversationMeta("testcid1", Instant.now(), Instant.now(), "", null, null, null),
new ConversationMeta("testcid2", Instant.now(), Instant.now(), "testname", null, null, null)
);
doAnswer(invocation -> {
ActionListener<List<ConversationMeta>> listener = invocation.getArgument(2);
Expand All @@ -132,9 +132,9 @@ public void testGetConversations() {
public void testPagination() {
List<ConversationMeta> testResult = List
.of(
new ConversationMeta("testcid1", Instant.now(), Instant.now(), "", null, null),
new ConversationMeta("testcid2", Instant.now(), Instant.now(), "testname", null, null),
new ConversationMeta("testcid3", Instant.now(), Instant.now(), "testname", null, null)
new ConversationMeta("testcid1", Instant.now(), Instant.now(), "", null, null, null),
new ConversationMeta("testcid2", Instant.now(), Instant.now(), "testname", null, null, null),
new ConversationMeta("testcid3", Instant.now(), Instant.now(), "testname", null, null, null)
);
doAnswer(invocation -> {
ActionListener<List<ConversationMeta>> listener = invocation.getArgument(2);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -315,7 +315,7 @@ public void testSearchInteractions_Future() {
}

public void testGetAConversation_Future() {
ConversationMeta response = new ConversationMeta("cid", Instant.now(), Instant.now(), "boring name", null, null);
ConversationMeta response = new ConversationMeta("cid", Instant.now(), Instant.now(), "boring name", null, null, null);
doAnswer(invocation -> {
ActionListener<ConversationMeta> listener = invocation.getArgument(1);
listener.onResponse(response);
Expand Down

0 comments on commit a841543

Please sign in to comment.