Update app.py

app.py

@@ -1,3 +1,4 @@
+import json
 import requests
 import streamlit as st
 import pdfplumber
@@ -5,13 +6,6 @@ import pandas as pd
 import sqlalchemy
 from typing import Any, Dict, List, Optional, Union
 from functools import lru_cache
-import json  # Explicit import
-import os
-
-# --- Constants ---
-DEFAULT_TEMPERATURE = 0.1
-DEFAULT_MAX_TOKENS = 2000
-API_TIMEOUT = 30
 
 # Provider clients with import guards
 try:
@@ -24,58 +18,41 @@ try:
 except ImportError:
     groq = None
 
-try:
-    import google.generativeai as genai
-    from google.generativeai import GenerativeModel, configure, Part
-except ImportError:
-    GenerativeModel = None
-    configure = None
-    genai = None
-    Part = None
-
 
 class SyntheticDataGenerator:
     """World's Most Advanced Synthetic Data Generation System"""
 
-    PROVIDER_CONFIG = {
+    PROVIDER_CONFIG: Dict[str, Dict[str, Union[str, List[str], Optional[str]]]] = {
         "Deepseek": {
             "base_url": "https://api.deepseek.com/v1",
             "models": ["deepseek-chat"],
             "requires_library": "openai",
-            "supports_json_output": True,  # Indicate that the provider reliably returns JSON
         },
         "OpenAI": {
             "base_url": "https://api.openai.com/v1",
             "models": ["gpt-4-turbo", "gpt-3.5-turbo"],
             "requires_library": "openai",
-            "supports_json_output": True,
         },
         "Groq": {
             "base_url": "https://api.groq.com/openai/v1",
             "models": ["mixtral-8x7b-32768", "llama2-70b-4096"],
             "requires_library": "groq",
-            "supports_json_output": True,
         },
         "HuggingFace": {
             "base_url": "https://api-inference.huggingface.co/models/",
             "models": ["gpt2", "llama-2-13b-chat"],
             "requires_library": None,
-            "supports_json_output": False,  # More likely to return text
-        },
-        "Google": {
-            "models": ["gemini-1.5-flash-latest", "gemini-1.5-pro-latest", "gemini-pro", "gemini-pro-vision"],
-            "requires_library": "google.generativeai",
-            "supports_json_output": True
         },
     }
 
-    def __init__(self):
+    def __init__(self) -> None:
+        """Initialize session state, input handlers, and providers."""
         self._init_session_state()
         self._setup_input_handlers()
         self._setup_providers()
 
-    def _init_session_state(self):
-        """Initialize enterprise-grade session management"""
+    def _init_session_state(self) -> None:
+        """Initialize enterprise-grade session management with default values."""
         defaults = {
             "active_provider": "OpenAI",
             "api_keys": {},
@@ -86,52 +63,24 @@ class SyntheticDataGenerator:
                 "tokens_used": 0,
                 "error_count": 0,
             },
+            "error_logs": [],
             "debug_mode": False,
-            "temperature": DEFAULT_TEMPERATURE,  # Add temperature control
-            "max_tokens": DEFAULT_MAX_TOKENS,  # Add max token control
-            "use_streaming": False,  # Control Streaming behavior
-            "prompt_template": None,  # Support prompt templates
-            "api_call_timeout": API_TIMEOUT,  # API call timeout
-            "image_parts": [],  # Store image parts for multimodal generation
-            "top_p": 0.95,  # Default top_p for Google
-            "top_k": 40,  # Default top_k for Google
-            "safety_settings": self._get_default_safety_settings(), #Default Safety Settings
         }
         for key, val in defaults.items():
             if key not in st.session_state:
                 st.session_state[key] = val
 
