lms chat Image Support Part I - Primitives and Rendering

src/subcommands/chat/react/Chat.tsx

@@ -2,7 +2,7 @@ import { produce } from "@lmstudio/immer-with-plugins";
 import { type Chat, type LLM, type LLMPredictionStats, type LMStudioClient } from "@lmstudio/sdk";
 import { Box, type DOMElement, useApp, Text } from "ink";
 import React, { useCallback, useMemo, useRef, useState } from "react";
-import { displayVerboseStats, getLargePastePlaceholderText } from "../util.js";
+import { displayVerboseStats, getChipPreviewText } from "../util.js";
 import { ChatInput } from "./ChatInput.js";
 import { ChatMessagesList } from "./ChatMessagesList.js";
 import { ChatSuggestions } from "./ChatSuggestions.js";
@@ -19,7 +19,15 @@ import {
 } from "./hooks.js";
 import { insertPasteAtCursor } from "./inputReducer.js";
 import { createSlashCommands } from "./slashCommands.js";
-import type { ChatUserInputState, InkChatMessage, Suggestion } from "./types.js";
+import type {
+  ChatUserInputState,
+  InkChatMessage,
+  Suggestion,
+  UserInputContentPart,
+} from "./types.js";
+import { buildUserMessageParts } from "./buildChat.js";
+import { ImagePreparationError } from "./images/imageErrors.js";
+import { ImageStore } from "./images/imageStore.js";
 
 // Freezes streaming content into static chunks at natural breaks to reduce re-renders.
 // Uses multiple boundaries to handle different content (best effort):
@@ -48,6 +56,8 @@ export const emptyChatInputState: ChatUserInputState = {
   cursorInSegmentOffset: 0,
 };
 
