Upload app.py

app.py

@@ -84,21 +84,40 @@ 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 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]]:
     """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."
+            return f"Database '{db_name}' already exists.", []
 
         # Create the database directory
         os.makedirs(db_path, exist_ok=True)
@@ -107,7 +126,7 @@ def create_new_database(file_content: str, db_name: str, password: str, progress
         # 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...")
@@ -128,13 +147,16 @@ def create_new_database(file_content: str, db_name: str, password: str, progress
 
         # Verify files were created
         if not os.listdir(db_path):
-            return f"Failed to save FAISS database files in: {db_path}"
+            return f"Failed to save FAISS database files in: {db_path}", []
         logging.info(f"FAISS database files: {os.listdir(db_path)}")
 
-        return f"Database '{db_name}' created successfully."
+        # 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
+
     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"""
@@ -212,7 +234,7 @@ with gr.Blocks() as app:
     def update_db_list():
         """Update the list of available databases"""
         if not os.path.exists(DATABASE_DIR):
-            return []
+            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))]
 
     with gr.Tab("Create Database"):
@@ -225,7 +247,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
@@ -233,15 +255,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,
@@ -280,10 +299,4 @@ with gr.Blocks() as app:
         )
 
 if __name__ == "__main__":
-    # Ensure DATABASE_DIR exists
-    if not os.path.exists(DATABASE_DIR):
-        logging.info(f"Creating persistent DATABASE_DIR: {DATABASE_DIR}")
-        os.makedirs(DATABASE_DIR, exist_ok=True)
-        logging.info(f"Contents of DATABASE_DIR: {os.listdir(DATABASE_DIR)}")
-
     app.launch(server_name="0.0.0.0", server_port=7860)