Update components/document_viewer.py

components/document_viewer.py

@@ -1,31 +1,31 @@
-# components/document_viewer.py
 import streamlit as st
 from typing import Dict, List
 import pandas as pd
 
+
 class DocumentViewer:
     def __init__(self, case_manager):
         self.case_manager = case_manager
 
     def render(self, case_id: str):
-        """Render document viewer"""
+        """Render document viewer with additional features."""
         case = self.case_manager.get_case(case_id)
         if not case:
             st.error("Case not found")
             return
 
         st.subheader("Documents")
-        
+
         # Document list
         if case['documents']:
             # Create DataFrame for documents
             docs_df = pd.DataFrame(case['documents'])
             docs_df['added_at'] = pd.to_datetime(docs_df['added_at'])
             docs_df = docs_df.sort_values('added_at', ascending=False)
-            
+
             # Display documents in tabs
             tabs = st.tabs([f"📄 {doc['title']}" for doc in docs_df.to_dict('records')])
-            
+
             for tab, doc in zip(tabs, docs_df.to_dict('records')):
                 with tab:
                     self._render_document_tab(doc)
@@ -33,17 +33,33 @@ class DocumentViewer:
             st.info("No documents uploaded yet")
 
     def _render_document_tab(self, doc: Dict):
-        """Render single document tab"""
+        """Render a single document tab with chunk-level references."""
         col1, col2 = st.columns([2, 1])
-        
+
         with col1:
             st.markdown(f"**{doc['title']}**")
             st.markdown(f"Added: {doc['added_at']}")
-            
-            # Display document content in expander
+
+            # Display document content in an expander
             with st.expander("View Content"):
-                st.markdown(doc['content'][:1000] + "..." if len(doc['content']) > 1000 else doc['content'])
-        
+                # Include chunk references if available
+                content_with_chunks = (
+                    "\n\n".join(
+                        [
+                            f"Chunk {idx + 1}:\n{chunk}"
+                            for idx, chunk in enumerate(doc.get('chunks', [doc['content']]))
+                        ]
+                    )
+                    if 'chunks' in doc
+                    else doc['content']
+                )
+                st.text_area(
+                    label="Document Content",
+                    value=content_with_chunks[:3000] + "..." if len(content_with_chunks) > 3000 else content_with_chunks,
+                    height=300,
+                    disabled=True,
+                )
+
         with col2:
             st.markdown("### Actions")
             if st.button("📥 Download", key=f"download_{doc['id']}"):
@@ -51,5 +67,15 @@ class DocumentViewer:
                     "Download Document",
                     doc['content'],
                     file_name=doc['title'],
-                    mime="text/plain"
+                    mime="text/plain",
                 )
+
+            # Add any additional actions (e.g., reprocess or annotate)
+
+    def render_chunk_references(self, references: List[Dict]):
+        """Render a section listing chunks referenced in responses."""
+        st.subheader("Referenced Chunks")
+        for ref in references:
+            st.markdown(
+                f"- **{ref['title']}** (Chunk {ref['chunk_id']}, Page {ref.get('page', 'N/A')}): {ref['snippet']}"
+            )