Update app.py

app.py

@@ -109,7 +109,7 @@ def get_docs(input_query, country = [], vulnerability_cat = []):
           filters = {'vulnerability_cat': {'$in': vulnerability_cat}}
     else:
       filters = {'country': {'$in': country},'vulnerability_cat': {'$in': vulnerability_cat}}
-  docs = retriever.retrieve(query=query, filters = filters, top_k = 10)
+  docs = retriever.retrieve(query=input_query, filters = filters, top_k = 10)
   # Break out the key fields and convert to pandas for filtering
   docs = [{**x.meta,"score":x.score,"content":x.content} for x in docs]
   df_docs = pd.DataFrame(docs)
@@ -154,11 +154,11 @@ def get_refs(docs, res):
   return result_str
 
 # define a special function for putting the prompt together (as we can't use haystack)
-def get_prompt(docs, query):
+def get_prompt(docs, input_query):
   base_prompt=prompt_template
   # Add the meta data for references
   context = ' - '.join(['&&& [ref. '+str(d.meta['ref_id'])+'] '+d.meta['document']+' &&&: '+d.content for d in docs])
-  prompt = base_prompt+"; Context: "+context+"; Question: "+query+"; Answer:"
+  prompt = base_prompt+"; Context: "+context+"; Question: "+input_query+"; Answer:"
   return(prompt)
 
 def run_query(input_query, country, model_sel):
@@ -167,13 +167,13 @@ def run_query(input_query, country, model_sel):
     # st.write('Selected country: ', country) # Debugging country
     if model_sel == "chatGPT":
         # res = pipe.run(query=input_text, documents=docs)
-        res = openai.ChatCompletion.create(model="gpt-3.5-turbo", messages=[{"role": "user", "content": get_prompt(docs, query=input_query)}])
+        res = openai.ChatCompletion.create(model="gpt-3.5-turbo", messages=[{"role": "user", "content": get_prompt(docs, input_query)}])
         output = res["results"][0]
         references = get_refs(docs, output)
-    else:
-        res = client.text_generation(get_prompt_llama2(docs, query=input_query), max_new_tokens=4000, temperature=0.01, model=model)
-        output = res
-        references = get_refs(docs, res)
+    # else:
+    #     res = client.text_generation(get_prompt_llama2(docs, query=input_query), max_new_tokens=4000, temperature=0.01, model=model)
+    #     output = res
+    #     references = get_refs(docs, res)
     st.write('Response')
     st.success(output)
     st.write('References')