fix: fixed computation of assembled_response (#648)

epam · Jan 21, 2025 · 299c4fc · 299c4fc
1 parent 0964e02
commit 299c4fc
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/server/src/main/java/com/epam/aidial/core/server/log/GfLogStore.java b/server/src/main/java/com/epam/aidial/core/server/log/GfLogStore.java
@@ -325,7 +325,20 @@ static String assembleStreamingResponse(@Nullable Buffer response) {
             }
 
             if (choices != null) {
-                MergeChunks.removeIndices(choices);
+                if (choices.isArray()) {
+                    for (JsonNode choice : choices) {
+                        MergeChunks.removeIndices(choice);
+                        if (choice.isObject()) {
+                            ObjectNode choiceObj = (ObjectNode) choice;
+                            JsonNode delta = choiceObj.get("delta");
+                            if (delta != null) {
+                                choiceObj.set("message", delta);
+                                choiceObj.remove("delta");
+                            }
+                        }
+                    }
+                }
+
                 result.set("choices", choices);
             }
             return ProxyUtil.convertToString(result);

diff --git a/server/src/test/java/com/epam/aidial/core/server/log/GfLogStoreTest.java b/server/src/test/java/com/epam/aidial/core/server/log/GfLogStoreTest.java
@@ -170,7 +170,7 @@ public void testAssembleStreamingResponse() {
         String res = GfLogStore.assembleStreamingResponse(Buffer.buffer(streamingResponse));
         assertNotNull(res);
         String expected = """
-                {"id":"1d84aa54-e476-405d-9713-386bdfc85993","object":"chat.completion","created":"1687222196","model":"gpt-35-turbo","usage":{"junk_string":"junk","junk_integer":1,"junk_float":1.0,"junk_null":null,"junk_true":true,"junk_false":false,"completion_tokens":10,"prompt_tokens":20,"total_tokens":30},"statistics":{"usage_per_model":[{"name":"text-embedding-ada-002","prompt_tokens":23,"total_tokens":23},{"name":"gpt-4","prompt_tokens":123,"completion_tokens":17,"total_tokens":140}]},"choices":[{"finish_reason":"stop","delta":{"role":"assistant","content":"As an AI language model, I don't have emotions, but I'm functioning perfectly well. How can I assist you today?","custom_content":{"attachments":[{"url":"url2"},{"url":"url1"}],"stages":[{"name":"stage1","status":"completed"},{"name":"stage2","status":"completed"}],"controls":[{"label":"label1"},{"label":"label2"}],"state":{"p1":1,"p2":1}}}}]}""";
+                {"id":"1d84aa54-e476-405d-9713-386bdfc85993","object":"chat.completion","created":"1687222196","model":"gpt-35-turbo","usage":{"junk_string":"junk","junk_integer":1,"junk_float":1.0,"junk_null":null,"junk_true":true,"junk_false":false,"completion_tokens":10,"prompt_tokens":20,"total_tokens":30},"statistics":{"usage_per_model":[{"name":"text-embedding-ada-002","prompt_tokens":23,"total_tokens":23},{"name":"gpt-4","prompt_tokens":123,"completion_tokens":17,"total_tokens":140}]},"choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"As an AI language model, I don't have emotions, but I'm functioning perfectly well. How can I assist you today?","custom_content":{"attachments":[{"url":"url2"},{"url":"url1"}],"stages":[{"name":"stage1","status":"completed"},{"name":"stage2","status":"completed"}],"controls":[{"label":"label1"},{"label":"label2"}],"state":{"p1":1,"p2":1}}}}]}""";
         assertEquals(expected, res);
     }
 
@@ -194,9 +194,10 @@ public void testAssembleStreamingResponse2() {
                 """;
 
         String res = GfLogStore.assembleStreamingResponse(Buffer.buffer(streamingResponse));
-        System.out.println(res);
-
-
+        assertNotNull(res);
+        String expected = """
+                {"id":"3c9c699a-d1ef-4ec2-82ff-47a07206fa99","object":"chat.completion","created":1724242846,"model":null,"choices":[{"index":0,"finish_reason":null,"message":{"role":"assistant","custom_content":{"attachments":[{"type":"text/markdown","title":"[0] 'Architecture'","data":"data","reference_url":"url1"},{"type":"text/markdown","title":"[1] 'User Guide'","data":"data","reference_url":"url2"},{"type":"text/markdown","title":"[2] 'Knowledge Base'","data":"data","reference_url":"url3"},{"type":"text/markdown","title":"[3] 'Documentation'","data":"you can pick one of three formats to copy its data: CSV, Markdown or Text.\\n\\n","reference_url":"url4"}]},"content":"A B C"}}]}""";
+        assertEquals(expected, res);
     }
 
     @Test