Update app.py

app.py

@@ -1,178 +1,274 @@
 import streamlit as st
 import pdfplumber
 import pytesseract
-from PIL import Image
-import json
 import pandas as pd
+import requests
+import json
+from PIL import Image
 from io import BytesIO
-import time
 from openai import OpenAI
 import groq
+import sqlalchemy
+from typing import Dict, Any
 
-class SyntheticDataFactory:
-    PROVIDER_CONFIG = {
-        "Deepseek": {
-            "client": lambda key: OpenAI(base_url="https://api.deepseek.com/v1", api_key=key),
-            "models": ["deepseek-chat"],
-            "key_label": "Deepseek API Key"
-        },
-        "OpenAI": {
-            "client": lambda key: OpenAI(api_key=key),
-            "models": ["gpt-4-turbo"],
-            "key_label": "OpenAI API Key"
-        },
-        "Groq": {
-            "client": lambda key: groq.Groq(api_key=key),
-            "models": ["mixtral-8x7b-32768", "llama2-70b-4096"],
-            "key_label": "Groq API Key"
-        }
-    }
-
+class SyntheticDataGenerator:
     def __init__(self):
-        self.init_session_state()
+        self.providers = {
+            "Deepseek": {
+                "client": lambda key: OpenAI(base_url="https://api.deepseek.com/v1", api_key=key),
+                "models": ["deepseek-chat"]
+            },
+            "OpenAI": {
+                "client": lambda key: OpenAI(api_key=key),
+                "models": ["gpt-4-turbo"]
+            },
+            "Groq": {
+                "client": lambda key: groq.Groq(api_key=key),
+                "models": ["mixtral-8x7b-32768"]
+            },
+            "HuggingFace": {
+                "client": lambda key: {"headers": {"Authorization": f"Bearer {key}"}},
+                "models": ["gpt2", "llama-2"]
+            }
+        }
+        
+        self.input_handlers = {
+            "pdf": self.handle_pdf,
+            "text": self.handle_text,
+            "csv": self.handle_csv,
+            "api": self.handle_api,
+            "db": self.handle_db
+        }
+        
+        self.init_session()
 
