Update app.py

app.py

@@ -5,7 +5,7 @@ import pandas as pd
 import sqlalchemy
 from typing import Any, Dict, List, Optional
 from functools import lru_cache
-import os # Import the 'os' module
+import os
 
 # Provider clients with import guards
 try:
@@ -20,17 +20,18 @@ except ImportError:
 
 try:
     import google.generativeai as genai
-    from google.generativeai import GenerativeModel, configure
+    from google.generativeai import GenerativeModel, configure, Part
 except ImportError:
     GenerativeModel = None
     configure = None
-    genai = None #Also set this to none
+    genai = None
+    Part = None
 
-import json  # Ensure json is explicitly imported for enhanced use
+import json
 
 class SyntheticDataGenerator:
     """World's Most Advanced Synthetic Data Generation System"""
-    
+
     PROVIDER_CONFIG = {
         "Deepseek": {
             "base_url": "https://api.deepseek.com/v1",
@@ -53,7 +54,7 @@ class SyntheticDataGenerator:
             "requires_library": None
         },
         "Google": {
-            "models": ["gemini-1.5-flash-latest", "gemini-1.5-pro-latest"],  # Include Gemini 2.0 Flash
+            "models": ["gemini-1.5-flash-latest", "gemini-1.5-pro-latest", "gemini-pro", "gemini-pro-vision"],
             "requires_library": "google.generativeai"
         }
     }
@@ -76,14 +77,15 @@ class SyntheticDataGenerator:
                 "error_count": 0
             },
             "debug_mode": False,
-            "google_configured": False,  # Track if Google API is configured
-            "advanced_options": {       # Store advanced generation options
-                "temperature": 0.7,    # Default temperature
-                "top_p": 0.95,          # Default top_p
-                "top_k": 40,            # Default top_k
-                "max_output_tokens": 2000 # Default max_output_tokens
+            "google_configured": False,
+            "advanced_options": {
+                "temperature": 0.7,
+                "top_p": 0.95,
+                "top_k": 40,
+                "max_output_tokens": 2000
             },
-            "generation_format": "json" # Default output format (json or text)
+            "generation_format": "json",
+            "csv_schema": ""
         }
         for key, val in defaults.items():
             if key not in st.session_state:
@@ -94,7 +96,7 @@ class SyntheticDataGenerator:
         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  # Skip providers with missing dependencies
+                continue
             self.available_providers.append(provider)
 
     def _setup_input_handlers(self):
@@ -106,12 +108,12 @@ class SyntheticDataGenerator:
             "api": self._process_api,
             "database": self._process_database,
             "web": self._process_web,
-            "image": self._process_image #Add Image
+            "image": self._process_image
         }
 
     # --- Core Generation Engine ---
     @lru_cache(maxsize=100)
-    def generate(self, provider: str, model: str, prompt: str) -> Dict[str, Any]:
+    def generate(self, provider: str, model: str, prompt: Any) -> Dict[str, Any]:  # Allow "prompt" to be a list or a string
         """Unified generation endpoint with failover support"""
         try:
             if provider not in self.available_providers:
@@ -132,7 +134,7 @@ class SyntheticDataGenerator:
         config = self.PROVIDER_CONFIG[provider]
         api_key = st.session_state.api_keys.get(provider, "")
 
-        if not api_key and provider != "Google": #Google API key is configured by configure()
+        if not api_key and provider != "Google":
             raise ValueError("API key required")
 
         try:
@@ -142,16 +144,18 @@ class SyntheticDataGenerator:
                 return {"headers": {"Authorization": f"Bearer {api_key}"}}
             elif provider == "Google":
                 if not st.session_state.google_configured:
-                   # Check if the API key is set as an environment variable
-                   if "GOOGLE_API_KEY" in os.environ:
-                       api_key = os.environ["GOOGLE_API_KEY"]
-                   else:
-                       # Use the API key from session state if available
-                       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.")
-                   configure(api_key=api_key) #Configure the Google API key. Only do once
-                   st.session_state.google_configured = True
+                    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.advanced_options["temperature"],
@@ -159,7 +163,25 @@ class SyntheticDataGenerator:
                     top_k=st.session_state.advanced_options["top_k"],
                     max_output_tokens=st.session_state.advanced_options["max_output_tokens"]
                 )
