Update app.py

app.py

@@ -19,13 +19,15 @@ import time
 import os
 import ssl
 from io import BytesIO
+from concurrent.futures import ThreadPoolExecutor
+import math
 
 # Initialize Dash app
 app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])
 server = app.server
 
 # Logging setup
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 # Thread-local storage for database connections
@@ -108,7 +110,7 @@ async def get_links(session, url, base_url):
         logger.error(f"Error getting links from {url}: {str(e)}")
         return []
 
-async def crawl_pages(base_url, max_depth):
+async def crawl_pages(base_url, max_depth, progress_callback):
     visited = set()
     to_visit = [(base_url, 0)]
     all_pages = []
@@ -140,6 +142,9 @@ async def crawl_pages(base_url, max_depth):
                 all_pages.append((current_url, content))
                 logger.info(f"Processed page: {current_url} at depth {depth} in {time.time() - start_time:.2f} seconds")
 
+                progress = len(all_pages) / (max_depth * 10)  # Rough estimate
+                progress_callback(f"Crawling pages... {progress:.0%}")
+
                 if depth < max_depth:
                     links = await get_links(session, current_url, base_url)
                     for link in links:
@@ -151,48 +156,75 @@ async def crawl_pages(base_url, max_depth):
 
     return all_pages
 
-def website_to_pdf(all_pages, progress_callback):
-    logger.info(f"Starting PDF generation for {len(all_pages)} pages")
-    
+def create_pdf_chunk(chunk, start_index):
     pdf = FPDF()
     pdf.set_auto_page_break(auto=True, margin=15)
     pdf.add_page()
     pdf.set_font("Arial", size=12)
 
-    batch_size = 100
-    for i in range(0, len(all_pages), batch_size):
-        batch = all_pages[i:i+batch_size]
-        for page_url, content in batch:
-            pdf.cell(0, 10, txt=page_url, ln=True)
-            pdf.ln(5)
-            for text in content:
-                try:
-                    pdf.multi_cell(0, 10, txt=text[:200])  # Limit text length to avoid issues
-                except Exception as e:
-                    logger.error(f"Error writing text to PDF: {str(e)}")
-            if pdf.get_y() > 250:  # Add a new page if the current page is almost full
-                pdf.add_page()
-        
-        progress = min((i + batch_size) / len(all_pages), 1.0)
-        progress_callback(f"Processing pages... {progress:.0%}")
-
-    return pdf.output(dest='S').encode('latin-1')  # Return bytes instead of BytesIO object
+    for i, (page_url, content) in enumerate(chunk, start=start_index):
+        pdf.cell(0, 10, txt=f"Page {i+1}: {page_url}", ln=True)
+        pdf.ln(5)
+        for text in content:
+            try:
+                pdf.multi_cell(0, 10, txt=text[:200])  # Limit text length to avoid issues
+            except Exception as e:
+                logger.error(f"Error writing text to PDF: {str(e)}")
+        pdf.add_page()
+
+    return pdf.output(dest='S').encode('latin-1')
+
+async def website_to_pdf(all_pages, progress_callback):
+    logger.info(f"Starting PDF generation for {len(all_pages)} pages")
+    
+    chunk_size = 100
+    num_chunks = math.ceil(len(all_pages) / chunk_size)
+    pdf_chunks = []
+
+    with ThreadPoolExecutor() as executor:
+        futures = []
+        for i in range(num_chunks):
+            start = i * chunk_size
+            end = min((i + 1) * chunk_size, len(all_pages))
+            chunk = all_pages[start:end]
+            future = executor.submit(create_pdf_chunk, chunk, start)
+            futures.append(future)
+
+        for i, future in enumerate(futures):
+            try:
+                pdf_chunk = await asyncio.wrap_future(future)
+                pdf_chunks.append(pdf_chunk)
+                progress = (i + 1) / num_chunks
+                progress_callback(f"Generating PDF... {progress:.0%}")
+            except Exception as e:
+                logger.error(f"Error generating PDF chunk {i}: {str(e)}")
+
+    # Combine PDF chunks
+    combined_pdf = FPDF()
+    for chunk in pdf_chunks:
+        combined_pdf.add_page()
+        combined_pdf.put_file(chunk)
+
+    return combined_pdf.output(dest='S').encode('latin-1')
 
 async def process_url(url, depth, progress_callback):
     try:
-        all_pages = await crawl_pages(url, depth)
+        all_pages = await asyncio.wait_for(crawl_pages(url, depth, progress_callback), timeout=3600)  # 1 hour timeout
         if not all_pages:
             return "No pages were successfully crawled. Please check the URL and try again."
