Update app.py

app.py

@@ -816,10 +816,23 @@ def submit_message(prompt, prompt_template, temperature, max_tokens, context_len
         print("Similarity_score:", similarity_score)
         print("")
 
-
+    highest_similarity_score = -1  # Initialize with a score lower than possible
+    selected_document = None  # To hold the document with the highest similarity score
+    
+    for document in completion['content']['source_documents']:
+        if document.state['query_similarity_score'] > highest_similarity_score:
+            highest_similarity_score = document.state['query_similarity_score']
+            selected_document = document
+    
+    if selected_document is not None:
+        # Remove the "/home/user/app/" part from the document name
+        modified_source = selected_document.metadata['source'].replace('/home/user/app/', '')
+        source_info = f"Primary source document: {modified_source}, Page: {selected_document.metadata['page']}")
+    else:
+        source_info = "Primary source document: not determined"
     
     #chat_messages = [(prompt_msg['content'], completion['content'])]    
-    chat_messages = [(prompt, completion['content']['answer'])] 
+    chat_messages = [(prompt, completion['content']['answer'], source_info)] 
     return '', chat_messages,  state    # total_tokens_used_msg,