AI Chat improvements #951

Author: Polleps

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

@@ -21,6 +21,8 @@ const AISettings: React.FC = () => {
     setOpenRouterApiKey,
     mcpServers,
     setMcpServers,
+    showTokenUsage,
+    setShowTokenUsage,
   } = useSettings();
 
   const [creditUsage, setCreditUsage] = useState<CreditUsage | undefined>();
@@ -71,6 +73,10 @@ const AISettings: React.FC = () => {
             )}
           </Column>
         </label>
+        <CheckboxLabel>
+          <Checkbox checked={showTokenUsage} onChange={setShowTokenUsage} />
+          Show token usage in chats
+        </CheckboxLabel>
         <Heading>MCP Servers</Heading>
         <MCPServersManager servers={mcpServers} setServers={setMcpServers} />
       </ConditionalSettings>

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

@@ -5,9 +5,7 @@ import toast from 'react-hot-toast';
 import { useAtomicMCPTools } from './useAtomicTools';
 import { AIChatMessage } from './AIChatMessage';
 import {
-  generateObject,
   InvalidToolArgumentsError,
-  NoSuchToolError,
   streamText,
   TypeValidationError,
   type CoreMessage,
@@ -112,7 +110,7 @@ export const SimpleAIChat: React.FC<
   const [userSelectedContextItems, setUserSelectedContextItems] = useState<
     AIMessageContext[]
   >([]);
-  const { openRouterApiKey } = useSettings() as { openRouterApiKey?: string };
+  const { openRouterApiKey, showTokenUsage } = useSettings();
   const openrouter = createOpenRouter({
     apiKey: openRouterApiKey,
     compatibility: 'strict',
@@ -287,38 +285,6 @@ export const SimpleAIChat: React.FC<
             });
           }
         },
