app.py

import os

import streamlit as st

import pandas as pd

import pickle

import base64

from io import BytesIO, StringIO

import sys

import operator

from typing import Literal, Sequence, TypedDict, Annotated, List, Dict, Tuple

import tempfile

import shutil

import plotly.io as pio

import io

import re

import json

import openai

# from fpdf import FPDF

import base64

from datetime import datetime

from reportlab.lib.pagesizes import letter

from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image

from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle

from reportlab.lib.units import inch

from PIL import Image as PILImage

 

# Import LangChain and LangGraph components

from langchain_core.messages import AIMessage, ToolMessage, HumanMessage, BaseMessage

from langchain_core.prompts import ChatPromptTemplate

from langchain_openai import ChatOpenAI

from langchain_experimental.utilities import PythonREPL

from langgraph.prebuilt import ToolInvocation, ToolExecutor

from langchain_core.tools import tool

from langgraph.prebuilt import InjectedState

from langgraph.graph import StateGraph, END

from reportlab.platypus import PageBreak

from PIL import Image as PILImage

 

# Initialize session state for AI provider settings

if 'ai_provider' not in st.session_state:

    st.session_state.ai_provider = "openai"

   

if 'api_key' not in st.session_state:

    st.session_state.api_key = ""

   

if 'selected_model' not in st.session_state:

    st.session_state.selected_model = "gpt-4"

   

# Define model options for each provider

OPENAI_MODELS = ["gpt-4", "gpt-4-turbo", "gpt-4-mini", "gpt-3.5-turbo"]

GROQ_MODELS = ["llama3.3-70b-versatile", "gemma2-9b-it", "llama-3-8b-8192"]

 

# Create temporary directory for file storage

if 'temp_dir' not in st.session_state:

    st.session_state.temp_dir = tempfile.mkdtemp()

    st.session_state.images_dir = os.path.join(st.session_state.temp_dir, "images/plotly_figures/pickle")

    os.makedirs(st.session_state.images_dir, exist_ok=True)

    print(f"Created temporary directory: {st.session_state.temp_dir}")

    print(f"Created images directory: {st.session_state.images_dir}")

 

# Define the system prompt

SYSTEM_PROMPT = """## Role

You are a professional data scientist helping a non-technical user understand, analyze, and visualize their data.

 

## Capabilities

1. **Execute python code** using the `complete_python_task` tool.

 

## Goals

1. Understand the user's objectives clearly.

2. Take the user on a data analysis journey, iterating to find the best way to visualize or analyse their data to solve their problems.

3. Investigate if the goal is achievable by running Python code via the `python_code` field.

4. Gain input from the user at every step to ensure the analysis is on the right track and to understand business nuances.

 

## Code Guidelines

- **ALL INPUT DATA IS LOADED ALREADY**, so use the provided variable names to access the data.

- **VARIABLES PERSIST BETWEEN RUNS**, so reuse previously defined variables if needed.

- **TO SEE CODE OUTPUT**, use `print()` statements. You won't be able to see outputs of `pd.head()`, `pd.describe()` etc. otherwise.

- **ONLY USE THE FOLLOWING LIBRARIES**:

  - `pandas`

  - `sklearn` (including all major ML models)

  - `plotly`

  - `numpy`

 

All these libraries are already imported for you.

 

## Machine Learning Guidelines

- For regression tasks:

  - Linear Regression: `LinearRegression`

  - Logistic Regression: `LogisticRegression`

  - Ridge Regression: `Ridge`

  - Lasso Regression: `Lasso`

  - Random Forest Regression: `RandomForestRegressor`

 

- For classification tasks:

  - Logistic Regression: `LogisticRegression`

  - Decision Trees: `DecisionTreeClassifier`

  - Random Forests: `RandomForestClassifier`

  - Support Vector Machines: `SVC`

  - K-Nearest Neighbors: `KNeighborsClassifier`

  - Naive Bayes: `GaussianNB`

 

- For clustering:

  - K-Means: `KMeans`

  - DBSCAN: `DBSCAN`

 

- For dimensionality reduction:

  - PCA: `PCA`

 

- Always preprocess data appropriately:

  - Scale numerical features with `StandardScaler` or `MinMaxScaler`

  - Encode categorical variables with `OneHotEncoder` when needed

  - Handle missing values with `SimpleImputer`

 

- Always split data into training and testing sets using `train_test_split`

- Evaluate models using appropriate metrics:

  - For regression: `mean_squared_error`, `mean_absolute_error`, `r2_score`

  - For classification: `accuracy_score`, `confusion_matrix`, `classification_report`

  - For clustering: `silhouette_score`

 

- Consider using `cross_val_score` for more robust evaluation

- Visualize ML results with plotly when possible

 

## Plotting Guidelines

- Always use the `plotly` library for plotting.

- Store all plotly figures inside a `plotly_figures` list, they will be saved automatically.

- Do not try and show the plots inline with `fig.show()`.

"""

 

# Define the State class

class AgentState(TypedDict):

    messages: Annotated[Sequence[BaseMessage], operator.add]

    input_data: Annotated[List[Dict], operator.add]

    intermediate_outputs: Annotated[List[dict], operator.add]

    current_variables: dict

    output_image_paths: Annotated[List[str], operator.add]

 

# Initialize session state variables

if 'in_memory_datasets' not in st.session_state:

    st.session_state.in_memory_datasets = {}

 

if 'persistent_vars' not in st.session_state:

    st.session_state.persistent_vars = {}

 

if 'dataset_metadata_list' not in st.session_state:

    st.session_state.dataset_metadata_list = []

 

if 'chat_history' not in st.session_state:

    st.session_state.chat_history = []

 

if 'dashboard_plots' not in st.session_state:

    st.session_state.dashboard_plots = [None, None, None, None]

 

if 'columns' not in st.session_state:

    st.session_state.columns = ["No columns available"]

 

if 'custom_plots_to_save' not in st.session_state:

    st.session_state.custom_plots_to_save = {}

 

# Set up the tools

repl = PythonREPL()

plotly_saving_code = """import pickle

 

import uuid

import os

for figure in plotly_figures:

    pickle_filename = f"{images_dir}/{uuid.uuid4()}.pickle"

    with open(pickle_filename, 'wb') as f:

        pickle.dump(figure, f)

"""

 

@tool

