feat: extend unified chat completion request with custom skills #1955

docs/open_api_core.yaml

@@ -14764,6 +14764,11 @@ components:
           description: |-
             The JSON schema describing a form that the assistant prompts the user to fill in.
             Given this schema, the user is expected to provide a JSON value in the next message in the `custom_content.form_value` field.
+        skills:
+          type: array
+          description: List of skills referenced by the model as part of the response.
+          items:
+            $ref: "#/components/schemas/RequestSkill"
       description: The custom content of the assistant message.
     ChatCompletionRequestCustomFields:
       type: object
@@ -15007,6 +15012,11 @@ components:
         form_value:
           type: object
           description: The JSON value corresponding to the JSON schema sent by the assistant in the previous message in the `custom_content.form_schema` field.
+        skills:
+          type: array
+          description: List of skills used to supply an additional input for the model.
+          items:
+            $ref: "#/components/schemas/RequestSkill"
       description: The custom content of the user message.
     ChatCompletionResponseAttachment:
       required:
@@ -16978,6 +16988,14 @@ components:
         reference_url:
           type: string
           description: "If `reference_type` is specified, the content of `reference_url` should follow the format described in the MIME standard for `reference_type`."
+    RequestSkill:
+      required:
+      - url
+      type: object
+      properties:
+        url:
+          type: string
+          description: "The URL of the skill resource, e.g. `skills/<bucket>/<path>`."
     ResourceAccessType:
       type: string
       enum:

openapi-generator/src/main/resources/schemas/ChatCompletionRequestAssistantMessageCustomContent.yaml

@@ -15,3 +15,8 @@ properties:
     description: |-
       The JSON schema describing a form that the assistant prompts the user to fill in.
       Given this schema, the user is expected to provide a JSON value in the next message in the `custom_content.form_value` field.
+  skills:
+    type: array
+    items:
+      $ref: "#/components/schemas/RequestSkill"
+    description: List of skills referenced by the model as part of the response.

openapi-generator/src/main/resources/schemas/ChatCompletionRequestUserMessageCustomContent.yaml

@@ -10,3 +10,8 @@ properties:
     type: object
     description: |-
       The JSON value corresponding to the JSON schema sent by the assistant in the previous message in the `custom_content.form_schema` field.
+  skills:
+    type: array
+    items:
+      $ref: "#/components/schemas/RequestSkill"
+    description: List of skills used to supply an additional input for the model.

openapi-generator/src/main/resources/schemas/RequestSkill.yaml

@@ -0,0 +1,7 @@
+type: object
+properties:
+  url:
+    type: string
+    description: The URL of the skill resource, e.g. `skills/<bucket>/<path>`.
+required:
+  - url

server/src/main/java/com/epam/aidial/core/server/controller/DeploymentPostController.java

@@ -22,6 +22,7 @@
 import com.epam.aidial.core.server.function.CollectChatCompletionUsageFn;
 import com.epam.aidial.core.server.function.CollectDeploymentsFn;
 import com.epam.aidial.core.server.function.CollectRequestApplicationFilesFn;
+import com.epam.aidial.core.server.function.CollectRequestSkillsFn;
 import com.epam.aidial.core.server.function.CollectRequestStandardAttachmentsFn;
 import com.epam.aidial.core.server.function.CollectResponseChatCompletionAttachmentsFn;
 import com.epam.aidial.core.server.function.StripUsagePerModelFn;
