Merge pull request #40 from pattern-tech/feat/accounting

feat: add query usage box

Author: mkermani144

app/(chat)/adapter.ts

@@ -8,6 +8,7 @@ import type {
   ApiSendMessageStreamedResponse,
   ApiGetAllConversationsResponse,
   ApiRenameConversationResponse,
+  ApiQueryUsageResponse,
 } from '@/app/(chat)/types';
 import config from '@/config';
 import { extractErrorMessageOrDefault } from '@/lib/utils';
@@ -359,3 +360,37 @@ export const renameConversation = async (
     );
   }
 };
+
+/**
+ * Get query usage
+ * @param accessToken
+ * @returns result containing query usage
+ */
+export const getQueryUsage = async (
+  accessToken: string,
+): Promise<Result<ApiQueryUsageResponse, string>> => {
+  try {
+    const response = await fetch(`${patternCoreEndpoint}/query-usage`, {
+      headers: {
+        Authorization: `Bearer ${accessToken}`,
+        'Content-Type': 'application/json',
+      },
+    });
+
+    if (response.ok) {
+      const queryUsage: ApiQueryUsageResponse = (await response.json()).data;
+      return Ok(queryUsage);
+    }
+
+    return Err(
+      `Fetching query usage failed with error code ${response.status}`,
+    );
+  } catch (error) {
+    return Err(
+      extractErrorMessageOrDefault(
+        error,
+        'Unknown error while fetching query usage',
+      ),
+    );
+  }
+};

app/(chat)/api/query-usage/route.ts

@@ -0,0 +1,26 @@
+import { auth } from '@/app/(auth)/auth';
+
+import { getQueryUsage } from '../../service';
+
+export async function GET() {
+  const session = await auth();
+
+  if (
+    !session ||
+    !session.chainId ||
+    !session.address ||
+    !session.accessToken
+  ) {
+    return Response.json('Unauthorized!', { status: 401 });
+  }
+
+  const queryUsageResult = await getQueryUsage(session.accessToken);
+
+  if (queryUsageResult.isErr()) {
+    return new Response(queryUsageResult.error, { status: 400 });
+  }
+
+  const queryUsage = queryUsageResult.value;
+
+  return Response.json(queryUsage);
+}

app/(chat)/service.ts

@@ -7,6 +7,7 @@ import {
   renameConversation,
   getConversation,
   getAllConversations,
+  getQueryUsage as getQueryUsageAdapter,
 } from './adapter';
 import type { Conversation } from './types';
 
@@ -95,6 +96,45 @@ export const getAllChats = async (
   return Ok(history);
 };
 
