feat: Display the total amount of tokens currently consumed at the menu.

CHANGE_LOG.md

@@ -2,7 +2,7 @@
 
 ## v0.6.1
 
-> 2023-06-11
+> 2023-06-13
 
 ### Fixed
 
@@ -11,6 +11,7 @@
 ### Add
 
 - Added backend capability to calculate token consumption
+- Display the total amount of tokens currently consumed at the menu.
 
 ### Changed
 

CHANGE_LOG.zh_CN.md

@@ -2,7 +2,7 @@
 
 ## v0.6.1
 
-> 2023-06-11
+> 2023-06-13
 
 ### 修复
 
@@ -11,6 +11,7 @@
 ### 新增
 
 - 新增后端计算 token 消耗的能力
+- 新增菜单处显示当前消耗的总 token 数量
 
 ### 修改
 

src/app/[locale]/(account)/account/page.tsx

@@ -42,8 +42,8 @@ const Account: React.FC = () => {
     };
 
     setLoading(true);
-    fetch("/api/user/update", {
-      method: "POST",
+    fetch("/api/user", {
+      method: "PUT",
       body: JSON.stringify(params),
     })
       .then(async (response) => {

src/app/api/user/update/route.ts → src/app/api/user/route.ts

@@ -2,8 +2,9 @@ import { NextResponse } from "next/server";
 import { getServerSession } from "next-auth/next";
 import { authOptions } from "@/app/api/auth/[...nextauth]/route";
 import { prisma } from "@/lib/prisma";
+import { LResponseError } from "@/lib";
 
-export async function POST(request: Request) {
+export async function PUT(request: Request) {
   const session = await getServerSession(authOptions);
 
   // check session
@@ -41,3 +42,27 @@ export async function POST(request: Request) {
 
   return NextResponse.json({ error: 0 }, { status: 200 });
 }
+
+export async function GET(request: Request) {
+  const session = await getServerSession(authOptions);
+
+  if (!session) return LResponseError("Please log in first");
+
+  try {
+    const user = await prisma.user.findUnique({
+      where: { id: session?.user.id },
+    });
+
+    if (!user) return LResponseError("User does not exist");
+
+    const response = {
+      costTokens: user.costTokens,
+      costUSD: user.costUSD,
+    };
+
+    return NextResponse.json({ error: 0, data: response }, { status: 200 });
+  } catch (error) {
+    console.log(error, "get user error");
+    return NextResponse.json({ error: -1, msg: "error" }, { status: 500 });
+  }
+}

src/components/auth/avatar.tsx

@@ -88,10 +88,6 @@ const Avatar: React.FC = () => {
     }
   };
 
-  React.useEffect(() => {
-    console.log(1234);
-  }, []);
-
   return (
     <Dropdown
       className="min-w-[8rem]"

src/components/chatSection/chatFooter/index.tsx

@@ -1,6 +1,7 @@
 import * as React from "react";
 import { useTranslations } from "next-intl";
 import { useRouter } from "next-intl/client";
+import { useSession } from "next-auth/react";
 import { v4 as uuidv4 } from "uuid";
 import type { ChatItem } from "@/hooks/useChannel";
 import {
@@ -16,6 +17,7 @@ import {
   useOpenAI,
   useStreamDecoder,
   useChat,
+  useTokens,
   useLLM,
   BASE_PROMPT,
 } from "@/hooks";
@@ -29,8 +31,10 @@ import Inputarea from "./inputArea";
 
 const ChatFooter: React.FC = () => {
   // data
+  const session = useSession();
   const [newOpenAI] = useOpenAI();
   const [channel, setChannel] = useChannel();
+  const [, setTokens] = useTokens();
   const { openai, azure } = useLLM();
   const {
     loadingResponseStart,
@@ -121,6 +125,8 @@ const ChatFooter: React.FC = () => {
         return channel;
       });
 
+      if (session.data) setTokens();
+
       return;
     }
 
@@ -381,6 +387,8 @@ const ChatFooter: React.FC = () => {
             return channel;
           });
 
+          if (session.data) setTokens();
+
           // get conversation title
           if (!channel_name) getChannelNameByGPT(channel_id, channel_chat_list);
         })
@@ -498,6 +506,8 @@ const ChatFooter: React.FC = () => {
 
         return channel;
       });
+
+      if (session.data) setTokens();
     });
   };
 
