Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 import pdfplumber
 import pytesseract
 import openai
+from openai import OpenAI
 import json
 import pandas as pd
 import numpy as np
@@ -11,227 +12,239 @@ import time
 import traceback
 import os
 import hashlib
+import groq
 
-# Configuration
-SUPPORTED_MODELS = {
-    "Deepseek": {
-        "base_url": "https://api.deepseek.com/v1",
-        "required_key": "DEEPSEEK_KEY"
-    },
-    "OpenAI": {
-        "base_url": "https://api.openai.com/v1",
-        "required_key": "OPENAI_KEY"
-    }
-}
-
-def initialize_session():
-    """Initialize session state with validation"""
-    defaults = {
-        'processing_stage': 'idle',
-        'document_data': [],
-        'qa_pairs': [],
-        'export_formats': ['JSON', 'CSV', 'Parquet'],
-        'model_settings': {
-            'current_model': 'Deepseek',
-            'temperature': 0.3
-        },
-        'api_keys': {}
-    }
-    
-    for key, val in defaults.items():
-        if key not in st.session_state:
-            st.session_state[key] = val
-
-def handle_image_errors(img_stream):
-    """Robust image processing with multiple fallbacks"""
-    try:
-        # First try standard RGB conversion
-        return Image.frombytes("RGB", (img_stream['width'], img_stream['height']), img_stream.get_data())
-    except:
+class SyntheticDataGenerator:
+    def __init__(self):
+        self.SUPPORTED_MODELS = {
+            "Deepseek": {
+                "client": lambda key: OpenAI(base_url="https://api.deepseek.com/v1", api_key=key),
+                "models": ["deepseek-chat"],
+                "key_name": "DEEPSEEK_KEY"
+            },
+            "OpenAI": {
+                "client": lambda key: OpenAI(api_key=key),
+                "models": ["gpt-4-turbo"],
+                "key_name": "OPENAI_KEY"
+            },
+            "Mistral-Groq": {
+                "client": lambda key: groq.Groq(api_key=key),
+                "models": ["mixtral-8x7b-32768", "llama2-70b-4096"],
+                "key_name": "GROQ_KEY"
+            }
+        }
+        self.init_session()
+
+    def init_session(self):
+        if 'qa_pairs' not in st.session_state:
+            st.session_state.qa_pairs = []
+        if 'doc_data' not in st.session_state:
+            st.session_state.doc_data = []
+        if 'processing' not in st.session_state:
+            st.session_state.processing = {
+                'stage': 'idle',
+                'errors': [],
+                'warnings': []
+            }
+
+    def process_pdf(self, uploaded_file):
+        """Robust PDF processing with advanced image handling"""
+        st.session_state.processing = {'stage': 'extracting', 'errors': [], 'warnings': []}
+        
         try:
-            # Fallback to grayscale conversion
-            return Image.frombytes("L", (img_stream['width'], img_stream['height']), img_stream.get_data()).convert("RGB")
+            with pdfplumber.load(uploaded_file) as pdf:
+                for page_num, page in enumerate(pdf.pages, 1):
+                    page_data = self._process_page(page, page_num)
+                    st.session_state.doc_data.append(page_data)
+            
+            if len(st.session_state.processing['errors']) > 0:
+                st.error(f"Processed with {len(st.session_state.processing['errors'])} errors")
+            return True
         except Exception as e:
-            st.error(f"Critical image error: {str(e)[:200]}")
-            return None
+            self._log_error(f"PDF loading failed: {str(e)}")
+            return False
 
