Update app.py

app.py

@@ -7,12 +7,22 @@ import json
 from PIL import Image
 from io import BytesIO
 from openai import OpenAI
+import google.generativeai as genai # Added Google GenAI
 import groq
 import sqlalchemy
 from typing import Dict, Any
 
+# Constants for Default Values and API URLs
+HF_API_URL = "https://api-inference.huggingface.co/models/"
+DEFAULT_TEMPERATURE = 0.3
+
 class SyntheticDataGenerator:
+    """
+    A class to generate synthetic Q&A data from various input sources using different LLM providers.
+    """
+
     def __init__(self):
+        """Initializes the SyntheticDataGenerator with supported providers, input handlers, and session state."""
         self.providers = {
             "Deepseek": {
                 "client": lambda key: OpenAI(base_url="https://api.deepseek.com/v1", api_key=key),
@@ -29,9 +39,13 @@ class SyntheticDataGenerator:
             "HuggingFace": {
                 "client": lambda key: {"headers": {"Authorization": f"Bearer {key}"}},
                 "models": ["gpt2", "llama-2"]
-            }
+            },
+             "Google": {
+                "client": lambda key: self._configure_google_genai(key), # Using a custom configure function
+                "models": ["gemini-2.0-pro"]  # Add supported Gemini models.  Consider adding "gemini-1.5-pro" when released.
+            },
         }
-        
+
         self.input_handlers = {
             "pdf": self.handle_pdf,
             "text": self.handle_text,
@@ -39,10 +53,20 @@ class SyntheticDataGenerator:
             "api": self.handle_api,
             "db": self.handle_db
         }
-        
+
         self.init_session()
 
+    def _configure_google_genai(self, api_key: str):
+        """Configures the Google Generative AI client."""
+        try:
+            genai.configure(api_key=api_key)
+            return genai.GenerativeModel # return the model class, not an instantiation
+        except Exception as e:
+            st.error(f"Error configuring Google GenAI: {e}")
+            return None # Important: Handle the case where configuration fails
+
     def init_session(self):
+        """Initializes the Streamlit session state with default values."""
         session_defaults = {
             'inputs': [],
             'qa_data': [],
@@ -54,34 +78,42 @@ class SyntheticDataGenerator:
             'config': {
                 'provider': "Deepseek",
                 'model': "deepseek-chat",
-                'temperature': 0.3
+                'temperature': DEFAULT_TEMPERATURE
             }
         }
-        
+
         for key, val in session_defaults.items():
             if key not in st.session_state:
                 st.session_state[key] = val
 
     # Input Processors
     def handle_pdf(self, file):
-        try:
+       """Extracts text and images from a PDF file."""
+       try:
             with pdfplumber.open(file) as pdf:
-                return [{
-                    "text": page.extract_text() or "",
-                    "images": self.process_images(page),
-                    "meta": {"type": "pdf", "page": i+1}
-                } for i, page in enumerate(pdf.pages)]
-        except Exception as e:
-            self.log_error(f"PDF Error: {str(e)}")
-            return []
+                extracted_data = []
+                for i, page in enumerate(pdf.pages):
+                    page_text = page.extract_text() or ""
+                    page_images = self.process_images(page)
+                    extracted_data.append({
+                        "text": page_text,
+                        "images": page_images,
+                        "meta": {"type": "pdf", "page": i + 1}
+                    })
+                return extracted_data
+       except Exception as e:
+           self.log_error(f"PDF Error: {str(e)}")
+           return []
 
     def handle_text(self, text):
+        """Handles manual text input."""
         return [{
             "text": text,
             "meta": {"type": "domain", "source": "manual"}
         }]
 
     def handle_csv(self, file):
