Enhance Streamlit app to support original JSON and YAML display

- Added buttons to show original document structures in JSON and YAML formats.
- Updated the UI to handle the display and download options for JSON and YAML content.
- Modified the document processing logic to include YAML format alongside JSON and markdown.
- Improved the layout for displaying document structures and added format information for better user guidance.

src/processing/document_processor.py

@@ -21,8 +21,10 @@ class DocumentResult:
     file_path: str
     structured_markdown: str
     structured_json: dict
+    structured_yaml: str  # Add YAML format
     redacted_markdown: str
     redacted_json: dict
+    raw_text: str  # Add raw text without preprocessing
 
 @dataclass
 class ProcessingResult:
@@ -30,7 +32,9 @@ class ProcessingResult:
     original_document_md: str
     redacted_document_md: str
     original_document_json: dict
+    original_document_yaml: str  # Add YAML format
     redacted_document_json: dict
+    raw_text: str  # Add raw text without preprocessing
     removed_indices: list  # Add the actual indices that were removed
     input_tokens: int
     output_tokens: int
@@ -103,7 +107,9 @@ def process_document_with_redaction(
         original_document_md=result.structured_markdown,
         redacted_document_md=result.redacted_markdown,
         original_document_json=result.structured_json,
+        original_document_yaml=result.structured_yaml,
         redacted_document_json=result.redacted_json,
+        raw_text=result.raw_text,
         removed_indices=removed_indices,
         input_tokens=total_input_tokens,
         output_tokens=total_output_tokens,
@@ -149,6 +155,11 @@ class DocumentProcessor:
         structured_md = conv_result.document.export_to_markdown()
         structured_text = conv_result.document.export_to_text()
         doc_json = conv_result.document.export_to_dict()
+        
+        # Convert JSON to YAML for display
+        import yaml
+        doc_yaml = yaml.dump(doc_json, default_flow_style=False, allow_unicode=True, sort_keys=False)
+        
         logger.info(f"Extracted document content (text length {len(structured_text)} characters)")
         
         # Use SectionExtractor to remove target sections if provided
@@ -197,8 +208,10 @@ class DocumentProcessor:
             file_path=file_path,
             structured_markdown=structured_md,
             structured_json=doc_json,
+            structured_yaml=doc_yaml,
             redacted_markdown=redacted_md,
-            redacted_json=redacted_json
+            redacted_json=redacted_json,
+            raw_text=structured_text  # Include the raw text
         )
         logger.info(f"Finished processing for file: {file_path}")
         return result

src/streamlit_app.py

@@ -433,7 +433,7 @@ if uploaded_files:
         uploaded_file = next(f for f in uploaded_files if f.name == selected_file)
         
         # Create buttons for different actions
-        col1, col2, col3 = st.columns(3)
+        col1, col2, col3, col4, col5 = st.columns(5)
         
         with col1:
             if st.button("📄 Show Original", type="primary"):
@@ -515,6 +515,54 @@ if uploaded_files:
                     st.session_state.show_original = True
                     st.session_state.show_processed = False
         
+        with col4:
+            if st.button("📄 Show Original JSON", type="secondary"):
+                # Process the document to get original structure (without redaction)
+                if selected_file not in st.session_state.original_structures:
+                    # Save uploaded file to a temporary location
+                    temp_path = save_uploaded_file(uploaded_file, selected_file)
+                    
+                    # Create a DocumentProcessor without section extraction (for original structure)
+                    processor = DocumentProcessor(section_extractor=None)
+                    
+                    # Process the document to get original structure
+                    result = processor.process(temp_path)
+                    st.session_state.original_structures[selected_file] = result.structured_json
+                    # Store the original markdown for comparison
+                    st.session_state.original_structures[f"{selected_file}_markdown"] = result.structured_markdown
+                    # Store the original YAML for comparison
+                    st.session_state.original_structures[f"{selected_file}_yaml"] = result.structured_yaml
+                
+                # Display the original JSON structure
+                st.session_state.show_original = True
+                st.session_state.show_processed = False
+                st.session_state.show_json = True
+                st.session_state.show_yaml = False
+        
+        with col5:
+            if st.button("📄 Show Original YAML", type="secondary"):
+                # Process the document to get original structure (without redaction)
+                if selected_file not in st.session_state.original_structures:
+                    # Save uploaded file to a temporary location
+                    temp_path = save_uploaded_file(uploaded_file, selected_file)
+                    
+                    # Create a DocumentProcessor without section extraction (for original structure)
+                    processor = DocumentProcessor(section_extractor=None)
+                    
+                    # Process the document to get original structure
+                    result = processor.process(temp_path)
+                    st.session_state.original_structures[selected_file] = result.structured_json
+                    # Store the original markdown for comparison
+                    st.session_state.original_structures[f"{selected_file}_markdown"] = result.structured_markdown
+                    # Store the original YAML for comparison
+                    st.session_state.original_structures[f"{selected_file}_yaml"] = result.structured_yaml
+                
+                # Display the original YAML structure
+                st.session_state.show_original = True
+                st.session_state.show_processed = False
+                st.session_state.show_json = False
+                st.session_state.show_yaml = True
+        
         # Show current view status
         if st.session_state.get("show_original", False):
             st.info("📄 Currently viewing: **Original Document Structure**")
