Update app.py

app.py

@@ -20,11 +20,13 @@ import requests # For requests.exceptions.HTTPError
 from requests.exceptions import HTTPError as RequestsHTTPError # Specific import for clarity
 
 import pdfplumber
-from pdf2image import convert_from_path, convert_from_bytes
+import pdf2image # <<<<<<<<<<<<<<<< CORRECTED: Added this import
+from pdf2image import convert_from_path, convert_from_bytes # Keep these for direct use too
+# from pdf2image.exceptions import ... # If you need to catch specific pdf2image errors
+
 import pytesseract
 from PIL import Image
 from huggingface_hub import HfApi, create_repo
-# from huggingface_hub.utils import HfHubHTTPError # This was the incorrect one
 
 # --- Flask App Initialization ---
 app = Flask(__name__)
@@ -69,34 +71,28 @@ def ensure_hf_dataset():
         logger.warning(msg)
         return "Error: " + msg
     try:
-        # create_repo can raise huggingface_hub.utils.RepositoryNotFoundError, 
-        # huggingface_hub.utils.HfHubHTTPError (which inherits from requests.HTTPError for some cases),
-        # or other requests.exceptions
         repo_id_obj = create_repo(repo_id=HF_DATASET_REPO_NAME, token=HF_TOKEN, repo_type="dataset", exist_ok=True)
         logger.info(f"Dataset repo ensured: {repo_id_obj.repo_id}")
         return repo_id_obj.repo_id
-    except RequestsHTTPError as e: # Catch HTTP errors from requests library directly
-        if e.response is not None and e.response.status_code == 409: # Conflict, repo already exists
+    except RequestsHTTPError as e: 
+        if e.response is not None and e.response.status_code == 409: 
              logger.info(f"Dataset repo '{HF_DATASET_REPO_NAME}' already exists (HTTP 409).")
-             # Attempt to construct the full repo_id (namespace/repo_name)
              try:
-                 user_info = hf_api.whoami(token=HF_TOKEN) # This call could also fail
+                 user_info = hf_api.whoami(token=HF_TOKEN) 
                  namespace = user_info.get('name') if user_info else None
                  if namespace:
                      return f"{namespace}/{HF_DATASET_REPO_NAME}"
-                 else: # Fallback if namespace cannot be determined
+                 else: 
                      logger.warning(f"Could not determine namespace for existing repo '{HF_DATASET_REPO_NAME}'. Using generic ID.")
-                     return HF_DATASET_REPO_NAME # Or f"{YOUR_DEFAULT_USERNAME_IF_KNOWN}/{HF_DATASET_REPO_NAME}"
+                     return HF_DATASET_REPO_NAME 
              except Exception as whoami_e:
                  logger.error(f"Could not determine namespace for existing repo via whoami due to: {whoami_e}. Using generic ID.")
-                 return HF_DATASET_REPO_NAME # Fallback
-        else: # Other HTTP errors
+                 return HF_DATASET_REPO_NAME 
+        else: 
             status_code = e.response.status_code if e.response is not None else "Unknown"
             logger.error(f"Hugging Face dataset HTTP error (Status: {status_code}): {str(e)}")
             return f"Error: Failed to access or create dataset '{HF_DATASET_REPO_NAME}' due to HTTP error: {str(e)}"
-    except Exception as e: # Catch other non-HTTP exceptions from huggingface_hub or general errors
-        # This could be Hf একাধিক RepoExistsError if exist_ok=False, or other utility errors.
-        # For exist_ok=True, a 409 is the more likely signal for existing repo.
+    except Exception as e: 
         logger.error(f"Hugging Face dataset general error: {str(e)}", exc_info=True)
         return f"Error: Failed to access or create dataset '{HF_DATASET_REPO_NAME}': {str(e)}"
 
@@ -112,6 +108,8 @@ def upload_image_to_hf_stream(image_pil, filename_base, page_num_for_log=""):
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
         repo_filename = f"images/{filename_base}_{page_num_for_log}_{timestamp}.png"
         
