Upload app.py

app.py

@@ -1,15 +1,11 @@
 import os
 import logging
-import numpy as np
-from typing import List, Optional, Tuple
+from typing import List, Tuple
 import torch
 import gradio as gr
-import spaces
 from sentence_transformers import SentenceTransformer
 from langchain_community.vectorstores import FAISS
 from langchain.embeddings.base import Embeddings
-from gradio_client import Client
-import requests
 from tqdm import tqdm
 
 # Configuration
@@ -17,7 +13,7 @@ QWEN_API_URL = "Qwen/Qwen2.5-Max-Demo"  # Gradio API for Qwen2.5 chat
 CHUNK_SIZE = 800
 TOP_K_RESULTS = 150
 SIMILARITY_THRESHOLD = 0.4
-PASSWORD_HASH = "abc12345"  # Replace with hashed password in production
+PASSWORD_HASH = os.getenv("PASSWORD_HASH", "default_password")  # Use environment variable for password
 
 BASE_SYSTEM_PROMPT = """
 Répondez en français selon ces règles :
@@ -83,30 +79,29 @@ def split_text_into_chunks(text: str) -> List[str]:
     
     return chunks
 
-def create_new_database(file_content: str, db_name: str, password: str, progress=gr.Progress()) -> str:
+def create_new_database(file_content: str, db_name: str, password: str, progress=gr.Progress()) -> Tuple[str, List[str]]:
     """Create a new FAISS database from uploaded file"""
     if password != PASSWORD_HASH:
-        return "Incorrect password. Database creation failed."
+        return "Incorrect password. Database creation failed.", []
 
     if not file_content.strip():
-        return "Uploaded file is empty. Database creation failed."
+        return "Uploaded file is empty. Database creation failed.", []
 
     if not db_name.isalnum():
-        return "Database name must be alphanumeric. Database creation failed."
+        return "Database name must be alphanumeric. Database creation failed.", []
 
     try:
-        # Define file names for the FAISS database
         faiss_file = f"{db_name}-index.faiss"
         pkl_file = f"{db_name}-index.pkl"
-        
+
         # Check if the database already exists
         if os.path.exists(faiss_file) or os.path.exists(pkl_file):
-            return f"Database '{db_name}' already exists."
+            return f"Database '{db_name}' already exists.", []
 
         # Initialize embeddings and split text
         chunks = split_text_into_chunks(file_content)
         if not chunks:
-            return "No valid chunks generated. Database creation failed."
+            return "No valid chunks generated. Database creation failed.", []
 
         logging.info(f"Creating {len(chunks)} chunks...")
         progress(0, desc="Starting embedding process...")
@@ -122,39 +117,34 @@ def create_new_database(file_content: str, db_name: str, password: str, progress
             text_embeddings=list(zip(chunks, embeddings_list)),
             embedding=embeddings
         )
-
-        # Save the FAISS database to the root directory
         vector_store.save_local(".")
         logging.info(f"FAISS database saved to: {faiss_file} and {pkl_file}")
 
-        # Rename the default FAISS files to match the desired naming convention
-        os.rename("index.faiss", faiss_file)
-        os.rename("index.pkl", pkl_file)
-
         # Verify files were created
         if not os.path.exists(faiss_file) or not os.path.exists(pkl_file):
-            return f"Failed to save FAISS database files: {faiss_file} or {pkl_file}"
-        logging.info(f"FAISS database files: {faiss_file}, {pkl_file}")
+            return f"Failed to save FAISS database files.", []
+        logging.info(f"FAISS database files created: {faiss_file}, {pkl_file}")
+
+        # Update the list of available databases
+        db_list = [os.path.splitext(f)[0].replace("-index", "") for f in os.listdir(".") if f.endswith(".faiss")]
+        return f"Database '{db_name}' created successfully.", db_list
 
-        return f"Database '{db_name}' created successfully."
     except Exception as e:
         logging.error(f"Database creation failed: {str(e)}")
-        return f"Error creating database: {str(e)}"
+        return f"Error creating database: {str(e)}", []
 
-def generate_response(user_input: str, db_name: str) -> Optional[str]:
+def generate_response(user_input: str, db_name: str) -> str:
     """Generate response using Qwen2.5 MAX"""
     try:
         if not db_name:
             return "Please select a database to chat with."
 
-        # Define file names for the FAISS database
         faiss_file = f"{db_name}-index.faiss"
         pkl_file = f"{db_name}-index.pkl"
 
         if not os.path.exists(faiss_file) or not os.path.exists(pkl_file):
             return f"Database '{db_name}' does not exist."
 
-        # Load the FAISS database
         vector_store = FAISS.load_local(".", embeddings, allow_dangerous_deserialization=True)
 
         # Contextual search
@@ -200,7 +190,7 @@ def generate_response(user_input: str, db_name: str) -> Optional[str]:
 
     except Exception as e:
         logging.error(f"Generation error: {str(e)}", exc_info=True)
-        return None
+        return "Erreur de génération - Veuillez réessayer."
 
 # Initialize models and vector store
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -216,11 +206,7 @@ with gr.Blocks() as app:
 
     def update_db_list():
         """Update the list of available databases"""
-        return [
-            name.replace("-index.faiss", "")  # Remove "-index.faiss" suffix for display
-            for name in os.listdir(".")
-            if name.endswith("-index.faiss")
-        ]
+        return [os.path.splitext(f)[0].replace("-index", "") for f in os.listdir(".") if f.endswith(".faiss")]
 
     with gr.Tab("Create Database"):
         gr.Markdown("## Create a New FAISS Database")
@@ -232,7 +218,7 @@ with gr.Blocks() as app:
         
         def handle_create(file, db_name, password, progress=gr.Progress()):
             if not file or not db_name or not password:
-                return "Please provide all required inputs."
+                return "Please provide all required inputs.", []
             
             # Check if the file is valid
             if isinstance(file, str):  # Gradio provides the file path as a string
@@ -240,15 +226,12 @@ with gr.Blocks() as app:
                     with open(file, "r", encoding="utf-8") as f:
                         file_content = f.read()
                 except Exception as e:
-                    return f"Error reading file: {str(e)}"
+                    return f"Error reading file: {str(e)}", []
             else:
-                return "Invalid file format. Please upload a .txt file."
+                return "Invalid file format. Please upload a .txt file.", []
             
-            result = create_new_database(file_content, db_name, password, progress)
-            if "created successfully" in result:
-                # Update the database list
-                return result, update_db_list()
-            return result, None
+            result, db_list = create_new_database(file_content, db_name, password, progress)
+            return result, db_list
         
         create_button.click(
             handle_create,
@@ -267,8 +250,8 @@ with gr.Blocks() as app:
             if not db_name:
                 return "", history + [("System", "Please select a database to chat with.")]
             response = generate_response(message, db_name)
-            return "", history + [(message, response or "Erreur de génération - Veuillez réessayer.")]
-        
+            return "", history + [(message, response)]
+
         msg.submit(
             chat_response,
             inputs=[msg, db_select, chatbot],
@@ -287,10 +270,4 @@ with gr.Blocks() as app:
         )
 
 if __name__ == "__main__":
-    # Log existing databases at startup
-    logging.info("Existing databases:")
-    for name in os.listdir("."):
-        if name.endswith("-index.faiss"):
-            logging.info(f"- {name}")
-
     app.launch(server_name="0.0.0.0", server_port=7860)