Update app.py

app.py

@@ -6,82 +6,43 @@ import io
 import fitz  # PyMuPDF
 import tempfile
 import os
-import shutil # Added for cleaning up temp dirs
 
 # --- OPENAI CLIENT SETUP ---
-# Use environment variable or textbox for API key for better security in deployed apps
-# client = OpenAI(
-#     base_url="https://openrouter.ai/api/v1",
-#     api_key=os.getenv("OPENROUTER_API_KEY") # Recommended approach
-# )
-# For this example, we'll get the key from the input field
+client = OpenAI(
+    base_url="https://openrouter.ai/api/v1",
+    api_key='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'
+)
 
-def get_openai_client(api_key):
-    """Initializes and returns the OpenAI client."""
-    if not api_key:
-        # Handle case where API key is missing (though Gradio will likely prevent this)
-        raise ValueError("API key is required.")
-        
-    return OpenAI(
-        base_url="https://openrouter.ai/api/v1",
-        api_key=api_key
-    )
-
-def convert_pdf_to_images(pdf_path):
-    """Convert PDF file path to list of PIL Images and return the images,
-    and a list of temporary image file paths."""
+def convert_pdf_to_images(pdf_file):
+    """Convert PDF to list of PIL Images"""
     images = []
-    temp_image_paths = []
-    temp_dir = None
     try:
-        pdf_document = fitz.open(pdf_path)
-        num_pages = len(pdf_document)
-        
-        # Create a temporary directory for images
-        temp_dir = tempfile.mkdtemp()
-        
-        for page_num in range(num_pages):
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
+            tmp_file.write(pdf_file.read())
+            tmp_file_path = tmp_file.name
+
+        pdf_document = fitz.open(tmp_file_path)
+        for page_num in range(len(pdf_document)):
             page = pdf_document.load_page(page_num)
-            # Render at a higher DPI for better clarity for VLM
-            pix = page.get_pixmap(dpi=300) 
+            pix = page.get_pixmap()
             img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
             images.append(img)
-            
-            # Save image to temp directory for Gradio preview/processing later
-            temp_img_path = os.path.join(temp_dir, f"page_{page_num+1}.png")
-            img.save(temp_img_path, format="PNG")
-            temp_image_paths.append(temp_img_path)
 
         pdf_document.close()
-        
-        return images, temp_image_paths, num_pages, temp_dir
-
+        os.unlink(tmp_file_path)
     except Exception as e:
-        print(f"Error converting PDF: {e}")
-        # Clean up temp dir if it was created
-        if temp_dir and os.path.exists(temp_dir):
-             shutil.rmtree(temp_dir)
-        return [], [], 0, None
+        return f"Error converting PDF: {e}"
+    return images
 
 def image_to_base64(image):
     """Convert PIL Image to base64 string"""
-    # Ensure image is RGB (some images might be RGBA, etc.)
-    if image.mode != 'RGB':
-        image = image.convert('RGB')
-        
     with io.BytesIO() as buffer:
-        # Using PNG as it's lossless and well-supported
         image.save(buffer, format="PNG")
         return base64.b64encode(buffer.getvalue()).decode("utf-8")
 
-def generate_summary(extracted_texts, api_key):
+def generate_summary(extracted_texts):
     """Generate a comprehensive summary of all extracted texts"""
-    if not extracted_texts:
-        return "No content extracted to summarize."
-        
     try:
-        client = get_openai_client(api_key)
-        
         summary_prompt = f"""
         You are an expert document analyst. Below are the extracted contents from multiple pages of a document.
         Please provide a comprehensive, detailed summary that:
@@ -91,357 +52,145 @@ def generate_summary(extracted_texts, api_key):
         4. Presents the information in a clear, structured format
         
         Extracted contents from pages:
-        ---
         {extracted_texts}
-        ---
         
         Comprehensive Summary:
         """
-            
+
         response = client.chat.completions.create(
-            model="opengvlab/internvl3-14b:free", # Ensure this model is available via OpenRouter
+            model="opengvlab/internvl3-14b:free",
             messages=[
                 {"role": "system", "content": "You are Dalton, an expert in analyzing and summarizing document contents."},
                 {"role": "user", "content": summary_prompt}
             ],
-            max_tokens=2048 # Adjust as needed
+            max_tokens=2048
         )
-        
+
         return response.choices[0].message.content
-            
     except Exception as e:
-        print(f"Error generating summary: {e}")
         return f"Error generating summary: {e}"
 