+        # Ensure UPLOAD_FOLDER exists before writing temp file
+        os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
         with tempfile.NamedTemporaryFile(delete=False, suffix=".png", dir=app.config['UPLOAD_FOLDER']) as tmp_file:
             temp_image_path = tmp_file.name
         image_pil.save(temp_image_path, format="PNG")
@@ -123,7 +121,7 @@ def upload_image_to_hf_stream(image_pil, filename_base, page_num_for_log=""):
         )
         logger.info(f"Successfully uploaded image: {file_url}")
         return file_url
-    except Exception as e: # Catch broadly here; specific HF errors could be caught if needed
+    except Exception as e: 
         logger.error(f"Image upload error for {filename_base}{page_num_for_log}: {str(e)}", exc_info=True)
         return f"Error uploading image {filename_base}{page_num_for_log}: {str(e)}"
     finally:
@@ -177,17 +175,17 @@ def generate_pdf_conversion_stream(pdf_input_source_path_or_url):
 
         if source_is_url:
             try:
-                response = requests.get(pdf_input_source_path_or_url, stream=False, timeout=60) # stream=False to get content
+                response = requests.get(pdf_input_source_path_or_url, stream=False, timeout=60) 
                 response.raise_for_status()
                 pdf_bytes_for_images = response.content 
                 pdf_handle_for_text = io.BytesIO(pdf_bytes_for_images) 
                 yield yield_message("status", {"message": f"PDF downloaded from URL ({len(pdf_bytes_for_images)/1024:.2f} KB)."})
                 time.sleep(0.01)
-            except RequestsHTTPError as e: # Catch HTTP errors specifically
+            except RequestsHTTPError as e: 
                 logger.error(f"URL fetch HTTP error for PDF processing: {str(e)} (Status: {e.response.status_code if e.response else 'N/A'})", exc_info=True)
                 yield yield_message("error", {"message": f"Error fetching PDF from URL (HTTP {e.response.status_code if e.response else 'N/A'}): {e.response.reason if e.response else str(e)}"})
                 return
-            except requests.RequestException as e: # Catch other network errors
+            except requests.RequestException as e: 
                 logger.error(f"URL fetch network error for PDF processing: {str(e)}", exc_info=True)
                 yield yield_message("error", {"message": f"Network error fetching PDF from URL: {str(e)}"})
                 return 
@@ -211,7 +209,7 @@ def generate_pdf_conversion_stream(pdf_input_source_path_or_url):
                     tables = page.extract_tables()
                     if tables:
                         for table_idx, table_data in enumerate(tables):
-                            if table_data and len(table_data) > 0 and len(table_data[0]) > 0 : # Check table has rows and columns
+                            if table_data and len(table_data) > 0 and table_data[0] is not None and len(table_data[0]) > 0 : 
                                 yield yield_message("status", {"message": f"  Processing table {table_idx+1} on page {i+1}..."})
                                 header_cells = table_data[0]
                                 header = [" | ".join(str(cell) if cell is not None else "" for cell in header_cells)]
@@ -229,9 +227,7 @@ def generate_pdf_conversion_stream(pdf_input_source_path_or_url):
         except Exception as e:
             logger.error(f"Error during PDF text/table extraction: {str(e)}", exc_info=True)
             yield yield_message("error", {"message": f"Error during text extraction: {str(e)}"})
-            # Decide if to return or continue to image extraction. Let's try to continue.
 
-        # 2. Image Extraction and OCR
         if not check_poppler():
             yield yield_message("error", {"message": "Poppler (for image extraction) not found or not working."})
         else:
@@ -241,97 +237,107 @@ def generate_pdf_conversion_stream(pdf_input_source_path_or_url):
                  yield yield_message("markdown_chunk", {"content": "**Note:** `HF_TOKEN` not set. Images will be described but not uploaded.\n\n"})
 
             time.sleep(0.01)
