update referces

app.py

@@ -23,6 +23,7 @@ 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
+import re
 
 # Load environment variables
 load_dotenv(".env")
@@ -125,15 +126,8 @@ When responding to queries, follow these guidelines:
    - 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] \
-    '
+   - ALWAYS USE INLINE CITATIONS with embed source URLs where users can verify information or explore further. 
+   - The citations should be in the format [1], [2], etc., in the response with links to reference sources.
 
 FOLLOW ALL THE GIVEN INSTRUCTIONS, FAILURE TO DO SO WILL RESULT IN TERMINATION OF THE CHAT.
 
@@ -192,14 +186,28 @@ async def websocket_endpoint(websocket: WebSocket):
             try:
                 # Define an async generator for streaming
                 async def stream_response():
+                    complete_response = ""
+                    context = {}
                     async for chunk in conversational_rag_chain.astream(
                         {"input": question, 'language': language},
                         config={"configurable": {"session_id": session_id}}
                     ):
+                        if "context" in chunk:
+                            context = chunk['context']
                         # Send each chunk to the client
                         if "answer" in chunk:
+                            complete_response += chunk['answer']
                             await websocket.send_json({'response': chunk['answer']})
 
+                    if context:
+                        citations = re.findall(r'\[(\d+)\]', complete_response)
+                        citation_numbers = list(map(int, citations))
+                        sources = dict()
+                        for index, doc in enumerate(context):
+                            if (index+1) in citation_numbers:
+                                sources[f"[{index+1}]"] = doc.metadata["source"]
+                        await websocket.send_json({'sources': sources})
+
                 await stream_response()
             except Exception as e:
                 print(f"Error during message handling: {e}")