-# --- Gradio App Functions ---
+def analyze_images(images, user_prompt, selected_pages=None):
+    if not images:
+        return "No images provided for analysis."
 
-def process_upload(file_obj):
-    """Handle file upload - converts PDF, prepares image previews, and updates state."""
-    if file_obj is None:
-        # Clear outputs
-        return None, None, [], [], "Please upload a document.", None, None, None
+    if isinstance(images, str):  # error message
+        return images
 
-    file_path = file_obj.name # Gradio's File component provides a path
-    file_type = file_obj.orig_name.split('.')[-1].lower() # Get extension from original name
+    if selected_pages is None:
+        selected_pages = list(range(1, len(images) + 1))
 
-    if file_type == "pdf":
-        images, temp_paths, num_pages, temp_dir = convert_pdf_to_images(file_path)
-        if not images:
-            return None, None, [], [], "Failed to convert PDF to images.", None, None, None
-        
-        page_options = [f"Page {i}" for i in range(1, num_pages + 1)]
-        # By default select all pages
-        default_selection = page_options 
-        
-        # Store original PIL images and temp dir in state
-        # State will hold (list of PIL images, list of temp file paths, temp directory path)
-        images_state = (images, temp_paths, temp_dir)
-        
-        status = f"PDF uploaded. {num_pages} pages detected. Select pages to analyze."
-        # Return selected pages (as names), image previews (as paths), page options, status
-        return images_state, default_selection, temp_paths, page_options, status, None, None, None # Also return None for results and summary
+    images_to_analyze = [images[i - 1] for i in selected_pages]
+    all_results = []
+    extracted_texts = []
 
-    elif file_type in ["jpg", "jpeg", "png"]:
+    for idx, image in enumerate(images_to_analyze, 1):
         try:
-            img = Image.open(file_path)
-            # Ensure it's RGB
-            if img.mode != 'RGB':
-                img = img.convert('RGB')
-                
-            # Save to a temp file for Gradio preview
-            temp_dir = tempfile.mkdtemp()
-            temp_img_path = os.path.join(temp_dir, "uploaded_image.png")
-            img.save(temp_img_path, format="PNG")
-            
-            # Store original PIL image and temp dir in state
-            # State will hold (list of PIL images, list of temp file paths, temp directory path)
-            images_state = ([img], [temp_img_path], temp_dir)
-            
-            status = "Image uploaded."
-            # Return empty selection/options for image, but provide the single image path for preview
-            return images_state, [], [temp_img_path], [], status, None, None, None # Also return None for results and summary
+            image_base64 = image_to_base64(image)
+
+            response = client.chat.completions.create(
+                model="opengvlab/internvl3-14b:free",
+                messages=[
+                    {"role": "system", "content": "You are Dalton, an expert in understanding images that can analyze images and provide detailed descriptions."},
+                    {"role": "user", "content": [
+                        {"type": "text", "text": user_prompt},
+                        {"type": "image_url", "image_url": {
+                            "url": f"data:image/png;base64,{image_base64}"
+                        }}
+                    ]}
+                ],
+                max_tokens=1024
+            )
+
+            result = response.choices[0].message.content
+            extracted_texts.append(f"=== Page {selected_pages[idx-1]} ===\n{result}\n")
+            all_results.append(f"### 📄 Page {selected_pages[idx-1]} Result:")
+            all_results.append(result)
+            all_results.append("---")
 
         except Exception as e:
-            print(f"Error loading image: {e}")
-            # Clean up temp dir if created
-            if temp_dir and os.path.exists(temp_dir):
-                 shutil.rmtree(temp_dir)
-            return None, None, [], [], f"Failed to load image: {e}", None, None, None
+            all_results.append(f"An error occurred analyzing page {selected_pages[idx-1]}: {e}")
+
+    full_result = "\n".join(all_results)
+
+    if len(extracted_texts) > 1:
+        full_extracted_text = "\n".join(extracted_texts)
+        summary = generate_summary(full_extracted_text)
+        full_result += "\n\n## 📝 Comprehensive Document Summary\n"
+        full_result += summary
+        return full_result, summary
+    elif len(extracted_texts) == 1:
+        return full_result, None
     else:
