Update app.py

app.py

@@ -3,10 +3,16 @@ import streamlit as st
 import pdfplumber
 import pandas as pd
 import sqlalchemy
-from typing import Any, Dict, List, Optional
+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:
     from openai import OpenAI
@@ -27,7 +33,6 @@ except ImportError:
     genai = None
     Part = None
 
-import json
 
 class SyntheticDataGenerator:
     """World's Most Advanced Synthetic Data Generation System"""
@@ -36,27 +41,32 @@ class SyntheticDataGenerator:
         "Deepseek": {
             "base_url": "https://api.deepseek.com/v1",
             "models": ["deepseek-chat"],
-            "requires_library": "openai"
+            "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"
+            "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"
+            "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
+            "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"
-        }
+            "requires_library": "google.generativeai",
+            "supports_json_output": True
+        },
     }
 
     def __init__(self):
@@ -74,29 +84,50 @@ class SyntheticDataGenerator:
             "system_metrics": {
                 "api_calls": 0,
                 "tokens_used": 0,
-                "error_count": 0
+                "error_count": 0,
             },
             "debug_mode": False,
-            "google_configured": False,
-            "advanced_options": {
-                "temperature": 0.7,
-                "top_p": 0.95,
-                "top_k": 40,
-                "max_output_tokens": 2000
-            },
-            "generation_format": "json",
-            "csv_schema": ""
+            "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 = []
         for provider, config in self.PROVIDER_CONFIG.items():
             if config["requires_library"] and not globals().get(config["requires_library"].split('.')[0].title()):
-                continue
+                continue  # Skip providers with missing dependencies
             self.available_providers.append(provider)
 
     def _setup_input_handlers(self):
@@ -108,12 +139,13 @@ class SyntheticDataGenerator:
             "api": self._process_api,
             "database": self._process_database,
             "web": self._process_web,
-            "image": self._process_image
+            "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]:  # Allow "prompt" to be a list or a string
+    def generate(self, provider: str, model: str, prompt: Any) -> Dict[str, Any]:
         """Unified generation endpoint with failover support"""
         try:
             if provider not in self.available_providers:
@@ -135,7 +167,7 @@ class SyntheticDataGenerator:
         api_key = st.session_state.api_keys.get(provider, "")
 
         if not api_key and provider != "Google":
-            raise ValueError("API key required")
+            raise ValueError(f"API key required for provider: {provider}")
 
         try:
             if provider == "Groq":