def complete_python_task(

    graph_state: Annotated[dict, InjectedState],

    thought: str,

    python_code: str

) -> Tuple[str, dict]:

    """Execute Python code for data analysis and visualization."""

 

    current_variables = graph_state.get("current_variables", {})

 

    # Load datasets from in-memory storage

    for input_dataset in graph_state.get("input_data", []):

        var_name = input_dataset.get("variable_name")

        if var_name and var_name not in current_variables and var_name in st.session_state.in_memory_datasets:

            print(f"Loading {var_name} from in-memory storage")

            current_variables[var_name] = st.session_state.in_memory_datasets[var_name]

    current_image_pickle_files = os.listdir(st.session_state.images_dir)

 

    try:

        # Capture stdout

        old_stdout = sys.stdout

        sys.stdout = StringIO()

      

        # Execute the code and capture the result

        exec_globals = globals().copy()

        exec_globals.update(st.session_state.persistent_vars)

        exec_globals.update(current_variables)

       

        # Add scikit-learn modules to execution environment

        import sklearn

        import numpy as np

       

        # Import scikit-learn components

        from sklearn.linear_model import LinearRegression, LogisticRegression, Ridge, Lasso # type: ignore

        from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor, GradientBoostingClassifier # type: ignore

        from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor

        from sklearn.svm import SVC, SVR

        from sklearn.naive_bayes import GaussianNB

        from sklearn.decomposition import PCA

        from sklearn.neighbors import KNeighborsClassifier, KNeighborsRegressor

        from sklearn.cluster import KMeans, DBSCAN

        from sklearn.preprocessing import StandardScaler, MinMaxScaler, OneHotEncoder

        from sklearn.model_selection import train_test_split, cross_val_score, GridSearchCV

        from sklearn.metrics import (

            accuracy_score, confusion_matrix, classification_report,

            mean_squared_error, r2_score, mean_absolute_error, silhouette_score

        )

        from sklearn.pipeline import Pipeline

        from sklearn.impute import SimpleImputer

       

        # Update execution globals with all ML components

        exec_globals.update({

            "plotly_figures": [],

            "images_dir": st.session_state.images_dir,

            "np": np,

            # Linear models

            "LinearRegression": LinearRegression,

            "LogisticRegression": LogisticRegression,

            "Ridge": Ridge,

            "Lasso": Lasso,

            # Tree-based models

            "DecisionTreeClassifier": DecisionTreeClassifier,

            "DecisionTreeRegressor": DecisionTreeRegressor,

            "RandomForestClassifier": RandomForestClassifier,

            "RandomForestRegressor": RandomForestRegressor,

            "GradientBoostingClassifier": GradientBoostingClassifier,

            # SVM models

            "SVC": SVC,

            "SVR": SVR,

            # Other models

            "GaussianNB": GaussianNB,

            "PCA": PCA,

            "KNeighborsClassifier": KNeighborsClassifier,

            "KNeighborsRegressor": KNeighborsRegressor,

            "KMeans": KMeans,

            "DBSCAN": DBSCAN,

            # Preprocessing

            "StandardScaler": StandardScaler,

            "MinMaxScaler": MinMaxScaler,

            "OneHotEncoder": OneHotEncoder,

            "SimpleImputer": SimpleImputer,

            # Model selection and evaluation

            "train_test_split": train_test_split,

            "cross_val_score": cross_val_score,

            "GridSearchCV": GridSearchCV,

            "accuracy_score": accuracy_score,

            "confusion_matrix": confusion_matrix,

            "classification_report": classification_report,

            "mean_squared_error": mean_squared_error,

            "r2_score": r2_score,

            "mean_absolute_error": mean_absolute_error,

            "silhouette_score": silhouette_score,

            # Pipeline

            "Pipeline": Pipeline

        })

       

        exec(python_code, exec_globals)

 

        st.session_state.persistent_vars.update({k: v for k, v in exec_globals.items() if k not in globals()})

 

        # Get the captured stdout

        output = sys.stdout.getvalue()

 

        # Restore stdout

        sys.stdout = old_stdout

 

        updated_state = {

            "intermediate_outputs": [{"thought": thought, "code": python_code, "output": output}],

            "current_variables": st.session_state.persistent_vars

        }

 

        if 'plotly_figures' in exec_globals and exec_globals['plotly_figures']:

            exec(plotly_saving_code, exec_globals)

          

            # Check if any images were created

            new_image_folder_contents = os.listdir(st.session_state.images_dir)

            new_image_files = [file for file in new_image_folder_contents if file not in current_image_pickle_files]

          

            if new_image_files:

                updated_state["output_image_paths"] = new_image_files

            st.session_state.persistent_vars["plotly_figures"] = []

        return output, updated_state

 

    except Exception as e:

        sys.stdout = old_stdout  # Restore stdout in case of error

        print(f"Error in complete_python_task: {str(e)}")

        return str(e), {"intermediate_outputs": [{"thought": thought, "code": python_code, "output": str(e)}]}

 

# Function to initialize the LLM based on selected provider and model

def initialize_llm():

    api_key = st.session_state.api_key

    model = st.session_state.selected_model

   

    if not api_key:

        return None

   

    try:

        if st.session_state.ai_provider == "openai":

            os.environ["OPENAI_API_KEY"] = api_key

            return ChatOpenAI(model=model, temperature=0)

        elif st.session_state.ai_provider == "groq":

            os.environ["GROQ_API_KEY"] = api_key

            # For Groq, set the base URL and use the model

            from langchain_groq import ChatGroq

            return ChatGroq(model=model, temperature=0)

    except Exception as e:

        print(f"Error initializing LLM: {str(e)}")

        return None

 

# Set up the tools

tools = [complete_python_task]

tool_executor = ToolExecutor(tools)

 

# Load the prompt template

chat_template = ChatPromptTemplate.from_messages([

    ("system", SYSTEM_PROMPT),

    ("placeholder", "{messages}"),

])

 

def create_data_summary(state: AgentState) -> str:

    summary = ""

    variables = []

  

    # Add sample data for each dataset

    for d in state.get("input_data", []):

        var_name = d.get("variable_name")

        if var_name:

          

            variables.append(var_name)

            summary += f"\n\nVariable: {var_name}\n"

            summary += f"Description: {d.get('data_description', 'No description')}\n"

 

            # Add sample data if available

            if var_name in st.session_state.in_memory_datasets:

                df = st.session_state.in_memory_datasets[var_name]

                summary += "\nSample Data (first 5 rows):\n"

                summary += df.head(5).to_string()

 

    if "current_variables" in state:

        remaining_variables = [v for v in state["current_variables"] if v not in variables and not v.startswith("_")]

      

        for v in remaining_variables:

          

            var_value = state["current_variables"].get(v)

 

            if isinstance(var_value, pd.DataFrame):

                summary += f"\n\nVariable: {v} (DataFrame with shape {var_value.shape})"

            else:

                summary += f"\n\nVariable: {v}"

    return summary

 

def route_to_tools(state: AgentState) -> Literal["tools", "__end__"]:

    """Determine if we should route to tools or end the chain"""

    if messages := state.get("messages", []):

        ai_message = messages[-1]

    else:

        raise ValueError(f"No messages found in input state to tool_edge: {state}")

 

    if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:

        return "tools"

  

    return "__end__"

 

