Support different kinds of images (#1116)

Author: logancyang

src/LLMProviders/chainRunner.ts

@@ -15,7 +15,9 @@ import {
   extractUniqueTitlesFromDocs,
   extractYoutubeUrl,
   formatDateTime,
-  imageToBase64,
+  ImageContent,
+  ImageProcessor,
+  MessageContent,
 } from "@/utils";
 import { Notice } from "obsidian";
 import ChainManager from "./chainManager";
@@ -251,6 +253,72 @@ class CopilotPlusChainRunner extends BaseChainRunner {
     return hasYoutubeCommand && youtubeUrl !== null && words.length === 1;
   }
 
+  private async processImageUrls(urls: string[]): Promise<ImageContent[]> {
+    try {
+      const imageUrls = await Promise.all(
+        urls.map(async (url) => {
+          if (await ImageProcessor.isImageUrl(url)) {
+            return await ImageProcessor.convertToBase64(url, this.chainManager.app.vault);
+          }
+          return null;
+        })
+      );
+
+      // Filter out null values and return valid image URLs
+      return imageUrls.filter((item): item is ImageContent => item !== null);
+    } catch (error) {
+      console.error("Error processing image URLs:", error);
+      return [];
+    }
+  }
+
+  private async processExistingImages(content: MessageContent[]): Promise<ImageContent[]> {
+    try {
+      const imageContent = await Promise.all(
+        content
+          .filter(
+            (item): item is ImageContent => item.type === "image_url" && !!item.image_url?.url
+          )
+          .map(async (item) => {
+            return await ImageProcessor.convertToBase64(
+              item.image_url.url,
+              this.chainManager.app.vault
+            );
+          })
+      );
+      return imageContent;
+    } catch (error) {
+      console.error("Error processing images:", error);
+      throw error;
+    }
+  }
+
+  private async buildMessageContent(
+    textContent: string,
+    userMessage: ChatMessage
+  ): Promise<MessageContent[]> {
+    const content: MessageContent[] = [
+      {
+        type: "text",
+        text: textContent,
+      },
+    ];
+
+    // Process URLs in the message to identify images
+    if (userMessage.context?.urls && userMessage.context.urls.length > 0) {
+      const imageContents = await this.processImageUrls(userMessage.context.urls);
+      content.push(...imageContents);
+    }
+
+    // Add existing image content if present
+    if (userMessage.content && userMessage.content.length > 0) {
+      const imageContents = await this.processExistingImages(userMessage.content);
+      content.push(...imageContents);
+    }
+
+    return content;
+  }
+
   private async streamMultimodalResponse(
     textContent: string,
     userMessage: ChatMessage,
@@ -289,39 +357,8 @@ class CopilotPlusChainRunner extends BaseChainRunner {
       messages.push({ role: "assistant", content: ai });
     }
 
-    // Create content array for current message
-    const content: Array<{ type: string } & ({ text: string } | { image_url: { url: string } })> = [
-      {
-        type: "text",
-        text: textContent,
-      },
-    ];
-
-    // Add image content if present
-    if (userMessage.content && userMessage.content.length > 0) {
-      try {
-        const imageContent = await Promise.all(
-          userMessage.content
-            .filter((item) => item.type === "image_url" && item.image_url?.url)
-            .map(async (item) => {
-              const base64Image = await imageToBase64(
-                item.image_url.url,
-                this.chainManager.app.vault
-              );
-              return {
-                type: "image_url",
-                image_url: {
-                  url: base64Image,
-                },
-              };
-            })
-        );
-        content.push(...imageContent);
-      } catch (error) {
-        console.error("Error processing images:", error);
-        throw error;
-      }
-    }
+    // Build message content with text and images
+    const content = await this.buildMessageContent(textContent, userMessage);
 
     // Add current user message
     messages.push({

src/components/Chat.tsx

@@ -165,13 +165,14 @@ const Chat: React.FC<ChatProps> = ({
     const urlContextAddition =
       currentChain === ChainType.COPILOT_PLUS_CHAIN
         ? await mention.processUrls(inputMessage || "")
-        : "";
+        : { urlContext: "", imageUrls: [] };
 
     // Add context notes
     const noteContextAddition = await processContextNotes(customPromptProcessor, fileParserManager);
 
     // Combine everything
-    processedUserMessage = processedUserMessage + urlContextAddition + noteContextAddition;
+    processedUserMessage =
+      processedUserMessage + urlContextAddition.urlContext + noteContextAddition;
 
     let messageWithToolCalls = inputMessage;
     // Add tool calls last
@@ -186,6 +187,13 @@ const Chat: React.FC<ChatProps> = ({
       isVisible: false,
       timestamp: timestamp,
       content: content,
+      context: {
+        notes,
+        urls:
+          currentChain === ChainType.COPILOT_PLUS_CHAIN
+            ? [...(urls || []), ...urlContextAddition.imageUrls]
+            : urls || [],
+      },
     };
 
     // Add hidden user message to chat history

src/mentions/Mention.ts

@@ -1,5 +1,5 @@
 import { BrevilabsClient, Url4llmResponse } from "@/LLMProviders/brevilabsClient";
-import { isYoutubeUrl } from "@/utils";
+import { ImageProcessor, isYoutubeUrl } from "@/utils";
 
 export interface MentionData {
   type: string;
@@ -50,17 +50,24 @@ export class Mention {
   }
 
   // For non-youtube URLs
-  async processUrls(text: string): Promise<string> {
+  async processUrls(text: string): Promise<{ urlContext: string; imageUrls: string[] }> {
     const urls = this.extractUrls(text);
     let urlContext = "";
+    const imageUrls: string[] = [];
 
     // Return empty string if no URLs to process
     if (urls.length === 0) {
-      return "";
+      return { urlContext: "", imageUrls: [] };
     }
 
     // Process all URLs concurrently
     const processPromises = urls.map(async (url) => {
+      // Check if it's an image URL
+      if (await ImageProcessor.isImageUrl(url)) {
+        imageUrls.push(url);
+        return null;
+      }
+
       if (!this.mentions.has(url)) {
         const processed = await this.processUrl(url);
         this.mentions.set(url, {
@@ -81,7 +88,7 @@ export class Mention {
       }
     });
 
-    return urlContext;
+    return { urlContext, imageUrls };
   }
 
   getMentions(): Map<string, MentionData> {