Update app.py

app.py

@@ -120,16 +120,140 @@ def handle_userinput(user_question):
             )
 
 
-def chat(slug):
+def get_user_chat_count(user_id):
+    """
+    Retrieves the chat count for the user from the MySQL database.
+    """
+    try:
+        conn = mysql.connector.connect(
+            user=os.getenv("SQL_USER"),
+            password=os.getenv("SQL_PWD"),
+            host=os.getenv("SQL_HOST"),
+            database="Birdseye_DB",
+        )
+        cursor = conn.cursor()
+
+        cursor.execute("SELECT count FROM Chat WHERE user_id = %s", (user_id,))
+        result = cursor.fetchone()
+        if result:
+            return result[0]
+        else:
+            # Insert a new row for the user if not found
+            cursor.execute(
+                "INSERT INTO Chat (user_id, count) VALUES (%s, %s)", (user_id, 0)
+            )
+            conn.commit()
+            return 0
+    except mysql.connector.Error as err:
+        st.error(f"Error: {err}")
+        return None
+    finally:
+        if conn.is_connected():
+            cursor.close()
+            conn.close()
+
+
+def increment_user_chat_count(user_id):
+    """
+    Increments the chat count for the user in the MySQL database.
+    """
+    try:
+        conn = mysql.connector.connect(
+            user=os.getenv("SQL_USER"),
+            password=os.getenv("SQL_PWD"),
+            host=os.getenv("SQL_HOST"),
+            database="Birdseye_DB",
+        )
+        cursor = conn.cursor()
+
+        cursor.execute(
+            "UPDATE Chat SET count = count + 1 WHERE user_id = %s", (user_id,)
+        )
+        conn.commit()
+    except mysql.connector.Error as err:
+        st.error(f"Error: {err}")
+    finally:
+        if conn.is_connected():
+            cursor.close()
+            conn.close()
+
+
+def is_user_in_unlimited_chat_group(user_id):
+    """
+    Checks if the user belongs to the 'Unlimited Chat' group.
+    """
+    try:
+        conn = mysql.connector.connect(
+            user=os.getenv("SQL_USER"),
+            password=os.getenv("SQL_PWD"),
+            host=os.getenv("SQL_HOST"),
+            database="Birdseye_DB",
+        )
+        cursor = conn.cursor()
+
+        cursor.execute(
+            """
+            SELECT 1
+            FROM auth_user_groups
+            JOIN auth_group ON auth_user_groups.group_id = auth_group.id
+            WHERE auth_user_groups.user_id = %s AND auth_group.name = 'Unlimited Chat'
+        """,
+            (user_id,),
+        )
+        return cursor.fetchone() is not None
+    except mysql.connector.Error as err:
+        st.error(f"Error: {err}")
+        return False
+    finally:
+        if conn.is_connected():
+            cursor.close()
+            conn.close()
+
+
+def chat(slug, user_id):
     """
     Manages the chat interface in the Streamlit application, handling the conversation
     flow and displaying the chat history.
+    Restricts chat based on user group and chat count.
     """
+    try:
+        conn = mysql.connector.connect(
+            user=os.getenv("SQL_USER"),
+            password=os.getenv("SQL_PWD"),
+            host=os.getenv("SQL_HOST"),
+            database="Birdseye_DB",
+        )
+        cursor = conn.cursor()
+
+        # Execute a query
+        cursor.execute("SELECT ocr_text FROM birdseye_temp WHERE slug = %s", (slug,))
+
+        # Fetch the results
+        rows = cursor.fetchall()
+        text = ""
+        for row in rows:
+            if row[0]:
+                text += row[0]
+
+    except mysql.connector.Error as err:
+        st.error(f"Error: {err}")
+        return
+    finally:
+        if conn.is_connected():
+            cursor.close()
+            conn.close()
 
     text_chunks = get_text_chunks(get_pdf_text(slug))
     vectorstore = get_vectorstore(text_chunks)
     st.session_state.conversation = get_conversation_chain(vectorstore)
 
+    # Check if the user can chat
+    if not is_user_in_unlimited_chat_group(user_id):
+        user_chat_count = get_user_chat_count(user_id)
+        if user_chat_count is None or user_chat_count >= 20:
+            st.write("You have reached your chat limit.")
+            return
+
     if len(st.session_state.messages) == 1:
         message = st.session_state.messages[0]
         with st.chat_message(message["role"]):