@@ -158,92 +190,84 @@ class SyntheticDataGenerator:
                         raise ValueError(f"Error configuring Google API: {e}")
 
                 generation_config = genai.GenerationConfig(
-                    temperature=st.session_state.advanced_options["temperature"],
-                    top_p=st.session_state.advanced_options["top_p"],
-                    top_k=st.session_state.advanced_options["top_k"],
-                    max_output_tokens=st.session_state.advanced_options["max_output_tokens"]
+                    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 = [
-                    {
-                        "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"
-                    },
-                ]
-                return GenerativeModel(model_name=model, generation_config=generation_config, safety_settings=safety_settings)
+                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:
                 return OpenAI(
                     base_url=config["base_url"],
                     api_key=api_key,
-                    timeout=30
+                    timeout=st.session_state["api_call_timeout"],  # Use session state timeout
                 )
         except Exception as e:
             self._log_error(f"Client Init Failed: {str(e)}")
             return None
 
-    def _execute_generation(self, client, provider: str, model: str, prompt: Any) -> Dict[str, Any]: # Use Any for prompt type
-
+    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
 
-        if provider == "HuggingFace":
-            response = requests.post(
-                self.PROVIDER_CONFIG[provider]["base_url"] + model,
-                headers=client["headers"],
-                json={"inputs": prompt},
-                timeout=30
-            )
-            response.raise_for_status()
-            return response.json()
-        elif provider == "Google":
-            try:
-                if isinstance(prompt, list): #Multimodal case
-
-                    response = client.generate_content(prompt)
+        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(prompt)
-
+                response = client.generate_content(parts)  # Send parts to Google
 
                 content = response.text
-
-                if st.session_state.generation_format == "json":
+                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 {"content": content} #Return raw text
 
-            except Exception as e:
-                self._log_error(f"Google Generation Error: {str(e)}")
-                return {"error": str(e), "content": ""}
-        else:
-            completion = client.chat.completions.create(
-                model=model,
-                messages=[{"role": "user", "content": prompt}],
-                temperature=st.session_state.advanced_options["temperature"],
-                max_tokens=st.session_state.advanced_options["max_output_tokens"]
-            )
-            st.session_state.system_metrics["tokens_used"] += completion.usage.total_tokens
-            try:
-                return json.loads(completion.choices[0].message.content)
-            except json.JSONDecodeError:
-                return {"content": completion.choices[0].message.content,
-                        "warning": "Could not parse response as valid JSON.  Returning 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 _failover_generation(self, prompt: str) -> Dict[str, Any]:
         """Enterprise failover to secondary providers"""
@@ -271,19 +295,20 @@ class SyntheticDataGenerator:
             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 Exception as e:
+        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."""
         try:
             df = pd.read_csv(file)
-            column_names = df.columns.tolist()
-            data_types = [str(df[col].dtype) for col in df.columns]
-            schema_prompt = f"Column Names: {column_names}\nData Types: {data_types}"
-            st.session_state.csv_schema = schema_prompt
+            # Add more sophisticated CSV processing here, e.g., schema inference
             return df.to_string()
         except Exception as e:
             self._log_error(f"CSV Processing Error: {str(e)}")
@@ -293,17 +318,40 @@ class SyntheticDataGenerator:
         """Simple text passthrough processor"""
         return text
 
+    def _process_prompt_template(self, file) -> str:
+        """Process prompt template file and store the content in session_state"""
+        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
+        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=10)
+                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=10)
+                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()
+            response.raise_for_status()  # Raise HTTPError for bad responses
 
             try:
                 return json.dumps(response.json(), indent=2)
@@ -312,6 +360,9 @@ class SyntheticDataGenerator:
         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."""
@@ -325,17 +376,6 @@ class SyntheticDataGenerator:
             self._log_error(f"Database Processing Error: {str(e)}")
             return ""
 
-    def _process_image(self, image_file) -> list: #Returns a 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 []
-
     # --- Enterprise Features ---
     def _log_error(self, message: str) -> None:
         """Centralized error logging with telemetry"""
@@ -353,7 +393,7 @@ class SyntheticDataGenerator:
                 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:
@@ -369,24 +409,26 @@ class SyntheticDataGenerator:
                 return response.status_code == 200
             elif provider == "Google":
                 try:
-                    if not st.session_state.google_configured: #Check if google has been configured
+                    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
+                        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.
+                        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
+                            return False  # Cant test API if no API Key
 
-                        configure(api_key=api_key) #Configure API Key
+                        configure(api_key=api_key)  # Configure API Key
                         st.session_state.google_configured = True
-                        #st.write("configuring key")
+                        # st.write("configuring key")
 
-                    genai.GenerativeModel(model_name=self.PROVIDER_CONFIG["Google"]["models"][0]).generate_content("test") #Test a generation
+                    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
+                except Exception as e:  # Catch any exceptions
                     print(e)
                     return False
 
@@ -406,7 +448,7 @@ def provider_config_ui(gen: SyntheticDataGenerator):
         provider = st.selectbox(
             "AI Provider",
             gen.available_providers,
-            help="Available providers based on system configuration"
+            help="Available providers based on system configuration",
         )
         st.session_state.active_provider = provider
 
@@ -415,7 +457,7 @@ def provider_config_ui(gen: SyntheticDataGenerator):
             f"{provider} API Key",
             type="password",
             value=st.session_state.api_keys.get(provider, ""),
-            help=f"Obtain API key from {provider} portal"
+            help=f"Obtain API key from {provider} portal",
         )
         st.session_state.api_keys[provider] = api_key
 
@@ -423,50 +465,49 @@ def provider_config_ui(gen: SyntheticDataGenerator):
         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
-        if provider == "Google" or provider == "OpenAI":
-            st.subheader("Advanced Generation Options")
-            st.session_state.advanced_options["temperature"] = st.slider("Temperature", min_value=0.0,
-                                                                           max_value=1.0,
-                                                                           value=st.session_state.advanced_options[
-                                                                               "temperature"], step=0.05,
-                                                                           help="Controls randomness. Lower values = more deterministic.")
-            if provider == "Google":
-                st.session_state.advanced_options["top_p"] = st.slider("Top P", min_value=0.0, max_value=1.0,
-                                                                        value=st.session_state.advanced_options["top_p"],
-                                                                        step=0.05,
-                                                                        help="Nucleus sampling: Considers the most probable tokens.")
-                st.session_state.advanced_options["top_k"] = st.slider("Top K", min_value=1, max_value=100,
-                                                                        value=st.session_state.advanced_options["top_k"],
-                                                                        step=1,
-                                                                        help="Considers the top K most probable tokens.")
-
-            st.session_state.advanced_options["max_output_tokens"] = st.number_input("Max Output Tokens",
-                                                                                      min_value=50, max_value=4096,
-                                                                                      value=st.session_state.advanced_options[
-                                                                                          "max_output_tokens"], step=50,
-                                                                                      help="Maximum number of tokens in the generated output.")
-
-        st.session_state.generation_format = st.selectbox("Output Format", ["json", "text"],
-                                                           help="Choose the desired output format.")
+        # 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", "Image",
-                                 "Structured Prompt (Advanced)"])  # Add Image input, Add Structured Prompt (Advanced)
+                                ["Text", "PDF", "Web URL", "CSV", "Prompt Template",
+                                 "Image"])  # Add Image input, Add Structured Prompt (Advanced)
+
     input_content = None