@@ -74,6 +75,7 @@ public DeploymentPostController(Proxy proxy, ProxyContext context) {
      */
     private List<BaseRequestFunction<RequestObject>> buildEnhancementFunctions() {
         return List.of(new CollectRequestStandardAttachmentsFn(proxy, context),
+                new CollectRequestSkillsFn(proxy, context),
                 new ApplyDefaultDeploymentSettingsFn(proxy, context, requestedInterface()),
                 new EnhanceDeploymentRequestFn(proxy, context),
                 new CollectRequestApplicationFilesFn(proxy, context),

server/src/main/java/com/epam/aidial/core/server/function/CollectRequestSkillsFn.java

@@ -0,0 +1,48 @@
+package com.epam.aidial.core.server.function;
+
+import com.epam.aidial.core.config.ResourceAccessType;
+import com.epam.aidial.core.server.Proxy;
+import com.epam.aidial.core.server.ProxyContext;
+import com.epam.aidial.core.server.data.AutoSharedData;
+import com.epam.aidial.core.server.function.request.RequestObject;
+import com.epam.aidial.core.server.security.AccessService;
+import com.epam.aidial.core.storage.http.HttpException;
+import com.epam.aidial.core.storage.http.HttpStatus;
+import com.epam.aidial.core.storage.resource.ResourceDescriptor;
+import com.epam.aidial.core.storage.resource.ResourceTypes;
+
+/**
+ * Collects skill resources referenced from {@code messages[*].custom_content.skills[*]} and auto-shares
+ * them to the invoked deployment's per-request API key, the same way {@link CollectRequestAttachmentsFn}
+ * auto-shares attachments.
+ */
+public class CollectRequestSkillsFn extends BaseRequestFunction<RequestObject> {
+    public CollectRequestSkillsFn(Proxy proxy, ProxyContext context) {
+        super(proxy, context);
+    }
+
+    @Override
+    public Boolean apply(RequestObject request) {
+        for (String url : request.collectSkills()) {
+            tryToAutoShareAttachedSkill(url);
+        }
+        return false;
+    }
+
+    private void tryToAutoShareAttachedSkill(String url) {
+        ResourceDescriptor resource = fromAnyUrl(url, proxy.getEncryptionService());
+        if (resource == null || resource.getType() != ResourceTypes.SKILL) {
+            throw new HttpException(HttpStatus.BAD_REQUEST, "Url must reference a skill resource: %s".formatted(url));
+        }
+        if (resource.isPublic()) {
+            return;
+        }
+        AccessService accessService = proxy.getAccessService();
+        if (accessService.hasReadAccess(resource, context)) {
+            context.getProxyApiKeyData().getAttachedSkills()
+                    .put(resource.getUrl(), new AutoSharedData(ResourceAccessType.READ_ONLY));
+        } else {
+            throw new HttpException(HttpStatus.FORBIDDEN, "Access denied to the skill %s".formatted(url));
+        }
+    }
+}

server/src/main/java/com/epam/aidial/core/server/function/request/ChatCompletionRequest.java

@@ -60,6 +60,11 @@ public Set<String> collectAppAttachments(List<String> paths) {
         return ChatUtil.collectAttachments(tree, paths);
     }
 
+    @Override
+    public Set<String> collectSkills() {
+        return ChatUtil.collectCustomSkills(tree, List.of("$.messages[*].custom_content.skills[*]"));
+    }
+
     @Override
     public List<CacheKey> buildCacheKeys(List<String> nodeOrder) {
         CacheKeyBuilder builder = new CacheKeyBuilder();

server/src/main/java/com/epam/aidial/core/server/function/request/RequestObject.java

@@ -38,6 +38,15 @@ public interface RequestObject {
      */
     Set<String> collectAttachments();
 
+    /**
+     * Collects skill resource URLs present in the request body.
+     *
+     * @return a set of skill URLs found in the request body, or an empty set if this request shape doesn't support it
+     */
+    default Set<String> collectSkills() {
+        return Set.of();
+    }
+
     /**
      * Collects image and file URLs located at the specified paths.
      *

server/src/main/java/com/epam/aidial/core/server/util/ChatUtil.java

@@ -22,6 +22,10 @@ public Set<String> collectCustomAttachments(JsonNode node, List<String> paths) {
         return JsonUtil.collectStrings(node, paths, ChatUtil::readCustomAttachment);
     }
 
+    public Set<String> collectCustomSkills(JsonNode node, List<String> paths) {
+        return JsonUtil.collectStrings(node, paths, ChatUtil::readCustomSkill);
+    }
+
     private String readAttachment(JsonNode node) {
         if (!node.isTextual()) {
             throw new IllegalArgumentException("Invalid attachment.");
@@ -49,6 +53,17 @@ public String readCustomAttachment(JsonNode node) {
         return url;
     }
 
+    public String readCustomSkill(JsonNode node) {
+        if (!node.isObject()) {
+            throw new IllegalArgumentException("Invalid skill attachment.");
+        }
+        String url = node.path("url").asText();
+        if (StringUtils.isBlank(url)) {
+            throw new IllegalArgumentException("Missing url in skill attachment.");
+        }
+        return url;
+    }
+
     public void removeInterceptorConfiguration(ObjectNode node) {
         ObjectNode customFields = (ObjectNode) node.get(CUSTOM_FIELDS_NODE);
         if (customFields != null) {

server/src/test/java/com/epam/aidial/core/server/controller/DeploymentPostControllerTest.java

@@ -70,6 +70,7 @@
 import static com.epam.aidial.core.server.Proxy.HEADER_APPLICATION_PROPERTIES;
 import static com.epam.aidial.core.server.Proxy.HEADER_CONTENT_TYPE_APPLICATION_JSON;
 import static com.epam.aidial.core.storage.http.HttpStatus.BAD_GATEWAY;
+import static com.epam.aidial.core.storage.http.HttpStatus.BAD_REQUEST;
 import static com.epam.aidial.core.storage.http.HttpStatus.FORBIDDEN;
 import static com.epam.aidial.core.storage.http.HttpStatus.NOT_FOUND;
 import static com.epam.aidial.core.storage.http.HttpStatus.UNSUPPORTED_MEDIA_TYPE;
@@ -576,6 +577,128 @@ public void testHandleRequestBody_OverrideModelName_Application() throws IOExcep
         assertEquals(tree.get("model").asText(), "overrideName");
     }
 
+    @Test
+    public void testHandleRequestBody_SkillAutoShared_WhenReadable() {
+        when(context.getRequest()).thenReturn(request);
+        UpstreamRoute upstreamRoute = mock(UpstreamRoute.class, RETURNS_DEEP_STUBS);
+        when(upstreamRoute.next()).thenReturn(new Upstream("endpoint", null, null, null, null, 0, 0, null, null, null));
+        when(context.getUpstreamRoute()).thenReturn(upstreamRoute);
+        HttpServerRequest request = mock(HttpServerRequest.class, RETURNS_DEEP_STUBS);
+        when(context.getRequest()).thenReturn(request);
+        when(request.path()).thenReturn("/openai/deployments/name/chat/completions");
+        when(proxy.getClient()).thenReturn(mock(HttpClient.class, RETURNS_DEEP_STUBS));
+        when(proxy.getApiKeyStore()).thenReturn(mock(ApiKeyStore.class));
+        when(proxy.getClientOptions()).thenReturn(new HttpClientOptions());
+        ApiKeyData proxyApiKeyData = new ApiKeyData();
+        proxyApiKeyData.setInterceptorIndex(0);
+        when(context.getProxyApiKeyData()).thenReturn(proxyApiKeyData);
+        when(proxy.getEncryptionService().decrypt("bucket")).thenReturn("location/");
+        when(proxy.getAccessService().hasReadAccess(any(), any())).thenReturn(true);
+
+        Model model = new Model();
+        model.setName("name");
+        model.setEndpoint("http://host/model");
+        when(context.getDeployment()).thenReturn(model);
+        String body = """
+                {
+                    "model": "name",
+                    "messages": [
+                        {
+                            "role": "user",
+                            "content": "use the summarizer skill",
+                            "custom_content": {
+                                "skills": [
+                                    {"url": "skills/bucket/summarizer"}
+                                ]
+                            }
+                        }
+                    ],
+                    "stream": false
+                }
+                """;
+        Buffer requestBody = Buffer.buffer(body);
+
+        controller.handleRequestBody(requestBody);
+
+        assertNotNull(proxyApiKeyData.getAttachedSkills().get("skills/bucket/summarizer"));
+    }
+
+    @Test
+    public void testHandleRequestBody_SkillAccessDenied() {
+        when(context.getRequest()).thenReturn(request);
+        HttpServerRequest request = mock(HttpServerRequest.class, RETURNS_DEEP_STUBS);
+        when(context.getRequest()).thenReturn(request);
+        when(request.path()).thenReturn("/openai/deployments/name/chat/completions");
+        ApiKeyData proxyApiKeyData = new ApiKeyData();
+        when(context.getProxyApiKeyData()).thenReturn(proxyApiKeyData);
+        when(proxy.getEncryptionService().decrypt("bucket")).thenReturn("location/");
+        // proxy.getAccessService().hasReadAccess(...) defaults to false (deep stub)
+
+        Model model = new Model();
+        model.setName("name");
+        model.setEndpoint("http://host/model");
+        when(context.getDeployment()).thenReturn(model);
+        String body = """
+                {
+                    "model": "name",
+                    "messages": [
+                        {
+                            "role": "user",
+                            "content": "use the summarizer skill",
+                            "custom_content": {
+                                "skills": [
+                                    {"url": "skills/bucket/summarizer"}
+                                ]
+                            }
+                        }
+                    ],
+                    "stream": false
+                }
+                """;
+        Buffer requestBody = Buffer.buffer(body);
+
+        controller.handleRequestBody(requestBody);
+
+        verify(context).respond(eq(FORBIDDEN), anyString());
+    }
+
+    @Test
+    public void testHandleRequestBody_SkillUrlIsNotSkillResource() {
+        when(context.getRequest()).thenReturn(request);
+        HttpServerRequest request = mock(HttpServerRequest.class, RETURNS_DEEP_STUBS);
+        when(context.getRequest()).thenReturn(request);
+        when(request.path()).thenReturn("/openai/deployments/name/chat/completions");
+        ApiKeyData proxyApiKeyData = new ApiKeyData();
+        when(context.getProxyApiKeyData()).thenReturn(proxyApiKeyData);
+
+        Model model = new Model();
+        model.setName("name");
+        model.setEndpoint("http://host/model");
+        when(context.getDeployment()).thenReturn(model);
+        String body = """
+                {
+                    "model": "name",
+                    "messages": [
+                        {
+                            "role": "user",
+                            "content": "use the summarizer skill",
+                            "custom_content": {
+                                "skills": [
+                                    {"url": "files/public/readme.md"}
+                                ]
+                            }
+                        }
+                    ],
+                    "stream": false
+                }
+                """;
+        Buffer requestBody = Buffer.buffer(body);
+
+        controller.handleRequestBody(requestBody);
+
+        verify(context).respond(eq(BAD_REQUEST), eq("Url must reference a skill resource: files/public/readme.md"));
+    }
+
     @Test
     public void testHandleRequestBody_UseUpstreamWithoutEndpoint() {
         when(context.getRequest()).thenReturn(request);