Reverted the usage-store (#380)

Zipstack · Jun 3, 2024 · 3f86e97 · 3f86e97
1 parent 979b1b0
commit 3f86e97
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/frontend/src/components/helpers/socket-messages/SocketMessages.js b/frontend/src/components/helpers/socket-messages/SocketMessages.js
@@ -7,6 +7,7 @@ import { useSocketLogsStore } from "../../../store/socket-logs-store";
 import { useSocketMessagesStore } from "../../../store/socket-messages-store";
 import { useSocketCustomToolStore } from "../../../store/socket-custom-tool";
 import { useSessionStore } from "../../../store/session-store";
+import { useUsageStore } from "../../../store/usage-store";
 
 function SocketMessages() {
   const [logId, setLogId] = useState("");
@@ -23,6 +24,7 @@ function SocketMessages() {
   const socket = useContext(SocketContext);
   const { setAlertDetails } = useAlertStore();
   const handleException = useExceptionHandler();
+  const { setLLMTokenUsage } = useUsageStore();
 
   useEffect(() => {
     setLogId(sessionDetails?.logEventsId || "");
@@ -44,6 +46,11 @@ function SocketMessages() {
       if (msg?.type === "LOG" && msg?.service === "prompt") {
         updateCusToolMessages(msg);
       }
+      if (msg?.type === "LOG" && msg?.service === "usage") {
+        const remainingTokens =
+          msg?.max_token_count_set - msg?.added_token_count;
+        setLLMTokenUsage(Math.max(remainingTokens, 0));
+      }
     } catch (err) {
       setAlertDetails(handleException(err, "Failed to process socket message"));
     }

diff --git a/frontend/src/store/usage-store.js b/frontend/src/store/usage-store.js
@@ -0,0 +1,10 @@
+import { create } from "zustand";
+
+const useUsageStore = create((setState) => ({
+  llmTokenUsage: null,
+  setLLMTokenUsage: (data) => {
+    setState({ llmTokenUsage: data });
+  },
+}));
+
+export { useUsageStore };