Spaces:
Running
Running
File size: 2,409 Bytes
5e433de |
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 |
import os
import bcrypt
from src.auth.db import get_db_connection
from src import config
from groq import Groq
def register_user(userid, password, api_key):
conn = get_db_connection()
cursor = conn.cursor()
password_hash = bcrypt.hashpw(password.encode(), bcrypt.gensalt())
try:
cursor.execute('INSERT INTO users (userid, password_hash, api_key) VALUES (?, ?, ?)',
(userid, password_hash, api_key))
conn.commit()
return True, "β
Registered successfully!"
except:
return False, "β User already exists."
finally:
conn.close()
def login_user(userid, password):
conn = get_db_connection()
cursor = conn.cursor()
cursor.execute('SELECT password_hash, api_key FROM users WHERE userid=?', (userid,))
result = cursor.fetchone()
conn.close()
if result:
stored_hash, api_key = result
if bcrypt.checkpw(password.encode(), stored_hash):
return True, api_key
return False, None
def verify_login(userid, password):
# Verify the user's login credentials
success, saved_api_key = login_user(userid, password)
if success:
config.api_key = saved_api_key
os.environ["GROQ_API_KEY"] = saved_api_key
return "β
Login successful!"
else:
return "β Incorrect userid or password."
def register_user_with_api_key(userid, password, user_api_key):
# Validate the API Key first
try:
client = Groq(api_key=user_api_key)
response = client.chat.completions.create(
messages=[{"role": "user", "content": "Hello"}],
model="llama3-70b-8192"
)
# If API key is valid, proceed to register the user
success, msg = register_user(userid, password, user_api_key)
if success:
config.api_key = user_api_key
os.environ["GROQ_API_KEY"] = user_api_key
return "β
API Key validated & registered!"
else:
return msg
except Exception as e:
# API key invalid
return f"β Invalid API Key: {str(e)}"
def handle_login(userid, password, user_api_key):
if user_api_key:
# Handle registration with API key validation
return register_user_with_api_key(userid, password, user_api_key)
else:
# Handle standard login
return verify_login(userid, password)
|