feat: supported metadata attachments (#314)

epam · Apr 23, 2024 · 248d363 · 248d363
1 parent 93ddf0f
commit 248d363
Show file tree

Hide file tree

Showing 8 changed files with 116 additions and 9 deletions.
diff --git a/src/main/java/com/epam/aidial/core/ProxyContext.java b/src/main/java/com/epam/aidial/core/ProxyContext.java
@@ -104,8 +104,13 @@ public Future<Void> respond(HttpStatus status) {
 
     @SneakyThrows
     public Future<Void> respond(HttpStatus status, Object object) {
+        return respond(status, Proxy.HEADER_CONTENT_TYPE_APPLICATION_JSON, object);
+    }
+
+    @SneakyThrows
+    public Future<Void> respond(HttpStatus status, String contentType, Object object) {
         String json = ProxyUtil.MAPPER.writeValueAsString(object);
-        response.putHeader(HttpHeaders.CONTENT_TYPE, Proxy.HEADER_CONTENT_TYPE_APPLICATION_JSON);
+        response.putHeader(HttpHeaders.CONTENT_TYPE, contentType);
         return respond(status, json);
     }
 

diff --git a/src/main/java/com/epam/aidial/core/controller/FileMetadataController.java b/src/main/java/com/epam/aidial/core/controller/FileMetadataController.java
@@ -7,6 +7,7 @@
 import com.epam.aidial.core.storage.ResourceDescription;
 import com.epam.aidial.core.util.HttpStatus;
 import io.vertx.core.Future;
+import io.vertx.core.http.HttpHeaders;
 import lombok.extern.slf4j.Slf4j;
 
 @Slf4j
@@ -16,6 +17,13 @@ public FileMetadataController(Proxy proxy, ProxyContext context) {
         super(proxy, context, false);
     }
 
+    private String getContentType() {
+        String acceptType = context.getRequest().getHeader(HttpHeaders.ACCEPT);
+        return acceptType != null && acceptType.contains(MetadataBase.MIME_TYPE)
+                ? MetadataBase.MIME_TYPE
+                : "application/json";
+    }
+
     @Override
     protected Future<?> handle(ResourceDescription resource) {
         BlobStorage storage = proxy.getStorage();
@@ -24,7 +32,7 @@ protected Future<?> handle(ResourceDescription resource) {
                 MetadataBase metadata = storage.listMetadata(resource);
                 if (metadata != null) {
                     proxy.getAccessService().filterForbidden(context, resource, metadata);
-                    context.respond(HttpStatus.OK, metadata);
+                    context.respond(HttpStatus.OK, getContentType(), metadata);
                 } else {
                     context.respond(HttpStatus.NOT_FOUND);
                 }

diff --git a/src/main/java/com/epam/aidial/core/controller/ResourceController.java b/src/main/java/com/epam/aidial/core/controller/ResourceController.java
@@ -3,6 +3,7 @@
 import com.epam.aidial.core.Proxy;
 import com.epam.aidial.core.ProxyContext;
 import com.epam.aidial.core.data.Conversation;
+import com.epam.aidial.core.data.MetadataBase;
 import com.epam.aidial.core.data.Prompt;
 import com.epam.aidial.core.data.ResourceLink;
 import com.epam.aidial.core.data.ResourceLinkCollection;
@@ -64,6 +65,13 @@ protected Future<?> handle(ResourceDescription descriptor) {
         return context.respond(HttpStatus.BAD_GATEWAY, "No route");
     }
 
+    private String getContentType() {
+        String acceptType = context.getRequest().getHeader(HttpHeaders.ACCEPT);
+        return acceptType != null && metadata && acceptType.contains(MetadataBase.MIME_TYPE)
+                ? MetadataBase.MIME_TYPE
+                : "application/json";
+    }
+
     private Future<?> getMetadata(ResourceDescription descriptor) {
         String token;
         int limit;
@@ -86,7 +94,7 @@ private Future<?> getMetadata(ResourceDescription descriptor) {
                         context.respond(HttpStatus.NOT_FOUND, "Not found: " + descriptor.getUrl());
                     } else {
                         proxy.getAccessService().filterForbidden(context, descriptor, result);
-                        context.respond(HttpStatus.OK, result);
+                        context.respond(HttpStatus.OK, getContentType(), result);
                     }
                 })
                 .onFailure(error -> {

diff --git a/src/main/java/com/epam/aidial/core/data/MetadataBase.java b/src/main/java/com/epam/aidial/core/data/MetadataBase.java
@@ -8,6 +8,8 @@
 @Data
 @NoArgsConstructor
 public abstract class MetadataBase {
+    public static final String MIME_TYPE = "application/vnd.dial.metadata+json";
+
     private String name;
     private String parentPath;
     private String bucket;

diff --git a/src/main/java/com/epam/aidial/core/storage/ResourceDescription.java b/src/main/java/com/epam/aidial/core/storage/ResourceDescription.java
@@ -202,7 +202,7 @@ private static ResourceDescription fromUrl(String url,
         }
 
         if (parts.length == 2 && !url.endsWith(BlobStorageUtil.PATH_SEPARATOR)) {
-            throw new IllegalArgumentException("Url must start resource/bucket/, but: " + BlobStorageUtil.PATH_SEPARATOR + ": " + url);
+            throw new IllegalArgumentException("Url must start with resource/bucket/, but: " + url);
         }
 
         ResourceType resourceType = ResourceType.of(UrlUtil.decodePath(parts[0]));

diff --git a/src/main/java/com/epam/aidial/core/util/ProxyUtil.java b/src/main/java/com/epam/aidial/core/util/ProxyUtil.java
@@ -1,6 +1,7 @@
 package com.epam.aidial.core.util;
 
 import com.epam.aidial.core.Proxy;
+import com.epam.aidial.core.data.MetadataBase;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.JsonMappingException;
@@ -99,14 +100,31 @@ public static void collectAttachedFiles(ObjectNode tree, Consumer<String> consum
             }
             for (int j = 0; j < attachments.size(); j++) {
                 JsonNode attachment = attachments.get(j);
-                JsonNode url = attachment.get("url");
-                if (url != null) {
-                    consumer.accept(url.textValue());
-                }
+                collectAttachedFile(attachment, consumer);
             }
         }
     }
 
+    private static void collectAttachedFile(JsonNode attachment, Consumer<String> consumer) {
+        JsonNode urlNode = attachment.get("url");
+        if (urlNode == null) {
+            return;
+        }
+
+        String url = urlNode.textValue();
+
+        JsonNode typeNode = attachment.get("type");
+        if (typeNode != null && typeNode.textValue().equals(MetadataBase.MIME_TYPE)) {
+            String prefix = "metadata/";
+            if (!url.startsWith(prefix)) {
+                throw new IllegalArgumentException("Url of metadata attachment must start with metadata/: " + url);
+            }
+            url = url.substring(prefix.length());
+        }
+
+        consumer.accept(url);
+    }
+
     public static <T> T convertToObject(Buffer json, Class<T> clazz) {
         try {
             String text = json.toString(StandardCharsets.UTF_8);

diff --git a/src/test/java/com/epam/aidial/core/FileApiTest.java b/src/test/java/com/epam/aidial/core/FileApiTest.java
@@ -3,6 +3,7 @@
 import com.epam.aidial.core.config.ApiKeyData;
 import com.epam.aidial.core.data.Bucket;
 import com.epam.aidial.core.data.FileMetadata;
+import com.epam.aidial.core.data.MetadataBase;
 import com.epam.aidial.core.data.ResourceFolderMetadata;
 import com.epam.aidial.core.data.ResourceType;
 import com.epam.aidial.core.util.ProxyUtil;
@@ -17,6 +18,7 @@
 import io.vertx.junit5.VertxExtension;
 import io.vertx.junit5.VertxTestContext;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.http.HttpHeaders;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 
@@ -122,6 +124,28 @@ public void testEmptyFilesList(Vertx vertx, VertxTestContext context) {
                 .send(context.succeeding(response -> {
                     context.verify(() -> {
                         assertEquals(200, response.statusCode());
+                        assertEquals("application/json", response.getHeader(HttpHeaders.CONTENT_TYPE));
+                        assertEquals(emptyBucketResponse, response.body());
+                        context.completeNow();
+                    });
+                }));
+    }
+
+    @Test
+    public void testMetadataContentType(Vertx vertx, VertxTestContext context) {
+        WebClient client = WebClient.create(vertx);
+
+        ResourceFolderMetadata emptyBucketResponse = new ResourceFolderMetadata(ResourceType.FILE, "7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt",
+                null, null, "files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/", List.of());
+
+        client.get(serverPort, "localhost", "/v1/metadata/files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/")
+                .putHeader("Api-key", "proxyKey2")
+                .putHeader(HttpHeaders.ACCEPT, MetadataBase.MIME_TYPE)
+                .as(BodyCodec.json(ResourceFolderMetadata.class))
+                .send(context.succeeding(response -> {
+                    context.verify(() -> {
+                        assertEquals(200, response.statusCode());
+                        assertEquals(MetadataBase.MIME_TYPE, response.getHeader(HttpHeaders.CONTENT_TYPE));
                         assertEquals(emptyBucketResponse, response.body());
                         context.completeNow();
                     });

diff --git a/src/test/java/com/epam/aidial/core/util/ProxyUtilTest.java b/src/test/java/com/epam/aidial/core/util/ProxyUtilTest.java
@@ -88,6 +88,11 @@ public void testCollectAttachedFiles_ChatRequest() throws IOException {
                             "type": "binary/octet-stream",
                             "title": "Dockerfile",
                             "url": "files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/b1/Dockerfile"
+                          },
+                          {
+                            "type": "application/vnd.dial.metadata+json",
+                            "title": ".dockerignore",
+                            "url": "metadata/files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/b1/.dockerignore"
                           }
                         ]
                       }
@@ -101,11 +106,48 @@ public void testCollectAttachedFiles_ChatRequest() throws IOException {
         ProxyUtil.collectAttachedFiles(tree, link -> apiKeyData.getAttachedFiles().add(link));
 
         assertEquals(
-                Set.of("files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/b1/Dockerfile", "files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/b1/LICENSE"),
+                Set.of(
+                        "files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/b1/Dockerfile",
+                        "files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/b1/LICENSE",
+                        "files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/b1/.dockerignore"
+                ),
                 apiKeyData.getAttachedFiles()
         );
     }
 
+    @Test
+    public void testCollectAttachedFiles_Fail() throws IOException {
+        String content = """
+                {
+                  "modelId": "model",
+                  "messages": [
+                    {
+                      "content": "test",
+                      "role": "user",
+                      "custom_content": {
+                        "attachments": [
+                          {
+                            "type": "application/vnd.dial.metadata+json",
+                            "title": ".dockerignore",
+                            "url": "metadatata/files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/b1/.dockerignore"
+                          }
+                        ]
+                      }
+                    }
+                  ],
+                  "id": "id"
+                }
+                """;
+
+        ObjectNode tree = (ObjectNode) ProxyUtil.MAPPER.readTree(content.getBytes());
+        ApiKeyData apiKeyData = new ApiKeyData();
+
+        IllegalArgumentException error = assertThrows(IllegalArgumentException.class,
+                () -> ProxyUtil.collectAttachedFiles(tree, link -> apiKeyData.getAttachedFiles().add(link)));
+
+        assertEquals("Url of metadata attachment must start with metadata/: metadatata/files/7G9WZNcoY26Vy9D7bEgbv6zqbJGfyDp9KZyEbJR4XMZt/b1/.dockerignore", error.getMessage());
+    }
+
     @Test
     public void testCollectAttachedFiles_EmbeddingRequest() throws IOException {
         String content = """