-            extracted_pil_images = []
+            extracted_pil_images_overall_count = 0 # Keep track of total images processed for numbering
             try:
                 image_source_for_convert = None
                 if source_is_url and pdf_bytes_for_images:
                     image_source_for_convert = pdf_bytes_for_images
                     logger.info("Using downloaded bytes for image conversion.")
                 elif not source_is_url: 
-                    image_source_for_convert = pdf_input_source_path_or_url # Local file path
+                    image_source_for_convert = pdf_input_source_path_or_url 
                     logger.info("Using local file path for image conversion.")
                 
                 if image_source_for_convert:
-                    # Attempt to get page count for more granular image processing if pdf2image is the bottleneck
                     try:
                         pdf_info = None
                         if isinstance(image_source_for_convert, bytes):
                             pdf_info = pdf2image.pdfinfo_from_bytes(image_source_for_convert, userpw=None, poppler_path=None)
-                        else: # path
+                        else: 
                             pdf_info = pdf2image.pdfinfo_from_path(image_source_for_convert, userpw=None, poppler_path=None)
                         
                         num_image_pages = pdf_info.get("Pages", 0)
                         yield yield_message("status", {"message": f"PDF has {num_image_pages} page(s) for potential image extraction."})
                         
-                        # Process images page by page (or small batches) to save memory
-                        batch_size = 1 # Process one page at a time for images
+                        batch_size = 1 
                         for page_idx_start in range(1, num_image_pages + 1, batch_size):
                             page_idx_end = min(page_idx_start + batch_size - 1, num_image_pages)
-                            yield yield_message("status", {"message": f"Extracting images from page(s) {page_idx_start}-{page_idx_end}..."})
+                            yield yield_message("status", {"message": f"Extracting images from PDF page(s) {page_idx_start}-{page_idx_end}..."})
                             time.sleep(0.01)
                             
                             page_images_pil = []
                             if isinstance(image_source_for_convert, bytes):
                                 page_images_pil = convert_from_bytes(image_source_for_convert, dpi=150, first_page=page_idx_start, last_page=page_idx_end)
-                            else: # path
+                            else: 
                                 page_images_pil = convert_from_path(image_source_for_convert, dpi=150, first_page=page_idx_start, last_page=page_idx_end)
                             
-                            extracted_pil_images.extend(page_images_pil) # Add to overall list for sequential numbering later
-                            
-                            # Process this batch of images immediately
-                            for img_pil in page_images_pil:
-                                current_image_index = len(extracted_pil_images) # Current overall index
-                                page_num_for_log = f"page_{page_idx_start + page_images_pil.index(img_pil)}"
-                                yield yield_message("status", {"message": f"Processing image {current_image_index} (from PDF page {page_num_for_log}) (OCR & Upload)..."})
+                            for img_idx_in_batch, img_pil in enumerate(page_images_pil):
+                                extracted_pil_images_overall_count += 1
+                                current_pdf_page_num = page_idx_start + img_idx_in_batch # Actual PDF page number
+                                page_num_for_log = f"pdfpage_{current_pdf_page_num}"
+
+                                yield yield_message("status", {"message": f"Processing image {extracted_pil_images_overall_count} (from PDF page {current_pdf_page_num}) (OCR & Upload)..."})
                                 time.sleep(0.01)
                                 
                                 ocr_text = ""
                                 try:
                                     ocr_text = pytesseract.image_to_string(img_pil).strip()
-                                    if ocr_text: yield yield_message("status", {"message": f"  OCR successful for image {current_image_index}."})
+                                    if ocr_text: yield yield_message("status", {"message": f"  OCR successful for image {extracted_pil_images_overall_count}."})
                                 except Exception as ocr_e:
-                                    logger.error(f"OCR error for image {current_image_index}: {str(ocr_e)}")
+                                    logger.error(f"OCR error for image {extracted_pil_images_overall_count}: {str(ocr_e)}")
                                     ocr_text = f"OCR failed: {str(ocr_e)}"
                                 
                                 image_md_chunk = ""
                                 if HF_TOKEN:
                                     image_url_or_error = upload_image_to_hf_stream(img_pil, "pdf_image", page_num_for_log)
                                     if isinstance(image_url_or_error, str) and not image_url_or_error.startswith("Error"):
