Save total token usage

nehabagdia · Apr 29, 2024 · Apr 3, 2024 · Apr 9, 2024 · Apr 17, 2024 · Apr 18, 2024
commit 2bc37c8d1b77a68c7c72702a2c425127b846d4d1
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 { useAlertStore } from "../../../store/alert-store";
 import { useSocketLogsStore } from "../../../store/socket-logs-store";
 import { useSocketMessagesStore } from "../../../store/socket-messages-store";
 import { useSocketCustomToolStore } from "../../../store/socket-custom-tool";
+import { useUsageStore } from "../../../store/usage-store";
 function SocketMessages({ logId }) {
   const {
     pushStagedMessage,
@@ -20,6 +21,7 @@ function SocketMessages({ logId }) {
   const socket = useContext(SocketContext);
   const { setAlertDetails } = useAlertStore();
   const handleException = useExceptionHandler();
+  const { setLLMTokenUsage } = useUsageStore();
 
   const onMessage = (data) => {
     try {
@@ -37,6 +39,9 @@ function SocketMessages({ logId }) {
       if (msg?.type === "LOG" && msg?.service === "prompt") {
         updateCusToolMessages(msg);
       }
+      if (msg?.type === "LOG" && msg?.service === "usage") {
+        setLLMTokenUsage(msg?.added_token_count);
+      }
     } catch (err) {
       setAlertDetails(handleException(err, "Failed to process socket message"));
     }

diff --git a/frontend/src/components/navigations/top-nav-bar/TopNavBar.jsx b/frontend/src/components/navigations/top-nav-bar/TopNavBar.jsx
@@ -1,4 +1,13 @@
-import { Alert, Button, Col, Dropdown, Image, Row, Typography } from "antd";
+import {
+  Alert,
+  Button,
+  Col,
+  Dropdown,
+  Image,
+  Row,
+  Space,
+  Typography,
+} from "antd";
 import { useEffect, useState } from "react";
 import { useNavigate } from "react-router-dom";
 import axios from "axios";
@@ -163,26 +172,28 @@ function TopNavBar() {
       </Col>
       <Col span={4}>
         <Row justify="end" align="middle">
-          {TrialDaysInfo && (
-            <TrialDaysInfo remainingTrialDays={remainingTrialDays} />
-          )}
-          <Dropdown menu={{ items }} placement="bottomLeft" arrow>
-            <div className="top-navbar-dp">
-              {sessionDetails?.picture ? (
-                <Image
-                  className="navbar-img"
-                  height="100%"
-                  width="100%"
-                  preview={false}
-                  src={sessionDetails?.picture}
-                />
-              ) : (
-                <Typography.Text className="initials">
-                  {getInitials(sessionDetails?.name)}
-                </Typography.Text>
-              )}
-            </div>
-          </Dropdown>
+          <Space>
+            {TrialDaysInfo && (
+              <TrialDaysInfo remainingTrialDays={remainingTrialDays} />
+            )}
+            <Dropdown menu={{ items }} placement="bottomLeft" arrow>
+              <div className="top-navbar-dp">
+                {sessionDetails?.picture ? (
+                  <Image
+                    className="navbar-img"
+                    height="100%"
+                    width="100%"
+                    preview={false}
+                    src={sessionDetails?.picture}
+                  />
+                ) : (
+                  <Typography.Text className="initials">
+                    {getInitials(sessionDetails?.name)}
+                  </Typography.Text>
+                )}
+              </div>
+            </Dropdown>
+          </Space>
         </Row>
       </Col>
     </Row>

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 };