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 index 848a1a806..cf6b11804 100644 --- 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 @@ -14,6 +14,7 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.common.annotations.VisibleForTesting; import io.netty.buffer.ByteBufInputStream; import io.vertx.core.Vertx; import io.vertx.core.buffer.Buffer; @@ -105,13 +106,10 @@ private void append(ProxyContext context, LogEntry entry) throws JsonProcessingE append(entry, context.getDeployment().getName(), true); append(entry, "\"", false); - String sourceDeployment = context.getSourceDeployment(); - if (sourceDeployment != null) { - String initialDeployment = context.getInitialDeployment(); + String parentDeployment = getParentDeployment(context); + if (parentDeployment != null) { append(entry, ",\"parent_deployment\":\"", false); - // if deployment(callee) is configured with interceptors the return initial deployment(which triggers interceptors) - // otherwise return source deployment(caller) - append(entry, initialDeployment != null ? initialDeployment : sourceDeployment, true); + append(entry, parentDeployment, true); append(entry, "\"", false); } @@ -341,4 +339,24 @@ static boolean isStreamingResponse(@Nullable Buffer response) { } return j == dataToken.length(); } + + @VisibleForTesting + static String getParentDeployment(ProxyContext context) { + List interceptors = context.getInterceptors(); + if (interceptors == null) { + return context.getSourceDeployment(); + } + // skip interceptors and return the deployment which called the current one + List executionPath = context.getExecutionPath(); + int i = executionPath.size() - 2; + for (int j = interceptors.size() - 1; i >= 0 && j >= 0; i--, j--) { + String deployment = executionPath.get(i); + String interceptor = interceptors.get(j); + if (!deployment.equals(interceptor)) { + log.warn("Can't find parent deployment because interceptor path doesn't match: expected - {}, actual - {}", interceptor, deployment); + return null; + } + } + return i < 0 ? null : executionPath.get(i); + } } 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 index a2208d3aa..8eec99f30 100644 --- 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 @@ -1,12 +1,18 @@ package com.epam.aidial.core.server.log; +import com.epam.aidial.core.server.ProxyContext; import io.vertx.core.buffer.Buffer; import org.junit.jupiter.api.Test; +import java.util.List; + import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; @SuppressWarnings("checkstyle:LineLength") public class GfLogStoreTest { @@ -192,4 +198,58 @@ public void testAssembleStreamingResponse2() { } + + @Test + public void testGetParentDeployment_NoInterceptors() { + ProxyContext context = mock(ProxyContext.class); + // app calls model without interceptors + when(context.getInterceptors()).thenReturn(null); + when(context.getSourceDeployment()).thenReturn("app"); + + String result = GfLogStore.getParentDeployment(context); + + assertEquals("app", result); + } + + @Test + public void testGetParentDeployment_DeploymentWithInterceptors1() { + ProxyContext context = mock(ProxyContext.class); + // app calls model with interceptors + List interceptors = List.of("interceptor1", "interceptor2"); + when(context.getInterceptors()).thenReturn(interceptors); + List executionPath = List.of("app", "interceptor1", "interceptor2", "model"); + when(context.getExecutionPath()).thenReturn(executionPath); + + String result = GfLogStore.getParentDeployment(context); + + assertEquals("app", result); + } + + @Test + public void testGetParentDeployment_DeploymentWithInterceptors2() { + ProxyContext context = mock(ProxyContext.class); + // chat calls model with interceptors + List interceptors = List.of("interceptor1", "interceptor2"); + when(context.getInterceptors()).thenReturn(interceptors); + List executionPath = List.of("interceptor1", "interceptor2", "model"); + when(context.getExecutionPath()).thenReturn(executionPath); + + String result = GfLogStore.getParentDeployment(context); + + assertNull(result); + } + + @Test + public void testGetParentDeployment_InterceptorPathMismatch() { + ProxyContext context = mock(ProxyContext.class); + // app calls model with interceptors but interceptor1 calls some dep1 in the middle using the same per request key + List interceptors = List.of("interceptor1", "interceptor2"); + when(context.getInterceptors()).thenReturn(interceptors); + List executionPath = List.of("app", "interceptor1", "dep1", "interceptor2", "model"); + when(context.getExecutionPath()).thenReturn(executionPath); + + String result = GfLogStore.getParentDeployment(context); + + assertNull(result); + } }