fix: Collect the last usage entry in the model response instead of th… (

#330) Co-authored-by: Aliaksandr Stsiapanay <aliaksandr_stsiapanay@epam.com>
epam · May 7, 2024 · 289644b · 289644b
1 parent ceeb8e2
commit 289644b
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/src/test/java/com/epam/aidial/core/token/TokenUsageParserTest.java b/src/test/java/com/epam/aidial/core/token/TokenUsageParserTest.java
@@ -139,6 +139,22 @@ void testValidStreamResponse() {
                 """, 10, 20, 30);
     }
 
+    @Test
+    void testValidStreamResponseWithMultipleUsages() {
+        valid("""
+                data: {"id":"eb69ae53-055b-4182-af8f-47f5f3ce810c","object":"chat.completion.chunk","created":1714665540,"model":"dbrx-instruct-032724","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}],"usage":{"prompt_tokens":226,"completion_tokens":1,"total_tokens":227}}
+                                
+                data: {"id":"eb69ae53-055b-4182-af8f-47f5f3ce810c","object":"chat.completion.chunk","created":1714665540,"model":"dbrx-instruct-032724","choices":[{"index":0,"delta":{"role":"assistant","content":" today"},"finish_reason":null}],"usage":{"prompt_tokens":226,"completion_tokens":25,"total_tokens":251}}
+                                
+                data: {"id":"eb69ae53-055b-4182-af8f-47f5f3ce810c","object":"chat.completion.chunk","created":1714665540,"model":"dbrx-instruct-032724","choices":[{"index":0,"delta":{"role":"assistant","content":"."},"finish_reason":null}],"usage":{"prompt_tokens":226,"completion_tokens":26,"total_tokens":252}}
+                                
+                data: {"id":"eb69ae53-055b-4182-af8f-47f5f3ce810c","object":"chat.completion.chunk","created":1714665540,"model":"dbrx-instruct-032724","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":"stop"}],"usage":{"prompt_tokens":226,"completion_tokens":26,"total_tokens":252}}
+                                
+                data: [DONE]
+                 
+                """, 26, 226, 252);
+    }
+
     private void valid(String body, long completion, long prompt, long total) {
         TokenUsage usage = TokenUsageParser.parse(Buffer.buffer(body));
         Assertions.assertNotNull(usage);