Spaces:
Sleeping
Sleeping
File size: 4,220 Bytes
1b2c3a0 0517f4f 1b2c3a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | import sqlite3
import os
import time
DB_PATH = "vqa.db"
def init_db():
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
# Users table
cursor.execute('''
CREATE TABLE IF NOT EXISTS users (
username TEXT PRIMARY KEY,
password TEXT
)
''')
# Logs table
cursor.execute('''
CREATE TABLE IF NOT EXISTS logs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
timestamp TEXT,
username TEXT,
model TEXT,
question TEXT
)
''')
# Gemini Usage table
cursor.execute('''
CREATE TABLE IF NOT EXISTS gemini_usage (
username TEXT,
date TEXT,
call_count INTEGER,
PRIMARY KEY (username, date)
)
''')
# Seed Admin
cursor.execute("INSERT OR IGNORE INTO users (username, password) VALUES (?, ?)", ("vignesh", "vignesh"))
conn.commit()
conn.close()
def login_user(username, password):
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
cursor.execute("SELECT password FROM users WHERE username=?", (username,))
row = cursor.fetchone()
conn.close()
if row and row[0] == password:
return True
return False
def signup_user(username, password):
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
try:
cursor.execute("INSERT INTO users (username, password) VALUES (?, ?)", (username, password))
conn.commit()
success = True
except sqlite3.IntegrityError:
success = False
conn.close()
return success
def log_request_db(username, model, question):
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
from datetime import datetime
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
cursor.execute("INSERT INTO logs (timestamp, username, model, question) VALUES (?, ?, ?, ?)",
(timestamp, username, model, question))
conn.commit()
conn.close()
def check_gemini_limit(username):
if username == "vignesh":
return True # Admin bypass
from datetime import datetime
today_str = datetime.now().strftime("%Y-%m-%d")
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
cursor.execute("SELECT call_count FROM gemini_usage WHERE username=? AND date=?", (username, today_str))
row = cursor.fetchone()
if row is None:
count = 0
else:
count = row[0]
conn.close()
if count >= 10:
return False
return True
def increment_gemini_limit(username):
if username == "vignesh":
return
from datetime import datetime
today_str = datetime.now().strftime("%Y-%m-%d")
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
cursor.execute("SELECT call_count FROM gemini_usage WHERE username=? AND date=?", (username, today_str))
row = cursor.fetchone()
if row is None:
cursor.execute("INSERT INTO gemini_usage (username, date, call_count) VALUES (?, ?, ?)", (username, today_str, 1))
else:
cursor.execute("UPDATE gemini_usage SET call_count = call_count + 1 WHERE username=? AND date=?", (username, today_str))
conn.commit()
conn.close()
def login_or_create_google_user(username):
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
# Check if user exists
cursor.execute("SELECT username FROM users WHERE username=?", (username,))
row = cursor.fetchone()
if not row:
# Create user with a dummy password since Google handles auth
cursor.execute("INSERT INTO users (username, password) VALUES (?, ?)", (username, "GOOGLE_AUTH_USER"))
conn.commit()
conn.close()
return True
def get_user_logs(username):
conn = sqlite3.connect(DB_PATH)
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
if username == "vignesh":
cursor.execute("SELECT * FROM logs ORDER BY id DESC")
else:
cursor.execute("SELECT * FROM logs WHERE username=? ORDER BY id DESC", (username,))
rows = cursor.fetchall()
conn.close()
return [dict(row) for row in rows]
|