feat: include query params when rewritePath is enabled (#598)

epam · Dec 3, 2024 · 6721d70 · 6721d70
1 parent 0d9c6db
commit 6721d70
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/server/src/main/java/com/epam/aidial/core/server/controller/RouteController.java b/server/src/main/java/com/epam/aidial/core/server/controller/RouteController.java
@@ -24,7 +24,9 @@
 import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.http.client.utils.URIBuilder;
+import org.apache.http.client.utils.URLEncodedUtils;
 
+import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
@@ -288,6 +290,10 @@ private String getEndpointUri(Upstream upstream) {
         URIBuilder uriBuilder = new URIBuilder(upstream.getEndpoint());
         if (context.getRoute().isRewritePath()) {
             uriBuilder.setPath(context.getRequest().path());
+            String query = context.getRequest().query();
+            if (query != null) {
+                uriBuilder.setParameters(URLEncodedUtils.parse(query, StandardCharsets.UTF_8));
+            }
         }
         return uriBuilder.toString();
     }

diff --git a/server/src/test/java/com/epam/aidial/core/server/RouteApiTest.java b/server/src/test/java/com/epam/aidial/core/server/RouteApiTest.java
@@ -62,6 +62,7 @@ private static List<Arguments> datasource() {
                 Arguments.of(HttpMethod.GET, "/v1/plain", "vstore_user_key", 200, "/"),
                 Arguments.of(HttpMethod.GET, "/v1/plain", "vstore_admin_key", 200, "/"),
                 Arguments.of(HttpMethod.GET, "/v1/vector_store/1", "vstore_user_key", 200, "/v1/vector_store/1"),
+                Arguments.of(HttpMethod.GET, "/v1/vector_store/1?q=p", "vstore_user_key", 200, "/v1/vector_store/1?q=p"),
                 Arguments.of(HttpMethod.GET, "/v1/vector_store/1", "vstore_admin_key", 200, "/v1/vector_store/1"),
                 Arguments.of(HttpMethod.HEAD, "/v1/vector_store/1", "vstore_user_key", 200, null),
                 Arguments.of(HttpMethod.HEAD, "/v1/vector_store/1", "vstore_admin_key", 200, null),