-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:
-            for page_num, page in enumerate(pdf.pages, 1):
-                page_data = {
-                    "page": page_num,
-                    "text": page.extract_text() or "",
-                    "images": []
-                }
-                
-                # 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)
-                
-                doc_data.append(page_data)
-                time.sleep(0.01)  # Yield for UI updates
-                
-        st.session_state.document_data = doc_data
-        return True
-    except Exception as e:
-        st.error(f"PDF processing failed: {str(e)}")
-        return False
-
-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_settings['current_model']]['base_url'],
-            api_key=st.session_state.api_keys.get(
-                SUPPORTED_MODELS[st.session_state.model_settings['current_model']]['required_key']
-            )
-        )
+    def _process_page(self, page, page_num):
+        """Process individual page with nested error handling"""
+        page_data = {"page": page_num, "text": "", "images": []}
+        
+        try:
+            page_data["text"] = page.extract_text() or ""
+        except Exception as e:
+            self._log_error(f"Page {page_num} text extraction failed: {str(e)}")
+
+        try:
+            for img_idx, img in enumerate(page.images):
+                img_data = self._process_image(img, page_num, img_idx)
+                if img_data:
+                    page_data["images"].append(img_data)
+        except Exception as e:
+            self._log_error(f"Page {page_num} image processing failed: {str(e)}")
         
-        for page in st.session_state.document_data:
-            text_content = page['text'] or " ".join([
-                pytesseract.image_to_string(img) for img in page['images']
-            ])
+        return page_data
+
+    def _process_image(self, img, page_num, img_idx):
+        """Advanced image processing with multiple fallbacks"""
+        try:
+            stream = img['stream']
+            width = self._get_dimension(stream, 'width')
+            height = self._get_dimension(stream, 'height')
             
-            response = client.chat.completions.create(
-                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 this financial content:\n{text_content}\nOutput JSON format with keys: question, answer_1, answer_2"
-                }],
-                response_format={"type": "json_object"},
-                temperature=st.session_state.model_settings['temperature']
-            )
+            if width <= 0 or height <= 0:
+                raise ValueError("Invalid image dimensions")
             
             try:
-                result = json.loads(response.choices[0].message.content)
-                qa_pairs.extend(result.get('qa_pairs', []))
-            except json.JSONDecodeError:
-                st.error("Failed to parse model response")
-                
-        st.session_state.qa_pairs = qa_pairs
-        return True
-    except Exception as e:
-        st.error(f"Generation failed: {str(e)}")
-        return False
-
-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 = {}
-    
-    try:
-        # JSON Export
-        if 'JSON' in formats:
-            export_package[f'{base_name}.json'] = json.dumps(
-                st.session_state.qa_pairs, 
-                indent=2
-            ).encode()
+                return Image.frombytes("RGB", (width, height), stream.get_data())
+            except:
+                return Image.frombytes("L", (width, height), stream.get_data()).convert("RGB")
+        except Exception as e:
+            self._log_error(f"Page {page_num} image {img_idx} failed: {str(e)}")
+            return None
+
+    def _get_dimension(self, stream, dimension):
+        """Safe dimension extraction with multiple fallbacks"""
+        try:
+            return int(stream[dimension])
+        except:
+            try:
+                return int(stream['stream'][dimension])
+            except:
+                try:
+                    return int(stream['data'][dimension])
+                except:
+                    return 0
+
+    def generate_qa(self, model_provider, model_name, temperature):
+        """Multi-model generation engine"""
+        st.session_state.processing = {'stage': 'generating', 'errors': []}
+        qa_pairs = []
         
-        # 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()
+        try:
+            client = self.SUPPORTED_MODELS[model_provider]["client"](
+                st.session_state[model_provider.lower() + "_key"]
+            )
+            
+            for page in st.session_state.doc_data:
+                content = self._get_page_content(page)
+                response = self._generate(client, model_name, content, temperature)
+                qa_pairs.extend(self._parse_response(response))
+            
+            st.session_state.qa_pairs = qa_pairs
+            return True
+        except Exception as e:
+            self._log_error(f"Generation failed: {str(e)}")
+            return False
+
+    def _generate(self, client, model, content, temp):
+        """Unified generation interface"""
+        if isinstance(client, groq.Groq):
+            return client.chat.completions.create(
+                messages=[{"role": "user", "content": content}],
+                model=model,
+                temperature=temp,
+                response_format={"type": "json_object"}
+            )
+        else:
+            return client.chat.completions.create(
+                model=model,
+                messages=[{"role": "user", "content": content}],
+                temperature=temp,
+                response_format={"type": "json_object"}
+            )
+
+    def _parse_response(self, response):
+        """Safe response parsing"""
+        try:
+            content = json.loads(response.choices[0].message.content)
+            return content.get('qa_pairs', [])
+        except Exception as e:
+            self._log_error(f"Response parsing failed: {str(e)}")
+            return []
+
+    def export_data(self, formats):
+        """Multi-format export system"""
+        exports = {}
+        df = pd.DataFrame(st.session_state.qa_pairs)
         
