app.py

import json
import requests
import streamlit as st
import pdfplumber
import pandas as pd
import sqlalchemy
from typing import Any, Dict, List

# Provider clients (ensure these are installed if you plan to use them)
try:
    from openai import OpenAI
except ImportError:
    OpenAI = None

try:
    import groq
except ImportError:
    groq = None

# Hugging Face Inference API endpoint
HF_API_URL = "https://api-inference.huggingface.co/models/"
DEFAULT_TEMPERATURE = 0.1
GROQ_MODEL = "mixtral-8x7b-32768"


class AdvancedSyntheticDataGenerator:
    """
    Advanced Synthetic Data Generator that supports multiple input types,
    customizable prompt templates, multiple LLM providers, and detailed debugging.
    """
    def __init__(self) -> None:
        self._setup_providers()
        self._setup_input_handlers()
        self._initialize_session_state()
        # Customizable prompt template with placeholders for data, instructions, and output format.
        self.custom_prompt_template = (
            "You are an expert synthetic data generator. "
            "Given the data below and following the instructions provided, generate high-quality, diverse synthetic data. "
            "Ensure the output adheres to the specified format.\n\n"
            "-------------------------\n"
            "Data:\n{data}\n\n"
            "Instructions:\n{instructions}\n\n"
            "Output Format: {format}\n"
            "-------------------------\n"
        )

    def _setup_providers(self) -> None:
        """Configure available LLM providers and their initialization routines."""
        self.providers: Dict[str, Dict[str, Any]] = {
            "Deepseek": {
                "client": lambda key: OpenAI(base_url="https://api.deepseek.com/v1", api_key=key) if OpenAI else None,
                "models": ["deepseek-chat"],
            },
            "OpenAI": {
                "client": lambda key: OpenAI(api_key=key) if OpenAI else None,
                "models": ["gpt-4-turbo", "gpt-3.5-turbo"],
            },
            "Groq": {
                "client": lambda key: groq.Groq(api_key=key) if groq else None,
                "models": [GROQ_MODEL],
            },
            "HuggingFace": {
                "client": lambda key: {"headers": {"Authorization": f"Bearer {key}"}},
                "models": ["gpt2", "llama-2"],
            },
        }

    def _setup_input_handlers(self) -> None:
        """Register handlers for different input data types."""
        self.input_handlers: Dict[str, Any] = {
            "text": self.handle_text,
            "pdf": self.handle_pdf,
            "csv": self.handle_csv,
            "api": self.handle_api,
            "db": self.handle_db,
        }

    def _initialize_session_state(self) -> None:
        """Initialize Streamlit session state with default configuration."""
        defaults = {
            "config": {
                "provider": "OpenAI",
                "model": "gpt-4-turbo",
                "temperature": DEFAULT_TEMPERATURE,
                "output_format": "plain_text",  # Options: plain_text, json, csv
            },
            "api_key": "",
            "inputs": [],         # List to store all input sources
            "instructions": "",   # Custom instructions for synthetic data generation
            "synthetic_data": "", # The generated output
            "error_logs": [],     # Logs for any errors during processing
        }
        for key, value in defaults.items():
            if key not in st.session_state:
                st.session_state[key] = value

    def log_error(self, message: str) -> None:
        """Log an error message both to session state and in the UI."""
        st.session_state.error_logs.append(message)
        st.error(message)

    # ===== Input Handlers =====
    def handle_text(self, text: str) -> Dict[str, Any]:
        return {"data": text, "source": "text"}

    def handle_pdf(self, file) -> Dict[str, Any]:
        try:
            with pdfplumber.open(file) as pdf:
                full_text = ""
                for page in pdf.pages:
                    page_text = page.extract_text() or ""
                    full_text += page_text + "\n"
                return {"data": full_text, "source": "pdf"}
        except Exception as e:
            self.log_error(f"PDF Processing Error: {e}")
            return {"data": "", "source": "pdf"}

    def handle_csv(self, file) -> Dict[str, Any]:
        try:
            df = pd.read_csv(file)
            # Convert the DataFrame to JSON for simplicity.
            return {"data": df.to_json(orient="records"), "source": "csv"}
        except Exception as e:
            self.log_error(f"CSV Processing Error: {e}")
            return {"data": "", "source": "csv"}

    def handle_api(self, config: Dict[str, str]) -> Dict[str, Any]:
        try:
            response = requests.get(config["url"], headers=config.get("headers", {}), timeout=10)
            response.raise_for_status()
            return {"data": json.dumps(response.json()), "source": "api"}
        except Exception as e:
            self.log_error(f"API Processing Error: {e}")
            return {"data": "", "source": "api"}

    def handle_db(self, config: Dict[str, str]) -> Dict[str, Any]:
        try:
            engine = sqlalchemy.create_engine(config["connection"])
            with engine.connect() as conn:
                result = conn.execute(sqlalchemy.text(config["query"]))
                rows = [dict(row) for row in result]
            return {"data": json.dumps(rows), "source": "db"}
        except Exception as e:
            self.log_error(f"Database Processing Error: {e}")
            return {"data": "", "source": "db"}

    def aggregate_inputs(self) -> str:
        """Combine all input sources into a single aggregated string."""
        aggregated_data = ""
        for item in st.session_state.inputs:
            aggregated_data += f"Source: {item.get('source', 'unknown')}\n"
            aggregated_data += item.get("data", "") + "\n\n"
        return aggregated_data.strip()

    def build_prompt(self) -> str:
        """
        Build the complete prompt using aggregated data, custom instructions,
        and the desired output format.
        """
        aggregated_data = self.aggregate_inputs()
        instructions = st.session_state.instructions or "Generate diverse, coherent synthetic data."
        output_format = st.session_state.config.get("output_format", "plain_text")
        prompt = self.custom_prompt_template.format(
            data=aggregated_data, instructions=instructions, format=output_format
        )
        st.write("### Built Prompt")
        st.write(prompt)
        return prompt

    def generate_synthetic_data(self) -> bool:
        """
        Generate synthetic data by sending the built prompt to the selected LLM provider.
        Returns True if generation succeeds.
        """
        api_key = st.session_state.api_key
        if not api_key:
            self.log_error("API key is missing!")
            return False

        provider_name = st.session_state.config["provider"]
        provider_cfg = self.providers.get(provider_name)
        if not provider_cfg:
            self.log_error(f"Provider {provider_name} is not configured.")
            return False

        client_initializer = provider_cfg["client"]
        client = client_initializer(api_key)
        model = st.session_state.config["model"]
        temperature = st.session_state.config["temperature"]
        prompt = self.build_prompt()

        st.info(f"Using **{provider_name}** with model **{model}** at temperature **{temperature:.2f}**")
        try:
            if provider_name == "HuggingFace":
                response = self._huggingface_inference(client, prompt, model)
            else:
                response = self._standard_inference(client, prompt, model, temperature)
            
            st.write("### Raw API Response")
            st.write(response)
            
            synthetic_data = self._parse_response(response, provider_name)
            st.write("### Parsed Synthetic Data")
            st.write(synthetic_data)
            
            st.session_state.synthetic_data = synthetic_data
            return True
        except Exception as e:
            self.log_error(f"Generation failed: {e}")
            return False

    def _standard_inference(self, client: Any, prompt: str, model: str, temperature: float) -> Any:
        """
        Inference for providers using an OpenAI-compatible API.
        """
        try:
            st.write("Sending prompt via standard inference...")
            result = client.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": prompt}],
                temperature=temperature,
            )
            st.write("Standard inference result received.")
            return result
        except Exception as e:
            self.log_error(f"Standard Inference Error: {e}")
            return None

    def _huggingface_inference(self, client: Dict[str, Any], prompt: str, model: str) -> Any:
        """
        Inference for the Hugging Face Inference API.
        """
        try:
            st.write("Sending prompt to HuggingFace API...")
            response = requests.post(
                HF_API_URL + model,
                headers=client["headers"],
                json={"inputs": prompt},
                timeout=30,
            )
            response.raise_for_status()
            st.write("HuggingFace API response received.")
            return response.json()
        except Exception as e:
            self.log_error(f"HuggingFace Inference Error: {e}")
            return None

    def _parse_response(self, response: Any, provider: str) -> str:
        """
        Parse the LLM response into a synthetic data string.
        """
        st.write("Parsing response for provider:", provider)
        try:
            if provider == "HuggingFace":
                if isinstance(response, list) and response and "generated_text" in response[0]:
                    return response[0]["generated_text"]
                else:
                    self.log_error("Unexpected HuggingFace response format.")
                    return ""
            else:
                # Expecting a structure similar to OpenAI's response.
                if response and hasattr(response, "choices") and response.choices:
                    return response.choices[0].message.content
                else:
                    self.log_error("Unexpected response format from provider.")
                    return ""
        except Exception as e:
            self.log_error(f"Response Parsing Error: {e}")
            return ""