@@ -622,6 +632,8 @@ const ChatFooter: React.FC = () => {
         return channel;
       });
 
+      if (session.data) setTokens();
+
       return;
     }
     setInputValue("");

src/components/menu/index.tsx

@@ -1,4 +1,5 @@
 import * as React from "react";
+import { useSession } from "next-auth/react";
 import { useParams } from "next/navigation";
 import { useRouter } from "next-intl/client";
 import { useTranslations } from "next-intl";
@@ -30,6 +31,7 @@ import type { ContextMenuOption } from "@/components/ui/ContextMenu";
 import Dropdown from "@/components/ui/Dropdown";
 import type { IDropdownItems } from "@/components/ui/Dropdown";
 import Logo from "@/components/logo";
+import Tokens from "@/components/tokens";
 import MenuIcon from "./icon";
 
 export const lans: IDropdownItems[] = [
@@ -46,6 +48,7 @@ export const lans: IDropdownItems[] = [
 ];
 
 const Menu: React.FC = () => {
+  const session = useSession();
   const { theme, setTheme } = useTheme();
   const t = useTranslations("menu");
   const { format } = useDateFormat();
@@ -180,7 +183,12 @@ const Menu: React.FC = () => {
       >
         {t("new-chat")}
       </Button>
-      <div className="h-[calc(100vh-16.75rem)] overflow-y-auto">
+      <div
+        className={cn("overflow-y-auto", {
+          "h-[calc(100vh-19.75rem)]": session.data,
+          "h-[calc(100vh-16.75rem)]": !session.data,
+        })}
+      >
         {channel.list.map((item) => (
           <ContextMenu
             key={item.channel_id}
@@ -255,7 +263,12 @@ const Menu: React.FC = () => {
           </ContextMenu>
         ))}
       </div>
-      <div className="border-t flex flex-col h-[9.25rem] pt-2 gap-1 dark:border-white/20">
+      <div
+        className={cn(
+          "border-t flex flex-col pt-2 gap-1 dark:border-white/20",
+          { "h-[12.25rem]": session.data, "h-[9.25rem]": !session.data }
+        )}
+      >
         <Confirm
           title={t("clear-all-conversation")}
           content={t("clear-conversation")}
@@ -282,6 +295,7 @@ const Menu: React.FC = () => {
         >
           <RiFeedbackLine size={16} /> {t("feedback")}
         </a>
+        {!!session.data && <Tokens type="pc" />}
         <div className="flex h-11 items-center justify-center">
           <div className="flex flex-1 justify-center">
             <div

src/components/menu/mobile.tsx

@@ -1,4 +1,5 @@
 import * as React from "react";
+import { useSession } from "next-auth/react";
 import { useParams } from "next/navigation";
 import { useRouter } from "next-intl/client";
 import { useTranslations } from "next-intl";
@@ -28,9 +29,11 @@ import {
   BASE_PROMPT,
 } from "@/hooks";
 import MenuIcon from "./icon";
+import Tokens from "@/components/tokens";
 import { lans } from "./index";
 
 const MobileMenu: React.FC = () => {
+  const session = useSession();
   const { theme, setTheme } = useTheme();
   const t = useTranslations("menu");
   const { format } = useDateFormat();
@@ -138,7 +141,7 @@ const MobileMenu: React.FC = () => {
       open={mobileMenuVisible}
       onClose={onClose}
     >
-      <div className="p-2 h-[calc(100%-3.5rem)] flex flex-col">
+      <div className="p-2 flex flex-col h-[calc(100%-3.5rem)]">
         <Button
           className="mb-2"
           type="primary"
@@ -223,7 +226,12 @@ const MobileMenu: React.FC = () => {
             </div>
           ))}
         </div>
-        <div className="h-[9rem] flex flex-col border-t gap-1 pt-1">
+        <div
+          className={cn("flex flex-col border-t gap-1 pt-1", {
+            "h-[12rem]": session.data,
+            "h-[9rem]": !session.data,
+          })}
+        >
           <Confirm
             title={t("clear-all-conversation")}
             content={t("clear-conversation")}
@@ -252,6 +260,7 @@ const MobileMenu: React.FC = () => {
           >
             <RiFeedbackLine size={16} /> {t("feedback")}
           </a>
+          {!!session.data && <Tokens type="mobile" />}
           <div className="h-11 items-center justify-center flex">
             <div className="flex-1 flex justify-center">
               <div

src/components/tokens/index.tsx

@@ -0,0 +1,40 @@
+import * as React from "react";
+import { cn } from "@/lib";
+import { AiOutlinePieChart, AiOutlineLoading } from "react-icons/ai";
+import { useTokens } from "@/hooks";
+
+interface TokenProps {
+  type: "pc" | "mobile";
+}
+
+const Tokens: React.FC<TokenProps> = ({ type }) => {
+  const [tokens, setTokens] = useTokens();
+
+  const { costTokens, loading } = tokens;
+
+  React.useEffect(() => {
+    setTokens();
+  }, []);
+
+  return (
+    <div
+      className={cn(
+        "h-11 text-sm cursor-pointer flex items-center gap-2 px-2 transition-colors",
+        "hover:bg-gray-200/60 dark:hover:bg-slate-700/70",
+        { "rounded-md text-black/90 dark:text-white/90": type === "mobile" }
+      )}
+    >
+      <AiOutlinePieChart size={16} />
+      <div className="flex items-center gap-2">
+        {!!(!costTokens && loading) ? (
+          <AiOutlineLoading className="animate-spin" />
+        ) : (
+          costTokens
+        )}
+        <span>Tokens</span>
+      </div>
+    </div>
+  );
+};
+
+export default Tokens;

src/hooks/index.ts

@@ -8,3 +8,4 @@ export * from "./useSetting";
 export * from "./useStream";
 export * from "./useLLM";
 export * from "./usePrompt";
+export * from "./useTokens";

src/hooks/useTokens.ts

@@ -0,0 +1,47 @@
+import { create } from "zustand";
+
+interface ITokensState {
+  loading: boolean;
+  costTokens: number;
+  costUSD: number;
+}
+
+interface ITokensAction {
+  update: () => void;
+}
+
+type useTokensReturn = [ITokensState, ITokensAction["update"]];
+
+let timeout: NodeJS.Timeout | null = null;
+
+const useStore = create<ITokensState & ITokensAction>((set) => ({
+  loading: false,
+  costTokens: 0,
+  costUSD: 0,
+  update: () => {
+    if (timeout) {
+      clearTimeout(timeout);
+      timeout = null;
+    }
+    timeout = setTimeout(async () => {
+      try {
+        set(() => ({ loading: true }));
+        const res = await fetch("/api/user").then((res) => res.json());
+        const { costTokens, costUSD } = res.data;
+        set(() => ({ costTokens, costUSD }));
+      } catch (error) {
+        set(() => ({ costTokens: 0, costUSD: 0 }));
+      } finally {
+        set(() => ({ loading: false }));
+      }
+    }, 3000);
+  },
+}));
+
+export const useTokens = (): useTokensReturn => {
+  const { costTokens, costUSD, loading } = useStore();
+
+  const update = useStore((state) => state.update);
+
+  return [{ costTokens, costUSD, loading }, update];
+};