Update app.py via AI Editor

app.py

@@ -1 +1,298 @@
--
+import dash
+import dash_bootstrap_components as dbc
+from dash import html, dcc, Input, Output, State, ctx
+import flask
+import uuid
+import os
+import tempfile
+import shutil
+import logging
+from flask import send_file, make_response
+import threading
+import pickle
+from PyPDF2 import PdfReader, PdfWriter
+import re
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+
+# Session storage and lock management
+SESSION_DATA = {}
+SESSION_LOCKS = {}
+
+def get_session_id():
+    session_id = flask.request.cookies.get('session-id')
+    if not session_id:
+        session_id = str(uuid.uuid4())
+    return session_id
+
+def get_session_dir(session_id):
+    base_tmp = tempfile.gettempdir()
+    path = os.path.join(base_tmp, f'dash_pdfsplit_{session_id}')
+    os.makedirs(path, exist_ok=True)
+    return path
+
+def clean_session(session_id):
+    try:
+        session_dir = get_session_dir(session_id)
+        if os.path.exists(session_dir):
+            shutil.rmtree(session_dir)
+        SESSION_DATA.pop(session_id, None)
+        SESSION_LOCKS.pop(session_id, None)
+    except Exception as e:
+        logging.error(f"Error cleaning session {session_id}: {e}")
+
+def get_session_lock(session_id):
+    if session_id not in SESSION_LOCKS:
+        SESSION_LOCKS[session_id] = threading.Lock()
+    return SESSION_LOCKS[session_id]
+
+def allowed_file(filename):
+    return '.' in filename and filename.lower().endswith('.pdf')
+
+def extract_text_headers(reader, page_num):
+    try:
+        page = reader.pages[page_num]
+        text = page.extract_text() or ""
+        # Extract the first non-blank line as a potential header
+        lines = [line.strip() for line in text.split('\n') if line.strip()]
+        header = lines[0] if lines else ""
+        return header
+    except Exception as e:
+        logging.warning(f"Failed extracting header from page {page_num}: {e}")
+        return ""
+
+def is_blank_page(reader, page_num):
+    try:
+        page = reader.pages[page_num]
+        text = (page.extract_text() or "").strip()
+        return len(text) == 0
+    except Exception as e:
+        logging.warning(f"Failed to check blank page at {page_num}: {e}")
+        return False
+
+def is_chapter_header(header):
+    patterns = [
+        r'^\s*chapter\b', r'^\s*section\b', r'^\s*part\b', r'^\s*appendix\b',
+        r'^\s*[ivxlcdm]+\.', r'^\s*\d+(\.\d+)*\s', r'^\s*introduction\b'
+    ]
+    for pat in patterns:
+        if re.match(pat, header, re.IGNORECASE):
+            return True
+    return False
+
+def estimate_writer_size(writer):
+    import io
+    f = io.BytesIO()
+    writer.write(f)
+    return f.tell()
+
+def intelligent_pdf_split(input_path, session_dir, max_mb=5, min_split_mb=4):
+    reader = PdfReader(input_path)
+    n_pages = len(reader.pages)
+    splits = []
+    current_writer = PdfWriter()
+    split_points = []
+    last_header = None
+    last_split_at = 0
+    for i in range(n_pages):
+        page = reader.pages[i]
+        current_writer.add_page(page)
+        size = estimate_writer_size(current_writer) / (1024 * 1024)
+        header = extract_text_headers(reader, i)
+        blank = is_blank_page(reader, i)
+        chapter = is_chapter_header(header)
+        split_here = False
+
+        # Force split if over max size
+        if size >= max_mb:
+            split_here = True
+        # Prefer to split between min_split_mb and max_mb at logical points
+        elif size >= min_split_mb:
+            if blank or chapter or (header and header != last_header):
+                split_here = True
+
+        if split_here:
+            splits.append((last_split_at, i+1))
+            last_split_at = i+1
+            current_writer = PdfWriter()
+        last_header = header
+
+    # Add final split if not already
+    if last_split_at < n_pages:
+        splits.append((last_split_at, n_pages))
+
+    # Write split files
+    split_files = []
+    for idx, (start, end) in enumerate(splits):
+        writer = PdfWriter()
+        for i in range(start, end):
+            writer.add_page(reader.pages[i])
+        out_path = os.path.join(session_dir, f'split_part_{idx+1}.pdf')
+        with open(out_path, 'wb') as f:
+            writer.write(f)
+        size = os.path.getsize(out_path) / (1024 * 1024)
+        split_files.append({'filename': os.path.basename(out_path), 'size': size, 'path': out_path})
+    return split_files
+
+# Dash app setup
+external_stylesheets = [dbc.themes.BOOTSTRAP]
+app = dash.Dash(__name__, external_stylesheets=external_stylesheets, suppress_callback_exceptions=True)
+server = app.server
+
+app.title = "Intelligent PDF Splitter"
+
+app.layout = dbc.Container(
+    [
+        dcc.Store(id='session-store', storage_type='session'),
+        html.Div(id='dummy-div', style={'display': 'none'}),
+        dbc.Row(
+            [
+                dbc.Col(
+                    dbc.Card(
+                        [
+                            dbc.CardHeader(html.H2("Intelligent PDF Splitter")),
+                            dbc.CardBody(
+                                [
+                                    html.P("Upload your PDF. The tool will split it into context-preserving sections, each under 5MB."),
+                                    dcc.Upload(
+                                        id='upload-pdf',
+                                        children=html.Div([
+                                            'Drag and Drop or ',
+                                            html.A('Select PDF File')
+                                        ]),
+                                        style={
+                                            'width': '100%', 'height': '80px', 'lineHeight': '80px',
+                                            'borderWidth': '1px', 'borderStyle': 'dashed', 'borderRadius': '5px',
+                                            'textAlign': 'center', 'margin': '10px 0'
+                                        },
+                                        multiple=False,
+                                        accept='.pdf'
+                                    ),
+                                    html.Div(id='file-info'),
+                                    dbc.Button("Clear Session", id='clear-session', color='secondary', className='mt-2'),
+                                    dcc.Loading(
+                                        id="loading", type="default",
+                                        children=[html.Div(id='split-results')]
+                                    )
+                                ]
+                            )
+                        ],
+                        className="mt-4"
+                    ),
+                    width=12
+                ),
+            ]
+        )
+    ],
+    fluid=True,
+    className="p-4"
+)
+
+@app.callback(
+    Output('file-info', 'children'),
+    Output('split-results', 'children'),
+    Output('session-store', 'data'),
+    Input('upload-pdf', 'contents'),
+    State('upload-pdf', 'filename'),
+    Input('clear-session', 'n_clicks'),
+    State('session-store', 'data'),
+    prevent_initial_call='initial_duplicate'
+)
+def handle_upload(contents, filename, clear_n, session_data):
+    trigger = ctx.triggered_id
+    session_id = get_session_id()
+    flask.g.session_id = session_id
+    session_dir = get_session_dir(session_id)
+    lock = get_session_lock(session_id)
+
+    if trigger == 'clear-session':
+        clean_session(session_id)
+        resp_data = {}
+        return "", "", resp_data
+
+    # If user returns, restore state
+    if not contents and session_data and 'split_files' in session_data:
+        split_files = session_data.get('split_files', [])
+        file_info = html.Div(f"Previous upload: {session_data.get('orig_filename', '')}")
+        results = [
+            html.H5("Split Files:"),
+            html.Ul([
+                html.Li([
+                    f"{fi['filename']} ({fi['size']:.2f} MB) ",
+                    dbc.Button("Download", id={'type': 'download-btn', 'index': idx}, href=f"/download/{session_id}/{fi['filename']}", color='primary', size='sm')
+                ]) for idx, fi in enumerate(split_files)
+            ])
+        ]
+        return file_info, results, session_data
+
+    if not contents:
+        return "", "", {}
+
+    if not allowed_file(filename):
+        return html.Div("Only .pdf files are allowed.", style={'color': 'red'}), "", {}
+
+    try:
+        # Save file
+        header, b64data = contents.split(',', 1)
+        import base64
+        pdf_bytes = base64.b64decode(b64data)
+        pdf_path = os.path.join(session_dir, filename)
+        with open(pdf_path, 'wb') as f:
+            f.write(pdf_bytes)
+        logging.info(f"PDF uploaded and saved to {pdf_path} for session {session_id}")
+
+        # Split PDF with lock
+        with lock:
+            split_files = intelligent_pdf_split(pdf_path, session_dir)
+        results = [
+            html.H5("Split Files:"),
+            html.Ul([
+                html.Li([
+                    f"{fi['filename']} ({fi['size']:.2f} MB) ",
+                    dbc.Button("Download", id={'type': 'download-btn', 'index': idx}, href=f"/download/{session_id}/{fi['filename']}", color='primary', size='sm')
+                ]) for idx, fi in enumerate(split_files)
+            ])
+        ]
+        file_info = html.Div(f"Uploaded: {filename} ({len(pdf_bytes)/1024/1024:.2f} MB)")
+        session_data = {
+            'orig_filename': filename,
+            'split_files': split_files,
+        }
+        logging.info(f"PDF split into {len(split_files)} chunks for session {session_id}")
+        return file_info, results, session_data
+    except Exception as e:
+        logging.error(f"Error processing PDF: {e}")
+        return html.Div(f"Error: {e}", style={'color': 'red'}), "", {}
+
+@app.server.route('/download/<session_id>/<filename>')
+def download_split_file(session_id, filename):
+    session_dir = get_session_dir(session_id)
+    file_path = os.path.join(session_dir, filename)
+    if os.path.exists(file_path):
+        logging.info(f"Serving file {file_path} for session {session_id}")
+        return send_file(file_path, mimetype='application/pdf', as_attachment=True, download_name=filename)
+    else:
+        logging.error(f"File not found for download: {file_path}")
+        return "File not found", 404
+
+@app.callback(
+    Output('dummy-div', 'children'),
+    Input('session-store', 'data'),
+    prevent_initial_call=True
+)
+def set_cookie_on_load(session_data):
+    session_id = get_session_id()
+    resp = flask.make_response("")
+    resp.set_cookie('session-id', session_id, max_age=60*60*24*3)
+    return ""
+
+@app.server.before_request
+def persist_session_cookie():
+    session_id = get_session_id()
+    flask.g.session_id = session_id
+
+if __name__ == '__main__':
+    print("Starting the Dash application...")
+    app.run(debug=True, host='0.0.0.0', port=7860, threaded=True)
+    print("Dash application has finished running.")