Update app.py

app.py

@@ -1,6 +1,7 @@
+import os
 import streamlit as st
 import requests
-import streamlit_cookies_manager
+from streamlit_cookies_manager import EncryptedCookieManager
 from openai import OpenAI
 import google.generativeai as genai
 
@@ -13,7 +14,7 @@ GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
 # -----------------------------------------------------
 # Initialize OpenAI & Gemini
 # -----------------------------------------------------
-client = OpenAI(api_key=OPENAI_API_KEY)  # OpenAI client
+client = OpenAI(api_key=OPENAI_API_KEY)
 genai.configure(api_key=GEMINI_API_KEY)
 gemini_model = genai.GenerativeModel("gemini-pro")
 
@@ -25,29 +26,28 @@ st.title("🤖 AI Model Comparator")
 st.subheader("Compare responses across multiple LLMs.")
 
 # -----------------------------------------------------
-# Cookie Manager for Persistent Request Limits
+# Secure Cookie Manager for Request Limits
 # -----------------------------------------------------
-cookies = streamlit_cookies_manager.CookieManager()
+cookies = EncryptedCookieManager(
+    prefix="ai-model-comparator/",
+    password=os.environ.get("COOKIES_PASSWORD", "MySecretPassword")  # Use a strong secret in deployment
+)
 
-def get_request_count():
-    """Retrieve the request count from cookies."""
-    if "request_count" in cookies:
-        return int(cookies["request_count"])
-    return 0
+if not cookies.ready():
+    st.stop()  # Wait until cookies are available
 
-def update_request_count():
-    """Increment the request count and save it to cookies."""
-    count = get_request_count() + 1
-    cookies["request_count"] = str(count)
-    cookies.save()
+# Initialize request count if not set
+if "request_count" not in cookies:
+    cookies["request_count"] = 0
+
+request_count = int(cookies["request_count"])
 
 # -----------------------------------------------------
 # Check Request Limit (Max 3 Requests Per User)
 # -----------------------------------------------------
-request_count = get_request_count()
 if request_count >= 3:
     st.warning("⚠️ Request limit reached (3 per user). Please wait before trying again.")
-    st.stop()  # Stop execution if limit is reached
+    st.stop()
 
 # -----------------------------------------------------
 # Sidebar: Model Settings and Future Works Section
@@ -107,7 +107,9 @@ def get_gemini_response(prompt):
 # Generate Responses when Button is Clicked
 # -----------------------------------------------------
 if st.button("🚀 Generate Responses"):
-    update_request_count()  # Increment the user's request count
+    request_count += 1
+    cookies["request_count"] = request_count  # Update cookie count
+    cookies.save()  # Persist changes
 
     with st.spinner("Fetching responses..."):
         openai_text = get_openai_response(user_prompt)