[OPIK-428] Detect images in the trace/span input by URL extension

comet-ml · Nov 15, 2024 · d3559f1 · d3559f1
1 parent f5cb060
commit d3559f1
Showing 1 changed file with 44 additions and 0 deletions.
diff --git a/.../opik-frontend/src/components/shared/TraceDetailsPanel/TraceDataViewer/InputOutputTab.tsx b/.../opik-frontend/src/components/shared/TraceDetailsPanel/TraceDataViewer/InputOutputTab.tsx
@@ -53,8 +53,45 @@ const BASE64_PREFIXES_MAP = {
   UklGR: "webp",
 } as const;
 
+const IMAGE_URL_EXTENSIONS = [
+  "apng",
+  "avif",
+  "bmp",
+  "cr2",
+  "djv",
+  "djvu",
+  "eps",
+  "gif",
+  "hdp",
+  "heic",
+  "heif",
+  "ico",
+  "j2k",
+  "jp2",
+  "jpeg",
+  "jpf",
+  "jpg",
+  "jpm",
+  "jxr",
+  "mj2",
+  "nef",
+  "orf",
+  "png",
+  "psd",
+  "raw",
+  "sr2",
+  "svg",
+  "tif",
+  "tiff",
+  "wdp",
+  "webp",
+] as const;
+
 const IMAGE_CHARS_REGEX = "[A-Za-z0-9+/]+={0,2}";
 const DATA_IMAGE_PREFIX = `"data:image/[^;]{3,4};base64,${IMAGE_CHARS_REGEX}"`;
+const IMAGE_URL_REGEX = `"https?:\\/\\/[^\\s"']+\\.(${IMAGE_URL_EXTENSIONS.join(
+  "|",
+)})(\\?[^"']*)?(#[^"']*)?"`;
 
 function extractInputImages(input?: object) {
   if (!input) return [];
@@ -84,6 +121,13 @@ function extractInputImages(input?: object) {
     images.push(...dataImageMatches.map((match) => match.replace(/"/g, "")));
   }
 
+  // Extract image URLs
+  const imageUrlRegex = new RegExp(IMAGE_URL_REGEX, "gi");
+  const imageUrlMatches = stringifiedInput.match(imageUrlRegex);
+  if (imageUrlMatches) {
+    images.push(...imageUrlMatches.map((match) => match.replace(/"/g, "")));
+  }
+
   return images;
 }