Update app.py

app.py

@@ -1,6 +1,11 @@
 import os
+import uuid
+import json
 import gradio as gr
 from huggingface_hub import InferenceClient
+import firebase_admin
+from firebase_admin import credentials, firestore
+import hashlib
 
 class XylariaChat:
     def __init__(self):
@@ -15,6 +20,11 @@ class XylariaChat:
             api_key=self.hf_token
         )
         
+        # Firebase initialization
+        self.firebase_cred = None
+        self.firebase_db = None
+        self.current_user_id = None
+        
         # Initialize conversation history and persistent memory
         self.conversation_history = []
         self.persistent_memory = {}
@@ -41,9 +51,153 @@ Capabilities:
 - Provide explanations and insights
 - Offer supportive and constructive guidance """
 
+    def generate_user_key(self):
+        """
+        Generate a unique user key that can be saved locally
+        """
+        # Generate a UUID and hash it for additional security
+        unique_key = str(uuid.uuid4())
+        hashed_key = hashlib.sha256(unique_key.encode()).hexdigest()
+        return hashed_key
+
+    def init_firebase(self):
+        """
+        Initialize Firebase using the service account from environment variable
+        """
+        try:
+            # Get Firebase service account from environment variable
+            firebase_cred_json = os.getenv("FIREBASE_SERVICE_ACCOUNT")
+            if not firebase_cred_json:
+                raise ValueError("Firebase service account not found in environment variables")
+            
+            # Parse the JSON credentials
+            cred_dict = json.loads(firebase_cred_json)
+            
+            # Write credentials to a temporary file (Firebase requirement)
+            with open('firebase_credentials.json', 'w') as f:
+                json.dump(cred_dict, f)
+            
+            # Initialize Firebase Admin SDK
+            cred = credentials.Certificate('firebase_credentials.json')
+            firebase_admin.initialize_app(cred)
+            
+            # Initialize Firestore
+            self.firebase_db = firestore.client()
+            
+            return True
+        except Exception as e:
+            print(f"Firebase initialization error: {e}")
+            return False
+
+    def save_user_chat_history(self):
+        """
+        Save chat history to Firebase for the current user
+        """
+        if not self.firebase_db or not self.current_user_id:
+            return False
+        
+        try:
+            # Create a document reference for the user
+            user_doc_ref = self.firebase_db.collection('user_chats').document(self.current_user_id)
+            
+            # Prepare chat history data
+            chat_data = {
+                'conversation_history': self.conversation_history,
+                'persistent_memory': self.persistent_memory,
+                'timestamp': firestore.SERVER_TIMESTAMP
+            }
+            
+            # Save or update the document
+            user_doc_ref.set(chat_data, merge=True)
+            return True
+        except Exception as e:
+            print(f"Error saving chat history: {e}")
+            return False
+
+    def load_user_chat_history(self, user_id):
+        """
+        Load chat history from Firebase for a specific user
+        """
+        if not self.firebase_db:
+            return False
+        
+        try:
+            # Retrieve user document
+            user_doc_ref = self.firebase_db.collection('user_chats').document(user_id)
+            user_doc = user_doc_ref.get()
+            
+            if user_doc.exists:
+                user_data = user_doc.to_dict()
+                # Restore conversation history and persistent memory
+                self.conversation_history = user_data.get('conversation_history', [])
+                self.persistent_memory = user_data.get('persistent_memory', {})
+                self.current_user_id = user_id
+                return True
+            return False
+        except Exception as e:
+            print(f"Error loading chat history: {e}")
+            return False
+
+    def authenticate_user(self, user_key):
+        """
+        Authenticate user based on their unique key
+        """
+        if not self.firebase_db:
+            return False, "Firebase not initialized"
+        
+        try:
+            # Check if user key exists in users collection
+            users_ref = self.firebase_db.collection('users')
+            query = users_ref.where('user_key', '==', user_key).limit(1)
+            docs = query.stream()
+            
+            for doc in docs:
+                # User found, set current user ID and load history
+                self.current_user_id = doc.id
+                self.load_user_chat_history(self.current_user_id)
+                return True, "Authentication successful"
+            
+            return False, "Invalid user key"
+        except Exception as e:
+            print(f"Authentication error: {e}")
+            return False, "Authentication error"
+
+    def register_user(self, user_key):
+        """
+        Register a new user with a unique key
+        """
+        if not self.firebase_db:
+            return False, "Firebase not initialized"
+        
+        try:
+            # Check if user key already exists
+            users_ref = self.firebase_db.collection('users')
+            query = users_ref.where('user_key', '==', user_key).limit(1)
+            existing_users = list(query.stream())
+            
+            if existing_users:
+                return False, "User key already exists"
+            
+            # Create new user document
+            new_user_ref = users_ref.document()
+            new_user_ref.set({
+                'user_key': user_key,
+                'created_at': firestore.SERVER_TIMESTAMP
+            })
+            
+            # Set current user ID
+            self.current_user_id = new_user_ref.id
+            return True, "User registered successfully"
+        except Exception as e:
+            print(f"Registration error: {e}")
+            return False, "Registration error"
+
     def store_information(self, key, value):
         """Store important information in persistent memory"""
         self.persistent_memory[key] = value
+        # Automatically save to Firebase if user is authenticated
+        if self.current_user_id:
+            self.save_user_chat_history()
 
     def retrieve_information(self, key):
         """Retrieve information from persistent memory"""
@@ -52,10 +206,13 @@ Capabilities:
     def reset_conversation(self):
         """
         Completely reset the conversation history and persistent memory
