From e78a5b0a050818454344020b1682c07c1051f932 Mon Sep 17 00:00:00 2001
From: Andrew Trokhymenko <atrokhym@gmail.com>
Date: Mon, 18 Nov 2024 19:55:28 -0500
Subject: [PATCH] image-upload

---
 app/components/chat/BaseChat.tsx          | 121 ++++++++++++++++------
 app/components/chat/Chat.client.tsx       |  42 +++++++-
 app/components/chat/FilePreview.tsx       |  40 +++++++
 app/components/chat/SendButton.client.tsx |   5 +-
 app/components/chat/UserMessage.tsx       |  27 ++++-
 app/components/sidebar/Menu.client.tsx    |   2 +-
 app/components/workbench/EditorPanel.tsx  |   9 +-
 app/components/workbench/FileTree.tsx     |   2 +-
 app/lib/.server/llm/stream-text.ts        |  41 ++++++--
 app/routes/api.chat.ts                    |  17 ++-
 10 files changed, 244 insertions(+), 62 deletions(-)
 create mode 100644 app/components/chat/FilePreview.tsx

diff --git a/app/components/chat/BaseChat.tsx b/app/components/chat/BaseChat.tsx
index 629c5cb..e0cd928 100644
--- a/app/components/chat/BaseChat.tsx
+++ b/app/components/chat/BaseChat.tsx
@@ -17,6 +17,8 @@ import Cookies from 'js-cookie';
 import styles from './BaseChat.module.scss';
 import type { ProviderInfo } from '~/utils/types';
 