@@ -526,13 +574,24 @@ if uploaded_files:
         # Display results based on button clicked
         if st.session_state.get("show_original", False):
             st.markdown("---")
-            st.subheader(f"Original Document Structure - {selected_file}")
+            
+            # Determine what to show based on button clicked
+            show_json = st.session_state.get("show_json", False)
+            show_yaml = st.session_state.get("show_yaml", False)
+            
+            if show_json:
+                st.subheader(f"Original Document Structure (JSON) - {selected_file}")
+            elif show_yaml:
+                st.subheader(f"Original Document Structure (YAML) - {selected_file}")
+            else:
+                st.subheader(f"Original Document Structure (Markdown) - {selected_file}")
             
             # Get the original structure
             original_json = st.session_state.original_structures[selected_file]
             original_markdown = st.session_state.original_structures.get(f"{selected_file}_markdown", "")
+            original_yaml = st.session_state.original_structures.get(f"{selected_file}_yaml", "")
             
-            # Display PDF viewer and original markdown side by side
+            # Display PDF viewer and content side by side
             col1, col2 = st.columns([1, 1])
             
             with col1:
@@ -547,30 +606,86 @@ if uploaded_files:
                 st.markdown(pdf_display, unsafe_allow_html=True)
             
             with col2:
-                st.subheader("📋 Original Document (Markdown)")
-                st.caption("Docling-generated markdown from the PDF")
-                # Use a text area for better readability and scrolling
-                st.text_area(
-                    label="Original markdown content",
-                    value=original_markdown,
-                    height=600,
-                    key="original_markdown_display",
-                    label_visibility="collapsed"
-                )
+                if show_json:
+                    st.subheader("📋 Original Document (JSON)")
+                    st.caption("Docling-generated JSON structure from the PDF")
+                    # Use a text area for better readability and scrolling
+                    st.text_area(
+                        label="Original JSON content",
+                        value=json.dumps(original_json, indent=2, ensure_ascii=False),
+                        height=600,
+                        key="original_json_display",
+                        label_visibility="collapsed"
+                    )
+                elif show_yaml:
+                    st.subheader("📋 Original Document (YAML)")
+                    st.caption("Docling-generated YAML structure from the PDF")
+                    # Use a text area for better readability and scrolling
+                    st.text_area(
+                        label="Original YAML content",
+                        value=original_yaml,
+                        height=600,
+                        key="original_yaml_display",
+                        label_visibility="collapsed"
+                    )
+                else:
+                    st.subheader("📋 Original Document (Markdown)")
+                    st.caption("Docling-generated markdown from the PDF")
+                    # Use a text area for better readability and scrolling
+                    st.text_area(
+                        label="Original markdown content",
+                        value=original_markdown,
+                        height=600,
+                        key="original_markdown_display",
+                        label_visibility="collapsed"
+                    )
             
-            # Add a download button for the original markdown
+            # Add download buttons for the original content
             st.markdown("---")
-            col1, col2 = st.columns(2)
+            col1, col2, col3 = st.columns(3)
             with col1:
-                st.download_button(
-                    label="📥 Download Original Markdown",
-                    data=original_markdown,
-                    file_name=f"{selected_file}_original.md",
-                    mime="text/markdown"
-                )
+                if show_json:
+                    st.download_button(
+                        label="📥 Download Original JSON",
+                        data=json.dumps(original_json, indent=2, ensure_ascii=False),
+                        file_name=f"{selected_file}_original.json",
+                        mime="application/json"
+                    )
+                elif show_yaml:
+                    st.download_button(
+                        label="📥 Download Original YAML",
+                        data=original_yaml,
+                        file_name=f"{selected_file}_original.yaml",
+                        mime="text/yaml"
+                    )
+                else:
+                    st.download_button(
+                        label="📥 Download Original Markdown",
+                        data=original_markdown,
+                        file_name=f"{selected_file}_original.md",
+                        mime="text/markdown"
+                    )
             with col2:
-                st.subheader("📊 JSON Structure")
-                st.json(original_json)
+                if show_json or show_yaml:
+                    st.subheader("📊 Document Structure")
+                    st.json(original_json)
+                else:
+                    st.subheader("📊 JSON Structure")
+                    st.json(original_json)
+            with col3:
+                if show_json or show_yaml:
+                    # Show format information
+                    st.subheader("📋 Format Info")
+                    if show_json:
+                        st.info("**JSON Format**: Structured data representation with key-value pairs")
+                        st.write("**Use case**: API integration, data processing, programmatic access")
+                    elif show_yaml:
+                        st.info("**YAML Format**: Human-readable data serialization")
+                        st.write("**Use case**: Configuration files, documentation, easy reading")
+                else:
+                    st.subheader("📋 Markdown Info")
+                    st.info("**Markdown Format**: Formatted text with headers, lists, and styling")
+                    st.write("**Use case**: Documentation, readable output, web display")
         
         elif st.session_state.get("show_processed", False):
             st.markdown("---")