Update app.py

app.py

@@ -1,4 +1,3 @@
-import json
 import requests
 import streamlit as st
 import pdfplumber
@@ -6,6 +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
 
 # Provider clients with import guards
 try:
@@ -18,6 +18,16 @@ try:
 except ImportError:
     groq = None
 
+try:
+    import google.generativeai as genai
+    from google.generativeai import GenerativeModel, configure
+except ImportError:
+    GenerativeModel = None
+    configure = None
+    genai = None #Also set this to none
+
+import json  # Ensure json is explicitly imported for enhanced use
+
 class SyntheticDataGenerator:
     """World's Most Advanced Synthetic Data Generation System"""
     
@@ -41,6 +51,10 @@ class SyntheticDataGenerator:
             "base_url": "https://api-inference.huggingface.co/models/",
             "models": ["gpt2", "llama-2-13b-chat"],
             "requires_library": None
+        },
+        "Google": {
+            "models": ["gemini-1.5-flash-latest", "gemini-1.5-pro-latest"],  # Include Gemini 2.0 Flash
+            "requires_library": "google.generativeai"
         }
     }
 
@@ -61,7 +75,15 @@ class SyntheticDataGenerator:
                 "tokens_used": 0,
                 "error_count": 0
             },
-            "debug_mode": False
+            "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
+            },
+            "generation_format": "json" # Default output format (json or text)
         }
         for key, val in defaults.items():
             if key not in st.session_state:
@@ -71,7 +93,7 @@ class SyntheticDataGenerator:
         """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"].title()):
+            if config["requires_library"] and not globals().get(config["requires_library"].split('.')[0].title()):
                 continue  # Skip providers with missing dependencies
             self.available_providers.append(provider)
 
@@ -83,7 +105,8 @@ class SyntheticDataGenerator:
             "csv": self._process_csv,
             "api": self._process_api,
             "database": self._process_database,
-            "web": self._process_web
+            "web": self._process_web,
+            "image": self._process_image #Add Image
         }
 
     # --- Core Generation Engine ---
@@ -108,8 +131,8 @@ class SyntheticDataGenerator:
         """Secure client initialization with connection pooling"""
         config = self.PROVIDER_CONFIG[provider]
         api_key = st.session_state.api_keys.get(provider, "")
-        
-        if not api_key:
+
+        if not api_key and provider != "Google": #Google API key is configured by configure()
             raise ValueError("API key required")
 
         try:
@@ -117,6 +140,26 @@ class SyntheticDataGenerator:
                 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:
+                   # 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
+
+                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"]
+                )
+                return GenerativeModel(model_name=model, generation_config=generation_config)  # Create the GenerativeModel with generation config
             else:
                 return OpenAI(
                     base_url=config["base_url"],
@@ -130,7 +173,7 @@ class SyntheticDataGenerator:
     def _execute_generation(self, client, provider: str, model: str, prompt: str) -> 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,
@@ -140,22 +183,41 @@ class SyntheticDataGenerator:
             )
             response.raise_for_status()
             return response.json()
+        elif provider == "Google":
+            try:
+                response = client.generate_content(prompt)
+                content = response.text
+
+                if st.session_state.generation_format == "json": # Check requested format
+                    try:
+                        return json.loads(content) # Attempt to parse as JSON
+                    except json.JSONDecodeError:
+                        return {"content": content, "warning": "Could not parse response as valid JSON.  Returning raw text."} #Return raw content with warning
+                else:
+                    return {"content": content} # Return raw content
+
+            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=0.1,
-                max_tokens=2000
+                temperature=st.session_state.advanced_options["temperature"], #Use temp from session
+                max_tokens=st.session_state.advanced_options["max_output_tokens"]
             )
             st.session_state.system_metrics["tokens_used"] += completion.usage.total_tokens
-            return json.loads(completion.choices[0].message.content)
+            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."}
 
     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, ...)
+                    return self.generate(backup_provider, ..., prompt=prompt)  # Corrected: include prompt
                 except Exception:
                     continue
         raise RuntimeError("All generation providers unavailable")
@@ -181,14 +243,76 @@ class SyntheticDataGenerator:
             self._log_error(f"Web Extraction Error: {str(e)}")
             return ""
 
-    # Additional processors follow similar patterns...
+    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 ""
+
+    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:
+        """Generic API endpoint processor with configurable methods and headers."""
+        try:
+            if method.upper() == "GET":
+                response = requests.get(url, headers=headers or {}, timeout=10)
+            elif method.upper() == "POST":
+                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)
+
+            try:
+                return json.dumps(response.json(), indent=2)  # Pretty print JSON if possible
+            except json.JSONDecodeError:
+                return response.text  # Otherwise, return raw text
+        except requests.exceptions.RequestException as e:
+            self._log_error(f"API Processing Error: {str(e)}")
+            return ""
+
+    def _process_database(self, connection_string: str, query: str) -> str:
+        """Database query processor using SQLAlchemy."""
+        try:
+            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()
+        except Exception as e:
+            self._log_error(f"Database Processing Error: {str(e)}")
+            return ""
+
+    def _process_image(self, image_file) -> str:
+        """Processes image files for multimodal generation"""
+        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
+
+        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"""
         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:
             st.error(f"[DEBUG] {message}")
 
@@ -214,6 +338,26 @@ class SyntheticDataGenerator:
                     timeout=5
                 )
                 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 api_key:
+                           return False
+
+                        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") #Send a test query
+                    return True #Connected if made it this far
+
+                except Exception as e:
+                    print(e)
+                    return False
+
             else:
                 client.models.list()
                 return True
@@ -225,14 +369,15 @@ def provider_config_ui(gen: SyntheticDataGenerator):
     """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"
         )
-        
+        st.session_state.active_provider = provider
+
         # API key management
         api_key = st.text_input(
             f"{provider} API Key",
@@ -241,19 +386,69 @@ def provider_config_ui(gen: SyntheticDataGenerator):
             help=f"Obtain API key from {provider} portal"
         )
         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"
         )
-        
+        st.session_state.active_model = model
+
+        # Advanced Options (for providers that support it)
+        if provider == "Google" or provider == "OpenAI": #Only add if 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.")
+
+        # Output format
+        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_content = None
+    additional_instructions = "" #For structured prompt
+
+    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
+            if "csv_schema" in st.session_state:
+                st.write("Inferred CSV Schema:")
+                st.write(st.session_state.csv_schema) #Display inferred schema
+
+    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 #Also return additional instructions
+
 def main():
     """Enterprise-grade user interface"""
     st.set_page_config(
@@ -261,18 +456,57 @@ def main():
         page_icon="🏭",
         layout="wide"
     )
-    
+
     gen = SyntheticDataGenerator()
-    
+
     st.title("🏭 Synthetic Data Factory Pro")
     st.markdown("""
-    **World's Most Advanced Synthetic Data Generation Platform**  
+    **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() #Get additonal instructions
+
+    if st.button("Generate Data"):
+        if input_content or input_method == "Structured Prompt (Advanced)": #Allow generation with *just* structured prompt
+            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)
+            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
+
+            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
+                    else:
+                        result = gen.generate(st.session_state.active_provider, st.session_state.active_model, processed_input) # Generic text case
+
+                    st.subheader("Generated Output:")
+                    st.json(result)  # Display the JSON output
+                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("Please provide input data.")
+
     # Input management and generation UI components...
-    
+
 if __name__ == "__main__":
     main()