changes made to manage session specific files and db for deletion of data

app.py

@@ -3,6 +3,7 @@ from flask import Flask,request, jsonify,session
 from flask import render_template
 from flask_session import Session
 from werkzeug.utils import secure_filename
+from apscheduler.schedulers.background import BackgroundScheduler
 
 from langchain.document_loaders import DirectoryLoader,PyPDFLoader,UnstructuredWordDocumentLoader,TextLoader,UnstructuredHTMLLoader,UnstructuredMarkdownLoader
 
@@ -10,12 +11,18 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 from langchain.embeddings.huggingface import HuggingFaceInferenceAPIEmbeddings
 from langchain_chroma import Chroma
+from chromadb import Client
 
 import google.generativeai as genai
 
 from dotenv import load_dotenv
 import os
 import shutil
+import logging
+
+logging.basicConfig(filename='log.txt',filemode='a', level=logging.DEBUG,
+                    format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger()
 
 load_dotenv()
 
@@ -23,6 +30,7 @@ HF_TOKEN = os.getenv('HF_TOKEN')
 GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY')
 CHROMA_PATH = "chroma"
 UPLOAD_FOLDER = "uploads"
+SESSION_TIMEOUT = 30 * 60  # 30 minutes
 
 # Initialize Hugging Face embedding
 hugging_face_ef = HuggingFaceInferenceAPIEmbeddings(
@@ -32,26 +40,21 @@ hugging_face_ef = HuggingFaceInferenceAPIEmbeddings(
 genai.configure(api_key=GOOGLE_API_KEY)
 llm_model = genai.GenerativeModel("gemini-1.5-flash")
 
-if not os.path.exists(UPLOAD_FOLDER):
-    os.makedirs(UPLOAD_FOLDER)
-    
 app = Flask(__name__)
 
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
+    
+
 app.config["SESSION_PERMANENT"] = True
 app.config["SESSION_TYPE"] = "filesystem"
 
 Session(app)
 
-def add_to_chroma(chunks,embedding_function,CHROMA_PATH=CHROMA_PATH):
-    # Load the existing database.
-    db = Chroma(
-        persist_directory=CHROMA_PATH, embedding_function=embedding_function
-    )
-
-    # Add or Update the documents.
-    db.add_documents(chunks)
+# Initialize ChromaDB client
+client = Client()
+active_sessions = dict()
 
-    return True
 
 def remove_file_from_chroma(file_id,CHROMA_PATH=CHROMA_PATH):
     db = Chroma(persist_directory=CHROMA_PATH)
@@ -61,90 +64,108 @@ def remove_file_from_chroma(file_id,CHROMA_PATH=CHROMA_PATH):
     return True
     
         
-def add_file_to_chroma(file_path,file_id):
-    # get extension of file
+def add_file_to_chroma(file_path, file_id, session_id):
+    """Add file chunks to ChromaDB."""
     extension = file_path.split(".")[-1]
-    match extension:
-        case "pdf":
-            loader = PyPDFLoader(file_path)
-        case "docx":
-            loader = UnstructuredWordDocumentLoader(file_path)
-        case "txt":
-            loader = TextLoader(file_path)
-        case "html":
-            loader = UnstructuredHTMLLoader(file_path)
-        case "md":
-            loader = UnstructuredMarkdownLoader(file_path)
-        case _:
-            raise ValueError(f"Unsupported file type: {extension}")
+    loader_map = {
+        "pdf": PyPDFLoader,
+        "docx": UnstructuredWordDocumentLoader,
+        "txt": TextLoader,
+        "html": UnstructuredHTMLLoader,
+        "md": UnstructuredMarkdownLoader,
+    }
+    if extension not in loader_map:
+        raise ValueError(f"Unsupported file type: {extension}")
+    
+    loader = loader_map[extension](file_path)
     documents = loader.load()
-    text_splitter = RecursiveCharacterTextSplitter(    chunk_size=1000,
-    chunk_overlap=100,
-    length_function=len,
-    add_start_index=True)
-    texts = text_splitter.split_documents(documents)
-    # add metadata file_id to documents
-    for text in texts:
-        text.metadata["file_id"] = file_id
     
-    add_to_chroma(texts,hugging_face_ef)
-    return True
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1000,
+        chunk_overlap=100,
+        length_function=len,
+        add_start_index=True
+    )
+    texts = text_splitter.split_documents(documents)
     
+    # Add metadata
+    for text in texts:
+        text.metadata.update({"file_id": file_id, "session_id": session_id})
     
+    # Save to ChromaDB
+    db = Chroma(persist_directory=CHROMA_PATH, embedding_function=hugging_face_ef)
+    db.add_documents(texts)
+    logger.info(f"Added file '{file_path}' to ChromaDB for session '{session_id}'.")
+
+
+# Clean up expired files and ChromaDB collections
+def cleanup_resources():
+    """Clean up expired files and ChromaDB collections."""
+    now = time.time()
+    for session_id, session_data in list(active_sessions.items()):
+        if now - session_data['last_accessed'] > SESSION_TIMEOUT:
+            # Remove files
+            files = session_data.get('files', {})
+            for file_id, (file_path, filename) in files.items():
+                if os.path.exists(file_path):
+                    os.remove(file_path)
+                    logger.info(f"Deleted file: {file_path}")
+            
+            # Remove ChromaDB chunks
+            db = Chroma(persist_directory=CHROMA_PATH)
+            db.delete(where={"session_id": session_id})
+            logger.info(f"Deleted ChromaDB chunks for session: {session_id}")
+            
+            # Remove session
+            del active_sessions[session_id]
+            
+# Start the scheduler
+scheduler = BackgroundScheduler()
+scheduler.add_job(cleanup_resources, 'interval', minutes=5)  # Run every 5 minutes
+scheduler.start()
 
 @app.route('/')
 def index():
+    session_id = session.sid
+    if session_id not in active_sessions:
+        active_sessions[session_id] = {
+            'last_accessed': time.time(),
+            'files': dict()
+        }
     return render_template('index.html')  # Serve the HTML file we created
 
-@app.route('/wait',methods=["POST"])
-def wait():
-    time.sleep(int(request.form.get("time")))
-    return jsonify({"status":"ok"}),200
-
-@app.route("/ai",methods=["POST"])
-def aiPost():
-    print("Post /ai called")
-    json_content = request.json
-    query = json_content.get("query")
-    
-    print("Query:",query)
-    
-    response_answer = llm_model.generate_content(query)
-    return response_answer.text
-
 # add files
 @app.route('/upload-file', methods=['POST'])
 def upload_file():
-    print(request.files)
+    """Handle file uploads."""
     if 'file' not in request.files:
-        return jsonify({'error': 'No files in request', 'status': 'error'}), 400
+        return jsonify({'error': 'No file in request', 'status': 'error'}), 400
     
     file = request.files['file']
     file_id = request.form.get('file_count')
+    session_id = session.sid
     
-    if not session.get("files"):
-        session["files"] = dict()
+    if not file or not file.filename:
+        return jsonify({'error': 'No file selected', 'status': 'error'}), 400
     
-    if file.filename:
-        # Secure the filename
-        filename = secure_filename(file.filename)
-        file_path = os.path.join(UPLOAD_FOLDER, filename)
-        file.save(file_path)
         
-        session['files'][file_id] = (file_path,filename)
-                        
-        # Here you can call your RAG pipeline processing function
-        # process_pdf(file_path)
-        add_file_to_chroma(file_path,file_id)
+    filename = secure_filename(file.filename)
+    file_path = os.path.join(UPLOAD_FOLDER, filename)
+    file.save(file_path)
     
-    return jsonify({
-        'message': 'Files uploaded successfully',
-        'status': 'success'
-    }), 200
+    # Update session data
+    active_sessions[session_id]['files'][file_id] = (file_path,file.filename)
+    active_sessions[session_id]['last_accessed'] = time.time()
+    
+    # Add file chunks to ChromaDB
+    add_file_to_chroma(file_path, file_id, session_id)
+    
+    return jsonify({'message': 'File uploaded successfully', 'status': 'success'}), 200
+
     
 @app.route('/get-files',methods=["GET"])
 def get_files():
-    return jsonify({"files":session.get("files")}),200
+    return jsonify({"files":active_sessions[session.sid]['files']}),200
 
 @app.route('/chroma-status',methods=["GET"])
 def chroma_status():
@@ -155,13 +176,29 @@ def chroma_status():
 @app.route('/remove-file',methods=["POST"])
 def remove_file():
     file_id = request.form.get('file_id')
-    session["files"].pop(file_id)
+    if file_id in active_sessions[session.sid]['files']:
+        file_path = active_sessions[session.sid]['files'][file_id][0]
+        # remove file from upload folder
+        if os.path.exists(file_path):
+            os.remove(file_path)
+            logger.info(f"Deleted file: {file_path}")
+        # Remove file from session
+        del active_sessions[session.sid]['files'][file_id]
+    else:
+        logger.info(f"File not found in session: {file_id}")
+    
     remove_file_from_chroma(file_id)
+    
     return jsonify({
         'message': 'File deleted successfully',
         'status': 'success'
     }), 200
 
+# Ensure scheduler stops on app exit
+@app.teardown_appcontext
+def shutdown_scheduler(exception=None):
+    scheduler.shutdown()
+
 def main():
     app.run(host="0.0.0.0",port=8000,debug=True)
 

log.txt

@@ -0,0 +1,98 @@
+2024-12-06 14:02:03,186 - INFO - Anonymized telemetry enabled. See                     https://docs.trychroma.com/telemetry for more information.
+2024-12-06 14:02:03,499 - DEBUG - Starting component System
+2024-12-06 14:02:03,499 - DEBUG - Starting component Posthog
+2024-12-06 14:02:03,499 - DEBUG - Starting component OpenTelemetryClient
+2024-12-06 14:02:03,499 - DEBUG - Starting component SqliteDB
+2024-12-06 14:02:03,501 - DEBUG - Starting component SimpleQuotaEnforcer
+2024-12-06 14:02:03,501 - DEBUG - Starting component SimpleRateLimitEnforcer
+2024-12-06 14:02:03,501 - DEBUG - Starting component LocalSegmentManager
+2024-12-06 14:02:03,501 - DEBUG - Starting component LocalExecutor
+2024-12-06 14:02:03,501 - DEBUG - Starting component SegmentAPI
+2024-12-06 14:02:03,519 - DEBUG - Looking up time zone info from registry
+2024-12-06 14:02:03,542 - INFO - Adding job tentatively -- it will be properly scheduled when the scheduler starts
+2024-12-06 14:02:03,542 - INFO - Added job "cleanup_resources" to job store "default"
+2024-12-06 14:02:03,542 - INFO - Scheduler started
+2024-12-06 14:02:03,542 - DEBUG - Looking for jobs to run
+2024-12-06 14:02:03,548 - DEBUG - Next wakeup is due at 2024-12-06 14:07:03.542688+05:30 (in 299.993982 seconds)
+2024-12-06 14:02:03,567 - WARNING -  * Debugger is active!
+2024-12-06 14:02:03,567 - INFO -  * Debugger PIN: 471-334-865
+2024-12-06 14:02:04,024 - DEBUG - Starting new HTTPS connection (1): us.i.posthog.com:443
+2024-12-06 14:02:05,409 - DEBUG - https://us.i.posthog.com:443 "POST /batch/ HTTP/11" 200 15
+2024-12-06 14:07:03,556 - DEBUG - Looking for jobs to run
+2024-12-06 14:07:03,556 - INFO - Running job "cleanup_resources (trigger: interval[0:05:00], next run at: 2024-12-06 14:07:03 IST)" (scheduled at 2024-12-06 14:07:03.542688+05:30)
+2024-12-06 14:07:03,556 - DEBUG - Next wakeup is due at 2024-12-06 14:12:03.542688+05:30 (in 299.985749 seconds)
+2024-12-06 14:07:03,565 - INFO - Deleted expired file: uploads\20_DSA_patterns_.pdf
+2024-12-06 14:07:03,565 - INFO - Deleted expired file: uploads\a16-ahn.pdf
+2024-12-06 14:07:03,565 - INFO - Deleted expired file: uploads\blank_page.pdf
+2024-12-06 14:07:03,565 - INFO - Deleted expired file: uploads\Complete_Git_Cheat_Sheet.pdf
+2024-12-06 14:07:03,571 - INFO - Deleted expired file: uploads\E150000182706715682136.PDF
+2024-12-06 14:07:03,573 - INFO - Deleted expired file: uploads\Prathamesh_Sable_Resume.pdf
+2024-12-06 14:07:03,573 - INFO - Deleted expired file: uploads\README.md
+2024-12-06 14:07:03,573 - INFO - Job "cleanup_resources (trigger: interval[0:05:00], next run at: 2024-12-06 14:12:03 IST)" executed successfully
+2024-12-06 14:12:03,550 - DEBUG - Looking for jobs to run
+2024-12-06 14:12:03,550 - DEBUG - Next wakeup is due at 2024-12-06 14:17:03.542688+05:30 (in 299.992333 seconds)
+2024-12-06 14:12:03,550 - INFO - Running job "cleanup_resources (trigger: interval[0:05:00], next run at: 2024-12-06 14:17:03 IST)" (scheduled at 2024-12-06 14:12:03.542688+05:30)
+2024-12-06 14:12:03,550 - INFO - Job "cleanup_resources (trigger: interval[0:05:00], next run at: 2024-12-06 14:17:03 IST)" executed successfully
+2024-12-06 14:17:03,550 - DEBUG - Looking for jobs to run
+2024-12-06 14:17:03,550 - DEBUG - Next wakeup is due at 2024-12-06 14:22:03.542688+05:30 (in 299.991993 seconds)
+2024-12-06 14:17:03,550 - INFO - Running job "cleanup_resources (trigger: interval[0:05:00], next run at: 2024-12-06 14:22:03 IST)" (scheduled at 2024-12-06 14:17:03.542688+05:30)
+2024-12-06 14:17:03,550 - INFO - Job "cleanup_resources (trigger: interval[0:05:00], next run at: 2024-12-06 14:22:03 IST)" executed successfully
+2024-12-06 14:22:03,548 - DEBUG - Looking for jobs to run
+2024-12-06 14:22:03,549 - DEBUG - Next wakeup is due at 2024-12-06 14:27:03.542688+05:30 (in 299.992908 seconds)
+2024-12-06 14:22:03,549 - INFO - Running job "cleanup_resources (trigger: interval[0:05:00], next run at: 2024-12-06 14:27:03 IST)" (scheduled at 2024-12-06 14:22:03.542688+05:30)
+2024-12-06 14:22:03,549 - INFO - Job "cleanup_resources (trigger: interval[0:05:00], next run at: 2024-12-06 14:27:03 IST)" executed successfully
+2024-12-06 14:27:03,559 - DEBUG - Looking for jobs to run
+2024-12-06 14:27:03,559 - DEBUG - Next wakeup is due at 2024-12-06 14:32:03.542688+05:30 (in 299.983279 seconds)
+2024-12-06 14:27:03,559 - INFO - Running job "cleanup_resources (trigger: interval[0:05:00], next run at: 2024-12-06 14:32:03 IST)" (scheduled at 2024-12-06 14:27:03.542688+05:30)
+2024-12-06 14:27:03,559 - INFO - Job "cleanup_resources (trigger: interval[0:05:00], next run at: 2024-12-06 14:32:03 IST)" executed successfully
+2024-12-06 14:32:03,558 - DEBUG - Looking for jobs to run
+2024-12-06 14:32:03,558 - DEBUG - Next wakeup is due at 2024-12-06 14:37:03.542688+05:30 (in 299.983849 seconds)
+2024-12-06 14:32:03,565 - INFO - Running job "cleanup_resources (trigger: interval[0:05:00], next run at: 2024-12-06 14:37:03 IST)" (scheduled at 2024-12-06 14:32:03.542688+05:30)
+2024-12-06 14:32:03,565 - INFO - Job "cleanup_resources (trigger: interval[0:05:00], next run at: 2024-12-06 14:37:03 IST)" executed successfully
+2024-12-06 14:37:03,558 - DEBUG - Looking for jobs to run
+2024-12-06 14:37:03,566 - DEBUG - Next wakeup is due at 2024-12-06 14:42:03.542688+05:30 (in 299.976199 seconds)
+2024-12-06 14:37:03,566 - INFO - Running job "cleanup_resources (trigger: interval[0:05:00], next run at: 2024-12-06 14:42:03 IST)" (scheduled at 2024-12-06 14:37:03.542688+05:30)
+2024-12-06 14:37:03,568 - INFO - Job "cleanup_resources (trigger: interval[0:05:00], next run at: 2024-12-06 14:42:03 IST)" executed successfully
+2024-12-06 14:40:58,432 - INFO -  * Detected change in 'D:\\Coding\\Multi-File-Chatting\\app.py', reloading
+2024-12-06 14:41:07,871 - INFO - Anonymized telemetry enabled. See                     https://docs.trychroma.com/telemetry for more information.
+2024-12-06 14:41:08,298 - DEBUG - Starting component System
+2024-12-06 14:41:08,298 - DEBUG - Starting component Posthog
+2024-12-06 14:41:08,298 - DEBUG - Starting component OpenTelemetryClient
+2024-12-06 14:41:08,298 - DEBUG - Starting component SqliteDB
+2024-12-06 14:41:08,322 - DEBUG - Starting component SimpleQuotaEnforcer
+2024-12-06 14:41:08,322 - DEBUG - Starting component SimpleRateLimitEnforcer
+2024-12-06 14:41:08,323 - DEBUG - Starting component LocalSegmentManager
+2024-12-06 14:41:08,323 - DEBUG - Starting component LocalExecutor
+2024-12-06 14:41:08,323 - DEBUG - Starting component SegmentAPI
+2024-12-06 14:41:08,327 - DEBUG - Looking up time zone info from registry
+2024-12-06 14:41:08,340 - INFO - Adding job tentatively -- it will be properly scheduled when the scheduler starts
+2024-12-06 14:41:08,340 - INFO - Added job "cleanup_resources" to job store "default"
+2024-12-06 14:41:08,341 - INFO - Scheduler started
+2024-12-06 14:41:08,341 - DEBUG - Looking for jobs to run
+2024-12-06 14:41:08,344 - DEBUG - Next wakeup is due at 2024-12-06 14:46:08.340975+05:30 (in 299.996005 seconds)
+2024-12-06 14:41:08,371 - WARNING -  * Debugger is active!
+2024-12-06 14:41:08,381 - INFO -  * Debugger PIN: 471-334-865
+2024-12-06 14:41:08,839 - DEBUG - Starting new HTTPS connection (1): us.i.posthog.com:443
+2024-12-06 14:41:10,082 - DEBUG - https://us.i.posthog.com:443 "POST /batch/ HTTP/11" 200 15
+2024-12-06 14:45:48,701 - INFO -  * Detected change in 'D:\\Coding\\Multi-File-Chatting\\app.py', reloading
+2024-12-06 14:45:56,500 - INFO - Anonymized telemetry enabled. See                     https://docs.trychroma.com/telemetry for more information.
+2024-12-06 14:45:56,772 - DEBUG - Starting component System
+2024-12-06 14:45:56,772 - DEBUG - Starting component Posthog
+2024-12-06 14:45:56,772 - DEBUG - Starting component OpenTelemetryClient
+2024-12-06 14:45:56,772 - DEBUG - Starting component SqliteDB
+2024-12-06 14:45:56,799 - DEBUG - Starting component SimpleQuotaEnforcer
+2024-12-06 14:45:56,799 - DEBUG - Starting component SimpleRateLimitEnforcer
+2024-12-06 14:45:56,799 - DEBUG - Starting component LocalSegmentManager
+2024-12-06 14:45:56,799 - DEBUG - Starting component LocalExecutor
+2024-12-06 14:45:56,799 - DEBUG - Starting component SegmentAPI
+2024-12-06 14:45:56,802 - DEBUG - Looking up time zone info from registry
+2024-12-06 14:45:56,802 - INFO - Adding job tentatively -- it will be properly scheduled when the scheduler starts
+2024-12-06 14:45:56,802 - INFO - Added job "cleanup_resources" to job store "default"
+2024-12-06 14:45:56,802 - INFO - Scheduler started
+2024-12-06 14:45:56,802 - DEBUG - Looking for jobs to run
+2024-12-06 14:45:56,802 - DEBUG - Next wakeup is due at 2024-12-06 14:50:56.802779+05:30 (in 300.000000 seconds)
+2024-12-06 14:45:56,818 - WARNING -  * Debugger is active!
+2024-12-06 14:45:56,834 - INFO -  * Debugger PIN: 471-334-865
+2024-12-06 14:45:57,307 - DEBUG - Starting new HTTPS connection (1): us.i.posthog.com:443
+2024-12-06 14:45:58,410 - INFO -  * Detected change in 'D:\\Coding\\Multi-File-Chatting\\utils.py', reloading
+2024-12-06 14:45:58,591 - DEBUG - https://us.i.posthog.com:443 "POST /batch/ HTTP/11" 200 15

requirements.txt

@@ -15,4 +15,5 @@ markdown
 python-docx 
 flask
 werkzeug
-Flask-Session
+Flask-Session
+apscheduler

utils.py

@@ -0,0 +1,16 @@
+
+@app.route('/wait',methods=["POST"])
+def wait():
+    time.sleep(int(request.form.get("time")))
+    return jsonify({"status":"ok"}),200
+
+@app.route("/ai",methods=["POST"])
+def aiPost():
+    print("Post /ai called")
+    json_content = request.json
+    query = json_content.get("query")
+    
+    print("Query:",query)
+    
+    response_answer = llm_model.generate_content(query)
+    return response_answer.text