Update app.py

app.py

@@ -79,23 +79,24 @@ def get_docs(input_query, country = None):
         ls_dict.append(doc)
     return(ls_dict)
 
-def get_refs(docs, res):
+def get_refs(res):
   '''
   Parse response for engineered reference ids (refer to prompt template)
   Extract documents using reference ids
   '''
-  # This pattern should be returned by gpt3.5 & llama2
+  # 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, res)]
   # extract
   result_str = ""  # Initialize an empty string to store the result
-  for i in range(len(res['documents'])):
+  for i in range(len(docs)):
       doc = docs[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'] + "'*<br> <br>" # Add <br> for a line break
-  return result_str 
+            result_str += "**Ref. " + str(ref_id) + " [" + doc['meta']['country'] + " " + doc['meta']['document_name'] + "]:** " + "*'" + doc['content'] + "'*<br> <br>" # Add <br> for a line break
+
+  return result_str
 
 def run_query(input_text, country):
     docs = get_docs(input_text, country)