From 9c4cdc0f6a731a240499e5904e7a40f3bff8c389 Mon Sep 17 00:00:00 2001 From: Ethan Liu Date: Tue, 13 Jun 2023 00:43:24 +0800 Subject: [PATCH] feat: Display the total amount of tokens currently consumed at the menu. --- CHANGE_LOG.md | 3 +- CHANGE_LOG.zh_CN.md | 3 +- src/app/[locale]/(account)/account/page.tsx | 4 +- src/app/api/user/{update => }/route.ts | 27 ++++++++++- src/components/auth/avatar.tsx | 4 -- .../chatSection/chatFooter/index.tsx | 12 +++++ src/components/menu/index.tsx | 18 ++++++- src/components/menu/mobile.tsx | 13 ++++- src/components/tokens/index.tsx | 40 ++++++++++++++++ src/hooks/index.ts | 1 + src/hooks/useTokens.ts | 47 +++++++++++++++++++ 11 files changed, 159 insertions(+), 13 deletions(-) rename src/app/api/user/{update => }/route.ts (58%) create mode 100644 src/components/tokens/index.tsx create mode 100644 src/hooks/useTokens.ts diff --git a/CHANGE_LOG.md b/CHANGE_LOG.md index 3867318..cf4699d 100644 --- a/CHANGE_LOG.md +++ b/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 diff --git a/CHANGE_LOG.zh_CN.md b/CHANGE_LOG.zh_CN.md index e26c552..d7b8f30 100644 --- a/CHANGE_LOG.zh_CN.md +++ b/CHANGE_LOG.zh_CN.md @@ -2,7 +2,7 @@ ## v0.6.1 -> 2023-06-11 +> 2023-06-13 ### 修复 @@ -11,6 +11,7 @@ ### 新增 - 新增后端计算 token 消耗的能力 +- 新增菜单处显示当前消耗的总 token 数量 ### 修改 diff --git a/src/app/[locale]/(account)/account/page.tsx b/src/app/[locale]/(account)/account/page.tsx index 9e29c63..985cdf4 100644 --- a/src/app/[locale]/(account)/account/page.tsx +++ b/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) => { diff --git a/src/app/api/user/update/route.ts b/src/app/api/user/route.ts similarity index 58% rename from src/app/api/user/update/route.ts rename to src/app/api/user/route.ts index c03bed9..e87ec0f 100644 --- a/src/app/api/user/update/route.ts +++ b/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 }); + } +} diff --git a/src/components/auth/avatar.tsx b/src/components/auth/avatar.tsx index a0660f6..c770e0f 100644 --- a/src/components/auth/avatar.tsx +++ b/src/components/auth/avatar.tsx @@ -88,10 +88,6 @@ const Avatar: React.FC = () => { } }; - React.useEffect(() => { - console.log(1234); - }, []); - return ( { // 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(""); diff --git a/src/components/menu/index.tsx b/src/components/menu/index.tsx index ef2e959..ff7e801 100644 --- a/src/components/menu/index.tsx +++ b/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")} -
+
{channel.list.map((item) => ( { ))}
-
+
{ > {t("feedback")} + {!!session.data && }
{ + 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} > -
+
))}
-
+
{ > {t("feedback")} + {!!session.data && }
= ({ type }) => { + const [tokens, setTokens] = useTokens(); + + const { costTokens, loading } = tokens; + + React.useEffect(() => { + setTokens(); + }, []); + + return ( +
+ +
+ {!!(!costTokens && loading) ? ( + + ) : ( + costTokens + )} + Tokens +
+
+ ); +}; + +export default Tokens; diff --git a/src/hooks/index.ts b/src/hooks/index.ts index 978828d..8f731a6 100644 --- a/src/hooks/index.ts +++ b/src/hooks/index.ts @@ -8,3 +8,4 @@ export * from "./useSetting"; export * from "./useStream"; export * from "./useLLM"; export * from "./usePrompt"; +export * from "./useTokens"; diff --git a/src/hooks/useTokens.ts b/src/hooks/useTokens.ts new file mode 100644 index 0000000..df02682 --- /dev/null +++ b/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((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]; +};