-    def init_session_state(self):
-        if 'qa_data' not in st.session_state:
-            st.session_state.qa_data = []
-        if 'processing' not in st.session_state:
-            st.session_state.processing = {
+    def init_session(self):
+        session_defaults = {
+            'inputs': [],
+            'qa_data': [],
+            'processing': {
                 'stage': 'idle',
-                'errors': [],
-                'progress': 0
+                'progress': 0,
+                'errors': []
+            },
+            'config': {
+                'provider': "Deepseek",
+                'model': "deepseek-chat",
+                'temperature': 0.3
             }
+        }
+        
+        for key, val in session_defaults.items():
+            if key not in st.session_state:
+                st.session_state[key] = val
 
-    def process_pdf(self, file):
-        """Process PDF with error handling"""
+    # Input Processors
+    def handle_pdf(self, file):
         try:
             with pdfplumber.open(file) as pdf:
-                pages = pdf.pages
-                for i, page in enumerate(pages):
-                    # Update progress
-                    st.session_state.processing['progress'] = (i+1)/len(pages)
-                    
-                    # Process page content
-                    text = page.extract_text() or ""
-                    images = self.process_images(page)
-                    
-                    # Store in session state
-                    st.session_state.qa_data.append({
-                        "page": i+1,
-                        "text": text,
-                        "images": images
-                    })
-                    time.sleep(0.1)  # Simulate processing
-            return True
+                return [{
+                    "text": page.extract_text() or "",
+                    "images": self.process_images(page),
+                    "meta": {"type": "pdf", "page": i+1}
+                } for i, page in enumerate(pdf.pages)]
         except Exception as e:
-            st.error(f"PDF processing failed: {str(e)}")
-            return False
+            self.log_error(f"PDF Error: {str(e)}")
+            return []
+
+    def handle_text(self, text):
+        return [{
+            "text": text,
+            "meta": {"type": "domain", "source": "manual"}
+        }]
+
+    def handle_csv(self, file):
+        try:
+            df = pd.read_csv(file)
+            return [{
+                "text": "\n".join([f"{col}: {row[col]}" for col in df.columns]),
+                "meta": {"type": "csv", "columns": list(df.columns)}
+            } for _, row in df.iterrows()]
+        except Exception as e:
+            self.log_error(f"CSV Error: {str(e)}")
+            return []
+
+    def handle_api(self, config):
+        try:
+            response = requests.get(config['url'], headers=config['headers'])
+            return [{
+                "text": json.dumps(response.json()),
+                "meta": {"type": "api", "endpoint": config['url']}
+            }]
+        except Exception as e:
+            self.log_error(f"API Error: {str(e)}")
+            return []
+
+    def handle_db(self, config):
+        try:
+            engine = sqlalchemy.create_engine(config['connection'])
+            with engine.connect() as conn:
+                result = conn.execute(sqlalchemy.text(config['query']))
+                return [{
+                    "text": "\n".join([f"{col}: {val}" for col, val in row._asdict().items()]),
+                    "meta": {"type": "db", "table": config.get('table', '')}
+                } for row in result]
+        except Exception as e:
+            self.log_error(f"DB Error: {str(e)}")
+            return []
 
     def process_images(self, page):
-        """Robust image processing"""
         images = []
         for img in page.images:
             try:
-                # Handle different PDF image formats
                 stream = img['stream']
-                width = int(stream.get('Width', stream.get('W', 0)))
-                height = int(stream.get('Height', stream.get('H', 0)))
-                
+                width = int(stream.get('Width', 0))
+                height = int(stream.get('Height', 0))
                 if width > 0 and height > 0:
-                    image = Image.frombytes(
-                        "RGB" if 'ColorSpace' in stream else "L",
-                        (width, height),
-                        stream.get_data()
-                    )
-                    images.append(image)
+                    images.append({
+                        "data": Image.frombytes("RGB", (width, height), stream.get_data()),
+                        "meta": {"dims": (width, height)}
+                    })
             except Exception as e:
-                st.warning(f"Image processing error: {str(e)[:100]}")
+                self.log_error(f"Image Error: {str(e)}")
         return images
 
-    def generate_qa(self, provider, api_key, model, temp):
-        """Generate Q&A pairs with selected provider"""
+    # Core Generation Engine
+    def generate(self, api_key: str) -> bool:
         try:
-            client = self.PROVIDER_CONFIG[provider]["client"](api_key)
+            provider_cfg = self.providers[st.session_state.config['provider']]
+            client = provider_cfg["client"](api_key)
             
-            for item in st.session_state.qa_data:
-                prompt = f"Generate 3 Q&A pairs from this financial content:\n{item['text']}\nOutput JSON format with keys: question, answer_1, answer_2"
+            for i, input_data in enumerate(st.session_state.inputs):
+                st.session_state.processing['progress'] = (i+1)/len(st.session_state.inputs)
                 
-                response = client.chat.completions.create(
-                    model=model,
-                    messages=[{"role": "user", "content": prompt}],
-                    temperature=temp,
-                    response_format={"type": "json_object"}
-                )
+                if st.session_state.config['provider'] == "HuggingFace":
+                    response = self._huggingface_inference(client, input_data)
+                else:
+                    response = self._standard_inference(client, input_data)
                 
-                try:
-                    result = json.loads(response.choices[0].message.content)
-                    item["qa_pairs"] = result.get("qa_pairs", [])
-                except json.JSONDecodeError:
-                    st.error("Failed to parse AI response")
+                if response:
+                    st.session_state.qa_data.extend(self._parse_response(response))
             
-            st.session_state.processing['stage'] = 'complete'
             return True
-            
         except Exception as e:
-            st.error(f"Generation failed: {str(e)}")
+            self.log_error(f"Generation Error: {str(e)}")
             return False
 
-def main():
-    st.set_page_config(
-        page_title="Enterprise Data Factory",
-        page_icon="🏭",
-        layout="wide"
-    )
-    
-    factory = SyntheticDataFactory()
-    
-    # Sidebar Configuration
+    def _standard_inference(self, client, input_data):
+        return client.chat.completions.create(
+            model=st.session_state.config['model'],
+            messages=[{
+                "role": "user",
+                "content": self._build_prompt(input_data)
+            }],
+            temperature=st.session_state.config['temperature'],
+            response_format={"type": "json_object"}
+        )
+
+    def _huggingface_inference(self, client, input_data):
+        API_URL = "https://api-inference.huggingface.co/models/"
+        response = requests.post(
+            API_URL + st.session_state.config['model'],
+            headers=client["headers"],
+            json={"inputs": self._build_prompt(input_data)}
+        )
+        return response.json()
+
+    def _build_prompt(self, input_data):
+        base = "Generate 3 Q&A pairs from this financial content:\n"
+        if input_data['meta']['type'] == 'csv':
+            return base + "Structured data:\n" + input_data['text']
+        elif input_data['meta']['type'] == 'api':
+            return base + "API response:\n" + input_data['text']
+        return base + input_data['text']
+
+    def _parse_response(self, response):
+        try:
+            if st.session_state.config['provider'] == "HuggingFace":
+                return response[0]['generated_text']
+            return json.loads(response.choices[0].message.content).get("qa_pairs", [])
+        except Exception as e:
+            self.log_error(f"Parse Error: {str(e)}")
+            return []
+
+    def log_error(self, message):
+        st.session_state.processing['errors'].append(message)
+        st.error(message)
+
+# Streamlit UI Components
+def input_sidebar(gen: SyntheticDataGenerator):
     with st.sidebar:
-        st.header("⚙️ AI Configuration")
-        provider = st.selectbox("Provider", list(factory.PROVIDER_CONFIG.keys()))
-        config = factory.PROVIDER_CONFIG[provider]
-        api_key = st.text_input(config["key_label"], type="password")
-        model = st.selectbox("Model", config["models"])
+        st.header("⚙️ Configuration")
+        
+        # AI Provider Settings
+        provider = st.selectbox("Provider", list(gen.providers.keys()))
+        provider_cfg = gen.providers[provider]
+        
+        api_key = st.text_input(f"{provider} API Key", type="password")
+        model = st.selectbox("Model", provider_cfg["models"])
         temp = st.slider("Temperature", 0.0, 1.0, 0.3)
+        
+        # Update session config
+        st.session_state.config.update({
+            "provider": provider,
+            "model": model,
+            "temperature": temp
+        })
+        
+        # Input Source Selection
+        st.header("🔗 Data Sources")
+        input_type = st.selectbox("Input Type", list(gen.input_handlers.keys()))
+        
+        if input_type == "text":
+            domain_input = st.text_area("Domain Knowledge", height=150)
+            if st.button("Add Domain Input"):
+                gen.input_handlers["text"](domain_input)
+        
+        elif input_type == "csv":
+            csv_file = st.file_uploader("Upload CSV", type=["csv"])
+            if csv_file:
+                gen.input_handlers["csv"](csv_file)
+        
+        elif input_type == "api":
+            api_url = st.text_input("API Endpoint")
+            if st.button("Connect API"):
+                gen.input_handlers["api"]({"url": api_url})
+        
+        return api_key
 
-    # Main Interface
+def main_display(gen: SyntheticDataGenerator):
     st.title("🚀 Enterprise Synthetic Data Factory")
     
-    uploaded_file = st.file_uploader("Upload Financial PDF", type=["pdf"])
+    # Input Processing
+    col1, col2 = st.columns([3, 1])
+    with col1:
+        pdf_file = st.file_uploader("Upload Document", type=["pdf"])
+        if pdf_file:
+            gen.input_handlers["pdf"](pdf_file)
     
-    if uploaded_file and api_key and st.button("Start Synthetic Generation"):
-        with st.status("Processing document...", expanded=True) as status:
-            # Process PDF
-            st.write("Extracting text and images...")
-            if factory.process_pdf(uploaded_file):
-                # Generate Q&A pairs
-                st.write("Generating synthetic data...")
-                if factory.generate_qa(provider, api_key, model, temp):
-                    status.update(label="Processing complete!", state="complete", expanded=False)
+    # Generation Controls
+    with col2:
+        if st.button("Start Generation"):
+            with st.status("Processing..."):
+                gen.generate(st.session_state.get('api_key'))
     
-    # Display Results
-    if st.session_state.processing.get('stage') == 'complete':
-        st.subheader("Generated Q&A Pairs")
-        
-        # Convert to DataFrame
-        all_qa = []
-        for item in st.session_state.qa_data:
-            for qa in item.get("qa_pairs", []):
-                qa["page"] = item["page"]
-                all_qa.append(qa)
+    # Results Display
+    if st.session_state.qa_data:
+        st.header("Generated Data")
+        df = pd.DataFrame(st.session_state.qa_data)
+        st.dataframe(df)
         
-        if len(all_qa) > 0:
-            df = pd.DataFrame(all_qa)
-            st.dataframe(df)
-            
-            # Export options
-            csv = df.to_csv(index=False).encode('utf-8')
-            st.download_button(
-                label="Download as CSV",
-                data=csv,
-                file_name="synthetic_data.csv",
-                mime="text/csv"
-            )
-        else:
-            st.warning("No Q&A pairs generated. Check your document content and API settings.")
+        # Export Options
+        st.download_button(
+            "Export CSV",
+            df.to_csv(index=False),
+            "synthetic_data.csv"
+        )
+
+def main():
+    gen = SyntheticDataGenerator()
+    api_key = input_sidebar(gen)
+    main_display(gen)
 
 if __name__ == "__main__":
     main()