FIX: Removed Polling to Fetch Token Usage API (#668)

* Removed polling for token usage * Persist the token usage details from the output manager API response * Refactored the code related to token usage * Removed the custom hook that's used to fetch the token usage
Zipstack · Sep 10, 2024 · 0621724 · 0621724
1 parent 469617e
commit 0621724
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 75 deletions.
diff --git a/frontend/src/components/custom-tools/prompt-card/PromptCard.jsx b/frontend/src/components/custom-tools/prompt-card/PromptCard.jsx
@@ -3,7 +3,6 @@ import { useEffect, useState } from "react";
 import { useParams } from "react-router-dom";
 
 import {
-  defaultTokenUsage,
   generateUUID,
   pollForCompletion,
 } from "../../../helpers/GetStaticData";
@@ -15,7 +14,6 @@ import { useSessionStore } from "../../../store/session-store";
 import { useSocketCustomToolStore } from "../../../store/socket-custom-tool";
 import { OutputForDocModal } from "../output-for-doc-modal/OutputForDocModal";
 import usePostHogEvents from "../../../hooks/usePostHogEvents";
-import useTokenUsage from "../../../hooks/useTokenUsage";
 import { useTokenUsageStore } from "../../../store/token-usage-store";
 import { PromptCardItems } from "./PromptCardItems";
 import "./PromptCard.css";
@@ -62,7 +60,7 @@ function PromptCard({
   const [openOutputForDoc, setOpenOutputForDoc] = useState(false);
   const [progressMsg, setProgressMsg] = useState({});
   const [docOutputs, setDocOutputs] = useState([]);
-  const [timers, setTimers] = useState({});
+  const [timers, setTimers] = useState({}); // Prompt run timer
   const [spsLoading, setSpsLoading] = useState({});
   const {
     getDropdownItems,
@@ -85,8 +83,7 @@ function PromptCard({
   const axiosPrivate = useAxiosPrivate();
   const handleException = useExceptionHandler();
   const { setPostHogCustomEvent } = usePostHogEvents();
-  const { tokenUsage, setTokenUsage } = useTokenUsageStore();
-  const { getTokenUsage } = useTokenUsage();
+  const { setTokenUsage } = useTokenUsageStore();
   const { id } = useParams();
 
   useEffect(() => {
@@ -235,6 +232,10 @@ function PromptCard({
     }));
   };
 
+  const getTokenUsageId = (promptId, docId, profileManagerId) => {
+    return promptId + "__" + docId + "__" + profileManagerId;
+  };
+
   // Generate the result for the currently selected document
   const handleRun = (
     profileManagerId,
@@ -544,7 +545,6 @@ function PromptCard({
     const runId = generateUUID();
     const maxWaitTime = 30 * 1000; // 30 seconds
     const pollingInterval = 5000; // 5 seconds
-    const tokenUsagepollingInterval = 5000;
 
     const body = {
       document_id: docId,
@@ -553,25 +553,13 @@ function PromptCard({
 
     if (profileManagerId) {
       body.profile_manager = profileManagerId;
-      let intervalId;
       let tokenUsageId;
       let url = `/api/v1/unstract/${sessionDetails?.orgId}/prompt-studio/fetch_response/${details?.tool_id}`;
       if (!isSimplePromptStudio) {
         body["run_id"] = runId;
-        // Update the token usage state with default token usage for a specific document ID
-        tokenUsageId = promptId + "__" + docId + "__" + profileManagerId;
-        setTokenUsage(tokenUsageId, defaultTokenUsage);
-
-        // Set up an interval to fetch token usage data at regular intervals
-        if (
-          profileManagerId === selectedLlmProfileId &&
-          docId === selectedDoc?.document_id
-        ) {
-          intervalId = setInterval(
-            () => getTokenUsage(runId, tokenUsageId),
-            tokenUsagepollingInterval // Fetch token usage data every 5000 milliseconds (5 seconds)
-          );
-        }
+        tokenUsageId = getTokenUsageId(promptId, docId, profileManagerId);
+
+        // Set prompt run timer
         setTimers((prev) => ({
           ...prev,
           [tokenUsageId]: 0,
@@ -611,8 +599,6 @@ function PromptCard({
         })
         .finally(() => {
           if (!isSimplePromptStudio) {
-            clearInterval(intervalId);
-            getTokenUsage(runId, tokenUsageId);
             stopTimer(tokenUsageId, timerIntervalId);
           }
         });
@@ -655,7 +641,6 @@ function PromptCard({
             profileManager: outputResult?.profile_manager,
             context: outputResult?.context,
             challengeData: outputResult?.challenge_data,
-            tokenUsage: outputResult?.token_usage,
             output: outputResult?.output,
             evalMetrics: getEvalMetrics(
               promptDetails?.evaluate,
@@ -748,16 +733,11 @@ function PromptCard({
             `${tokenUsageId}__challenge_data`,
             usage?.challenge_data
           );
-          if (usage) {
-            setTokenUsage(tokenUsageId, usage?.token_usage);
-          }
+          setTokenUsage(tokenUsageId, usage?.token_usage);
         } else {
           data?.forEach((item) => {
             const tokenUsageId = `${item?.prompt_id}__${item?.document_manager}__${item?.profile_manager}`;
-
-            if (tokenUsage[tokenUsageId] === undefined) {
-              setTokenUsage(tokenUsageId, item?.token_usage);
-            }
+            setTokenUsage(tokenUsageId, item?.token_usage);
           });
         }
         return res;

diff --git a/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx b/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx
@@ -206,7 +206,14 @@ function PromptOutput({
                     <Typography.Text className="prompt-cost-item">
                       Tokens:{" "}
                       {!singlePassExtractMode && (
-                        <TokenUsage tokenUsageId={tokenUsageId} />
+                        <TokenUsage
+                          tokenUsageId={tokenUsageId}
+                          isLoading={
+                            isRunLoading[
+                              `${selectedDoc?.document_id}_${profileId}`
+                            ]
+                          }
+                        />
                       )}
                     </Typography.Text>
                     <Typography.Text className="prompt-cost-item">

diff --git a/frontend/src/components/custom-tools/token-usage/TokenUsage.jsx b/frontend/src/components/custom-tools/token-usage/TokenUsage.jsx
@@ -9,7 +9,7 @@ import { formatNumberWithCommas } from "../../../helpers/GetStaticData";
  * @param {string} tokenUsageId - The token usage ID to fetch token usage for.
  * @return {JSX.Element} - The TokenUsage component.
  */
-function TokenUsage({ tokenUsageId }) {
+function TokenUsage({ tokenUsageId, isLoading }) {
   const [tokens, setTokens] = useState({});
   const { tokenUsage } = useTokenUsageStore();
 
@@ -24,15 +24,16 @@ function TokenUsage({ tokenUsageId }) {
   }, [tokenUsage, tokenUsageId]);
 
   // If no tokens data is available, render nothing
-  if (!tokens || !Object.keys(tokens)?.length) {
-    return 0;
+  if (!tokens || !Object.keys(tokens)?.length || isLoading) {
+    return "NA";
   }
 
   return formatNumberWithCommas(tokens?.total_tokens);
 }
 
 TokenUsage.propTypes = {
   tokenUsageId: PropTypes.string.isRequired,
+  isLoading: PropTypes.bool.isRequired,
 };
 
 export { TokenUsage };
diff --git a/frontend/src/hooks/useTokenUsage.js b/frontend/src/hooks/useTokenUsage.js