-    def _get_default_safety_settings(self):
-      """Provides a default safety setting configuration for the Google provider"""
-      return [
-          {
-              "category": "HARM_CATEGORY_HARASSMENT",
-              "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-          },
-          {
-              "category": "HARM_CATEGORY_HATE_SPEECH",
-              "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-          },
-          {
-              "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
-              "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-          },
-          {
-              "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
-              "threshold": "BLOCK_MEDIUM_AND_ABOVE"
-          },
-      ]
-
-    def _setup_providers(self):
-        """Configure available providers with health checks"""
-        self.available_providers = []
+    def _setup_providers(self) -> None:
+        """Configure available providers with health checks."""
+        self.available_providers: List[str] = []
         for provider, config in self.PROVIDER_CONFIG.items():
-            if config["requires_library"] and not globals().get(config["requires_library"].split('.')[0].title()):
+            required_lib = config.get("requires_library")
+            if required_lib and not globals().get(required_lib.title()):
                 continue  # Skip providers with missing dependencies
             self.available_providers.append(provider)
 
-    def _setup_input_handlers(self):
-        """Register enterprise input processors"""
+    def _setup_input_handlers(self) -> None:
+        """Register enterprise input processors."""
         self.input_processors = {
             "text": self._process_text,
             "pdf": self._process_pdf,
@@ -139,149 +88,117 @@ class SyntheticDataGenerator:
             "api": self._process_api,
             "database": self._process_database,
             "web": self._process_web,
-            "prompt_template": self._process_prompt_template,
-            "image": self._process_image,
         }
 
-    # --- Core Generation Engine ---
     @lru_cache(maxsize=100)
-    def generate(self, provider: str, model: str, prompt: Any) -> Dict[str, Any]:
-        """Unified generation endpoint with failover support"""
+    def generate(self, provider: str, model: str, prompt: str) -> Dict[str, Any]:
+        """
+        Unified generation endpoint with failover support.
+        Caches results to improve performance.
+        """
         try:
             if provider not in self.available_providers:
-                raise ValueError(f"Provider {provider} not available")
+                raise ValueError(f"Provider {provider} is not available.")
 
             client = self._get_client(provider)
             if not client:
-                raise ConnectionError("Client initialization failed")
+                raise ConnectionError("Client initialization failed.")
 
             return self._execute_generation(client, provider, model, prompt)
 
         except Exception as e:
-            self._log_error(f"Generation Error: {str(e)}")
-            return self._failover_generation(prompt)
+            self._log_error(f"Generation Error with provider '{provider}': {str(e)}")
+            return self._failover_generation(provider, model, prompt)
 
     def _get_client(self, provider: str) -> Any:
-        """Secure client initialization with connection pooling"""
+        """
+        Secure client initialization with connection pooling.
+        Raises ValueError if API key is missing.
+        """
         config = self.PROVIDER_CONFIG[provider]
-        api_key = st.session_state.api_keys.get(provider, "")
+        api_key = st.session_state["api_keys"].get(provider, "")
 
-        if not api_key and provider != "Google":
-            raise ValueError(f"API key required for provider: {provider}")
+        if not api_key:
+            raise ValueError(f"API key required for provider {provider}.")
 
         try:
             if provider == "Groq":
                 return groq.Groq(api_key=api_key)
             elif provider == "HuggingFace":
                 return {"headers": {"Authorization": f"Bearer {api_key}"}}
-            elif provider == "Google":
-                if not st.session_state.google_configured:
-                    if "GOOGLE_API_KEY" in os.environ:
-                        api_key = os.environ["GOOGLE_API_KEY"]
-                    else:
-                        api_key = st.session_state.api_keys.get("Google", "")
-                        if not api_key:
-                            raise ValueError(
-                                "Google API key is required. Please set it in the app or as the GOOGLE_API_KEY environment variable.")
-                    try:
-                        configure(api_key=api_key)  # Moved configure into try block
-                        st.session_state.google_configured = True
-                    except Exception as e:
-                        raise ValueError(f"Error configuring Google API: {e}")
-
-                generation_config = genai.GenerationConfig(
-                    temperature=st.session_state["temperature"],
-                    top_p=st.session_state["top_p"],
-                    top_k=st.session_state["top_k"],
-                    max_output_tokens=st.session_state["max_tokens"],
-                )
-                safety_settings = st.session_state["safety_settings"] #Get Safety Settings
-
-                return GenerativeModel(model_name=model, generation_config=generation_config,
-                                       safety_settings=safety_settings)  # Use all settings
             else:
+                # Initialize OpenAI client for OpenAI and Deepseek
                 return OpenAI(
                     base_url=config["base_url"],
                     api_key=api_key,
-                    timeout=st.session_state["api_call_timeout"],  # Use session state timeout
+                    timeout=30
                 )
         except Exception as e:
-            self._log_error(f"Client Init Failed: {str(e)}")
+            self._log_error(f"Client Initialization Failed for {provider}: {str(e)}")
             return None
 
-    def _execute_generation(self, client, provider: str, model: str, prompt: Any) -> Dict[str, Any]:
-        """Execute provider-specific generation with circuit breaker"""
-        st.session_state.system_metrics["api_calls"] += 1
-
-        try:
-            if provider == "HuggingFace":
-                response = requests.post(
-                    self.PROVIDER_CONFIG[provider]["base_url"] + model,
-                    headers=client["headers"],
-                    json={"inputs": prompt},
-                    timeout=st.session_state["api_call_timeout"]
-                )
-                response.raise_for_status()  # Raise HTTPError for bad responses
-                return response.json()
-            elif provider == "Google":
-                # Construct parts list.  If prompt is already a list, assume it contains Parts and text
-
-                if isinstance(prompt, str):
-                  parts = [prompt] #If plain text
-                else:
-                    parts = prompt #Multimodal prompt
-
-                response = client.generate_content(parts)  # Send parts to Google
-
-                content = response.text
-                if self.PROVIDER_CONFIG[provider]["supports_json_output"]:
-                    try:
-                        return json.loads(content)
-                    except json.JSONDecodeError:
-                        return {"content": content,
-                                "warning": "Could not parse response as valid JSON.  Returning raw text."}
-                else:
-                    return {"content": content} #Return raw text
-
-            else:
-                completion = client.chat.completions.create(
-                    model=model,
-                    messages=[{"role": "user", "content": prompt}],
-                    temperature=st.session_state["temperature"],  # Get temperature from session
-                    max_tokens=st.session_state["max_tokens"],  # Get max_tokens from session
-                    stream=st.session_state["use_streaming"],  # Use streaming bool from session
-                )
-                st.session_state.system_metrics["tokens_used"] += completion.usage.total_tokens
-                content = completion.choices[0].message.content
-                # Attempt to parse JSON if supported, otherwise return text
-                if self.PROVIDER_CONFIG[provider]["supports_json_output"]:
-                    try:
-                        return json.loads(content)
-                    except json.JSONDecodeError:
-                        return {"content": content,
-                                "warning": "Could not parse response as valid JSON.  Returning raw text."}
-                else:
-                    return {"content": content}  # return raw text
-        except requests.exceptions.RequestException as e:
-            self._log_error(f"API Request Error: {str(e)}")
-            return {"error": str(e), "content": ""}
-        except Exception as e:
-            self._log_error(f"Generation Error: {str(e)}")
-            return {"error": str(e), "content": ""}
+    def _execute_generation(self, client: Any, provider: str, model: str, prompt: str) -> Dict[str, Any]:
+        """
+        Execute provider-specific generation with circuit breaker.
+        Increments API call and token usage metrics.
+        """
+        st.session_state["system_metrics"]["api_calls"] += 1
+
+        if provider == "HuggingFace":
+            url = self.PROVIDER_CONFIG[provider]["base_url"] + model
+            response = requests.post(
+                url,
+                headers=client["headers"],
+                json={"inputs": prompt},
+                timeout=30
+            )
+            response.raise_for_status()
+            return response.json()
+        else:
+            completion = client.chat.completions.create(
+                model=model,
+                messages=[{"role": "user", "content": prompt}],
+                temperature=0.1,
+                max_tokens=2000
+            )
+            # Update token usage if available
+            if hasattr(completion.usage, "total_tokens"):
+                st.session_state["system_metrics"]["tokens_used"] += completion.usage.total_tokens
+            try:
+                result = json.loads(completion.choices[0].message.content)
+            except json.JSONDecodeError:
+                result = {"response": completion.choices[0].message.content}
+            return result
 
-    def _failover_generation(self, prompt: str) -> Dict[str, Any]:
-        """Enterprise failover to secondary providers"""
+    def _failover_generation(self, original_provider: str, model: str, prompt: str) -> Dict[str, Any]:
+        """
+        Enterprise failover to secondary providers.
+        Tries available providers (excluding the original) until one succeeds.
+        """
         for backup_provider in self.available_providers:
-            if backup_provider != st.session_state.active_provider:
-                try:
-                    return self.generate(backup_provider, ..., prompt=prompt)
-                except Exception:
-                    continue
-        raise RuntimeError("All generation providers unavailable")
+            if backup_provider == original_provider:
+                continue
+            # Determine model to use: if the desired model is available, use it; otherwise use the first available model.
+            backup_models = self.PROVIDER_CONFIG[backup_provider]["models"]
+            backup_model = model if model in backup_models else backup_models[0]
+            try:
+                st.session_state["active_provider"] = backup_provider
+                result = self.generate(backup_provider, backup_model, prompt)
+                # Optionally, log the failover event
+                self._log_error(f"Failover succeeded with provider '{backup_provider}' using model '{backup_model}'.")
+                return result
+            except Exception as e:
+                self._log_error(f"Failover attempt with provider '{backup_provider}' failed: {str(e)}")
+                continue
+        raise RuntimeError("All generation providers are unavailable.")
 
     # --- Input Processors ---
+    def _process_text(self, text: str) -> str:
+        """Process plain text input by stripping unnecessary whitespace."""
+        return text.strip()
+
     def _process_pdf(self, file) -> str:
-        """Advanced PDF processing with OCR fallback"""
+        """Advanced PDF processing with OCR fallback."""
         try:
             with pdfplumber.open(file) as pdf:
                 return "\n".join(page.extract_text() or "" for page in pdf.pages)
@@ -289,268 +206,129 @@ class SyntheticDataGenerator:
             self._log_error(f"PDF Processing Error: {str(e)}")
             return ""
 
-    def _process_web(self, url: str) -> str:
-        """Web content extraction with anti-bot measures"""
-        try:
-            response = requests.get(url, headers={
-                "User-Agent": "Mozilla/5.0 (compatible; SyntheticBot/1.0)"
-            }, timeout=10)
-            response.raise_for_status()  # Raises HTTPError for bad responses (4xx or 5xx)
-            return response.text
-        except requests.exceptions.RequestException as e:
-            self._log_error(f"Web Extraction Error: {str(e)}")
-            return ""
-        except Exception as e:
-            self._log_error(f"Unexpected Web Extraction Error: {str(e)}")
-            return ""
-
     def _process_csv(self, file) -> str:
-        """Process CSV files and return as a string representation."""
+        """Process CSV input by reading it into a DataFrame and converting to CSV string."""
         try:
             df = pd.read_csv(file)
-            # Add more sophisticated CSV processing here, e.g., schema inference
-            return df.to_string()
+            return df.to_csv(index=False)
         except Exception as e:
             self._log_error(f"CSV Processing Error: {str(e)}")
             return ""
 
-    def _process_text(self, text: str) -> str:
-        """Simple text passthrough processor"""
-        return text
-
-    def _process_prompt_template(self, file) -> str:
-        """Process prompt template file and store the content in session_state"""
+    def _process_api(self, api_url: str) -> str:
+        """Process API input by fetching JSON data from an endpoint."""
         try:
-            template_content = file.read().decode("utf-8")  # Read file content
-            st.session_state["prompt_template"] = template_content  # Store in session_state
-            return "Prompt template uploaded and stored."  # Inform the user
+            response = requests.get(api_url, timeout=10)
+            response.raise_for_status()
+            data = response.json()
+            return json.dumps(data, indent=2)
         except Exception as e:
-            self._log_error(f"Prompt Template Processing Error: {str(e)}")
-            return ""
-
-    def _process_image(self, image_file) -> list:
-        """Processes image files for multimodal generation (Google Gemini)"""
-        try:
-            image_data = image_file.read()
-            image_part = Part.from_data(image_data, mime_type=image_file.type)  # Use Part for google
-            return [image_part]  # Return a list with the image part as a Google Part object
-
-        except Exception as e:
-            self._log_error(f"Image Processing Error: {str(e)}")
-            return []
-
-    def _process_api(self, url: str, method="GET", headers: Optional[Dict[str, str]] = None,
-                     data: Optional[Dict[str, Any]] = None) -> str:
-        """Generic API endpoint processor with configurable methods and headers."""
-        try:
-            if method.upper() == "GET":
-                response = requests.get(url, headers=headers or {},
-                                        timeout=st.session_state["api_call_timeout"])
-            elif method.upper() == "POST":
-                response = requests.post(url, headers=headers or {}, json=data,
-                                         timeout=st.session_state["api_call_timeout"])
-            else:
-                raise ValueError("Unsupported HTTP method.")
-            response.raise_for_status()  # Raise HTTPError for bad responses
-
-            try:
-                return json.dumps(response.json(), indent=2)
-            except json.JSONDecodeError:
-                return response.text
-        except requests.exceptions.RequestException as e:
             self._log_error(f"API Processing Error: {str(e)}")
             return ""
-        except Exception as e:
-            self._log_error(f"Unexpected API Processing Error: {str(e)}")
-            return ""
 
-    def _process_database(self, connection_string: str, query: str) -> str:
-        """Database query processor using SQLAlchemy."""
+    def _process_database(self, config: Dict[str, str]) -> str:
+        """
+        Process database input by executing a query.
+        Expects a configuration dictionary with 'connection_string' and 'query' keys.
+        """
         try:
+            connection_string = config.get("connection_string", "")
+            query = config.get("query", "")
+            if not connection_string or not query:
+                raise ValueError("Missing connection string or query.")
             engine = sqlalchemy.create_engine(connection_string)
             with engine.connect() as connection:
-                result = connection.execute(sqlalchemy.text(query))
-                df = pd.DataFrame(result.fetchall(), columns=result.keys())
-                return df.to_string()
+                result = pd.read_sql(query, connection)
+            return result.to_csv(index=False)
         except Exception as e:
             self._log_error(f"Database Processing Error: {str(e)}")
             return ""
 
+    def _process_web(self, url: str) -> str:
+        """Web content extraction with anti-bot measures."""
+        try:
+            response = requests.get(url, headers={
+                "User-Agent": "Mozilla/5.0 (compatible; SyntheticBot/1.0)"
+            }, timeout=10)
+            response.raise_for_status()
+            return response.text
+        except Exception as e:
+            self._log_error(f"Web Extraction Error: {str(e)}")
+            return ""
+
     # --- Enterprise Features ---
     def _log_error(self, message: str) -> None:
-        """Centralized error logging with telemetry"""
-        st.session_state.system_metrics["error_count"] += 1
-        st.session_state.error_logs = st.session_state.get("error_logs", []) + [message]
-
-        if st.session_state.debug_mode:
+        """Centralized error logging with telemetry."""
+        st.session_state["system_metrics"]["error_count"] += 1
+        st.session_state["error_logs"].append(message)
+        if st.session_state.get("debug_mode"):
             st.error(f"[DEBUG] {message}")
 
     def health_check(self) -> Dict[str, Any]:
-        """Comprehensive system diagnostics"""
+        """Comprehensive system diagnostics."""
         return {
             "providers_available": self.available_providers,
             "api_connectivity": {
                 provider: self._test_provider_connectivity(provider)
                 for provider in self.available_providers
             },
-            "system_metrics": st.session_state.system_metrics,
+            "system_metrics": st.session_state["system_metrics"],
         }
 
     def _test_provider_connectivity(self, provider: str) -> bool:
-        """Provider-specific connectivity test"""
+        """Provider-specific connectivity test."""
         try:
             client = self._get_client(provider)
             if provider == "HuggingFace":
-                response = requests.get(
-                    self.PROVIDER_CONFIG[provider]["base_url"],
-                    headers=client["headers"],
-                    timeout=5
-                )
+                url = self.PROVIDER_CONFIG[provider]["base_url"]
+                response = requests.get(url, headers=client["headers"], timeout=5)
                 return response.status_code == 200
-            elif provider == "OpenAI":
-                try:
-                    client.models.list()
-                    return True
-                except Exception:
-                    return False
-            elif provider == "Groq":
-                try:
-                    client.models.list()
-                    return True
-                except Exception:
-                    return False
-            elif provider == "Google":
-                try:
-                    if not st.session_state.google_configured:  # Check if google has been configured
-
-                        api_key = st.session_state.api_keys.get("Google",
-                                                               "")  # Get Key from session state
-
-                        if not api_key:  # If that is not set, check environment variable.
-                            api_key = os.environ.get("GOOGLE_API_KEY")
-
-                        if not api_key:
-                            return False  # Cant test API if no API Key
-
-                        configure(api_key=api_key)  # Configure API Key
-                        st.session_state.google_configured = True
-
-                    genai.GenerativeModel(model_name=self.PROVIDER_CONFIG["Google"]["models"][0]).generate_content(
-                        "test")  # Test a generation
-                    return True
-
-                except Exception as e:  # Catch any exceptions
-                    print(e)
-                    return False
-
             else:
-                return False  # Unknown provider
-
+                # Assuming client has a models.list() method for other providers.
+                client.models.list()
+                return True
         except Exception:
             return False
 
+
 # --- Enterprise UI Components ---
-def provider_config_ui(gen: SyntheticDataGenerator):
-    """Advanced provider configuration interface"""
+def provider_config_ui(gen: SyntheticDataGenerator) -> None:
+    """Advanced provider configuration interface."""
     with st.sidebar:
         st.header("⚙️ AI Engine Configuration")
-
-        # Provider selection with availability checks
         provider = st.selectbox(
             "AI Provider",
             gen.available_providers,
             help="Available providers based on system configuration",
+            index=gen.available_providers.index(st.session_state.get("active_provider", "OpenAI"))
         )
-        st.session_state.active_provider = provider
+        st.session_state["active_provider"] = provider
 
-        # API key management
         api_key = st.text_input(
             f"{provider} API Key",
             type="password",
-            value=st.session_state.api_keys.get(provider, ""),
-            help=f"Obtain API key from {provider} portal",
+            value=st.session_state["api_keys"].get(provider, ""),
+            help=f"Obtain API key from {provider}'s portal"
         )
-        st.session_state.api_keys[provider] = api_key
+        st.session_state["api_keys"][provider] = api_key
 
-        # Model selection
         model = st.selectbox(
             "Model",
             gen.PROVIDER_CONFIG[provider]["models"],
-            help="Select model version based on your API plan",
+            help="Select model version based on your API plan"
         )
-        st.session_state.active_model = model
-
-        # Advanced options
-        st.subheader("Advanced Options")
-        st.session_state["temperature"] = st.slider("Temperature", 0.0, 1.0, DEFAULT_TEMPERATURE, 0.05)
-        st.session_state["max_tokens"] = st.number_input("Max Tokens", 50, 4000, DEFAULT_MAX_TOKENS, 50)
-        st.session_state["use_streaming"] = st.checkbox("Enable Streaming")
-        st.session_state["api_call_timeout"] = st.slider("API Call Timeout (seconds)", 5, 60, API_TIMEOUT, 5)
-
-        # Google Specific Options
-        if provider == "Google":
-          st.subheader("Google Specific Settings")
-          st.session_state["top_p"] = st.slider("Top P", 0.0, 1.0, 0.95, 0.05, help="Nucleus sampling: Considers the most probable tokens.")
-          st.session_state["top_k"] = st.slider("Top K", 1, 100, 40, 1, help="Considers the top K most probable tokens.")
-
-        # Safety Settings Configuration
-          st.subheader("Safety Settings")
-          safety_categories = ["HARM_CATEGORY_HARASSMENT", "HARM_CATEGORY_HATE_SPEECH", "HARM_CATEGORY_SEXUALLY_EXPLICIT", "HARM_CATEGORY_DANGEROUS_CONTENT"]
-          threshold_options = ["BLOCK_NONE", "BLOCK_LOW_AND_ABOVE", "BLOCK_MEDIUM_AND_ABOVE", "BLOCK_ONLY_HIGH",]
-
-          for category in safety_categories:
-              threshold = st.selectbox(f"Threshold for {category}", options=threshold_options, index=2, key=f"{category}_threshold")  # Start with Medium and Above
-              #Update Threshold
-              for setting in st.session_state["safety_settings"]:
-                  if setting["category"] == category:
-                      setting["threshold"] = threshold
-                      break
-
-
-        # System monitoring
+
         if st.button("Run Health Check"):
             report = gen.health_check()
             st.json(report)
 
-def input_ui():
-    """Creates the input method UI"""
-    input_method = st.selectbox("Input Method",
-                                ["Text", "PDF", "Web URL", "CSV", "Prompt Template",
-                                 "Image"])  # Add Image input, Add Structured Prompt (Advanced)
-
-    input_content = None
-
-    if input_method == "Text":
-        input_content = st.text_area("Enter Text", height=200)
-    elif input_method == "PDF":
-        uploaded_file = st.file_uploader("Upload a PDF file", type=["pdf"])
-        if uploaded_file is not None:
-            input_content = uploaded_file
-    elif input_method == "Web URL":
-        url = st.text_input("Enter Web URL")
-        input_content = url
-    elif input_method == "CSV":
-        uploaded_file = st.file_uploader("Upload a CSV file", type=["csv"])
-        if uploaded_file is not None:
-            input_content = uploaded_file
-    elif input_method == "Prompt Template":
-        uploaded_file = st.file_uploader("Upload a Prompt Template file", type=["txt", "j2"])
-        if uploaded_file is not None:
-            input_content = uploaded_file
-    elif input_method == "Image":
-        uploaded_file = st.file_uploader("Upload an Image file", type=["png", "jpg", "jpeg"])
-        if uploaded_file is not None:
-            input_content = uploaded_file
-
-    return input_method, input_content
-
-
-def main():
-    """Enterprise-grade user interface"""
+
+def main() -> None:
+    """Enterprise-grade user interface."""
     st.set_page_config(
         page_title="Synthetic Data Factory Pro",
         page_icon="🏭",
-        layout="wide",
+        layout="wide"
     )
 
     gen = SyntheticDataGenerator()
@@ -558,56 +336,55 @@ def main():
     st.title("🏭 Synthetic Data Factory Pro")
     st.markdown(
         """
-    **World's Most Advanced Synthetic Data Generation Platform**
-    *Multi-provider AI Engine | Enterprise Input Processors | Real-time Monitoring*
-    """
+        **World's Most Advanced Synthetic Data Generation Platform**  
+        *Multi-provider AI Engine | Enterprise Input Processors | Real-time Monitoring*
+        """
     )
 
     provider_config_ui(gen)
 
-    input_method, input_content = input_ui()
+    # Example: Input management and processing
+    st.subheader("Input Data")
+    input_type = st.selectbox("Select Input Type", list(gen.input_processors.keys()))
+    if input_type == "text":
+        user_input = st.text_area("Enter your text here:")
+    elif input_type == "pdf":
+        user_input = st.file_uploader("Upload a PDF file", type=["pdf"])
+    elif input_type == "csv":
+        user_input = st.file_uploader("Upload a CSV file", type=["csv"])
+    elif input_type == "api":
+        user_input = st.text_input("Enter API URL:")
+    elif input_type == "database":
+        user_input = st.text_area("Enter Database Config as JSON (with 'connection_string' and 'query'):")
+    elif input_type == "web":
+        user_input = st.text_input("Enter Website URL:")
+
+    processed_input = ""
+    if st.button("Process Input"):
+        processor = gen.input_processors.get(input_type)
+        if processor:
+            if input_type in ["pdf", "csv"]:
+                processed_input = processor(user_input)
+            elif input_type == "database":
+                try:
+                    db_config = json.loads(user_input)
+                    processed_input = processor(db_config)
+                except json.JSONDecodeError:
+                    st.error("Invalid JSON for database configuration.")
+            else:
+                processed_input = processor(user_input)
+            st.text_area("Processed Input", value=processed_input, height=200)
+
+    # Generation section
+    st.subheader("Generate Synthetic Data")
+    prompt = st.text_area("Enter your prompt for data generation:")
+    if st.button("Generate"):
+        active_provider = st.session_state.get("active_provider", "OpenAI")
+        # Allow model selection for the generation step
+        model = st.selectbox("Select Generation Model", gen.PROVIDER_CONFIG[active_provider]["models"])
+        result = gen.generate(active_provider, model, prompt)
+        st.json(result)
 
-    if st.button("Generate Data"):
-        if input_content:
-            try:
-                if input_method == "Text":
-                    processed_input = gen._process_text(input_content)
-                elif input_method == "PDF":
-                    processed_input = gen._process_pdf(input_content)
-                elif input_method == "Web URL":
-                    processed_input = gen._process_web(input_content)
-                elif input_method == "CSV":
-                    processed_input = gen._process_csv(input_content)
-                elif input_method == "Prompt Template":
-                    processed_input = gen._process_prompt_template(
-                        input_content)  # Process the uploaded template
-                elif input_method == "Image":
-                    processed_input = gen._process_image(input_content)  # Returns a List of Parts
-
-                # If a prompt template is loaded, use it.
-                if st.session_state["prompt_template"] is not None and input_method != "Prompt Template":
-                    try:
-                        from jinja2 import Template  # Conditionally import it.
-
-                        template = Template(st.session_state["prompt_template"])  # Load Jinja2 Template
-                        processed_input = template.render(
-                            input=processed_input)  # Render the template - Overwrites the Input, Google needs parts, not text
-
-                    except Exception as e:
-                        st.error(f"Error rendering prompt template: {e}")
-                        st.stop()  # Stop the app if template rendering fails
-
-                if processed_input:
-                    result = gen.generate(st.session_state.active_provider, st.session_state.active_model,
-                                           processed_input)
-                    st.subheader("Generated Output:")
-                    st.json(result)
-                else:
-                    st.warning("No data to process. Please check your input.")
-            except Exception as e:
-                st.error(f"An unexpected error occurred: {e}")
-        else:
-            st.warning("Please provide input data.")
 
 if __name__ == "__main__":
-    main()
+    main()