Prompt images (#25)

* image upload feature mostly finished

* Updated UI

* Fixed backend logic error + frontend lint issue

* fix: apply frontend code formatting fixes

---------

Co-authored-by: Erina <erinakibria@gmail.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: 3 people

configs/models.txt

@@ -1,2 +1,2 @@
-llama3.2:1b
-qwen2.5:0.5b
+qwen2.5:0.5b
+llava-llama3:8b

internal/controllers/post_conversation_chat.go

@@ -17,9 +17,10 @@ import (
 )
 
 type postConversationChatRequest struct {
-	Prompt         string `json:"prompt"`
-	Model          string `json:"model"`
-	ConversationID string `json:"conversationId"`
+	Prompt         string   `json:"prompt"`
+	Images         []string `json:"images"`
+	Model          string   `json:"model"`
+	ConversationID string   `json:"conversationId"`
 }
 
 func (i *Injection) PostConversationChat(w http.ResponseWriter, r *http.Request) {
@@ -59,7 +60,7 @@ func (i *Injection) PostConversationChat(w http.ResponseWriter, r *http.Request)
 	// Construct request to Ollama API
 	ollamaReq := models.OllamaChatRequest{
 		Model:    req.Model,
-		Messages: utils.BuildOllamaMessages(systemPrompt, req.Prompt, prevChatRecords),
+		Messages: utils.BuildOllamaMessages(systemPrompt, req.Prompt, req.Images, prevChatRecords),
 		Stream:   true,
 	}
 
@@ -110,10 +111,6 @@ func (i *Injection) PostConversationChat(w http.ResponseWriter, r *http.Request)
 		return
 	}
 
-	// Images for now
-	// TODO
-	promptImages := []string{}
-
 	// Aggregate saving the chat record in the database and cache
 	wg := sync.WaitGroup{}
 	wg.Add(2)
@@ -123,7 +120,7 @@ func (i *Injection) PostConversationChat(w http.ResponseWriter, r *http.Request)
 	go func() {
 		defer wg.Done()
 
-		if _, err := db.SaveChatRecord(i.Sdb, req.Prompt, promptImages, replyBuilder.String(), creationTime, user.ID, req.ConversationID); err != nil {
+		if _, err := db.SaveChatRecord(i.Sdb, req.Prompt, req.Images, replyBuilder.String(), creationTime, user.ID, req.ConversationID); err != nil {
 			errorChan <- err
 		}
 	}()
@@ -135,7 +132,7 @@ func (i *Injection) PostConversationChat(w http.ResponseWriter, r *http.Request)
 		record := models.ClientChatRecord{
 			CreatedAt:    creationTime,
 			Prompt:       req.Prompt,
-			PromptImages: promptImages,
+			PromptImages: req.Images,
 			Reply:        replyBuilder.String(),
 		}
 

internal/controllers/post_guest_chat.go

@@ -14,8 +14,9 @@ import (
 )
 
 type postGuestChatRequest struct {
-	Prompt string `json:"prompt"`
-	Model  string `json:"model"`
+	Prompt string   `json:"prompt"`
+	Images []string `json:"images"`
+	Model  string   `json:"model"`
 }
 
 // Endpoint for sending a prompt on a guest chat.
@@ -57,7 +58,7 @@ func (i *Injection) PostGuestChat(w http.ResponseWriter, r *http.Request) {
 	// Construct request to Ollama API.
 	ollamaReq := models.OllamaChatRequest{
 		Model:    req.Model,
-		Messages: utils.BuildOllamaMessages(systemPrompt, req.Prompt, prevChatRecords),
+		Messages: utils.BuildOllamaMessages(systemPrompt, req.Prompt, req.Images, prevChatRecords),
 		Stream:   true,
 	}
 
@@ -108,15 +109,11 @@ func (i *Injection) PostGuestChat(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// Images for now.
-	// TODO.
-	promptImages := []string{}
-
 	// Cache chat in Redis.
 	record := models.ClientChatRecord{
 		CreatedAt:    creationTime,
 		Prompt:       req.Prompt,
-		PromptImages: promptImages,
+		PromptImages: req.Images,
 		Reply:        replyBuilder.String(),
 	}
 

internal/utils/build_ollama_messages.go

@@ -12,7 +12,7 @@ import "tokenbase/internal/models"
 //
 // Returns:
 // - A list of Ollama compatible chat messages.
-func BuildOllamaMessages(systemPrompt string, newPrompt string, records []models.ClientChatRecord) []models.OllamaChatMessage {
+func BuildOllamaMessages(systemPrompt string, newPrompt string, newImages []string, records []models.ClientChatRecord) []models.OllamaChatMessage {
 	// Allocate enough space for the messages.
 	// 2 messages per record + 2 messages for the system prompt and the new prompt.
 	n := (len(records) * 2) + 2
@@ -39,6 +39,7 @@ func BuildOllamaMessages(systemPrompt string, newPrompt string, records []models
 		newMessage := models.OllamaChatMessage{
 			Role:    models.UserRole,
 			Content: newPrompt,
+			Images:  newImages,
 		}
 
 		messages = append(messages, newMessage)

web/src/components/Chat.tsx

@@ -3,12 +3,18 @@ import ButtonColor from "@/models/ButtonColor";
 import IconButton from "./IconButton";
 import TypesetRenderer from "./TypesetRenderer";
 import TypeCursor from "./TypeCursor";
-import { faCopy, faTrash } from "@fortawesome/free-solid-svg-icons";
 import { useState } from "react";
+import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
+import {
+  faCopy,
+  faPaperclip,
+  faTrash,
+} from "@fortawesome/free-solid-svg-icons";
 
 type ChatProps = {
   createdAt: number;
   prompt: string;
+  images: string[];
   reply: string;
   isComplete: boolean;
   onDelete?: (createdAt: number) => void;
@@ -17,6 +23,7 @@ type ChatProps = {
 const Chat: React.FC<ChatProps> = ({
   createdAt,
   prompt,
+  images,
   reply,
   isComplete,
   onDelete,
@@ -35,6 +42,12 @@ const Chat: React.FC<ChatProps> = ({
   return (
     <div className={styles.container}>
       <div className={styles.promptContainer}>
+        {images.length > 0 && (
+          <div className={styles.imageCount}>
+            <FontAwesomeIcon icon={faPaperclip} />
+            {images.length}
+          </div>
+        )}
         <TypesetRenderer>{prompt}</TypesetRenderer>
       </div>
       {reply.length > 0 && (

web/src/components/ChatContainer.tsx

@@ -23,7 +23,10 @@ type ChatContainerProps = {
   promptEndpoint: string;
   deleteEndpoint: string;
   suggestions: string[];
-  constructPromptRequest: (prompt: string) => Promise<Result<HttpChatReq>>;
+  constructPromptRequest: (
+    prompt: string,
+    promptImages: string[],
+  ) => Promise<Result<HttpChatReq>>;
   constructDeleteRequest: (createdAt: number) => Promise<Result<HttpChatReq>>;
 };
 
@@ -43,13 +46,14 @@ const ChatContainer: React.FC<ChatContainerProps> = ({
   const loadingIndicatorRef = useRef<HTMLDivElement | null>(null);
 
   const onPromptSend = useCallback(
-    async (prompt: string) => {
+    async (prompt: string, promptImages: string[]) => {
       // Clear any previous error
       setError(null);
 
       // Initialize a new chat entry
       let newChat: ChatRecord = {
-        prompt: prompt,
+        prompt,
+        promptImages,
         reply: "",
       };
 
@@ -58,7 +62,11 @@ const ChatContainer: React.FC<ChatContainerProps> = ({
       setStreamingChat(newChat);
 
       // Construct the request using the provided callback
-      const { ok, error } = await constructPromptRequest(prompt);
+
+      const { ok, error } = await constructPromptRequest(
+        prompt,
+        promptImages.map((img) => img.substring(img.search(/,/) + 1)),
+      );
 
       if (error) {
         setError(error);
@@ -197,7 +205,10 @@ const ChatContainer: React.FC<ChatContainerProps> = ({
           <div className={styles.suggestionsContainer}>
             {/* Render some random chat suggestions */}
             {suggestions?.map((suggestion, i) => (
-              <StandardButton key={i} onClick={() => onPromptSend(suggestion)}>
+              <StandardButton
+                key={i}
+                onClick={() => onPromptSend(suggestion, [])}
+              >
                 {suggestion}
               </StandardButton>
             ))}
@@ -212,6 +223,7 @@ const ChatContainer: React.FC<ChatContainerProps> = ({
                 key={i}
                 createdAt={chat.createdAt ?? -1}
                 prompt={chat.prompt}
+                images={chat.promptImages ?? []}
                 reply={chat.reply}
                 isComplete={true}
                 onDelete={onChatDelete}
@@ -223,6 +235,7 @@ const ChatContainer: React.FC<ChatContainerProps> = ({
               <Chat
                 createdAt={streamingChat.createdAt ?? -1}
                 prompt={streamingChat.prompt}
+                images={streamingChat.promptImages ?? []}
                 reply={streamingChat.reply}
                 isComplete={false}
               />

web/src/components/GuestChat.tsx

@@ -26,7 +26,10 @@ const GuestChat: React.FC<GuestChatProps> = ({ chatSuggestions }) => {
   const { bearer, setBearer } = useBearerContext();
   const { availableModels, selectedModelIndex } = useModelsContext();
 
-  const constructGuestPromptRequest = async (prompt: string) => {
+  const constructGuestPromptRequest = async (
+    prompt: string,
+    images: string[],
+  ) => {
     const constructReq = (token: string) => {
       if (availableModels.length === 0) {
         return {
@@ -41,6 +44,7 @@ const GuestChat: React.FC<GuestChatProps> = ({ chatSuggestions }) => {
         },
         body: JSON.stringify({
           prompt,
+          images,
           model: availableModels[selectedModelIndex].tag,
         }),
       };

web/src/components/PromptArea.tsx

@@ -1,7 +1,7 @@
 import styles from "@/styles/components/PromptArea.module.scss";
 import { useState, useRef, useEffect, useCallback } from "react";
 import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
-import { merriweather500 } from "@/utils/fonts";
+import { merriweather400, merriweather500 } from "@/utils/fonts";
 import {
   faArrowUp,
   faPaperclip,
@@ -12,7 +12,7 @@ type PromptAreaProps = {
   isDisabled: boolean;
   canCancel: boolean;
   canAttach: boolean;
-  onSend: (prompt: string) => void;
+  onSend: (prompt: string, promptImages: string[]) => void;
   onCancel: () => void;
 };
 
@@ -24,16 +24,19 @@ const PromptArea: React.FC<PromptAreaProps> = ({
   onCancel,
 }) => {
   const [prompt, setPrompt] = useState("");
+  const [attachedImages, setAttachedImages] = useState<string[]>([]);
   const textareaRef = useRef<HTMLTextAreaElement | null>(null);
 
   const handleSend = () => {
-    onSend(prompt.trim());
+    onSend(prompt.trim(), attachedImages);
+    setAttachedImages([]);
     setPrompt("");
   };
 
   const handleKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
     if (e.key === "Enter" && !e.shiftKey) {
       e.preventDefault();
+
       if (prompt.trim().length > 0 && !isDisabled) {
         handleSend();
       }
@@ -44,6 +47,33 @@ const PromptArea: React.FC<PromptAreaProps> = ({
     setPrompt(e.target.value);
   };
 
+  const handleImageUpload = (e: React.ChangeEvent<HTMLInputElement>) => {
+    const files = e.target.files;
+
+    if (files) {
+      const uploadedImages: string[] = [];
+
+      Array.from(files).forEach((file) => {
+        const reader = new FileReader();
+        reader.onload = (event) => {
+          const target = event.target;
+
+          if (target === null) {
+            return;
+          }
+
+          uploadedImages.push(target.result as string);
+          setAttachedImages((prev) => [...prev, target.result as string]);
+        };
+
+        reader.readAsDataURL(file);
+      });
+    }
+
+    // Reset the file input
+    e.target.value = "";
+  };
+
   const adjustTextareaHeight = useCallback(() => {
     if (textareaRef.current) {
       textareaRef.current.style.height = "auto";
@@ -77,8 +107,22 @@ const PromptArea: React.FC<PromptAreaProps> = ({
           spellCheck={false}
         />
         <div className={styles.buttonContainer}>
-          <button disabled={isDisabled || !canAttach}>
+          <button type="button" disabled={isDisabled || !canAttach}>
+            {attachedImages.length > 0 && (
+              <span
+                className={`${styles.imageCount} ${merriweather400.className}`}
+              >
+                {Math.min(attachedImages.length, 99)}
+              </span>
+            )}
             <FontAwesomeIcon icon={faPaperclip} />
+            <input
+              id="file-upload"
+              type="file"
+              multiple
+              accept="image/*"
+              onChange={handleImageUpload}
+            />
           </button>
           {canCancel ? (
             <button onClick={onCancel}>

web/src/components/UserChat.tsx

@@ -28,7 +28,10 @@ const UserChat: React.FC<UserChatProps> = ({ chatSuggestions }) => {
     throw new Error("Bearer token is not available");
   }
 
-  const constructUserPromptRequest = async (prompt: string) => {
+  const constructUserPromptRequest = async (
+    prompt: string,
+    images: string[],
+  ) => {
     if (conversationRecords === null) {
       return {
         error: "Cannot create prompt request without conversation records",
@@ -49,6 +52,7 @@ const UserChat: React.FC<UserChatProps> = ({ chatSuggestions }) => {
         },
         body: JSON.stringify({
           prompt,
+          images,
           model: availableModels[selectedModelIndex].tag,
           conversationId,
         }),
@@ -60,6 +64,7 @@ const UserChat: React.FC<UserChatProps> = ({ chatSuggestions }) => {
       const newConversation = await reqNewConversation(
         availableModels[selectedModelIndex].tag,
         prompt,
+        images,
         bearer.token,
       );
 

web/src/pages/_app.tsx

@@ -21,8 +21,10 @@ class RootApp extends App<RootAppProps> {
   static async getInitialProps(appContext: AppContext) {
     const appProps = await App.getInitialProps(appContext);
 
+    // Only fetch available models on the home page
     return {
-      availableModels: await getAvailableModels(),
+      availableModels:
+        appContext.router.pathname === "/" ? await getAvailableModels() : [],
       ...appProps,
     };
   }