From f1d175a5ea2cfd3098f8d7f484a98931d8495e74 Mon Sep 17 00:00:00 2001
From: Marcus Schiesser <mail@marcusschiesser.de>
Date: Tue, 14 Nov 2023 14:17:29 +0700
Subject: [PATCH] feat: support other image file types

---
 app/client/fetch/file.ts           | 28 +++++++++++++---------------
 app/client/fetch/url.ts            |  6 +++++-
 app/components/chat/chat-input.tsx | 11 +++++++----
 app/components/chat/chat.tsx       | 18 ++++++++++--------
 app/constant.ts                    | 19 ++++++++++++++++++-
 app/locales/en.ts                  |  2 +-
 app/store/session.ts               |  8 ++++++--
 7 files changed, 60 insertions(+), 32 deletions(-)

diff --git a/app/client/fetch/file.ts b/app/client/fetch/file.ts
index 53e7093f..4fd71a9f 100644
--- a/app/client/fetch/file.ts
+++ b/app/client/fetch/file.ts
@@ -1,24 +1,19 @@
 import { URLDetailContent } from "./url";
 import { FileWrap } from "../../utils/file";
+import {
+  ALLOWED_IMAGE_EXTENSIONS,
+  IMAGE_TYPES,
+  ImageType,
+} from "@/app/constant";
 
 export async function getDetailContentFromFile(
   file: FileWrap,
 ): Promise<URLDetailContent> {
-  switch (file.extension) {
-    case "pdf": {
-      return await getPDFFileDetail(file);
-    }
-    case "txt": {
-      return await getTextFileDetail(file);
-    }
-    case "jpg":
-    case "jpeg": {
-      return await getImageFileDetail(file);
-    }
-    default: {
-      throw new Error("Not supported file type");
-    }
-  }
+  if (file.extension === "pdf") return await getPDFFileDetail(file);
+  if (file.extension === "txt") return await getTextFileDetail(file);
+  if (ALLOWED_IMAGE_EXTENSIONS.includes(file.extension))
+    return await getImageFileDetail(file);
+  throw new Error("Not supported file type");
 }
 
 async function getPDFFileDetail(file: FileWrap): Promise<URLDetailContent> {
@@ -70,3 +65,6 @@ async function getImageFileDetail(file: FileWrap) {
   console.log(data);
   return data as URLDetailContent;
 }
+
+export const isImageFileType = (type: string) =>
+  IMAGE_TYPES.includes(type as ImageType);
diff --git a/app/client/fetch/url.ts b/app/client/fetch/url.ts
index 28e5dbf4..812cd69d 100644
--- a/app/client/fetch/url.ts
+++ b/app/client/fetch/url.ts
@@ -1,12 +1,16 @@
+import { DocumentType, ImageType } from "@/app/constant";
+
 export type Embedding = {
   text: string;
   embedding: number[];
 };
 
+export type UrlDetailType = DocumentType | ImageType;
+
 export type URLDetail = {
   url: string;
   size: number;
-  type: "text/html" | "application/pdf" | "text/plain" | "image/jpeg";
+  type: UrlDetailType;
   embeddings?: Embedding[];
 };
 
diff --git a/app/components/chat/chat-input.tsx b/app/components/chat/chat-input.tsx
index d966e591..0198783b 100644
--- a/app/components/chat/chat-input.tsx
+++ b/app/components/chat/chat-input.tsx
@@ -1,4 +1,7 @@
-import { getDetailContentFromFile } from "@/app/client/fetch/file";
+import {
+  getDetailContentFromFile,
+  isImageFileType,
+} from "@/app/client/fetch/file";
 import { URLDetail, URLDetailContent, isURL } from "@/app/client/fetch/url";
 import { Button } from "@/app/components/ui/button";
 import { Textarea } from "@/app/components/ui/textarea";
@@ -119,13 +122,13 @@ export default function ChatInput(props: ChatInputProps) {
     action: () => Promise<void>,
   ): Promise<void> => {
     let tempUrl: string;
-    if (file.type === "image/jpeg") {
+    if (isImageFileType(file.type)) {
       tempUrl = URL.createObjectURL(file);
       setTemporaryBlobUrl(tempUrl);
     }
 
     return action().finally(() => {
-      if (file.type === "image/jpeg") {
+      if (isImageFileType(file.type)) {
         URL.revokeObjectURL(tempUrl);
         setTemporaryBlobUrl(undefined);
       }
@@ -136,7 +139,7 @@ export default function ChatInput(props: ChatInputProps) {
     try {
       await manageTemporaryBlobUrl(fileInput.file, async () => {
         const fileDetail = await getDetailContentFromFile(fileInput);
-        if (fileInput.file.type === "image/jpeg") {
+        if (isImageFileType(fileInput.file.type)) {
           setImageFile(fileDetail);
         } else {
           callLLM({ fileDetail });
diff --git a/app/components/chat/chat.tsx b/app/components/chat/chat.tsx
index e2487500..c91beb8a 100644
--- a/app/components/chat/chat.tsx
+++ b/app/components/chat/chat.tsx
@@ -24,6 +24,7 @@ import { ChatAction } from "./chat-action";
 import ChatHeader from "./chat-header";
 import ChatInput from "./chat-input";
 import { ClearContextDivider } from "./clear-context-divider";
+import { isImageFileType } from "@/app/client/fetch/file";
 
 const Markdown = dynamic(
   async () => (await import("../ui/markdown")).Markdown,
@@ -94,7 +95,7 @@ export function Chat() {
   const renderMessages = useMemo(() => {
     const getFrontendMessages = (messages: ChatMessage[]) => {
       return messages.map((message) => {
-        if (!message.urlDetail || message.urlDetail.type === "image/jpeg")
+        if (!message.urlDetail || isImageFileType(message.urlDetail.type))
           return message;
         const urlTypePrefix = getUrlTypePrefix(message.urlDetail.type);
         const sizeInKB = Math.round(message.urlDetail.size / 1024);
@@ -256,13 +257,14 @@ export function Chat() {
                               : "bg-muted",
                           )}
                         >
-                          {message.urlDetail?.type === "image/jpeg" && (
-                            <img
-                              src={message.urlDetail.url}
-                              alt="Message image"
-                              className="object-contain w-full h-52 rounded-lg mb-2"
-                            />
-                          )}
+                          {message.urlDetail?.type &&
+                            isImageFileType(message.urlDetail.type) && (
+                              <img
+                                src={message.urlDetail.url}
+                                alt="Message image"
+                                className="object-contain w-full h-52 rounded-lg mb-2"
+                              />
+                            )}
                           <Markdown
                             content={message.content}
                             loading={
diff --git a/app/constant.ts b/app/constant.ts
index 46ab0667..6d9ff264 100644
--- a/app/constant.ts
+++ b/app/constant.ts
@@ -16,10 +16,27 @@ export const REQUEST_TIMEOUT_MS = 60000;
 export const CHAT_PAGE_SIZE = 15;
 export const MAX_RENDER_MSG_COUNT = 45;
 
-export const ALLOWED_IMAGE_EXTENSIONS = ["jpeg", "jpg"];
+export const ALLOWED_IMAGE_EXTENSIONS = ["jpeg", "jpg", "png", "gif", "webp"];
 export const ALLOWED_TEXT_EXTENSIONS = ["pdf", "txt"];
 export const ALLOWED_DOCUMENT_EXTENSIONS = [
   ...ALLOWED_TEXT_EXTENSIONS,
   ...ALLOWED_IMAGE_EXTENSIONS,
 ];
 export const DOCUMENT_FILE_SIZE_LIMIT = 1024 * 1024 * 10; // 10 MB
+
+export const DOCUMENT_TYPES = [
+  "text/html",
+  "application/pdf",
+  "text/plain",
+] as const;
+
+export type DocumentType = (typeof DOCUMENT_TYPES)[number];
+
+export const IMAGE_TYPES = [
+  "image/jpeg",
+  "image/png",
+  "image/gif",
+  "image/webp",
+] as const;
+
+export type ImageType = (typeof IMAGE_TYPES)[number];
diff --git a/app/locales/en.ts b/app/locales/en.ts
index c7970b7b..d5f97467 100644
--- a/app/locales/en.ts
+++ b/app/locales/en.ts
@@ -137,7 +137,7 @@ const en = {
   },
 
   Welcome: {
-    Title: "LlamaIndex Chat",
+    Title: "Chat LlamaIndex",
     SubTitle: "Create chat bots that know your data",
     Quote:
       "“This tool has saved me countless hours of work and helped me apply AI features to my work faster than ever before.”",
diff --git a/app/store/session.ts b/app/store/session.ts
index 728748f9..b3c5c43a 100644
--- a/app/store/session.ts
+++ b/app/store/session.ts
@@ -16,6 +16,7 @@ import {
 } from "../client/platforms/llm";
 import { prettyObject } from "../utils/format";
 import { Bot } from "./bot";
+import { isImageFileType } from "../client/fetch/file";
 
 export type ChatMessage = {
   role: MessageRole;
@@ -157,7 +158,7 @@ export async function callSession(
 
   let embeddings: Embedding[] | undefined;
   let message;
-  if (userMessage.urlDetail && userMessage.urlDetail.type !== "image/jpeg") {
+  if (userMessage.urlDetail && !isImageFileType(userMessage.urlDetail.type)) {
     // if the user sends document, let the LLM summarize the content of the URL and just use the document's embeddings
     message = "Summarize the given context briefly in 200 words or less";
     embeddings = userMessage.urlDetail?.embeddings;
@@ -168,7 +169,10 @@ export async function callSession(
       .flatMap((message: ChatMessage) => message.urlDetail?.embeddings)
       .filter((m) => m !== undefined) as Embedding[];
     embeddings = embeddings.length > 0 ? embeddings : undefined;
-    if (userMessage.urlDetail?.type === "image/jpeg") {
+    if (
+      userMessage.urlDetail?.type &&
+      isImageFileType(userMessage.urlDetail?.type)
+    ) {
       message = [
         {
           type: "text",