-        pdf_content = website_to_pdf(all_pages, progress_callback)
-        return pdf_content  # This is now bytes, not BytesIO
+        pdf_content = await asyncio.wait_for(website_to_pdf(all_pages, progress_callback), timeout=3600)  # 1 hour timeout for PDF generation
+        return pdf_content
+    except asyncio.TimeoutError:
+        logger.error("Process timed out after 1 hour")
+        return "The process timed out after 1 hour. Please try again with a smaller depth or a more specific URL."
     except Exception as e:
         logger.error(f"Error in process_url: {str(e)}")
         return f"An error occurred: {str(e)}"
 
-# App layout
 # App layout
 app.layout = dbc.Container([
-    dcc.Store(id='pdf-store'),  # Add this line at the beginning of the layout
+    dcc.Store(id='pdf-store'),
+    dcc.Store(id='progress-store'),
     dbc.Navbar(
         dbc.Container([
             html.A(
@@ -218,6 +250,7 @@ app.layout = dbc.Container([
             dbc.Input(id="url-input", type="text", placeholder="Enter website URL (e.g., https://www.gradio.app/docs)", className="mb-3"),
             dcc.Slider(id="depth-slider", min=1, max=10, step=1, value=3, marks={i: str(i) for i in range(1, 11)}, className="mb-3"),
             dbc.Button("Convert to PDF", id="submit-button", color="primary", className="mb-3 w-100"),
+            dbc.Progress(id="progress-bar", animated=True, striped=True, className="mb-3"),
             dbc.Spinner(html.Div(id="output-area"), color="primary", type="grow"),
             dcc.Interval(id="progress-interval", interval=1000, n_intervals=0, disabled=True),
         ]),
@@ -228,13 +261,16 @@ app.layout = dbc.Container([
 @app.callback(
     Output("output-area", "children"),
     Output("progress-interval", "disabled"),
+    Output("progress-bar", "value"),
+    Output("progress-bar", "label"),
     Input("submit-button", "n_clicks"),
     Input("progress-interval", "n_intervals"),
     State("url-input", "value"),
     State("depth-slider", "value"),
+    State("progress-store", "data"),
     prevent_initial_call=True
 )
-def update_output(n_clicks, n_intervals, url, depth):
+def update_output(n_clicks, n_intervals, url, depth, progress):
     ctx = dash.callback_context
     if not ctx.triggered:
         raise PreventUpdate
@@ -243,17 +279,19 @@ def update_output(n_clicks, n_intervals, url, depth):
 
     if triggered_id == "submit-button":
         if not url:
-            return "Please enter a valid URL.", True
+            return "Please enter a valid URL.", True, 0, ""
 
-        return dcc.Store(id='pdf-store', data='processing'), False
+        return "Processing... Please wait.", False, 0, "0%"
 
     elif triggered_id == "progress-interval":
         store = dash.callback_context.inputs.get('pdf-store', None)
-        if store is None or store == 'processing':
-            return "Processing... Please wait.", False
+        if store is None:
+            if progress:
+                return "Processing... Please wait.", False, int(progress.split('%')[0]), progress
+            return "Processing... Please wait.", False, 0, "0%"
 
         if isinstance(store, str) and store.startswith("Error"):
-            return store, True
+            return store, True, 100, "100%"
 
         try:
             encoded = base64.b64encode(store).decode()
@@ -264,15 +302,16 @@ def update_output(n_clicks, n_intervals, url, depth):
                     href=f"data:application/pdf;base64,{encoded}",
                     download="website_content.pdf"
                 )
-            ]), True
+            ]), True, 100, "100%"
         except Exception as e:
             logger.error(f"Error creating download link: {str(e)}")
-            return f"An error occurred while creating the download link: {str(e)}", True
+            return f"An error occurred while creating the download link: {str(e)}", True, 100, "100%"
 
     raise PreventUpdate
 
 @app.callback(
     Output('pdf-store', 'data'),
+    Output('progress-store', 'data'),
     Input('submit-button', 'n_clicks'),
     State('url-input', 'value'),
     State('depth-slider', 'value'),
@@ -280,9 +319,9 @@ def update_output(n_clicks, n_intervals, url, depth):
 )
 def generate_pdf(n_clicks, url, depth):
     if not url:
-        return "Please enter a valid URL."
+        return "Please enter a valid URL.", "0%"
 
-    progress_store = {'progress': 0}
+    progress_store = {'progress': "0%"}
 
     def progress_callback(message):
         progress_store['progress'] = message
@@ -290,9 +329,9 @@ def generate_pdf(n_clicks, url, depth):
     pdf_content = asyncio.run(process_url(url, depth, progress_callback))
 
     if isinstance(pdf_content, str):
-        return pdf_content  # This is an error message
+        return pdf_content, "100%"  # This is an error message
 
-    return pdf_content
+    return pdf_content, "100%"
 
 if __name__ == '__main__':
     print("Starting the Dash application...")