Update app.py

app.py

@@ -81,47 +81,50 @@ def get_refs(res):
   Parse response for engineered reference ids (refer to prompt template)
   Extract documents using reference ids
   '''
+  import re
   text = res["results"][0]
   # This pattern should be returned by gpt3.5
   # pattern = r'ref\. (\d+)\]\.'
   pattern = r'ref\. (\d+)'
   ref_ids = [int(match) for match in re.findall(pattern, text)]
   # extract
+  result_str = ""  # Initialize an empty string to store the result
   for i in range(len(res['documents'])):
-    doc = res['documents'][i].to_dict()
-    ref_id = doc['meta']['ref_id']
-    if ref_id in ref_ids:
-      print("Ref. "+str(ref_id)+" ["+doc['meta']['country']+" "+doc['meta']['document_name']+"]: "+doc['content'])
-
-# def run_query(input_text):
-#     docs = get_docs(input_text)
-#     res = pipe.run(query=input_text, documents=docs)
-#     output = res["results"][0]
-#     st.write('Response')
-#     st.success(output)
-#     references = get_refs(res)
-#     st.write('References')
-#     st.success(str(references))
+      doc = res['documents'][i].to_dict()
+      ref_id = doc['meta']['ref_id']
+      if ref_id in ref_ids:
+          result_str += "Ref. " + str(ref_id) + " [" + doc['meta']['country'] + " " + doc['meta']['document_name'] + "]: " + "'" + doc['content'] + "' \n "
+  return result_str 
 
 def run_query(input_text):
     docs = get_docs(input_text)
-    st.write('Debug: Documents', docs) # Debugging output
-
     res = pipe.run(query=input_text, documents=docs)
-    st.write('Debug: Results', res) # Debugging output
-
     output = res["results"][0]
     st.write('Response')
-    st.success(str(output))
-
+    st.success(output)
     references = get_refs(res)
-    st.write('Debug: References', references) # Debugging output
-
     st.write('References')
-    if references:
-        st.success(str(references))
-    else:
-        st.warning('No references found.')
+    st.success(str(references))
+
+# def run_query(input_text):
+#     docs = get_docs(input_text)
+#     st.write('Debug: Documents', docs) # Debugging output
+
+#     res = pipe.run(query=input_text, documents=docs)
+#     st.write('Debug: Results', res) # Debugging output
+
+#     output = res["results"][0]
+#     st.write('Response')
+#     st.success(str(output))
+
+#     references = get_refs(res)
+#     st.write('Debug: References', references) # Debugging output
+
+#     st.write('References')
+#     if references:
+#         st.success(str(references))
+#     else:
+#         st.warning('No references found.')