feat: setup premium feature backend (#1467)

Issue: https://github.com/StanGirard/quivr/issues/1468

.backend_env.example

@@ -21,3 +21,8 @@ RESEND_CONTACT_SALES_FROM=contact_sales@resend.dev
 RESEND_CONTACT_SALES_TO=<change-me>
 
 CRAWL_DEPTH=1
+
+
+PREMIUM_MAX_BRAIN_NUMBER=30
+PREMIUM_MAX_BRAIN_SIZE=10000000
+PREMIUM_DAILY_CHAT_CREDIT=100

backend/models/databases/supabase/user_usage.py

@@ -1,3 +1,4 @@
+import os
 from datetime import datetime
 from uuid import UUID
 
@@ -25,31 +26,74 @@ class UserUsage(Repository):
             .execute()
         )
 
+    def check_if_is_premium_user(self, user_id: UUID):
+        """
+        Check if the user is a premium user
+        """
+        try:
+            user_email_customer = (
+                self.db.from_("users")
+                .select("*")
+                .filter("id", "eq", str(user_id))
+                .execute()
+            ).data
+
+            if len(user_email_customer) == 0:
+                return False
+
+            matching_customers = (
+                self.db.table("customers")
+                .select("email")
+                .filter("email", "eq", user_email_customer[0]["email"])
+                .execute()
+            ).data
+        except Exception as e:
+            logger.error("Error while checking if user is a premium user")
+            logger.error(e)
+            return False
+
+        return len(matching_customers) > 0
+
     def get_user_settings(self, user_id):
         """
         Fetch the user settings from the database
         """
