WIP AI stuff

Author: joepio

browser/data-browser/src/components/AI/AIChatMessage.tsx

@@ -1,7 +1,7 @@
 import styled from 'styled-components';
 import Markdown from '../datatypes/Markdown';
 import type { CoreAssistantMessage, CoreToolMessage } from 'ai';
-import { FaCircleExclamation, FaTrash } from 'react-icons/fa6';
+import { FaCircleExclamation, FaRetweet, FaTrash } from 'react-icons/fa6';
 import {
   isAIErrorMessage,
   isMessageWithContext,
@@ -14,6 +14,7 @@ import { IconButton } from '../IconButton/IconButton';
 interface MessageProps {
   message: AIChatDisplayMessage;
   onDeleteMessage?: (message: AIChatDisplayMessage) => void;
+  onRegenerateMessage?: (message: AIChatDisplayMessage) => void;
 }
 
 function isToolMessage(
@@ -31,14 +32,19 @@ function isAssistantMessage(
 export const AIChatMessage = ({
   message: messageIn,
   onDeleteMessage,
+  onRegenerateMessage,
 }: MessageProps) => {
   const [message, context] = isMessageWithContext(messageIn)
     ? [messageIn.message, messageIn.context]
     : [messageIn];
 
   if (message.role === 'user') {
     return (
-      <MessageActionWrapper message={message} onDeleteMessage={onDeleteMessage}>
+      <MessageActionWrapper
+        message={message}
+        onDeleteMessage={onDeleteMessage}
+        onRegenerateMessage={onRegenerateMessage}
+      >
         <UserMessage message={message} context={context} />
       </MessageActionWrapper>
     );
@@ -116,20 +122,30 @@ const MessageActionWrapper: React.FC<React.PropsWithChildren<MessageProps>> = ({
   children,
   message,
   onDeleteMessage,
+  onRegenerateMessage,
 }) => {
   return (
     <MessageTopWrapper>
-      {onDeleteMessage && (
-        <FloatingActionRow>
+      <FloatingActionRow>
+        {onDeleteMessage && (
           <IconButton
             color='textLight'
             onClick={() => onDeleteMessage(message)}
             title='Delete Message'
           >
             <FaTrash />
           </IconButton>
-        </FloatingActionRow>
-      )}
+        )}
+        {onRegenerateMessage && (
+          <IconButton
+            color='textLight'
+            onClick={() => onRegenerateMessage(message)}
+            title='Regenerate response'
+          >
+            <FaRetweet />
+          </IconButton>
+        )}
+      </FloatingActionRow>
       {children}
     </MessageTopWrapper>
   );

browser/data-browser/src/components/AI/AISettings.tsx

@@ -11,6 +11,13 @@ import { OpenRouterLoginButton } from './OpenRouterLoginButton';
 import { TabPanel, Tabs } from '../Tabs';
 import { effectFetch } from '../../helpers/effectFetch';
 
+const intl = new Intl.NumberFormat('default', {
+  style: 'currency',
+  currency: 'USD',
+  minimumFractionDigits: 2,
+  maximumFractionDigits: 2,
+});
+
 interface CreditUsage {
   total: number;
   used: number;
@@ -99,8 +106,8 @@ const AISettings: React.FC = () => {
                 </Row>
                 {creditUsage && (
                   <Subtle as='p'>
-                    Credits used: {creditUsage.used} / Total:{' '}
-                    {creditUsage.total}
+                    Credits used: {intl.format(creditUsage.used)} /{' '}
+                    {intl.format(creditUsage.total)}
                   </Subtle>
                 )}
                 {!openRouterApiKey && (

browser/data-browser/src/components/AI/AgentConfig.tsx

@@ -368,6 +368,7 @@ const AgentForm = ({ agent, onChange }: AgentFormProps) => {
         <Label htmlFor='agent-name'>Name</Label>
         <Input
           id='agent-name'
+          required
           max={50}
           value={agent.name}
           onChange={e => handleChange('name', e.target.value)}

browser/data-browser/src/components/AI/SimpleAIChat.tsx

@@ -47,6 +47,7 @@ import { NoKeyOverlay } from './NoKeyOverlay';
 import { useAutoAgentSelect } from './useAgentAutoSelect';
 import { useOpenRouterModels } from './useOpenRouterModels';
 import type { MentionItem } from '../../chunks/MarkdownEditor/AIChatInput/types';
+import { flushSync } from 'react-dom';
 
 const AIChatInput = React.lazy(
   () => import('../../chunks/MarkdownEditor/AIChatInput/AsyncAIChatInput'),
@@ -82,6 +83,7 @@ interface SimpleAIChatProps {
     React.SetStateAction<AIMessageContext[]>
   >;
   onDeleteMessage: (message: AIChatDisplayMessage) => void;
+  onRegenerateMessage: (message: AIChatDisplayMessage) => Promise<void>;
 }
 
 export const SimpleAIChat: React.FC<
@@ -94,6 +96,7 @@ export const SimpleAIChat: React.FC<
   setExternalContextItems,
   onNewMessage,
   onDeleteMessage,
+  onRegenerateMessage,
   children,
 }) => {
   const abortSignalRef = useRef<AbortController>(null);
@@ -140,11 +143,12 @@ export const SimpleAIChat: React.FC<
   const [tokensUsed, setTokensUsed] = useState<[number, number]>([0, 0]);
 
   const getToolsForAgent = useTools();
-  const [hasToolResultFollowUp, setHasToolResultFollowUp] = useState(false);
 
   const [agentConfigOpen, setAgentConfigOpen] = useState(false);
   const pickAgent = useAutoAgentSelect();
 
+  const [shouldRegenMessages, setShouldRegenMessages] = useState(false);
+
   const normalizeAndApplyContext = useProcessMessages();
 
   const { tools: atomicTools } = useAtomicMCPTools({
@@ -225,7 +229,12 @@ export const SimpleAIChat: React.FC<
     setUserSelectedContextItems(newContextItems);
   };
 
-  const sendMessage = async (isFollowUp = false) => {
+  const sendMessage = async ({
+    regenerate = false,
+  }: {
+    /** Whether to regenerate the response */
+    regenerate?: boolean;
+  } = {}) => {
     if (readonly) {
       toast.error('You do not have the permissions to edit this chat.');
 
@@ -281,21 +290,28 @@ export const SimpleAIChat: React.FC<
       ...userSelectedContextItems,
     ];
 
-    if (!isFollowUp) {
+    if (regenerate) {
+      messagesToUse = messages;
+    } else {
       const userMessage = prepareUserMessage(
         userInput,
         attachedFile,
         allContextItems,
       );
       messagesToUse = [...messages, userMessage];
       onNewMessage(userMessage);
-    } else {
-      messagesToUse = messages;
     }
 
     // Filter message to only include non-error messages, error messages are only intended for the user.
     const filteredMessages = await normalizeAndApplyContext(messagesToUse);
 
+    console.log(
+      'last message sent as context: "',
+      filteredMessages.at(-1).content[0].text,
+      'other mgss',
+      filteredMessages,
+    );
+
     // Update messages with the user message first
     setExternalContextItems([]);
     setUserSelectedContextItems([]);
@@ -545,13 +561,6 @@ export const SimpleAIChat: React.FC<
     }
   };
 
-  useEffect(() => {
-    if (hasToolResultFollowUp) {
-      setHasToolResultFollowUp(false);
-      sendMessage(true);
-    }
-  }, [hasToolResultFollowUp]);
-
   // Combine both context item lists when needed
   const allContextItems = [
     ...externalContextItems,
@@ -564,18 +573,37 @@ export const SimpleAIChat: React.FC<
     });
   };
 
+  const regenerateMessage = async (message: AIChatDisplayMessage) => {
+    flushSync(async () => {
+      // Removes the following messages
+      await onRegenerateMessage(message);
+
+      setShouldRegenMessages(true);
+    });
+  };
+
+  useEffect(() => {
+    if (shouldRegenMessages) {
+      sendMessage({
+        regenerate: true,
+      });
+      setShouldRegenMessages(false);
+    }
+  }, [shouldRegenMessages]);
+
   return (
     <ChatWindow fullView={fullView}>
       {children}
       <ChatMessagesContainer
         enableAutoScroll={aiState !== AIState.Stopped}
         fullView={fullView}
       >
-        {messages.filter(cleanMessages).map(message => (
+        {messages.filter(cleanMessages).map((message, index) => (
           <AIChatMessage
-            key={JSON.stringify(message)}
+            key={`${JSON.stringify(message)}-${index}`}
             message={message}
             onDeleteMessage={onDeleteMessage}
+            onRegenerateMessage={regenerateMessage}
           />
         ))}
         {ongoingMessage.text && (

browser/data-browser/src/components/AI/chatConversionUtils.ts

@@ -17,6 +17,7 @@ import { newContextItem } from './AISidebarContext';
 import {
   type AIAtomicResourceMessageContext,
   type AIChatDisplayMessage,
+  AIChatErrorMessage,
   type AIMCPResourceMessageContext,
   type AIMessageContext,
   isAtomicResource,
@@ -157,7 +158,15 @@ export const messageResourcesToDisplayMessages = async (
 
   for (const resource of resources) {
     if (resource.error) {
-      throw resource.error;
+      console.error(resource.error);
+      messages.set(
+        {
+          role: 'error',
+          content: resource.error.message,
+        } satisfies AIChatErrorMessage,
+        resource,
+      );
+      continue;
     }
 
     const role = tagToRole(resource.props.role);

browser/data-browser/src/components/ErrorLook.tsx

@@ -4,6 +4,7 @@ import { styled, css } from 'styled-components';
 import { FaExclamationTriangle } from 'react-icons/fa';
 
 import type { JSX } from 'react';
+import { getMessageForErrorType } from '@tomic/react';
 
 export const errorLookStyle = css`
   color: ${props => props.theme.colors.alert};
@@ -25,7 +26,7 @@ export function ErrorBlock({ error, showTrace }: ErrorBlockProps): JSX.Element {
     <ErrorLookBig>
       <BiggerText>
         <FaExclamationTriangle />
-        Something went wrong
+        {getMessageForErrorType(error)}
       </BiggerText>
       <Pre>
         <code>{error.message}</code>

browser/data-browser/src/helpers/formatCurrency.ts

@@ -1,6 +1,7 @@
 import { useEffect, useState } from 'react';
 import { SimpleAIChat } from '../../components/AI/SimpleAIChat';
 import {
+  Ai,
   ai,
   useArray,
   useCanWrite,
@@ -23,8 +24,11 @@ import {
   messageResourcesToDisplayMessages,
 } from '../../components/AI/chatConversionUtils';
 import { TagBar } from '../../components/Tag/TagBar';
+import { flushSync } from 'react-dom';
 
-export const AIChatPage: React.FC<ResourcePageProps> = ({ resource }) => {
+export const AIChatPage: React.FC<ResourcePageProps<Ai.AiChat>> = ({
+  resource,
+}) => {
   const store = useStore();
   const [messages, setMessages] = useState<AIChatDisplayMessage[]>([]);
   const [contextItems, setContextItems] = useState<AIMessageContext[]>([]);
@@ -87,6 +91,54 @@ export const AIChatPage: React.FC<ResourcePageProps> = ({ resource }) => {
     });
   };
 
+  const removeFollowingMessages = async (message: AIChatDisplayMessage) => {
+    const nextMessages = messages.slice(messages.indexOf(message) + 1);
+
+    // We need to destroy the resources server side as well as in the internal state.
+    // We also need to update the `messages` prop in the chat resource.
+    const destroySubjects: string[] = [];
+
+    for (const m of nextMessages) {
+      const r = messageToResourceMap.get(m);
+
+      if (r) {
+        destroySubjects.push(r.subject);
+
+        try {
+          await r.destroy();
+        } catch (error) {
+          console.error('Error removing message:', error);
+        }
+      } else {
+        throw new Error('Resource not found for message:', m);
+      }
+    }
+
+    try {
+      // Set chat resource on server with new message array
+      await resource.set(
+        ai.properties.messages,
+        resource.props.messages?.filter(x => !destroySubjects.includes(x)),
+      );
+      await resource.save();
+      // Set internal message state
+      setMessages(prev => {
+        const newMessages = prev.slice(0, prev.indexOf(message) + 1);
+
+        console.log(
+          'setting messages',
+          newMessages,
+          'last message',
+          newMessages.at(-1).content[0].text,
+        );
+
+        return newMessages;
+      });
+    } catch (error) {
+      console.error('Error removing messages:', error);
+    }
+  };
+
   // On load create AIChatDisplayMessages from the resource's messages.
   useEffect(() => {
     messageResourcesToDisplayMessages(messageSubjects, store).then(map => {
@@ -110,6 +162,7 @@ export const AIChatPage: React.FC<ResourcePageProps> = ({ resource }) => {
       externalContextItems={contextItems}
       setExternalContextItems={setContextItems}
       onDeleteMessage={handleDeleteMessage}
+      onRegenerateMessage={removeFollowingMessages}
     >
       <Column gap='0.5rem'>
         <Row>