mirror of
https://github.com/StanGirard/quivr.git
synced 2024-09-11 21:17:36 +03:00
feat: setup premium feature backend (#1467)
Issue: https://github.com/StanGirard/quivr/issues/1468
This commit is contained in:
parent
cee050a2e4
commit
c47548d3cd
@ -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
|
||||
|
@ -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):
|
||||
"""
|
||||
|
@ -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
|
||||
|
@ -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"],
|
||||
}
|
||||
|
||||
|
||||
|
@ -12,4 +12,5 @@ export interface UserStats {
|
||||
requests_stats: RequestStat[];
|
||||
date: string;
|
||||
models: string[];
|
||||
is_premium: boolean;
|
||||
}
|
||||
|
@ -22,4 +22,5 @@ WHERE NOT EXISTS (
|
||||
SELECT 1 FROM migrations WHERE name = '202309307004032_change_user_settings'
|
||||
);
|
||||
|
||||
COMMIT;
|
||||
COMMIT;
|
||||
|
||||
|
86
scripts/20231023140000_add_stripe_wrapper.sql
Normal file
86
scripts/20231023140000_add_stripe_wrapper.sql
Normal file
@ -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;
|
14
scripts/20231023160000_copy_auth_users_to_public_users.sql
Normal file
14
scripts/20231023160000_copy_auth_users_to_public_users.sql
Normal file
@ -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;
|
27
scripts/new.sql
Normal file
27
scripts/new.sql
Normal file
@ -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;
|
||||
|
||||
|
@ -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'
|
||||
);
|
||||
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user