Update src/app.py

src/app.py

@@ -7,7 +7,7 @@ from retriever import get_retriever
 from transformers import pipeline
 import os
 
-# Load local HuggingFace pipeline (avoids .post() bug)
+# Load local HuggingFace pipeline
 pipe = pipeline(
     "text-generation",
     model="tiiuae/falcon-7b-instruct",
@@ -20,10 +20,9 @@ llm = HuggingFacePipeline(pipeline=pipe)
 
 retriever = get_retriever()
 
-# Prompt template for consistent answers & citations
+# Prompt template
 template = """
-You are a legal assistant. Answer the question using the provided context. 
-Always include the source reference(s) from the context in your answer.
+You are a legal assistant. Use the provided context to answer the question.
 If language mode is Nigerian Pidgin, respond in Nigerian Pidgin.
 
 Question: {question}
@@ -43,13 +42,22 @@ qa_chain = RetrievalQA.from_chain_type(
     llm=llm,
     retriever=retriever,
     chain_type="stuff",
+    return_source_documents=True,  # Needed to list references
     chain_type_kwargs={"prompt": prompt}
 )
 
 def answer_question(user_input, lang_choice):
     if lang_choice == "pidgin":
         user_input = f"Respond in Nigerian Pidgin: {user_input}"
-    return qa_chain.run(user_input)
+
+    result = qa_chain(user_input)
+    answer_text = result["result"]
+
+    # Collect unique source file names
+    sources = list({doc.metadata.get("source", "Unknown") for doc in result["source_documents"]})
+    sources_list = "\n".join(f"- {src}" for src in sources)
+
+    return f"{answer_text}\n\nReferences:\n{sources_list}"
 
 def launch_interface():
     iface = gr.Interface(
@@ -60,7 +68,7 @@ def launch_interface():
         ],
         outputs=gr.Textbox(label="Answer"),
         title="KnowYourRight Bot",
-        description="Ask legal rights questions in English or Nigerian Pidgin with sources"
+        description="Ask legal rights questions in English or Nigerian Pidgin with references"
     )
     iface.launch()