-                                        image_md_chunk += f"![Image {current_image_index}]({image_url_or_error})\n"
-                                        yield yield_message("status", {"message": f"  Image {current_image_index} uploaded."})
+                                        image_md_chunk += f"![Image {extracted_pil_images_overall_count}]({image_url_or_error})\n"
+                                        yield yield_message("status", {"message": f"  Image {extracted_pil_images_overall_count} uploaded."})
                                     else:
-                                        image_md_chunk += f"**Image {current_image_index} (Upload Error):** {str(image_url_or_error)}\n\n"
-                                        yield yield_message("error", {"message": f"Failed to upload image {current_image_index}: {str(image_url_or_error)}"})
+                                        image_md_chunk += f"**Image {extracted_pil_images_overall_count} (Upload Error):** {str(image_url_or_error)}\n\n"
+                                        yield yield_message("error", {"message": f"Failed to upload image {extracted_pil_images_overall_count}: {str(image_url_or_error)}"})
                                 else:
-                                    image_md_chunk += f"**Image {current_image_index} (not uploaded due to missing HF_TOKEN)**\n"
+                                    image_md_chunk += f"**Image {extracted_pil_images_overall_count} (not uploaded due to missing HF_TOKEN)**\n"
 
                                 if ocr_text:
-                                    image_md_chunk += f"**Image {current_image_index} OCR Text:**\n```\n{ocr_text}\n```\n\n"
+                                    image_md_chunk += f"**Image {extracted_pil_images_overall_count} OCR Text:**\n```\n{ocr_text}\n```\n\n"
                                 
                                 yield yield_message("image_md", {"content": image_md_chunk})
                                 time.sleep(0.01)
                     except Exception as e_img_info:
                         logger.error(f"Could not get PDF info for image batching or during batched conversion: {e_img_info}", exc_info=True)
                         yield yield_message("error", {"message": f"Error preparing for image extraction: {e_img_info}. Trying bulk."})
-                        # Fallback to bulk conversion if pdfinfo or batching fails (original behavior)
+                        # Fallback to bulk conversion
+                        bulk_images_pil = []
                         if isinstance(image_source_for_convert, bytes):
-                            extracted_pil_images = convert_from_bytes(image_source_for_convert, dpi=150)
-                        else: # path
-                            extracted_pil_images = convert_from_path(image_source_for_convert, dpi=150)
-                        # Process these bulk images (copy-paste the loop from above, adjust indexing)
-                        for i, img_pil in enumerate(extracted_pil_images):
-                            page_num_for_log = f"bulk_image_{i+1}"
-                            yield yield_message("status", {"message": f"Processing image {i+1}/{len(extracted_pil_images)} (OCR & Upload)..."}) # ... (rest of loop) ...
-                            # (omitted rest of duplicated loop for brevity, but it would be the same as the inner loop above)
-                            ocr_text = pytesseract.image_to_string(img_pil).strip() # Simplified for brevity
-                            image_md_chunk = f"![Image {i+1} Fallback]\n**OCR:** {ocr_text}\n\n"
+                            bulk_images_pil = convert_from_bytes(image_source_for_convert, dpi=150)
+                        else: 
+                            bulk_images_pil = convert_from_path(image_source_for_convert, dpi=150)
+                        
+                        yield yield_message("status", {"message": f"Fallback: Extracted {len(bulk_images_pil)} images in bulk."})
+                        for i, img_pil in enumerate(bulk_images_pil):
+                            extracted_pil_images_overall_count +=1
+                            page_num_for_log = f"bulk_image_{i+1}" # Less precise page info in fallback
+                            yield yield_message("status", {"message": f"Processing image {extracted_pil_images_overall_count} (bulk) (OCR & Upload)..."})
+                            ocr_text = ""
+                            try: ocr_text = pytesseract.image_to_string(img_pil).strip()
+                            except Exception as e: ocr_text = f"OCR Error: {e}"
+                            
+                            image_md_chunk = f"![Image {extracted_pil_images_overall_count} (Fallback)]\n"
+                            if HF_TOKEN:
+                                image_url_or_error = upload_image_to_hf_stream(img_pil, "pdf_image_fallback", page_num_for_log)
+                                if isinstance(image_url_or_error, str) and not image_url_or_error.startswith("Error"):
+                                    image_md_chunk = f"![Image {extracted_pil_images_overall_count} (Fallback)]({image_url_or_error})\n"
+                                else:
+                                    image_md_chunk += f"**Upload Error:** {str(image_url_or_error)}\n"
+                            if ocr_text: image_md_chunk += f"**OCR Text:**\n```\n{ocr_text}\n```\n\n"
+                            else: image_md_chunk += "\n"
                             yield yield_message("image_md", {"content": image_md_chunk})
                             time.sleep(0.01)
 