def call_model(state: AgentState):

    """Call the LLM to get a response"""

    current_data_template = """The following data is available:\n{data_summary}"""

    current_data_message = HumanMessage(

        content=current_data_template.format(data_summary=create_data_summary(state))

    )

    messages = [current_data_message] + state["messages"]

   

    # Get the initialized LLM

    llm = initialize_llm()

    if llm is None:

        return {"messages": [AIMessage(content="Please configure a valid API key and model in the settings tab.")]}

   

    # Create the model with bound tools

    model = llm.bind_tools(tools)

    model = chat_template | model

   

    llm_outputs = model.invoke({"messages": messages})

    return {"messages": [llm_outputs], "intermediate_outputs": [current_data_message.content]}

 

def call_tools(state: AgentState):

    """Execute tools called by the LLM"""

    last_message = state["messages"][-1]

    tool_invocations = []

 

    if isinstance(last_message, AIMessage) and hasattr(last_message, 'tool_calls'):

        tool_invocations = [

            ToolInvocation(

                tool=tool_call["name"],

                tool_input={**tool_call["args"], "graph_state": state}

            ) for tool_call in last_message.tool_calls

        ]

    responses = tool_executor.batch(tool_invocations, return_exceptions=True)

 

    tool_messages = []

    state_updates = {}

 

    for tc, response in zip(last_message.tool_calls, responses):

        if isinstance(response, Exception):

            print(f"Exception in tool execution: {str(response)}")

            tool_messages.append(ToolMessage(               

                content=f"Error: {str(response)}",

                name=tc["name"],

                tool_call_id=tc["id"]

            ))

            continue

 

        message, updates = response

        tool_messages.append(ToolMessage(

            content=str(message),

            name=tc["name"],

            tool_call_id=tc["id"]

        ))

 

        # Merge updates instead of overwriting

        for key, value in updates.items():

            if key in state_updates:

                if isinstance(value, list) and isinstance(state_updates[key], list):

                    state_updates[key].extend(value)

                elif isinstance(value, dict) and isinstance(state_updates[key], dict):

                    state_updates[key].update(value)

                else:

                    state_updates[key] = value

            else:

                state_updates[key] = value

 

    if 'messages' not in state_updates:

        state_updates["messages"] = []

 

    state_updates["messages"] = tool_messages

    return state_updates

 

# Set up the graph

workflow = StateGraph(AgentState)

workflow.add_node("agent", call_model)

workflow.add_node("tools", call_tools)

workflow.add_conditional_edges(

    "agent",

    route_to_tools,

    {

        "tools": "tools",

        "__end__": END

    }

)

workflow.add_edge("tools", "agent")

workflow.set_entry_point("agent")

 

chain = workflow.compile()

 

def process_file_upload(files):

    """Process uploaded files and return dataframe previews and column names"""

    st.session_state.in_memory_datasets = {}  # Clear previous datasets

    st.session_state.dataset_metadata_list = []  # Clear previous metadata

    st.session_state.persistent_vars.clear()  # Clear persistent variables for new session

 

    if not files:

        return "No files uploaded.", [], ["No columns available"]

 

    results = []

    all_columns = []  # Track all columns from all datasets

 

    for file in files:

        try:

            # Use file object directly

            if file.name.endswith('.csv'):

                df = pd.read_csv(file)

            elif file.name.endswith(('.xls', '.xlsx')):

                df = pd.read_excel(file)

            else:

                results.append(f"Unsupported file format: {file.name}. Please upload CSV or Excel files.")

                continue

 

            var_name = file.name.split('.')[0].replace('-', '_').replace(' ', '_').lower()

            st.session_state.in_memory_datasets[var_name] = df

 

            # Collect all columns

            all_columns.extend(df.columns.tolist())

 

            # Create dataset metadata

            dataset_metadata = {

                "variable_name": var_name,

                "data_path": "in_memory",

                "data_description": f"Dataset containing {df.shape[0]} rows and {df.shape[1]} columns. Columns: {', '.join(df.columns.tolist())}",

                "original_filename": file.name

            }

 

            st.session_state.dataset_metadata_list.append(dataset_metadata)

 

            # Return preview of the dataset

            preview = f"### Dataset: {file.name}\nVariable name: `{var_name}`\n\n"

            preview += df.head(10).to_markdown()

            results.append(preview)

            print(f"Successfully processed {file.name}")

 

        except Exception as e:

            print(f"Error processing {file.name}: {str(e)}")

            results.append(f"Error processing {file.name}: {str(e)}")

 

    # Get unique columns

    unique_columns = []

    seen = set()

 

    for col in all_columns:

        if col not in seen:

            seen.add(col)

            unique_columns.append(col)

 

    if not unique_columns:

        unique_columns = ["No columns available"]

 

    print(f"Found {len(unique_columns)} unique columns across datasets")

    return "\n\n".join(results), st.session_state.dataset_metadata_list, unique_columns

 

def get_columns():

    """Directly gets columns from in-memory datasets"""

    all_columns = []

 

    for var_name, df in st.session_state.in_memory_datasets.items():

        if isinstance(df, pd.DataFrame):

            all_columns.extend(df.columns.tolist())

 

    # Remove duplicates while preserving order

    unique_columns = []

    seen = set()

  

    for col in all_columns:

        if col not in seen:

            seen.add(col)

            unique_columns.append(col)

 

    if not unique_columns:

        unique_columns = ["No columns available"]

 

    print(f"Populating dropdowns with {len(unique_columns)} columns")

    return unique_columns

 

# === FUNCTIONS ===

import openai

import pandas as pd

import json

import re

 

def standard_clean(df):

    df.columns = [re.sub(r'\W+', '_', col).strip().lower() for col in df.columns]

    df.drop_duplicates(inplace=True)

    df.dropna(axis=1, how='all', inplace=True)

    df.dropna(axis=0, how='all', inplace=True)

    for col in df.select_dtypes(include='object').columns:

        df[col] = df[col].astype(str).str.strip()

    return df

 

def query_openai(prompt):

    try:

        # Use the configured API key and model from session state

        api_key = st.session_state.api_key

        model = st.session_state.selected_model

       

        if st.session_state.ai_provider == "openai":

            client = openai.OpenAI(api_key=api_key)

            response = client.chat.completions.create(

                model=model,

                messages=[{"role": "user", "content": prompt}],

                temperature=0.7

            )

            return response.choices[0].message.content

        elif st.session_state.ai_provider == "groq":

            from groq import Groq

            client = Groq(api_key=api_key)

            response = client.chat.completions.create(

                model=model,

                messages=[{"role": "user", "content": prompt}],

                temperature=0.7

            )

            return response.choices[0].message.content

    except Exception as e:

        print(f"API Error: {e}")

        return "{}"

 

