feat: integrate AWS BE

packages/frontend/src/layouts/Gen/MainForm/Job.tsx

@@ -1,72 +1,37 @@
-import { Spinner } from "@/components/spinner";
 import { Button } from "@/components/ui/button";
-import { Card, CardContent, CardHeader } from "@/components/ui/card";
-import { Progress } from "@/components/ui/progress";
+import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
 import { Tables } from "@/repo/database.types";
-import React, { useEffect } from "react";
+import React from "react";
 import { useNavigate } from "react-router-dom";
 
 export interface JobProps {
   doc: Tables<"doc_meta">;
 }
 
 const Job: React.FC<JobProps> = ({ doc }) => {
-  const [progress, setProgress] = React.useState(0);
   const navigate = useNavigate();
 
-  const handleResolveProgress = () => {
-    switch (doc?.status) {
-      case "started":
-        return 5;
-      case "searching":
-        return 20;
-      case "persona":
-        return 30;
-      case "refining":
-        return 50;
-      case "writing":
-        return 60;
-      case "polishing":
-        return 80;
-      case "uploading":
-        return 90;
-      case "completed":
-        return 100;
-      case "error":
-        return 0;
-      default:
-        return 0;
-    }
-  };
-
-  useEffect(() => {
-    const timer = setTimeout(() => setProgress(handleResolveProgress()), 500);
-    return () => clearTimeout(timer);
-  }, [doc]);
-
   const handleNavigateToDocViewer = () => {
     navigate("/viewer/" + encodeURIComponent(doc?.file_name));
   };
 
   return (
     <Card>
       <CardHeader>
-        <Progress value={progress} />
+        <CardTitle>Result</CardTitle>
       </CardHeader>
       <CardContent>
         <div className="flex items-center justify-between">
           <div className="flex-[2] font-semibold text-xl">{doc?.title}</div>
 
           <div className="flex items-center gap-2">
             <div className="p-2 flex-1 bg-slate-100 rounded-full flex gap-1 items-center font-semibold text-sm">
-              {doc.status !== "completed" ? <Spinner /> : null}
-              {doc.status?.toUpperCase() || "UNKNOWN"}
+              COMPLETED
             </div>
 
             <Button
               className="flex items-center gap-1"
               onClick={() => handleNavigateToDocViewer()}
-              disabled={doc?.status !== "completed"}
             >
               Open
             </Button>

packages/frontend/src/layouts/Gen/MainForm/index.tsx

@@ -17,16 +17,15 @@ import {
 } from "@/components/ui/form";
 import { Input } from "@/components/ui/input";
 import { Separator } from "@/components/ui/separator";
+import { useToast } from "@/components/ui/use-toast";
 import { Tables } from "@/repo/database.types";
 import { handleGetLatestDoc } from "@/repo/docMeta";
 import { supabase } from "@/supabase";
 import { zodResolver } from "@hookform/resolvers/zod";
-import React, { useEffect, useRef, useState } from "react";
+import React, { useState } from "react";
 import { useForm } from "react-hook-form";
 import { z } from "zod";
 import Job from "./Job";
-import { useQueryClient } from "@tanstack/react-query";
-import { useToast } from "@/components/ui/use-toast";
 
 interface MainFormProps {}
 
@@ -44,14 +43,10 @@ const MainForm: React.FC<MainFormProps> = ({}) => {
     },
   });
   const [doc, setDoc] = useState<Tables<"doc_meta">>();
-  const [needRefetch, setNeedRefetch] = useState(false);
   const [isLoading, setIsLoading] = useState(false);
   const tokens = useTokensQuery();
-  const intervalRef = useRef<number>();
-  const queryClient = useQueryClient();
   const { toast } = useToast();
 
-  useEffect(() => {}, []);
   const onSubmit = (values: z.infer<typeof GenSchema>) => {
     handleSubmitGenRequest(values.topic);
   };
@@ -78,7 +73,7 @@ const MainForm: React.FC<MainFormProps> = ({}) => {
       }
 
       const baseUrl = import.meta.env.VITE_SUPABASE_URL || "";
