From 289644be6c8c75652ff6b850a74ff8c876e69263 Mon Sep 17 00:00:00 2001 From: Aliaksandr Stsiapanay Date: Tue, 7 May 2024 11:48:52 +0200 Subject: [PATCH] =?UTF-8?q?fix:=20Collect=20the=20last=20usage=20entry=20i?= =?UTF-8?q?n=20the=20model=20response=20instead=20of=20th=E2=80=A6=20(#330?= =?UTF-8?q?)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Aliaksandr Stsiapanay --- .../aidial/core/token/TokenUsageParserTest.java | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/test/java/com/epam/aidial/core/token/TokenUsageParserTest.java b/src/test/java/com/epam/aidial/core/token/TokenUsageParserTest.java index 14a603936..f2191d64f 100644 --- 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);