Update src/visualization.py

src/visualization.py

@@ -131,28 +131,29 @@ 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():
-            with st.expander(f"{source.capitalize()} ({len(docs)})", expanded=source == "ontology_context"):
-                for i, doc in enumerate(docs):
-                    # Add separator between documents
-                    if i > 0:
-                        st.markdown("---")
-                    
-                    # Display document content
-                    if hasattr(doc, 'page_content'):
-                        st.markdown(f"**Content:**")
-                        
-                        # Format depending on source
-                        if source in ["ontology", "ontology_context"]:
-                            st.markdown(doc.page_content)
-                        else:
-                            st.code(doc.page_content)
+            st.subheader(f"{entity_type} ({len(entity_list)})")
+            
+            for entity in entity_list:
+                # Add separator between documents
+                if i > 0:
+                    st.markdown("---")
+                
+                # Display document content
+                if hasattr(doc, 'page_content'):
+                    st.markdown(f"**Content:**")
                     
-                    # Display metadata if present
-                    if hasattr(doc, 'metadata') and doc.metadata:
-                        st.markdown("**Metadata:**")
-                        for key, value in doc.metadata.items():
-                            if key != 'source':  # Already shown in section title
-                                st.markdown(f"- **{key}**: {value}")
+                    # Format depending on source
+                    if source in ["ontology", "ontology_context"]:
+                        st.markdown(doc.page_content)
+                    else:
+                        st.code(doc.page_content)
+                
+                # Display metadata if present
+                if hasattr(doc, 'metadata') and doc.metadata:
+                    st.markdown("**Metadata:**")
+                    for key, value in doc.metadata.items():
+                        if key != 'source':  # Already shown in section title
+                            st.markdown(f"- **{key}**: {value}")
     
     with tab3:
         # Show the reasoning flow from query to answer