-      await fetch(`${baseUrl}/functions/v1/backend/gen/emit`, {
+      await fetch(`${baseUrl}/functions/v1/backend/v2/gen`, {
         method: "POST",
         headers: {
           "Content-Type": "application/json",
@@ -91,64 +86,27 @@ const MainForm: React.FC<MainFormProps> = ({}) => {
         }),
       });
 
-      queryClient.removeQueries({
-        queryKey: ["get-doc", doc?.file_name],
-      });
-      setDoc(undefined);
-      setNeedRefetch(!needRefetch);
+      const doc = await handleGetLatestDoc();
+
+      if (!doc) {
+        throw new Error("Doc not found");
+      }
+
+      setDoc(doc);
       setIsLoading(false);
     } catch (error) {
       setIsLoading(false);
     }
   };
 
-  const handleGetDoc = async (): Promise<Tables<"doc_meta"> | undefined> => {
-    const docRes = await handleGetLatestDoc();
-    if (docRes) {
-      setDoc(docRes);
-      return docRes;
-    }
-  };
-
-  useEffect(() => {
-    // handle ongoing job
-    setIsLoading(true);
-    intervalRef.current = window.setInterval(async () => {
-      const res = await handleGetDoc();
-      if (res?.status === "completed") {
-        // Handle timeout
-        // If doc is created more than 10 minutes and not completed -> cancel interval
-        if (
-          new Date(res.created_at).getTime() + 10 * 60 * 1000 <
-          new Date().getTime()
-        ) {
-          clearInterval(intervalRef.current);
-          setDoc(undefined);
-          setIsLoading(false);
-          return;
-        }
-
-        await queryClient.invalidateQueries({
-          queryKey: ["list-docs"],
-        });
-        clearInterval(intervalRef.current);
-      }
-      setIsLoading(false);
-    }, 3000);
-
-    return () => {
-      clearInterval(intervalRef.current);
-    };
-  }, [needRefetch]);
-
   return (
     <Card className="rounded-none w-full h-screen">
       <CardHeader>
         <CardTitle>ThinkForce (Private Early Access v0.0.1)</CardTitle>
         <CardDescription>
           ThinkForce replicate the process of a human researcher, hence, once
           you pressed the "Generate" button, it will take some time to process
-          the information <b>(5-8 minutes)</b>.
+          the information <b>(around 1 minute)</b>.
           <br />
           <br />
           <i>
@@ -180,11 +138,7 @@ const MainForm: React.FC<MainFormProps> = ({}) => {
               )}
             />
 