-    additional_instructions = ""  # For structured prompt
 
     if input_method == "Text":
         input_content = st.text_area("Enter Text", height=200)
@@ -481,22 +522,16 @@ def input_ui():
         uploaded_file = st.file_uploader("Upload a CSV file", type=["csv"])
         if uploaded_file is not None:
             input_content = uploaded_file
-            if "csv_schema" in st.session_state:
-                st.write("Inferred CSV Schema:")
-                st.write(st.session_state.csv_schema)
-
+    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
 
-    elif input_method == "Structured Prompt (Advanced)":
-        st.subheader("Structured Prompt")
-        input_content = st.text_area("Enter the base prompt/instructions", height=100)
-        additional_instructions = st.text_area("Specify constraints, data format, or other requirements:",
-                                                height=100)
-
-    return input_method, input_content, additional_instructions
+    return input_method, input_content
 
 
 def main():
@@ -504,56 +539,62 @@ def main():
     st.set_page_config(
         page_title="Synthetic Data Factory Pro",
         page_icon="🏭",
-        layout="wide"
+        layout="wide",
     )
 
     gen = SyntheticDataGenerator()
 
     st.title("🏭 Synthetic Data Factory Pro")
-    st.markdown("""
+    st.markdown(
+        """
     **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, additional_instructions = input_ui()
+    input_method, input_content = input_ui()
 
     if st.button("Generate Data"):
-        if input_content or input_method == "Structured Prompt (Advanced)":
-            processed_input = None
-
-            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 == "Image":
-                processed_input = gen._process_image(input_content) #This is a list now
-                if not processed_input: #If something went wrong with image processing, don't proceed
-                    st.error("Error processing image.")
-                    return
-
-            elif input_method == "Structured Prompt (Advanced)":
-                processed_input = input_content + "\n" + additional_instructions
-
-            if processed_input:
-                try:
-                    if st.session_state.active_provider == "Google" and input_method == "Image":
-                        prompt_parts = [input_content] + processed_input #Keeps text and images separate for google
-                        result = gen.generate(st.session_state.active_provider, st.session_state.active_model, prompt_parts)
-                    else:
-                        result = gen.generate(st.session_state.active_provider, st.session_state.active_model, processed_input)
+        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)
-                except Exception as e:
-                    st.error(f"Error during generation: {e}")
-            else:
-                st.warning("No data to process. Please check your input.")
+                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.")