# ===== Advanced UI Components =====

def advanced_config_ui(generator: AdvancedSyntheticDataGenerator):
    """Display advanced configuration options in the sidebar."""
    with st.sidebar:
        st.header("Advanced Configuration")
        provider = st.selectbox("Select Provider", list(generator.providers.keys()))
        st.session_state.config["provider"] = provider
        provider_cfg = generator.providers[provider]

        model = st.selectbox("Select Model", provider_cfg["models"])
        st.session_state.config["model"] = model

        temperature = st.slider("Temperature", 0.0, 1.0, DEFAULT_TEMPERATURE)
        st.session_state.config["temperature"] = temperature

        output_format = st.radio("Output Format", ["plain_text", "json", "csv"])
        st.session_state.config["output_format"] = output_format

        api_key = st.text_input(f"{provider} API Key", type="password")
        st.session_state.api_key = api_key

        instructions = st.text_area("Custom Instructions",
                                    "Generate diverse, coherent synthetic data based on the input sources.",
                                    height=100)
        st.session_state.instructions = instructions


def advanced_input_ui(generator: AdvancedSyntheticDataGenerator):
    """Display input data source options using tabs."""
    st.subheader("Add Input Data")
    tabs = st.tabs(["Text", "PDF", "CSV", "API", "Database"])
    
    with tabs[0]:
        text_input = st.text_area("Enter text input", height=150)
        if st.button("Add Text Input", key="text_input"):
            if text_input.strip():
                st.session_state.inputs.append(generator.handle_text(text_input))
                st.success("Text input added!")
            else:
                st.warning("Empty text input.")
    
    with tabs[1]:
        pdf_file = st.file_uploader("Upload PDF", type=["pdf"])
        if pdf_file is not None:
            st.session_state.inputs.append(generator.handle_pdf(pdf_file))
            st.success("PDF input added!")
    
    with tabs[2]:
        csv_file = st.file_uploader("Upload CSV", type=["csv"])
        if csv_file is not None:
            st.session_state.inputs.append(generator.handle_csv(csv_file))
            st.success("CSV input added!")
    
    with tabs[3]:
        api_url = st.text_input("API Endpoint URL")
        api_headers = st.text_area("API Headers (JSON format, optional)", height=100)
        if st.button("Add API Input", key="api_input"):
            headers = {}
            try:
                if api_headers:
                    headers = json.loads(api_headers)
            except Exception as e:
                generator.log_error(f"Invalid JSON for API Headers: {e}")
            st.session_state.inputs.append(generator.handle_api({"url": api_url, "headers": headers}))
            st.success("API input added!")
    
    with tabs[4]:
        db_conn = st.text_input("Database Connection String")
        db_query = st.text_area("Database Query", height=100)
        if st.button("Add Database Input", key="db_input"):
            st.session_state.inputs.append(generator.handle_db({"connection": db_conn, "query": db_query}))
            st.success("Database input added!")