-        This helps prevent exposing previous users' conversations
         """
         self.conversation_history = []
         self.persistent_memory = {}
+        
+        # If user is authenticated, update Firebase
+        if self.current_user_id:
+            self.save_user_chat_history()
 
     def get_response(self, user_input):
         # Prepare messages with conversation context and persistent memory
@@ -88,7 +245,28 @@ Capabilities:
             return f"Error generating response: {str(e)}"
 
     def create_interface(self):
+        # Initialize Firebase
+        firebase_initialized = self.init_firebase()
+
+        def generate_key():
+            """Generate a new user key"""
+            return self.generate_user_key()
+
+        def authenticate(user_key):
+            """Authenticate user and load chat history"""
+            success, message = self.authenticate_user(user_key)
+            return message
+
+        def register(user_key):
+            """Register a new user"""
+            success, message = self.register_user(user_key)
+            return message
+
         def streaming_response(message, chat_history):
+            # Check if user is authenticated
+            if not self.current_user_id:
+                return "", chat_history + [[message, "Please authenticate first."]]
+            
             # Clear input textbox
             response_stream = self.get_response(message)
             
@@ -121,6 +299,9 @@ Capabilities:
             # Limit conversation history to prevent token overflow
             if len(self.conversation_history) > 10:
                 self.conversation_history = self.conversation_history[-10:]
+            
+            # Save chat history to Firebase
+            self.save_user_chat_history()
 
         # Custom CSS for Inter font
         custom_css = """
@@ -142,6 +323,41 @@ Capabilities:
         """
 
         with gr.Blocks(theme='soft', css=custom_css) as demo:
+            # User Authentication Section
+            with gr.Column():
+                # Generate Key Section
+                with gr.Row():
+                    generate_key_btn = gr.Button("Generate User Key")
+                    user_key_output = gr.Textbox(label="Your User Key", interactive=False)
+                generate_key_btn.click(
+                    fn=generate_key, 
+                    inputs=None, 
+                    outputs=[user_key_output]
+                )
+
+                # Authentication Section
+                with gr.Row():
+                    auth_key_input = gr.Textbox(label="Enter User Key")
+                    auth_btn = gr.Button("Authenticate")
+                    register_btn = gr.Button("Register")
+                
+                # Authentication result
+                auth_result = gr.Textbox(label="Authentication Result", interactive=False)
+                
+                # Authenticate button click
+                auth_btn.click(
+                    fn=authenticate, 
+                    inputs=[auth_key_input], 
+                    outputs=[auth_result]
+                )
+                
+                # Register button click
+                register_btn.click(
+                    fn=register, 
+                    inputs=[auth_key_input], 
+                    outputs=[auth_result]
+                )
+
             # Chat interface with improved styling
             with gr.Column():
                 chatbot = gr.Chatbot(