-            <Button
-              type="submit"
-              className="mt-4"
-              disabled={isLoading || (doc && doc?.status !== "completed")}
-            >
+            <Button type="submit" className="mt-4" disabled={isLoading}>
               {isLoading ? (
                 <div className="flex gap-2 items-center">
                   Generate <Spinner />
@@ -199,10 +153,15 @@ const MainForm: React.FC<MainFormProps> = ({}) => {
           <Separator />
         </div>
         <div className="flex justify-between items-center">
-          <CardTitle className="">Current jobs</CardTitle>
+          <CardTitle className="">Current generation</CardTitle>
         </div>
         <div className="mt-3">
-          {isLoading && <Spinner />}
+          {isLoading && (
+            <div className="flex gap-1 items-center">
+              <Spinner />
+              <span>Please wait...</span>
+            </div>
+          )}
           {isLoading ? null : doc && <Job doc={doc} />}
         </div>
       </CardContent>

packages/frontend/src/repo/database.types.ts

@@ -11,6 +11,8 @@ export type Database = {
     Tables: {
       doc_meta: {
         Row: {
+          cost: number | null
+          cost_details: Json | null
           created_at: string
           file_name: string
           id: number
@@ -20,6 +22,8 @@ export type Database = {
           user_id: string
         }
         Insert: {
+          cost?: number | null
+          cost_details?: Json | null
           created_at?: string
           file_name: string
           id?: number
@@ -29,6 +33,8 @@ export type Database = {
           user_id: string
         }
         Update: {
+          cost?: number | null
+          cost_details?: Json | null
           created_at?: string
           file_name?: string
           id?: number

packages/frontend/supabase/functions/backend/index.ts

@@ -1,6 +1,10 @@
 import { Hono } from "jsr:@hono/hono";
 import { cors } from "jsr:@hono/hono/cors";
-import { GetObjectCommand, S3Client } from "npm:@aws-sdk/client-s3";
+import {
+  GetObjectCommand,
+  PutObjectCommand,
+  S3Client,
+} from "npm:@aws-sdk/client-s3";
 import { getSignedUrl } from "npm:@aws-sdk/s3-request-presigner";
 import { createClient } from "jsr:@supabase/supabase-js@2";
 
@@ -221,4 +225,100 @@ app.get("/doc", async (c) => {
   return c.json({ signedUrl });
 });
 
+app.post("/v2/gen", async (c) => {
+  const body = await c.req.json();
+  const userId = body.userId;
+  const title = body.title;
+  const refinedTitle = title.replaceAll(" ", "_") + "_" + new Date().getTime();
+  const fileName = `${userId}/${refinedTitle}.md`;
+
+  const supabase = createClient(
+    Deno.env.get("X_SUPABASE_URL") ?? "",
+    Deno.env.get("X_SUPABASE_SERVICE_KEY") ?? "",
+    {
+      global: {
+        headers: { Authorization: c.req.header("Authorization") || "" },
+      },
+    },
+  );
+
+  // Check tokens
+  const preTokenCheck = await supabase.from("gen_usage").select(
+    "tokens",
+  ).eq("user_id", userId).single();
+
+  if (preTokenCheck.error) {
+    throw new Error(preTokenCheck.error?.message ?? "");
+  }
+
+  if ((preTokenCheck!.data!.tokens as unknown as number) <= 0) {
+    throw new Error("Not enough tokens");
+  }
+
+  // Call Storm
+  const baseUrl = new URL(Deno.env.get("AWS_BE_V2_URL") || "");
+  baseUrl.pathname = "/storm";
+  baseUrl.searchParams.set("topic", title);
+  console.log(baseUrl.toString());
+
+  const req = await fetch(baseUrl.toString(), {
+    method: "POST",
+  });
+  const res = await req.text();
+
+  console.log(res);
+
+  // Upload to R2
+  const S3 = new S3Client({
+    region: "auto",
+    endpoint: `https://${
+      Deno.env.get("R2_ACCOUNT_ID") || ""
+    }.r2.cloudflarestorage.com`,
+    credentials: {
+      accessKeyId: Deno.env.get("R2_ACCESS_KEY_ID") || "",
+      secretAccessKey: Deno.env.get("R2_SECRET_KEY") || "",
+    },
+  });
+
+  const putCommand = new PutObjectCommand({
+    Bucket: Deno.env.get("R2_BUCKET_NAME") || "",
+    Key: fileName,
+    Body: res,
+  });
+
+  await S3.send(putCommand);
+
+  // Update doc meta
+  const docMetaReq = await supabase.from("doc_meta").insert({
+    title,
+    file_name: fileName,
+    user_id: userId,
+    run_id: refinedTitle,
+    status: "started",
+  });
+
+  if (docMetaReq.error) {
+    throw new Error(docMetaReq.error?.message ?? "");
+  }
+
+  const tokens = await supabase.from("gen_usage").select(
+    "tokens",
+  ).eq("user_id", userId).single();
+
+  if (tokens.error) {
+    throw new Error(tokens.error?.message ?? "");
+  }
+
+  const deductToken = await supabase.from("gen_usage").update({
+    tokens: (tokens!.data!.tokens as unknown as number) - 10,
+  }).eq("user_id", userId);
+
+  if (deductToken.error) {
+    console.log(deductToken.error);
+    throw new Error("Failed to deduct token");
+  }
+
+  return c.json({ message: "Success" });
+});
+
 Deno.serve(app.fetch);