working RAG pipeline

app.py

@@ -10,7 +10,6 @@ from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings
 import google.generativeai as genai
 
 from langchain_chroma import Chroma
-from chromadb import Client
 
 import os
 from dotenv import load_dotenv
@@ -29,39 +28,56 @@ GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY')
 CHROMA_PATH = "chroma"
 UPLOAD_FOLDER = "uploads"
 SESSION_TIMEOUT = 30 * 60  # 30 minutes
+PROMPT_TEMPLATE = """
+Answer the given query based only on the context given below.
+context:
+{context}
 
-# Initialize Hugging Face embedding
+---
+based on the context above answer the following query:
+{query}
+"""
+
+# make required folder for files to upload
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)    
+
+# Initialize Hugging Face embeddings Model
 hugging_face_ef = HuggingFaceInferenceAPIEmbeddings(
     api_key=HF_TOKEN,
     model_name="sentence-transformers/all-MiniLM-L6-v2"
 )
+# initialize LLM
 genai.configure(api_key=GOOGLE_API_KEY)
 llm_model = genai.GenerativeModel("gemini-1.5-flash")
 
 app = Flask(__name__)
 
-os.makedirs(UPLOAD_FOLDER, exist_ok=True)
-app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
-    
-
+# initialize session 
 app.config["SESSION_PERMANENT"] = True
 app.config["SESSION_TYPE"] = "filesystem"
-
 Session(app)
 
 # Initialize ChromaDB client
-client = Client()
+db = Chroma(persist_directory=CHROMA_PATH, embedding_function=hugging_face_ef)
 active_sessions = dict()
 
 
-def remove_file_from_chroma(file_id,CHROMA_PATH=CHROMA_PATH):
-    db = Chroma(persist_directory=CHROMA_PATH)
+def remove_file_from_chroma(file_id,session_id):    
+    # Get chunks for session
+    session_chunks = db.get(where={"session_id": session_id})
     
-    # delete chunks from db where metadata file_id is equal to file_id
-    db.delete(ids=db.get(where={"file_id": file_id},include=[])['ids'])
-    return True
-    
-        
+    # Further filter by file_id
+    ids_to_del = []
+    for i in range(len(session_chunks['ids'])):
+        if session_chunks['metadatas'][i]['file_id'] == str(file_id):
+            ids_to_del.append(session_chunks['ids'][i])
+
+    # delete chunks from db where metadata file_id is equal to file_id if there are ;)
+    if len(ids_to_del) > 0:
+        db.delete(ids=ids_to_del)
+        return True
+    return False
+
 def add_file_to_chroma(file_path, file_id, session_id):
     """Add file chunks to ChromaDB."""
     extension = file_path.split(".")[-1]
@@ -91,45 +107,48 @@ def add_file_to_chroma(file_path, file_id, session_id):
         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)
+    db.add_documents(texts,embedding=hugging_face_ef)
     logger.info(f"Added file '{file_path}' to ChromaDB for session '{session_id}'.")
 
+def generate_query_response(query):
+    response = dict()
+    top_related = db.similarity_search_with_relevance_scores(query,filter={"session_id": session.sid},k=4)
+    
+    response['is_relevant'] = top_related[0][1] >= 0.6
+    
+    # filter chunks with score > 0.3
+    # top_related = [chunk for chunk in top_related if chunk[1] > 0.3]
+        
+    context = "\n".join([chunk[0].page_content for chunk in top_related])
+    
+    prompt = PROMPT_TEMPLATE.format(context = context,query = query)
+    
+    print(top_related)
+    
+    response['answer'] = llm_model.generate_content(prompt).text
+    response['sources'] = [{
+                               "page_content":chunk[0].page_content,
+                               "score" : chunk[1],
+                               "metadata":chunk[0].metadata 
+                           } for chunk in top_related]
+    
+    return response
 
-# 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():
+@app.before_request
+def update_session():
     session_id = session.sid
     if session_id not in active_sessions:
         active_sessions[session_id] = {
             'last_accessed': time.time(),
             'files': dict()
         }
+        logger.info(f"CREATED NEW SESSION with ID {session_id}")
+    else:
+        active_sessions[session_id]['last_accessed'] = time.time()
+    
+
+@app.route('/')
+def index():
     return render_template('index.html')  # Serve the HTML file we created
 
 # add files
@@ -165,39 +184,80 @@ def upload_file():
 def get_files():
     return jsonify({"files":active_sessions[session.sid]['files']}),200
 
-@app.route('/chroma-status',methods=["GET"])
-def chroma_status():
+@app.route('/status',methods=["GET"])
+def status():
     # return all data from chroma db
-    db = Chroma(persist_directory=CHROMA_PATH)
-    return jsonify({"data":db.get()}),200
+    return jsonify({
+        "Active_sessions":active_sessions,
+        "chroma_data":db.get()
+    }),200
 
 @app.route('/remove-file',methods=["POST"])
 def remove_file():
     file_id = request.form.get('file_id')
-    if file_id in active_sessions[session.sid]['files']:
-        file_path = active_sessions[session.sid]['files'][file_id][0]
+    session_id = session.sid
+    if file_id in active_sessions[session_id]['files']:
+        file_path = active_sessions[session_id]['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]
+        del active_sessions[session_id]['files'][file_id]
     else:
         logger.info(f"File not found in session: {file_id}")
     
-    remove_file_from_chroma(file_id)
+    if remove_file_from_chroma(file_id,session_id):
+        return jsonify({
+            'message': 'File deleted successfully',
+            'status': 'success'
+        }), 200
+    else:
+        return jsonify({
+            'message': 'File Not Found',
+            'status': 'fail'
+        }), 404
+
+# 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.delete(where={"session_id": session_id})
+            logger.info(f"Deleted ChromaDB chunks for session: {session_id}")
+            
+            # Remove session
+            del active_sessions[session_id]
+            
+@app.route("/ask_query", methods=['POST'])
+def ask_query():
+    query = request.form.get("query")
     
-    return jsonify({
-        'message': 'File deleted successfully',
-        'status': 'success'
-    }), 200
+    resp = generate_query_response(query)
+    
+    return jsonify(resp),200
+            
+# Start the scheduler
+scheduler = BackgroundScheduler()
+scheduler.add_job(cleanup_resources, 'interval', minutes=5)  # Run every 5 minutes
+scheduler.start()
 
 # Ensure scheduler stops on app exit
 @app.teardown_appcontext
 def shutdown_scheduler(exception=None):
-    scheduler.shutdown()
+    if scheduler.running:
+        scheduler.shutdown()
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0",port=8000,debug=True)
+    app.run(host="0.0.0.0",port=8000,debug=True,threaded=True)
     
     

templates/index.html

@@ -22,7 +22,8 @@
 <body class="bg-gray-100 dark:bg-gray-900 text-gray-900 dark:text-gray-100">
 
 
-    <div id="loader" class="fixed top-0 left-0 w-full h-full" style="background-color: rgba(50, 50, 50, 0.5);z-index: 1000;">
+    <div id="loader" class="fixed top-0 left-0 w-full h-full"
+        style="background-color: rgba(50, 50, 50, 0.5);z-index: 1000;">
         <div class="text-center" style="position: absolute;top: 50%;left: 50%;transform: translate(-50%, -50%);">
             <div role="status">
                 <svg aria-hidden="true"
@@ -115,6 +116,13 @@
                 <!-- Chat Messages -->
                 <div class="flex-1 overflow-y-auto p-4 space-y-4" id="chat-messages">
                     <div class="flex items-start space-x-2">
+                        <div class="bg-gray-200 dark:bg-gray-700 p-3 rounded-lg">
+                            <p>
+                                Upload your files, And you are ready to ask questions about them🫡.
+                            </p>
+                        </div>
+                    </div>
+                    <!-- <div class="flex items-start space-x-2">
                         <div class="bg-gray-200 dark:bg-gray-700 p-3 rounded-lg">
                             <p>
                                 Hello! How can I assist you today?
@@ -131,7 +139,7 @@
                                 I need help with my account.
                             </p>
                         </div>
-                    </div>
+                    </div> -->
                 </div>
                 <!-- Message Input -->
                 <form id="message-form">
@@ -166,6 +174,7 @@
         const dropZone = document.getElementById('dropzone'); // label of input
         const fileInput = document.getElementById('file-upload'); // input element
         const fileList = document.getElementById('file-list'); // div to display files uploaded
+        const chat = document.getElementById('chat-messages');
 
         var file_list = [];
         var file_count = 0;
@@ -187,6 +196,46 @@
             }
         });
 
+        function add_user_message(msg) {
+            var item = document.createElement('div');
+            item.className = "flex items-start space-x-2 flex-row-reverse";
+            item.innerHTML = `<div class="bg-blue-500 text-white p-3 rounded-lg">
+                            <p>
+                                ${msg}
+                            </p>
+                        </div>`;
+            chat.appendChild(item);
+        }
+
+        function add_ai_response(resp) {
+            sources = ""
+            if (resp.is_relevant) {
+                sources += "Sources: "
+                resp.sources.forEach(function (source) {
+                    sources += '<a class="text-blue-500" href="#" style="background: #00000085;border-radius: 4px;padding: 1px 3px;margin: 2px;">'
+                    sources += source.metadata.source.split("\\")[1]
+                    if (source.metadata.page != undefined) {
+                        sources += ", Page:" + source.metadata.page
+                    }
+                    sources += '</a>,'
+                });
+
+            } else {
+                sources = "The Documents provided are not sufficient to answer the query provided, So response might not be accurate."
+            }
+
+            var item = document.createElement('div');
+            item.className = "flex items-start space-x-2";
+            item.innerHTML = `<div class="bg-gray-200 dark:bg-gray-700 p-3 rounded-lg">
+                            <p>
+                                ${resp.answer}
+                            </p>
+                            <p class="text-sm text-gray-500 dark:text-gray-400 mt-2">
+                                ${sources}
+                            </p>
+                        </div>`;
+            chat.appendChild(item);
+        }
 
 
         // message send 
@@ -194,9 +243,26 @@
             e.preventDefault();
             const messageInput = document.getElementById('message-input');
             const message = messageInput.value.trim();
+
+            add_user_message(message)
+
             if (message) {
-                console.log('Sending message:', message);
-                messageInput.value = '';
+                $.ajax({
+                    url: "/ask_query",
+                    method: "POST",
+                    data: {
+                        query: message
+                    },
+                    success: function (response) {
+                        resp = JSON.stringify(response)
+                        resp = JSON.parse(resp)
+                        console.log(resp);
+                        add_ai_response(resp);
+                    },
+                    error: function (xhr, status, error) {
+                        console.log(error);
+                    }
+                })
             }
         });
 
@@ -263,7 +329,7 @@
                         <i class="fas fa-times"></i>
                     </button>`;
                     fileList.appendChild(listItem);
-                    set_loading("Uploading File "+ file.name);
+                    set_loading("Uploading File " + file.name);
 
                     var formData = new FormData();
                     formData.append('file', file);