Upload app.py

app.py

@@ -13,7 +13,6 @@ import requests
 from tqdm import tqdm
 
 # Configuration
-DATABASE_DIR = "semantic_memory"  # Persistent directory for FAISS databases
 QWEN_API_URL = "Qwen/Qwen2.5-Max-Demo"  # Gradio API for Qwen2.5 chat
 CHUNK_SIZE = 800
 TOP_K_RESULTS = 150
@@ -84,49 +83,30 @@ def split_text_into_chunks(text: str) -> List[str]:
     
     return chunks
 
-def initialize_vector_store(embeddings: Embeddings, db_name: str) -> FAISS:
-    """Initialize or load a FAISS vector store"""
-    db_path = os.path.join(DATABASE_DIR, db_name)
-    if os.path.exists(db_path):
-        try:
-            logging.info(f"Loading existing database: {db_name}")
-            return FAISS.load_local(
-                db_path,
-                embeddings,
-                allow_dangerous_deserialization=True
-            )
-        except Exception as e:
-            logging.error(f"FAISS load error: {str(e)}")
-            raise
-
-    logging.info(f"Creating new vector database: {db_name}")
-    os.makedirs(db_path, exist_ok=True)
-    return None
-
-def create_new_database(file_content: str, db_name: str, password: str, progress=gr.Progress()) -> Tuple[str, List[str]]:
+def create_new_database(file_content: str, db_name: str, password: str, progress=gr.Progress()) -> 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:
-        db_path = os.path.join(DATABASE_DIR, db_name)
-        if os.path.exists(db_path):
-            return f"Database '{db_name}' already exists.", []
-
-        # Create the database directory
-        os.makedirs(db_path, exist_ok=True)
-        logging.info(f"Created database directory: {db_path}")
+        # 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."
 
         # 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...")
@@ -142,21 +122,24 @@ def create_new_database(file_content: str, db_name: str, password: str, progress
             text_embeddings=list(zip(chunks, embeddings_list)),
             embedding=embeddings
         )
-        vector_store.save_local(db_path)
-        logging.info(f"FAISS database saved to: {db_path}")
 
-        # Verify files were created
-        if not os.listdir(db_path):
-            return f"Failed to save FAISS database files in: {db_path}", []
-        logging.info(f"FAISS database files: {os.listdir(db_path)}")
+        # Save the FAISS database to the root directory
+        vector_store.save_local(".")
+        logging.info(f"FAISS database saved to: {faiss_file} and {pkl_file}")
 
-        # Update the list of available databases
-        db_list = [name for name in os.listdir(DATABASE_DIR) if os.path.isdir(os.path.join(DATABASE_DIR, name))]
-        return f"Database '{db_name}' created successfully.", db_list
+        # 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"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]:
     """Generate response using Qwen2.5 MAX"""
@@ -164,15 +147,15 @@ def generate_response(user_input: str, db_name: str) -> Optional[str]:
         if not db_name:
             return "Please select a database to chat with."
 
-        db_path = os.path.join(DATABASE_DIR, db_name)
-        if not os.path.exists(db_path):
+        # 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."
 
-        vector_store = FAISS.load_local(
-            db_path,
-            embeddings,
-            allow_dangerous_deserialization=True
-        )
+        # Load the FAISS database
+        vector_store = FAISS.load_local(".", embeddings, allow_dangerous_deserialization=True)
 
         # Contextual search
         docs_scores = vector_store.similarity_search_with_score(
@@ -233,9 +216,11 @@ with gr.Blocks() as app:
 
     def update_db_list():
         """Update the list of available databases"""
-        if not os.path.exists(DATABASE_DIR):
-            os.makedirs(DATABASE_DIR, exist_ok=True)
-        return [name for name in os.listdir(DATABASE_DIR) if os.path.isdir(os.path.join(DATABASE_DIR, name))]
+        return [
+            name.replace("-index.faiss", "")  # Remove "-index.faiss" suffix for display
+            for name in os.listdir(".")
+            if name.endswith("-index.faiss")
+        ]
 
     with gr.Tab("Create Database"):
         gr.Markdown("## Create a New FAISS Database")
@@ -247,7 +232,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
@@ -255,12 +240,15 @@ 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, db_list = create_new_database(file_content, db_name, password, progress)
-            return result, db_list
+            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
         
         create_button.click(
             handle_create,
@@ -299,4 +287,10 @@ 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)