Update src/load_file.py

src/load_file.py

@@ -1,368 +1,394 @@
 import streamlit as st
 import pymongo
 from datetime import datetime, date, timezone
-import os
-from typing import List, Tuple, Optional
+from typing import List, Tuple, Optional, Dict
 import requests
-from PIL import Image
+from PIL import Image, ImageFile
 import io
 from dotenv import load_dotenv
 import concurrent.futures
 import threading
 from functools import lru_cache
-
+import hashlib
+import contextlib
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
 
 load_dotenv()
+ImageFile.LOAD_TRUNCATED_IMAGES = True
 
+@st.cache_data(show_spinner=False, ttl=60 * 30)
+def _download_bytes_cached(url: str, timeout_s: float = 12.0) -> Optional[bytes]:
+    """Download image bytes from URL with Streamlit cache (30 min TTL)."""
+    try:
+        r = requests.get(url, timeout=(3.05, timeout_s), stream=True)
+        r.raise_for_status()
+        return r.content
+    except Exception:
+        return None
 
 class ImageGalleryApp:
     def __init__(self, mongo_uri: str, db_name: str, collection_name: str):
-        """Initialize the MongoDB connection"""
+        """Initialize MongoDB and HTTP session with retry pooling."""
         self.client = pymongo.MongoClient(mongo_uri)
         self.db = self.client[db_name]
         self.collection = self.db[collection_name]
-        self._categories_cache = None
-        self._filenames_cache = None
         self._cache_lock = threading.Lock()
 
+        self.session = requests.Session()
+        retries = Retry(
+            total=3,
+            connect=3,
+            read=3,
+            backoff_factor=0.4,
+            status_forcelist=[429, 500, 502, 503, 504],
+            allowed_methods=["GET", "HEAD"]
+        )
+        adapter = HTTPAdapter(pool_connections=64, pool_maxsize=64, max_retries=retries)
+        self.session.mount("http://", adapter)
+        self.session.mount("https://", adapter)
+
+        self.thumb_max_size = (768, 768)
+
     @lru_cache(maxsize=128)
     def get_unique_categories(self) -> List[str]:
-        """Get all unique categories from the database with caching"""
+        """Return all unique categories from DB."""
         try:
-            # Only get categories for completed records to improve performance
-            pipeline = [
-                {"$match": {"status": "completed", "lob": "leadgen_vivek"}},
-                {"$group": {"_id": "$category"}},
-                {"$sort": {"_id": 1}}
-            ]
-            categories = [doc["_id"] for doc in self.collection.aggregate(pipeline) if doc["_id"]]
-            return ["All"] + categories
+            categories = self.collection.distinct(
+                "category",
+                {"status": "completed", "lob": "leadgen_vivek", "category": {"$ne": None}}
+            )
+            return ["All"] + sorted(categories)
         except Exception as e:
             print(f"Error fetching categories: {e}")
             return ["All"]
 
     @lru_cache(maxsize=128)
     def get_unique_filenames(self) -> List[str]:
-        """Get all unique file names from the database with caching"""
+        """Return all unique file names from DB."""
         try:
-            # Only get filenames for completed records
-            pipeline = [
-                {"$match": {"status": "completed", "lob": "leadgen_vivek"}},
-                {"$group": {"_id": "$file_name"}},
-                {"$sort": {"_id": 1}}
-            ]
-            filenames = [doc["_id"] for doc in self.collection.aggregate(pipeline) if doc["_id"]]
-            return ["All"] + filenames
+            filenames = self.collection.distinct(
+                "file_name",
+                {"status": "completed", "lob": "leadgen_vivek", "file_name": {"$ne": None}}
+            )
+            return ["All"] + sorted(filenames)
         except Exception as e:
             print(f"Error fetching filenames: {e}")
             return ["All"]
 