-        return None, None, [], [], "Unsupported file type. Please upload JPG, PNG, or PDF.", None, None, None
-
-def analyze_document(api_key, user_prompt, images_state, selected_page_names):
-    """Analyze selected images using the VLM and generate summary."""
-    if not api_key:
-        return None, None, "Please enter your Open Router API Key."
-        
-    if not images_state or not images_state[0]: # Check if images_state exists and contains images
-        return None, None, "No document uploaded or converted."
-
-    all_pil_images = images_state[0]
-    temp_dir = images_state[2] # Get the temp directory path
-
-    images_to_analyze = []
-    extracted_texts = []
-    all_results = []
-    
-    # Determine which images to process based on selection (or process all if image file)
-    if selected_page_names: # This indicates PDF and pages were selected
-        selected_indices = [int(name.split(" ")[1]) - 1 for name in selected_page_names]
-        images_to_analyze = [(idx + 1, all_pil_images[idx]) for idx in selected_indices if idx < len(all_pil_images)]
-    elif all_pil_images: # This indicates a single image file
-         images_to_analyze = [(1, all_pil_images[0])]
-         
-    if not images_to_analyze:
-         # Clean up temp dir as analysis failed or no pages selected
-         if temp_dir and os.path.exists(temp_dir):
-             shutil.rmtree(temp_dir)
-         return None, None, "No pages selected for analysis."
-
-
-    try:
-        client = get_openai_client(api_key)
-        
-        for page_num, image in images_to_analyze:
-            status_message = f"Analyzing page {page_num}..."
-            yield None, None, status_message # Update status message during processing
-
-            try:
-                image_base64_data = image_to_base64(image)
-                
-                response = client.chat.completions.create(
-                    model="opengvlab/internvl3-14b:free", # Ensure this model is available via OpenRouter
-                    messages=[
-                        {"role": "system", "content": "You are Dalton, an expert in understanding images that can analyze images and provide detailed descriptions."},
-                        {"role": "user", "content": [
-                            {"type": "text", "text": user_prompt},
-                            {"type": "image_url", "image_url": {
-                                "url": f"data:image/png;base64,{image_base64_data}"
-                            }}
-                        ]}
-                    ],
-                    max_tokens=1024 # Adjust as needed
-                )
-
-                result = response.choices[0].message.content
-                extracted_texts.append(f"=== Page {page_num} ===\n{result}\n")
-                
-                if len(images_to_analyze) > 1:
-                    all_results.append(f"### 📄 Page {page_num} Result:")
-                else:
-                    all_results.append("### ✅ Analysis Result:")
-                all_results.append(result)
-                all_results.append("---")
-
-            except Exception as e:
-                error_msg = f"An error occurred analyzing page {page_num}: {e}"
-                print(error_msg)
-                all_results.append(f"### ❌ Error on Page {page_num}:")
-                all_results.append(error_msg)
-                all_results.append("---")
-                # Don't stop, try other pages
-
-        # Combine individual results
-        individual_results_markdown = "\n".join(all_results) if all_results else "No results generated."
-        
-        # Generate and display comprehensive summary if multiple pages were processed
-        summary_text = ""
-        if len(images_to_analyze) > 1 and extracted_texts:
-            yield individual_results_markdown, None, "Generating comprehensive summary..."
-            full_extracted_text = "\n".join(extracted_texts)
-            summary_text = generate_summary(full_extracted_text, api_key)
-            status_message = "Analysis complete. Summary generated."
-        elif extracted_texts: # Single page case
-             summary_text = "Summary not generated for single page analysis. See analysis result above."
-             status_message = "Analysis complete."
-        else:
-            summary_text = "No content extracted for summary."
-            status_message = "Analysis complete, but no text extracted."
-
-        # Clean up the temporary directory used for images
-        if temp_dir and os.path.exists(temp_dir):
-            shutil.rmtree(temp_dir)
-            
-        return individual_results_markdown, summary_text, status_message
-
-    except Exception as e:
-        # Clean up the temporary directory in case of error
-        if temp_dir and os.path.exists(temp_dir):
-            shutil.rmtree(temp_dir)
-            
-        error_msg = f"An unhandled error occurred during analysis: {e}"
-        print(error_msg)
-        return None, None, error_msg
-
-
-# Function to clean up temp dir when session ends or is closed
-def clean_temp_dir(temp_dir):
-    if temp_dir and os.path.exists(temp_dir):
-        print(f"Cleaning up temporary directory: {temp_dir}")
-        shutil.rmtree(temp_dir)
-
+        return "No valid results generated.", None
 
-# --- Gradio Interface Layout ---
+def process_input(file, user_prompt, page_numbers):
+    if file is None:
+        return "Please upload a file.", None
 
