Merge pull request #2090 from quadratichq/jim-ai-fixes-misc

fix: ai tweaks

.vscode/settings.json

@@ -92,5 +92,8 @@
   },
   "[javascript]": {
     "editor.defaultFormatter": "esbenp.prettier-vscode"
+  },
+  "[html]": {
+    "editor.defaultFormatter": "esbenp.prettier-vscode"
   }
 }

quadratic-api/src/routes/ai/anthropic.ts

@@ -28,12 +28,12 @@ anthropic_router.post('/anthropic/chat', validateAccessToken, ai_rate_limiter, a
     });
     response.json(result.content);
   } catch (error: any) {
-    if (error.response) {
-      response.status(error.response.status).json(error.response.data);
-      console.log(error.response.status, error.response.data);
+    if (error instanceof Anthropic.APIError) {
+      response.status(error.status ?? 400).json(error.message);
+      console.log(error.status, error.message);
     } else {
-      response.status(400).json(error.message);
-      console.log(error.message);
+      response.status(400).json(error);
+      console.log(error);
     }
   }
 });
@@ -75,14 +75,15 @@ anthropic_router.post(
       }
     } catch (error: any) {
       if (!response.headersSent) {
-        if (error.response) {
-          response.status(error.response.status).json(error.response.data);
-          console.log(error.response.status, error.response.data);
+        if (error instanceof Anthropic.APIError) {
+          response.status(error.status ?? 400).json(error.message);
+          console.log(error.status, error.message);
         } else {
-          response.status(400).json(error.message);
-          console.log(error.message);
+          response.status(400).json(error);
+          console.log(error);
         }
       } else {
+        response.status(500).json('Error occurred after headers were sent');
         console.error('Error occurred after headers were sent:', error);
       }
     }

quadratic-api/src/routes/ai/bedrock.ts

@@ -1,4 +1,5 @@
 import { AnthropicBedrock } from '@anthropic-ai/bedrock-sdk';
+import Anthropic from '@anthropic-ai/sdk';
 import { BedrockRuntimeClient, ConverseCommand, ConverseStreamCommand } from '@aws-sdk/client-bedrock-runtime';
 import express from 'express';
 import { MODEL_OPTIONS } from 'quadratic-shared/AI_MODELS';
@@ -125,12 +126,12 @@ bedrock_router.post('/bedrock/anthropic/chat', validateAccessToken, ai_rate_limi
     });
     response.json(result.content);
   } catch (error: any) {
-    if (error.response) {
-      response.status(error.response.status).json(error.response.data);
-      console.log(error.response.status, error.response.data);
+    if (error instanceof Anthropic.APIError) {
+      response.status(error.status ?? 400).json(error.message);
+      console.log(error.status, error.message);
     } else {
-      response.status(400).json(error.message);
-      console.log(error.message);
+      response.status(400).json(error);
+      console.log(error);
     }
   }
 });
