diff --git a/server/src/main/java/com/epam/aidial/core/server/function/CollectChatCompletionUsageFn.java b/server/src/main/java/com/epam/aidial/core/server/function/CollectChatCompletionUsageFn.java index 45c824c01..d1f65281f 100644 --- a/server/src/main/java/com/epam/aidial/core/server/function/CollectChatCompletionUsageFn.java +++ b/server/src/main/java/com/epam/aidial/core/server/function/CollectChatCompletionUsageFn.java @@ -26,6 +26,10 @@ public CollectChatCompletionUsageFn(Proxy proxy, ProxyContext context) { @Override public Future apply(JsonNode tree) { + if (tree == null) { + return Future.succeededFuture(tree); + } + usage = MergeChunks.merge(usage, tree.get("usage")); if (tree.get("service_tier") != null) { serviceTier = tree.get("service_tier"); diff --git a/server/src/main/java/com/epam/aidial/core/server/function/CollectResponseAttachmentsFn.java b/server/src/main/java/com/epam/aidial/core/server/function/CollectResponseAttachmentsFn.java index 2ec5afc39..3041ebf93 100644 --- a/server/src/main/java/com/epam/aidial/core/server/function/CollectResponseAttachmentsFn.java +++ b/server/src/main/java/com/epam/aidial/core/server/function/CollectResponseAttachmentsFn.java @@ -38,7 +38,7 @@ public Future apply(JsonNode tree) { processAttachedFile(attachment, permittedAttachments); } if (permittedAttachments.isEmpty()) { - return Future.succeededFuture(); + return Future.succeededFuture(tree); } return proxy.getTaskExecutor().submit(() -> { proxy.getApiKeyStore().updatePerRequestApiKey(perRequestKey, json -> updateAutoSharedAttachments(json, permittedAttachments, perRequestKey)); diff --git a/server/src/main/java/com/epam/aidial/core/server/vertx/stream/BufferingReadStream.java b/server/src/main/java/com/epam/aidial/core/server/vertx/stream/BufferingReadStream.java index aa528251b..223fa60cf 100644 --- a/server/src/main/java/com/epam/aidial/core/server/vertx/stream/BufferingReadStream.java +++ b/server/src/main/java/com/epam/aidial/core/server/vertx/stream/BufferingReadStream.java @@ -267,10 +267,10 @@ public void onComplete() { @SneakyThrows private Future handle(SseEvent event) { Future result; - if (functions.isEmpty() || skipEvent(event)) { + String data = event.getData(); + if (functions.isEmpty() || skipEvent(event) || data == null || data.isBlank()) { result = Future.succeededFuture(); } else { - String data = event.getData(); try { JsonNode tree = ProxyUtil.MAPPER.readTree(data); result = Future.succeededFuture(tree); diff --git a/server/src/test/java/com/epam/aidial/core/server/function/ChatCompletionFunctionChainTest.java b/server/src/test/java/com/epam/aidial/core/server/function/ChatCompletionFunctionChainTest.java new file mode 100644 index 000000000..f80b2ada2 --- /dev/null +++ b/server/src/test/java/com/epam/aidial/core/server/function/ChatCompletionFunctionChainTest.java @@ -0,0 +1,75 @@ +package com.epam.aidial.core.server.function; + +import com.epam.aidial.core.server.ProxyContext; +import com.epam.aidial.core.server.data.ApiKeyData; +import com.epam.aidial.core.server.util.ProxyUtil; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import io.vertx.core.Future; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doCallRealMethod; +import static org.mockito.Mockito.when; + +/** + * Exercises the real function chain wired by DeploymentPostController for streaming + * {@code /chat/completions} responses (StripUsagePerModelFn -> CollectResponseChatCompletionAttachmentsFn + * -> CollectChatCompletionUsageFn), composed the same way {@code BufferingReadStream.BaseEventListener} + * chains it via {@code Future.compose}. Regression test for a null tree silently produced by one + * function and blindly dereferenced by the next. + */ +@ExtendWith(MockitoExtension.class) +class ChatCompletionFunctionChainTest { + + @Mock + private ProxyContext context; + + @Test + public void testChainToleratesChunksWithNoAttachmentsAndCollectsUsage() throws JsonProcessingException { + ApiKeyData apiKeyData = new ApiKeyData(); + apiKeyData.setPerRequestKey("per-request-key"); + when(context.getApiKeyData()).thenReturn(apiKeyData); + when(context.isStreamingRequest()).thenReturn(true); + doCallRealMethod().when(context).setPricingUsageNode(any()); + doCallRealMethod().when(context).getPricingUsageNode(); + + List functions = List.of( + new StripUsagePerModelFn(null, context), + new CollectResponseChatCompletionAttachmentsFn(null, context), + new CollectChatCompletionUsageFn(null, context)); + + Future contentChunk = process(functions, """ + { "choices": [{ "index": 0, "delta": { "content": "hi" } }] } + """); + assertTrue(contentChunk.succeeded(), () -> "chain failed: " + contentChunk.cause()); + + Future usageChunk = process(functions, """ + { + "choices": [{ "index": 0, "delta": {} }], + "usage": { "completion_tokens": 33, "prompt_tokens": 19, "total_tokens": 52 } + } + """); + assertTrue(usageChunk.succeeded(), () -> "chain failed: " + usageChunk.cause()); + + JsonNode pricingUsageNode = context.getPricingUsageNode(); + assertNotNull(pricingUsageNode); + assertEquals(52, pricingUsageNode.path("usage").path("total_tokens").asLong()); + } + + private static Future process(List functions, String json) throws JsonProcessingException { + Future result = Future.succeededFuture(ProxyUtil.MAPPER.readTree(json)); + for (BaseResponseFunction fn : functions) { + result = result.compose(fn); + } + return result; + } +} diff --git a/server/src/test/java/com/epam/aidial/core/server/function/CollectChatCompletionUsageFnTest.java b/server/src/test/java/com/epam/aidial/core/server/function/CollectChatCompletionUsageFnTest.java index 5b5de82a5..9d917e8b1 100644 --- a/server/src/test/java/com/epam/aidial/core/server/function/CollectChatCompletionUsageFnTest.java +++ b/server/src/test/java/com/epam/aidial/core/server/function/CollectChatCompletionUsageFnTest.java @@ -11,6 +11,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doCallRealMethod; @@ -67,6 +68,15 @@ public void testPassesEventThroughUnchanged() throws JsonProcessingException { assertEquals(input, output); } + @Test + public void testTolerantOfNullTree() { + CollectChatCompletionUsageFn fn = new CollectChatCompletionUsageFn(null, context); + + JsonNode output = fn.apply(null).result(); + + assertNull(output); + } + private static JsonNode tree(String json) throws JsonProcessingException { return ProxyUtil.MAPPER.readTree(json); } diff --git a/server/src/test/java/com/epam/aidial/core/server/function/CollectResponseChatCompletionAttachmentsFnTest.java b/server/src/test/java/com/epam/aidial/core/server/function/CollectResponseChatCompletionAttachmentsFnTest.java index dd0246911..f32717717 100644 --- a/server/src/test/java/com/epam/aidial/core/server/function/CollectResponseChatCompletionAttachmentsFnTest.java +++ b/server/src/test/java/com/epam/aidial/core/server/function/CollectResponseChatCompletionAttachmentsFnTest.java @@ -1,8 +1,10 @@ package com.epam.aidial.core.server.function; import com.epam.aidial.core.server.ProxyContext; +import com.epam.aidial.core.server.data.ApiKeyData; import com.epam.aidial.core.server.util.ProxyUtil; import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.ObjectNode; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -12,6 +14,7 @@ import java.util.Set; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.mockito.Mockito.when; @@ -218,4 +221,21 @@ public void testCollectAttachmentsFromResponse_ChatStreamingResponse() throws Js "files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/b1/file2.txt"), files); } + + @Test + public void testApplyPassesTreeThroughUnchangedWhenNoAttachments() throws JsonProcessingException { + ApiKeyData apiKeyData = new ApiKeyData(); + apiKeyData.setPerRequestKey("per-request-key"); + when(context.getApiKeyData()).thenReturn(apiKeyData); + + JsonNode tree = ProxyUtil.MAPPER.readTree(""" + { + "choices": [{ "index": 0, "delta": { "content": "hi" } }] + } + """); + + JsonNode result = new CollectResponseChatCompletionAttachmentsFn(null, context).apply(tree).result(); + + assertSame(tree, result); + } } \ No newline at end of file