Update app.py

app.py

@@ -46,11 +46,10 @@ examples = [
 
 def get_docs(input_query, country = None):
     # Construct a hacky query to focus the retriever on the target country (see notes below)
-    query = "For the country of "+country+", "+input_query
-    # if country:
-    #     query = "For the country of "+country+", "+input_query
-    # else:
-    #     query = input_query
+    if country:
+        query = "For the country of "+country+", "+input_query
+    else:
+        query = input_query
     
     # Get top 150 because we want to make sure we have 10 pertaining to the selected country
     # TEMP SOLUTION: not ideal, but FAISS document store doesnt allow metadata filtering. Needs to be tested with the full dataset
@@ -98,7 +97,7 @@ def get_refs(res):
              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):
+def run_query(input_text, country):
     docs = get_docs(input_text)
     st.write('Selected country: ', country) # Debugging country
     res = pipe.run(query=input_text, documents=docs)
@@ -170,4 +169,4 @@ else:
 
 
 if st.button('Submit'):
-    run_query(text)
+    run_query(text, country)