def llm_suggest_cleaning(df):

    sample = df.head(10).to_csv(index=False)

    prompt = f"""

You are a professional data wrangler. Below is a sample of a messy dataset.

 

Return a Python dictionary with the following keys:

 

1. rename_columns – fix unclear or inconsistent column names

2. convert_types – correct datatypes: int, float, str, or date

3. fill_missing – use 'mean', 'median', 'mode', or a constant like 'Unknown' or 0

4. value_map – map inconsistent values (e.g., yes/Yes/Y → Yes)

 

Do not drop any rows or columns. Your output must be a valid Python dict.

 

Example:

{{

  "rename_columns": {{"dob": "date_of_birth"}},

  "convert_types": {{"age": "int", "salary": "float", "signup_date": "date"}},

  "fill_missing": {{"gender": "mode", "salary": -1}},

  "value_map": {{

    "gender": {{"M": "Male", "F": "Female"}},

    "subscribed": {{"Y": "Yes", "N": "No"}}

  }}

}}

Apart from these mentioned steps, study the data and also do whatever things are good and needed for that particular dataset and do the cleaning.

Sample data:

{sample}

"""

    raw_response = query_openai(prompt)

    try:

        suggestions = eval(raw_response)

        return suggestions

    except:

        print("Could not parse suggestions.")

        return {

            "rename_columns": {},

            "convert_types": {},

            "fill_missing": {},

            "value_map": {}

        }

 

def apply_suggestions(df, suggestions):

    df.rename(columns=suggestions.get("rename_columns", {}), inplace=True)

 

    for col, dtype in suggestions.get("convert_types", {}).items():

        if col not in df.columns:

            continue

        try:

            if dtype == "int":

                df[col] = pd.to_numeric(df[col], errors='coerce').astype("Int64")

            elif dtype == "float":

                df[col] = pd.to_numeric(df[col], errors='coerce')

            elif dtype == "str":

                df[col] = df[col].astype(str)

            elif dtype == "date":

                df[col] = pd.to_datetime(df[col], errors='coerce')

        except:

            print(f"Failed to convert {col} to {dtype}")

 

    for col, method in suggestions.get("fill_missing", {}).items():

        if col not in df.columns:

            continue

        try:

            if method == "mean":

                df[col].fillna(df[col].mean(), inplace=True)

            elif method == "median":

                df[col].fillna(df[col].median(), inplace=True)

            elif method == "mode":

                df[col].fillna(df[col].mode().iloc[0], inplace=True)

            elif isinstance(method, str):

                df[col].fillna(method, inplace=True)

        except:

            print(f"Could not fill missing values for {col}")

 

    for col, mapping in suggestions.get("value_map", {}).items():

        if col in df.columns:

            try:

                df[col] = df[col].replace(mapping)

            except:

                print(f"Could not map values in {col}")

 

    return df

 