-# Custom CSS (simplified from Streamlit CSS)
-css = """
-body {
-    font-family: 'Inter', sans-serif;
-}
-.gradio-container {
-    max-width: 800px !important;
-    margin: auto;
-    padding: 20px;
-    background-color: #f9fafb; /* Light gray background */
-}
-h1, h2, h3, h4 {
-    color: #111827; /* Darker text for headers */
-}
-.subtitle {
-    font-size: 1rem;
-    color: #6b7280; /* Gray text for subtitle */
-    margin-bottom: 2rem;
-}
-.summary-box {
-    background-color: #e0f2fe; /* Light blue background */
-    padding: 1.5rem;
-    border-radius: 8px;
-    margin-top: 1rem; /* Reduced margin-top */
-    border: 1px solid #bfdbfe; /* Light blue border */
-}
-.summary-box p {
-    margin: 0; /* Remove paragraph margin */
-}
-.file-upload-label .wrap {
-    text-align: center !important;
-}
-.gr-button {
-    margin-top: 1rem !important;
-}
-/* Style for the status message */
-#status_message_id { 
-    margin-top: 1rem;
-    font-weight: bold;
-    color: #1f2937;
-}
-"""
+    mime_type = file.type
+    images = []
 
-with gr.Blocks(css=css, title="DocSum - Document Summarizer", theme=gr.themes.Soft()) as demo:
-    
-    # State to hold images and temp paths after PDF conversion
-    # Structure: (list of PIL images, list of temp file paths for preview/analysis, temp directory path)
-    images_state = gr.State(None)
-    # State to hold the temp dir path for cleanup
-    current_temp_dir = gr.State(None)
+    if mime_type == "application/pdf":
+        images = convert_pdf_to_images(file)
+        if isinstance(images, str):  # error message
+            return images, None
+        page_options = list(range(1, len(images) + 1))
+        if not page_numbers or len(page_numbers) == 0:
+            page_numbers = page_options
+        return analyze_images(images, user_prompt, page_numbers)
+    elif mime_type.startswith("image/"):
+        images = [Image.open(file)]
+        return analyze_images(images, user_prompt)
+    else:
+        return "Unsupported file type. Please upload a JPG/PNG/PDF.", None
 