@@ -173,12 +174,12 @@ bedrock_router.post(
       }
     } catch (error: any) {
       if (!response.headersSent) {
-        if (error.response) {
-          response.status(error.response.status).json(error.response.data);
-          console.log(error.response.status, error.response.data);
+        if (error instanceof Anthropic.APIError) {
+          response.status(error.status ?? 400).json(error.message);
+          console.log(error.status, error.message);
         } else {
-          response.status(400).json(error.message);
-          console.log(error.message);
+          response.status(400).json(error);
+          console.log(error);
         }
       } else {
         console.error('Error occurred after headers were sent:', error);

quadratic-api/src/routes/ai/openai.ts

@@ -26,12 +26,12 @@ openai_router.post('/openai/chat', validateAccessToken, ai_rate_limiter, async (
     });
     response.json(result.choices[0].message);
   } catch (error: any) {
-    if (error.response) {
-      response.status(error.response.status).json(error.response.data);
-      console.log(error.response.status, error.response.data);
+    if (error instanceof OpenAI.APIError) {
+      response.status(error.status ?? 400).json(error.message);
+      console.log(error.status, error.message);
     } else {
-      response.status(400).json(error.message);
-      console.log(error.message);
+      response.status(400).json(error);
+      console.log(error);
     }
   }
 });
@@ -65,12 +65,12 @@ openai_router.post('/openai/chat/stream', validateAccessToken, ai_rate_limiter,
     }
   } catch (error: any) {
     if (!response.headersSent) {
-      if (error.response) {
-        response.status(error.response.status).json(error.response.data);
-        console.log(error.response.status, error.response.data);
+      if (error instanceof OpenAI.APIError) {
+        response.status(error.status ?? 400).json(error.message);
+        console.log(error.status, error.message);
       } else {
-        response.status(400).json(error.message);
-        console.log(error.message);
+        response.status(400).json(error);
+        console.log(error);
       }
     } else {
       console.error('Error occurred after headers were sent:', error);

quadratic-client/src/app/ai/hooks/useAIRequestToAPI.tsx

@@ -356,17 +356,16 @@ export function useAIRequestToAPI() {
         });
 
         if (!response.ok) {
+          const data = await response.json();
           const error =
             response.status === 429
               ? 'You have exceeded the maximum number of requests. Please try again later.'
-              : `Looks like there was a problem. Status Code: ${response.status}`;
+              : `Looks like there was a problem. Error: ${data}`;
           setMessages?.((prev) => [
             ...prev.slice(0, -1),
             { role: 'assistant', content: error, contextType: 'userPrompt', model, toolCalls: [] },
           ]);
-          if (response.status !== 429) {
-            console.error(`Error retrieving data from AI API: ${response.status}`);
-          }
+          console.error(`Error retrieving data from AI API. Error: ${data}`);
           return { error: true, content: error, toolCalls: [] };
         }
 

quadratic-client/src/app/ai/hooks/useCodeCellContextMessages.tsx

@@ -68,6 +68,7 @@ Use any functions that are part of the ${language} library.\n
 A code cell can return only one type of value as specified in the Quadratic documentation.\n
 A code cell cannot display both a chart and return a data frame at the same time.\n
 A code cell cannot display multiple charts at the same time.\n
+Do not use conditional returns in code cells.\n
 Do not use any markdown syntax besides triple backticks for ${language} code blocks.\n
 Do not reply code blocks in plain text, use markdown with triple backticks and language name ${language}.`
 }

quadratic-client/src/app/ai/hooks/useQuadraticContextMessages.tsx

@@ -26,7 +26,7 @@ ${
     : 'Choose the language of your response based on the context and user prompt.'
 }
 Provide complete code blocks with language syntax highlighting. Don't provide small code snippets of changes.
-Respond in minimum number of words with direct answer. Include a concise explanation of the answer.
+Respond in minimum number of words and include a concise explanation of the actions you are taking. Don't guess the answer itself, just the actions you are taking to respond to the user prompt and what the user can do next.
 `,
         contextType: 'quadraticDocs',
       },

quadratic-client/src/app/ai/hooks/useToolUseMessages.tsx

@@ -10,7 +10,7 @@ export function useToolUseMessages() {
         content: `Note: This is an internal message for context. Do not quote it in your response.\n\n
 Following are the tools you should use to do actions in the spreadsheet, use them to respond to the user prompt.\n
 
-Include a concise explanation of the actions you are taking to respond to the user prompt.
+Include a concise explanation of the actions you are taking to respond to the user prompt. Never guess the answer itself, just the actions you are taking to respond to the user prompt and what the user can do next.\n
 
 Don't include tool details in your response. Reply in layman's terms what actions you are taking.\n
 

quadratic-client/src/app/ai/hooks/useVisibleContextMessages.tsx

@@ -68,9 +68,10 @@ Add imports to the top of the code cell and do not use any libraries or function
 Use any functions that are part of the code cell language library.\n
 A code cell can return only one type of value as specified in the Quadratic documentation.\n
 A code cell cannot display both a chart and return a data frame at the same time.\n
+Do not use conditional returns in code cells.\n
 A code cell cannot display multiple charts at the same time.\n
 Do not use any markdown syntax besides triple backticks for code cell language code blocks.\n
-Do not reply code blocks in plain text, use markdown with triple backticks and language name code cell language.
+Do not reply code blocks in plain text, use markdown with triple backticks and language name code cell language.\n
 
 ${erroredCodeCells[0].map(({ x, y, language, code_string, std_out, std_err }) => {
   const consoleOutput = {

quadratic-client/src/app/ai/tools/aiToolsSpec.ts

@@ -223,6 +223,7 @@ The required location (x,y) for this code cell is one which satisfies the follow
  - If there are multiple tables or data sources being referenced, place the code cell in a location that provides a good balance between proximity to all referenced data and maintaining readability of the currently open sheet.\n
  - Consider the overall layout and organization of the currently open sheet when placing the code cell, ensuring it doesn't disrupt existing data or interfere with other code cells.\n
  - A plot returned by the code cell occupies just one cell, the plot overlay is larger but the code cell is always just one cell.\n
+ - Do not use conditional returns in code cells.\n
  `,
   },
   [AITool.MoveCells]: {

quadratic-client/src/app/atoms/aiAnalystAtom.ts

@@ -36,11 +36,14 @@ export const aiAnalystAtom = atom<AIAnalystState>({
   effects: [
     async ({ getPromise, setSelf, trigger }) => {
       if (trigger === 'get') {
+        // Determine if we want to override the default showAIAnalyst value on initialization
+        const aiAnalystOpenCount = getAiAnalystOpenCount();
         const isSheetEmpty = sheets.sheet.bounds.type === 'empty';
-        if (isSheetEmpty) {
+        const showAIAnalyst = aiAnalystOpenCount <= 3 ? true : isSheetEmpty;
+        if (showAIAnalyst) {
           setSelf({
             ...defaultAIAnalystState,
-            showAIAnalyst: true,
+            showAIAnalyst,
           });
         }
 
@@ -52,7 +55,7 @@ export const aiAnalystAtom = atom<AIAnalystState>({
             const chats = await aiAnalystOfflineChats.loadChats();
             setSelf({
               ...defaultAIAnalystState,
-              showAIAnalyst: isSheetEmpty,
+              showAIAnalyst,
               chats,
             });
           } catch (error) {
@@ -271,3 +274,14 @@ export const aiAnalystCurrentChatMessagesCountAtom = selector<number>({
   key: 'aiAnalystCurrentChatMessagesCountAtom',
   get: ({ get }) => getPromptMessages(get(aiAnalystCurrentChatAtom).messages).length,
 });
+
+const STORAGE_KEY = 'aiAnalystOpenCount';
+export function getAiAnalystOpenCount() {
+  const count = window.localStorage.getItem(STORAGE_KEY);
+  return count ? parseInt(count) : 0;
+}
+export function incrementAiAnalystOpenCount() {
+  const count = getAiAnalystOpenCount();
+  const newCount = count + 1;
+  window.localStorage.setItem(STORAGE_KEY, newCount.toString());
+}

quadratic-client/src/app/keyboard/defaults.ts

@@ -358,7 +358,7 @@ export const defaultShortcuts: ActionShortcut = {
     windows: [[WindowsModifiers.Ctrl, WindowsModifiers.Shift, Keys.Semicolon]],
   },
   [Action.ToggleAIAnalyst]: {
-    mac: [[MacModifiers.Cmd, MacModifiers.Shift, Keys.L]],
-    windows: [[WindowsModifiers.Ctrl, WindowsModifiers.Shift, Keys.L]],
+    mac: [[MacModifiers.Cmd, Keys.K]],
+    windows: [[WindowsModifiers.Ctrl, Keys.K]],
   },
 };

quadratic-client/src/app/ui/QuadraticSidebar.tsx

@@ -1,7 +1,7 @@
 import { isAvailableBecauseCanEditFile, isAvailableBecauseFileLocationIsAccessibleAndWriteable } from '@/app/actions';
 import { Action } from '@/app/actions/actions';
 import { defaultActionSpec } from '@/app/actions/defaultActionsSpec';
-import { showAIAnalystAtom } from '@/app/atoms/aiAnalystAtom';
+import { incrementAiAnalystOpenCount, showAIAnalystAtom } from '@/app/atoms/aiAnalystAtom';
 import { codeEditorShowCodeEditorAtom } from '@/app/atoms/codeEditorAtom';
 import {
   editorInteractionStateShowCommandPaletteAtom,
@@ -75,7 +75,19 @@ export const QuadraticSidebar = () => {
       <div className="mt-2 flex flex-col items-center gap-1">
         {canEditFile && isAuthenticated && (
           <SidebarTooltip label={toggleAIChat.label} shortcut={keyboardShortcutEnumToDisplay(Action.ToggleAIAnalyst)}>
-            <SidebarToggle pressed={showAIAnalyst} onPressedChange={() => setShowAIAnalyst((prev) => !prev)}>
+            <SidebarToggle
+              pressed={showAIAnalyst}
+              onPressedChange={() => {
+                setShowAIAnalyst((prevShowAiAnalyst) => {
+                  // if it's hidden and therefore being opened by the user, count it!
+                  if (prevShowAiAnalyst === false) {
+                    incrementAiAnalystOpenCount();
+                  }
+
+                  return !prevShowAiAnalyst;
+                });
+              }}
+            >
               <AIIcon />
             </SidebarToggle>
           </SidebarTooltip>

quadratic-client/src/app/ui/components/AIUserMessageForm.tsx

@@ -14,7 +14,7 @@ import { SetterOrUpdater } from 'recoil';
 
 export type AIUserMessageFormWrapperProps = {
   textareaRef: React.RefObject<HTMLTextAreaElement>;
-  autoFocus?: boolean;
+  autoFocusRef?: React.RefObject<boolean>;
   initialPrompt?: string;
   messageIndex?: number;
 };
@@ -36,7 +36,7 @@ export const AIUserMessageForm = forwardRef<HTMLTextAreaElement, Props>((props:
   const {
     initialPrompt,
     ctx,
-    autoFocus,
+    autoFocusRef,
     textareaRef: bottomTextareaRef,
     abortController,
     loading,
@@ -58,12 +58,10 @@ export const AIUserMessageForm = forwardRef<HTMLTextAreaElement, Props>((props:
 
   // Focus the input when relevant & the tab comes into focus
   useEffect(() => {
-    if (autoFocus) {
-      window.requestAnimationFrame(() => {
-        textareaRef.current?.focus();
-      });
+    if (autoFocusRef?.current) {
+      textareaRef.current?.focus();
     }
-  }, [autoFocus, textareaRef]);
+  }, [autoFocusRef, textareaRef]);
 
   useEffect(() => {
     if (loading && initialPrompt !== undefined) {
@@ -115,6 +113,7 @@ export const AIUserMessageForm = forwardRef<HTMLTextAreaElement, Props>((props:
 
             if (event.key === 'Enter' && !(event.ctrlKey || event.shiftKey)) {
               event.preventDefault();
+              if (loading) return;
 
               if (prompt.trim().length === 0) return;
 
@@ -128,6 +127,9 @@ export const AIUserMessageForm = forwardRef<HTMLTextAreaElement, Props>((props:
                 bottomTextareaRef.current?.focus();
               }
             }
+
+            if (loading) return;
+
             if (formOnKeyDown) {
               formOnKeyDown(event);
             }

quadratic-client/src/app/ui/menus/AIAnalyst/AIAnalyst.tsx

@@ -9,7 +9,7 @@ import { AIAnalystMessages } from '@/app/ui/menus/AIAnalyst/AIAnalystMessages';
 import { AIAnalystUserMessageForm } from '@/app/ui/menus/AIAnalyst/AIAnalystUserMessageForm';
 import { useAIAnalystPanelWidth } from '@/app/ui/menus/AIAnalyst/hooks/useAIAnalystPanelWidth';
 import { cn } from '@/shared/shadcn/utils';
-import { useCallback, useRef } from 'react';
+import { useCallback, useEffect, useRef } from 'react';
 import { useRecoilValue } from 'recoil';
 
 export const AIAnalyst = () => {
@@ -20,6 +20,16 @@ export const AIAnalyst = () => {
   const textareaRef = useRef<HTMLTextAreaElement>(null);
   const { panelWidth, setPanelWidth } = useAIAnalystPanelWidth();
 
+  const initialLoadRef = useRef(true);
+  const autoFocusRef = useRef(false);
+  useEffect(() => {
+    if (initialLoadRef.current) {
+      initialLoadRef.current = false;
+    } else {
+      autoFocusRef.current = true;
+    }
+  }, [showAIAnalyst]);
+
   const handleResize = useCallback(
     (event: MouseEvent) => {
       const panel = aiPanelRef.current;
@@ -67,7 +77,7 @@ export const AIAnalyst = () => {
               <AIAnalystMessages textareaRef={textareaRef} />
 
               <div className="px-2 py-0.5">
-                <AIAnalystUserMessageForm ref={textareaRef} autoFocus={true} textareaRef={textareaRef} />
+                <AIAnalystUserMessageForm ref={textareaRef} autoFocusRef={autoFocusRef} textareaRef={textareaRef} />
                 <AIUserMessageFormDisclaimer />
               </div>
             </>