def capture_dashboard_screenshot():

    """Capture the entire dashboard as a single image"""

    try:

        # Create a figure that combines all dashboard plots

        import plotly.graph_objects as go

        from plotly.subplots import make_subplots

      

        # Create a 2x2 subplot

        fig = make_subplots(rows=2, cols=2,

                           subplot_titles=["Visualization 1", "Visualization 2",

                                          "Visualization 3", "Visualization 4"])

 

        # Add each plot from the dashboard to the combined figure

        for i, plot in enumerate(st.session_state.dashboard_plots):

            if plot is not None:

                row = (i // 2) + 1

                col = (i % 2) + 1

 

                # Extract traces from the original figure and add to our subplot

                for trace in plot.data:

                    fig.add_trace(trace, row=row, col=col)

              

                # Copy layout properties for each subplot

                for axis_type in ['xaxis', 'yaxis']:

                    axis_name = f"{axis_type}{i+1 if i > 0 else ''}"

                    subplot_name = f"{axis_type}{row}{col}"

 

                    # Copy axis properties if they exist

                    if hasattr(plot.layout, axis_name):

                        axis_props = getattr(plot.layout, axis_name)

                        fig.update_layout({subplot_name: axis_props})

      

        # Update layout for better appearance

        fig.update_layout(

            height=800,

            width=1000,

            title_text="Dashboard Overview",

            showlegend=False,

        )

 

        # Save to a temporary file

        dashboard_path = f"{st.session_state.temp_dir}/dashboard_combined.png"

        fig.write_image(dashboard_path, scale=2)  # Higher scale for better resolution

        return dashboard_path

 

    except Exception as e:

        import traceback

        print(f"Error capturing dashboard: {str(e)}")

        print(traceback.format_exc())

        return None

   

def generate_enhanced_pdf_report():

    """Generate an enhanced PDF report with proper handling of base64 image data"""

    try:

        # Create a buffer for the PDF

        buffer = io.BytesIO()

       

        # Create the PDF document

        doc = SimpleDocTemplate(buffer, pagesize=letter,

                               leftMargin=36, rightMargin=36,

                               topMargin=36, bottomMargin=36)

 

        # Create custom styles with better formatting

        styles = getSampleStyleSheet()

       

        # Add custom styles with improved formatting

        styles.add(ParagraphStyle(

            name='ReportTitle',

            parent=styles['Heading1'],

            fontSize=24,

            alignment=1,  # Centered

            spaceAfter=20,

            textColor='#2C3E50'  # Dark blue color

        ))

       

        styles.add(ParagraphStyle(

            name='SectionHeader',

            parent=styles['Heading2'],

            fontSize=16,

            spaceBefore=15,

            spaceAfter=10,

            textColor='#2C3E50',

            borderWidth=1,

            borderColor='#95A5A6',

            borderPadding=5,

            borderRadius=5

        ))

       

        styles.add(ParagraphStyle(

            name='SubHeader',

            parent=styles['Heading3'],

            fontSize=14,

            spaceBefore=10,

            spaceAfter=8,

            textColor='#34495E',

            fontWeight='bold'

        ))

        styles.add(ParagraphStyle(

                name='UserMessage',

                parent=styles['Normal'],

                fontSize=11,

                leftIndent=10,

                spaceBefore=8,

                spaceAfter=4

            ))

           

        styles.add(ParagraphStyle(

            name='AssistantMessage',

            parent=styles['Normal'],

            fontSize=11,

            leftIndent=10,

            spaceBefore=4,

            spaceAfter=12,

            textColor='#2980B9'

        ))

       

        styles.add(ParagraphStyle(

            name='Timestamp',

            parent=styles['Italic'],

            fontSize=10,

            textColor='#7F8C8D',

            alignment=2  # Right aligned

        ))

 

        # Create the document content

        elements = []

       

        # Add title

        elements.append(Paragraph('Data Analysis Report', styles['ReportTitle']))

       

        # Add timestamp

        elements.append(Paragraph(f'Generated on: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}',

                              styles['Timestamp']))

        elements.append(Spacer(1, 0.5*inch))

       

        # Add conversation history with better formatting

        elements.append(Paragraph('Analysis Conversation History', styles['SectionHeader']))

       

        if st.session_state.chat_history:

            for i, (user_msg, assistant_msg) in enumerate(st.session_state.chat_history):

                # Format user message with proper styling

                elements.append(Paragraph(f'<b>You:</b>', styles['SubHeader']))

                user_msg_formatted = user_msg.replace('\n', '<br/>')

                elements.append(Paragraph(user_msg_formatted, styles['UserMessage']))

               

                # Process assistant message to handle visualization

                # Look for markdown image syntax with base64 data

                base64_pattern = r'!\[Visualization\]\(data:image\/png;base64,([^\)]+)\)'

               

                # Check if the message contains visualizations

                if '### Visualizations' in assistant_msg or re.search(base64_pattern, assistant_msg):

                    # Split the message at the Visualizations header if it exists

                    if '### Visualizations' in assistant_msg:

                        parts = assistant_msg.split('### Visualizations', 1)

                        text_part = parts[0]

                        viz_part = "### Visualizations" + parts[1] if len(parts) > 1 else ""

                    else:

                        # If no header but still has visualization

                        match = re.search(base64_pattern, assistant_msg)

                        text_part = assistant_msg[:match.start()]

                        viz_part = assistant_msg[match.start():]

                   

                    # Format the text part

                    elements.append(Paragraph(f'<b>Assistant:</b>', styles['SubHeader']))

                    text_part = text_part.replace('\n', '<br/>')

                    elements.append(Paragraph(text_part, styles['AssistantMessage']))

                   

                    # Process visualizations

                    matches = re.findall(base64_pattern, viz_part)

                    for j, base64_data in enumerate(matches):

                        try:

                            # Decode the base64 image

                            image_data = base64.b64decode(base64_data)

                           

                            # Create a temporary file for the image

                            temp_img_path = f"{st.session_state.temp_dir}/chat_viz_{i}_{j}.png"

                           

                            with open(temp_img_path, 'wb') as f:

                                f.write(image_data)

                           

                            # Add the image to the PDF

                            elements.append(Paragraph(f'<b>Visualization:</b>', styles['SubHeader']))

                            elements.append(Spacer(1, 0.1*inch))

                            img = Image(temp_img_path, width=6*inch, height=4*inch)

                            elements.append(img)

                            elements.append(Spacer(1, 0.2*inch))

                        except Exception as e:

                            print(f"Error processing base64 image: {str(e)}")

                            elements.append(Paragraph(f"[Error displaying visualization: {str(e)}]",

                                                    styles['Normal']))

                else:

                    # No visualizations, just format the text

                    elements.append(Paragraph(f'<b>Assistant:</b>', styles['SubHeader']))

                    assistant_msg_formatted = assistant_msg.replace('\n', '<br/>')

                    if len(assistant_msg_formatted) > 1500:

                        assistant_msg_formatted = assistant_msg_formatted[:1500] + '...'

                    elements.append(Paragraph(assistant_msg_formatted, styles['AssistantMessage']))

               

                elements.append(Spacer(1, 0.2*inch))

        else:

            elements.append(Paragraph('No conversation history available.', styles['Normal']))

       

        # Force a page break before the dashboard

        elements.append(PageBreak())

       

        # Add dashboard section header

        elements.append(Paragraph('Dashboard Overview', styles['SectionHeader']))

        elements.append(Spacer(1, 0.2*inch))

        

        # Capture the dashboard as a single image

        dashboard_img_path = capture_dashboard_screenshot()

       

        if dashboard_img_path:

            # Calculate available width

            available_width = doc.width

           

            # Create PIL image to get dimensions

            pil_img = PILImage.open(dashboard_img_path)

            img_width, img_height = pil_img.size

           

            # Calculate scaling factor to fit within page width

            scale_factor = available_width / img_width

           

            # Calculate new height based on aspect ratio

            new_height = img_height * scale_factor

           

            # Add the image with scaled dimensions

            img = Image(dashboard_img_path, width=available_width, height=new_height)

            elements.append(img)

        else:

            # Fallback: Add individual plots if combined dashboard fails

            plot_count = 0

            for i, plot in enumerate(st.session_state.dashboard_plots):

                if plot is not None:

                    plot_count += 1

                   

                    # Convert plotly figure to image

                    img_bytes = io.BytesIO()

                    plot.write_image(img_bytes, format='png', width=500, height=300)

                    img_bytes.seek(0)

                   

                    # Create a temporary file for the image

                    temp_img_path = f"{st.session_state.temp_dir}/plot_{i}.png"

                   

                    with open(temp_img_path, 'wb') as f:

                        f.write(img_bytes.getvalue())

                    

                    # Add to PDF with appropriate caption and formatting

                    elements.append(Paragraph(f'Dashboard Visualization {i+1}', styles['SubHeader']))

                    elements.append(Spacer(1, 0.1*inch))

                    

                    # Add the image with proper scaling

                    img = Image(temp_img_path, width=6.5*inch, height=4*inch)

                    elements.append(img)

                    elements.append(Spacer(1, 0.3*inch))

            

            if plot_count == 0:

                elements.append(Paragraph('No visualizations have been added to the dashboard.',

                                        styles['Normal']))

       

        # Build the PDF with improved formatting

        doc.build(elements)

       

        # Get the value of the buffer

        pdf_value = buffer.getvalue()

        buffer.close()

       

        return pdf_value

       

    except Exception as e:

        import traceback

        print(f"Error generating enhanced PDF report: {str(e)}")

        print(traceback.format_exc())

        return None

 

def chat_with_workflow(message, history, dataset_info):

    """Send user query to the workflow and get response"""

   

    if not dataset_info:

        return "Please upload at least one dataset before asking questions."

 

    # Check if we have a valid API key and model

    if not st.session_state.api_key:

        return "Please set up your API key and model in the Settings tab before chatting."

 

    print(f"Chat with workflow called with {len(dataset_info)} datasets")

 

    try:

        # Extract chat history for context (last 3 exchanges)

        max_history = 3

        previous_messages = []

       

        if history:

            start_idx = max(0, len(history) - max_history)

            recent_history = history[start_idx:]

           

            for exchange in recent_history:

                if exchange[0]:  # User message

                    previous_messages.append(HumanMessage(content=exchange[0]))

                if exchange[1]:  # AI response

                    previous_messages.append(AIMessage(content=exchange[1]))

 

        # Initialize the workflow state

        state = AgentState(

            messages=previous_messages + [HumanMessage(content=message)],

            input_data=dataset_info,

            intermediate_outputs=[],

            current_variables=st.session_state.persistent_vars,

            output_image_paths=[]

        )

 

        # Execute the workflow

        print("Executing workflow...")

        result = chain.invoke(state)

        print("Workflow execution completed")

 

        # Extract messages from the result

        messages = result["messages"]

 

        # Format the response - only get the latest response

        response = ""

        if messages:

            latest_message = messages[-1]  # Get only the last message

            if hasattr(latest_message, "content"):

                content = latest_message.content

               

                # Clean up the response

                # Remove any instances where the user's message is repeated

                if message in content:

                    content = content.split(message)[-1].strip()

               

                # Remove any chat history markers

                content_lines = content.split('\n')

                filtered_lines = [line for line in content_lines

                                if not line.strip().startswith(("You:", "User:", "Human:", "Assistant:"))]

                content = '\n'.join(filtered_lines)

               

                response = content.strip() + "\n\n"

 

        # Handle visualizations

        if "output_image_paths" in result and result["output_image_paths"]:

            response += "### Visualizations\n\n"

            for img_path in result["output_image_paths"]:

                try:

                    full_path = os.path.join(st.session_state.images_dir, img_path)

                    with open(full_path, 'rb') as f:

                        fig = pickle.load(f)

 

                    # Convert plotly figure to image

                    img_bytes = BytesIO()

                    fig.update_layout(width=800, height=500)

                    pio.write_image(fig, img_bytes, format='png')

                    img_bytes.seek(0)

 

                    # Convert to base64 for markdown image

                    b64_img = base64.b64encode(img_bytes.read()).decode()

                    response += f"![Visualization](data:image/png;base64,{b64_img})\n\n"

                except Exception as e:                    

                    response += f"Error loading visualization: {str(e)}\n\n"

                   

        return response

 

    except Exception as e:

        import traceback

        print(f"Error in chat_with_workflow: {str(e)}")

        print(traceback.format_exc())

        return f"Error executing workflow: {str(e)}"

 

def auto_generate_dashboard(dataset_info):

    """Generate an automatic dashboard with four plots"""

 

    if not dataset_info:

        return "Please upload a dataset first.", [None, None, None, None]

 

    prompt = """

    You are a data visualization expert. Given a dataset, identify the top 4 most insightful plots using statistical reasoning or patterns (correlation, distribution, trends).

 

    Use plotly and store the plots in a list named plotly_figures.

 

    Include multivariate plots using color/size/facets when helpful.

    """

 

    state = AgentState(

        messages=[HumanMessage(content=prompt)],

        input_data=dataset_info,

        intermediate_outputs=[],

        current_variables=st.session_state.persistent_vars,

        output_image_paths=[]

    )

 

    result = chain.invoke(state)

    figures = []

 

    if "output_image_paths" in result:

        for img_path in result["output_image_paths"][:4]:

            try:

                full_path = os.path.join(st.session_state.images_dir, img_path)

                with open(full_path, 'rb') as f:

                    fig = pickle.load(f)

                    figures.append(fig)

            except Exception as e:

                print(f"Error loading figure: {e}")

 

    while len(figures) < 4:

        figures.append(None)

 

    st.session_state.dashboard_plots = figures

    return "Dashboard generated!", figures

 

def generate_custom_plots_with_llm(dataset_info, x_col, y_col, facet_col):

    """Generate custom plots based on user-selected columns"""

 

    if not dataset_info or not x_col or not y_col:

        return [None, None, None]

 

    prompt = f"""

    You are a data visualization expert.

 

    Create 3 insightful visualizations using Plotly based on:

 

    - X-axis: {x_col}

    - Y-axis: {y_col}

    - Facet (optional): {facet_col if facet_col != 'None' else 'None'}

 

    Try to find interesting relationships, trends, or clusters using appropriate chart types.

 

    Use `plotly_figures` list and avoid using fig.show().

    """

 

    state = AgentState(

        messages=[HumanMessage(content=prompt)],

        input_data=dataset_info,

        intermediate_outputs=[],

        current_variables=st.session_state.persistent_vars,

        output_image_paths=[]

    )

 

    result = chain.invoke(state)

    figures = []

 

    if "output_image_paths" in result:

        for img_path in result["output_image_paths"][:3]:

            try:

                full_path = os.path.join(st.session_state.images_dir, img_path)

                with open(full_path, 'rb') as f:

                    fig = pickle.load(f)

                    figures.append(fig)

            except Exception as e:

                print(f"Error loading figure: {e}")

              

    while len(figures) < 3:

        figures.append(None)

    return figures

 

def remove_plot(index):

    """Remove a plot from the dashboard"""

    if 0 <= index < len(st.session_state.dashboard_plots):

        st.session_state.dashboard_plots[index] = None

 

def respond(message):

    """Handle chat message response"""

    if not st.session_state.dataset_metadata_list:

        bot_message = "Please upload at least one dataset before asking questions."

    else:

        bot_message = chat_with_workflow(message, st.session_state.chat_history, st.session_state.dataset_metadata_list)

 

    st.session_state.chat_history.append((message, bot_message))

    st.rerun()

 

def save_plot_to_dashboard(plot_index):

    """Callback for the Add Plot button"""

    for i in range(len(st.session_state.dashboard_plots)):

        if st.session_state.dashboard_plots[i] is None:

            # Found an empty slot

            st.session_state.dashboard_plots[i] = st.session_state.custom_plots_to_save[plot_index]

            return

 

# Streamlit UI

st.set_page_config(page_title="QueryMind 🧠", layout="wide")

st.title("QueryMind 🧠 - Data Assistant")

st.markdown("Upload your datasets, ask questions, and generate visualizations to gain insights.")

 

# Create tabs

tab1, tab2, tab3, tab4, tab5, tab6 = st.tabs(["Upload Datasets", "Data Cleaning", "Chat with AI Assistant", "Auto Dashboard Generator", "Generate Report", "Settings"])

 

with tab1:

    st.header("Upload Datasets")

    uploaded_files = st.file_uploader("Upload CSV or Excel Files",

                                    accept_multiple_files=True,

                                    type=['csv', 'xlsx', 'xls'])

 

    if uploaded_files and st.button("Process Uploaded Files"):

        with st.spinner("Processing files..."):

            preview, metadata_list, columns = process_file_upload(uploaded_files)

            st.session_state.columns = columns

          

            # Display basic information about processed files

            st.success(f"✅ Successfully processed {len(uploaded_files)} file(s)")

          

            # Show detailed preview for each dataset

            st.subheader("Dataset Previews")

          

            for dataset_name, df in st.session_state.in_memory_datasets.items():

                with st.expander(f"Preview: {dataset_name}"):

                    # Display dataset info

                    st.write(f"**Rows:** {df.shape[0]} | **Columns:** {df.shape[1]}")

                  

                    # Display column information

                    col_info = pd.DataFrame({

                        'Column Name': df.columns,

                        'Data Type': df.dtypes.astype(str),

                        'Non-Null Count': df.count().values,

                        'Sample Values': [', '.join(df[col].dropna().astype(str).head(3).tolist()) for col in df.columns]

                    })

                  

                    # Show column information in a compact table

                    st.write("**Column Information:**")

                    st.dataframe(col_info, use_container_width=True)

                  

                    # Show actual data preview

                    st.write("**Data Preview (First 10 rows):**")

                    st.dataframe(df.head(10), use_container_width=True)

          

            # Provide hint for the next steps

            st.info("👆 Click on the dataset names above to see detailed previews. Then proceed to the Data Cleaning tab to clean your data or Chat with AI Assistant to analyze it.")

 

with tab2:

    st.header("Data Cleaning")

 

    if 'cleaning_done' not in st.session_state:

        st.session_state.cleaning_done = False

 

    if 'cleaned_datasets' not in st.session_state:

        st.session_state.cleaned_datasets = {}

 

    if 'cleaning_summaries' not in st.session_state:

        st.session_state.cleaning_summaries = {}

 

    if st.session_state.get("in_memory_datasets"):

        if not st.session_state.cleaning_done:

            if st.button("Run Data Cleaning"):

                with st.spinner("Running LLM-assisted cleaning..."):

                    for name, df in st.session_state.in_memory_datasets.items():

                        raw_df = df.copy()

                        df_std = standard_clean(raw_df.copy())

                        suggestions = llm_suggest_cleaning(df_std.copy())

                        df_clean = apply_suggestions(df_std.copy(), suggestions)

                        st.session_state.cleaned_datasets[name] = df_clean

                        st.session_state.cleaning_summaries[name] = suggestions

                    st.session_state.cleaning_done = True

                    st.rerun()

            else:

                st.info("Click Run Data Cleaning to clean your datasets using the LLM.")

        else:

            for name, df_clean in st.session_state.cleaned_datasets.items():

                raw_df = st.session_state.in_memory_datasets[name]

 

                st.subheader(f"Dataset: {name}")

                col1, col2 = st.columns(2)

 

                with col1:

                    st.markdown("Original Data (First 5 Rows)")

                    st.dataframe(raw_df.head())

 

                with col2:

                    st.markdown("Cleaned Data (First 5 Rows)")

                    st.dataframe(df_clean.head())

 

                st.markdown("Summary of Cleaning Actions")

                suggestions = st.session_state.cleaning_summaries[name]

                summary_text = ""

 

                if suggestions:

                    for key, value in suggestions.items():

                        summary_text += f"**{key}**: {json.dumps(value, indent=2)}\n\n"

                    st.markdown(summary_text)

 

                st.markdown("Refine the Cleaning (Natural Language Instructions)")

                user_input = st.text_input("Example: Convert 'dob' to datetime and fill missing with '2000-01-01'",

                                           key=f"user_input_{name}")

 

                if f'corrections_{name}' not in st.session_state:

                    st.session_state[f'corrections_{name}'] = []

 

                if st.button("Apply Correction", key=f'apply_correction_{name}'):

                    if user_input.strip():

                        correction_prompt = f"""

You are a data cleaning expert. Below is a previously cleaned dataset with these actions:

 

{summary_text}

 

The user now wants the following additional instruction:

\"{user_input.strip()}\"

 

Write only the Python code that modifies the pandas DataFrame `df` accordingly.

Do not include explanations or markdown.

"""

                        correction_code = query_openai(correction_prompt)

 

                        try:

                            df = st.session_state.cleaned_datasets[name].copy()

                            local_vars = {"df": df}

                            exec(correction_code, {}, local_vars)

                            df_updated = local_vars["df"]

 

                            st.session_state.cleaned_datasets[name] = df_updated

                            st.session_state[f'corrections_{name}'].append((user_input, correction_code))

                            st.success("Correction applied.")

                            st.rerun()

 

                        except Exception as e:

                            st.error(f"Failed to apply correction: {str(e)}")

 

                if st.session_state[f'corrections_{name}']:

                    st.markdown("Applied Corrections")

                    for i, (msg, code) in enumerate(st.session_state[f'corrections_{name}']):

                        st.markdown(f"**Instruction:** {msg}")

                        st.code(code, language='python')

 

            col1, col2 = st.columns([1, 2])

            with col1:

                if st.button("Reset Cleaning and Re-run"):

                    st.session_state.cleaning_done = False

                    st.rerun()

 

            with col2:

                if st.button("Finalize and Proceed to Visualizations"):

                    st.session_state.cleaning_finalized = True

                    st.rerun()

    else:

        st.info("Please upload and process datasets first.")

 

with tab3:

    st.header("Chat with AI Assistant")

   

    # Show API warning if not set

    if not st.session_state.api_key:

        st.warning("⚠️ Please set up your API key and model in the Settings tab before using the chat.")

       

    st.markdown("""

    ## Example Questions

    - "What analysis can you perform on this dataset?"

    - "Show me basic statistics for all columns"

    - "Create a correlation heatmap"

    - "Plot the distribution of a specific column"

    - "What is the relationship between two columns?"

    """)

 

    # Display chat history

    for exchange in st.session_state.chat_history:

        with st.chat_message("user"):

            st.write(exchange[0])

        with st.chat_message("assistant"):

            st.write(exchange[1])

 

    # Chat input

    if prompt := st.chat_input("Your question"):

        with st.spinner("Thinking..."):

            respond(prompt)

 

with tab4:

    st.header("Auto Dashboard Generator")

 

    # Dashboard controls

    dashboard_title = st.text_input("Dashboard Title", placeholder="Enter your dashboard title")


    col1, col2 = st.columns(2)

 

    with col1:

        if st.button("Generate Suggested Dashboard (Auto)"):

            if not st.session_state.api_key:

                st.warning("⚠️ Please set up your API key and model in the Settings tab first.")

            else:

                with st.spinner("Generating dashboard..."):

                    message, figures = auto_generate_dashboard(st.session_state.dataset_metadata_list)

                    st.success(message)

 

    with col2:

        if st.button("Refresh Column Options"):

            st.session_state.columns = get_columns()

            st.rerun()

 

    # Dashboard display

    st.subheader("Dashboard")

 

    # Row 1

    col1, col2 = st.columns(2)

 

    with col1:

        if st.session_state.dashboard_plots[0]:

            st.plotly_chart(st.session_state.dashboard_plots[0], use_container_width=True)

            if st.button("Remove Plot 1"):

                remove_plot(0)

                st.rerun()

 

    with col2:

        if st.session_state.dashboard_plots[1]:

            st.plotly_chart(st.session_state.dashboard_plots[1], use_container_width=True)

            if st.button("Remove Plot 2"):

                remove_plot(1)

                st.rerun()

 

    # Row 2

    col3, col4 = st.columns(2)

 

    with col3:

        if st.session_state.dashboard_plots[2]:

            st.plotly_chart(st.session_state.dashboard_plots[2], use_container_width=True)

            if st.button("Remove Plot 3"):

                remove_plot(2)

                st.rerun()

 

    with col4:

        if st.session_state.dashboard_plots[3]:

            st.plotly_chart(st.session_state.dashboard_plots[3], use_container_width=True)

            if st.button("Remove Plot 4"):

                remove_plot(3)

                st.rerun()

 

    # Custom plot generator

    st.subheader("Custom Plot Generator")

 

    # Column selection

    col1, col2, col3 = st.columns(3)

 

    with col1:

        x_axis = st.selectbox("X-axis Column", options=st.session_state.columns)

 

    with col2:

        y_axis = st.selectbox("Y-axis Column", options=st.session_state.columns)

 

    with col3:

        facet = st.selectbox("Facet (optional)", options=["None"] + st.session_state.columns)


    if st.button("Generate Custom Visualizations"):

        if not st.session_state.api_key:

            st.warning("⚠️ Please set up your API key and model in the Settings tab first.")

        else:

            with st.spinner("Generating custom visualizations..."):

                custom_plots = generate_custom_plots_with_llm(st.session_state.dataset_metadata_list, x_axis, y_axis, facet)

                # Store plots in session state

                for i, plot in enumerate(custom_plots):

                    if plot:

                        st.session_state.custom_plots_to_save[i] = plot

 

                # Display custom plots with add buttons

                for i, plot in enumerate(custom_plots):

                    if plot:

                        st.plotly_chart(plot, use_container_width=True)

                        st.button(

                            f"Add Plot {i+1} to Dashboard",

                            key=f"add_plot_{i}",

                            on_click=save_plot_to_dashboard,

                            args=(i,)

                        )

 

with tab5:

    st.header("Generate Analysis Report")

 

    st.markdown("""

    Generate a PDF report containing:

    - Dashboard visualizations

    - Chat conversation history

    """)

 

    report_title = st.text_input("Report Title (Optional)", "Data Analysis Report")

  

    if st.button("Generate PDF Report"):

        if not st.session_state.api_key:

            st.warning("⚠️ Please set up your API key and model in the Settings tab first.")

        else:

            with st.spinner("Generating report..."):

                pdf_data = generate_enhanced_pdf_report()

                if pdf_data:

                    # Create download button for PDF

                    b64_pdf = base64.b64encode(pdf_data).decode('utf-8')

                    # Create download link

                    pdf_download_link = f'<a href="data:application/pdf;base64,{b64_pdf}" download="data_analysis_report.pdf">Download PDF Report</a>'

                    st.markdown("### Your report is ready!")

                    st.markdown(pdf_download_link, unsafe_allow_html=True)

                    # Preview option (simplified)

                    with st.expander("Preview Report"):

                        st.warning("PDF preview is not available in Streamlit, please download the report to view it.")

                else:

                    st.error("Failed to generate the report. Please try again.")

 

with tab6:

    st.header("AI Provider Settings")

   

    # AI Provider selection

    provider = st.radio("Select AI Provider",

                         options=["OpenAI", "Groq"],

                         index=0 if st.session_state.ai_provider == "openai" else 1,

                         horizontal=True)

   

    # Update session state based on selection

    st.session_state.ai_provider = provider.lower()

   

    # API Key input

    api_key = st.text_input("Enter API Key",

                           value=st.session_state.api_key,

                           type="password",

                           help="Your API key for the selected provider")

   

    # Display different model options based on provider

    if st.session_state.ai_provider == "openai":

        model_options = OPENAI_MODELS

        model_help = "GPT-4 provides the best results but is slower. GPT-3.5-Turbo is faster but less capable."

    else:  # groq

        model_options = GROQ_MODELS

        model_help = "Llama 3.3 70B is most capable. Gemma 2 9B offers good balance. Llama 3 8B is fastest."

   

    # Model selection

    selected_model = st.selectbox("Select Model",

                                 options=model_options,

                                 index=model_options.index(st.session_state.selected_model) if st.session_state.selected_model in model_options else 0,

                                 help=model_help)

   

    # Save button

    if st.button("Save Settings"):

        st.session_state.api_key = api_key

        st.session_state.selected_model = selected_model

       

        # Test the API key and model

        try:

            # Initialize LLM using the provided settings

            test_llm = initialize_llm()

            if test_llm:

                st.success(f"✅ Successfully configured {provider} with model: {selected_model}")

            else:

                st.error("Failed to initialize the AI provider. Please check your API key and model selection.")

        except Exception as e:

            st.error(f"Error testing settings: {str(e)}")

   

    # Display current settings

    st.subheader("Current Settings")

    settings_info = f"""

    - **Provider**: {st.session_state.ai_provider.upper()}

    - **Model**: {st.session_state.selected_model}

    - **API Key**: {'✅ Set' if st.session_state.api_key else '❌ Not Set'}

    """

    st.markdown(settings_info)

   

    # Provider-specific information

    if st.session_state.ai_provider == "openai":

        st.info("""

        **OpenAI Models Information:**

        - **GPT-4**: Most powerful model, best for complex analysis and detailed explanations

        - **GPT-4-Turbo**: Faster than GPT-4 with similar capabilities

        - **GPT-4-Mini**: Economical option with good performance for standard tasks

        - **GPT-3.5-Turbo**: Fastest option, suitable for basic analysis and visualization

        """)

    else:

        st.info("""

        **Groq Models Information:**

        - **llama3.3-70b-versatile**: Most powerful model for comprehensive analysis

        - **gemma2-9b-it**: Good balance of speed and capabilities

        - **llama-3-8b-8192**: Fastest option for basic analysis tasks

        """)

   

    # Integration instructions

    with st.expander("How to get API Keys"):

        if st.session_state.ai_provider == "openai":

            st.markdown("""

            ### Getting an OpenAI API Key

           

            1. Go to [OpenAI's platform](https://platform.openai.com)

            2. Sign up or log in to your account

            3. Navigate to the API section

            4. Create a new API key

            5. Copy the key and paste it above

           

            Note: OpenAI API usage incurs charges based on tokens used.

            """)

        else:

            st.markdown("""

            ### Getting a Groq API Key

           

            1. Go to [Groq's website](https://console.groq.com/keys)

            2. Sign up or log in to your account

            3. Navigate to API Keys section

            4. Create a new API key

            5. Copy the key and paste it above

           

            Note: Check Groq's pricing page for current rates.

            """)

 

# Cleanup on app exit

def cleanup():

    try:

        shutil.rmtree(st.session_state.temp_dir)

        print(f"Cleaned up temporary directory: {st.session_state.temp_dir}")

    except Exception as e:

        print(f"Error cleaning up: {e}")

 

import atexit

atexit.register(cleanup)