From 14dd21b62bcb1afbc9499530bd65d6b1c4e0c6fa Mon Sep 17 00:00:00 2001 From: Aliaksandr Stsiapanay Date: Wed, 20 Nov 2024 16:24:05 +0300 Subject: [PATCH 1/2] chore: refactor base request functions --- .../controller/DeploymentPostController.java | 16 +++++----- .../controller/InterceptorController.java | 17 ++++++----- .../server/function/BaseRequestFunction.java | 2 +- .../function/CollectRequestAttachmentsFn.java | 24 ++++----------- .../server/function/CollectRequestDataFn.java | 4 +-- .../ApplyDefaultDeploymentSettingsFn.java | 18 +----------- .../EnhanceAssistantRequestFn.java | 29 ++++--------------- .../enhancement/EnhanceModelRequestFn.java | 23 ++++----------- .../aidial/core/server/util/ProxyUtil.java | 10 +++---- .../function/CollectRequestDataFnTest.java | 12 ++++---- .../ApplyDefaultDeploymentSettingsFnTest.java | 10 +++---- 11 files changed, 56 insertions(+), 109 deletions(-) diff --git a/server/src/main/java/com/epam/aidial/core/server/controller/DeploymentPostController.java b/server/src/main/java/com/epam/aidial/core/server/controller/DeploymentPostController.java index 7d7ab98eb..4ee6c7819 100644 --- a/server/src/main/java/com/epam/aidial/core/server/controller/DeploymentPostController.java +++ b/server/src/main/java/com/epam/aidial/core/server/controller/DeploymentPostController.java @@ -267,14 +267,16 @@ void handleRequestBody(Buffer requestBody) { try (InputStream stream = new ByteBufInputStream(requestBody.getByteBuf())) { ObjectNode tree = (ObjectNode) ProxyUtil.MAPPER.readTree(stream); - Throwable error = ProxyUtil.processChain(tree, enhancementFunctions); - if (error != null) { - finalizeRequest(); - return; + if (ProxyUtil.processChain(tree, enhancementFunctions)) { + context.setRequestBody(Buffer.buffer(ProxyUtil.MAPPER.writeValueAsBytes(tree))); } - } catch (IOException e) { - respond(HttpStatus.BAD_REQUEST); - log.warn("Can't parse JSON request body. Trace: {}. Span: {}. Error:", + } catch (Throwable e) { + if (e instanceof HttpException httpException) { + respond(httpException.getStatus(), httpException); + } else { + respond(HttpStatus.BAD_REQUEST); + } + log.warn("Can't process JSON request body. Trace: {}. Span: {}. Error:", context.getTraceId(), context.getSpanId(), e); return; } diff --git a/server/src/main/java/com/epam/aidial/core/server/controller/InterceptorController.java b/server/src/main/java/com/epam/aidial/core/server/controller/InterceptorController.java index 260db525c..611c4255d 100644 --- a/server/src/main/java/com/epam/aidial/core/server/controller/InterceptorController.java +++ b/server/src/main/java/com/epam/aidial/core/server/controller/InterceptorController.java @@ -10,6 +10,7 @@ import com.epam.aidial.core.server.function.CollectResponseAttachmentsFn; import com.epam.aidial.core.server.util.ProxyUtil; import com.epam.aidial.core.server.vertx.stream.BufferingReadStream; +import com.epam.aidial.core.storage.http.HttpException; import com.epam.aidial.core.storage.http.HttpStatus; import com.fasterxml.jackson.databind.node.ObjectNode; import io.netty.buffer.ByteBufInputStream; @@ -69,14 +70,16 @@ private void handleRequestBody(Buffer requestBody) { context.setRequestBodyTimestamp(System.currentTimeMillis()); try (InputStream stream = new ByteBufInputStream(requestBody.getByteBuf())) { ObjectNode tree = (ObjectNode) ProxyUtil.MAPPER.readTree(stream); - Throwable error = ProxyUtil.processChain(tree, enhancementFunctions); - if (error != null) { - finalizeRequest(); - return; + if (ProxyUtil.processChain(tree, enhancementFunctions)) { + context.setRequestBody(Buffer.buffer(ProxyUtil.MAPPER.writeValueAsBytes(tree))); } - } catch (IOException e) { - respond(HttpStatus.BAD_REQUEST); - log.warn("Can't parse JSON request body. Trace: {}. Span: {}. Error:", + } catch (Throwable e) { + if (e instanceof HttpException httpException) { + respond(httpException.getStatus(), httpException); + } else { + respond(HttpStatus.BAD_REQUEST); + } + log.warn("Can't process JSON request body. Trace: {}. Span: {}. Error:", context.getTraceId(), context.getSpanId(), e); return; } diff --git a/server/src/main/java/com/epam/aidial/core/server/function/BaseRequestFunction.java b/server/src/main/java/com/epam/aidial/core/server/function/BaseRequestFunction.java index 65ef352a5..90d452a43 100644 --- a/server/src/main/java/com/epam/aidial/core/server/function/BaseRequestFunction.java +++ b/server/src/main/java/com/epam/aidial/core/server/function/BaseRequestFunction.java @@ -3,7 +3,7 @@ import com.epam.aidial.core.server.Proxy; import com.epam.aidial.core.server.ProxyContext; -public abstract class BaseRequestFunction extends BaseFunction { +public abstract class BaseRequestFunction extends BaseFunction { public BaseRequestFunction(Proxy proxy, ProxyContext context) { diff --git a/server/src/main/java/com/epam/aidial/core/server/function/CollectRequestAttachmentsFn.java b/server/src/main/java/com/epam/aidial/core/server/function/CollectRequestAttachmentsFn.java index b165ba7b2..420a538fa 100644 --- a/server/src/main/java/com/epam/aidial/core/server/function/CollectRequestAttachmentsFn.java +++ b/server/src/main/java/com/epam/aidial/core/server/function/CollectRequestAttachmentsFn.java @@ -26,24 +26,12 @@ public CollectRequestAttachmentsFn(Proxy proxy, ProxyContext context) { } @Override - public Throwable apply(ObjectNode tree) { - try { - ProxyUtil.collectAttachedFilesFromRequest(tree, this::processAttachedFile); - // assign api key data after processing attachments - ApiKeyData destApiKeyData = context.getProxyApiKeyData(); - proxy.getApiKeyStore().assignPerRequestApiKey(destApiKeyData); - return null; - } catch (HttpException e) { - context.respond(e.getStatus(), e.getMessage()); - log.warn("Can't collect attached files. Trace: {}. Span: {}. Error: {}", - context.getTraceId(), context.getSpanId(), e.getMessage()); - return e; - } catch (Throwable e) { - context.respond(HttpStatus.BAD_REQUEST); - log.warn("Can't collect attached files. Trace: {}. Span: {}. Error: {}", - context.getTraceId(), context.getSpanId(), e.getMessage()); - return e; - } + public Boolean apply(ObjectNode tree) { + ProxyUtil.collectAttachedFilesFromRequest(tree, this::processAttachedFile); + // assign api key data after processing attachments + ApiKeyData destApiKeyData = context.getProxyApiKeyData(); + proxy.getApiKeyStore().assignPerRequestApiKey(destApiKeyData); + return false; } private void processAttachedFile(String url) { diff --git a/server/src/main/java/com/epam/aidial/core/server/function/CollectRequestDataFn.java b/server/src/main/java/com/epam/aidial/core/server/function/CollectRequestDataFn.java index d9797a61f..e791b1b3d 100644 --- a/server/src/main/java/com/epam/aidial/core/server/function/CollectRequestDataFn.java +++ b/server/src/main/java/com/epam/aidial/core/server/function/CollectRequestDataFn.java @@ -11,10 +11,10 @@ public CollectRequestDataFn(Proxy proxy, ProxyContext context) { } @Override - public Throwable apply(ObjectNode tree) { + public Boolean apply(ObjectNode tree) { JsonNode stream = tree.get("stream"); boolean result = stream != null && stream.asBoolean(false); context.setStreamingRequest(result); - return null; + return false; } } diff --git a/server/src/main/java/com/epam/aidial/core/server/function/enhancement/ApplyDefaultDeploymentSettingsFn.java b/server/src/main/java/com/epam/aidial/core/server/function/enhancement/ApplyDefaultDeploymentSettingsFn.java index 1f62ff99f..bcfa290e4 100644 --- a/server/src/main/java/com/epam/aidial/core/server/function/enhancement/ApplyDefaultDeploymentSettingsFn.java +++ b/server/src/main/java/com/epam/aidial/core/server/function/enhancement/ApplyDefaultDeploymentSettingsFn.java @@ -5,10 +5,8 @@ import com.epam.aidial.core.server.ProxyContext; import com.epam.aidial.core.server.function.BaseRequestFunction; import com.epam.aidial.core.server.util.ProxyUtil; -import com.epam.aidial.core.storage.http.HttpStatus; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.ObjectNode; -import io.vertx.core.buffer.Buffer; import lombok.extern.slf4j.Slf4j; import java.util.Map; @@ -20,21 +18,7 @@ public ApplyDefaultDeploymentSettingsFn(Proxy proxy, ProxyContext context) { } @Override - public Throwable apply(ObjectNode tree) { - try { - if (applyDefaults(context, tree)) { - context.setRequestBody(Buffer.buffer(ProxyUtil.MAPPER.writeValueAsBytes(tree))); - } - return null; - } catch (Throwable e) { - context.respond(HttpStatus.BAD_REQUEST); - log.warn("Can't apply default parameters to deployment {}. Trace: {}. Span: {}. Error: {}", - context.getDeployment().getName(), context.getTraceId(), context.getSpanId(), e.getMessage()); - return e; - } - } - - private static boolean applyDefaults(ProxyContext context, ObjectNode tree) { + public Boolean apply(ObjectNode tree) { Deployment deployment = context.getDeployment(); boolean applied = false; for (Map.Entry e : deployment.getDefaults().entrySet()) { diff --git a/server/src/main/java/com/epam/aidial/core/server/function/enhancement/EnhanceAssistantRequestFn.java b/server/src/main/java/com/epam/aidial/core/server/function/enhancement/EnhanceAssistantRequestFn.java index 004a4e883..76569f1c2 100644 --- a/server/src/main/java/com/epam/aidial/core/server/function/enhancement/EnhanceAssistantRequestFn.java +++ b/server/src/main/java/com/epam/aidial/core/server/function/enhancement/EnhanceAssistantRequestFn.java @@ -9,13 +9,11 @@ import com.epam.aidial.core.server.ProxyContext; import com.epam.aidial.core.server.controller.DeploymentController; import com.epam.aidial.core.server.function.BaseRequestFunction; -import com.epam.aidial.core.server.util.ProxyUtil; import com.epam.aidial.core.storage.http.HttpException; import com.epam.aidial.core.storage.http.HttpStatus; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.ArrayNode; import com.fasterxml.jackson.databind.node.ObjectNode; -import io.vertx.core.buffer.Buffer; import lombok.extern.slf4j.Slf4j; import java.util.HashMap; @@ -30,30 +28,16 @@ public EnhanceAssistantRequestFn(Proxy proxy, ProxyContext context) { } @Override - public Throwable apply(ObjectNode tree) { + public Boolean apply(ObjectNode tree) { Deployment deployment = context.getDeployment(); if (deployment instanceof Assistant) { - try { - Map.Entry> enhancedRequest = enhanceAssistantRequest(context, tree); - context.setRequestBody(enhancedRequest.getKey()); - context.setRequestHeaders(enhancedRequest.getValue()); - } catch (HttpException e) { - context.respond(e.getStatus(), e.getMessage()); - log.warn("Can't enhance assistant request. Trace: {}. Span: {}. Error: {}", - context.getTraceId(), context.getSpanId(), e.getMessage()); - return e; - } catch (Throwable e) { - context.respond(HttpStatus.BAD_REQUEST); - log.warn("Can't enhance assistant request. Trace: {}. Span: {}. Error: {}", - context.getTraceId(), context.getSpanId(), e.getMessage()); - return e; - } + enhanceAssistantRequest(context, tree); + return true; } - return null; + return false; } - private static Map.Entry> enhanceAssistantRequest(ProxyContext context, ObjectNode tree) - throws Exception { + private static void enhanceAssistantRequest(ProxyContext context, ObjectNode tree) { Config config = context.getConfig(); Assistant assistant = (Assistant) context.getDeployment(); @@ -108,8 +92,7 @@ private static Map.Entry> enhanceAssistantRequest(Pr throw new HttpException(HttpStatus.FORBIDDEN, "Forbidden model: " + name); } - Buffer updatedBody = Buffer.buffer(ProxyUtil.MAPPER.writeValueAsBytes(tree)); - return Map.entry(updatedBody, headers); + context.setRequestHeaders(headers); } private static void deletePrompt(ArrayNode messages) { diff --git a/server/src/main/java/com/epam/aidial/core/server/function/enhancement/EnhanceModelRequestFn.java b/server/src/main/java/com/epam/aidial/core/server/function/enhancement/EnhanceModelRequestFn.java index ce96621d4..a586ecd03 100644 --- a/server/src/main/java/com/epam/aidial/core/server/function/enhancement/EnhanceModelRequestFn.java +++ b/server/src/main/java/com/epam/aidial/core/server/function/enhancement/EnhanceModelRequestFn.java @@ -5,10 +5,7 @@ import com.epam.aidial.core.server.Proxy; import com.epam.aidial.core.server.ProxyContext; import com.epam.aidial.core.server.function.BaseRequestFunction; -import com.epam.aidial.core.server.util.ProxyUtil; -import com.epam.aidial.core.storage.http.HttpStatus; import com.fasterxml.jackson.databind.node.ObjectNode; -import io.vertx.core.buffer.Buffer; import lombok.extern.slf4j.Slf4j; @Slf4j @@ -18,33 +15,25 @@ public EnhanceModelRequestFn(Proxy proxy, ProxyContext context) { } @Override - public Throwable apply(ObjectNode tree) { + public Boolean apply(ObjectNode tree) { Deployment deployment = context.getDeployment(); if (deployment instanceof Model) { - try { - context.setRequestBody(enhanceModelRequest(context, tree)); - } catch (Throwable e) { - context.respond(HttpStatus.BAD_REQUEST); - log.warn("Can't enhance model request. Trace: {}. Span: {}. Error: {}", - context.getTraceId(), context.getSpanId(), e.getMessage()); - return e; - } + return enhanceModelRequest(context, tree); } - return null; + return false; } - private static Buffer enhanceModelRequest(ProxyContext context, ObjectNode tree) throws Exception { + private static boolean enhanceModelRequest(ProxyContext context, ObjectNode tree) { Model model = (Model) context.getDeployment(); String overrideName = model.getOverrideName(); - Buffer requestBody = context.getRequestBody(); if (overrideName == null) { - return requestBody; + return false; } tree.remove("model"); tree.put("model", overrideName); - return Buffer.buffer(ProxyUtil.MAPPER.writeValueAsBytes(tree)); + return true; } } diff --git a/server/src/main/java/com/epam/aidial/core/server/util/ProxyUtil.java b/server/src/main/java/com/epam/aidial/core/server/util/ProxyUtil.java index 33842de6c..2207c0392 100644 --- a/server/src/main/java/com/epam/aidial/core/server/util/ProxyUtil.java +++ b/server/src/main/java/com/epam/aidial/core/server/util/ProxyUtil.java @@ -301,14 +301,14 @@ public static String convertToString(Object data) { } } - public static Throwable processChain(T item, List> chain) { + public static boolean processChain(T item, List> chain) { + boolean result = false; for (BaseRequestFunction fn : chain) { - Throwable error = fn.apply(item); - if (error != null) { - return error; + if (fn.apply(item)) { + result = true; } } - return null; + return result; } public static EtagHeader etag(HttpServerRequest request) { diff --git a/server/src/test/java/com/epam/aidial/core/server/function/CollectRequestDataFnTest.java b/server/src/test/java/com/epam/aidial/core/server/function/CollectRequestDataFnTest.java index 3ea4c6b5a..86bae989f 100644 --- a/server/src/test/java/com/epam/aidial/core/server/function/CollectRequestDataFnTest.java +++ b/server/src/test/java/com/epam/aidial/core/server/function/CollectRequestDataFnTest.java @@ -35,9 +35,9 @@ public void test_01() throws JsonProcessingException { doCallRealMethod().when(context).setStreamingRequest(anyBoolean()); when(context.isStreamingRequest()).thenCallRealMethod(); - Throwable error = fn.apply((ObjectNode) ProxyUtil.MAPPER.readTree("{\"stream\": true}")); + boolean result = fn.apply((ObjectNode) ProxyUtil.MAPPER.readTree("{\"stream\": true}")); - assertNull(error); + assertFalse(result); assertTrue(context.isStreamingRequest()); } @@ -46,9 +46,9 @@ public void test_02() throws JsonProcessingException { doCallRealMethod().when(context).setStreamingRequest(anyBoolean()); when(context.isStreamingRequest()).thenCallRealMethod(); - Throwable error = fn.apply((ObjectNode) ProxyUtil.MAPPER.readTree("{\"stream\": false}")); + boolean result = fn.apply((ObjectNode) ProxyUtil.MAPPER.readTree("{\"stream\": false}")); - assertNull(error); + assertFalse(result); assertFalse(context.isStreamingRequest()); } @@ -57,9 +57,9 @@ public void test_03() throws JsonProcessingException { doCallRealMethod().when(context).setStreamingRequest(anyBoolean()); when(context.isStreamingRequest()).thenCallRealMethod(); - Throwable error = fn.apply((ObjectNode) ProxyUtil.MAPPER.readTree("{}")); + boolean result = fn.apply((ObjectNode) ProxyUtil.MAPPER.readTree("{}")); - assertNull(error); + assertFalse(result); assertFalse(context.isStreamingRequest()); } } diff --git a/server/src/test/java/com/epam/aidial/core/server/function/enhancement/ApplyDefaultDeploymentSettingsFnTest.java b/server/src/test/java/com/epam/aidial/core/server/function/enhancement/ApplyDefaultDeploymentSettingsFnTest.java index 6305bfb10..6e0b5e8e2 100644 --- a/server/src/test/java/com/epam/aidial/core/server/function/enhancement/ApplyDefaultDeploymentSettingsFnTest.java +++ b/server/src/test/java/com/epam/aidial/core/server/function/enhancement/ApplyDefaultDeploymentSettingsFnTest.java @@ -5,6 +5,7 @@ import com.epam.aidial.core.server.ProxyContext; 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 io.vertx.core.buffer.Buffer; import org.junit.jupiter.api.Test; @@ -18,6 +19,7 @@ import java.util.Map; 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; @@ -42,12 +44,8 @@ public void test() throws JsonProcessingException { Map defaults = Map.of("key1", true, "key2", 123, "key3", 0.45, "key4", "str"); model.setDefaults(defaults); when(context.getDeployment()).thenReturn(model); - Mockito.doCallRealMethod().when(context).setRequestBody(any(Buffer.class)); - when(context.getRequestBody()).thenCallRealMethod(); - Throwable error = fn.apply((ObjectNode) ProxyUtil.MAPPER.readTree("{}")); - assertNull(error); - String json = context.getRequestBody().toString(StandardCharsets.UTF_8); - ObjectNode result = (ObjectNode) ProxyUtil.MAPPER.readTree(json); + JsonNode result = ProxyUtil.MAPPER.readTree("{}"); + assertTrue(fn.apply((ObjectNode) result)); assertNotNull(result); assertEquals(123, result.get("key2").asInt()); assertEquals(0.45, result.get("key3").asDouble()); From 69cabd3989ba35532490a3ac48a0d5095acb8a70 Mon Sep 17 00:00:00 2001 From: Aliaksandr Stsiapanay Date: Wed, 20 Nov 2024 16:32:34 +0300 Subject: [PATCH 2/2] fix: respond with error message --- .../aidial/core/server/controller/DeploymentPostController.java | 2 +- .../aidial/core/server/controller/InterceptorController.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/server/src/main/java/com/epam/aidial/core/server/controller/DeploymentPostController.java b/server/src/main/java/com/epam/aidial/core/server/controller/DeploymentPostController.java index 4ee6c7819..cfe7750da 100644 --- a/server/src/main/java/com/epam/aidial/core/server/controller/DeploymentPostController.java +++ b/server/src/main/java/com/epam/aidial/core/server/controller/DeploymentPostController.java @@ -272,7 +272,7 @@ void handleRequestBody(Buffer requestBody) { } } catch (Throwable e) { if (e instanceof HttpException httpException) { - respond(httpException.getStatus(), httpException); + respond(httpException.getStatus(), httpException.getMessage()); } else { respond(HttpStatus.BAD_REQUEST); } diff --git a/server/src/main/java/com/epam/aidial/core/server/controller/InterceptorController.java b/server/src/main/java/com/epam/aidial/core/server/controller/InterceptorController.java index 611c4255d..f9c2ba0a6 100644 --- a/server/src/main/java/com/epam/aidial/core/server/controller/InterceptorController.java +++ b/server/src/main/java/com/epam/aidial/core/server/controller/InterceptorController.java @@ -75,7 +75,7 @@ private void handleRequestBody(Buffer requestBody) { } } catch (Throwable e) { if (e instanceof HttpException httpException) { - respond(httpException.getStatus(), httpException); + respond(httpException.getStatus(), httpException.getMessage()); } else { respond(HttpStatus.BAD_REQUEST); }