+const imageStore = new ImageStore();
+
 export const ChatComponent = React.memo(
   ({ client, llm, chat, onExit, stats, ttl, shouldFetchModelCatalog }: ChatComponentProps) => {
     const { exit } = useApp();
@@ -181,7 +191,7 @@ export const ChatComponent = React.memo(
         return [];
       }
       const firstSegment = userInputState.segments[0];
-      const inputText = firstSegment.content;
+      const inputText = firstSegment?.type === "text" ? firstSegment.content : "";
       return commandHandler.getSuggestions({
         input: inputText,
         shouldShowSuggestions: !isConfirmationActive && !isPredicting && inputText.startsWith("/"),
@@ -292,11 +302,7 @@ export const ChatComponent = React.memo(
 
     const handlePaste = useCallback((content: string) => {
       const normalizedContent = content.replaceAll("\r\n", "\n").replaceAll("\r", "\n");
-
-      if (normalizedContent.length === 0) {
-        return;
-      }
-
+      if (normalizedContent.trim().length === 0) return;
       setUserInputState(previousState =>
         insertPasteAtCursor({
           state: previousState,
@@ -307,27 +313,44 @@ export const ChatComponent = React.memo(
     }, []);
 
     const handleSubmit = useCallback(async () => {
+      const inputStateSnapshot = userInputState;
+      const inputSegments = userInputState.segments;
+      const hasImageChip = inputSegments.some(
+        segment => segment.type === "chip" && segment.data.kind === "image",
+      );
+
       // Collect the full text input from the userInputState
-      const userInputText = userInputState.segments
-        .map(segment => segment.content)
+      const userInputText = inputSegments
+        .map(segment => {
+          if (segment.type === "text") {
+            return segment.content;
+          }
+          return segment.data.kind === "largePaste" ? segment.data.content : "";
+        })
         .join("")
         .trim();
-      // Clear the input state
-      setUserInputState(emptyChatInputState);
-      const confirmationResponse = await handleConfirmationResponse(userInputText);
-      if (confirmationResponse === "handled") {
-        return;
-      }
-      if (confirmationResponse === "invalid") {
-        logInChat("Please answer 'yes' or 'no'");
-        return;
+
+      if (isConfirmationActive) {
+        const confirmationResponse = await handleConfirmationResponse(userInputText);
+        if (confirmationResponse === "handled") {
+          setUserInputState(emptyChatInputState);
+          return;
+        }
+        if (confirmationResponse === "invalid") {
+          logInChat("Please answer 'yes' or 'no'");
+          return;
+        }
       }
 
-      if (userInputText.length === 0) {
+      if (userInputText.length === 0 && hasImageChip === false) {
         return;
       }
 
-      if (userInputText.startsWith("/") && userInputState.segments.length === 1) {
+      if (
+        userInputText.startsWith("/") &&
+        inputSegments.length === 1 &&
+        inputSegments[0]?.type === "text"
+      ) {
         const selectedSuggestion =
           normalizedSelectedSuggestionIndex !== null
             ? suggestions[normalizedSelectedSuggestionIndex]
@@ -344,6 +367,7 @@ export const ChatComponent = React.memo(
         // Check if the command is in exception list, if not,
         // execute it.
         if (commandHandler.commandIsIgnored(command) === false) {
+          setUserInputState(emptyChatInputState);
           const wasCommandHandled = await commandHandler.execute(command, argumentsText);
           if (wasCommandHandled === false) {
             logInChat(`Unknown command: ${userInputText}`);
@@ -362,14 +386,35 @@ export const ChatComponent = React.memo(
         return;
       }
 
+      if (hasImageChip === true && llmRef.current.vision === false) {
+        logErrorInChat("The current model does not support image input (vision).");
+        return;
+      }
+
       if (isPredicting) {
         logInChat(
           "A prediction is already in progress. Please wait for it to finish or press CTRL+C to abort it.",
         );
         return;
       }
 
+      const userMessageContent = inputSegments.map<UserInputContentPart>(segment => {
+        if (segment.type === "text") {
+          return { type: "text", text: segment.content };
+        }
+
+        const displayText = getChipPreviewText(segment.data);
+        return { type: "chip", kind: segment.data.kind, displayText };
+      });
+
       // If nothing else, proceed with normal message submission
+      setUserInputState(emptyChatInputState);
+      // Render the user message immediately (before any async image preparation) to avoid a
+      // transient frame where the input clears but the message hasn't appeared yet ("flash").
+      addMessage({
+        type: "user",
+        content: userMessageContent,
+      });
       setIsPredicting(true);
       setPromptProcessingProgress(null);
       setShowPredictionSpinner(true);
@@ -382,20 +427,30 @@ export const ChatComponent = React.memo(
       const signal = abortControllerRef.current.signal;
 
       try {
-        chatRef.current.append("user", userInputText);
-        addMessage({
-          type: "user",
-          content: userInputState.segments.map(segment => {
-            if (segment.type === "largePaste") {
-              const placeholder = getLargePastePlaceholderText(segment.content);
-              return {
-                type: "largePaste",
-                text: placeholder,
-              };
-            }
-            return { type: segment.type, text: segment.content };
-          }),
-        });
+        let parts: Awaited<ReturnType<typeof buildUserMessageParts>>;
+        try {
+          parts = await buildUserMessageParts({
+            client,
+            inputSegments,
+            imageStore,
+          });
+        } catch (error) {
+          if (error instanceof ImagePreparationError) {
+            setMessages(
+              produce(draftMessages => {
+                const lastMessage = draftMessages.at(-1);
+                if (lastMessage?.type === "user") {
+                  draftMessages.pop();
+                }
+              }),
+            );
+            setUserInputState(inputStateSnapshot);
+            logErrorInChat(`Failed to attach image: ${error.message}`);
+            return;
+          }
+          throw error;
+        }
+        chatRef.current.append({ role: "user", content: parts });
         const result = await llmRef.current.respond(chatRef.current, {
           onFirstToken() {
             setShowPredictionSpinner(false);
@@ -577,20 +632,21 @@ export const ChatComponent = React.memo(
         abortControllerRef.current = null;
       }
     }, [
-      userInputState.segments,
-      handleConfirmationResponse,
+      userInputState,
+      isConfirmationActive,
       isPredicting,
+      addMessage,
+      handleConfirmationResponse,
       logInChat,
       normalizedSelectedSuggestionIndex,
       suggestions,
       commandHandler,
       handleExit,
       logErrorInChat,
-      addMessage,
+      client,
       stats,
       promptProcessingProgress,
       requestConfirmation,
-      client.llm,
       ttl,
     ]);
 

src/subcommands/chat/react/ChatInput.tsx

@@ -11,7 +11,9 @@ import {
 } from "./inputReducer.js";
 import { renderInputWithCursor } from "./inputRenderer.js";
 import { type ChatUserInputState } from "./types.js";
-import { getLargePastePlaceholderText } from "../util.js";
+import { getChipPreviewText } from "../util.js";
+
+type ChipRange = { start: number; end: number; kind: "largePaste" | "image" };
 
 interface ChatInputProps {
   inputState: ChatUserInputState;
@@ -150,10 +152,18 @@ export const ChatInput = ({
         return;
       }
 
-      const currentText = inputState.segments.map(segment => segment.content).join("");
+      const currentText = inputState.segments
+        .map(segment => (segment.type === "text" ? segment.content : ""))
+        .join("");
 
       // Check if input is a slash command without arguments that has suggestions
-      if (currentText.startsWith("/") && currentText.includes(" ") === false) {
+      // Only auto-insert space if input is a single text segment (no chips)
+      if (
+        currentText.startsWith("/") &&
+        currentText.includes(" ") === false &&
+        inputState.segments.length === 1 &&
+        inputState.segments[0]?.type === "text"
+      ) {
         const commandName = currentText.slice(1);
         if (commandHasSuggestions(commandName)) {
           setUserInputState(previousState =>
@@ -197,16 +207,16 @@ export const ChatInput = ({
     }
   });
 
-  const { fullText, cursorPosition, pasteRanges } = useMemo(() => {
+  const { fullText, cursorPosition, chipRanges } = useMemo(() => {
     let fullText = "";
     let cursorPosition = 0;
-    const pasteRanges: Array<{ start: number; end: number }> = [];
+    const chipRanges: ChipRange[] = [];
 
     for (let index = 0; index < inputState.segments.length; index++) {
       const segment = inputState.segments[index];
 
-      if (segment.type === "largePaste") {
-        const placeholder = getLargePastePlaceholderText(segment.content);
+      if (segment.type === "chip") {
+        const placeholder = getChipPreviewText(segment.data);
         const startPos = fullText.length;
 
         if (index < inputState.cursorOnSegmentIndex) {
@@ -216,7 +226,7 @@ export const ChatInput = ({
         }
 
         fullText += placeholder;
-        pasteRanges.push({ start: startPos, end: fullText.length });
+        chipRanges.push({ start: startPos, end: fullText.length, kind: segment.data.kind });
       } else {
         if (index < inputState.cursorOnSegmentIndex) {
           cursorPosition += segment.content.length;
@@ -228,7 +238,7 @@ export const ChatInput = ({
       }
     }
 
-    return { fullText, cursorPosition, pasteRanges };
+    return { fullText, cursorPosition, chipRanges };
   }, [inputState]);
 
   const lines = fullText.split("\n");
@@ -266,7 +276,7 @@ export const ChatInput = ({
                   {renderInputWithCursor({
                     fullText: lineText,
                     cursorPosition: isCursorLine ? cursorColumnIndex : -1,
-                    pasteRanges,
+                    chipRanges,
                     lineStartPos,
                   })}
                 </Text>

src/subcommands/chat/react/ChatMessages.tsx

@@ -13,7 +13,12 @@ export const ChatMessage = memo(({ message, modelName }: ChatMessageProps) => {
   const type = message.type;
   switch (type) {
     case "user": {
-      const contentParts = message.content.map(part => ({ ...part }));
+      const contentParts = message.content.map(part => {
+        if (part.type === "text") {
+          return { type: "text", text: part.text };
+        }
+        return { type: "chip", kind: part.kind, text: part.displayText };
+      });
 
       // Trim only edge newlines on the first/last non-empty parts to keep internal newlines intact.
       let startIndex = 0;
@@ -38,8 +43,11 @@ export const ChatMessage = memo(({ message, modelName }: ChatMessageProps) => {
       // After edge trimming, apply large paste coloring and join into a single string.
       const formattedContent = contentParts
         .map(contentPart => {
-          const text = contentPart.text;
-          return contentPart.type === "largePaste" ? chalk.blue(text) : text;
+          if (contentPart.type === "text") {
+            return contentPart.text;
+          }
+
+          return chalk.blue(contentPart.text);
         })
         .join("");