added support for image recognition models

src/main/main.ts

@@ -10,11 +10,12 @@
  */
 require('dotenv').config();
 import path from 'path';
-import { app, BrowserWindow, shell, ipcMain } from 'electron';
+import { app, BrowserWindow, shell, ipcMain, dialog } from 'electron';
 import { autoUpdater } from 'electron-updater';
 import log from 'electron-log';
 import MenuBuilder from './menu';
 import { resolveHtmlPath } from './util';
+const fs = require('fs');
 
 class AppUpdater {
   constructor() {
@@ -32,6 +33,44 @@ ipcMain.on('ipc-example', async (event, arg) => {
   // event.reply('ipc-example', msgTemplate('pong'));
 });
 
+ipcMain.on('select-image', async (event) => {
+  const result = await dialog.showOpenDialog({
+    properties: ['openFile'],
+    filters: [{ name: 'Images', extensions: ['jpg', 'png', 'gif'] }]
+  });
+
+  if (!result.canceled) {
+
+    const imagePath = result.filePaths[0];
+
+    // Read the image file
+    fs.readFile(imagePath, (err: any, data: any) => {
+      if (err) {
+        console.error('Error reading image file:', err);
+        return;
+      }
+
+      // Convert the image data to a base64 string
+      const base64Image = Buffer.from(data).toString('base64');
+
+      event.reply('image-selected', { base64: base64Image, path: imagePath });
+    });
+
+
+  }
+});
+
+// ipcMain.handle('select-image', async (event) => {
+//   const result = await dialog.showOpenDialog({
+//     properties: ['openFile'],
+//     filters: [{ name: 'Images', extensions: ['jpg', 'png', 'gif'] }]
+//   })
+
+//   if (!result.canceled) {
+//     return result.filePaths[0]
+//   }
+// })
+
 if (process.env.NODE_ENV === 'production') {
   const sourceMapSupport = require('source-map-support');
   sourceMapSupport.install();

src/renderer/App.css

@@ -97,4 +97,17 @@
 *::-webkit-scrollbar-thumb {
   background-color: #2b2b2b;
   border-radius: 5px;
+}
+
+.fade-in {
+  animation: fadeIn 0.3s ease-in-out;
+}
+
+@keyframes fadeIn {
+  from {
+    opacity: 0;
+  }
+  to {
+    opacity: 1;
+  }
 }

src/renderer/components/Chat.tsx

@@ -1,5 +1,5 @@
 import { useEffect, useRef, useState } from 'react';
-import { MESSAGE } from '../utils/interfaces';
+import { MESSAGE, SELECTED_IMAGE } from '../utils/interfaces';
 import MessageContainer from './MessageContainer';
 import TopBar from './TopBar';
 import { useAtom } from 'jotai';
@@ -16,6 +16,8 @@ import {
 } from '../utils/managers/conversationManager';
 import { BounceLoader } from 'react-spinners';
 import { getSystemPrompt } from '../utils/utils';
+import { Image } from 'react-feather';
+import SelectedImages from './SelectedImages';
 let window: any = global;
 
 export default function Chat() {
@@ -29,7 +31,7 @@ export default function Chat() {
   const [input, setInput] = useState<string>('');
   const [pending, setPending] = useState<boolean>(false);
   const [_, forceUpdate] = useState<number>(0);
-  const [checkedSettings, setCheckedSettings] = useState<boolean>(false);
+  const [selectedImages, setSelectedImages] = useState<SELECTED_IMAGE[]>([]);
 
   const messagesRef = useRef<MESSAGE[]>([]);
   const containerRef = useRef<HTMLDivElement>(null);
@@ -67,14 +69,6 @@ export default function Chat() {
     pending ? 100 : null,
   );
 
-  // const checkSettings = () => {
-  //   const settings = loadSettings() as SETTINGS;
-  //   if (Object.entries(settings).length > 0 && settings.lastConversation) {
-  //     setCurrentConversation(settings.lastConversation);
-  //   }
-  //   setCheckedSettings(true);
-  // };
-
   const scrollToBottom = () => {
     containerRef.current?.scrollTo(0, containerRef.current?.scrollHeight);
   };
@@ -116,14 +110,20 @@ export default function Chat() {
     }
   };
 
+  const addSelectedImages = () => {
+    if (selectedImages.length === 0) return {};
+    else
+      return {
+        images: selectedImages.map((image) => image.base64),
+      };
+  };
+
   const formatRequest = async () => {
     // get the server address from the environment variables
     const addr =
       window.envVars.OLLAMA_SERVER_ADDRESS || 'http://localhost:11434';
     // store the user input before clearing
     const userInput = input;
-    // clear input
-    setInput('');
     // setup request body based on chat type
     let body = {};
     if (chatType === 'chat') {
@@ -138,6 +138,7 @@ export default function Chat() {
               currentModelName as string,
               message.content,
             ),
+            ...addSelectedImages(),
           };
         }),
       };
@@ -147,8 +148,12 @@ export default function Chat() {
         prompt: userInput,
         options: modelOptions,
         template: getSystemPrompt(currentModelName as string, userInput),
+        ...addSelectedImages(),
       };
     }
+    // clear input
+    setInput('');
+    setSelectedImages([]);
     // request a response from the assistant
     const response = await fetch(`${addr}/api/${chatType}`, {
       method: 'POST',
@@ -168,6 +173,7 @@ export default function Chat() {
         role: 'user',
         content: input,
         timestamp: Date.now(),
+        images: selectedImages,
       };
       // add the new message to the messages array
       messagesRef.current = [
@@ -187,7 +193,6 @@ export default function Chat() {
           currentModelName as string,
         );
         setCurrentConversation(newConversation);
-        console.log(newConversation);
         conversationUid = newConversation.uid;
       } else {
         conversationUid = currentConversation?.uid;
@@ -214,26 +219,73 @@ export default function Chat() {
     }
   };
 
+  const isUsingImageModel = () => {
+    return currentModelName === 'llava:latest';
+  };
+
+  const renderImageSelection = () => {
+    // prompt the user to select an image from their computer
+    const selectImage = async () => {
+      // setup listener for a response from the main process
+      window.electron.ipcRenderer.once('image-selected', (newImageData: {
+        base64: string, 
+        path: string
+      }) => {
+        if (newImageData) {
+          setSelectedImages([...selectedImages, newImageData]);
+        }
+      });
+
+      // send a message to the main process to open the file dialog
+      await window.electron.ipcRenderer.sendMessage('select-image', [
+        'select-image',
+      ]);
+    };
+
+    if (isUsingImageModel()) {
+      return (
+        <div
+          className="cursor-pointer text-white opacity-50 hover:opacity-100 transition-all"
+          onClick={selectImage}
+        >
+          <Image
+            className="absolute left-4 top-1/2 transform -translate-y-1/2"
+            size={18}
+          />
+        </div>
+      );
+    } else return null;
+  };
+
   const renderInput = () => {
+    // set left padding based on if an image model is being used
+    const leftPadding = isUsingImageModel() ? 'pl-11' : '';
     return (
-      <div className="relative w-full h-auto flex">
-        <textarea
-          className="w-full h-12 border-[1px] border-zinc-700 min-h-10 px-3 rounded-md bg-zinc-800 text-white outline-zinc-900 pt-[10px] pr-10"
-          disabled={!currentModelName}
-          placeholder={`${
-            !currentModelName ? 'No model selected' : 'Type a message...'
-          }`}
-          value={input}
-          onChange={(e) => setInput(e.target.value)}
-          onKeyDown={keyDown}
+      <div className="flex flex-col">
+        <SelectedImages 
+          selectedImages={selectedImages}
+          setSelectedImages={setSelectedImages}
         />
-        {pending && (
-          <BounceLoader
-            className="!absolute right-3 top-1/2 transform -translate-y-1/2"
-            size={20}
-            color="rgb(96 165 250)"
+        <div className="relative w-full h-auto flex">
+          {renderImageSelection()}
+          <textarea
+            className={`w-full h-12 border-[1px] border-zinc-700 min-h-10 px-3 ${leftPadding} rounded-md bg-zinc-800 text-white outline-zinc-900 pt-[10px] pr-10`}
+            disabled={!currentModelName}
+            placeholder={`${
+              !currentModelName ? 'No model selected' : 'Type a message...'
+            }`}
+            value={input}
+            onChange={(e) => setInput(e.target.value)}
+            onKeyDown={keyDown}
           />
-        )}
+          {pending && (
+            <BounceLoader
+              className="!absolute right-3 top-1/2 transform -translate-y-1/2"
+              size={20}
+              color="rgb(96 165 250)"
+            />
+          )}
+        </div>
       </div>
     );
   };

src/renderer/components/MessageContainer.tsx

@@ -1,22 +1,24 @@
 import { MESSAGE } from '../utils/interfaces';
-import { Monitor, User, Zap } from 'react-feather';
+import { User, Monitor } from 'react-feather';
 import ReactMarkdown from 'react-markdown';
 import { BeatLoader } from 'react-spinners';
+import { useState } from 'react';
 
 export default function MessageContainer(props: { message: MESSAGE }) {
   const { message } = props;
+  const [selectedImage, setSelectedImage] = useState<number>(-1);
 
   const renderIcon = () => {
     // is the message from the user or the assistant
     const isUser = message.role === 'user';
     // icon sizing and style
     const iconSize = 15;
-    const iconStyle = 'transform translate-y-[5px] text-blue-400';
+    const iconStyle = 'transform translate-y-[5px] text-blue-400 z-1';
     // render the icon based on the role
     const icon = isUser ? (
       <User size={iconSize} className={iconStyle} />
     ) : (
-      <Zap size={iconSize} className={iconStyle} />
+      <Monitor size={iconSize} className={iconStyle} />
     );
     return <div className="">{icon}</div>;
   };
@@ -37,14 +39,64 @@ export default function MessageContainer(props: { message: MESSAGE }) {
     else return regularMessage;
   };
 
+  const renderAttachedImages = () => {
+    // if there are no images attached to the message return null
+    if (!message.images || message.images.length === 0) return null
+    return (
+      <div className="min-w-full flex gap-2 justify-start overflow-x-auto">
+        {message.images.map((image, index) => {
+          return (
+            <img
+              key={index}
+              src={`data:image/jpeg;base64,${image.base64}`}
+              alt="attached image"
+              className="w-20 h-20 object-cover rounded-md mb-2 cursor-pointer"
+              onClick={() => setSelectedImage(index)}
+            />
+          );
+        })}
+      </div>
+    );
+  }
+
+  // render a full screen view of the selected image
+  const renderFullScreenImage = () => {
+    // check if there is an image selected
+    if (
+      selectedImage === -1 ||
+      !message.images
+    ) return null
+    return (
+      <div 
+        className="fixed left-0 top-0 w-full h-full flex justify-center items-center z-10 fade-in"
+        onClick={() => {
+          setSelectedImage(-1);
+        }}
+      >
+        <div
+          className="fixed left-0 top-0 w-full h-full bg-black opacity-90 z-[-1]"
+        />
+        <img
+          src={`data:image/jpeg;base64,${message.images[selectedImage].base64}`}
+          alt="full screen image"
+          className="w-[90%] h-auto object-cover"
+        />
+      </div>
+    )
+  }
+
   const renderMessage = () => {
     return (
       <div className="flex gap-3 items-start">
         {renderIcon()}
-        <div className="flex flex-col markdown-content pr-2">
+        <div className="flex flex-col pr-2 w-full">
           <span className="font-bold capitalize">{message.role}</span>
-          {messageLoadingLogic()}
+          <div className="markdown-content">
+            {messageLoadingLogic()}
+          </div>
+          {renderAttachedImages()}
         </div>
+        {renderFullScreenImage()}
       </div>
     );
   };