refetch token count on generation finished (BloopAI#952)

aileague · Sep 13, 2023 · bfdcff4 · bfdcff4
1 parent 8f1d73c
commit bfdcff4
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 13 deletions.
diff --git a/client/src/pages/StudioTab/Content.tsx b/client/src/pages/StudioTab/Content.tsx
@@ -58,14 +58,23 @@ const ContentContainer = ({
   const { leftPanelRef, rightPanelRef, dividerRef, containerRef } =
     useResizeableSplitPanel();
 
-  const refetchCodeStudio = useCallback(async () => {
-    if (tab.key) {
-      const resp = await getCodeStudio(tab.key);
-      setCurrentState((prev) =>
-        JSON.stringify(resp) === JSON.stringify(prev) ? prev : resp,
-      );
-    }
-  }, [tab.key]);
+  const refetchCodeStudio = useCallback(
+    async (keyToUpdate?: keyof CodeStudioType) => {
+      if (tab.key) {
+        const resp = await getCodeStudio(tab.key);
+        setCurrentState((prev) => {
+          if (JSON.stringify(resp) === JSON.stringify(prev)) {
+            return prev;
+          }
+          if (keyToUpdate) {
+            return { ...prev, [keyToUpdate]: resp[keyToUpdate] };
+          }
+          return resp;
+        });
+      }
+    },
+    [tab.key],
+  );
 
   useEffect(() => {
     if (isActive) {

diff --git a/client/src/pages/StudioTab/RightPanel/Conversation/index.tsx b/client/src/pages/StudioTab/RightPanel/Conversation/index.tsx
@@ -19,7 +19,7 @@ import {
 import Button from '../../../../components/Button';
 import { ArrowRefresh, TrashCanFilled } from '../../../../icons';
 import KeyboardChip from '../../KeyboardChip';
-import { CodeStudioMessageType } from '../../../../types/api';
+import { CodeStudioMessageType, CodeStudioType } from '../../../../types/api';
 import { patchCodeStudio } from '../../../../services/api';
 import useKeyboardNavigation from '../../../../hooks/useKeyboardNavigation';
 import { DeviceContext } from '../../../../context/deviceContext';
@@ -34,7 +34,7 @@ type Props = {
   setIsHistoryOpen: Dispatch<SetStateAction<boolean>>;
   messages: CodeStudioMessageType[];
   studioId: string;
-  refetchCodeStudio: () => Promise<void>;
+  refetchCodeStudio: (keyToUpdate?: keyof CodeStudioType) => Promise<void>;
   isTokenLimitExceeded: boolean;
   isPreviewing: boolean;
   isActiveTab: boolean;
@@ -69,6 +69,7 @@ const Conversation = ({
   isPreviewing,
   handleRestore,
   isActiveTab,
+  refetchCodeStudio,
 }: Props) => {
   const { t } = useTranslation();
   const { inputValue } = useContext(StudioContext.Input);
@@ -224,6 +225,7 @@ const Conversation = ({
       if (ev.data === '[DONE]') {
         eventSource.close();
         setIsLoading(false);
+        refetchCodeStudio('token_counts');
         return;
       }
       try {
@@ -312,6 +314,8 @@ const Conversation = ({
           .filter((m, j) => (andSubsequent ? i > j : i !== j));
       await patchCodeStudio(studioId, {
         messages,
+      }).then(() => {
+        refetchCodeStudio('token_counts');
       });
     },
     [conversation],
@@ -336,7 +340,7 @@ const Conversation = ({
         } else {
           if (inputValue && !isTokenLimitExceeded && requestsLeft) {
             onSubmit();
-          } else {
+          } else if (!requestsLeft) {
             setUpgradePopupOpen(true);
           }
         }

diff --git a/client/src/pages/StudioTab/RightPanel/index.tsx b/client/src/pages/StudioTab/RightPanel/index.tsx
@@ -3,7 +3,7 @@ import { Trans, useTranslation } from 'react-i18next';
 import TokensUsageProgress from '../TokensUsageProgress';
 import { TOKEN_LIMIT } from '../../../consts/codeStudio';
 import { StudioLeftPanelDataType } from '../../../types/general';
-import { CodeStudioMessageType } from '../../../types/api';
+import { CodeStudioMessageType, CodeStudioType } from '../../../types/api';
 import Conversation from './Conversation';
 
 type Props = {
@@ -12,7 +12,7 @@ type Props = {
   setIsHistoryOpen: Dispatch<SetStateAction<boolean>>;
   messages: CodeStudioMessageType[];
   studioId: string;
-  refetchCodeStudio: () => Promise<void>;
+  refetchCodeStudio: (keyToUpdate?: keyof CodeStudioType) => Promise<void>;
   handleRestore: () => void;
   isPreviewing: boolean;
   hasContextError: boolean;