Update app.py

app.py

@@ -77,12 +77,32 @@ def get_docs(input_query, country = None):
         ls_dict.append(doc)
     return(ls_dict)
 
+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
+  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)
+    st.write('References')
+    st.success(get_refs(res))
 
 
 # Setup retriever, pulling from local faiss datastore