Update app.py

app.py

@@ -1,10 +1,13 @@
 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
+from retrival2 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
@@ -13,6 +16,7 @@ import re
 import glob
 import shutil
 from werkzeug.utils import secure_filename
+import asyncio
 
 app = Flask(__name__)
 
@@ -22,16 +26,19 @@ app.secret_key = os.urandom(24)
 # Configurations
 UPLOAD_FOLDER = "uploads/"
 VECTOR_DB_FOLDER = "VectorDB/"
-NLTK_FOLDER = "nltk_data/"
+#TABLE_DB_FOLDER = "TableDB/"
 
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
-os.environ["MPLCONFIGDIR"] = "/tmp"
+
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 os.makedirs(VECTOR_DB_FOLDER, exist_ok=True)
-os.makedirs(NLTK_FOLDER, exist_ok=True)
+#os.makedirs(TABLE_DB_FOLDER, exist_ok=True)
 
 # Global variables
 CHROMA_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. Do not include any external information.
 
@@ -45,8 +52,28 @@ Question:
 
 Response:
 """
-HFT = os.getenv('HF_TOKEN')
-client = InferenceClient(api_key=HFT)
+
+# PROMPT_TEMPLATE = """
+# You are working with a retrieval-augmented generation (RAG) setup. Your task is to generate a response based on the provided context, table data, and the question asked. Consider only the given inputs strictly and use them to answer the question. Do not include any external information.
+
+# If the table variable contains tabular data, analyze and extract all relevant details from it. Provide a structured response summarizing the table data if it is relevant to the question. If the table data is not relevant, base your answer only on the context.
+
+# Context:
+# {context}
+
+# Table:
+# {table}
+
+# ---
+
+# Question:
+# {question}
+
+# Response:
+# """
+
+#HFT = os.getenv('HF_TOKEN')
+#client = InferenceClient(api_key=HFT)
 
 @app.route('/', methods=['GET'])
 def home():
@@ -60,12 +87,14 @@ def chat():
     print("sessionhist1",session['history'])
     
     global CHROMA_PATH
+    #global TABLE_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'] = [] 
+    # if old_db != None:
+    #     if CHROMA_PATH != old_db:
+    #         session['history'] = [] 
         
     #print("sessionhist1",session['history'])
     
@@ -74,25 +103,46 @@ def chat():
         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")
+        # 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 = 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])
+        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, question=query_text)
+        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)
-        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
+        
+        
+        #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 = response.choices[0].message.content
 
+        data= llm(prompt)
+        #data = response.choices[0].message.content
+
+        # 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."
             
@@ -159,12 +209,44 @@ def list_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)
+    app.run(debug=False, use_reloader=False)
+    
+