Update src/streamlit_app.py

src/streamlit_app.py

@@ -62,10 +62,16 @@ def retrieve_documents(query, top_k=10):
 
 def generate_response(user_query, docs):
     context = "\n\n---\n\n".join(d['metadata']['text'] for d in docs)
-    messages = [{"role": "system", "content":
-                 "You are a helpful legal assistant. Use provided context from documents. Answer only using the context."}]
+    sources = sorted({d['metadata']['source'] for d in docs if 'source' in d['metadata']})
+    messages = [
+        {"role": "system", "content":
+         "You are a helpful legal assistant. Use the provided context from the documents to answer the user's question. "
+         "At the end of your answer, write a single line starting with 'Source: ' and list the sources of the documents you used. "
+         "If multiple sources are used, separate them with commas. The user should be able to clearly understand where the information came from."}
+    ]
     messages.extend(st.session_state.history)
     messages.append({"role": "user", "content": f"Context:\n{context}\n\nQuestion:\n{user_query}"})
+
     try:
         resp = client.chat.completions.create(
             model="gpt-4o-mini",
@@ -78,23 +84,17 @@ def generate_response(user_query, docs):
         st.error(f"Response error: {e}")
         reply = "Sorry, I encountered an error generating the answer."
 
-
-    # unique_sources = sorted({d['metadata']['source'] for d in docs})
-    # if unique_sources:
-    #     reply += "\n\n---\n\n**Sources used:**\n"
-    #     for src in unique_sources:
-    #         reply += f"> “{src}”\n"
-
-    # unique_sources_txt = sorted({d['metadata']['text'] for d in docs})
-    # if unique_sources_txt:
-    #     reply += "\n\n---\n\n**Document(s) retrieved:**\n"
-    #     for src in unique_sources_txt:
-    #         reply += f"> “{src}”\n"
+    # Optional: force clean source line if LLM misses it
+    if sources:
+        clean_sources = ", ".join(sources)
+        if "Source:" not in reply:
+            reply += f"\n\nSource: {clean_sources}"
 
     st.session_state.history.append({"role": "assistant", "content": reply})
     return reply
 
 
+
 # Chat UI
 with st.form("chat_input", clear_on_submit=True):
     user_input = st.text_input("You:", "")