added refereced functionality

.gitignore

@@ -1,2 +1,3 @@
 .env
-*.ipynb
+*.ipynb
+__pycache__/*

Dockerfile

@@ -13,4 +13,4 @@ COPY --chown=user ./requirements.txt requirements.txt
 RUN pip install --no-cache-dir --upgrade -r requirements.txt
 
 COPY --chown=user . /app
-CMD ["gunicorn", "-b", "0.0.0.0:7860", "-k", "geventwebsocket.gunicorn.workers.GeventWebSocketWorker", "-w", "1", "app:app"]
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,18 +1,13 @@
-from gevent import monkey
-monkey.patch_all()
-
-import nltk
-nltk.download('punkt_tab')
-
 import nltk
 nltk.download('punkt_tab')
 
 import os
 from dotenv import load_dotenv
 import asyncio
-from flask import Flask, request, render_template
-from flask_cors import CORS
-from flask_socketio import SocketIO, emit, join_room, leave_room
+from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
+from fastapi.responses import HTMLResponse
+from fastapi.templating import Jinja2Templates
+from fastapi.middleware.cors import CORSMiddleware
 from langchain.chains import create_history_aware_retriever, create_retrieval_chain
 from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain_community.chat_message_histories import ChatMessageHistory
@@ -27,6 +22,7 @@ from langchain.retrievers import ContextualCompressionRetriever
 from langchain_community.chat_models import ChatPerplexity
 from langchain.retrievers.document_compressors import CrossEncoderReranker
 from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+from langchain_core.prompts import PromptTemplate
 
 # Load environment variables
 load_dotenv(".env")
@@ -41,14 +37,19 @@ os.environ['USER_AGENT'] = USER_AGENT
 os.environ["GROQ_API_KEY"] = GROQ_API_KEY
 os.environ["TOKENIZERS_PARALLELISM"] = 'true'
 
-# Initialize Flask app and SocketIO with CORS
-app = Flask(__name__)
-CORS(app)
-socketio = SocketIO(app, async_mode='gevent', cors_allowed_origins="*")
-app.config['SESSION_COOKIE_SECURE'] = True  # Use HTTPS
-app.config['SESSION_COOKIE_HTTPONLY'] = True
-app.config['SESSION_COOKIE_SAMESITE'] = 'Lax'
-app.config['SECRET_KEY'] = SECRET_KEY
+# Initialize FastAPI app and CORS
+app = FastAPI()
+origins = ["*"]  # Adjust as needed
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+templates = Jinja2Templates(directory="templates")
 
 # Function to initialize Pinecone connection
 def initialize_pinecone(index_name: str):
@@ -59,7 +60,6 @@ def initialize_pinecone(index_name: str):
         print(f"Error initializing Pinecone: {e}")
         raise
 
-
 ##################################################
 ##          Change down here
 ##################################################
@@ -71,8 +71,6 @@ bm25 = BM25Encoder().load("./mbzuai-policies.json")
 ##################################################
 ##################################################
 
-# old_embed_model = HuggingFaceEmbeddings(model_name="sentence-transformers/gte-multilingual-base")
-
 # Initialize models and retriever
 embed_model = HuggingFaceEmbeddings(model_name="jinaai/jina-embeddings-v3", model_kwargs={"trust_remote_code":True})
 retriever = PineconeHybridSearchRetriever(
@@ -81,20 +79,15 @@ retriever = PineconeHybridSearchRetriever(
     index=pinecone_index, 
     top_k=20, 
     alpha=0.5,
-    
 )
 
 # Initialize LLM
-# llm = ChatGroq(model="llama-3.1-70b-versatile", temperature=0, max_tokens=1024, max_retries=2)
-llm = ChatPerplexity(temperature=0, pplx_api_key=GROQ_API_KEY, model="llama-3.1-sonar-large-128k-chat", max_tokens=1024, max_retries=2)
-
+llm = ChatPerplexity(temperature=0, pplx_api_key=GROQ_API_KEY, model="llama-3.1-sonar-large-128k-chat", max_tokens=512, max_retries=2)
 
 # Initialize Reranker
-# compressor = FlashrankRerank()
 model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base")
 compressor = CrossEncoderReranker(model=model, top_n=20)
 
-
 compression_retriever = ContextualCompressionRetriever(
     base_compressor=compressor, base_retriever=retriever
 )
@@ -115,33 +108,35 @@ contextualize_q_prompt = ChatPromptTemplate.from_messages(
 history_aware_retriever = create_history_aware_retriever(llm, compression_retriever, contextualize_q_prompt)
 
 # QA system prompt and chain
-qa_system_prompt = """You are a highly skilled information retrieval assistant. Use the following context to answer questions effectively. \
-If you don't know the answer, simply state that you don't know. \
-Your answer should be in {language} language. \
-Provide answers in proper HTML format and keep them concise. \
-
-When responding to queries, follow these guidelines: \
-
-    1. Provide Clear Answers: \
-        - Based on the language of the question, you have to answer in that language. E.g. if the question is in English language then answer in the English language or if the question is in Arabic language then you should answer in Arabic language. /
-        - Ensure the response directly addresses the query with accurate and relevant information.\
-
-    2. Include Detailed References: \
-        - Links to Sources: Include URLs to credible sources where users can verify information or explore further. \
-        - Reference Sites: Mention specific websites or platforms that offer additional information. \
-        - Downloadable Materials: Provide links to any relevant downloadable resources if applicable. \
-    
-    3. Formatting for Readability: \
-        - The answer should be in a proper HTML format with appropriate tags. \
-        - For arabic language response align the text to right and convert numbers also.
-        - Double check if the language of answer is correct or not.
-        - Use bullet points or numbered lists where applicable to present information clearly. \
-        - Highlight key details using bold or italics. \
-        - Provide proper and meaningful abbreviations for urls. Do not include naked urls. \
-    
-    4. Organize Content Logically: \
-        - Structure the content in a logical order, ensuring easy navigation and understanding for the user. \
-        
+qa_system_prompt = """ You are a highly skilled information retrieval assistant. Use the following context to answer questions effectively. 
+If you don't know the answer, simply state that you don't know. 
+Your answer should be in {language} language. 
+
+When responding to queries, follow these guidelines:
+
+1. Provide Clear Answers: 
+   - Based on the language of the question, you have to answer in that language. E.g., if the question is in English, then answer in English; if the question is in Arabic, you should answer in Arabic.
+   - Ensure the response directly addresses the query with accurate and relevant information.
+   - Do not give long answers. Provide detailed but concise responses.
+
+2. Formatting for Readability: 
+   - Provide the entire response in proper markdown format.
+   - Use structured Maekdown elements such as headings, subheading, lists, tables, and links.
+   - Use emaphsis on headings, important texts and phrases.
+
+3. Proper Citations and References:
+   - ALWAYS INCLUDE SOURCES URLs where users can verify information or explore further. 
+   - Use inline citations with embed referenced source link in the format [1], [2], etc., in the response to reference sources.
+   - ALWAYS PROVIDE "References" SECTION AT THE END OF RESPONSE.
+   - In the "References" section, list the referenced sources with their urls in the following format 
+    'References
+        [1] Heading 1[Source 1 url] \
+        [2] Heading 2[Source 2 url] \
+        [3] Heading 3[Source 2 url] \
+    '
+
+FOLLOW ALL THE GIVEN INSTRUCTIONS, FAILURE TO DO SO WILL RESULT IN TERMINATION OF THE CHAT.
+
 {context}
 """
 qa_prompt = ChatPromptTemplate.from_messages(
@@ -151,7 +146,9 @@ qa_prompt = ChatPromptTemplate.from_messages(
         ("human", "{input}")
     ]
 )
-question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
+
+document_prompt = PromptTemplate(input_variables=["page_content", "source"], template="{page_content} \n\n Source: {source}")
+question_answer_chain = create_stuff_documents_chain(llm, qa_prompt, document_prompt=document_prompt)
 
 # Retrieval and Generative (RAG) Chain
 rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
@@ -159,9 +156,6 @@ rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chai
 # Chat message history storage
 store = {}
 
-def clean_temporary_data():
-    store.clear()
-
 def get_session_history(session_id: str) -> BaseChatMessageHistory:
     if session_id not in store:
         store[session_id] = ChatMessageHistory()
@@ -177,53 +171,45 @@ conversational_rag_chain = RunnableWithMessageHistory(
     output_messages_key="answer",
 )
 
-# Function to handle WebSocket connection
-@socketio.on('connect')
-def handle_connect():
-    print(f"Client connected: {request.sid}")
-    emit('connection_response', {'message': 'Connected successfully.'})
-
-# Function to handle WebSocket disconnection
-@socketio.on('disconnect')
-def handle_disconnect():
-    print(f"Client disconnected: {request.sid}")
-    clean_temporary_data()
-
-# Function to handle WebSocket messages
-@socketio.on('message')
-def handle_message(data):
-    question = data.get('question')
-    language = data.get('language')
-    if "en" in language:
-        language = "English"
-    else:
-        language = "Arabic"
-    session_id = data.get('session_id', SESSION_ID_DEFAULT)
-    # chain = conversational_rag_chain.pick("answer")
-
-    # try:
-    #     for chunk in conversational_rag_chain.stream(
-    #             {"input": question, 'language': language},
-    #             config={"configurable": {"session_id": session_id}},
-    #         ):
-    #         emit('response', chunk, room=request.sid)
-    # except Exception as e:
-    #     print(f"Error during message handling: {e}")
-    #     emit('response', "An error occurred while processing your request." + str(e), room=request.sid)
 
+# WebSocket endpoint with streaming
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    print(f"Client connected: {websocket.client}")
+    session_id = None
     try:
-        response = conversational_rag_chain.invoke({"input": question, 'language': language}, config={"configurable": {"session_id": session_id}})
-        emit('response', response, room=request.sid)
-    except Exception as e:
-        print(f"Error during message handling: {e}")
-        emit('response', "An error occurred while processing your request." + str(e), room=request.sid)
-
+        while True:
+            data = await websocket.receive_json()
+            question = data.get('question')
+            language = data.get('language')
+            if "en" in language:
+                language = "English"
+            else:
+                language = "Arabic"
+            session_id = data.get('session_id', SESSION_ID_DEFAULT)
+            # Process the question
+            try:
+                # Define an async generator for streaming
+                async def stream_response():
+                    async for chunk in conversational_rag_chain.astream(
+                        {"input": question, 'language': language},
+                        config={"configurable": {"session_id": session_id}}
+                    ):
+                        # Send each chunk to the client
+                        if "answer" in chunk:
+                            await websocket.send_json({'response': chunk['answer']})
+
+                await stream_response()
+            except Exception as e:
+                print(f"Error during message handling: {e}")
+                await websocket.send_json({'response': "Something went wrong, Please try again.."})
+    except WebSocketDisconnect:
+        print(f"Client disconnected: {websocket.client}")
+        if session_id:
+            store.pop(session_id, None)
 
 # Home route
-@app.route("/")
-def index_view():
-    return render_template('chat.html') 
-
-# Main function to run the app
-if __name__ == '__main__': 
-    socketio.run(app, debug=True)
+@app.get("/", response_class=HTMLResponse)
+async def read_index(request: Request):
+    return templates.TemplateResponse("chat.html", {"request": request})

requirements.txt

@@ -5,11 +5,9 @@ langchain-huggingface
 pinecone
 pinecone-text
 flashrank
-flask
-flask-cors
-flask-socketio
-gunicorn
-gevent 
-gevent-websocket
+fastapi>=0.68.0
+uvicorn[standard]>=0.15.0
+websockets>=10.0
+python-multipart>=0.0.5
 openai
 einops