-                return GenerativeModel(model_name=model, generation_config=generation_config)  # Create the GenerativeModel with generation config
+                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)
             else:
                 return OpenAI(
                     base_url=config["base_url"],
@@ -170,7 +192,8 @@ class SyntheticDataGenerator:
             self._log_error(f"Client Init Failed: {str(e)}")
             return None
 
-    def _execute_generation(self, client, provider: str, model: str, prompt: str) -> Dict[str, Any]:
+    def _execute_generation(self, client, provider: str, model: str, prompt: Any) -> Dict[str, Any]: # Use Any for prompt type
+
         """Execute provider-specific generation with circuit breaker"""
         st.session_state.system_metrics["api_calls"] += 1
 
@@ -185,16 +208,25 @@ class SyntheticDataGenerator:
             return response.json()
         elif provider == "Google":
             try:
-                response = client.generate_content(prompt)
+                if isinstance(prompt, list): #Multimodal case
+
+                    response = client.generate_content(prompt)
+
+                else:
+
+                    response = client.generate_content(prompt)
+
+
                 content = response.text
 
-                if st.session_state.generation_format == "json": # Check requested format
+                if st.session_state.generation_format == "json":
                     try:
-                        return json.loads(content) # Attempt to parse as JSON
+                        return json.loads(content)
                     except json.JSONDecodeError:
-                        return {"content": content, "warning": "Could not parse response as valid JSON.  Returning raw text."} #Return raw content with warning
+                        return {"content": content,
+                                "warning": "Could not parse response as valid JSON.  Returning raw text."}
                 else:
-                    return {"content": content} # Return raw content
+                    return {"content": content}
 
             except Exception as e:
                 self._log_error(f"Google Generation Error: {str(e)}")
@@ -203,21 +235,22 @@ class SyntheticDataGenerator:
             completion = client.chat.completions.create(
                 model=model,
                 messages=[{"role": "user", "content": prompt}],
-                temperature=st.session_state.advanced_options["temperature"], #Use temp from session
+                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."}
+                return {"content": completion.choices[0].message.content,
+                        "warning": "Could not parse response as valid JSON.  Returning raw text."}
 
     def _failover_generation(self, prompt: str) -> Dict[str, Any]:
         """Enterprise failover to secondary providers"""
         for backup_provider in self.available_providers:
             if backup_provider != st.session_state.active_provider:
                 try:
-                    return self.generate(backup_provider, ..., prompt=prompt)  # Corrected: include prompt
+                    return self.generate(backup_provider, ..., prompt=prompt)
                 except Exception:
                     continue
         raise RuntimeError("All generation providers unavailable")
@@ -244,26 +277,24 @@ class SyntheticDataGenerator:
             return ""
 
     def _process_csv(self, file) -> str:
-      """Process CSV files and return as a string representation."""
-      try:
-          df = pd.read_csv(file)
-
-          # Attempt to infer a schema for the synthetic data generation
-          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 # Store the schema
-
-          return df.to_string()  # Convert DataFrame to string
-      except Exception as e:
-          self._log_error(f"CSV Processing Error: {str(e)}")
-          return ""
+        """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
+            return df.to_string()
+        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_api(self, url: str, method="GET", headers: Optional[Dict[str, str]] = None, data: Optional[Dict[str, Any]] = None) -> str:
+    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":
@@ -272,12 +303,12 @@ class SyntheticDataGenerator:
                 response = requests.post(url, headers=headers or {}, json=data, timeout=10)
             else:
                 raise ValueError("Unsupported HTTP method.")
-            response.raise_for_status()  # Raise HTTPError for bad responses (4xx or 5xx)
+            response.raise_for_status()
 
             try:
-                return json.dumps(response.json(), indent=2)  # Pretty print JSON if possible
+                return json.dumps(response.json(), indent=2)
             except json.JSONDecodeError:
-                return response.text  # Otherwise, return raw text
+                return response.text
         except requests.exceptions.RequestException as e:
             self._log_error(f"API Processing Error: {str(e)}")
             return ""
@@ -294,18 +325,16 @@ class SyntheticDataGenerator:
             self._log_error(f"Database Processing Error: {str(e)}")
             return ""
 
-    def _process_image(self, image_file) -> str:
-        """Processes image files for multimodal generation"""
+    def _process_image(self, image_file) -> list: #Returns a list
+        """Processes image files for multimodal generation (Google Gemini)"""
         try:
-            # For Google's Gemini, you need to prepare the image in a specific format
             image_data = image_file.read()
-            image_part = {"mime_type": image_file.type, "data": image_data}
-            st.session_state.image_part = image_part #Store image part
-            return "Image uploaded.  Include instructions for processing the image in your prompt." # Basic instruction to the LLM
+            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 ""
+            return []
 
     # --- Enterprise Features ---
     def _log_error(self, message: str) -> None:
@@ -340,21 +369,24 @@ class SyntheticDataGenerator:
                 return response.status_code == 200
             elif provider == "Google":
                 try:
-                    #Need to initialize before listing models
-                    if not st.session_state.google_configured:
-                        api_key = st.session_state.api_keys.get("Google", "")
-                        if not api_key:
-                           api_key = os.environ.get("GOOGLE_API_KEY") #Check env variables
+                    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
+                            return False #Cant test API if no API Key
 
                         configure(api_key=api_key) #Configure API Key
                         st.session_state.google_configured = True
+                        #st.write("configuring key")
 
-                    genai.GenerativeModel(model_name=self.PROVIDER_CONFIG["Google"]["models"][0]).generate_content("test") #Send a test query
-                    return True #Connected if made it this far
+                    genai.GenerativeModel(model_name=self.PROVIDER_CONFIG["Google"]["models"][0]).generate_content("test") #Test a generation
+                    return True
 
-                except Exception as e:
+                except Exception as e: #Catch any exceptions
                     print(e)
                     return False
 
@@ -395,30 +427,46 @@ def provider_config_ui(gen: SyntheticDataGenerator):
         )
         st.session_state.active_model = model
 
-        # Advanced Options (for providers that support it)
-        if provider == "Google" or provider == "OpenAI": #Only add if OpenAI
+        # 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.")
-
+            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.")
-
-        # Output format
-        st.session_state.generation_format = st.selectbox("Output Format", ["json", "text"], help="Choose the desired output format.")
+                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.")
 
         # 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)
+    input_method = st.selectbox("Input Method",
+                                ["Text", "PDF", "Web URL", "CSV", "Image",
+                                 "Structured Prompt (Advanced)"])  # Add Image input, Add Structured Prompt (Advanced)
     input_content = None
-    additional_instructions = "" #For structured prompt
+    additional_instructions = ""  # For structured prompt
 
     if input_method == "Text":
         input_content = st.text_area("Enter Text", height=200)
@@ -435,7 +483,7 @@ def input_ui():
             input_content = uploaded_file
             if "csv_schema" in st.session_state:
                 st.write("Inferred CSV Schema:")
-                st.write(st.session_state.csv_schema) #Display inferred schema
+                st.write(st.session_state.csv_schema)
 
     elif input_method == "Image":
         uploaded_file = st.file_uploader("Upload an Image file", type=["png", "jpg", "jpeg"])
@@ -445,9 +493,11 @@ def input_ui():
     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)
+        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, additional_instructions #Also return additional instructions
 
 def main():
     """Enterprise-grade user interface"""
@@ -467,10 +517,10 @@ def main():
 
     provider_config_ui(gen)
 
-    input_method, input_content, additional_instructions = input_ui() #Get additonal instructions
+    input_method, input_content, additional_instructions = input_ui()
 
     if st.button("Generate Data"):
-        if input_content or input_method == "Structured Prompt (Advanced)": #Allow generation with *just* structured prompt
+        if input_content or input_method == "Structured Prompt (Advanced)":
             processed_input = None
 
             if input_method == "Text":
@@ -482,23 +532,24 @@ def main():
             elif input_method == "CSV":
                 processed_input = gen._process_csv(input_content)
             elif input_method == "Image":
-                processed_input = gen._process_image(input_content)
+                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 #Combine instructions and constraints
-                #st.write("Combined Prompt:")
-                #st.write(processed_input) #Debug
+                processed_input = input_content + "\n" + additional_instructions
 
             if processed_input:
                 try:
-                    #Handle Google image case - requires a list of content.  Other providers just use the text
                     if st.session_state.active_provider == "Google" and input_method == "Image":
-                        prompt_parts = [processed_input, st.session_state.image_part] # Image part already stored
-                        result = gen.generate(st.session_state.active_provider, st.session_state.active_model, prompt_parts) # Process Google Images
+                        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) # Generic text case
+                        result = gen.generate(st.session_state.active_provider, st.session_state.active_model, processed_input)
 
                     st.subheader("Generated Output:")
-                    st.json(result)  # Display the JSON output
+                    st.json(result)
                 except Exception as e:
                     st.error(f"Error during generation: {e}")
             else:
@@ -506,7 +557,5 @@ def main():
         else:
             st.warning("Please provide input data.")
 
-    # Input management and generation UI components...
-
 if __name__ == "__main__":
     main()