From c80ae8fe6d03d69532ca4a0b2c2f510249cf79fa Mon Sep 17 00:00:00 2001
From: Lancelot Owczarczak <lancelot@owczarczak.fr>
Date: Thu, 17 Aug 2023 17:42:20 +0200
Subject: [PATCH] :sparkles: Add usage & email

---
 lib/index.ts | 14 +++++++++++--
 lib/user.ts  | 55 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 2 deletions(-)
 create mode 100644 lib/user.ts

diff --git a/lib/index.ts b/lib/index.ts
index dc5fab5..54df820 100644
--- a/lib/index.ts
+++ b/lib/index.ts
@@ -26,6 +26,7 @@ import memoryClient, {
     MemoryClient,
 } from "./memory";
 import { splitString, tokenCount } from "./split";
+import userClient, { usage, UserClient } from "./user";
 import { InputClientOptions } from "./clientOpts";
 import kvClient, { get as KVGet, set as KVSet, KVClient } from "./kv";
 
@@ -58,13 +59,15 @@ export {
     Chat,
     Memory,
     kv,
+    usage,
 };
 
 type Client = GenerationClient &
     GenerationWithTypeClient &
     TranscribeClient &
     MemoryClient &
-    ChatClient & { kv: KVClient };
+    ChatClient &
+    UserClient & { kv: KVClient };
 
 function client(co: InputClientOptions): Client {
     return {
@@ -73,6 +76,7 @@ function client(co: InputClientOptions): Client {
         ...transcribeClient(co),
         ...memoryClient(co),
         ...chatClient(co),
+        ...userClient(co),
         kv: kvClient(co),
     };
 }
@@ -245,6 +249,7 @@ type Provider = "github" | "google";
 export function usePolyfact({ project, endpoint }: { project: string; endpoint?: string }): {
     polyfact: Client | undefined;
     login: ((input: { provider: Provider }) => Promise<void>) | undefined;
+    email?: string;
     loading: boolean;
 } {
     if (typeof window === "undefined") {
@@ -253,6 +258,7 @@ export function usePolyfact({ project, endpoint }: { project: string; endpoint?:
 
     const react = require("react"); // eslint-disable-line
     const [polyfact, setPolyfact] = react.useState();
+    const [email, setEmail] = react.useState();
     const [loading, setLoading] = react.useState(true);
     const [login, setLogin] = react.useState();
 
@@ -298,6 +304,10 @@ export function usePolyfact({ project, endpoint }: { project: string; endpoint?:
                         data.session?.refresh_token,
                     );
                 }
+
+                const { data } = await supabase.auth.getUser(token);
+
+                setEmail(data.user?.email);
                 const p = await Polyfact.endpoint(endpoint || "https://api2.polyfact.com")
                     .project(project)
                     .signInWithToken(token);
@@ -316,5 +326,5 @@ export function usePolyfact({ project, endpoint }: { project: string; endpoint?:
         })();
     }, []);
 
-    return { polyfact, login, loading };
+    return { polyfact, login, loading, email };
 }
diff --git a/lib/user.ts b/lib/user.ts
new file mode 100644
index 0000000..e878445
--- /dev/null
+++ b/lib/user.ts
@@ -0,0 +1,55 @@
+import axios, { AxiosError } from "axios";
+import * as t from "polyfact-io-ts";
+import { InputClientOptions, defaultOptions } from "./clientOpts";
+import { ApiError, ErrorData } from "./helpers/error";
+
+const UsageResultType = t.intersection([
+    t.type({
+        usage: t.number,
+    }),
+    t.partial({
+        rate_limit: t.union([t.undefined, t.null, t.number]),
+    }),
+]);
+
+export async function usage(
+    clientOptions: InputClientOptions = {},
+): Promise<{ usage: number; rateLimit?: number }> {
+    const { token, endpoint } = await defaultOptions(clientOptions);
+
+    try {
+        const res = await axios.get(`${endpoint}/usage`, {
+            headers: {
+                "Content-Type": "application/json",
+                "X-Access-Token": token,
+            },
+        });
+
+        if (!UsageResultType.is(res.data)) {
+            throw new ApiError({
+                code: "mismatched_response",
+                message: "The response from the API does not match the expected format",
+            });
+        }
+
+        return {
+            usage: res.data.usage,
+            rateLimit: res.data.rate_limit ?? undefined,
+        };
+    } catch (e: unknown) {
+        if (e instanceof AxiosError) {
+            throw new ApiError(e?.response?.data as ErrorData);
+        }
+        throw e;
+    }
+}
+
+export type UserClient = {
+    usage: () => Promise<{ usage: number; rateLimit?: number }>;
+};
+
+export default function client(clientOptions: InputClientOptions = {}): UserClient {
+    return {
+        usage: () => usage(clientOptions),
+    };
+}