def advanced_output_ui(generator: AdvancedSyntheticDataGenerator):
    """Display the generated synthetic data with output options."""
    st.subheader("Synthetic Data Output")
    if st.session_state.synthetic_data:
        output_format = st.session_state.config.get("output_format", "plain_text")
        if output_format == "json":
            try:
                json_output = json.loads(st.session_state.synthetic_data)
                st.json(json_output)
            except Exception:
                st.text_area("Output", st.session_state.synthetic_data, height=300)
        else:
            st.text_area("Output", st.session_state.synthetic_data, height=300)
        st.download_button("Download Output", st.session_state.synthetic_data,
                           file_name="synthetic_data.txt", mime="text/plain")
    else:
        st.info("No synthetic data generated yet.")


def advanced_logs_ui():
    """Display error logs and debug information in an expandable section."""
    with st.expander("Error Logs & Debug Info", expanded=False):
        if st.session_state.error_logs:
            for log in st.session_state.error_logs:
                st.write(log)
        else:
            st.write("No logs yet.")


# ===== Main Application =====

def main() -> None:
    st.set_page_config(page_title="Advanced Synthetic Data Generator", layout="wide")
    
    # Sidebar for advanced configuration
    generator = AdvancedSyntheticDataGenerator()
    advanced_config_ui(generator)
    
    st.title("Advanced Synthetic Data Generator")
    st.markdown(
        """
        Welcome! This application allows you to generate synthetic data from multiple input sources.
        Use the sections below to add inputs, generate data, view outputs, and review logs.
        """
    )

    # Input Data Section
    with st.container():
        st.header("1. Input Data Sources")
        advanced_input_ui(generator)
        if st.button("Clear All Inputs"):
            st.session_state.inputs = []
            st.success("All inputs have been cleared!")

    # Generation Section with a clearly visible button
    with st.container():
        st.header("2. Generate Synthetic Data")
        if st.button("Generate Synthetic Data", key="generate_button"):
            with st.spinner("Generating synthetic data..."):
                if generator.generate_synthetic_data():
                    st.success("Synthetic data generated successfully!")
                else:
                    st.error("Data generation failed. Check logs for details.")

    # Output Section
    with st.container():
        st.header("3. Synthetic Data Output")
        advanced_output_ui(generator)

    # Logs Section
    with st.container():
        st.header("4. Error Logs & Debug Information")
        advanced_logs_ui()


if __name__ == "__main__":
    main()