-        response = (
+
+        user_settings_response = (
             self.db.from_("user_settings")
             .select("*")
             .filter("user_id", "eq", str(user_id))
             .execute()
         ).data
 
-        if len(response) == 0:
+        if len(user_settings_response) == 0:
             # Create the user settings
-            result = (
+            user_settings_response = (
                 self.db.table("user_settings")
                 .insert({"user_id": str(user_id)})
                 .execute()
+            ).data
+
+        if len(user_settings_response) == 0:
+            raise ValueError("User settings could not be created")
+
+        user_settings = user_settings_response[0]
+        user_settings["is_premium"] = False
+        is_premium_user = self.check_if_is_premium_user(user_id)
+
+        if is_premium_user:
+            user_settings["is_premium"] = True
+            user_settings["max_brains"] = int(
+                os.environ.get("PREMIUM_MAX_BRAIN_NUMBER", 30)
             )
-            if result:
-                return self.get_user_settings(user_id)
-            else:
-                raise ValueError("User settings could not be created")
-        if response and len(response) > 0:
-            return response[0]
-        return None
+            user_settings["max_brain_size"] = int(
+                os.environ.get("PREMIUM_MAX_BRAIN_SIZE", 10000000)
+            )
+            user_settings["daily_chat_credit"] = int(
+                os.environ.get("PREMIUM_DAILY_CHAT_CREDIT", 100)
+            )
+
+        return user_settings
 
     def get_user_usage(self, user_id):
         """

backend/requirements.txt

@@ -7,7 +7,7 @@ GitPython==3.1.36
 pdf2image==1.16.3
 pypdf==3.9.0
 StrEnum==0.4.15
-supabase==1.0.3
+supabase==1.1.0
 tiktoken==0.4.0
 unstructured==0.6.7
 fastapi==0.95.2

backend/routes/user_routes.py

@@ -29,19 +29,19 @@ async def get_user_endpoint(
     information about the user's API usage.
     """
 
-    userDailyUsage = UserUsage(
+    user_daily_usage = UserUsage(
         id=current_user.id,
         email=current_user.email,
         openai_api_key=current_user.openai_api_key,
     )
-    userSettings = userDailyUsage.get_user_settings()
-    max_brain_size = userSettings.get("max_brain_size", 1000000000)
+    user_settings = user_daily_usage.get_user_settings()
+    max_brain_size = user_settings.get("max_brain_size", 1000000000)
 
     date = time.strftime("%Y%m%d")
-    daily_chat_credit = userSettings.get("daily_chat_credit", 10)
+    daily_chat_credit = user_settings.get("daily_chat_credit", 10)
 
-    userDailyUsage = UserUsage(id=current_user.id)
-    requests_stats = userDailyUsage.get_user_usage()
+    user_daily_usage = UserUsage(id=current_user.id)
+    requests_stats = user_daily_usage.get_user_usage()
     default_brain = get_user_default_brain(current_user.id)
 
     if default_brain:
@@ -55,9 +55,10 @@ async def get_user_endpoint(
         "current_brain_size": defaul_brain_size,
         "daily_chat_credit": daily_chat_credit,
         "requests_stats": requests_stats,
-        "models": userSettings.get("models", []),
+        "models": user_settings.get("models", []),
         "date": date,
         "id": current_user.id,
+        "is_premium": user_settings["is_premium"],
     }
 
 

frontend/lib/types/User.ts

@@ -12,4 +12,5 @@ export interface UserStats {
   requests_stats: RequestStat[];
   date: string;
   models: string[];
+  is_premium: boolean;
 }

scripts/202309307004032_change_user_settings.sql

@@ -22,4 +22,5 @@ WHERE NOT EXISTS (
     SELECT 1 FROM migrations WHERE name = '202309307004032_change_user_settings'
 );
 
-COMMIT;
+COMMIT;
+

scripts/20231023140000_add_stripe_wrapper.sql

@@ -0,0 +1,86 @@
+-- Create extension 'wrappers' if it doesn't exist
+CREATE EXTENSION IF NOT EXISTS wrappers;
+
+-- Create foreign data wrapper 'stripe_wrapper' if it doesn't exist
+DO $$ 
+BEGIN
+  IF NOT EXISTS (
+    SELECT 1 
+    FROM information_schema.foreign_data_wrappers 
+    WHERE foreign_data_wrapper_name = 'stripe_wrapper'
+  ) THEN
+    CREATE FOREIGN DATA WRAPPER stripe_wrapper
+      HANDLER stripe_fdw_handler;
+  END IF;
+END $$;
+
+-- Check if the server 'stripe_server' exists before creating it
+DO $$ 
+BEGIN
+  IF NOT EXISTS (SELECT 1 FROM pg_foreign_server WHERE srvname = 'stripe_server') THEN
+    CREATE SERVER stripe_server
+      FOREIGN DATA WRAPPER stripe_wrapper
+      OPTIONS (
+        api_key 'your_stripe_api_key'  -- Replace with your Stripe API key
+      );
+  END IF;
+END $$;
+
+-- Create foreign table 'public.customers' if it doesn't exist
+DO $$ 
+BEGIN
+  IF NOT EXISTS (
+    SELECT 1 
+    FROM information_schema.tables 
+    WHERE table_name = 'customers' 
+  ) THEN
+    CREATE FOREIGN TABLE public.customers (
+      id text,
+      email text,
+      name text,
+      description text,
+      created timestamp,
+      attrs jsonb
+    )
+    SERVER stripe_server
+    OPTIONS (
+      OBJECT 'customers',
+      ROWID_COLUMN 'id'
+    );
+  END IF;
+END $$;
+
+-- Create table 'users' if it doesn't exist
+CREATE TABLE IF NOT EXISTS public.users (
+  id uuid REFERENCES auth.users NOT NULL PRIMARY KEY,
+  email text
+);
+
+-- Create or replace function 'public.handle_new_user' 
+CREATE OR REPLACE FUNCTION public.handle_new_user() 
+RETURNS TRIGGER AS $$
+BEGIN
+  INSERT INTO public.users (id, email)
+  VALUES (NEW.id, NEW.email);
+  RETURN NEW;
+END;
+$$ LANGUAGE plpgsql SECURITY DEFINER;
+
+-- Check if the trigger 'on_auth_user_created' exists before creating it
+DO $$ 
+BEGIN
+  IF NOT EXISTS (SELECT 1 FROM pg_trigger WHERE tgname = 'on_auth_user_created') THEN
+    CREATE TRIGGER on_auth_user_created
+      AFTER INSERT ON auth.users
+      FOR EACH ROW EXECUTE FUNCTION public.handle_new_user();
+  END IF;
+END $$;
+
+-- Update migrations table
+INSERT INTO migrations (name) 
+SELECT '20231023140000_add_stripe_wrapper'
+WHERE NOT EXISTS (
+    SELECT 1 FROM migrations WHERE name = '20231023140000_add_stripe_wrapper'
+);
+
+COMMIT;

scripts/20231023160000_copy_auth_users_to_public_users.sql

@@ -0,0 +1,14 @@
+-- Auth users to public users -- 
+
+INSERT INTO public.users (id, email)
+SELECT id, email
+FROM auth.users
+ON CONFLICT (id) DO NOTHING;
+
+INSERT INTO migrations (name) 
+SELECT '20231023160000_copy_auth_users_to_public_users'
+WHERE NOT EXISTS (
+    SELECT 1 FROM migrations WHERE name = '20231023160000_copy_auth_users_to_public_users'
+);
+
+COMMIT;

scripts/new.sql

@@ -0,0 +1,27 @@
+DO $$
+BEGIN
+    -- Check if the column max_requests_number exists
+    IF EXISTS (
+        SELECT 1 
+        FROM information_schema.columns 
+        WHERE table_name='user_settings' AND column_name='max_requests_number'
+    ) THEN
+        -- Rename the column
+        ALTER TABLE user_settings RENAME COLUMN max_requests_number TO daily_chat_credit;
+    END IF;
+    
+    -- Modify default values
+    ALTER TABLE user_settings ALTER COLUMN daily_chat_credit SET DEFAULT 20;
+    ALTER TABLE user_settings ALTER COLUMN max_brains SET DEFAULT 3;
+END $$;
+
+-- Update migrations table
+INSERT INTO migrations (name) 
+SELECT '202309307004032_change_user_settings'
+WHERE NOT EXISTS (
+    SELECT 1 FROM migrations WHERE name = '202309307004032_change_user_settings'
+);
+
+COMMIT;
+
+

scripts/tables.sql

@@ -287,6 +287,86 @@ CREATE TABLE IF NOT EXISTS onboardings (
   PRIMARY KEY (user_id)
 );
 
+
+-- Stripe settings --
+-- Create extension 'wrappers' if it doesn't exist
+CREATE EXTENSION IF NOT EXISTS wrappers;
+
+-- Create foreign data wrapper 'stripe_wrapper' if it doesn't exist
+DO $$ 
+BEGIN
+  IF NOT EXISTS (
+    SELECT 1 
+    FROM information_schema.foreign_data_wrappers 
+    WHERE foreign_data_wrapper_name = 'stripe_wrapper'
+  ) THEN
+    CREATE FOREIGN DATA WRAPPER stripe_wrapper
+      HANDLER stripe_fdw_handler;
+  END IF;
+END $$;
+
+-- Check if the server 'stripe_server' exists before creating it
+DO $$ 
+BEGIN
+  IF NOT EXISTS (SELECT 1 FROM pg_foreign_server WHERE srvname = 'stripe_server') THEN
+    CREATE SERVER stripe_server
+      FOREIGN DATA WRAPPER stripe_wrapper
+      OPTIONS (
+        api_key 'your_stripe_api_key'  -- Replace with your Stripe API key
+      );
+  END IF;
+END $$;
+
+-- Create foreign table 'public.customers' if it doesn't exist
+DO $$ 
+BEGIN
+  IF NOT EXISTS (
+    SELECT 1 
+    FROM information_schema.tables 
+    WHERE table_name = 'customers' 
+  ) THEN
+    CREATE FOREIGN TABLE public.customers (
+      id text,
+      email text,
+      name text,
+      description text,
+      created timestamp,
+      attrs jsonb
+    )
+    SERVER stripe_server
+    OPTIONS (
+      OBJECT 'customers',
+      ROWID_COLUMN 'id'
+    );
+  END IF;
+END $$;
+
+-- Create table 'users' if it doesn't exist
+CREATE TABLE IF NOT EXISTS public.users (
+  id uuid REFERENCES auth.users NOT NULL PRIMARY KEY,
+  email text
+);
+
+-- Create or replace function 'public.handle_new_user' 
+CREATE OR REPLACE FUNCTION public.handle_new_user() 
+RETURNS TRIGGER AS $$
+BEGIN
+  INSERT INTO public.users (id, email)
+  VALUES (NEW.id, NEW.email);
+  RETURN NEW;
+END;
+$$ LANGUAGE plpgsql SECURITY DEFINER;
+
+-- Check if the trigger 'on_auth_user_created' exists before creating it
+DO $$ 
+BEGIN
+  IF NOT EXISTS (SELECT 1 FROM pg_trigger WHERE tgname = 'on_auth_user_created') THEN
+    CREATE TRIGGER on_auth_user_created
+      AFTER INSERT ON auth.users
+      FOR EACH ROW EXECUTE FUNCTION public.handle_new_user();
+  END IF;
+END $$;
+
 insert into
   storage.buckets (id, name)
 values
@@ -301,9 +381,9 @@ CREATE POLICY "Access Quivr Storage 1jccrwz_2" ON storage.objects FOR UPDATE TO
 CREATE POLICY "Access Quivr Storage 1jccrwz_3" ON storage.objects FOR DELETE TO anon USING (bucket_id = 'quivr');
 
 INSERT INTO migrations (name) 
-SELECT '20231012150000_add_creation_time_to_onboardings_table'
+SELECT '20231023160000_copy_auth_users_to_public_users'
 WHERE NOT EXISTS (
-    SELECT 1 FROM migrations WHERE name = '20231012150000_add_creation_time_to_onboardings_table'
+    SELECT 1 FROM migrations WHERE name = '20231023160000_copy_auth_users_to_public_users'
 );