Skip to content

Commit

Permalink
feat: Add feature flag accessibleByPerRequestKey #523 (#526)
Browse files Browse the repository at this point in the history
  • Loading branch information
@astsiapanay
astsiapanay authored Oct 16, 2024
1 parent fba5741 commit fbc983a
Show file tree
Hide file tree
Showing 9 changed files with 69 additions and 18 deletions.
3 changes: 0 additions & 3 deletions src/main/java/com/epam/aidial/core/config/ApiKeyData.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,14 @@

import com.epam.aidial.core.ProxyContext;
import com.epam.aidial.core.data.AutoSharedData;
import com.epam.aidial.core.data.ResourceAccessType;
import com.epam.aidial.core.security.ExtractedClaims;
import com.fasterxml.jackson.annotation.JsonIgnore;
import lombok.Data;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
* The container keeps data associated with API key.
Expand Down
1 change: 1 addition & 0 deletions src/main/java/com/epam/aidial/core/config/Features.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,4 +18,5 @@ public class Features {
private Boolean urlAttachmentsSupported;
private Boolean folderAttachmentsSupported;
private Boolean allowResume;
private Boolean accessibleByPerRequestKey;
}
3 changes: 3 additions & 0 deletions src/main/java/com/epam/aidial/core/config/FileConfigStore.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -180,5 +180,8 @@ private static void setMissingFeatures(Deployment model, Features features) {
if (modelFeatures.getAllowResume() == null) {
modelFeatures.setAllowResume(features.getAllowResume());
}
if (modelFeatures.getAccessibleByPerRequestKey() == null) {
modelFeatures.setAccessibleByPerRequestKey(features.getAccessibleByPerRequestKey());
}
}
}
4 changes: 4 additions & 0 deletions src/main/java/com/epam/aidial/core/controller/DeploymentController.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,10 @@ static FeaturesData createFeatures(Features features) {
data.setAllowResume(features.getAllowResume());
}

if (features.getAccessibleByPerRequestKey() != null) {
data.setAccessibleByPerRequestKey(features.getAccessibleByPerRequestKey());
}

