From f2f037e10a6900a7685b2669ca1a6afcc777252c Mon Sep 17 00:00:00 2001
From: Fernando Carril <fernando@comet.com>
Date: Fri, 15 Nov 2024 15:38:00 +0100
Subject: [PATCH] [OPIK-428] Detect images in the trace/span input by URL
 extension (#647)

---
 .../TraceDataViewer/InputOutputTab.tsx        | 44 +++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/apps/opik-frontend/src/components/shared/TraceDetailsPanel/TraceDataViewer/InputOutputTab.tsx b/apps/opik-frontend/src/components/shared/TraceDetailsPanel/TraceDataViewer/InputOutputTab.tsx
index 948dd0148d..86aa6c3fbf 100644
--- a/apps/opik-frontend/src/components/shared/TraceDetailsPanel/TraceDataViewer/InputOutputTab.tsx
+++ b/apps/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;
 }