refactor(api): remove AI worker, add ai api route, add usage tiers

- Remove separate AI worker service - Added generation limits: FREE: 1000/month (For the beta version) PRO: 500/month ENTERPRISE: 1000/month - Integrate AI functionality into main API routes - Added monthly generations reset and usage tier upgrade API routes - Upgrade tier page to be added along with profile page section
2024-11-23 20:31:24 -05:00
parent 91a4a54f24
commit a25097108d
29 changed files with 590 additions and 4049 deletions
--- a/frontend/components/editor/AIChat/lib/chatUtils.ts
+++ b/frontend/components/editor/AIChat/lib/chatUtils.ts
@ -89,7 +89,8 @@ export const handleSend = async (
  setIsGenerating: React.Dispatch<React.SetStateAction<boolean>>,
  setIsLoading: React.Dispatch<React.SetStateAction<boolean>>,
  abortControllerRef: React.MutableRefObject<AbortController | null>,
-  activeFileContent: string
+  activeFileContent: string,
+  isEditMode: boolean = false
 ) => {
  // Return if input is empty and context is null
  if (input.trim() === "" && !context) return
@ -129,17 +130,17 @@ export const handleSend = async (
    }))

    // Fetch AI response for chat message component
-    const response = await fetch(
-      `${process.env.NEXT_PUBLIC_AI_WORKER_URL}/api`,
-      {
-        method: "POST",
-        headers: {
-          "Content-Type": "application/json",
+    const response = await fetch("/api/ai", 
+    {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
        },
        body: JSON.stringify({
          messages: anthropicMessages,
          context: context || undefined,
          activeFileContent: activeFileContent,
+          isEditMode: isEditMode,
        }),
        signal: abortControllerRef.current.signal,
      }
@ -147,7 +148,8 @@ export const handleSend = async (

    // Throw error if response is not ok
    if (!response.ok) {
-      throw new Error("Failed to get AI response")
+      const error = await response.text()
+      throw new Error(error)
    }

    // Get reader for chat message component
@ -197,7 +199,7 @@ export const handleSend = async (
      console.error("Error fetching AI response:", error)
      const errorMessage = {
        role: "assistant" as const,
-        content: "Sorry, I encountered an error. Please try again.",
+        content: error.message || "Sorry, I encountered an error. Please try again.",
      }
      setMessages((prev) => [...prev, errorMessage])
    }
--- a/frontend/components/editor/generate.tsx
+++ b/frontend/components/editor/generate.tsx
@ -5,14 +5,12 @@ import { Editor } from "@monaco-editor/react"
 import { Check, Loader2, RotateCw, Sparkles, X } from "lucide-react"
 import { usePathname, useRouter } from "next/navigation"
 import { useCallback, useEffect, useRef, useState } from "react"
-import { Socket } from "socket.io-client"
 import { toast } from "sonner"
 import { Button } from "../ui/button"
 // import monaco from "monaco-editor"

 export default function GenerateInput({
  user,
-  socket,
  width,
  data,
  editor,
@ -21,7 +19,6 @@ export default function GenerateInput({
  onClose,
 }: {
  user: User
-  socket: Socket
  width: number
  data: {
    fileName: string
@ -59,32 +56,54 @@ export default function GenerateInput({
  }: {
    regenerate?: boolean
  }) => {
-    if (user.generations >= 1000) {
-      toast.error("You reached the maximum # of generations.")
-      return
-    }
+    try {
+      setLoading({ generate: !regenerate, regenerate })
+      setCurrentPrompt(input)

-    setLoading({ generate: !regenerate, regenerate })
-    setCurrentPrompt(input)
-    socket.emit(
-      "generateCode",
-      {
-        fileName: data.fileName,
-        code: data.code,
-        line: data.line,
-        instructions: regenerate ? currentPrompt : input,
-      },
-      (res: { response: string; success: boolean }) => {
-        console.log("Generated code", res.response, res.success)
-        // if (!res.success) {
-        //   toast.error("Failed to generate code.");
-        //   return;
-        // }
+      const response = await fetch("/api/ai", {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify({
+          messages: [{
+            role: "user",
+            content: regenerate ? currentPrompt : input
+          }],
+          context: null,
+          activeFileContent: data.code,
+          isEditMode: true,
+          fileName: data.fileName,
+          line: data.line
+        }),
+      })

-        setCode(res.response)
-        router.refresh()
+      if (!response.ok) {
+        const error = await response.text()
+        toast.error(error)
+        return
      }
-    )
+
+      const reader = response.body?.getReader()
+      const decoder = new TextDecoder()
+      let result = ""
+
+      if (reader) {
+        while (true) {
+          const { done, value } = await reader.read()
+          if (done) break
+          result += decoder.decode(value, { stream: true })
+        }
+      }
+
+      setCode(result.trim())
+      router.refresh()
+    } catch (error) {
+      console.error("Generation error:", error)
+      toast.error("Failed to generate code")
+    } finally {
+      setLoading({ generate: false, regenerate: false })
+    }
  }
  const handleGenerateForm = useCallback(
    (e: React.FormEvent<HTMLFormElement>) => {
--- a/frontend/components/editor/index.tsx
+++ b/frontend/components/editor/index.tsx
@ -949,7 +949,6 @@ export default function CodeEditor({
          {generate.show ? (
            <GenerateInput
              user={userData}
-              socket={socket!}
              width={generate.width - 90}
              data={{
                fileName: tabs.find((t) => t.id === activeFileId)?.name ?? "",
--- a/frontend/components/ui/userButton.tsx
+++ b/frontend/components/ui/userButton.tsx
@ -9,22 +9,78 @@ import {
 } from "@/components/ui/dropdown-menu"
 import { User } from "@/lib/types"
 import { useClerk } from "@clerk/nextjs"
-import { LogOut, Sparkles } from "lucide-react"
+import { Crown, LogOut, Sparkles } from "lucide-react"
 import { useRouter } from "next/navigation"
+import { useEffect, useState } from "react"
 import Avatar from "./avatar"
+import { Button } from "./button"
+import { TIERS } from "@/lib/tiers"

-export default function UserButton({ userData }: { userData: User }) {
-  if (!userData) return null
+// TODO: Remove this once we have a proper tier system
+const TIER_INFO = {
+  FREE: {
+    color: "text-gray-500",
+    icon: Sparkles,
+    limit: TIERS.FREE.generations,
+  },
+  PRO: {
+    color: "text-blue-500",
+    icon: Crown,
+    limit: TIERS.PRO.generations,
+  },
+  ENTERPRISE: {
+    color: "text-purple-500",
+    icon: Crown,
+    limit: TIERS.ENTERPRISE.generations,
+  },
+} as const

+export default function UserButton({ userData: initialUserData }: { userData: User }) {
+  const [userData, setUserData] = useState<User>(initialUserData)
+  const [isOpen, setIsOpen] = useState(false)
  const { signOut } = useClerk()
  const router = useRouter()

+  const fetchUserData = async () => {
+    try {
+      const res = await fetch(
+        `${process.env.NEXT_PUBLIC_DATABASE_WORKER_URL}/api/user?id=${userData.id}`,
+        {
+          headers: {
+            Authorization: `${process.env.NEXT_PUBLIC_WORKERS_KEY}`,
+          },
+          cache: 'no-store'
+        }
+      )
+      if (res.ok) {
+        const updatedUserData = await res.json()
+        setUserData(updatedUserData)
+      }
+    } catch (error) {
+      console.error("Failed to fetch user data:", error)
+    }
+  }
+
+  useEffect(() => {
+    if (isOpen) {
+      fetchUserData()
+    }
+  }, [isOpen])
+
+  const tierInfo = TIER_INFO[userData.tier as keyof typeof TIER_INFO] || TIER_INFO.FREE
+  const TierIcon = tierInfo.icon
+  const usagePercentage = Math.floor((userData.generations || 0) * 100 / tierInfo.limit)
+
+  const handleUpgrade = async () => {
+    router.push('/upgrade')
+  }
+
  return (
-    <DropdownMenu>
+    <DropdownMenu open={isOpen} onOpenChange={setIsOpen}>
      <DropdownMenuTrigger>
        <Avatar name={userData.name} avatarUrl={userData.avatarUrl} />
      </DropdownMenuTrigger>
-      <DropdownMenuContent className="w-48" align="end">
+      <DropdownMenuContent className="w-64" align="end">
        <div className="py-1.5 px-2 w-full">
          <div className="font-medium">{userData.name}</div>
          <div className="text-sm w-full overflow-hidden text-ellipsis whitespace-nowrap text-muted-foreground">
@ -33,20 +89,48 @@ export default function UserButton({ userData }: { userData: User }) {
        </div>

        <DropdownMenuSeparator />
-        <div className="py-1.5 px-2 w-full flex flex-col items-start text-sm">
-          <div className="flex items-center">
-            <Sparkles className={`h-4 w-4 mr-2 text-indigo-500`} />
-            AI Usage: {userData.generations}/1000
+
+        <div className="py-1.5 px-2 w-full">
+          <div className="flex items-center justify-between">
+            <div className="flex items-center gap-1.5">
+              <TierIcon className={`h-4 w-4 ${tierInfo.color}`} />
+              <span className="text-sm font-medium">{userData.tier || "FREE"} Plan</span>
+            </div>
+            {/* {(userData.tier === "FREE" || userData.tier === "PRO") && (
+              <Button
+                variant="ghost"
+                size="sm"
+                className="h-7 text-xs border-b hover:border-b-foreground"
+                onClick={handleUpgrade}
+              >
+                Upgrade
+              </Button>
+            )} */}
          </div>
-          <div className="rounded-full w-full mt-2 h-2 overflow-hidden bg-secondary">
+        </div>
+
+        <DropdownMenuSeparator />
+
+        <div className="py-1.5 px-2 w-full">
+          <div className="flex items-center justify-between text-sm text-muted-foreground mb-2">
+            <span>AI Usage</span>
+            <span>{userData.generations}/{tierInfo.limit}</span>
+          </div>
+
+          <div className="rounded-full w-full h-2 overflow-hidden bg-secondary">
            <div
-              className="h-full bg-indigo-500 rounded-full"
+              className={`h-full rounded-full transition-all duration-300 ${
+                usagePercentage > 90 ? 'bg-red-500' : 
+                usagePercentage > 75 ? 'bg-yellow-500' : 
+                tierInfo.color.replace('text-', 'bg-')
+              }`}
              style={{
-                width: `${(userData.generations * 100) / 1000}%`,
+                width: `${Math.min(usagePercentage, 100)}%`,
              }}
            />
          </div>
        </div>
+
        <DropdownMenuSeparator />

        {/* <DropdownMenuItem className="cursor-pointer">
@ -64,3 +148,4 @@ export default function UserButton({ userData }: { userData: User }) {
    </DropdownMenu>
  )
 }
+