feat: Add customViewState field to conversation api (#599)

epam · Dec 4, 2024 · 99cfcab · 99cfcab
1 parent 6721d70
commit 99cfcab
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 1 deletion.
diff --git a/server/src/main/java/com/epam/aidial/core/server/data/Conversation.java b/server/src/main/java/com/epam/aidial/core/server/data/Conversation.java
@@ -6,6 +6,7 @@
 import lombok.Data;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 @Data
@@ -32,6 +33,7 @@ public class Conversation {
     ModelId model;
     Set<String> selectedAddons;
     List<Message> messages;
+    Map<String, Object> customViewState;
 
     @JsonCreator
     public Conversation(@JsonProperty(value = "id", required = true) String id,
@@ -42,7 +44,8 @@ public Conversation(@JsonProperty(value = "id", required = true) String id,
                         @JsonProperty(value = "lastActivityDate", required = true) long lastActivityDate,
                         @JsonProperty(value = "model", required = true) ModelId model,
                         @JsonProperty(value = "selectedAddons", required = true) Set<String> selectedAddons,
-                        @JsonProperty(value = "messages", required = true) List<Message> messages) {
+                        @JsonProperty(value = "messages", required = true) List<Message> messages,
+                        @JsonProperty(value = "customViewState") Map<String, Object> customViewState) {
         this.id = id;
         this.folderId = folderId;
         this.name = name;
@@ -52,6 +55,7 @@ public Conversation(@JsonProperty(value = "id", required = true) String id,
         this.model = model;
         this.selectedAddons = selectedAddons;
         this.messages = messages;
+        this.customViewState = customViewState;
     }
 
     @Data

diff --git a/server/src/test/java/com/epam/aidial/core/server/util/ProxyUtilTest.java b/server/src/test/java/com/epam/aidial/core/server/util/ProxyUtilTest.java
@@ -461,4 +461,51 @@ public void testCollectAttachmentsFromResponse_ChatStreamingResponse() throws Js
 
     }
 
+    @Test
+    public void testCustomViewStateValidation() {
+        String validConversationJson = """
+                {
+                "id": "conversation_id",
+                "name": "display_name",
+                "model": {
+                  "id": "model_id"
+                  },
+                "prompt": "system prompt",
+                "temperature": 1,
+                "folderId": "folder1",
+                "messages": [
+                  {
+                  "role": "user",
+                  "content": "content",
+                  "custom_content": {"attachment_url": "some_url"},
+                  "model": {"id": "model_id"},
+                  "settings":
+                    {
+                    "prompt": "sysPrompt",
+                    "temperature": 5,
+                    "selectedAddons": ["A", "B", "C"],
+                    "assistantModelId": "assistantId"
+                    }
+                  }
+                ],
+                "replay": {
+                  "isReplay": true,
+                  "replayUserMessagesStack": [],
+                  "activeReplayIndex": 0
+                  },
+                "selectedAddons": ["R", "T", "G"],
+                "assistantModelId": "assistantId",
+                "lastActivityDate": 4848683153,
+                "customViewState": {
+                    "a": ["A"],
+                    "b": {
+                        "c": ["C"],
+                        "d": 5.12
+                    }
+                }
+                }
+                """;
+
+        assertDoesNotThrow(() -> ProxyUtil.convertToObject(validConversationJson, Conversation.class));
+    }
 }