return data;
}
}
6 changes: 6 additions & 0 deletions src/main/java/com/epam/aidial/core/controller/DeploymentPostController.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import com.epam.aidial.core.config.ApiKeyData;
import com.epam.aidial.core.config.Config;
import com.epam.aidial.core.config.Deployment;
import com.epam.aidial.core.config.Features;
import com.epam.aidial.core.config.Interceptor;
import com.epam.aidial.core.config.Model;
import com.epam.aidial.core.config.ModelType;
Expand Down Expand Up @@ -91,6 +92,11 @@ public Future<?> handle(String deploymentId, String deploymentApi) {
private Future<?> handleDeployment(String deploymentId, String deploymentApi) {
return DeploymentController.selectDeployment(context, deploymentId)
.map(dep -> {
Features features = dep.getFeatures();
boolean isPerRequestKey = context.getApiKeyData().getPerRequestKey() != null;
if (features != null && Boolean.FALSE.equals(features.getAccessibleByPerRequestKey()) && isPerRequestKey) {
throw new PermissionDeniedException(String.format("Deployment %s is not accessible by %s", deploymentId, context.getApiKeyData().getSourceDeployment()));
}
context.setDeployment(dep);
return dep;
})
Expand Down
1 change: 1 addition & 0 deletions src/main/java/com/epam/aidial/core/data/FeaturesData.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,4 +20,5 @@ public class FeaturesData {
private boolean urlAttachments = false;
private boolean folderAttachments = false;
private boolean allowResume = true;
private boolean accessibleByPerRequestKey = true;
}
33 changes: 22 additions & 11 deletions src/test/java/com/epam/aidial/core/CustomApplicationApiTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -333,7 +333,8 @@ void testApplicationSharing() {
"seed" : false,
"url_attachments" : false,
"folder_attachments" : false,
"allow_resume":true
"allow_resume":true,
"accessible_by_per_request_key": true
},
"defaults" : { },
"description_keywords":[]
Expand Down Expand Up @@ -367,7 +368,8 @@ void testApplicationSharing() {
"seed":false,
"url_attachments":false,
"folder_attachments":false,
"allow_resume":true
"allow_resume":true,
"accessible_by_per_request_key": true
},
"defaults":{},
"description_keywords":[]
Expand Down Expand Up @@ -395,7 +397,8 @@ void testApplicationSharing() {
"seed" : false,
"url_attachments" : false,
"folder_attachments" : false,
"allow_resume":true
"allow_resume":true,
"accessible_by_per_request_key": true
},
"defaults" : { },
"description_keywords":[]
Expand Down Expand Up @@ -432,7 +435,8 @@ void testApplicationSharing() {
"seed":false,
"url_attachments":false,
"folder_attachments":false,
"allow_resume":true
"allow_resume":true,
"accessible_by_per_request_key": true
},
"defaults":{},
"description_keywords":[]
Expand Down Expand Up @@ -460,7 +464,8 @@ void testApplicationSharing() {
"seed" : false,
"url_attachments" : false,
"folder_attachments" : false,
"allow_resume":true
"allow_resume":true,
"accessible_by_per_request_key": true
},
"defaults" : { },
"description_keywords":[]
Expand Down Expand Up @@ -570,7 +575,8 @@ void testApplicationPublication() {
"seed":false,
"url_attachments":false,
"folder_attachments":false,
"allow_resume":true
"allow_resume":true,
"accessible_by_per_request_key": true
},
"defaults":{},
"description_keywords":[]
Expand Down Expand Up @@ -598,7 +604,8 @@ void testApplicationPublication() {
"seed" : false,
"url_attachments" : false,
"folder_attachments" : false,
"allow_resume":true
"allow_resume":true,
"accessible_by_per_request_key": true
},
"defaults" : { },
"description_keywords":[]
Expand Down Expand Up @@ -638,7 +645,8 @@ void testOpenAiApplicationListing() {
"seed":false,
"url_attachments":false,
"folder_attachments":false,
"allow_resume": true
"allow_resume": true,
"accessible_by_per_request_key": true
},
"defaults":{},
"description_keywords":[]
Expand Down Expand Up @@ -690,7 +698,8 @@ void testOpenAiApplicationListing() {
"seed":false,
"url_attachments":false,
"folder_attachments":false,
"allow_resume": true
"allow_resume": true,
"accessible_by_per_request_key": true
},
"defaults":{},
"description_keywords":[]
Expand Down Expand Up @@ -724,7 +733,8 @@ void testOpenAiApplicationListing() {
"seed":false,
"url_attachments":false,
"folder_attachments":false,
"allow_resume": true
"allow_resume": true,
"accessible_by_per_request_key": true
},
"defaults":{},
"description_keywords":[]
Expand Down Expand Up @@ -752,7 +762,8 @@ void testOpenAiApplicationListing() {
"seed" : false,
"url_attachments" : false,
"folder_attachments" : false,
"allow_resume": true
"allow_resume": true,
"accessible_by_per_request_key": true
},
"defaults" : { },
"description_keywords":[]
Expand Down
8 changes: 4 additions & 4 deletions src/test/java/com/epam/aidial/core/ListingTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ void testFeaturesEmbedding(Vertx vertx, VertxTestContext context) {
{ "rate": false, "tokenize": false, "truncate_prompt": false
, "system_prompt": true, "tools": false, "seed": false
, "url_attachments": false, "folder_attachments": false
, "configuration": false, "allow_resume": true
, "configuration": false, "allow_resume": true, "accessible_by_per_request_key": true
}
"""));
}
Expand All @@ -62,7 +62,7 @@ void testFeaturesModel(Vertx vertx, VertxTestContext context) {
{ "rate": true, "tokenize": true, "truncate_prompt": true
, "system_prompt": true, "tools": true, "seed": true
, "url_attachments": true, "folder_attachments": false
, "configuration": false, "allow_resume": true
, "configuration": false, "allow_resume": true, "accessible_by_per_request_key": true
}
"""));
}
Expand All @@ -73,7 +73,7 @@ void testFeaturesApplication(Vertx vertx, VertxTestContext context) {
{ "rate": true, "tokenize": false, "truncate_prompt": false
, "system_prompt": false, "tools": false, "seed": false
, "url_attachments": false, "folder_attachments": false
, "configuration": true, "allow_resume": true
, "configuration": true, "allow_resume": true, "accessible_by_per_request_key": true
}
"""));
}
Expand All @@ -84,7 +84,7 @@ void testFeaturesAssistant(Vertx vertx, VertxTestContext context) {
{ "rate": true, "tokenize": false, "truncate_prompt": false
, "system_prompt": true, "tools": false, "seed": false
, "url_attachments": false, "folder_attachments": false
, "configuration": false, "allow_resume": true
, "configuration": false, "allow_resume": true, "accessible_by_per_request_key": true
}
"""));
}
Expand Down
28 changes: 28 additions & 0 deletions src/test/java/com/epam/aidial/core/controller/DeploymentPostControllerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import com.epam.aidial.core.config.ApiKeyData;
import com.epam.aidial.core.config.Application;
import com.epam.aidial.core.config.Config;
import com.epam.aidial.core.config.Features;
import com.epam.aidial.core.config.Model;
import com.epam.aidial.core.limiter.RateLimiter;
import com.epam.aidial.core.log.LogStore;
Expand Down Expand Up @@ -139,9 +140,36 @@ public void testDeploymentNotFound() {
verify(context).respond(eq(NOT_FOUND), anyString());
}

@Test
public void testDeploymentIsNotAccessible() {
when(context.getRequest()).thenReturn(request);
when(request.getHeader(eq(HttpHeaders.CONTENT_TYPE))).thenReturn(HEADER_CONTENT_TYPE_APPLICATION_JSON);
Config config = new Config();
config.setApplications(new HashMap<>());
Application app = new Application();
Features features = new Features();
features.setAccessibleByPerRequestKey(false);
app.setFeatures(features);
ApiKeyData apiKeyData = new ApiKeyData();
apiKeyData.setPerRequestKey("perRequestKey");
when(context.getApiKeyData()).thenReturn(apiKeyData);
config.getApplications().put("app1", app);
when(context.getConfig()).thenReturn(config);
when(proxy.getVertx()).thenReturn(vertx);
when(proxy.getTokenStatsTracker()).thenReturn(tokenStatsTracker);
when(vertx.executeBlocking(any(Callable.class), eq(false)))
.thenReturn(Future.succeededFuture(app));
when(context.getProxy()).thenReturn(proxy);

controller.handle("unknown-app", "chat/completions");

verify(context).respond(eq(FORBIDDEN), anyString());
}

@Test
public void testNoRoute() {
when(context.getRequest()).thenReturn(request);
when(context.getApiKeyData()).thenReturn(new ApiKeyData());
when(request.getHeader(eq(HttpHeaders.CONTENT_TYPE))).thenReturn(HEADER_CONTENT_TYPE_APPLICATION_JSON);
Config config = new Config();
config.setApplications(new HashMap<>());
Expand Down

0 comments on commit fbc983a

Please sign in to comment.