-        # Parquet Export
+        if 'JSON' in formats:
+            exports['synthetic_data.json'] = df.to_json(orient='records').encode()
+        if 'CSV' in formats:
+            exports['synthetic_data.csv'] = df.to_csv(index=False).encode()
         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()
+            exports['synthetic_data.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",
+        return exports
+
+    def _log_error(self, message):
+        """Centralized error logging"""
+        st.session_state.processing['errors'].append(message)
+        st.error(message)
+
+    def _get_page_content(self, page):
+        """Multimodal content extraction"""
+        text = page["text"]
+        if not text:
+            text = " ".join([pytesseract.image_to_string(img) for img in page["images"]])
+        return text
+
+def ui_setup():
+    """Enterprise-grade UI configuration"""
+    st.set_page_config(
+        page_title="Synthetic Data Factory Pro",
+        page_icon="🏭",
+        layout="wide",
+        initial_sidebar_state="expanded"
+    )
+    
+    with st.sidebar:
+        st.header("🔑 API Key Management")
+        for provider in ["Deepseek", "OpenAI", "Mistral-Groq"]:
+            st.text_input(
+                f"{provider} API Key",
                 type="password",
-                key=f"key_{model}"
+                key=f"{provider.lower()}_key"
             )
-            if key:
-                st.session_state.api_keys[SUPPORTED_MODELS[model]['required_key']] = key
+        
+        st.header("🧠 AI Configuration")
+        provider = st.selectbox("Model Provider", ["Deepseek", "OpenAI", "Mistral-Groq"])
+        model = st.selectbox("Model", generator.SUPPORTED_MODELS[provider]["models"])
+        temp = st.slider("Temperature", 0.0, 1.0, 0.3)
+    
+    return provider, model, temp
 
-def main_interface():
-    """Core application interface"""
-    st.title("Global Wealth Report Analyzer")
-    st.write("Advanced financial document processing with multi-model AI support")
+def main():
+    """Main application flow"""
+    provider, model, temp = ui_setup()
+    generator = SyntheticDataGenerator()
     
-    # File Upload
-    uploaded_file = st.file_uploader(
-        "Upload PDF Report", 
-        type=["pdf"],
-        accept_multiple_files=False
-    )
+    st.title("🏭 Synthetic Data Factory Pro")
+    st.write("Enterprise-grade document processing with multi-modal AI")
     
-    # 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'
+    uploaded_file = st.file_uploader("Upload PDF Document", type=["pdf"])
     
-    # 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(
-                    "Download Results",
-                    data=export_data()['wealth_report.json'],
-                    file_name="wealth_report.zip",
-                    mime="application/zip",
-                    disabled=not st.session_state.qa_pairs
-                )
-        
-        # 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 uploaded_file and st.button("Start Generation"):
+        if generator.process_pdf(uploaded_file):
+            if generator.generate_qa(provider, model, temp):
+                st.success("Generation completed successfully!")
+                
+                with st.expander("📊 Results Preview"):
+                    st.dataframe(pd.DataFrame(st.session_state.qa_pairs))
+                
+                with st.expander("📦 Advanced Export"):
+                    formats = st.multiselect(
+                        "Select formats",
+                        ["JSON", "CSV", "Parquet"],
+                        default=["JSON", "CSV"]
+                    )
+                    exports = generator.export_data(formats)
+                    
+                    if st.download_button("Export Package", 
+                                        data=json.dumps(exports),
+                                        file_name="synthetic_data.zip",
+                                        mime="application/zip"):
+                        st.success("Export package generated!")
 
 if __name__ == "__main__":
-    initialize_session()
-    api_key_manager()
-    model_settings()
-    main_interface()
+    main()