Merge branch 'main' into copyMyFix

app/components/chat/AssistantMessage.tsx

@@ -1,13 +1,30 @@
 import { memo } from 'react';
 import { Markdown } from './Markdown';
+import type { JSONValue } from 'ai';
 
 interface AssistantMessageProps {
   content: string;
+  annotations?: JSONValue[];
 }
 
-export const AssistantMessage = memo(({ content }: AssistantMessageProps) => {
+export const AssistantMessage = memo(({ content, annotations }: AssistantMessageProps) => {
+  const filteredAnnotations = (annotations?.filter(
+    (annotation: JSONValue) => annotation && typeof annotation === 'object' && Object.keys(annotation).includes('type'),
+  ) || []) as { type: string; value: any }[];
+
+  const usage: {
+    completionTokens: number;
+    promptTokens: number;
+    totalTokens: number;
+  } = filteredAnnotations.find((annotation) => annotation.type === 'usage')?.value;
+
   return (
     <div className="overflow-hidden w-full">
+      {usage && (
+        <div className="text-sm text-bolt-elements-textSecondary mb-2">
+          Tokens: {usage.totalTokens} (prompt: {usage.promptTokens}, completion: {usage.completionTokens})
+        </div>
+      )}
       <Markdown html>{content}</Markdown>
     </div>
   );

app/components/chat/BaseChat.tsx

@@ -77,7 +77,8 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
       input = '',
       enhancingPrompt,
       handleInputChange,
-      promptEnhanced,
+
+      // promptEnhanced,
       enhancePrompt,
       sendMessage,
       handleStop,
@@ -490,10 +491,7 @@ export const BaseChat = React.forwardRef<HTMLDivElement, BaseChatProps>(
                       <IconButton
                         title="Enhance prompt"
                         disabled={input.length === 0 || enhancingPrompt}
-                        className={classNames(
-                          'transition-all',
-                          enhancingPrompt ? 'opacity-100' : '',
-                        )}
+                        className={classNames('transition-all', enhancingPrompt ? 'opacity-100' : '')}
                         onClick={() => {
                           enhancePrompt?.();
                           toast.success('Prompt enhanced!');

app/components/chat/Chat.client.tsx

@@ -93,7 +93,7 @@ export const ChatImpl = memo(
     const [uploadedFiles, setUploadedFiles] = useState<File[]>([]); // Move here
     const [imageDataList, setImageDataList] = useState<string[]>([]); // Move here
     const files = useStore(workbenchStore.files);
-    const { activeProviders } = useSettings();
+    const { activeProviders, promptId } = useSettings();
 
     const [model, setModel] = useState(() => {
       const savedModel = Cookies.get('selectedModel');
@@ -115,14 +115,24 @@ export const ChatImpl = memo(
       body: {
         apiKeys,
         files,
+        promptId,
       },
+      sendExtraMessageFields: true,
       onError: (error) => {
         logger.error('Request failed\n\n', error);
         toast.error(
           'There was an error processing your request: ' + (error.message ? error.message : 'No details were returned'),
         );
       },
-      onFinish: () => {
+      onFinish: (message, response) => {
+        const usage = response.usage;
+
+        if (usage) {
+          console.log('Token usage:', usage);
+
+          // You can now use the usage data as needed
+        }
+
         logger.debug('Finished streaming');
       },
       initialMessages,

app/components/chat/Messages.client.tsx

@@ -65,12 +65,16 @@ export const Messages = React.forwardRef<HTMLDivElement, MessagesProps>((props:
                   </div>
                 )}
                 <div className="grid grid-col-1 w-full">
-                  {isUserMessage ? <UserMessage content={content} /> : <AssistantMessage content={content} />}
+                  {isUserMessage ? (
+                    <UserMessage content={content} />
+                  ) : (
+                    <AssistantMessage content={content} annotations={message.annotations} />
+                  )}
                 </div>
                 {!isUserMessage && (
                   <div className="flex gap-2 flex-col lg:flex-row">
-                    <WithTooltip tooltip="Revert to this message">
-                      {messageId && (
+                    {messageId && (
+                      <WithTooltip tooltip="Revert to this message">
                         <button
                           onClick={() => handleRewind(messageId)}
                           key="i-ph:arrow-u-up-left"
@@ -79,8 +83,8 @@ export const Messages = React.forwardRef<HTMLDivElement, MessagesProps>((props:
                             'text-xl text-bolt-elements-textSecondary hover:text-bolt-elements-textPrimary transition-colors',
                           )}
                         />
-                      )}
-                    </WithTooltip>
+                      </WithTooltip>
+                    )}
 
                     <WithTooltip tooltip="Fork chat from this message">
                       <button

app/components/chat/UserMessage.tsx

@@ -12,42 +12,36 @@ interface UserMessageProps {
 export function UserMessage({ content }: UserMessageProps) {
   if (Array.isArray(content)) {
     const textItem = content.find((item) => item.type === 'text');
-    const textContent = sanitizeUserMessage(textItem?.text || '');
+    const textContent = stripMetadata(textItem?.text || '');
     const images = content.filter((item) => item.type === 'image' && item.image);
 
     return (
       <div className="overflow-hidden pt-[4px]">
-        <div className="flex items-start gap-4">
-          <div className="flex-1">
-            <Markdown limitedMarkdown>{textContent}</Markdown>
-          </div>
-          {images.length > 0 && (
-            <div className="flex-shrink-0 w-[160px]">
-              {images.map((item, index) => (
-                <div key={index} className="relative">
-                  <img
-                    src={item.image}
-                    alt={`Uploaded image ${index + 1}`}
-                    className="w-full h-[160px] rounded-lg object-cover border border-bolt-elements-borderColor"
-                  />
-                </div>
-              ))}
-            </div>
-          )}
+        <div className="flex flex-col gap-4">
+          {textContent && <Markdown html>{textContent}</Markdown>}
+          {images.map((item, index) => (
+            <img
+              key={index}
+              src={item.image}
+              alt={`Image ${index + 1}`}
+              className="max-w-full h-auto rounded-lg"
+              style={{ maxHeight: '512px', objectFit: 'contain' }}
+            />
+          ))}
         </div>
       </div>
     );
   }
 
-  const textContent = sanitizeUserMessage(content);
+  const textContent = stripMetadata(content);
 
   return (
     <div className="overflow-hidden pt-[4px]">
-      <Markdown limitedMarkdown>{textContent}</Markdown>
+      <Markdown html>{textContent}</Markdown>
     </div>
   );
 }
 
-function sanitizeUserMessage(content: string) {
+function stripMetadata(content: string) {
   return content.replace(MODEL_REGEX, '').replace(PROVIDER_REGEX, '');
 }