-    gr.HTML("""
-        <div style="text-align: center;">
-            <img src='https://raw.githubusercontent.com/KoshurAI/DocSum/main/blob.png' width='100'>
-            <h1>DocSum</h1>
-            <p class="subtitle">Document Summarizer Powered by VLM • Developed by <a href="https://koshurai.com" target="_blank">Koshur AI</a></p>
-        </div>
+# --- GRADIO INTERFACE ---
+with gr.Blocks(title="DocSum - Document Summarizer") as demo:
+    gr.Markdown("""
+    <h1 style="text-align:center;">🧾 DocSum</h1>
+    <p style="text-align:center;">Document Summarizer Powered by VLM • Developed by <a href='https://koshurai.com' target='_blank'>Koshur AI</a></p>
     """)
     
     with gr.Row():
-        user_prompt_input = gr.Textbox(
-            label="📝 Enter Your Prompt",
-            value="Extract all content structurally",
-            lines=2,
-            interactive=True,
-            container=True,
-            scale=2
-        )
-        api_key_input = gr.Textbox(
-            label="🔒 OpenRouter API Key",
-            type="password",
-            interactive=True,
-            container=True,
-            scale=1,
-            info="Your key is not stored."
-            # Consider adding value=os.getenv("OPENROUTER_API_KEY", "") for easier local testing
-        )
-
-    file_upload = gr.File(
-        label="Upload a document (JPG/PNG/PDF)", 
-        file_types=[".jpg", ".jpeg", ".png", ".pdf"],
-        interactive=True
-    )
-
-    # Components for PDF page selection and preview (initially hidden)
-    page_selector = gr.Checkboxgroup(
-        label="Select PDF Pages to Analyze", 
-        choices=[], 
-        value=[], 
-        visible=False,
-        interactive=True
-    )
-    preview_gallery = gr.Gallery(
-        label="Selected Page Previews",
-        visible=False,
-        container=True,
-        preview=True, # Show previews
-        columns=3,
-        rows=1,
-        object_fit="contain",
-        height="auto"
-    )
-    
-    status_message = gr.Markdown(elem_id="status_message_id") # Use a Markdown element for status updates
-
-    analyze_button = gr.Button("🔍 Analyze Document")
-
-    # Outputs
-    individual_results_output = gr.Markdown(label="Page-by-Page Analysis Results")
-    summary_output = gr.Markdown(label="Comprehensive Document Summary", elem_classes="summary-box") # Apply CSS class
+        with gr.Column():
+            file_upload = gr.File(label="Upload a document (JPG/PNG/PDF)", file_types=[".jpg", ".jpeg", ".png", ".pdf"])
+            prompt = gr.Textbox(label="📝 Enter Your Prompt", value="Extract all content structurally")
+            page_selector = gr.CheckboxGroup(label="Select Pages (for PDFs only)", choices=[], visible=False)
+
+            def update_page_selector(file):
+                if file and file.type == "application/pdf":
+                    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
+                        tmp_file.write(file.read())
+                        tmp_file_path = tmp_file.name
+                    doc = fitz.open(tmp_file_path)
+                    num_pages = len(doc)
+                    doc.close()
+                    os.unlink(tmp_file_path)
+                    return gr.update(choices=list(range(1, num_pages + 1)), visible=True)
+                else:
+                    return gr.update(choices=[], visible=False)
 
-    # --- Event Handling ---
+            file_upload.change(fn=update_page_selector, inputs=file_upload, outputs=page_selector)
 
-    # When a file is uploaded, process it (convert PDF, show previews, update state)
-    file_upload.change(
-        fn=process_upload,
-        inputs=[file_upload],
-        outputs=[images_state, page_selector, preview_gallery, page_selector.choices, status_message, individual_results_output, summary_output, current_temp_dir],
-        show_progress=True # Show Gradio's built-in progress indicator
-    )
-    
-    # When page selection changes (for PDF), update the preview gallery
-    # Note: This requires saving the temp image paths in the state from process_upload
-    page_selector.change(
-        fn=lambda selected_pages, images_state: [images_state[1][int(name.split(" ")[1]) - 1] for name in selected_pages] if images_state and images_state[1] else [],
-        inputs=[page_selector, images_state],
-        outputs=[preview_gallery],
-        show_progress=False # No need for progress bar here
-    ).then( # Chain another event to update status message
-         fn=lambda num_selected: f"{num_selected} pages selected." if num_selected > 0 else "No pages selected.",
-         inputs=[page_selector],
-         outputs=[status_message],
-         show_progress=False
-    )
+            submit_btn = gr.Button("🔍 Analyze Document")
 
+        with gr.Column():
+            output_box = gr.Markdown(label="Analysis Output")
+            summary_download = gr.File(label="Download Summary", visible=False)
 
-    # When the Analyze button is clicked, run the analysis function
-    analyze_button.click(
-        fn=analyze_document,
-        inputs=[api_key_input, user_prompt_input, images_state, page_selector],
-        outputs=[individual_results_output, summary_output, status_message],
-        show_progress=False # We handle progress manually with status_message yield
-    )
-    
-    # --- Footer ---
-    gr.HTML("<footer style='text-align: center; margin-top: 3rem; color: #9ca3af; font-size: 0.875rem;'>© 2025 Koshur AI. All rights reserved.</footer>")
+    def handle_submit(file, prompt, pages):
+        result, summary = process_input(file, prompt, pages)
+        summary_file = None
+        if summary:
+            with tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".txt") as tmpfile:
+                tmpfile.write(summary)
+                summary_file = tmpfile.name
+        return result, summary_file
 
-# Clean up temp directory when the Gradio app finishes or encounters a critical error
-# Note: This might not catch all termination scenarios, especially if the server crashes unexpectedly.
-# A more robust solution for production might involve monitoring temp dirs periodically.
-# Using demo.load() to clean up at startup and demo.close() to clean up at exit.
-demo.load(fn=lambda: clean_temp_dir(current_temp_dir.value), inputs=[], outputs=[], every=10, show_progress=False) # Check & cleanup periodically (adjust interval)
-# The close event handler is tricky for cleanup; rely more on periodic check or OS cleanup.
+    submit_btn.click(fn=handle_submit, inputs=[file_upload, prompt, page_selector], outputs=[output_box, summary_download])
 
-# --- Launch App ---
-if __name__ == "__main__":
-    # The share=True option creates a public URL (useful for testing)
-    # The debug=True option provides more detailed error messages
-    demo.launch(share=False, debug=True) 
+    gr.Markdown("<footer>© 2025 Koshur AI. All rights reserved.</footer>")
 
-    # You might want to add cleanup here if running locally and not sharing
-    # clean_temp_dir(current_temp_dir.value) # This won't run if the app is killed externally
+# Launch Gradio App
+demo.launch()