-                else: # No valid source for image conversion
-                    yield yield_message("status", {"message": "No valid source (URL download failed or no file path) for image extraction."})
+                else: 
+                    yield yield_message("status", {"message": "No valid source for image extraction."})
 
-            except Exception as e: # Catch errors from the image extraction block
+            except Exception as e: 
                 logger.error(f"Error during image extraction/processing: {str(e)}", exc_info=True)
                 yield yield_message("error", {"message": f"Error during image extraction: {str(e)}"})
         
@@ -353,15 +359,10 @@ def process_pdf_stream():
     pdf_file = request.files.get('pdf_file')
     pdf_url = request.form.get('pdf_url', '').strip()
     
-    # Use a list to hold temp_pdf_path so it can be modified in the inner function
-    # and accessed in finally. Or pass it around.
-    # For simplicity, we'll rely on the generator's finally block if it's created within.
-    # Here, temp_pdf_path is primarily for the *uploaded* file before passing its path.
-    
-    outer_temp_pdf_path = None # For uploaded file cleanup
+    outer_temp_pdf_path = None 
 
     def stream_processor():
-        nonlocal outer_temp_pdf_path # Make it accessible in this inner function for cleanup
+        nonlocal outer_temp_pdf_path 
         pdf_input_source_for_generator = None
 
         try:
@@ -375,7 +376,7 @@ def process_pdf_stream():
                 fd, temp_path = tempfile.mkstemp(suffix=".pdf", prefix="upload_", dir=app.config['UPLOAD_FOLDER'])
                 os.close(fd) 
                 pdf_file.save(temp_path)
-                outer_temp_pdf_path = temp_path # Store for cleanup
+                outer_temp_pdf_path = temp_path 
                 logger.info(f"Uploaded PDF saved to temporary path: {outer_temp_pdf_path}")
                 pdf_input_source_for_generator = outer_temp_pdf_path
                 yield yield_message("status", {"message": f"Processing uploaded PDF: {filename}"})
@@ -400,17 +401,11 @@ def process_pdf_stream():
         except Exception as e:
             logger.error(f"Error setting up stream or in initial validation: {str(e)}", exc_info=True)
             yield yield_message("error", {"message": f"Setup error: {str(e)}"})
-        # The 'finally' block for cleaning outer_temp_pdf_path will be outside this generator,
-        # in the main route function after the Response is fully generated.
-        # However, with stream_with_context, the 'finally' here is better.
         finally:
             if outer_temp_pdf_path and os.path.exists(outer_temp_pdf_path):
                 try:
                     os.remove(outer_temp_pdf_path)
                     logger.info(f"Cleaned up temporary PDF: {outer_temp_pdf_path}")
-                    # Yielding from finally inside a generator that's part of a streamed response can be tricky.
-                    # It's better if status messages about cleanup are logged or handled differently.
-                    # For this case, logging is sufficient.
                 except OSError as ose:
                     logger.error(f"Error removing temporary PDF {outer_temp_pdf_path}: {ose}")