From 06217249997c5a5c33b887a55a6812929e553933 Mon Sep 17 00:00:00 2001 From: Tahier Hussain <89440263+tahierhussain@users.noreply.github.com> Date: Tue, 10 Sep 2024 16:17:48 +0530 Subject: [PATCH] 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 --- .../custom-tools/prompt-card/PromptCard.jsx | 42 +++++-------------- .../custom-tools/prompt-card/PromptOutput.jsx | 9 +++- .../custom-tools/token-usage/TokenUsage.jsx | 7 ++-- frontend/src/hooks/useTokenUsage.js | 40 ------------------ 4 files changed, 23 insertions(+), 75 deletions(-) delete mode 100644 frontend/src/hooks/useTokenUsage.js diff --git a/frontend/src/components/custom-tools/prompt-card/PromptCard.jsx b/frontend/src/components/custom-tools/prompt-card/PromptCard.jsx index 8cceaa993..1304a4df7 100644 --- 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 index 936f10509..2dc3b4dfa 100644 --- 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({ Tokens:{" "} {!singlePassExtractMode && ( - + )} diff --git a/frontend/src/components/custom-tools/token-usage/TokenUsage.jsx b/frontend/src/components/custom-tools/token-usage/TokenUsage.jsx index 2f242dedb..513f1633b 100644 --- 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,8 +24,8 @@ 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); @@ -33,6 +33,7 @@ function TokenUsage({ tokenUsageId }) { 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 deleted file mode 100644 index 527c1fdd8..000000000 --- a/frontend/src/hooks/useTokenUsage.js +++ /dev/null @@ -1,40 +0,0 @@ -import { useAlertStore } from "../store/alert-store"; -import { useSessionStore } from "../store/session-store"; -import { useTokenUsageStore } from "../store/token-usage-store"; -import { useAxiosPrivate } from "./useAxiosPrivate"; -import { useExceptionHandler } from "./useExceptionHandler"; - -/** - * Custom hook to fetch token usage data. - * - * @return {Object} - An object containing the getTokenUsage function. - */ -const useTokenUsage = () => { - const axiosPrivate = useAxiosPrivate(); - const { sessionDetails } = useSessionStore(); - const { setAlertDetails } = useAlertStore(); - const handleException = useExceptionHandler(); - const { setTokenUsage } = useTokenUsageStore(); - - const getTokenUsage = (runId, tokenUsageId) => { - const requestOptions = { - method: "GET", - url: `/api/v1/unstract/${sessionDetails?.orgId}/usage/get_token_usage/?run_id=${runId}`, - }; - - // Make the API request using the axios instance - axiosPrivate(requestOptions) - .then((res) => { - const tokens = res?.data; - // Update the token usage state with token usage data for a specific tokenUsageId - setTokenUsage(tokenUsageId, tokens); - }) - .catch((err) => { - setAlertDetails(handleException(err, "Failed to get token usage")); - }); - }; - - return { getTokenUsage }; -}; - -export default useTokenUsage;