Update src/visualization.py

src/visualization.py

@@ -131,9 +131,9 @@ def display_reasoning_trace(query: str, retrieved_docs: List[Dict], answer: str,
         
         # Display retrieved document details in expandable sections
         for source, docs in docs_by_source.items():
-            st.subheader(f"{entity_type} ({len(entity_list)})")
+            st.subheader(f"{source} ({len(docs)})")
             
-            for entity in entity_list:
+            for i, doc in enumerate(docs):
                 # Add separator between documents
                 if i > 0:
                     st.markdown("---")