Update app.py

app.py

@@ -7,21 +7,23 @@ from oauth2client.service_account import ServiceAccountCredentials
 import PyPDF2
 import io
 
-# Set up OpenAI client
-client = OpenAI(api_key=st.secrets["OPENAI_API_KEY"])
-
 # Constants
 WORD_LIMIT = 8000
+DAILY_API_LIMIT = 5  # Set your desired limit per user per day
 
-# Google Sheets setup remains the same
+# Set up OpenAI client
+client = OpenAI(api_key=st.secrets["OPENAI_API_KEY"])
+
+# Google Sheets setup
 scope = ["https://spreadsheets.google.com/feeds", "https://www.googleapis.com/auth/drive"]
 creds = ServiceAccountCredentials.from_json_keyfile_name("genexam-2c8c645ecc0d.json", scope)
 client_gs = gspread.authorize(creds)
 sheet = client_gs.open("GeneXam user").sheet1
 
 def check_user_in_sheet(username):
+    """Check if user exists in sheet"""
     try:
-        users_list = sheet.col_values(1)
+        users_list = sheet.col_values(1)  # UserID column
         if username in users_list:
             return True
         return False
@@ -29,15 +31,59 @@ def check_user_in_sheet(username):
         st.error(f"Error checking user: {str(e)}")
         return False
 
+def get_user_stats(username):
+    """Get user's current API usage statistics"""
+    try:
+        users_list = sheet.col_values(1)  # UserID column
+        row_number = users_list.index(username) + 1
+        
+        daily_count = int(sheet.cell(row_number, 2).value)    # DailyAPICount
+        total_count = int(sheet.cell(row_number, 3).value)    # TotalAPICount
+        last_used = sheet.cell(row_number, 4).value           # LastUsedDate
+        
+        return {
+            'daily_count': daily_count,
+            'total_count': total_count,
+            'last_used': last_used
+        }
+    except Exception as e:
+        st.error(f"Error getting user stats: {str(e)}")
+        return None
+
 def update_api_usage(username):
+    """Update both daily and total API usage counts"""
     try:
         users_list = sheet.col_values(1)
         row_number = users_list.index(username) + 1
-        api_usage = int(sheet.cell(row_number, 2).value)
-        api_usage += 1
-        sheet.update_cell(row_number, 2, api_usage)
+        today = datetime.now().strftime('%Y-%m-%d')
+        
+        # Get current values
+        stats = get_user_stats(username)
+        if not stats:
+            return False, "Error retrieving user statistics"
+        
+        # Reset daily count if it's a new day
+        daily_count = stats['daily_count']
+        if stats['last_used'] != today:
+            daily_count = 0
+        
+        # Check daily limit
+        if daily_count >= DAILY_API_LIMIT:
+            return False, f"You have reached your daily limit of {DAILY_API_LIMIT} generations. Please try again tomorrow."
+        
+        # Update counts
+        new_daily_count = daily_count + 1
+        new_total_count = stats['total_count'] + 1
+        
+        # Update all values in sheet
+        sheet.update_cell(row_number, 2, new_daily_count)  # Update DailyAPICount
+        sheet.update_cell(row_number, 3, new_total_count)  # Update TotalAPICount
+        sheet.update_cell(row_number, 4, today)            # Update LastUsedDate
+        
+        return True, None
+        
     except Exception as e:
-        st.error(f"Error updating API usage: {str(e)}")
+        return False, f"Error updating API usage: {str(e)}"
 
 def extract_text_from_pdf(pdf_file):
     """Simple PDF text extraction with word limit check"""
@@ -55,6 +101,14 @@ def extract_text_from_pdf(pdf_file):
         return None, f"Error processing PDF: {str(e)}"
 
 def generate_questions_with_retry(knowledge_material, question_type, cognitive_level, extra_instructions, case_based, num_choices=None, max_retries=3):
+    """Generate questions and update API usage"""
+    
+    # Check and update API usage before generating
+    can_generate, error_message = update_api_usage(username)
+    if not can_generate:
+        st.error(error_message)
+        return None
+    
     # Adjust number of questions based on type
     if question_type == "Multiple Choice":
         num_questions = 3
@@ -171,8 +225,16 @@ if 'username' not in st.session_state:
         if username_input:
             if check_user_in_sheet(username_input):
                 st.session_state['username'] = username_input
-                st.success(f"Welcome, {username_input}!")
-                update_api_usage(username_input)
+                
+                # Show user stats on login
+                stats = get_user_stats(username_input)
+                if stats:
+                    st.success(f"Welcome, {username_input}!")
+                    st.info(f"""
+                    📊 Your API Usage Statistics:
+                    - Today's Usage: {stats['daily_count']}/{DAILY_API_LIMIT} generations
+                    - Total All-Time Usage: {stats['total_count']} generations
+                    """)
             else:
                 st.warning("Username not found. Please try again.")
         else:
@@ -180,6 +242,16 @@ if 'username' not in st.session_state:
 else:
     st.title(f"Welcome, {st.session_state['username']}! Generate your exam questions")
     
+    # Show current usage stats
+    stats = get_user_stats(st.session_state['username'])
+    if stats:
+        remaining = DAILY_API_LIMIT - stats['daily_count']
+        st.info(f"""
+        📊 Usage Statistics:
+        - Daily Generations Remaining: {remaining}/{DAILY_API_LIMIT}
+        - Total All-Time Generations: {stats['total_count']}
+        """)
+    
     # Create tabs for input methods
     tab1, tab2 = st.tabs(["Text Input", "PDF Upload"])
     
@@ -226,8 +298,8 @@ else:
     if st.button("Generate Questions"):
         if 'knowledge_material' in locals() and knowledge_material.strip():
             with st.spinner("Generating questions..."):
-                # Your existing generate_questions_with_retry function call here
                 questions = generate_questions_with_retry(
+                    st.session_state['username'],
                     knowledge_material,
                     question_type,
                     cognitive_level,
@@ -240,6 +312,16 @@ else:
                     st.write("### Generated Exam Questions:")
                     st.write(questions)
                     
+                    # Update displayed stats after generation
+                    new_stats = get_user_stats(st.session_state['username'])
+                    if new_stats:
+                        remaining = DAILY_API_LIMIT - new_stats['daily_count']
+                        st.info(f"""
+                        📊 Updated Usage Statistics:
+                        - Daily Generations Remaining: {remaining}/{DAILY_API_LIMIT}
+                        - Total All-Time Generations: {new_stats['total_count']}
+                        """)
+                    
                     # Download button
                     st.download_button(
                         label="Download Questions",