Update app.py

app.py

@@ -298,15 +298,20 @@ def retrieve_bot_answer(query, loaded_documents):
     qa_bot_instance = create_retrieval_qa_bot(loaded_documents)
     bot_response = qa_bot_instance({"question": query})
     # Check if the 'answer' key exists in the bot_response dictionary
-    if 'answer' in bot_response:
-        # answer = bot_response['answer']
-        return bot_response
-    else:
-        raise KeyError("Expected 'answer' key in bot_response, but it was not found.")
-
+    # if 'answer' in bot_response:
+    #     # answer = bot_response['answer']
+    #     return bot_response
+    # else:
+    #     raise KeyError("Expected 'answer' key in bot_response, but it was not found.")
+    result = bot_response['result']
+    sources = []
+    for source in bot_response["source_documents"]:
+        sources.append(source.metadata['source'])
+    return result, sources
 
 # from your_module import load_model, set_custom_prompt, set_custom_prompt_condense, create_vector_database, retrieve_bot_answer
 
+
 def main():
    
     st.title("Docuverse")
@@ -353,11 +358,12 @@ def main():
                     CONDENSE_QUESTION_PROMPT = set_custom_prompt_condense()
                     db = create_vector_database(loaded_documents)
                     # st.write(f"db: {db}") 
-                    response = retrieve_bot_answer(query,loaded_documents)
+                    result, sources = retrieve_bot_answer(query,loaded_documents)
                     # st.write(f"response: {response}") 
                     # Display bot response
                     st.write("Bot Response:")
-                    st.write(response['answer'])
+                    st.write(result)
+                    st.write(sources)
                 except Exception as e:
                     st.error(f"An error occurred: {str(e)}")
             else: