app.py

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 retrival import generate_data_store #,add_document_to_existing_db, delete_chunks_by_source
from langchain_community.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.prompts import ChatPromptTemplate
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
from langchain_huggingface import HuggingFaceEndpoint
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
import asyncio

import nltk
nltk.download('punkt_tab')

import nltk
nltk.download('averaged_perceptron_tagger_eng')

app = Flask(__name__)

# Set the secret key for session management
app.secret_key = os.urandom(24)

# Configurations
UPLOAD_FOLDER = "uploads/"
VECTOR_DB_FOLDER = "VectorDB/"
#TABLE_DB_FOLDER = "TableDB/"

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(VECTOR_DB_FOLDER, exist_ok=True)
#os.makedirs(TABLE_DB_FOLDER, exist_ok=True)

# Global variables
CHROMA_PATH = None
TEMP_PATH = None
#TABLE_PATH = None

#System prompt

'''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. If the question cannot be answered using the context, respond with: "The information requested is not mentioned in the context."

Context:
{context}

---

Question:
{question}

Response:
"""
'''

PROMPT_TEMPLATE = """
You are working as a retrieval-augmented generation (RAG) assistant specializing in providing precise and accurate responses. Generate a response based only on the provided context and question, following these concrete instructions:

- **Adhere strictly to the context:** Use only the information in the context to answer the question. Do not add any external details or assumptions.
- **Handle multiple chunks:** The context is divided into chunks, separated by "###". Query-related information may be present in any chunk.
- **Focus on relevance:** Identify and prioritize chunks relevant to the question while ignoring unrelated chunks.
- **Answer concisely and factually:** Provide clear, direct, and structured responses based on the retrieved 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
    #global TABLE_PATH
    
    #old_db = session.get('old_db', None) 
    #print(f"Selected DB: {CHROMA_PATH}")
    
    #if TEMP_PATH is not None and TEMP_PATH != CHROMA_PATH:
    #    session['history'] = []
    #TEMP_PATH = CHROMA_PATH
    
    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 Document Database 
        embedding_function = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
        #embedding_function = HuggingFaceEmbeddings(model_name="mixedbread-ai/mxbai-embed-large-v1")
        db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)
        results_document = db.similarity_search_with_relevance_scores(query_text, k=3)
        
        print("results------------------->",results_document)       
        context_text_document = "\n\n---\n\n".join([doc.page_content for doc, _score in results_document])
        
        
        # # Load the selected Table Database
        # #embedding_function = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
        # embedding_function = HuggingFaceEmbeddings(model_name="mixedbread-ai/mxbai-embed-large-v1")
        # tdb = Chroma(persist_directory=TABLE_PATH, embedding_function=embedding_function)
        # results_table = tdb.similarity_search_with_relevance_scores(query_text, k=2)
        
        # print("results------------------->",results_table)       
        # context_text_table = "\n\n---\n\n".join([doc.page_content for doc, _score in results_table])

        # Prepare the prompt and query the model
        prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
        prompt = prompt_template.format(context=context_text_document,question=query_text)
        #prompt = prompt_template.format(context=context_text_document,table=context_text_table, question=query_text)
        print("results------------------->",prompt)
        
        
        #Model Defining and its use 
        repo_id = "mistralai/Mistral-7B-Instruct-v0.3"
        HFT = os.environ["HF_TOKEN"]
        llm = HuggingFaceEndpoint(
            repo_id=repo_id,
            max_tokens=3000,
            temperature=0.8,
            huggingfacehub_api_token=HFT,
        )

        data= llm(prompt)
        #data = response.choices[0].message.content
        print("LLM response------------------>",data)
        # filtering the uneccessary context.
        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'])

    return render_template('chat.html', history=session['history'])

'''
@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('/create-db', methods=['GET', 'POST'])
def create_db():
    if request.method == 'POST':
        db_name = request.form['db_name']

        # Ensure the upload folder exists
        os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

        # Define the base upload path
        upload_base_path = os.path.join(app.config['UPLOAD_FOLDER'], secure_filename(db_name))
        os.makedirs(upload_base_path, exist_ok=True)

        # Check for uploaded folder or files
        folder_files = request.files.getlist('folder')
        single_files = request.files.getlist('file')

        if folder_files and any(file.filename for file in folder_files):
            # Process folder files
            for file in folder_files:
                file_path = os.path.join(upload_base_path, secure_filename(file.filename))
                os.makedirs(os.path.dirname(file_path), exist_ok=True)
                file.save(file_path)

        elif single_files and any(file.filename for file in single_files):
            # Process single files
            for file in single_files:
                file_path = os.path.join(upload_base_path, secure_filename(file.filename))
                file.save(file_path)

        else:
            return "No files uploaded", 400

        # Generate datastore
        generate_data_store(upload_base_path, db_name)

        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):
    
    #Selecting the Documnet Vector DB
    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}")
    
    #Selecting the Table Vector DB
    # global TABLE_PATH
    # print(f"Selected DB: {TABLE_PATH}")
    # TABLE_PATH = os.path.join(TABLE_DB_FOLDER, db_name)
    # TABLE_PATH = TABLE_PATH.replace("\\", "/")
    # print(f"Selected DB: {TABLE_PATH}")
    
    
    return redirect(url_for('chat'))

@app.route('/update-dbs/<db_name>', methods=['GET','POST'])
def update_db(db_name):
    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
        print(f"Selected DB: {db_name}")
        DB_PATH = os.path.join(VECTOR_DB_FOLDER, db_name)
        DB_PATH = DB_PATH.replace("\\", "/")
        print(f"Selected DB: {DB_PATH}")

        generate_data_store(DB_PATH, db_name)

        return redirect(url_for('list_dbs'))
    return render_template('update_db.html')

if __name__ == "__main__":
    app.run(debug=False, use_reloader=False)