+        """Reads a CSV file and prepares data for Q&A generation."""
         try:
             df = pd.read_csv(file)
             return [{
@@ -93,17 +125,21 @@ class SyntheticDataGenerator:
             return []
 
     def handle_api(self, config):
+        """Fetches data from an API endpoint."""
         try:
             response = requests.get(config['url'], headers=config['headers'])
+            response.raise_for_status()  # Raise HTTPError for bad responses (4xx or 5xx)
             return [{
                 "text": json.dumps(response.json()),
                 "meta": {"type": "api", "endpoint": config['url']}
             }]
-        except Exception as e:
+        except requests.exceptions.RequestException as e:
             self.log_error(f"API Error: {str(e)}")
             return []
 
+
     def handle_db(self, config):
+        """Connects to a database and executes a query."""
         try:
             engine = sqlalchemy.create_engine(config['connection'])
             with engine.connect() as conn:
@@ -117,6 +153,7 @@ class SyntheticDataGenerator:
             return []
 
     def process_images(self, page):
+        """Extracts and processes images from a PDF page."""
         images = []
         for img in page.images:
             try:
@@ -134,130 +171,237 @@ class SyntheticDataGenerator:
 
     # Core Generation Engine
     def generate(self, api_key: str) -> bool:
+        """
+        Generates Q&A pairs using the selected LLM provider.
+
+        Args:
+            api_key (str): The API key for the selected LLM provider.
+
+        Returns:
+            bool: True if generation was successful, False otherwise.
+        """
         try:
             provider_cfg = self.providers[st.session_state.config['provider']]
-            client = provider_cfg["client"](api_key)
-            
+            client_initializer = provider_cfg["client"] #Get the client init function.
+
+            # Check that the key is not an empty string
+            if not api_key:
+                st.error("API Key cannot be empty.")
+                return False
+
+            # Initialize the client
+            if st.session_state.config['provider'] == "Google":
+                client = client_initializer(api_key) # Client is the class
+                if not client:
+                    return False  # Google config failed
+            else:
+                client = client_initializer(api_key)
+
             for i, input_data in enumerate(st.session_state.inputs):
                 st.session_state.processing['progress'] = (i+1)/len(st.session_state.inputs)
-                
+
                 if st.session_state.config['provider'] == "HuggingFace":
                     response = self._huggingface_inference(client, input_data)
+                elif st.session_state.config['provider'] == "Google":
+                   response = self._google_inference(client, input_data)
                 else:
                     response = self._standard_inference(client, input_data)
-                
+
                 if response:
-                    st.session_state.qa_data.extend(self._parse_response(response))
-            
+                    # Check if the parsing function needs access to the provider
+                    st.session_state.qa_data.extend(self._parse_response(response, st.session_state.config['provider']))
+
             return True
         except Exception as e:
             self.log_error(f"Generation Error: {str(e)}")
             return False
 
     def _standard_inference(self, client, input_data):
-        return client.chat.completions.create(
-            model=st.session_state.config['model'],
-            messages=[{
-                "role": "user",
-                "content": self._build_prompt(input_data)
-            }],
-            temperature=st.session_state.config['temperature'],
-            response_format={"type": "json_object"}
-        )
+         """Performs inference using standard OpenAI-compatible API."""
+         try:
+            return client.chat.completions.create(
+                model=st.session_state.config['model'],
+                messages=[{
+                    "role": "user",
+                    "content": self._build_prompt(input_data)
+                }],
+                temperature=st.session_state.config['temperature'],
+                response_format={"type": "json_object"} #Request json
+            )
+         except Exception as e:
+             self.log_error(f"OpenAI Inference Error: {e}")
+             return None
 
     def _huggingface_inference(self, client, input_data):
-        API_URL = "https://api-inference.huggingface.co/models/"
-        response = requests.post(
-            API_URL + st.session_state.config['model'],
-            headers=client["headers"],
-            json={"inputs": self._build_prompt(input_data)}
-        )
-        return response.json()
+        """Performs inference using Hugging Face Inference API."""
+        try:
+            response = requests.post(
+                HF_API_URL + st.session_state.config['model'],
+                headers=client["headers"],
+                json={"inputs": self._build_prompt(input_data)}
+            )
+            response.raise_for_status() #Check for HTTP errors
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            self.log_error(f"Hugging Face Inference Error: {e}")
+            return None
+
+    def _google_inference(self, client, input_data):
+        """Performs inference using Google Generative AI API."""
+        try:
+
+            model = client(st.session_state.config['model'])  # Instantiate the model with the selected model name
+            response = model.generate_content(
+                self._build_prompt(input_data),
+                generation_config = genai.types.GenerationConfig(temperature=st.session_state.config['temperature'])
+
+            )
+            return response
+        except Exception as e:
+            self.log_error(f"Google GenAI Inference Error: {e}")
+            return None
 
     def _build_prompt(self, input_data):
-        base = "Generate 3 Q&A pairs from this financial content:\n"
+        """Builds the prompt for the LLM based on the input data type."""
+        base = "Generate 3 Q&A pairs from this financial content, formatted as a JSON list of dictionaries with 'question' and 'answer' keys:\n"
         if input_data['meta']['type'] == 'csv':
             return base + "Structured data:\n" + input_data['text']
         elif input_data['meta']['type'] == 'api':
             return base + "API response:\n" + input_data['text']
         return base + input_data['text']
 
-    def _parse_response(self, response):
+    def _parse_response(self, response, provider):
+        """Parses the response from the LLM into a list of Q&A pairs."""
         try:
-            if st.session_state.config['provider'] == "HuggingFace":
+            if provider == "HuggingFace":
                 return response[0]['generated_text']
-            return json.loads(response.choices[0].message.content).get("qa_pairs", [])
+            elif provider == "Google":
+                # Expecting a text response from Gemini
+                try:
+                    json_string = response.text.strip()  # Removes surrounding whitespace that can cause errors
+                    qa_pairs = json.loads(json_string).get("qa_pairs", []) # Extract the qa_pairs
+
+                    # Validate the structure of qa_pairs
+                    if not isinstance(qa_pairs, list):
+                        raise ValueError("Expected a list of QA pairs.")
+
+                    for pair in qa_pairs:
+                        if not isinstance(pair, dict) or "question" not in pair or "answer" not in pair:
+                            raise ValueError("Each item in the list must be a dictionary with 'question' and 'answer' keys.")
+                    return qa_pairs  # Return the extracted and validated list
+                except (json.JSONDecodeError, ValueError) as e:
+                     self.log_error(f"Google JSON Parse Error: {e}.  Raw Response: {response.text}")
+                     return [] # Return empty in case of parsing failure
+            else:
+                # Assuming JSON response from other providers (OpenAI, Deepseek, Groq)
+                json_output = json.loads(response.choices[0].message.content) # load the JSON data
+                return json_output.get("qa_pairs", []) # Return the qa_pairs
         except Exception as e:
-            self.log_error(f"Parse Error: {str(e)}")
+            self.log_error(f"Parse Error: {e}. Raw Response: {response}")
             return []
 
     def log_error(self, message):
+        """Logs an error message to the Streamlit session state and displays it in the UI."""
         st.session_state.processing['errors'].append(message)
         st.error(message)
 
 # Streamlit UI Components
 def input_sidebar(gen: SyntheticDataGenerator):
+    """
+    Creates the input sidebar in the Streamlit UI.
+
+    Args:
+        gen (SyntheticDataGenerator): The SyntheticDataGenerator instance.
+
+    Returns:
+        str: The API key entered by the user.
+    """
     with st.sidebar:
         st.header("⚙️ Configuration")
-        
+
         # AI Provider Settings
         provider = st.selectbox("Provider", list(gen.providers.keys()))
         provider_cfg = gen.providers[provider]
-        
+
         api_key = st.text_input(f"{provider} API Key", type="password")
+        st.session_state['api_key'] = api_key  #Store API Key
+
         model = st.selectbox("Model", provider_cfg["models"])
-        temp = st.slider("Temperature", 0.0, 1.0, 0.3)
-        
+        temp = st.slider("Temperature", 0.0, 1.0, DEFAULT_TEMPERATURE)
+
         # Update session config
         st.session_state.config.update({
             "provider": provider,
             "model": model,
             "temperature": temp
         })
-        
+
         # Input Source Selection
         st.header("🔗 Data Sources")
         input_type = st.selectbox("Input Type", list(gen.input_handlers.keys()))
-        
+
         if input_type == "text":
             domain_input = st.text_area("Domain Knowledge", height=150)
             if st.button("Add Domain Input"):
-                gen.input_handlers["text"](domain_input)
-        
+                st.session_state.inputs.append(gen.input_handlers["text"](domain_input)[0])
+
         elif input_type == "csv":
             csv_file = st.file_uploader("Upload CSV", type=["csv"])
             if csv_file:
-                gen.input_handlers["csv"](csv_file)
-        
+                 st.session_state.inputs.extend(gen.input_handlers["csv"](csv_file))
+
         elif input_type == "api":
             api_url = st.text_input("API Endpoint")
-            if st.button("Connect API"):
-                gen.input_handlers["api"]({"url": api_url})
-        
+            api_headers = st.text_area("API Headers (JSON format, optional)", height=50)
+            headers = {}
+            try:
+                if api_headers:
+                    headers = json.loads(api_headers)
+            except json.JSONDecodeError:
+                st.error("Invalid JSON format for API headers.")
+            if st.button("Add API Input"):
+                st.session_state.inputs.extend(gen.input_handlers["api"]({"url": api_url, "headers": headers}))
+
+        elif input_type == "db":
+            db_connection = st.text_input("Database Connection String")
+            db_query = st.text_area("Database Query")
+            db_table = st.text_input("Table Name (optional)")
+            if st.button("Add DB Input"):
+                 st.session_state.inputs.extend(gen.input_handlers["db"]({"connection": db_connection, "query": db_query, "table": db_table}))
+
         return api_key
 
 def main_display(gen: SyntheticDataGenerator):
+    """
+    Creates the main display area in the Streamlit UI.
+
+    Args:
+        gen (SyntheticDataGenerator): The SyntheticDataGenerator instance.
+    """
     st.title("🚀 Enterprise Synthetic Data Factory")
-    
+
     # Input Processing
     col1, col2 = st.columns([3, 1])
     with col1:
         pdf_file = st.file_uploader("Upload Document", type=["pdf"])
         if pdf_file:
-            gen.input_handlers["pdf"](pdf_file)
-    
+           st.session_state.inputs.extend(gen.input_handlers["pdf"](pdf_file))
+
     # Generation Controls
     with col2:
         if st.button("Start Generation"):
             with st.status("Processing..."):
-                gen.generate(st.session_state.get('api_key'))
-    
+                if not st.session_state.get('api_key'):
+                     st.error("Please provide an API Key.")
+                else:
+                    gen.generate(st.session_state.get('api_key'))
+
     # Results Display
     if st.session_state.qa_data:
         st.header("Generated Data")
         df = pd.DataFrame(st.session_state.qa_data)
         st.dataframe(df)
-        
+
         # Export Options
         st.download_button(
             "Export CSV",
@@ -266,6 +410,7 @@ def main_display(gen: SyntheticDataGenerator):
         )
 
 def main():
+    """Main function to run the Streamlit application."""
     gen = SyntheticDataGenerator()
     api_key = input_sidebar(gen)
     main_display(gen)