-    def load_image_from_url(self, url: str) -> Optional[Image.Image]:
-        """Load image from URL with better error handling"""
-        try:
-            full_url = url
-
-            response = requests.get(full_url, timeout=5, stream=True)  # Reduced timeout, added streaming
-            response.raise_for_status()
-
-            # Limit image size to prevent memory issues
-            image = Image.open(io.BytesIO(response.content))
-
-            # Resize large images to improve performance
-            max_size = (800, 800)
-            if image.size[0] > max_size[0] or image.size[1] > max_size[1]:
-                image.thumbnail(max_size, Image.Resampling.LANCZOS)
-
-            return image
-        except Exception as e:
-            print(f"Error loading image from {url}: {e}")
-            return None
-
-    def load_images_parallel(self, urls: List[str], max_workers: int = 5) -> List[Image.Image]:
-        """Load multiple images in parallel"""
-        images = []
-        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
-            future_to_url = {executor.submit(self.load_image_from_url, url): url for url in
-                             urls[:]}  # Limit to first 10 URLs
-            for future in concurrent.futures.as_completed(future_to_url):
-                try:
-                    image = future.result()
-                    if image:
-                        images.append(image)
-                except Exception as e:
-                    print(f"Error in parallel loading: {e}")
-        return images
-    
     def parse_date_input(self, date_input) -> Optional[date]:
-        """Convert string date input to date object"""
+        """Convert string or date input to date object."""
         if not date_input or date_input == "":
             return None
-        
         if isinstance(date_input, date):
             return date_input
-        
         if isinstance(date_input, str):
             try:
-                # Handle YYYY-MM-DD format (most common from Calendar component)
                 if date_input.count('-') == 2:
                     return datetime.strptime(date_input, '%Y-%m-%d').date()
-                # Handle MM/DD/YYYY format
                 elif date_input.count('/') == 2:
                     return datetime.strptime(date_input, '%m/%d/%Y').date()
             except ValueError as e:
                 print(f"Error parsing date string '{date_input}': {e}")
                 return None
-        
         return None
 
-    def search_images(self,
-                      category: str = "All",
-                      file_name: str = "All",
-                      start_date: Optional[date] = datetime.now(timezone.utc).date(),
-                      end_date: Optional[date] = datetime.now(timezone.utc).date(),
-                      lob: str = "leadgen_vivek") -> Tuple[List[Image.Image], str]:
-        """Search images based on filters with optimizations"""
-        
-        # Convert string inputs to date objects
-        start_date_obj = self.parse_date_input(start_date)
-        end_date_obj = self.parse_date_input(end_date)
+    def load_image_from_url(self, url: str) -> Optional[Image.Image]:
+        """Load and thumbnail image from URL with caching and pooling."""
+        try:
+            data = _download_bytes_cached(url)
+            if data is None:
+                r = self.session.get(url, timeout=(3.05, 12), stream=True)
+                r.raise_for_status()
+                data = r.content
 
