Update app.py

app.py

@@ -1094,10 +1094,10 @@ def generate_answer(message, choice, retrieval_mode):
             chain_type_kwargs={"prompt": prompt_template}
         )
         response = qa_chain({"query": message})
-        return response['output'], extract_addresses(response['output'])
+        return response['result'], extract_addresses(response['result'])
     elif retrieval_mode == "Knowledge-Graph":
-        response = chain_neo4j.invoke({"context": "", "question": message})
-        return response['output'], extract_addresses(response['output'])
+        response = chain_neo4j.invoke({"question": message})
+        return response, extract_addresses(response)
     else:
         return "Invalid retrieval mode selected.", []
 
@@ -1680,3 +1680,4 @@ demo.queue()
 demo.launch(share=True)
 
 
+