Update app.py

app.py

@@ -7,10 +7,16 @@ import pdfplumber
 from docx import Document as DocxDocument
 import json
 from datetime import datetime
+import hashlib
 
 # Global variables
 chat_engine = None
 conversation_history = []
+current_api_key_hash = None
+
+# Function to hash API key for identification
+def hash_api_key(api_key):
+    return hashlib.sha256(api_key.encode()).hexdigest()
 
 # Function to read PDF files
 def read_pdf(file_path):
@@ -30,7 +36,7 @@ def read_docx(file_path):
 
 # Function to load and index documents
 def load_data(files, api_key):
-    global chat_engine
+    global chat_engine, current_api_key_hash
     
     if not files or not api_key:
         return "Please provide both API key and files to proceed."
@@ -48,6 +54,9 @@ def load_data(files, api_key):
         # Set OpenAI API key
         os.environ["OPENAI_API_KEY"] = api_key
         
+        # Store hashed API key
+        current_api_key_hash = hash_api_key(api_key)
+        
         Settings.llm = OpenAI(
             model="gpt-3.5-turbo", 
             temperature=0.5,
@@ -64,11 +73,15 @@ def load_data(files, api_key):
 
 # Function to handle chat
 def chat_with_docs(message, history, api_key):
-    global chat_engine, conversation_history
+    global chat_engine, conversation_history, current_api_key_hash
     
     if not api_key:
         return history + [{"role": "assistant", "content": "Please enter your OpenAI API key first."}]
     
+    # Set current API key hash if not set
+    if current_api_key_hash is None:
+        current_api_key_hash = hash_api_key(api_key)
+    
     if chat_engine is None:
         return history + [
             {"role": "user", "content": message},
@@ -91,41 +104,54 @@ def chat_with_docs(message, history, api_key):
         ]
 
 # Function to save conversation
-def save_conversation():
-    global conversation_history
+def save_conversation(api_key):
+    global conversation_history, current_api_key_hash
+    
+    if not api_key:
+        return "Please enter your OpenAI API key first.", gr.update()
     
     if not conversation_history:
         return "No conversation to save.", gr.update()
     
     try:
+        # Set current API key hash if not set
+        if current_api_key_hash is None:
+            current_api_key_hash = hash_api_key(api_key)
+        
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         with open("conversations.json", "a") as f:
             conv_data = {
                 "timestamp": timestamp,
+                "api_key_hash": current_api_key_hash,
                 "messages": conversation_history
             }
             json.dump(conv_data, f)
             f.write("\n")
         
         # Update the dropdown choices
-        choices = get_conversation_choices()
+        choices = get_conversation_choices(api_key)
         return "Conversation saved successfully!", gr.update(choices=choices)
     except Exception as e:
         return f"Error saving conversation: {str(e)}", gr.update()
 
 # Function to get conversation choices for dropdown
-def get_conversation_choices():
+def get_conversation_choices(api_key):
+    if not api_key:
+        return []
+    
     if os.path.exists("conversations.json"):
         try:
+            api_key_hash = hash_api_key(api_key)
             with open("conversations.json", "r") as f:
                 conversations = [json.loads(line) for line in f]
             
             choices = []
             for i, conv in enumerate(conversations):
-                timestamp = conv.get("timestamp", "Unknown time")
-                # Format: "1. 2024-01-13 14:30:45"
-                formatted_time = datetime.strptime(timestamp, "%Y%m%d_%H%M%S").strftime("%Y-%m-%d %H:%M:%S")
-                choices.append(f"{i + 1}. {formatted_time}")
+                # Only show conversations for this API key
+                if conv.get("api_key_hash") == api_key_hash:
+                    timestamp = conv.get("timestamp", "Unknown time")
+                    formatted_time = datetime.strptime(timestamp, "%Y%m%d_%H%M%S").strftime("%Y-%m-%d %H:%M:%S")
+                    choices.append(f"{i + 1}. {formatted_time}")
             
             return choices
         except Exception as e:
@@ -133,36 +159,51 @@ def get_conversation_choices():
     return []
 
 # Function to load previous conversations
-def load_conversations():
+def load_conversations(api_key):
+    if not api_key:
+        return "Please enter your OpenAI API key to view conversations."
+    
     if os.path.exists("conversations.json"):
         try:
+            api_key_hash = hash_api_key(api_key)
             with open("conversations.json", "r") as f:
                 conversations = [json.loads(line) for line in f]
             
             conv_text = ""
+            conv_count = 0
             for i, conv in enumerate(conversations):
-                conv_text += f"\n{'='*50}\nConversation {i + 1}\n{'='*50}\n"
-                timestamp = conv.get("timestamp", "Unknown time")
-                formatted_time = datetime.strptime(timestamp, "%Y%m%d_%H%M%S").strftime("%Y-%m-%d %H:%M:%S")
-                conv_text += f"Timestamp: {formatted_time}\n\n"
-                
-                messages = conv.get("messages", conv)
-                for message in messages:
-                    role = message.get('role', 'unknown')
-                    content = message.get('content', '')
-                    conv_text += f"{role.upper()}: {content}\n\n"
+                # Only show conversations for this API key
+                if conv.get("api_key_hash") == api_key_hash:
+                    conv_count += 1
+                    conv_text += f"\n{'='*50}\nConversation {i + 1}\n{'='*50}\n"
+                    timestamp = conv.get("timestamp", "Unknown time")
+                    formatted_time = datetime.strptime(timestamp, "%Y%m%d_%H%M%S").strftime("%Y-%m-%d %H:%M:%S")
+                    conv_text += f"Timestamp: {formatted_time}\n\n"
+                    
+                    messages = conv.get("messages", conv)
+                    for message in messages:
+                        role = message.get('role', 'unknown')
+                        content = message.get('content', '')
+                        conv_text += f"{role.upper()}: {content}\n\n"
+            
+            if conv_count == 0:
+                return "No conversations found for your API key."
             
-            return conv_text if conv_text else "No previous conversations found."
+            return conv_text
         except Exception as e:
             return f"Error loading conversations: {str(e)}"
     return "No previous conversations found."
 
 # Function to view selected conversation
-def view_selected_conversation(selection):
+def view_selected_conversation(selection, api_key):
+    if not api_key:
+        return "Please enter your OpenAI API key first."
+    
     if not selection or selection == "Select a conversation":
         return "Please select a conversation to view."
     
     try:
+        api_key_hash = hash_api_key(api_key)
         # Extract conversation number from selection
         conv_number = int(selection.split(".")[0]) - 1
         
@@ -172,6 +213,11 @@ def view_selected_conversation(selection):
             
             if conv_number < len(conversations):
                 conv = conversations[conv_number]
+                
+                # Check if this conversation belongs to the current API key
+                if conv.get("api_key_hash") != api_key_hash:
+                    return "This conversation does not belong to your API key."
+                
                 conv_text = f"{'='*50}\n"
                 timestamp = conv.get("timestamp", "Unknown time")
                 formatted_time = datetime.strptime(timestamp, "%Y%m%d_%H%M%S").strftime("%Y-%m-%d %H:%M:%S")
@@ -191,11 +237,15 @@ def view_selected_conversation(selection):
         return f"Error viewing conversation: {str(e)}"
 
 # Function to delete selected conversation
-def delete_selected_conversation(selection):
+def delete_selected_conversation(selection, api_key):
+    if not api_key:
+        return "Please enter your OpenAI API key first.", gr.update(), "", ""
+    
     if not selection or selection == "Select a conversation":
-        return "Please select a conversation to delete.", gr.update(), ""
+        return "Please select a conversation to delete.", gr.update(), "", ""
     
     try:
+        api_key_hash = hash_api_key(api_key)
         # Extract conversation number from selection
         conv_number = int(selection.split(".")[0]) - 1
         
@@ -204,6 +254,12 @@ def delete_selected_conversation(selection):
                 conversations = [json.loads(line) for line in f]
             
             if conv_number < len(conversations):
+                conv = conversations[conv_number]
+                
+                # Check if this conversation belongs to the current API key
+                if conv.get("api_key_hash") != api_key_hash:
+                    return "You can only delete your own conversations.", gr.update(), "", ""
+                
                 # Remove the selected conversation
                 del conversations[conv_number]
                 
@@ -214,18 +270,18 @@ def delete_selected_conversation(selection):
                         f.write("\n")
                 
                 # Update dropdown choices
-                choices = get_conversation_choices()
+                choices = get_conversation_choices(api_key)
                 
                 # Get updated all conversations display
-                all_convs = load_conversations()
+                all_convs = load_conversations(api_key)
                 
-                return f"Conversation {conv_number + 1} deleted successfully!", gr.update(choices=choices, value="Select a conversation"), all_convs
+                return f"Conversation {conv_number + 1} deleted successfully!", gr.update(choices=choices, value="Select a conversation"), all_convs, ""
             else:
-                return "Conversation not found.", gr.update(), ""
+                return "Conversation not found.", gr.update(), "", ""
         
-        return "No conversations file found.", gr.update(), ""
+        return "No conversations file found.", gr.update(), "", ""
     except Exception as e:
-        return f"Error deleting conversation: {str(e)}", gr.update(), ""
+        return f"Error deleting conversation: {str(e)}", gr.update(), "", ""
 
 # Function to clear current conversation
 def clear_conversation():
@@ -233,25 +289,54 @@ def clear_conversation():
     conversation_history = []
     return []
 
-# Function to delete all conversations
-def delete_all_conversations():
+# Function to delete all conversations for current API key
+def delete_all_conversations(api_key):
+    if not api_key:
+        return "Please enter your OpenAI API key first.", gr.update(), "", ""
+    
     try:
+        api_key_hash = hash_api_key(api_key)
+        
         if os.path.exists("conversations.json"):
-            os.remove("conversations.json")
-            return "All conversations deleted successfully!", gr.update(choices=[], value="Select a conversation"), "No previous conversations found."
-        return "No conversations to delete.", gr.update(), "No previous conversations found."
+            with open("conversations.json", "r") as f:
+                conversations = [json.loads(line) for line in f]
+            
+            # Keep only conversations that don't belong to this API key
+            remaining_conversations = [conv for conv in conversations if conv.get("api_key_hash") != api_key_hash]
+            
+            deleted_count = len(conversations) - len(remaining_conversations)
+            
+            if deleted_count == 0:
+                return "No conversations to delete for your API key.", gr.update(), "No previous conversations found.", ""
+            
+            # Write back remaining conversations
+            with open("conversations.json", "w") as f:
+                for conv in remaining_conversations:
+                    json.dump(conv, f)
+                    f.write("\n")
+            
+            # If no conversations left, remove the file
+            if len(remaining_conversations) == 0:
+                os.remove("conversations.json")
+            
+            return f"All {deleted_count} conversation(s) deleted successfully!", gr.update(choices=[], value="Select a conversation"), "No previous conversations found.", ""
+        
+        return "No conversations to delete.", gr.update(), "No previous conversations found.", ""
     except Exception as e:
-        return f"Error deleting conversations: {str(e)}", gr.update(), ""
+        return f"Error deleting conversations: {str(e)}", gr.update(), "", ""
 
 # Function to refresh dropdown
-def refresh_dropdown():
-    choices = get_conversation_choices()
+def refresh_dropdown(api_key):
+    if not api_key:
+        return gr.update(choices=[])
+    choices = get_conversation_choices(api_key)
     return gr.update(choices=choices)
 
 # Create Gradio interface
-with gr.Blocks(title="Chat with Documents 💬 📚") as demo:
+with gr.Blocks(title="Chat with Documents 💬 📚", theme=gr.themes.Ocean()) as demo:
     gr.Markdown("# Chat with Documents 💬 📚")
     gr.Markdown("Upload PDF or DOCX files and chat with them using AI!")
+    gr.Markdown("**Note:** Your conversations are private and tied to your API key. Only you can view and manage your conversations.")
     
     with gr.Row():
         with gr.Column(scale=2):
@@ -297,6 +382,7 @@ with gr.Blocks(title="Chat with Documents 💬 📚") as demo:
         
         with gr.Column(scale=1):
             gr.Markdown("### Conversation Management")
+            gr.Markdown("_Only your conversations are shown below_")
             
             with gr.Row():
                 load_convs_btn = gr.Button("Load All Conversations")
@@ -312,7 +398,7 @@ with gr.Blocks(title="Chat with Documents 💬 📚") as demo:
             
             conversation_dropdown = gr.Dropdown(
                 label="Select Conversation",
-                choices=get_conversation_choices(),
+                choices=[],
                 value="Select a conversation",
                 interactive=True
             )
@@ -326,7 +412,7 @@ with gr.Blocks(title="Chat with Documents 💬 📚") as demo:
             
             with gr.Row():
                 delete_selected_btn = gr.Button("Delete Selected", variant="stop")
-                delete_all_btn = gr.Button("Delete All", variant="stop")
+                delete_all_btn = gr.Button("Delete All Mine", variant="stop")
             
             delete_status = gr.Textbox(label="Delete Status", interactive=False)
     
@@ -356,34 +442,38 @@ with gr.Blocks(title="Chat with Documents 💬 📚") as demo:
     
     save_btn.click(
         fn=save_conversation,
+        inputs=api_key_input,
         outputs=[save_status, conversation_dropdown]
     )
     
     load_convs_btn.click(
         fn=load_conversations,
+        inputs=api_key_input,
         outputs=convs_display
     )
     
     refresh_btn.click(
         fn=refresh_dropdown,
+        inputs=api_key_input,
         outputs=conversation_dropdown
     )
     
     view_btn.click(
         fn=view_selected_conversation,
-        inputs=conversation_dropdown,
+        inputs=[conversation_dropdown, api_key_input],
         outputs=selected_conv_display
     )
     
     delete_selected_btn.click(
         fn=delete_selected_conversation,
-        inputs=conversation_dropdown,
-        outputs=[delete_status, conversation_dropdown, convs_display]
+        inputs=[conversation_dropdown, api_key_input],
+        outputs=[delete_status, conversation_dropdown, convs_display, selected_conv_display]
     )
     
     delete_all_btn.click(
         fn=delete_all_conversations,
-        outputs=[delete_status, conversation_dropdown, convs_display]
+        inputs=api_key_input,
+        outputs=[delete_status, conversation_dropdown, convs_display, selected_conv_display]
     )
 
 if __name__ == "__main__":