-        experimental_repairToolCall: async ({
-          toolCall,
-          parameterSchema,
-          error,
-        }) => {
-          if (NoSuchToolError.isInstance(error)) {
-            console.log('No such tool error', error);
-
-            return null; // do not attempt to fix invalid tool names
-          }
-
-          console.log('arg error', error);
-          console.log('Repairing tool call', toolCall.toolName);
-
-          const { object: repairedArgs } = await generateObject({
-            model: openrouter('qwen/qwq-32b:free'),
-            output: 'no-schema',
-            mode: 'json',
-            prompt: [
-              `The model tried to call the tool "${toolCall.toolName}"` +
-                ` with the following arguments:`,
-              JSON.stringify(toolCall.args),
-              `The tool accepts the following schema:`,
-              JSON.stringify(parameterSchema(toolCall)),
-              'Please fix the arguments.',
-            ].join('\n'),
-          });
-
-          console.log('Repaired tool call', repairedArgs);
-
-          return { ...toolCall, args: JSON.stringify(repairedArgs) };
-        },
       });
     } catch (err) {
       console.error(err);
@@ -354,7 +320,6 @@ export const SimpleAIChat: React.FC<
     try {
       for await (const part of textStream.fullStream) {
         // Update ongoing message with streamed chunks
-        console.log('Part', part);
 
         if (part.type === 'tool-call') {
           const toolCallMessage: AIChatDisplayMessage = {
@@ -379,13 +344,15 @@ export const SimpleAIChat: React.FC<
           pendingToolCalls.push(part);
         }
 
+        // @ts-expect-error 'tool-result' is not in the types but there are definitely parts that have this type, removing this would break tool use.
         if (part.type === 'tool-result') {
           onNewMessage({
             role: 'tool',
             content: [part],
           });
 
           pendingToolCalls = pendingToolCalls.filter(
+            // @ts-expect-error part has type never do to incorrect types.
             call => call.toolCallId !== part.toolCallId,
           );
         }
@@ -672,14 +639,14 @@ export const SimpleAIChat: React.FC<
               </Column>
               <NoKeyOverlay />
             </ChatInputWrapper>
-            <TokensUsed>
-              Tokens used: {tokensUsed[0]} input, {tokensUsed[1]} output
-            </TokensUsed>
+            {showTokenUsage && (
+              <TokensUsed>
+                Tokens used: {tokensUsed[0]} input, {tokensUsed[1]} output
+              </TokensUsed>
+            )}
           </>
         )}
       </Column>
-
-      {/* Agent configuration dialog */}
       <AgentConfig
         open={agentConfigOpen}
         onOpenChange={setAgentConfigOpen}

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

@@ -5,6 +5,7 @@ import {
 } from 'react';
 import { styled, DefaultTheme, keyframes } from 'styled-components';
 import { transition } from '../../helpers/transition';
+import { adjustHue } from 'polished';
 
 export enum IconButtonVariant {
   Simple,
@@ -197,7 +198,7 @@ const cycle = keyframes`
     filter: hue-rotate(0deg) brightness(var(--brightness-factor)) saturate(1.5);
   }
   to {
-    filter: hue-rotate(360deg) brightness(var(--brightness-factor)) saturate(1.5);
+    filter: hue-rotate(360deg) brightness(var(--brightness-factor)) saturate(2);
   }
 `;
 
@@ -225,11 +226,17 @@ const MagicIconButton = styled(IconButtonBase)<ButtonStyleProps>`
     inset: 0;
     opacity: 0;
     z-index: -2;
-    background: linear-gradient(
-      45deg,
-      ${p => p.theme.colors.main},
-      ${p => p.theme.colors.complementary}
-    );
+    background: radial-gradient(ellipse at top right, #365ccd, transparent),
+      radial-gradient(
+        ellipse at bottom left,
+        ${adjustHue(-45, '#365ccd')},
+        transparent
+      ),
+      radial-gradient(
+        ellipse at bottom right,
+        ${adjustHue(180, '#365ccd')},
+        transparent
+      );
   }
   &::after {
     content: '';

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

@@ -64,6 +64,7 @@ export const MCPServersManager: React.FC<MCPServersManagerProps> = ({
         )}
       </ServerList>
       <Column gap='0.5rem'>
+        <h4>Add Server</h4>
         <Row gap='1rem' center>
           <Column gap='0.5rem' style={{ flex: 1 }}>
             <label htmlFor='server-name'>Server Name</label>
@@ -107,7 +108,6 @@ export const MCPServersManager: React.FC<MCPServersManagerProps> = ({
   );
 };
 
-// Styled components for MCP servers UI
 const ServerList = styled.div`
   display: flex;
   flex-direction: column;

browser/data-browser/src/helpers/AppSettings.tsx

@@ -59,6 +59,10 @@ export const AppSettingsContextProvider = (
   const [drive, innerSetDrive] = useLocalStorage('drive', baseURL);
 
   const [enableAI, setEnableAI] = useLocalStorage('atomic.ai.enabled', true);
+  const [showTokenUsage, setShowTokenUsage] = useLocalStorage(
+    'atomic.ai.showTokenUsage',
+    true,
+  );
 
   const setDrive = useCallback(
     (newDrive: string) => {
@@ -112,6 +116,8 @@ export const AppSettingsContextProvider = (
       setMcpServers,
       enableAI,
       setEnableAI,
+      showTokenUsage,
+      setShowTokenUsage,
     }),
     [
       drive,
@@ -141,6 +147,8 @@ export const AppSettingsContextProvider = (
       setMcpServers,
       enableAI,
       setEnableAI,
+      showTokenUsage,
+      setShowTokenUsage,
     ],
   );
 
@@ -195,6 +203,9 @@ export interface AppSettings {
   /** Enable all AI features in the app */
   enableAI: boolean;
   setEnableAI: (b: boolean) => void;
+  /** Whether to show the token usage in AI chats */
+  showTokenUsage: boolean;
+  setShowTokenUsage: (b: boolean) => void;
 }
 
 const initialState: AppSettings = {
@@ -225,6 +236,8 @@ const initialState: AppSettings = {
   setMcpServers: () => undefined,
   enableAI: true,
   setEnableAI: () => undefined,
+  showTokenUsage: true,
+  setShowTokenUsage: () => undefined,
 };
 
 /** Hook for using App Settings, such as theme and darkmode */