From 64af65ac6166009e9f744dbe398a5f6a956672ef Mon Sep 17 00:00:00 2001
From: zaidmukaddam <zaidaltafmukaddam@gmail.com>
Date: Mon, 23 Dec 2024 11:34:55 +0530
Subject: [PATCH] chore: add image URL validation

---
 app/api/chat/route.ts | 54 ++++++++++++++++++++++++++++++++++---------
 1 file changed, 43 insertions(+), 11 deletions(-)

diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts
index 5e73d5e..a1147db 100644
--- a/app/api/chat/route.ts
+++ b/app/api/chat/route.ts
@@ -114,6 +114,24 @@ function sanitizeUrl(url: string): string {
   return url.replace(/\s+/g, '%20')
 }
 
+async function isValidImageUrl(url: string): Promise<boolean> {
+  try {
+    const controller = new AbortController();
+    const timeout = setTimeout(() => controller.abort(), 5000);
+
+    const response = await fetch(url, {
+      method: 'HEAD',
+      signal: controller.signal
+    });
+
+    clearTimeout(timeout);
+
+    return response.ok && (response.headers.get('content-type')?.startsWith('image/') ?? false);
+  } catch {
+    return false;
+  }
+}
+
 const defaultsystemPrompt = `
 You are an expert AI web search engine called MiniPerplx, that helps users find information on the internet with no bullshit talks.
 Always start with running the tool(s) and then and then only write your response AT ALL COSTS!!
@@ -285,18 +303,32 @@ export async function POST(req: Request) {
                 published_date: topics[index] === "news" ? obj.published_date : undefined,
               })),
               images: includeImageDescriptions
-                ? data.images
-                  .map(({ url, description }: { url: string; description?: string }) => ({
-                    url: sanitizeUrl(url),
-                    description: description ?? ''
-                  }))
-                  .filter(
-                    (image: { url: string; description: string }): image is { url: string; description: string } =>
-                      typeof image === 'object' &&
-                      image.description !== undefined &&
-                      image.description !== ''
+                ? await Promise.all(
+                  data.images
+                    .map(async ({ url, description }: { url: string; description?: string }) => {
+                      const sanitizedUrl = sanitizeUrl(url);
+                      const isValid = await isValidImageUrl(sanitizedUrl);
+
+                      return isValid ? {
+                        url: sanitizedUrl,
+                        description: description ?? ''
+                      } : null;
+                    })
+                ).then(results =>
+                  results.filter((image): image is { url: string; description: string } =>
+                    image !== null &&
+                    typeof image === 'object' &&
+                    typeof image.description === 'string' &&
+                    image.description !== ''
                   )
-                : data.images.map(({ url }: { url: string }) => sanitizeUrl(url))
+                )
+                : await Promise.all(
+                  data.images
+                    .map(async ({ url }: { url: string }) => {
+                      const sanitizedUrl = sanitizeUrl(url);
+                      return await isValidImageUrl(sanitizedUrl) ? sanitizedUrl : null;
+                    })
+                ).then(results => results.filter((url): url is string => url !== null))
             };
           });