Update src/streamlit_app.py

src/streamlit_app.py

@@ -10,7 +10,7 @@ import openai
 from collections import deque
 from sentence_transformers import SentenceTransformer
 from pinecone import Pinecone
-
+import re
 
 # Setup (exact hardcoded keys you provided)
 client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
@@ -56,38 +56,56 @@ def retrieve_documents(query, top_k=10):
         st.error(f"Retrieve error: {e}")
         return []
 
+
+def clean_chunk_id(cid: str) -> str:
+    """Beautify chunk_id by replacing underscores/dashes with spaces and capitalizing words."""
+    # Remove any trailing '_chunk_xxx' stuff
+    cid = re.sub(r'_chunk.*$', '', cid)
+    # Replace _ and - with spaces
+    cid = cid.replace("_", " ").replace("-", " ")
+    # Capitalize each word
+    cid = " ".join(word.capitalize() for word in cid.split())
+    return cid
+
+
+
 def generate_response(user_query, docs):
     context = "\n\n---\n\n".join(d['metadata']['text'] for d in docs)
     # sources = sorted({d['metadata']['chunk_id'] for d in docs if 'source' in d['metadata']})
+    # --- Build human-friendly sources ---
     readable_sources = []
     for d in docs:
         meta = d['metadata']
-        src = meta.get("source", "unknown")
-        cid = meta.get("chunk_id", "N/A")
+        src = meta.get("source", "unknown").lower()
+        cid = meta.get("chunk_id", "")
+        text_preview = " ".join(meta.get("text", "").split()[:30])
+
+        if src in ["constitution"]:
+            readable_sources.append(f"Constitution ({clean_chunk_id(cid)})")
 
-        if src.lower() in ["constitution", "fbr_ordinance", "ordinance"]:
-            # For constitution and ordinances, chunk_id is enough
-            readable_sources.append(f"{src.title()} (Chunk {cid})")
+        elif src in ["fbr_ordinance", "ordinance", "tax_ordinance"]:
+            readable_sources.append(f"Tax Ordinance ({clean_chunk_id(cid)})")
 
-        elif src.lower() in ["case_law", "case", "tax_case"]:
-            # For case law, add first ~30 words of text
-            text_preview = " ".join(meta.get("text", "").split()[:30])
-            readable_sources.append(f"Case Law (Chunk {cid}): {text_preview}...")
+        elif src in ["case_law", "case", "tax_case"]:
+            # Use first ~30 words of the actual text
+            readable_sources.append(f"Case Law: {text_preview}...")
 
         else:
-            readable_sources.append(f"{src.title()} (Chunk {cid})")
+            readable_sources.append(f"{src.title()} ({clean_chunk_id(cid)})")
 
-    # Deduplicate sources
+    # Deduplicate and sort
     readable_sources = sorted(set(readable_sources))
-    
+
+    # --- System prompt ---
     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. "
-         "Sources should be written in a human-friendly way using metadata:\n"
-         "- Constitution / Ordinances → just mention source name and chunk ID.\n"
-         "- Case law → mention source name, chunk ID, and first ~30 words of the text as a preview.\n"
-         "Do not invent sources. If multiple are used, separate them with commas."}
+         "You are a helpful legal assistant. Use the provided context from documents to answer the user's question. "
+         "At the end of your answer, write a single line starting with 'Source: ' followed by the sources used. "
+         "Formatting rules:\n"
+         "- For Constitution / Ordinances: show the clean chunk id, no underscores/dashes, capitalized words.\n"
+         "- For Case law: ignore chunk id, instead show first ~30 words of the case text.\n"
+         "- Do not use technical terms like 'chunk'. Present sources in a human-friendly way.\n"
+         "If multiple are used, separate them with commas."}
     ]