From 8d5824bc2e2b19aa31955e2951a5bcacf391db35 Mon Sep 17 00:00:00 2001
From: Arik Chakma <arikchangma@gmail.com>
Date: Fri, 22 Nov 2024 13:47:15 +0600
Subject: [PATCH] feat: course ai token limit

---
 src/components/CourseAI/CourseAI.tsx        |  2 +-
 src/components/CourseAI/CourseAILimit.tsx   | 24 ++++++++++++++++
 src/components/CourseAI/CourseAIPopover.tsx | 32 +++++++++++++++------
 src/hooks/use-course.ts                     | 18 ++++++++++++
 4 files changed, 66 insertions(+), 10 deletions(-)
 create mode 100644 src/components/CourseAI/CourseAILimit.tsx

diff --git a/src/components/CourseAI/CourseAI.tsx b/src/components/CourseAI/CourseAI.tsx
index 04cf1e086..a76d5cc4d 100644
--- a/src/components/CourseAI/CourseAI.tsx
+++ b/src/components/CourseAI/CourseAI.tsx
@@ -33,7 +33,7 @@ export function CourseAI(props: CourseAIProps) {
   return (
     <div className="relative">
       <button
-        className="flex items-center gap-1 rounded-lg border border-black pl-3 pr-4 py-2 text-sm leading-none disabled:opacity-60 hover:bg-black hover:text-white transition-colors"
+        className="flex items-center gap-1 rounded-lg border border-black py-2 pl-3 pr-4 text-sm leading-none transition-colors hover:bg-black hover:text-white disabled:opacity-60"
         onClick={() => setIsOpen(!isOpen)}
       >
         <Sparkles className="size-4 stroke-[2]" />
diff --git a/src/components/CourseAI/CourseAILimit.tsx b/src/components/CourseAI/CourseAILimit.tsx
new file mode 100644
index 000000000..9ae5633e0
--- /dev/null
+++ b/src/components/CourseAI/CourseAILimit.tsx
@@ -0,0 +1,24 @@
+import { getPercentage } from '../../helper/number';
+import { useCourseAILimit } from '../../hooks/use-course';
+
+export function CourseAILimit() {
+  const { data: tokenUsage, isLoading } = useCourseAILimit();
+
+  if (isLoading || !tokenUsage) {
+    return (
+      <div className="h-5 w-1/4 animate-pulse rounded-md bg-zinc-600"></div>
+    );
+  }
+
+  const percentageUsed = getPercentage(
+    tokenUsage.usedTokenCount,
+    tokenUsage.maxTokenCount,
+  );
+
+  return (
+    <p className="text-sm text-yellow-500">
+      <strong className="font-medium">{percentageUsed}%</strong> of daily limit
+      used
+    </p>
+  );
+}
diff --git a/src/components/CourseAI/CourseAIPopover.tsx b/src/components/CourseAI/CourseAIPopover.tsx
index 90c4f7fde..783a7d622 100644
--- a/src/components/CourseAI/CourseAIPopover.tsx
+++ b/src/components/CourseAI/CourseAIPopover.tsx
@@ -12,6 +12,8 @@ import { flushSync } from 'react-dom';
 import type { AIChatHistoryType, AllowedAIChatRole } from './CourseAI';
 import { useToast } from '../../hooks/use-toast';
 import { removeAuthToken } from '../../lib/jwt';
+import { useCourseAILimit } from '../../hooks/use-course';
+import { CourseAILimit } from './CourseAILimit';
 
 type CourseAIPopoverProps = {
   courseId: string;
@@ -42,17 +44,26 @@ export function CourseAIPopover(props: CourseAIPopoverProps) {
   const containerRef = useRef<HTMLDivElement | null>(null);
   const scrollareaRef = useRef<HTMLDivElement | null>(null);
 
-  const [isLoading, setIsLoading] = useState(false);
+  const {
+    data: tokenUsage,
+    isLoading,
+    refetch: refreshLimit,
+  } = useCourseAILimit();
+
+  const [isStreamingMessage, setIsStreamingMessage] = useState(false);
   const [message, setMessage] = useState('');
   const [streamedMessage, setStreamedMessage] = useState('');
 
   useOutsideClick(containerRef, onOutsideClick);
 
+  const isLimitExceeded =
+    (tokenUsage?.maxTokenCount || 0) <= (tokenUsage?.usedTokenCount || 0);
+
   const handleChatSubmit = (e: FormEvent<HTMLFormElement>) => {
     e.preventDefault();
 
     const trimmedMessage = message.trim();
-    if (!trimmedMessage || isLoading) {
+    if (!trimmedMessage || isStreamingMessage || isLoading || isLimitExceeded) {
       return;
     }
 
@@ -81,7 +92,7 @@ export function CourseAIPopover(props: CourseAIPopoverProps) {
   };
 
   const completeCourseAIChat = async (messages: AIChatHistoryType[]) => {
-    setIsLoading(true);
+    setIsStreamingMessage(true);
 
     const response = await fetch(
       `${import.meta.env.PUBLIC_API_URL}/v1-course-ai/${courseId}`,
@@ -105,7 +116,7 @@ export function CourseAIPopover(props: CourseAIPopoverProps) {
 
       toast.error(data?.message || 'Something went wrong');
       setCourseAIChatHistory([...messages].slice(0, messages.length - 1));
-      setIsLoading(false);
+      setIsStreamingMessage(false);
 
       // Logout user if token is invalid
       if (data.status === 401) {
@@ -117,7 +128,7 @@ export function CourseAIPopover(props: CourseAIPopoverProps) {
     const reader = response.body?.getReader();
 
     if (!reader) {
-      setIsLoading(false);
+      setIsStreamingMessage(false);
       toast.error('Something went wrong');
       return;
     }
@@ -141,7 +152,7 @@ export function CourseAIPopover(props: CourseAIPopoverProps) {
 
         flushSync(() => {
           setStreamedMessage('');
-          setIsLoading(false);
+          setIsStreamingMessage(false);
           setCourseAIChatHistory(newMessages);
         });
 
@@ -149,7 +160,8 @@ export function CourseAIPopover(props: CourseAIPopoverProps) {
       },
     });
 
-    setIsLoading(false);
+    refreshLimit();
+    setIsStreamingMessage(false);
   };
 
   useEffect(() => {
@@ -163,6 +175,8 @@ export function CourseAIPopover(props: CourseAIPopoverProps) {
     >
       <div className="flex items-center justify-between gap-2 border-b border-zinc-700 px-4 py-2 text-sm">
         <h4 className="text-base font-medium">Roadmap AI</h4>
+
+        <CourseAILimit />
       </div>
 
       <div
@@ -182,7 +196,7 @@ export function CourseAIPopover(props: CourseAIPopoverProps) {
                 );
               })}
 
-              {isLoading && !streamedMessage && (
+              {isStreamingMessage && !streamedMessage && (
                 <AIChatCard role="assistant" content="Thinking..." />
               )}
 
@@ -207,7 +221,7 @@ export function CourseAIPopover(props: CourseAIPopoverProps) {
         />
         <button
           type="submit"
-          disabled={isLoading}
+          disabled={isStreamingMessage || isLoading || isLimitExceeded}
           className="flex aspect-square h-full items-center justify-center text-zinc-500 hover:text-zinc-50"
         >
           <Send className="size-4 stroke-[2.5]" />
diff --git a/src/hooks/use-course.ts b/src/hooks/use-course.ts
index cd15febf2..0e0e9f58c 100644
--- a/src/hooks/use-course.ts
+++ b/src/hooks/use-course.ts
@@ -56,3 +56,21 @@ export function useCompleteLessonMutation(courseId: string) {
     queryClient,
   );
 }
+
+export type CourseAILimitResponse = {
+  maxTokenCount: number;
+  usedTokenCount: number;
+};
+
+export function useCourseAILimit() {
+  return useQuery(
+    {
+      queryKey: ['course-ai-limit'],
+      queryFn: async () => {
+        return httpGet<CourseAILimitResponse>('/v1-course-ai-limit');
+      },
+      enabled: isLoggedIn(),
+    },
+    queryClient,
+  );
+}