+import FilePreview from './FilePreview';
+
 const EXAMPLE_PROMPTS = [
   { text: 'Build a todo app in React using Tailwind' },
   { text: 'Build a simple blog using Astro' },
@@ -33,7 +35,7 @@ const ModelSelector = ({ model, setModel, provider, setProvider, modelList, prov
       <select
         value={provider?.name}
         onChange={(e) => {
-          setProvider(providerList.find((p) => p.name === e.target.value));
+          setProvider(providerList.find(p => p.name === e.target.value));
           const firstModel = [...modelList].find((m) => m.provider == e.target.value);
           setModel(firstModel ? firstModel.name : '');
         }}
@@ -49,7 +51,7 @@ const ModelSelector = ({ model, setModel, provider, setProvider, modelList, prov
         key={provider?.name}
         value={model}
         onChange={(e) => setModel(e.target.value)}
-        style={{ maxWidth: '70%' }}
+        style={{ maxWidth: "70%" }}
         className="flex-1 p-2 rounded-lg border border-bolt-elements-borderColor bg-bolt-elements-prompt-background text-bolt-elements-textPrimary focus:outline-none focus:ring-2 focus:ring-bolt-elements-focus transition-all"
       >
         {[...modelList]
@@ -85,32 +87,38 @@ interface BaseChatProps {
   sendMessage?: (event: React.UIEvent, messageInput?: string) => void;
   handleInputChange?: (event: React.ChangeEvent<HTMLTextAreaElement>) => void;
   enhancePrompt?: () => void;
+  uploadedFiles?: File[];
+  setUploadedFiles?: (files: File[]) => void;
+  imageDataList?: string[];
+  setImageDataList?: (dataList: string[]) => void;
 }
-
 export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
-  (
-    {
-      textareaRef,
-      messageRef,
-      scrollRef,
-      showChat = true,
-      chatStarted = false,
-      isStreaming = false,
-      enhancingPrompt = false,
-      promptEnhanced = false,
-      messages,
-      input = '',
-      model,
-      setModel,
-      provider,
-      setProvider,
-      sendMessage,
-      handleInputChange,
-      enhancePrompt,
-      handleStop,
-    },
-    ref,
-  ) => {
+  ({
+    textareaRef,
+    messageRef,
+    scrollRef,
+    showChat,
+    chatStarted = false,
+    isStreaming = false,
+    model,
+    setModel,
+    provider,
+    setProvider,
+    input = '',
+    enhancingPrompt,
+    handleInputChange,
+    promptEnhanced,
+    enhancePrompt,
+    sendMessage,
+    handleStop,
+    uploadedFiles,
+    setUploadedFiles,
+    imageDataList,
+    setImageDataList,
+    messages,
+    children,  // Add this
+  }, ref) => {
+    console.log(provider);
     const TEXTAREA_MAX_HEIGHT = chatStarted ? 400 : 200;
     const [apiKeys, setApiKeys] = useState<Record<string, string>>({});
     const [modelList, setModelList] = useState(MODEL_LIST);
@@ -131,7 +139,7 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
         Cookies.remove('apiKeys');
       }
 
-      initializeModelList().then((modelList) => {
+      initializeModelList().then(modelList => {
         setModelList(modelList);
       });
     }, []);
@@ -152,6 +160,32 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
       }
     };
 
+    const handleRemoveFile = () => {
+      setUploadedFiles([]);
+      setImageDataList([]);
+    };
+
+    const handleFileUpload = () => {
+      const input = document.createElement('input');
+      input.type = 'file';
+      input.accept = 'image/*';
+
+      input.onchange = async (e) => {
+        const file = (e.target as HTMLInputElement).files?.[0];
+        if (file) {
+          const reader = new FileReader();
+          reader.onload = (e) => {
+            const base64Image = e.target?.result as string;
+            setUploadedFiles?.([...uploadedFiles, file]);
+            setImageDataList?.([...imageDataList, base64Image]);
+          };
+          reader.readAsDataURL(file);
+        }
+      };
+
+      input.click();
+    };
+
     return (
       <div
         ref={ref}
@@ -205,13 +239,22 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
                   setProvider={setProvider}
                   providerList={PROVIDER_LIST}
                 />
-                {provider && (
+                {provider &&
                   <APIKeyManager
                     provider={provider}
                     apiKey={apiKeys[provider.name] || ''}
                     setApiKey={(key) => updateApiKey(provider.name, key)}
-                  />
-                )}
+                  />}
+
+                <FilePreview
+                  files={uploadedFiles}
+                  imageDataList={imageDataList}
+                  onRemove={(index) => {
+                    setUploadedFiles?.(uploadedFiles.filter((_, i) => i !== index));
+                    setImageDataList?.(imageDataList.filter((_, i) => i !== index));
+                  }}
+                />
+
                 <div
                   className={classNames(
                     'shadow-lg border border-bolt-elements-borderColor bg-bolt-elements-prompt-background backdrop-filter backdrop-blur-[8px] rounded-lg overflow-hidden transition-all',
@@ -245,21 +288,30 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
                   <ClientOnly>
                     {() => (
                       <SendButton
-                        show={input.length > 0 || isStreaming}
+                        show={input.length > 0 || isStreaming || uploadedFiles.length > 0}
                         isStreaming={isStreaming}
                         onClick={(event) => {
                           if (isStreaming) {
                             handleStop?.();
                             return;
                           }
-
-                          sendMessage?.(event);
+                          if (input.length > 0 || uploadedFiles.length > 0) {
+                            sendMessage?.(event);
+                          }
                         }}
                       />
                     )}
                   </ClientOnly>
                   <div className="flex justify-between items-center text-sm p-4 pt-2">
                     <div className="flex gap-1 items-center">
+                      <IconButton
+                        title="Upload file"
+                        className="transition-all"
+                        onClick={() => handleFileUpload()}
+                      >
+                        <div className="i-ph:upload text-xl"></div>
+                      </IconButton>
+
                       <IconButton
                         title="Enhance prompt"
                         disabled={input.length === 0 || enhancingPrompt}
@@ -322,3 +374,6 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
     );
   },
 );
+
+
+
diff --git a/app/components/chat/Chat.client.tsx b/app/components/chat/Chat.client.tsx
index 31748b3..d62a5bd 100644
--- a/app/components/chat/Chat.client.tsx
+++ b/app/components/chat/Chat.client.tsx
@@ -73,8 +73,11 @@ export const ChatImpl = memo(({ initialMessages, storeMessageHistory }: ChatProp
   useShortcuts();
 
   const textareaRef = useRef<HTMLTextAreaElement>(null);
-
   const [chatStarted, setChatStarted] = useState(initialMessages.length > 0);
+  const [uploadedFiles, setUploadedFiles] = useState<File[]>([]);  // Move here
+  const [imageDataList, setImageDataList] = useState<string[]>([]); // Move here
+
+
   const [model, setModel] = useState(() => {
     const savedModel = Cookies.get('selectedModel');
     return savedModel || DEFAULT_MODEL;
@@ -196,7 +199,19 @@ export const ChatImpl = memo(({ initialMessages, storeMessageHistory }: ChatProp
        * manually reset the input and we'd have to manually pass in file attachments. However, those
        * aren't relevant here.
        */
-      append({ role: 'user', content: `[Model: ${model}]\n\n[Provider: ${provider.name}]\n\n${diff}\n\n${_input}` });
+      append({
+        role: 'user',
+        content: [
+          {
+            type: 'text',
+            text: `[Model: ${model}]\n\n[Provider: ${provider.name}]\n\n${diff}\n\n${_input}`
+          },
+          ...(imageDataList.map(imageData => ({
+            type: 'image',
+            image: imageData
+          })))
+        ]
+      });
 
       /**
        * After sending a new message we reset all modifications since the model
@@ -204,16 +219,30 @@ export const ChatImpl = memo(({ initialMessages, storeMessageHistory }: ChatProp
        */
       workbenchStore.resetAllFileModifications();
     } else {
-      append({ role: 'user', content: `[Model: ${model}]\n\n[Provider: ${provider.name}]\n\n${_input}` });
+      append({
+        role: 'user',
+        content: [
+          {
+            type: 'text',
+            text: `[Model: ${model}]\n\n[Provider: ${provider.name}]\n\n${_input}`
+          },
+          ...(imageDataList.map(imageData => ({
+            type: 'image',
+            image: imageData
+          })))
+        ]
+      });
     }
 
     setInput('');
+    // Add file cleanup here
+    setUploadedFiles([]);
+    setImageDataList([]);
 
     resetEnhancer();
 
     textareaRef.current?.blur();
   };
-
   const [messageRef, scrollRef] = useSnapScroll();
 
   useEffect(() => {
@@ -274,6 +303,11 @@ export const ChatImpl = memo(({ initialMessages, storeMessageHistory }: ChatProp
           apiKeys
         );
       }}
+      uploadedFiles={uploadedFiles}
+      setUploadedFiles={setUploadedFiles}
+      imageDataList={imageDataList}
+      setImageDataList={setImageDataList}
     />
   );
 });
+
diff --git a/app/components/chat/FilePreview.tsx b/app/components/chat/FilePreview.tsx
new file mode 100644
index 0000000..378ada6
--- /dev/null
+++ b/app/components/chat/FilePreview.tsx
@@ -0,0 +1,40 @@
+// FilePreview.tsx
+import React from 'react';
+import { X } from 'lucide-react';
+
+interface FilePreviewProps {
+    files: File[];
+    imageDataList: string[]; // or imagePreviews: string[]
+    onRemove: (index: number) => void;
+}
+
+const FilePreview: React.FC<FilePreviewProps> = ({ files, imageDataList, onRemove }) => {
+    if (!files || files.length === 0) {
+        return null; // Or render a placeholder if desired
+    }
+
+    return (
+        <div className="flex flex-row overflow-x-auto"> {/* Add horizontal scrolling if needed */}
+            {files.map((file, index) => (
+                <div key={file.name + file.size} className="mr-2 relative">
+                    {/* Display image preview or file icon */}
+                    {imageDataList[index] && (
+                        <div className="relative">
+                            <img src={imageDataList[index]} alt={file.name} className="max-h-20" />
+                            <button
+                                onClick={() => onRemove(index)}
+                                className="absolute -top-2 -right-2 z-10 bg-white rounded-full p-1 shadow-md hover:bg-gray-100"
+                            >
+                                <div className="bg-black rounded-full p-1">
+                                    <X className="w-3 h-3 text-gray-400" strokeWidth={2.5} />
+                                </div>
+                            </button>
+                        </div>
+                    )}
+                </div>
+            ))}
+        </div>
+    );
+};
+
+export default FilePreview;
diff --git a/app/components/chat/SendButton.client.tsx b/app/components/chat/SendButton.client.tsx
index 99770fc..59e5e52 100644
--- a/app/components/chat/SendButton.client.tsx
+++ b/app/components/chat/SendButton.client.tsx
@@ -4,11 +4,12 @@ interface SendButtonProps {
   show: boolean;
   isStreaming?: boolean;
   onClick?: (event: React.MouseEvent<HTMLButtonElement, MouseEvent>) => void;
+  onImagesSelected?: (images: File[]) => void;
 }
 
 const customEasingFn = cubicBezier(0.4, 0, 0.2, 1);
 
-export function SendButton({ show, isStreaming, onClick }: SendButtonProps) {
+export const SendButton = ({ show, isStreaming, onClick }: SendButtonProps) => {
   return (
     <AnimatePresence>
       {show ? (
@@ -30,4 +31,4 @@ export function SendButton({ show, isStreaming, onClick }: SendButtonProps) {
       ) : null}
     </AnimatePresence>
   );
-}
+};
diff --git a/app/components/chat/UserMessage.tsx b/app/components/chat/UserMessage.tsx
index 803d2cd..a57d4fa 100644
--- a/app/components/chat/UserMessage.tsx
+++ b/app/components/chat/UserMessage.tsx
@@ -9,13 +9,34 @@ interface UserMessageProps {
 }
 
 export function UserMessage({ content }: UserMessageProps) {
+  const sanitizedContent = sanitizeUserMessage(content);
+  const textContent = Array.isArray(sanitizedContent) 
+    ? sanitizedContent.find(item => item.type === 'text')?.text || ''
+    : sanitizedContent;
+
   return (
     <div className="overflow-hidden pt-[4px]">
-      <Markdown limitedMarkdown>{sanitizeUserMessage(content)}</Markdown>
+      <Markdown limitedMarkdown>{textContent}</Markdown>
     </div>
   );
 }
 
-function sanitizeUserMessage(content: string) {
-  return content.replace(modificationsRegex, '').replace(MODEL_REGEX, 'Using: $1').replace(PROVIDER_REGEX, ' ($1)\n\n').trim();
+// function sanitizeUserMessage(content: string) {
+//   return content.replace(modificationsRegex, '').replace(MODEL_REGEX, 'Using: $1').replace(PROVIDER_REGEX, ' ($1)\n\n').trim();
+// }
+function sanitizeUserMessage(content: string | Array<{type: string, text?: string, image_url?: {url: string}}>) {
+  if (Array.isArray(content)) {
+    return content.map(item => {
+      if (item.type === 'text') {
+        return {
+          type: 'text',
+          text: item.text?.replace(/\[Model:.*?\]\n\n/, '').replace(/\[Provider:.*?\]\n\n/, '')
+        };
+      }
+      return item; // Keep image_url items unchanged
+    });
+  }
+  
+  // Handle legacy string content
+  return content.replace(/\[Model:.*?\]\n\n/, '').replace(/\[Provider:.*?\]\n\n/, '');
 }
diff --git a/app/components/sidebar/Menu.client.tsx b/app/components/sidebar/Menu.client.tsx
index cf6d978..5b1bc6f 100644
--- a/app/components/sidebar/Menu.client.tsx
+++ b/app/components/sidebar/Menu.client.tsx
@@ -33,7 +33,7 @@ const menuVariants = {
 
 type DialogContent = { type: 'delete'; item: ChatHistoryItem } | null;
 
-export function Menu() {
+export const Menu = () => {
   const menuRef = useRef<HTMLDivElement>(null);
   const [list, setList] = useState<ChatHistoryItem[]>([]);
   const [open, setOpen] = useState(false);
diff --git a/app/components/workbench/EditorPanel.tsx b/app/components/workbench/EditorPanel.tsx
index 046eaf8..e789f1d 100644
--- a/app/components/workbench/EditorPanel.tsx
+++ b/app/components/workbench/EditorPanel.tsx
@@ -23,7 +23,6 @@ import { isMobile } from '~/utils/mobile';
 import { FileBreadcrumb } from './FileBreadcrumb';
 import { FileTree } from './FileTree';
 import { Terminal, type TerminalRef } from './terminal/Terminal';
-import React from 'react';
 
 interface EditorPanelProps {
   files?: FileMap;
@@ -204,7 +203,7 @@ export const EditorPanel = memo(
                   const isActive = activeTerminal === index;
 
                   return (
-                    <React.Fragment key={index}>
+                    <>
                       {index == 0 ? (
                         <button
                           key={index}
@@ -223,7 +222,7 @@ export const EditorPanel = memo(
                           Bolt Terminal
                         </button>
                       ) : (
-                        <React.Fragment>
+                        <>
                           <button
                             key={index}
                             className={classNames(
@@ -239,9 +238,9 @@ export const EditorPanel = memo(
                             <div className="i-ph:terminal-window-duotone text-lg" />
                             Terminal {terminalCount > 1 && index}
                           </button>
-                        </React.Fragment>
+                        </>
                       )}
-                    </React.Fragment>
+                    </>
                   );
                 })}
                 {terminalCount < MAX_TERMINALS && <IconButton icon="i-ph:plus" size="md" onClick={addTerminal} />}
diff --git a/app/components/workbench/FileTree.tsx b/app/components/workbench/FileTree.tsx
index bfd8302..71f40af 100644
--- a/app/components/workbench/FileTree.tsx
+++ b/app/components/workbench/FileTree.tsx
@@ -111,7 +111,7 @@ export const FileTree = memo(
     };
 
     return (
-      <div className={classNames('text-sm', className ,'overflow-y-auto')}>
+      <div className={classNames('text-sm', className)}>
         {filteredFileList.map((fileOrFolder) => {
           switch (fileOrFolder.kind) {
             case 'file': {
diff --git a/app/lib/.server/llm/stream-text.ts b/app/lib/.server/llm/stream-text.ts
index 8bfd797..a50c28e 100644
--- a/app/lib/.server/llm/stream-text.ts
+++ b/app/lib/.server/llm/stream-text.ts
@@ -25,26 +25,39 @@ export type Messages = Message[];
 export type StreamingOptions = Omit<Parameters<typeof _streamText>[0], 'model'>;
 
 function extractPropertiesFromMessage(message: Message): { model: string; provider: string; content: string } {
+  const textContent = Array.isArray(message.content)
+    ? message.content.find(item => item.type === 'text')?.text || ''
+    : message.content;
+
+  const modelMatch = textContent.match(MODEL_REGEX);
+  const providerMatch = textContent.match(PROVIDER_REGEX);
+
   // Extract model
-  const modelMatch = message.content.match(MODEL_REGEX);
+  // const modelMatch = message.content.match(MODEL_REGEX);
   const model = modelMatch ? modelMatch[1] : DEFAULT_MODEL;
 
   // Extract provider
-  const providerMatch = message.content.match(PROVIDER_REGEX);
+  // const providerMatch = message.content.match(PROVIDER_REGEX);
   const provider = providerMatch ? providerMatch[1] : DEFAULT_PROVIDER;
 
-  // Remove model and provider lines from content
-  const cleanedContent = message.content
-    .replace(MODEL_REGEX, '')
-    .replace(PROVIDER_REGEX, '')
-    .trim();
+  const cleanedContent = Array.isArray(message.content)
+    ? message.content.map(item => {
+      if (item.type === 'text') {
+        return {
+          type: 'text',
+          text: item.text?.replace(/\[Model:.*?\]\n\n/, '').replace(/\[Provider:.*?\]\n\n/, '')
+        };
+      }
+      return item; // Preserve image_url and other types as is
+    })
+    : textContent.replace(/\[Model:.*?\]\n\n/, '').replace(/\[Provider:.*?\]\n\n/, '');
 
   return { model, provider, content: cleanedContent };
 }
 
 export function streamText(
-  messages: Messages, 
-  env: Env, 
+  messages: Messages,
+  env: Env,
   options?: StreamingOptions,
   apiKeys?: Record<string, string>
 ) {
@@ -67,6 +80,16 @@ export function streamText(
     return message; // No changes for non-user messages
   });
 
+  // const modelConfig = getModel(currentProvider, currentModel, env, apiKeys);
+  // const coreMessages = convertToCoreMessages(processedMessages);
+  
+  // console.log('Debug streamText:', JSON.stringify({
+  //   model: modelConfig,
+  //   messages: processedMessages,
+  //   coreMessages: coreMessages,
+  //   system: getSystemPrompt()
+  // }, null, 2));
+
   return _streamText({
     model: getModel(currentProvider, currentModel, env, apiKeys),
     system: getSystemPrompt(),
diff --git a/app/routes/api.chat.ts b/app/routes/api.chat.ts
index 47666c7..284fccb 100644
--- a/app/routes/api.chat.ts
+++ b/app/routes/api.chat.ts
@@ -30,13 +30,15 @@ function parseCookies(cookieHeader) {
 }
 
 async function chatAction({ context, request }: ActionFunctionArgs) {
-  const { messages } = await request.json<{
-    messages: Messages
+  // console.log('=== API CHAT LOGGING START ===');
+  // console.log('Request received:', request.url);
+  
+  const { messages, imageData } = await request.json<{
+    messages: Messages,
+    imageData?: string[]
   }>();
 
   const cookieHeader = request.headers.get("Cookie");
-
-  // Parse the cookie's value (returns an object or null if no cookie exists)
   const apiKeys = JSON.parse(parseCookies(cookieHeader).apiKeys || "{}");
 
   const stream = new SwitchableStream();
@@ -69,6 +71,13 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
 
     const result = await streamText(messages, context.cloudflare.env, options, apiKeys);
 
+    // console.log('=== API CHAT LOGGING START ===');
+    // console.log('StreamText:', JSON.stringify({
+    //   messages,
+    //   result,
+    // }, null, 2));
+    // console.log('=== API CHAT LOGGING END ===');
+
     stream.switchSource(result.toAIStream());
 
     return new Response(stream.readable, {