-        # Build query with default status filter
-        query = {
+            img = Image.open(io.BytesIO(data))
+            with contextlib.suppress(Exception):
+                img = img.convert("RGB")
+
+            if img.size[0] > self.thumb_max_size[0] or img.size[1] > self.thumb_max_size[1]:
+                img.thumbnail(self.thumb_max_size, Image.Resampling.LANCZOS)
+
+            return img
+        except Exception as e:
+            print(f"Error loading image from {url}: {e}")
+            return None
+
+    def load_images_parallel(self, urls: List[str], max_workers: int = 8) -> List[Tuple[str, Optional[Image.Image]]]:
+        """Load multiple images in parallel; returns list of (url, image)."""
+        results: List[Tuple[str, Optional[Image.Image]]] = []
+        if not urls:
+            return results
+        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as ex:
+            future_map = {ex.submit(self.load_image_from_url, u): u for u in urls}
+            for fut in concurrent.futures.as_completed(future_map):
+                u = future_map[fut]
+                try:
+                    img = fut.result()
+                except Exception as e:
+                    print(f"Error in parallel loading: {e}")
+                    img = None
+                results.append((u, img))
+        return results
+
+    def search_images_page(
+        self,
+        category: str = "All",
+        file_name: str = "All",
+        start_date: Optional[date] = None,
+        end_date: Optional[date] = None,
+        lob: str = "leadgen_vivek",
+        page: int = 0,
+        page_size: int = 24,
+    ) -> Tuple[List[Dict], int]:
+        """
+        Paginated search for images by filters.
+        Returns:
+          - docs: list of {url, category, file_name, created_at, prompt, status}
+          - total: total number of image URLs matching the filters
+        """
+        match = {
             "lob": lob,
-            "status": "completed",  # Default filter for completed status
-            "urls": {"$exists": True, "$ne": []}  # Only get records with URLs
+            "status": "completed",
+            "urls": {"$exists": True, "$ne": []}
         }
-
         if category != "All":
-            query["category"] = category
-
+            match["category"] = category
         if file_name != "All":
-            query["file_name"] = file_name
+            match["file_name"] = file_name
 
-        # Handle date filtering with proper date objects
-        if start_date or end_date:
+        start_date_obj = self.parse_date_input(start_date)
+        end_date_obj = self.parse_date_input(end_date)
+        if start_date_obj or end_date_obj:
             date_query = {}
             if start_date_obj:
-                try:
-                    # Convert date to datetime at start of day
-                    start_dt = datetime.combine(start_date_obj, datetime.min.time())
-                    date_query["$gte"] = start_dt
-                except Exception as e:
-                    print(f"Error parsing start date: {e}")
-
+                date_query["$gte"] = datetime.combine(start_date_obj, datetime.min.time())
             if end_date_obj:
-                try:
-                    # Convert date to datetime at end of day
-                    end_dt = datetime.combine(end_date_obj, datetime.max.time())
-                    date_query["$lte"] = end_dt
-                except Exception as e:
-                    print(f"Error parsing end date: {e}")
+                date_query["$lte"] = datetime.combine(end_date_obj, datetime.max.time())
             if date_query:
-                query["created_at"] = date_query
+                match["created_at"] = date_query
 
         try:
-            # Use projection to only fetch needed fields
-            projection = {
+            count_pipeline = [
+                {"$match": match},
+                {"$unwind": "$urls"},
+                {"$count": "n"}
+            ]
+            count_doc = list(self.collection.aggregate(count_pipeline))
+            total = count_doc[0]["n"] if count_doc else 0
+        except Exception as e:
+            print(f"Count error: {e}")
+            total = 0
+
+        if total == 0:
+            return [], 0
+
+        pipeline = [
+            {"$match": match},
+            {"$unwind": "$urls"},
+            {"$sort": {"created_at": -1}},
+            {"$skip": max(0, page) * max(1, page_size)},
+            {"$limit": max(1, page_size)},
+            {"$project": {
+                "_id": 0,
+                "url": "$urls",
                 "category": 1,
                 "file_name": 1,
                 "created_at": 1,
-                "urls": 1,
-                "status": 1,
-                "prompt": 1
-            }
-
-            # Execute query with optimizations
-            cursor = self.collection.find(query, projection).sort("created_at", -1)
-            documents = list(cursor)
-
-            if not documents:
-                return [], "No records found matching the criteria."
-
-            # Collect all URLs for parallel loading
-            all_urls = []
-            url_to_doc = {}
-
-            for doc in documents:
-                urls = doc.get("urls", [])
-                if urls:
-                    
-                    for url in urls:
-                        first_url = url
-                        all_urls.append(first_url)
-                        url_to_doc[first_url] = doc
-
-            # Load images in parallel
-            print(f"Loading {len(all_urls)} images...")
-            images = self.load_images_parallel(all_urls, max_workers=8)
-
-            # Build info text
-            info_text = f"Found {len(documents)} records (showing {len(images)} images)\n"
-            info_text += f"Filter: Status = completed, LOB = {lob}\n"
-            if start_date_obj:
-                info_text += f"Start Date: {start_date_obj}\n"
-            if end_date_obj:
-                info_text += f"End Date: {end_date_obj}\n"
-            info_text += "\n"
-
-            for i, doc in enumerate(documents[:len(images)]):
-                info_text += f"#{i + 1}\n"
-                info_text += f"Category: {doc.get('category', 'N/A')}\n"
-                info_text += f"File: {doc.get('file_name', 'N/A')}\n"
-                info_text += f"Prompt: {doc.get('prompt', 'N/A')}\n"
-                info_text += f"Created: {doc.get('created_at', 'N/A')}\n"
-                info_text += f"URLs: {len(doc.get('urls', []))} image(s)\n"
-                info_text += f"Status: {doc.get('status', 'N/A')}\n"
-                info_text += "-" * 30 + "\n"
-
-            return images, info_text
-
+                "prompt": 1,
+                "status": 1
+            }}
+        ]
+        try:
+            docs = list(self.collection.aggregate(pipeline, allowDiskUse=True))
         except Exception as e:
-            error_msg = f"Error searching database: {str(e)}"
-            print(error_msg)
-            return [], error_msg
-        
+            print(f"Aggregation error: {e}")
+            docs = []
+        return docs, total
 
 def create_streamlit_app(mongo_uri: str, db_name: str, collection_name: str):
+    """Main Streamlit UI for image gallery."""
     app = ImageGalleryApp(mongo_uri, db_name, collection_name)
+
     def get_filter_choices():
+        """Fetch filter choices for category and filename."""
         try:
             categories = app.get_unique_categories()
             filenames = app.get_unique_filenames()
             return categories, filenames
-        except Exception as e:
+        except Exception:
             return ["All"], ["All"]
-    # Session State
-    if "categories" not in st.session_state:
-        st.session_state["categories"], st.session_state["filenames"] = get_filter_choices()
-    if "category" not in st.session_state:
-        st.session_state["category"] = "All"
-    if "filename" not in st.session_state:
-        st.session_state["filename"] = "All"
-    if "lob" not in st.session_state:
-        st.session_state["lob"] = "leadgen_vivek"
-    if "start_date" not in st.session_state:
-        st.session_state["start_date"] = datetime.now(timezone.utc).date()
-    if "end_date" not in st.session_state:
-        st.session_state["end_date"] = datetime.now(timezone.utc).date()
-    if "images" not in st.session_state or "info" not in st.session_state:
-        imgs, info = app.search_images()
-        st.session_state["images"] = imgs
-        st.session_state["info"] = info
-
-    
+
+    # Session state defaults
+    if "categories_list" not in st.session_state:
+        st.session_state["categories_list"], st.session_state["filenames_list"] = get_filter_choices()
+
+    st.session_state.setdefault("selected_category", "All")
+    st.session_state.setdefault("selected_filename", "All")
+    st.session_state.setdefault("selected_lob", "leadgen_vivek")
+
+    today = datetime.now(timezone.utc).date()
+    st.session_state.setdefault("use_date_filter", True)
+    st.session_state.setdefault("selected_start_date", today)
+    st.session_state.setdefault("selected_end_date", today)
+
+    st.session_state.setdefault("page", 0)
+    st.session_state.setdefault("page_size", 24)
+    st.session_state.setdefault("last_query_total", 0)
+    st.session_state.setdefault("did_search", False)
+
+    # Custom styles for UI
     st.markdown("""
     <style>
-    .tab-bar {
-        display: flex;
-        margin-bottom: 1rem;
-        font-size: 20px;
-    }
-    .tab {
-        margin-right: 32px;
-        padding-bottom: 6px;
-        color: #02050d;
-        cursor: pointer;
-    }
-    .tab.selected {
-        color: #FF7300;
-        border-bottom: 3px solid #FF7300;
-        font-weight: bold;
-    }
-    /* Button Styles */
-    .orange-btn > button {
-        background-color: #FF7300 !important;
-        color: white !important;
-        font-weight: bold !important;
-        width: 100% !important;
-        border-radius: 8px !important;
-        font-size: 20px !important;
-        height: 48px !important;
-        margin-top: 8px !important;
-    }
-    .grey-btn > button {
-        background-color: #ECECEC !important;
-        color: #010203!important;
-        font-weight: 500 !important;
-        border-radius: 8px !important;
-        font-size: 18px !important;
-        height: 48px !important;
-    }
-    /* Card style */
-    .filter-card {
-        background: #fff;
+      .orange-btn > button { background:#FF7300 !important; color:#fff !important; font-weight:600 !important; }
+      .skel {
+        width: 100%;
+        aspect-ratio: 1 / 1;
         border-radius: 10px;
-        padding: 32px 24px 24px 24px;
-        box-shadow: 0 1px 6px 0 rgba(32,33,36,.08);
-        margin-bottom: 24px;
-    }
+        background: linear-gradient(90deg, #eee 25%, #f5f5f5 37%, #eee 63%);
+        background-size: 400% 100%;
+        animation: shimmer 1.4s ease infinite;
+      }
+      @keyframes shimmer {
+        0% { background-position: 0% 0; }
+        100% { background-position: -135% 0; }
+      }
+    </style>
     """, unsafe_allow_html=True)
-   
-    category = st.selectbox(
-        "Category",
-        options=st.session_state["categories"],
-        index=st.session_state["categories"].index(st.session_state["category"])
-            if st.session_state["category"] in st.session_state["categories"] else 0,
-        key="category"
-    )
-    filename = st.selectbox(
-        "File Name",
-        options=st.session_state["filenames"],
-        index=st.session_state["filenames"].index(st.session_state["filename"])
-            if st.session_state["filename"] in st.session_state["filenames"] else 0,
-        key="filename"
-    )
-    lob = st.text_input("LOB (Line of Business)", value=st.session_state["lob"], key="lob")
-    col_date1, col_date2 = st.columns(2)
-    with col_date1:
-        start_date = st.date_input("Start Date", value=st.session_state["start_date"], key="start_date")
-    with col_date2:
-        end_date = st.date_input("End Date", value=st.session_state["end_date"], key="end_date")
-    st.markdown("</div>", unsafe_allow_html=True)
-    
-    
-    col_btn1, col_btn2 = st.columns([3, 3])
+
+    # Filter controls
+    col1, col2, col3 = st.columns([1,1,1])
+    with col1:
+        category = st.selectbox(
+            "Category",
+            options=st.session_state["categories_list"],
+            index=st.session_state["categories_list"].index(st.session_state["selected_category"])
+                if st.session_state["selected_category"] in st.session_state["categories_list"] else 0
+        )
+    with col2:
+        filename = st.selectbox(
+            "File Name",
+            options=st.session_state["filenames_list"],
+            index=st.session_state["filenames_list"].index(st.session_state["selected_filename"])
+                if st.session_state["selected_filename"] in st.session_state["filenames_list"] else 0
+        )
+    with col3:
+        lob = st.text_input("LOB (Line of Business)", value=st.session_state["selected_lob"])
+
+    coldf = st.columns([1,1,1])
+    with coldf[0]:
+        use_date_filter = st.checkbox("Filter by date", value=st.session_state["use_date_filter"])
+    with coldf[1]:
+        start_date = st.date_input(
+            "Start Date",
+            value=st.session_state["selected_start_date"],
+            disabled=not use_date_filter
+        )
+    with coldf[2]:
+        end_date = st.date_input(
+            "End Date",
+            value=st.session_state["selected_end_date"],
+            disabled=not use_date_filter
+        )
+
+    col_misc = st.columns([1,1,1,2])
+    with col_misc[0]:
+        page_size = st.selectbox("Images per page", [8, 12, 16, 24, 32, 48],
+                                 index=[8, 12, 16, 24, 32, 48].index(st.session_state["page_size"]))
+
+    col_btn1, col_btn2, col_btn3 = st.columns([2,2,2])
     with col_btn1:
-        with st.container():
-            search_clicked = st.button("🔍 Search Images", key="search", use_container_width=True)
+        search_clicked = st.button("🔍 Search", use_container_width=True)
     with col_btn2:
-        with st.container():
-            refresh_clicked = st.button("🔄 Refresh Filters", key="refresh", use_container_width=True)
-   
-
-    col_gal, col_info = st.columns([2.2, 1])
-    with col_gal:
-        st.markdown("#### Images (Status: completed)")
-        images = st.session_state.get("images", [])
-        if images:
-            grid_cols = st.columns(4)
-            for idx, img in enumerate(images):
-                with grid_cols[idx % 4]:
-                    st.image(img, use_container_width=True)
-        else:
-            st.info("No images to display.")
-    with col_info:
-        st.text_area(
-            "Document Information",
-            value=st.session_state.get("info", ""),
-            height=450,
-            disabled=True,
-            key="info_text"
-        )
-    # --- Event handlers ---
+        refresh_clicked = st.button("🔄 Refresh Filters", use_container_width=True)
+    with col_btn3:
+        reset_clicked = st.button("♻️ Reset Page", use_container_width=True)
+
+    # Button events
+    if refresh_clicked:
+        st.session_state["categories_list"], st.session_state["filenames_list"] = get_filter_choices()
+        st.session_state["selected_category"] = "All"
+        st.session_state["selected_filename"] = "All"
+        st.session_state["page"] = 0
+        st.rerun()
+
+    if reset_clicked:
+        st.session_state["page"] = 0
+        st.rerun()
+
     if search_clicked:
-        imgs, info = app.search_images(
-            category, filename, start_date, end_date, lob
+        st.session_state["selected_category"] = category
+        st.session_state["selected_filename"] = filename
+        st.session_state["selected_lob"] = lob
+        st.session_state["use_date_filter"] = use_date_filter
+        st.session_state["selected_start_date"] = start_date
+        st.session_state["selected_end_date"] = end_date
+        st.session_state["page_size"] = page_size
+        st.session_state["page"] = 0
+        st.session_state["did_search"] = True
+        st.rerun()
+
+    # Results display
+    if st.session_state["did_search"]:
+        _start = st.session_state["selected_start_date"] if st.session_state["use_date_filter"] else None
+        _end = st.session_state["selected_end_date"] if st.session_state["use_date_filter"] else None
+
+        docs, total = app.search_images_page(
+            category=st.session_state["selected_category"],
+            file_name=st.session_state["selected_filename"],
+            start_date=_start,
+            end_date=_end,
+            lob=st.session_state["selected_lob"],
+            page=st.session_state["page"],
+            page_size=st.session_state["page_size"]
         )
-        st.session_state["images"] = imgs
-        st.session_state["info"] = info
-        st.session_state["category"] = category
-        st.session_state["filename"] = filename
-        st.session_state["lob"] = lob
-        st.session_state["start_date"] = start_date
-        st.session_state["end_date"] = end_date
-    if refresh_clicked:
-        st.session_state["categories"], st.session_state["filenames"] = get_filter_choices()
-        st.session_state["category"] = "All"
-        st.session_state["filename"] = "All"
-        st.rerun()
+        st.session_state["last_query_total"] = total
+
+        total_pages = max(1, (total + st.session_state["page_size"] - 1) // st.session_state["page_size"])
+
+        nav1, nav2, nav3 = st.columns([1,2,1])
+        with nav1:
+            if st.button("⬅️ Prev", disabled=(st.session_state["page"] <= 0)):
+                st.session_state["page"] -= 1
+                st.rerun()
+        with nav2:
+            st.markdown(
+                f"<div style='text-align:center'>Page <b>{st.session_state['page']+1}</b> of <b>{total_pages}</b>"
+                f" &nbsp;&middot;&nbsp; <b>{total}</b> images total</div>",
+                unsafe_allow_html=True
+            )
+        with nav3:
+            if st.button("Next ➡️", disabled=(st.session_state["page"] >= total_pages - 1)):
+                st.session_state["page"] += 1
+                st.rerun()
+
+        st.divider()
+
+        if total == 0 or not docs:
+            st.info("No images found for the current filters.")
+            return
+
+        st.markdown("#### Images")
+        cols = st.columns(4)
+
+        # Show skeletons while loading images
+        placeholders = []
+        for i, _ in enumerate(docs):
+            ph = cols[i % 4].empty()
+            ph.markdown("<div class='skel'></div>", unsafe_allow_html=True)
+            placeholders.append(ph)
+
+        urls = [d["url"] for d in docs]
+        loaded = app.load_images_parallel(urls, max_workers=8)
+        url_to_img = {u: img for (u, img) in loaded}
+
+        # Fill placeholders with loaded images
+        for i, d in enumerate(docs):
+            img = url_to_img.get(d["url"])
+            meta = f"{d.get('category','N/A')} | {d.get('file_name','N/A')} | {d.get('created_at','')}"
+            if img:
+                placeholders[i].image(img, use_container_width=True, caption=meta)
+            else:
+                placeholders[i].warning("Failed to load image")
+    else:
+        st.info("Set your filters and click **Search** to load images.")