Update app.py

app.py

@@ -1,231 +1,237 @@
-# app.py
 import streamlit as st
 import pdfplumber
 import pytesseract
-from PIL import Image
-import os
-import json
 import openai
+import json
 import pandas as pd
 import numpy as np
+from PIL import Image
 from io import BytesIO
-from concurrent.futures import ThreadPoolExecutor
-import hashlib
 import time
 import traceback
+import os
+import hashlib
 
 # Configuration
-MAX_THREADS = 4
 SUPPORTED_MODELS = {
     "Deepseek": {
-        "model": "deepseek-chat",
-        "base_url": "https://api.deepseek.com/v1"
+        "base_url": "https://api.deepseek.com/v1",
+        "required_key": "DEEPSEEK_KEY"
+    },
+    "OpenAI": {
+        "base_url": "https://api.openai.com/v1",
+        "required_key": "OPENAI_KEY"
     }
 }
 
-def debug_log(message):
-    """Enhanced logging system"""
-    if st.session_state.get("debug_mode"):
-        st.toast(f"DEBUG: {message}", icon="🐛")
-
-def initialize_session_state():
-    """Initialize all session state variables with validation"""
-    required_keys = {
+def initialize_session():
+    """Initialize session state with validation"""
+    defaults = {
+        'processing_stage': 'idle',
         'document_data': [],
         'qa_pairs': [],
-        'processing_complete': False,
-        'current_stage': 'idle',
-        'api_keys': {},
-        'model_choice': "Deepseek",
-        'temperature': 0.3,
-        'debug_mode': True
+        'export_formats': ['JSON', 'CSV', 'Parquet'],
+        'model_settings': {
+            'current_model': 'Deepseek',
+            'temperature': 0.3
+        },
+        'api_keys': {}
     }
     
-    for key, value in required_keys.items():
+    for key, val in defaults.items():
         if key not in st.session_state:
-            st.session_state[key] = value
-
-def show_processing_status():
-    """Visual feedback system"""
-    status_messages = {
-        'idle': "🟢 Ready to process",
-        'extracting': "🔍 Extracting document content...",
-        'generating': "🧠 Generating Q&A pairs...",
-        'evaluating': "📊 Evaluating results...",
-        'error': "❌ Processing failed"
-    }
-    
-    status = st.session_state.current_stage
-    debug_log(f"Status update: {status}")
-    st.sidebar.markdown(f"**System Status:** {status_messages.get(status, 'Unknown')}")
+            st.session_state[key] = val
 
-def process_image(img_data, page_num, img_idx):
-    """Robust image processing with validation"""
+def handle_image_errors(img_stream):
+    """Robust image processing with multiple fallbacks"""
     try:
-        img = img_data["stream"]
-        width = int(img_data["width"])
-        height = int(img_data["height"])
-        
-        debug_log(f"Processing image {img_idx} on page {page_num}")
-        
-        # Convert image to RGB
+        # First try standard RGB conversion
+        return Image.frombytes("RGB", (img_stream['width'], img_stream['height']), img_stream.get_data())
+    except:
         try:
-            return Image.frombytes("RGB", (width, height), img.get_data())
-        except:
-            return Image.frombytes("L", (width, height), img.get_data()).convert("RGB")
-            
-    except Exception as e:
-        st.error(f"Image processing failed (Page {page_num}, Image {img_idx}): {str(e)}")
-        return None
+            # Fallback to grayscale conversion
+            return Image.frombytes("L", (img_stream['width'], img_stream['height']), img_stream.get_data()).convert("RGB")
+        except Exception as e:
+            st.error(f"Critical image error: {str(e)[:200]}")
+            return None
 
-def pdf_processing_workflow(uploaded_file):
-    """PDF processing with real-time feedback"""
-    st.session_state.current_stage = 'extracting'
+def process_pdf(uploaded_file):
+    """PDF processing with enhanced error recovery"""
+    st.session_state.processing_stage = 'extracting'
+    doc_data = []
     
     try:
         with pdfplumber.open(uploaded_file) as pdf:
-            total_pages = len(pdf.pages)
-            progress_bar = st.progress(0)
-            status_text = st.empty()
-            
             for page_num, page in enumerate(pdf.pages, 1):
-                status_text.text(f"Processing page {page_num}/{total_pages}")
-                progress_bar.progress(page_num/total_pages)
+                page_data = {
+                    "page": page_num,
+                    "text": page.extract_text() or "",
+                    "images": []
+                }
                 
-                try:
-                    text = page.extract_text() or ""
-                    images = [process_image(img, page_num, idx) 
-                            for idx, img in enumerate(page.images)]
-                    
-                    st.session_state.document_data.append({
-                        "page": page_num,
-                        "text": text.strip(),
-                        "images": [img for img in images if img is not None]
-                    })
-                except Exception as e:
-                    st.error(f"Page {page_num} error: {str(e)}")
+                # Process images with error containment
+                for img_idx, img in enumerate(page.images):
+                    processed_img = handle_image_errors(img['stream'])
+                    if processed_img:
+                        page_data["images"].append(processed_img)
                 
-                time.sleep(0.1)  # Simulate processing
+                doc_data.append(page_data)
+                time.sleep(0.01)  # Yield for UI updates
                 
-            progress_bar.empty()
-            status_text.success("Document processing complete!")
-            return True
-            
+        st.session_state.document_data = doc_data
+        return True
     except Exception as e:
-        st.session_state.current_stage = 'error'
         st.error(f"PDF processing failed: {str(e)}")
-        debug_log(traceback.format_exc())
         return False
 
-def generate_qa_pairs():
-    """Q&A generation with validation"""
-    st.session_state.current_stage = 'generating'
+def generate_qa_content():
+    """Model-agnostic content generation"""
+    st.session_state.processing_stage = 'generating'
     qa_pairs = []
     
     try:
         client = openai.OpenAI(
-            base_url=SUPPORTED_MODELS[st.session_state.model_choice]["base_url"],
-            api_key=st.secrets["DEEPSEEK_API_KEY"]
+            base_url=SUPPORTED_MODELS[st.session_state.model_settings['current_model']]['base_url'],
+            api_key=st.session_state.api_keys.get(
+                SUPPORTED_MODELS[st.session_state.model_settings['current_model']]['required_key']
+            )
         )
         
-        for idx, entry in enumerate(st.session_state.document_data):
-            text_content = entry["text"] or " ".join([
-                pytesseract.image_to_string(img) for img in entry["images"]
+        for page in st.session_state.document_data:
+            text_content = page['text'] or " ".join([
+                pytesseract.image_to_string(img) for img in page['images']
             ])
             
             response = client.chat.completions.create(
-                model=SUPPORTED_MODELS[st.session_state.model_choice]["model"],
+                model="gpt-4-turbo" if st.session_state.model_settings['current_model'] == "OpenAI" else "deepseek-chat",
                 messages=[{
                     "role": "user",
-                    "content": f"Generate 3 Q&A pairs from:\n{text_content}\nReturn JSON format: {{'qa_pairs': [{{'question': '...', 'answer_1': '...', 'answer_2': '...'}}]}}"
+                    "content": f"Generate 3 Q&A pairs from this financial content:\n{text_content}\nOutput JSON format with keys: question, answer_1, answer_2"
                 }],
-                max_tokens=2048,
                 response_format={"type": "json_object"},
-                temperature=st.session_state.temperature
+                temperature=st.session_state.model_settings['temperature']
             )
             
             try:
                 result = json.loads(response.choices[0].message.content)
-                qa_pairs.extend(result.get("qa_pairs", []))
-                debug_log(f"Generated {len(result.get('qa_pairs', []))} pairs for page {entry['page']}")
+                qa_pairs.extend(result.get('qa_pairs', []))
             except json.JSONDecodeError:
-                st.error(f"Invalid response format from API for page {entry['page']}")
+                st.error("Failed to parse model response")
                 
         st.session_state.qa_pairs = qa_pairs
-        st.session_state.current_stage = 'evaluating'
         return True
-        
     except Exception as e:
-        st.session_state.current_stage = 'error'
-        st.error(f"Q&A generation failed: {str(e)}")
-        debug_log(traceback.format_exc())
+        st.error(f"Generation failed: {str(e)}")
         return False
 
-def main():
-    """Main application interface"""
-    st.set_page_config(
-        page_title="Synthetic Data Generator",
-        page_icon="🧪",
-        layout="wide"
-    )
-    
-    initialize_session_state()
+def export_data():
+    """Multi-format export handler"""
+    formats = st.session_state.export_formats
+    timestamp = time.strftime("%Y%m%d-%H%M%S")
+    base_name = f"wealth_report_{timestamp}"
+    export_package = {}
     
-    # Debug panel
-    with st.sidebar:
-        st.header("⚙️ Configuration")
-        st.session_state.model_choice = st.selectbox(
-            "AI Model", list(SUPPORTED_MODELS.keys())
-        )
-        st.session_state.temperature = st.slider(
-            "Creativity Level", 0.0, 1.0, 0.3
-        )
-        st.session_state.debug_mode = st.checkbox("Debug Mode", True)
-        show_processing_status()
+    try:
+        # JSON Export
+        if 'JSON' in formats:
+            export_package[f'{base_name}.json'] = json.dumps(
+                st.session_state.qa_pairs, 
+                indent=2
+            ).encode()
+        
+        # CSV Export
+        if 'CSV' in formats:
+            df = pd.DataFrame(st.session_state.qa_pairs)
+            export_package[f'{base_name}.csv'] = df.to_csv(index=False).encode()
+        
+        # Parquet Export
+        if 'Parquet' in formats:
+            df = pd.DataFrame(st.session_state.qa_pairs)
+            buffer = BytesIO()
+            df.to_parquet(buffer)
+            export_package[f'{base_name}.parquet'] = buffer.getvalue()
+        
+        return export_package
+    except Exception as e:
+        st.error(f"Export failed: {str(e)}")
+        return None
+
+def api_key_manager():
+    """Secure API key management UI"""
+    with st.sidebar.expander("🔑 API Key Management", expanded=True):
+        for model in SUPPORTED_MODELS:
+            key = st.text_input(
+                f"{model} API Key",
+                type="password",
+                key=f"key_{model}"
+            )
+            if key:
+                st.session_state.api_keys[SUPPORTED_MODELS[model]['required_key']] = key
+
+def main_interface():
+    """Core application interface"""
+    st.title("Global Wealth Report Analyzer")
+    st.write("Advanced financial document processing with multi-model AI support")
     
-    st.title("🧪 Synthetic Data Generator")
+    # File Upload
+    uploaded_file = st.file_uploader(
+        "Upload PDF Report", 
+        type=["pdf"],
+        accept_multiple_files=False
+    )
     
-    # File upload section
-    uploaded_file = st.file_uploader("Upload PDF Document", type=["pdf"])
+    # Processing Controls
+    if uploaded_file and st.button("Start Analysis"):
+        if process_pdf(uploaded_file) and generate_qa_content():
+            st.session_state.processing_stage = 'complete'
     
-    if uploaded_file and st.button("Start Processing"):
-        if pdf_processing_workflow(uploaded_file):
-            if generate_qa_pairs():
-                st.success("Processing completed successfully!")
-                
-                # Show results
-                st.header("Generated Q&A Pairs")
-                for idx, pair in enumerate(st.session_state.qa_pairs[:10]):
-                    with st.expander(f"Q{idx+1}: {pair['question']}"):
-                        st.write(f"**Answer 1:** {pair['answer_1']}")
-                        st.write(f"**Answer 2:** {pair['answer_2']}")
-                
-                # Data export
-                st.header("Data Export")
-                df = pd.DataFrame(st.session_state.qa_pairs)
+    # Results Display
+    if st.session_state.processing_stage == 'complete':
+        st.success("Analysis Complete!")
+        
+        # Data Export
+        with st.expander("📦 Export Results", expanded=True):
+            cols = st.columns(3)
+            with cols[0]:
+                st.multiselect(
+                    "Export Formats",
+                    ['JSON', 'CSV', 'Parquet'],
+                    default=['JSON', 'CSV'],
+                    key='export_formats'
+                )
+            with cols[1]:
                 st.download_button(
-                    label="Download as CSV",
-                    data=df.to_csv(index=False).encode('utf-8'),
-                    file_name="synthetic_data.csv",
-                    mime="text/csv"
+                    "Download Results",
+                    data=export_data()['wealth_report.json'],
+                    file_name="wealth_report.zip",
+                    mime="application/zip",
+                    disabled=not st.session_state.qa_pairs
                 )
-    
-    # Debug information
-    if st.session_state.debug_mode:
-        with st.expander("Debug Information"):
-            st.write("### Session State")
-            st.json(st.session_state)
-            
-            if st.session_state.get("document_data"):
-                st.write("### Document Data Summary")
-                st.write(f"Pages processed: {len(st.session_state.document_data)}")
-                st.write(f"Total images extracted: {sum(len(p['images']) for p in st.session_state.document_data)}")
-            
-            if st.session_state.get("qa_pairs"):
-                st.write("### Q&A Statistics")
-                st.write(f"Total pairs generated: {len(st.session_state.qa_pairs)}")
-                st.write("Sample Q&A pairs:")
-                st.table(pd.DataFrame(st.session_state.qa_pairs[:3]))
+        
+        # Results Preview
+        with st.expander("🔍 View Generated Content"):
+            st.dataframe(
+                pd.DataFrame(st.session_state.qa_pairs),
+                use_container_width=True,
+                height=400
+            )
+
+def model_settings():
+    """Model configuration panel"""
+    with st.sidebar.expander("🧠 AI Settings", expanded=True):
+        st.selectbox(
+            "AI Model",
+            list(SUPPORTED_MODELS.keys()),
+            key='model_settings.current_model'
+        )
+        st.slider(
+            "Creativity Level",
+            0.0, 1.0, 0.3,
+            key='model_settings.temperature'
+        )
 
 if __name__ == "__main__":
-    main()
+    initialize_session()
+    api_key_manager()
+    model_settings()
+    main_interface()