+/**
+ * Get query usage
+ * @param accessToken
+ * @returns result containing query usage
+ */
+export const getQueryUsage = async (
+  accessToken: string,
+): Promise<
+  Result<
+    {
+      todayQueryCount: number;
+      remainingQueriesToday: number;
+      maxQueryAllowancePerDay: number;
+      nextResetTime: string;
+    },
+    string
+  >
+> => {
+  const result = await getQueryUsageAdapter(accessToken);
+
+  if (result.isErr()) {
+    return Err(result.error);
+  }
+
+  const {
+    today_query_count,
+    remaining_queries_today,
+    max_query_allowance_per_day,
+    next_reset_time,
+  } = result.value;
+
+  return Ok({
+    todayQueryCount: today_query_count,
+    remainingQueriesToday: remaining_queries_today,
+    maxQueryAllowancePerDay: max_query_allowance_per_day,
+    nextResetTime: next_reset_time,
+  });
+};
+
 export {
   sendMessage,
   sendMessageStreamed,

app/(chat)/types.ts

@@ -11,6 +11,13 @@ export interface Message {
   content: string;
 }
 
+export interface QueryUsage {
+  todayQueryCount: number;
+  remainingQueriesToday: number;
+  maxQueryAllowancePerDay: number;
+  nextResetTime: string;
+}
+
 export type ApiGetConversationResponse = Conversation | null;
 export type ApiGetConversationMessagesResponse = Message[];
 export type ApiCreateConversationResponse = Conversation;
@@ -20,3 +27,9 @@ export type ApiGetAllConversationsResponse = Conversation[];
 export interface ApiRenameConversationResponse {
   title: string;
 }
+export interface ApiQueryUsageResponse {
+  today_query_count: number;
+  remaining_queries_today: number;
+  max_query_allowance_per_day: number;
+  next_reset_time: string;
+}

app/layout.tsx

@@ -1,4 +1,5 @@
 import type { Metadata } from 'next';
+import { SessionProvider } from 'next-auth/react';
 import { headers } from 'next/headers';
 import { Toaster } from 'sonner';
 import { cookieToInitialState } from 'wagmi';
@@ -59,15 +60,17 @@ export default async function RootLayout({
       </head>
       <body className="antialiased">
         <ContextProvider initialState={initialState}>
-          <ThemeProvider
-            attribute="class"
-            defaultTheme="system"
-            enableSystem
-            disableTransitionOnChange
-          >
-            <Toaster position="top-center" />
-            {children}
-          </ThemeProvider>
+          <SessionProvider>
+            <ThemeProvider
+              attribute="class"
+              defaultTheme="system"
+              enableSystem
+              disableTransitionOnChange
+            >
+              <Toaster position="top-center" />
+              {children}
+            </ThemeProvider>
+          </SessionProvider>
         </ContextProvider>
       </body>
     </html>

components/chat-header.tsx

@@ -9,6 +9,7 @@ import { Button } from '@/components/ui/button';
 
 import AppkitButton from './appkit-button';
 import { PlusIcon } from './icons';
+import { QueryUsage } from './query-usage';
 import { useSidebar } from './ui/sidebar';
 import { Tooltip, TooltipContent, TooltipTrigger } from './ui/tooltip';
 
@@ -41,7 +42,8 @@ function PureChatHeader() {
         </Tooltip>
       )}
 
-      <div className="order-4 m-4 md:ml-auto">
+      <div className="order-4 m-4 md:ml-auto flex items-center gap-4">
+        <QueryUsage />
         <AppkitButton />
       </div>
     </header>

components/chat.tsx

@@ -46,6 +46,7 @@ export function Chat({
     generateId: generateUUID,
     onFinish: () => {
       mutate('/api/history');
+      mutate('/api/query-usage');
     },
     onError: (error) => {
       toast.error(error.message);

components/query-usage.tsx

@@ -0,0 +1,68 @@
+'use client';
+
+import { formatDistanceToNow } from 'date-fns';
+import { useSession } from 'next-auth/react';
+import useSWR from 'swr';
+
+import type { QueryUsage as QueryUsageType } from '@/app/(chat)/types';
+import { fetcher } from '@/lib/utils';
+
+import { Tooltip, TooltipContent, TooltipTrigger } from './ui/tooltip';
+
+export function QueryUsage() {
+  const { status } = useSession();
+  const { data, isLoading } = useSWR<QueryUsageType>(
+    status === 'authenticated' ? '/api/query-usage' : null,
+    fetcher,
+    { revalidateOnFocus: false, revalidateOnReconnect: false },
+  );
+
+  if (status !== 'authenticated') return null;
+
+  if (isLoading) {
+    return (
+      <div className="flex items-center gap-2 text-sm text-muted-foreground">
+        <div className="size-4 animate-pulse rounded-full bg-muted" />
+        <span>Loading usage...</span>
+      </div>
+    );
+  }
+
+  if (!data) return null;
+
+  const {
+    todayQueryCount,
+    remainingQueriesToday,
+    maxQueryAllowancePerDay,
+    nextResetTime,
+  } = data;
+
+  const resetTime = formatDistanceToNow(new Date(nextResetTime), {
+    addSuffix: true,
+  });
+
+  return (
+    <Tooltip>
+      <TooltipTrigger asChild>
+        <div className="flex items-center justify-between text-sm cursor-pointer border rounded-md px-2 py-1 dark:border-zinc-700">
+          <span className="text-muted-foreground mr-2">Credits</span>
+          <span className="font-medium">
+            {remainingQueriesToday} / {maxQueryAllowancePerDay}
+          </span>
+        </div>
+      </TooltipTrigger>
+      <TooltipContent>
+        <div className="flex flex-col gap-2">
+          <div className="flex items-center gap-2">
+            <span className="text-muted-foreground">Used today:</span>
+            <span className="font-medium">{todayQueryCount}</span>
+          </div>
+          <div className="flex items-center gap-2">
+            <span className="text-muted-foreground">Resets:</span>
+            <span className="font-medium">{resetTime}</span>
+          </div>
+        </div>
+      </TooltipContent>
+    </Tooltip>
+  );
+}