Update app.py

app.py

@@ -88,7 +88,7 @@ def retrieve(query):
             break
         elif i == len(contexts) - 1:
             prompt = prompt_start + "\n\n---\n\n".join(contexts) + prompt_end
-    return prompt
+    return prompt, contexts
 
 
 # first let's make it simpler to get answers
@@ -109,8 +109,8 @@ def complete(prompt):
 
 def query(question, top_k_reader, top_k_retriever):
     # first we retrieve relevant items from Pinecone
-    query_with_contexts = retrieve(question)
-    return complete(query_with_contexts)
+    query_with_contexts, contexts = retrieve(question)
+    return complete(query_with_contexts), contexts
 
 
 indexing_pipeline_with_classification = Pipeline()
@@ -312,28 +312,28 @@ if st.session_state.results:
 
     st.write("## Results:")
 
-    for result in st.session_state.results:
-        # answer, context = result.answer, result.context
-        # start_idx = context.find(answer)
-        # end_idx = start_idx + len(answer)
-        # Hack due to this bug: https://github.com/streamlit/streamlit/issues/3190
-        try:
-            # source = f"[{result.meta['Title']}]({result.meta['link']})"
-            # st.write(
-            # markdown(f'**Source:** {source} \n {context[:start_idx] } {str(annotation(answer, "ANSWER", "#8ef"))} {context[end_idx:]} \n '),
-            # unsafe_allow_html=True,
-            #  )
-            # all_contexts = '\n'.join(contexts)
-            st.write(markdown(f"# Answer: {result}"),#"\n ## Extracted from context: \n {all_contexts}"),
-                unsafe_allow_html=True,
-            )
-        except:
-            # filename = result.meta.get('filename', "")
-            # st.write(
-            # markdown(f'From file: {filename} \n {context[:start_idx] } {str(annotation(answer, "ANSWER", "#8ef"))} {context[end_idx:]} \n '),
-            # unsafe_allow_html=True,
-            # )
-            st.write(
-                markdown(f"Answer: {result}"),
-                unsafe_allow_html=True,
-            )
+    result, contexts = st.session_state.results:
+    # answer, context = result.answer, result.context
+    # start_idx = context.find(answer)
+    # end_idx = start_idx + len(answer)
+    # Hack due to this bug: https://github.com/streamlit/streamlit/issues/3190
+    try:
+        # source = f"[{result.meta['Title']}]({result.meta['link']})"
+        # st.write(
+        # markdown(f'**Source:** {source} \n {context[:start_idx] } {str(annotation(answer, "ANSWER", "#8ef"))} {context[end_idx:]} \n '),
+        # unsafe_allow_html=True,
+        #  )
+        all_contexts = '\n'.join(contexts)
+        st.write(markdown(f"# Answer: {result} \n ## Extracted from context: \n {all_contexts}"),
+            unsafe_allow_html=True,
+        )
+    except:
+        # filename = result.meta.get('filename', "")
+        # st.write(
+        # markdown(f'From file: {filename} \n {context[:start_idx] } {str(annotation(answer, "ANSWER", "#8ef"))} {context[end_idx:]} \n '),
+        # unsafe_allow_html=True,
+        # )
+        st.write(
+            markdown(f"Answer: {result}"),
+            unsafe_allow_html=True,
+        )