Update app.py

app.py

@@ -1,166 +1,166 @@
-from flask import Flask, render_template, request, redirect, url_for, session
-import os
-from werkzeug.utils import secure_filename
-from retrival import generate_data_store
-from langchain_community.vectorstores import Chroma
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.prompts import ChatPromptTemplate
-from huggingface_hub import InferenceClient
-from langchain.schema import Document 
-from langchain_core.documents import Document
-from dotenv import load_dotenv
-import re
-import glob
-import shutil
-from werkzeug.utils import secure_filename
-
-app = Flask(__name__)
-
-# Set the secret key for session management
-app.secret_key = os.urandom(24)
-
-# Configurations
-UPLOAD_FOLDER = "uploads/"
-VECTOR_DB_FOLDER = "VectorDB/"
-app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
-os.makedirs(UPLOAD_FOLDER, exist_ok=True)
-os.makedirs(VECTOR_DB_FOLDER, exist_ok=True)
-
-# Global variables
-CHROMA_PATH = None
-PROMPT_TEMPLATE = """
-You are working with a retrieval-augmented generation (RAG) setup. Your task is to generate a response based on the context provided and the question asked. Consider only the following context strictly, and use it to answer the question. Do not include any external information.
-
-Context:
-{context}
-
----
-
-Question:
-{question}
-
-Response:
-"""
-HFT = os.getenv('HF_TOKEN')
-client = InferenceClient(api_key=HFT)
-
-@app.route('/', methods=['GET'])
-def home():
-    return render_template('home.html')
-
-@app.route('/chat', methods=['GET', 'POST'])
-def chat():
-    
-    if 'history' not in session:
-        session['history'] = []
-    print("sessionhist1",session['history'])
-    
-    global CHROMA_PATH
-    old_db = session.get('old_db', None) 
-    print(f"Selected DB: {CHROMA_PATH}")
-    
-    if old_db != None:
-        if CHROMA_PATH != old_db:
-            session['history'] = [] 
-        
-    #print("sessionhist1",session['history'])
-    
-    if request.method == 'POST':
-        query_text = request.form['query_text']
-        if CHROMA_PATH is None:
-            return render_template('chat.html', error="No vector database selected!", history=[])
-
-        # Load the selected database
-        embedding_function = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
-        db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)
-        results = db.similarity_search_with_relevance_scores(query_text, k=3)        
-        context_text = "\n\n---\n\n".join([doc.page_content for doc, _score in results])
-
-        # Prepare the prompt and query the model
-        prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
-        prompt = prompt_template.format(context=context_text, question=query_text)
-        print("results------------------->",prompt)
-        response = client.chat.completions.create(
-            model="mistralai/Mistral-7B-Instruct-v0.3",
-            messages=[{"role": "system", "content": "You are an assistant specifically designed to generate responses based on the context provided. Your task is to answer questions strictly using the context without adding any external knowledge or information. Please ensure that your responses are relevant, accurate, and based solely on the given context."},
-                      {"role": "user", "content": prompt}],
-            max_tokens=500,
-            temperature=0.3
-        )
-        data = response.choices[0].message.content
-
-        if re.search(r'\bmention\b|\bnot mention\b|\bnot mentioned\b|\bnot contain\b|\bnot include\b|\bnot provide\b|\bdoes not\b|\bnot explicitly\b|\bnot explicitly mentioned\b', data, re.IGNORECASE):
-            data = "We do not have information related to your query on our end."
-            
-        # Save the query and answer to the session history
-        session['history'].append((query_text, data))
-        
-        # Mark the session as modified to ensure it gets saved
-        session.modified = True
-        print("sessionhist2",session['history'])
-
-        return render_template('chat.html', query_text=query_text, answer=data, history=session['history'],old_db=CHROMA_PATH)
-
-    return render_template('chat.html', history=session['history'], old_db=CHROMA_PATH)
-
-@app.route('/create-db', methods=['GET', 'POST'])
-def create_db():
-    if request.method == 'POST':
-        db_name = request.form['db_name']
-
-        # Get all files from the uploaded folder
-        files = request.files.getlist('folder')
-        if not files:
-            return "No files uploaded", 400
-
-        # Define the base upload path
-        upload_base_path = os.path.join(app.config['UPLOAD_FOLDER'], secure_filename(db_name))
-        #upload_base_path = upload_base_path.replace("\\", "/")
-        print(f"Base Upload Path: {upload_base_path}")
-        os.makedirs(upload_base_path, exist_ok=True)
-
-        # Save each file and recreate folder structure
-        for file in files:
-            print("file , files",files,file)
-            #relative_path = file.filename  # This should contain the subfolder structure
-            file_path = os.path.join(upload_base_path)
-            #file_path = file_path.replace("\\", "/")
-
-            # Ensure the directory exists before saving the file
-            print(f"Saving to: {file_path}")
-            os.makedirs(os.path.dirname(file_path), exist_ok=True)
-            
-            
-            # Get the file path and save it
-            file_path = os.path.join(upload_base_path, secure_filename(file.filename))
-            file.save(file_path)
-            # with open(file_path, 'wb') as f:
-            #     shutil.copyfileobj(file.stream, f)
-        # Generate datastore
-        generate_data_store(upload_base_path, db_name)
-
-        # # Clean up uploaded files (if needed)
-        # uploaded_files = glob.glob(os.path.join(app.config['UPLOAD_FOLDER'], '*'))
-        # for f in uploaded_files:
-        #     os.remove(f)
-
-        return redirect(url_for('list_dbs'))
-
-    return render_template('create_db.html')
-
-@app.route('/list-dbs', methods=['GET'])
-def list_dbs():
-    vector_dbs = [name for name in os.listdir(VECTOR_DB_FOLDER) if os.path.isdir(os.path.join(VECTOR_DB_FOLDER, name))]
-    return render_template('list_dbs.html', vector_dbs=vector_dbs)
-
-@app.route('/select-db/<db_name>', methods=['POST'])
-def select_db(db_name):
-    global CHROMA_PATH
-    print(f"Selected DB: {CHROMA_PATH}")
-    CHROMA_PATH = os.path.join(VECTOR_DB_FOLDER, db_name)
-    CHROMA_PATH = CHROMA_PATH.replace("\\", "/")
-    print(f"Selected DB: {CHROMA_PATH}")
-    return redirect(url_for('chat'))
-
-if __name__ == "__main__":
+from flask import Flask, render_template, request, redirect, url_for, session
+import os
+from werkzeug.utils import secure_filename
+from retrival import generate_data_store
+from langchain_community.vectorstores import Chroma
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.prompts import ChatPromptTemplate
+from huggingface_hub import InferenceClient
+from langchain.schema import Document 
+from langchain_core.documents import Document
+from dotenv import load_dotenv
+import re
+import glob
+import shutil
+from werkzeug.utils import secure_filename
+
+app = Flask(__name__)
+
+# Set the secret key for session management
+app.secret_key = os.urandom(24)
+
+# Configurations
+UPLOAD_FOLDER = "uploads/"
+VECTOR_DB_FOLDER = "VectorDB/"
+app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+os.makedirs(VECTOR_DB_FOLDER, exist_ok=True)
+
+# Global variables
+CHROMA_PATH = None
+PROMPT_TEMPLATE = """
+You are working with a retrieval-augmented generation (RAG) setup. Your task is to generate a response based on the context provided and the question asked. Consider only the following context strictly, and use it to answer the question. Do not include any external information.
+
+Context:
+{context}
+
+---
+
+Question:
+{question}
+
+Response:
+"""
+HFT = os.getenv('HF_TOKEN')
+client = InferenceClient(api_key=HFT)
+
+@app.route('/', methods=['GET'])
+def home():
+    return render_template('home.html')
+
+@app.route('/chat', methods=['GET', 'POST'])
+def chat():
+    
+    if 'history' not in session:
+        session['history'] = []
+    print("sessionhist1",session['history'])
+    
+    global CHROMA_PATH
+    old_db = session.get('old_db', None) 
+    print(f"Selected DB: {CHROMA_PATH}")
+    
+    if old_db != None:
+        if CHROMA_PATH != old_db:
+            session['history'] = [] 
+        
+    #print("sessionhist1",session['history'])
+    
+    if request.method == 'POST':
+        query_text = request.form['query_text']
+        if CHROMA_PATH is None:
+            return render_template('chat.html', error="No vector database selected!", history=[])
+
+        # Load the selected database
+        embedding_function = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+        db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)
+        results = db.similarity_search_with_relevance_scores(query_text, k=3)        
+        context_text = "\n\n---\n\n".join([doc.page_content for doc, _score in results])
+
+        # Prepare the prompt and query the model
+        prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
+        prompt = prompt_template.format(context=context_text, question=query_text)
+        print("results------------------->",prompt)
+        response = client.chat.completions.create(
+            model="mistralai/Mistral-7B-Instruct-v0.3",
+            messages=[{"role": "system", "content": "You are an assistant specifically designed to generate responses based on the context provided. Your task is to answer questions strictly using the context without adding any external knowledge or information. Please ensure that your responses are relevant, accurate, and based solely on the given context."},
+                      {"role": "user", "content": prompt}],
+            max_tokens=1000,
+            temperature=0.3
+        )
+        data = response.choices[0].message.content
+
+        if re.search(r'\bmention\b|\bnot mention\b|\bnot mentioned\b|\bnot contain\b|\bnot include\b|\bnot provide\b|\bdoes not\b|\bnot explicitly\b|\bnot explicitly mentioned\b', data, re.IGNORECASE):
+            data = "We do not have information related to your query on our end."
+            
+        # Save the query and answer to the session history
+        session['history'].append((query_text, data))
+        
+        # Mark the session as modified to ensure it gets saved
+        session.modified = True
+        print("sessionhist2",session['history'])
+
+        return render_template('chat.html', query_text=query_text, answer=data, history=session['history'],old_db=CHROMA_PATH)
+
+    return render_template('chat.html', history=session['history'], old_db=CHROMA_PATH)
+
+@app.route('/create-db', methods=['GET', 'POST'])
+def create_db():
+    if request.method == 'POST':
+        db_name = request.form['db_name']
+
+        # Get all files from the uploaded folder
+        files = request.files.getlist('folder')
+        if not files:
+            return "No files uploaded", 400
+
+        # if not exist
+        os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+        # Define the base upload path
+        upload_base_path = os.path.join(app.config['UPLOAD_FOLDER'], secure_filename(db_name))
+        #upload_base_path = upload_base_path.replace("\\", "/")
+        print(f"Base Upload Path: {upload_base_path}")
+        os.makedirs(upload_base_path, exist_ok=True)
+
+        # Save each file and recreate folder structure
+        for file in files:
+            print("file , files",files,file)
+            #relative_path = file.filename  # This should contain the subfolder structure
+            file_path = os.path.join(upload_base_path)
+            #file_path = file_path.replace("\\", "/")
+
+            # Ensure the directory exists before saving the file
+            print(f"Saving to: {file_path}")
+            os.makedirs(os.path.dirname(file_path), exist_ok=True)
+            
+            
+            # Get the file path and save it
+            file_path = os.path.join(upload_base_path, secure_filename(file.filename))
+            file.save(file_path)
+                    
+        # Generate datastore
+        generate_data_store(upload_base_path, db_name)
+
+        # # Clean up uploaded files (if needed)
+        #if os.path.exists(app.config['UPLOAD_FOLDER']):
+        #    shutil.rmtree(app.config['UPLOAD_FOLDER'])
+
+        return redirect(url_for('list_dbs'))
+
+    return render_template('create_db.html')
+
+@app.route('/list-dbs', methods=['GET'])
+def list_dbs():
+    vector_dbs = [name for name in os.listdir(VECTOR_DB_FOLDER) if os.path.isdir(os.path.join(VECTOR_DB_FOLDER, name))]
+    return render_template('list_dbs.html', vector_dbs=vector_dbs)
+
+@app.route('/select-db/<db_name>', methods=['POST'])
+def select_db(db_name):
+    global CHROMA_PATH
+    print(f"Selected DB: {CHROMA_PATH}")
+    CHROMA_PATH = os.path.join(VECTOR_DB_FOLDER, db_name)
+    CHROMA_PATH = CHROMA_PATH.replace("\\", "/")
+    print(f"Selected DB: {CHROMA_PATH}")
+    return redirect(url_for('chat'))
+
+if __name__ == "__main__":
     app.run(debug=False, use_reloader=False)