From a62b5d803aa73f4c36c71f938762c10fe4f8181f Mon Sep 17 00:00:00 2001
From: Tianzhou Chen <t@bytebase.com>
Date: Wed, 7 Jun 2023 23:56:19 +0800
Subject: [PATCH] chore: display token in the prompt drawer

---
 src/components/SchemaDrawer.tsx | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/components/SchemaDrawer.tsx b/src/components/SchemaDrawer.tsx
index 3c72077..a334733 100644
--- a/src/components/SchemaDrawer.tsx
+++ b/src/components/SchemaDrawer.tsx
@@ -1,9 +1,9 @@
+import { encode } from "@nem035/gpt-3-encoder";
 import { Drawer } from "@mui/material";
 import { useEffect, useState } from "react";
 import Icon from "./Icon";
 import { getAssistantById, getPromptGeneratorOfAssistant, useConnectionStore, useConversationStore, useSettingStore } from "@/store";
 import { getModel, generateDbPromptFromContext } from "@/utils";
-import toast from "react-hot-toast";
 import { CodeBlock } from "./CodeBlock";
 import { useTranslation } from "react-i18next";
 
@@ -16,8 +16,8 @@ const SchemaDrawer = (props: Props) => {
   const conversationStore = useConversationStore();
   const connectionStore = useConnectionStore();
   const settingStore = useSettingStore();
-
   const currentConversation = conversationStore.getConversationById(conversationStore.currentConversationId);
+  const maxToken = getModel(settingStore.setting.openAIApiConfig?.model || "").max_token;
   const [prompt, setPrompt] = useState<string>("");
 
   const getPrompt = async () => {
@@ -25,7 +25,6 @@ const SchemaDrawer = (props: Props) => {
     const promptGenerator = getPromptGeneratorOfAssistant(getAssistantById(currentConversation.assistantId)!);
     let dbPrompt = promptGenerator();
     if (connectionStore.currentConnectionCtx?.database) {
-      const maxToken = getModel(settingStore.setting.openAIApiConfig?.model || "").max_token;
       const schemaList = await connectionStore.getOrFetchDatabaseSchema(connectionStore.currentConnectionCtx?.database);
       dbPrompt = generateDbPromptFromContext(
         promptGenerator,
@@ -49,11 +48,16 @@ const SchemaDrawer = (props: Props) => {
   const close = () => props.close();
   return (
     <Drawer open={true} anchor="right" className="w-full" onClose={close}>
-      <div className="dark:text-gray-300 w-screen sm:w-[calc(50vw)] max-w-full flex flex-col justify-start items-start p-4">
+      <div className="dark:text-gray-300 w-screen sm:w-[calc(75vw)] max-w-full flex flex-col  p-4">
         <button className="w-8 h-8 p-1 bg-zinc-600 text-gray-100 rounded-full hover:opacity-80" onClick={close}>
           <Icon.IoMdClose className="w-full h-auto" />
         </button>
-        <h3 className="font-bold text-2xl mt-4">{t("prompt.current-conversation")}</h3>
+        <div className="flex justify-between items-center">
+          <h3 className="font-bold text-2xl mt-4">{t("prompt.current-conversation")}</h3>
+          <div className="text-base mt-4">
+            {t("connection.total-token")} {encode(prompt).length}/{maxToken}
+          </div>
+        </div>
         <div>
           <CodeBlock language="Prompt" value={prompt} messageId={currentConversation?.id || ""} wrapLongLines={true} />
         </div>