Fix 'Process Document Again' button to only appear for the current document

app.py

@@ -58,7 +58,7 @@ def convert_pdf_to_images(pdf_bytes, dpi=150, rotation=0):
         return []
 
 # Cache preprocessed images for better performance
-@st.cache_data(ttl=24*3600, show_spinner=False)  # Cache for 24 hours
+@st.cache_data(ttl=24*3600, show_spinner=False, hash_funcs={dict: lambda x: str(sorted(x.items()))})  # Cache for 24 hours
 def preprocess_image(image_bytes, preprocessing_options):
     """Preprocess image with selected options optimized for historical document OCR quality"""
     # Setup basic console logging
@@ -175,7 +175,7 @@ def preprocess_image(image_bytes, preprocessing_options):
 
 # Cache OCR results in memory to speed up repeated processing
 @st.cache_data(ttl=24*3600, max_entries=20, show_spinner=False)
-def process_file_cached(file_path, file_type, use_vision, file_size_mb, cache_key):
+def process_file_cached(file_path, file_type, use_vision, file_size_mb, cache_key, preprocessing_options_hash=None):
     """Cached version of OCR processing to reuse results"""
     # Initialize OCR processor
     processor = StructuredOCR()
@@ -241,6 +241,8 @@ def process_file(uploaded_file, use_vision=True, preprocessing_options=None, pro
         with tempfile.NamedTemporaryFile(delete=False, suffix=file_ext) as tmp:
             tmp.write(file_bytes)
             temp_path = tmp.name
+            # Track temporary file for cleanup
+            st.session_state.temp_file_paths.append(temp_path)
         
         # Get PDF rotation value if available and file is a PDF
         pdf_rotation_value = pdf_rotation if 'pdf_rotation' in locals() and file_type == "pdf" else 0
@@ -284,11 +286,34 @@ def process_file(uploaded_file, use_vision=True, preprocessing_options=None, pro
                 # Generate cache key
                 import hashlib
                 file_hash = hashlib.md5(file_bytes).hexdigest()
-                cache_key = f"{file_hash}_{file_type}_{use_vision}_{pdf_rotation_value}"
+                
+                # Include preprocessing options in cache key if available
+                preprocessing_options_hash = ""
+                if 'preprocessing_options' in locals() and preprocessing_options:
+                    # Add pdf_rotation to preprocessing options to ensure it's part of the cache key
+                    if pdf_rotation_value != 0:
+                        preprocessing_options_with_rotation = preprocessing_options.copy()
+                        preprocessing_options_with_rotation['pdf_rotation'] = pdf_rotation_value
+                        preprocessing_str = str(sorted(preprocessing_options_with_rotation.items()))
+                    else:
+                        preprocessing_str = str(sorted(preprocessing_options.items()))
+                    preprocessing_options_hash = hashlib.md5(preprocessing_str.encode()).hexdigest()
+                elif pdf_rotation_value != 0:
+                    # If no preprocessing options but we have rotation, include that in the hash
+                    preprocessing_options_hash = hashlib.md5(f"pdf_rotation_{pdf_rotation_value}".encode()).hexdigest()
+                
+                cache_key = f"{file_hash}_{file_type}_{use_vision}_{preprocessing_options_hash}"
+                
+                # Check if we have custom prompt to include in cache key
+                has_custom_prompt = 'custom_prompt' in locals() and custom_prompt and len(str(custom_prompt).strip()) > 0
+                if has_custom_prompt:
+                    # Update cache key to include custom prompt hash
+                    custom_prompt_hash = hashlib.md5(str(custom_prompt).encode()).hexdigest()
+                    cache_key = f"{cache_key}_{custom_prompt_hash}"
                 
                 # Process with cached function if possible
                 try:
-                    result = process_file_cached(temp_path, file_type, use_vision, file_size_mb, cache_key)
+                    result = process_file_cached(temp_path, file_type, use_vision, file_size_mb, cache_key, preprocessing_options_hash)
                     progress_bar.progress(90)
                     status_text.markdown('<div class="processing-status-container">Finalizing results...</div>', unsafe_allow_html=True)
                 except Exception as e:
@@ -343,7 +368,12 @@ def process_file(uploaded_file, use_vision=True, preprocessing_options=None, pro
                     # Clean up original temp file and use the processed one
                     if os.path.exists(temp_path):
                         os.unlink(temp_path)
+                        # Remove original temp path from tracking list
+                        if temp_path in st.session_state.temp_file_paths:
+                            st.session_state.temp_file_paths.remove(temp_path)
                     temp_path = proc_tmp.name
+                    # Track new temporary file for cleanup
+                    st.session_state.temp_file_paths.append(temp_path)
                 progress_bar.progress(30)
             else:
                 progress_bar.progress(30)
@@ -377,19 +407,37 @@ def process_file(uploaded_file, use_vision=True, preprocessing_options=None, pro
             # Add pdf_rotation to cache key if present
             pdf_rotation_value = pdf_rotation if 'pdf_rotation' in locals() else 0
             file_hash = hashlib.md5(open(temp_path, 'rb').read()).hexdigest()
-            cache_key = f"{file_hash}_{file_type}_{use_vision}_{pdf_rotation_value}"
+            
+            # Include preprocessing options in cache key to ensure reprocessing when options change
+            preprocessing_options_hash = ""
+            if preprocessing_options:
+                # Add pdf_rotation to preprocessing options to ensure it's part of the cache key
+                if pdf_rotation_value != 0:
+                    preprocessing_options_with_rotation = preprocessing_options.copy()
+                    preprocessing_options_with_rotation['pdf_rotation'] = pdf_rotation_value
+                    preprocessing_str = str(sorted(preprocessing_options_with_rotation.items()))
+                else:
+                    preprocessing_str = str(sorted(preprocessing_options.items()))
+                preprocessing_options_hash = hashlib.md5(preprocessing_str.encode()).hexdigest()
+            
+            cache_key = f"{file_hash}_{file_type}_{use_vision}_{preprocessing_options_hash}"
             
             progress_bar.progress(50)
             # Check if we have custom instructions
             has_custom_prompt = 'custom_prompt' in locals() and custom_prompt and len(str(custom_prompt).strip()) > 0
+            
+            # If we have custom instructions, include them in cache key
             if has_custom_prompt:
                 status_text.markdown('<div class="processing-status-container">Processing document with custom instructions...</div>', unsafe_allow_html=True)
+                # Update cache key to include custom prompt hash
+                custom_prompt_hash = hashlib.md5(str(custom_prompt).encode()).hexdigest()
+                cache_key = f"{cache_key}_{custom_prompt_hash}"
             else:
                 status_text.markdown('<div class="processing-status-container">Processing document with OCR...</div>', unsafe_allow_html=True)
             
             # Process the file using cached function if possible
             try:
-                result = process_file_cached(temp_path, file_type, use_vision, file_size_mb, cache_key)
+                result = process_file_cached(temp_path, file_type, use_vision, file_size_mb, cache_key, preprocessing_options_hash)
                 progress_bar.progress(80)
                 status_text.markdown('<div class="processing-status-container">Analyzing document structure...</div>', unsafe_allow_html=True)
                 progress_bar.progress(90)
@@ -471,6 +519,40 @@ except ImportError:
 if 'previous_results' not in st.session_state:
     st.session_state.previous_results = []
 
+# Initialize temp file tracking
+if 'temp_file_paths' not in st.session_state:
+    st.session_state.temp_file_paths = []
+    
+# Initialize last processed file tracking to fix "Process Document Again" button
+if 'last_processed_file' not in st.session_state:
+    st.session_state.last_processed_file = None
+
+# Check if we need to perform a complete reset (coming from "Close Document" button)
+if 'perform_reset' in st.session_state and st.session_state.perform_reset:
+    # List of all session state keys that should be reset, except previous_results
+    reset_keys = [key for key in list(st.session_state.keys()) 
+                 if key != 'previous_results']
+    
+    # Remove all keys except previous_results
+    for key in reset_keys:
+        if key == 'perform_reset':
+            st.session_state[key] = False  # Clear this flag
+        else:
+            st.session_state.pop(key, None)
+    
+    # Reinitialize required session state variables
+    st.session_state.auto_process_sample = False
+    st.session_state.sample_just_loaded = False
+    st.session_state.processed_document_active = False
+    st.session_state.sample_document_processed = False
+    st.session_state.last_processed_file = None
+    
+    # Explicitly reset document-related variables
+    st.session_state.sample_document = None
+    st.session_state.original_sample_bytes = None
+    st.session_state.original_sample_name = None
+    st.session_state.is_sample_document = False
+
 # Create main layout with tabs and columns
 main_tab1, main_tab2, main_tab3 = st.tabs(["Document Processing", "Previous Results", "About"])
 
@@ -1243,32 +1325,47 @@ with main_tab1:
     </style>
     """, unsafe_allow_html=True)
         
-    # Use uploaded_file or sample_document if available
+    # Check if we're using a sample document (either newly loaded or from session state)
+    using_sample_document = False
+    sample_document_name = None
+    
+    # Check for newly loaded sample document
     if 'sample_document' in st.session_state and st.session_state.sample_document is not None:
         # Use the sample document
         uploaded_file = st.session_state.sample_document
-        # Add a notice about using sample document with better style
+        using_sample_document = True
+        sample_document_name = uploaded_file.name
+        
+        # Set auto-process flag in session state if this is a newly loaded sample
+        if st.session_state.sample_just_loaded:
+            st.session_state.auto_process_sample = True
+            # Mark that this is a sample document being processed
+            st.session_state.sample_document_processed = True
+            st.session_state.sample_just_loaded = False
+            
+        # Store sample document bytes in a separate session state variable for potential reprocessing
+        st.session_state.original_sample_bytes = uploaded_file.getvalue()
+        st.session_state.original_sample_name = uploaded_file.name
+        st.session_state.is_sample_document = True
+    
+    # Check for reprocessing of previously loaded sample
+    elif 'is_sample_document' in st.session_state and st.session_state.is_sample_document:
+        using_sample_document = True
+        sample_document_name = st.session_state.original_sample_name if 'original_sample_name' in st.session_state else "Sample Document"
+    
+    # Display sample document notice if using a sample document
+    if using_sample_document:
         st.markdown(
             f"""
             <div style="background-color: #D4EDDA; color: #155724; padding: 10px; 
                  border-radius: 4px; border-left: 5px solid #155724; margin-bottom: 10px;">
                 <div style="display: flex; justify-content: space-between; align-items: center;">
-                    <span style="font-weight: bold;">Sample Document: {uploaded_file.name}</span>
+                    <span style="font-weight: bold;">Sample Document: {sample_document_name}</span>
                 </div>
             </div>
             """, 
             unsafe_allow_html=True
         )
-        
-        # Set auto-process flag in session state if this is a newly loaded sample
-        if st.session_state.sample_just_loaded:
-            st.session_state.auto_process_sample = True
-            # Mark that this is a sample document being processed
-            st.session_state.sample_document_processed = True
-            st.session_state.sample_just_loaded = False
-            
-        # Clear sample document after use to avoid interference with future uploads
-        st.session_state.sample_document = None
     
     if uploaded_file is not None:
         # Check file size (cap at 50MB)
@@ -1278,6 +1375,12 @@ with main_tab1:
             with left_col:
                 st.error(f"File too large ({file_size_mb:.1f} MB). Maximum file size is 50MB.")
             st.stop()
+            
+        # Check if this is a new file (different from the last processed file)
+        current_file_identifier = f"{uploaded_file.name}_{len(uploaded_file.getvalue())}"
+        if st.session_state.last_processed_file != current_file_identifier:
+            # Reset processed_document_active if a new file is uploaded
+            st.session_state.processed_document_active = False
         
         file_ext = Path(uploaded_file.name).suffix.lower()
         
@@ -1288,6 +1391,32 @@ with main_tab1:
             # Make the button more clear about its function
             if st.session_state.processed_document_active:
                 process_button = st.button("Process Document Again")
+                
+                # If process button is clicked and we're dealing with a sample document that was cleared
+                if process_button and uploaded_file is None and 'original_sample_bytes' in st.session_state:
+                    # Recreate the uploaded file from stored bytes
+                    from io import BytesIO
+                    import mimetypes
+                    
+                    # Determine mime type based on file extension
+                    file_ext = os.path.splitext(st.session_state.original_sample_name)[1].lower()
+                    if file_ext == '.pdf':
+                        mime_type = 'application/pdf'
+                    elif file_ext in ['.jpg', '.jpeg']:
+                        mime_type = 'image/jpeg'
+                    elif file_ext == '.png':
+                        mime_type = 'image/png'
+                    else:
+                        mime_type = mimetypes.guess_type(st.session_state.original_sample_name)[0] or 'application/octet-stream'
+                    
+                    # Create a synthetic file-like object with the same interface as UploadedFile
+                    uploaded_file = type('obj', (object,), {
+                        'name': st.session_state.original_sample_name,
+                        'getvalue': lambda: st.session_state.original_sample_bytes,
+                        'read': lambda: st.session_state.original_sample_bytes,
+                        'seek': lambda x: None,
+                        'type': mime_type
+                    })
             else:
                 process_button = st.button("Process Document")
             
@@ -1392,6 +1521,8 @@ with main_tab1:
                                 with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
                                     tmp.write(uploaded_file.getvalue())
                                     temp_path = tmp.name
+                                    # Track temporary file for cleanup
+                                    st.session_state.temp_file_paths.append(temp_path)
                                 
                                 # Apply PDF rotation if specified
                                 pdf_rotation_value = pdf_rotation if 'pdf_rotation' in locals() else 0
@@ -1451,6 +1582,9 @@ with main_tab1:
                                 # Clean up temp file
                                 if os.path.exists(temp_path):
                                     os.unlink(temp_path)
+                                    # Remove from tracking list
+                                    if temp_path in st.session_state.temp_file_paths:
+                                        st.session_state.temp_file_paths.remove(temp_path)
                                     
                             except Exception as e:
                                 # If anything fails, revert to standard processing
@@ -1460,6 +1594,9 @@ with main_tab1:
                         # For non-PDF files, use normal processing with custom prompt
                         # Save the uploaded file to a temporary file with preprocessing
                         with tempfile.NamedTemporaryFile(delete=False, suffix=Path(uploaded_file.name).suffix) as tmp:
+                            temp_path = tmp.name
+                            # Track temporary file for cleanup
+                            st.session_state.temp_file_paths.append(temp_path)
                             # Apply preprocessing if any options are selected
                             if any(preprocessing_options.values()):
                                 # Apply performance mode settings
@@ -1473,7 +1610,6 @@ with main_tab1:
                                 tmp.write(processed_bytes)
                             else:
                                 tmp.write(uploaded_file.getvalue())
-                            temp_path = tmp.name
                         
                         # Show progress
                         with progress_placeholder.container():
@@ -1522,6 +1658,9 @@ with main_tab1:
                     # Clean up temporary file
                     if os.path.exists(temp_path):
                         try:
+                            # Remove from tracking list
+                            if temp_path in st.session_state.temp_file_paths:
+                                st.session_state.temp_file_paths.remove(temp_path)
                             os.unlink(temp_path)
                         except:
                             pass
@@ -2104,6 +2243,10 @@ with main_tab1:
                 # Set processed_document_active to True when a new document is processed
                 st.session_state.processed_document_active = True
                 
+                # Store information about this processed file to track when new files are uploaded
+                if uploaded_file is not None:
+                    st.session_state.last_processed_file = f"{uploaded_file.name}_{len(uploaded_file.getvalue())}"
+                
                 # Button styling is now handled by global CSS
                 
                 # Display success message with close button for dismissing processed documents
@@ -2114,12 +2257,18 @@ with main_tab1:
                     # Close button styling is now handled by global CSS
                     
                     if st.button("✕ Close Document", key="close_document_button", help="Clear current document and start over"):
-                        # Clear the session state
-                        st.session_state.processed_document_active = False
-                        # Reset any active document data
-                        if 'current_result' in st.session_state:
-                            del st.session_state.current_result
-                        # Rerun to reset the page
+                        # Create a special flag to signal a complete reset on the next rerun
+                        st.session_state.perform_reset = True
+                        
+                        # Clean up any temporary files
+                        if 'temp_file_paths' in st.session_state:
+                            for temp_path in st.session_state.temp_file_paths:
+                                try:
+                                    if os.path.exists(temp_path):
+                                        os.remove(temp_path)
+                                except Exception:
+                                    pass  # Ignore errors in cleanup
+                        
                         st.rerun()
                     
                 # Store the result in the previous results list