From 47abbdf46d8db0dc8fa63998a1c6f4754b6cab1e Mon Sep 17 00:00:00 2001 From: Madhu Date: Thu, 2 Jan 2025 04:39:58 +0530 Subject: [PATCH 1/7] Added AI Data visualization folder with togetherAI --- .../ai_data_visualisation_agent/README.md | 0 .../ai_data_visualisation_agent.py | 220 ++++++++++++++++++ .../requirements.txt | 3 + 3 files changed, 223 insertions(+) create mode 100644 ai_agent_tutorials/ai_data_visualisation_agent/README.md create mode 100644 ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py create mode 100644 ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/README.md b/ai_agent_tutorials/ai_data_visualisation_agent/README.md new file mode 100644 index 0000000..e69de29 diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py b/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py new file mode 100644 index 0000000..7e70ec9 --- /dev/null +++ b/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py @@ -0,0 +1,220 @@ +import streamlit as st +import pandas as pd +import tempfile +import os +import re +from together import Together +import csv +import uuid +from dotenv import load_dotenv +from e2b_code_interpreter import Sandbox +from typing import Optional, Union, List + +# Load environment variables +load_dotenv() + +# Function to preprocess and save the uploaded file to a temporary file +def preprocess_and_save(file): + try: + # Read the uploaded file into a DataFrame + if file.name.endswith('.csv'): + df = pd.read_csv(file, encoding='utf-8', na_values=['NA', 'N/A', 'missing']) + elif file.name.endswith('.xlsx'): + df = pd.read_excel(file, na_values=['NA', 'N/A', 'missing']) + else: + st.error("Unsupported file format. Please upload a CSV or Excel file.") + return None, None, None + + # Ensure string columns are properly quoted + for col in df.select_dtypes(include=['object']): + df[col] = df[col].astype(str).replace({r'"': '""'}, regex=True) + + # Parse dates and numeric columns + for col in df.columns: + if 'date' in col.lower(): + df[col] = pd.to_datetime(df[col], errors='coerce') + elif df[col].dtype == 'object': + try: + df[col] = pd.to_numeric(df[col]) + except (ValueError, TypeError): + # Keep as is if conversion fails + pass + + # Create a temporary file to save the preprocessed data + with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as temp_file: + temp_path = temp_file.name + # Save the DataFrame to the temporary CSV file + df.to_csv(temp_path, index=False, quoting=csv.QUOTE_ALL) + + return temp_path, df.columns.tolist(), df # Return the DataFrame as well + except Exception as e: + st.error(f"Error processing file: {e}") + return None, None, None + +# Function to execute Python code in E2B sandbox +def code_interpret(code: str) -> str: + """ + Execute Python code in E2B sandbox. + + Args: + code: Python code to execute + + Returns: + String containing stdout from code execution + """ + print("Running code in E2B sandbox...") + + sbx = Sandbox(api_key=st.session_state.e2b_api_key) + + try: + execution = sbx.run_code("code") + # Convert list output to string if needed + stdout = execution.logs.stdout + if isinstance(stdout, list): + return '\n'.join(map(str, stdout)) + return stdout if stdout else "" + except Exception as e: + return f"Error executing code: {str(e)}" + +# Function to communicate with LLM +def chat_with_llm(user_message, file_path, columns): + print(f"\n{'='*50}\nUser message: {user_message}\n{'='*50}") + + # Update the system prompt with the file path, columns, and plot path + system_prompt = SYSTEM_PROMPT.format( + file_path=file_path, + columns=columns, + ) + + # Add a hint to include a plot if the user asks for visualization + if "plot" in user_message.lower(): + system_prompt += " Include a plot in your response and output the base64 string of the plot image." + + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_message}, + ] + + # Use the Together API key from session state + response = client.chat.completions.create( + model="meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo", + messages=messages, + ) + + response_message = response.choices[0].message.content + print("LLM Response:", response_message) # Debug: Print the LLM's response + python_code = match_code_blocks(response_message) + print("Extracted Python Code:", python_code) # Debug: Print the extracted code + + if python_code: + # Modify the code to handle the 'approx_cost(for two people)' column correctly + python_code = python_code.replace( + "df['approx_cost(for two people)'] = df['approx_cost(for two people)'].str.replace(',', '')", + "df['approx_cost(for two people)'] = df['approx_cost(for two people)'].astype(str).str.replace(',', '')" + ) + stdout = code_interpret(python_code) + return response_message, stdout, None + else: + print(f"Failed to match any Python code in model's response {response_message}") + return response_message, None, None + +# Set up Streamlit app +st.title("AI Data Scientist") + +# Sidebar for API keys and file upload +st.sidebar.header("API Keys") +together_api_key = st.sidebar.text_input("Together AI API Key", type="password") +e2b_api_key = st.sidebar.text_input("E2B API Key", type="password") + +# Store API keys in session state +if 'together_api_key' not in st.session_state: + st.session_state.together_api_key = None +if 'e2b_api_key' not in st.session_state: + st.session_state.e2b_api_key = None + +uploaded_file = st.sidebar.file_uploader("Upload CSV or Excel File", type=['csv', 'xlsx']) + +# System prompt (dynamic based on the uploaded file) +SYSTEM_PROMPT = """ +You are a Python data scientist. You have access to a CSV file located at '{file_path}'. +The dataset has the following columns: {columns}. +You can read this file into a DataFrame using `df = pd.read_csv('{file_path}')` and perform data analysis tasks based on user queries. +Make sure to handle missing values and data type inconsistencies. When generating plots, +use matplotlib or seaborn and output the plot as a base64 string. +Always respond with the Python code to answer the user's query, and include visualizations only if explicitly requested. +""" + +# Function to match Python code blocks +pattern = re.compile(r"```python\n(.*?)\n```", re.DOTALL) + +def match_code_blocks(llm_response): + match = pattern.search(llm_response) + if match: + code = match.group(1) + # Remove comments and extra text + code = "\n".join([line for line in code.split("\n") if not line.strip().startswith("#")]) + return code + return "" + +# Function to extract base64 image from stdout +def extract_base64_image(stdout: Optional[Union[str, List[str]]]) -> Optional[str]: + """ + Extract base64 image from stdout content. + + Args: + stdout: String or list of strings containing output + + Returns: + Base64 encoded image string if found, None otherwise + """ + if stdout is None: + return None + + # Convert list to string if needed + if isinstance(stdout, list): + stdout = '\n'.join(map(str, stdout)) + elif not isinstance(stdout, str): + stdout = str(stdout) + + # Look for base64 image data in the stdout + image_pattern = re.compile(r'base64_image:\s*(.*?)\n', re.DOTALL) + match = image_pattern.search(stdout) + if match: + return match.group(1) + return None + +# Main app logic +if 'together_api_key' in st.session_state and 'e2b_api_key' in st.session_state and uploaded_file: + # Preprocess and save the uploaded file to a temporary file + temp_path, columns, df = preprocess_and_save(uploaded_file) + if temp_path: + # Initialize Together AI client using the API key from session state + client = Together(api_key=st.session_state.together_api_key) + + # User query input + user_query = st.text_input("Ask a query about the data:") + if st.button("Submit Query"): + # Chat with LLM + response_message, stdout, stderr = chat_with_llm(user_query, temp_path, columns) + # Display AI's response + st.write("AI's Response:") + st.write(response_message) + + # Display any printed output + if stdout: + st.write("Code Output:") + st.write(stdout) + else: + st.write("No output produced by the code.") + + # Extract base64 image from stdout + base64_image = extract_base64_image(stdout) + if base64_image: + # Display the image + st.image(base64_image, use_container_width=True) + else: + st.write("No plot generated.") + else: + st.error("Failed to preprocess and save the data.") +else: + st.warning("Please provide API keys and upload a file.") \ No newline at end of file diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt b/ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt new file mode 100644 index 0000000..9740fa5 --- /dev/null +++ b/ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt @@ -0,0 +1,3 @@ +e2b-code-interpreter==1.0.3 +togetherai==1.3.10 +streamlit==1.41.1 \ No newline at end of file From d3cfeaab2e9144bdaf712a0856c7fe3d7ca0a634 Mon Sep 17 00:00:00 2001 From: Madhu Date: Sat, 4 Jan 2025 23:47:56 +0530 Subject: [PATCH 2/7] new issues --- .../ai_data_visualisation_agent/.gitignore | 1 + .../ai_data_visualisation_agent.py | 189 ++++++++-------- .../requirements.txt | 7 +- .../ai_data_visualisation_agent/test.py | 202 ++++++++++++++++++ 4 files changed, 296 insertions(+), 103 deletions(-) create mode 100644 ai_agent_tutorials/ai_data_visualisation_agent/.gitignore create mode 100644 ai_agent_tutorials/ai_data_visualisation_agent/test.py diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/.gitignore b/ai_agent_tutorials/ai_data_visualisation_agent/.gitignore new file mode 100644 index 0000000..2eea525 --- /dev/null +++ b/ai_agent_tutorials/ai_data_visualisation_agent/.gitignore @@ -0,0 +1 @@ +.env \ No newline at end of file diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py b/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py index 7e70ec9..adf1b53 100644 --- a/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py +++ b/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py @@ -1,14 +1,14 @@ import streamlit as st import pandas as pd import tempfile -import os import re from together import Together import csv -import uuid from dotenv import load_dotenv -from e2b_code_interpreter import Sandbox -from typing import Optional, Union, List +import base64 +import matplotlib.pyplot as plt +import io +import seaborn as sns # Load environment variables load_dotenv() @@ -25,6 +25,10 @@ def preprocess_and_save(file): st.error("Unsupported file format. Please upload a CSV or Excel file.") return None, None, None + # Log the data types of columns before preprocessing + st.write("Data types before preprocessing:") + st.write(df.dtypes) + # Ensure string columns are properly quoted for col in df.select_dtypes(include=['object']): df[col] = df[col].astype(str).replace({r'"': '""'}, regex=True) @@ -35,11 +39,24 @@ def preprocess_and_save(file): df[col] = pd.to_datetime(df[col], errors='coerce') elif df[col].dtype == 'object': try: - df[col] = pd.to_numeric(df[col]) + # Handle columns with values like "4.1/5" + if df[col].str.contains('/').any(): + # Split the values and take the first part (e.g., "4.1/5" -> 4.1) + df[col] = df[col].str.split('/').str[0] + # Convert to numeric, coerce errors to NaN + df[col] = pd.to_numeric(df[col], errors='coerce') except (ValueError, TypeError): # Keep as is if conversion fails + st.warning(f"Could not convert column '{col}' to numeric. Keeping as string.") pass + # Drop rows with all NaN values + df.dropna(how='all', inplace=True) + + # Log the data types of columns after preprocessing + st.write("Data types after preprocessing:") + st.write(df.dtypes) + # Create a temporary file to save the preprocessed data with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as temp_file: temp_path = temp_file.name @@ -51,33 +68,38 @@ def preprocess_and_save(file): st.error(f"Error processing file: {e}") return None, None, None -# Function to execute Python code in E2B sandbox -def code_interpret(code: str) -> str: - """ - Execute Python code in E2B sandbox. - - Args: - code: Python code to execute - - Returns: - String containing stdout from code execution - """ - print("Running code in E2B sandbox...") - - sbx = Sandbox(api_key=st.session_state.e2b_api_key) - +# Function to execute Python code and generate plots +def execute_code(code: str, df): try: - execution = sbx.run_code("code") - # Convert list output to string if needed - stdout = execution.logs.stdout - if isinstance(stdout, list): - return '\n'.join(map(str, stdout)) - return stdout if stdout else "" + # Define locals with necessary imports and the DataFrame + local_env = { + 'pd': pd, + 'df': df, + 'plt': plt, + 'sns': sns # if seaborn is needed + } + # Execute the code in the local environment + exec(code, globals(), local_env) + + # Check if a plot was generated + if 'plt' in local_env: + # Save the plot to a BytesIO object + buf = io.BytesIO() + plt.savefig(buf, format='png') + plt.close() + buf.seek(0) + # Encode the plot as base64 + base64_image = base64.b64encode(buf.read()).decode('utf-8') + return base64_image + else: + st.warning("No plot generated. Ensure the data being plotted is numeric.") + return None except Exception as e: - return f"Error executing code: {str(e)}" + st.error(f"Error executing code: {e}") + return None -# Function to communicate with LLM -def chat_with_llm(user_message, file_path, columns): +# Function to communicate with Together AI +def chat_with_llm(user_message, file_path, columns, df): print(f"\n{'='*50}\nUser message: {user_message}\n{'='*50}") # Update the system prompt with the file path, columns, and plot path @@ -107,16 +129,12 @@ def chat_with_llm(user_message, file_path, columns): print("Extracted Python Code:", python_code) # Debug: Print the extracted code if python_code: - # Modify the code to handle the 'approx_cost(for two people)' column correctly - python_code = python_code.replace( - "df['approx_cost(for two people)'] = df['approx_cost(for two people)'].str.replace(',', '')", - "df['approx_cost(for two people)'] = df['approx_cost(for two people)'].astype(str).str.replace(',', '')" - ) - stdout = code_interpret(python_code) - return response_message, stdout, None + # Execute the code and generate the plot + base64_image = execute_code(python_code, df) + return response_message, base64_image else: print(f"Failed to match any Python code in model's response {response_message}") - return response_message, None, None + return response_message, None # Set up Streamlit app st.title("AI Data Scientist") @@ -124,13 +142,10 @@ st.title("AI Data Scientist") # Sidebar for API keys and file upload st.sidebar.header("API Keys") together_api_key = st.sidebar.text_input("Together AI API Key", type="password") -e2b_api_key = st.sidebar.text_input("E2B API Key", type="password") -# Store API keys in session state +# Store API key in session state if 'together_api_key' not in st.session_state: st.session_state.together_api_key = None -if 'e2b_api_key' not in st.session_state: - st.session_state.e2b_api_key = None uploaded_file = st.sidebar.file_uploader("Upload CSV or Excel File", type=['csv', 'xlsx']) @@ -141,6 +156,7 @@ The dataset has the following columns: {columns}. You can read this file into a DataFrame using `df = pd.read_csv('{file_path}')` and perform data analysis tasks based on user queries. Make sure to handle missing values and data type inconsistencies. When generating plots, use matplotlib or seaborn and output the plot as a base64 string. +Always check if the data being plotted is numeric. If the data is not numeric, preprocess it to convert it to numeric values. Always respond with the Python code to answer the user's query, and include visualizations only if explicitly requested. """ @@ -156,65 +172,38 @@ def match_code_blocks(llm_response): return code return "" -# Function to extract base64 image from stdout -def extract_base64_image(stdout: Optional[Union[str, List[str]]]) -> Optional[str]: - """ - Extract base64 image from stdout content. - - Args: - stdout: String or list of strings containing output - - Returns: - Base64 encoded image string if found, None otherwise - """ - if stdout is None: - return None - - # Convert list to string if needed - if isinstance(stdout, list): - stdout = '\n'.join(map(str, stdout)) - elif not isinstance(stdout, str): - stdout = str(stdout) - - # Look for base64 image data in the stdout - image_pattern = re.compile(r'base64_image:\s*(.*?)\n', re.DOTALL) - match = image_pattern.search(stdout) - if match: - return match.group(1) - return None - # Main app logic -if 'together_api_key' in st.session_state and 'e2b_api_key' in st.session_state and uploaded_file: - # Preprocess and save the uploaded file to a temporary file - temp_path, columns, df = preprocess_and_save(uploaded_file) - if temp_path: - # Initialize Together AI client using the API key from session state - client = Together(api_key=st.session_state.together_api_key) - - # User query input - user_query = st.text_input("Ask a query about the data:") - if st.button("Submit Query"): - # Chat with LLM - response_message, stdout, stderr = chat_with_llm(user_query, temp_path, columns) - # Display AI's response - st.write("AI's Response:") - st.write(response_message) - - # Display any printed output - if stdout: - st.write("Code Output:") - st.write(stdout) - else: - st.write("No output produced by the code.") - - # Extract base64 image from stdout - base64_image = extract_base64_image(stdout) - if base64_image: - # Display the image - st.image(base64_image, use_container_width=True) - else: - st.write("No plot generated.") +if uploaded_file: + if not together_api_key: + st.warning("Please provide the Together AI API key.") else: - st.error("Failed to preprocess and save the data.") + # Update session state with API key + st.session_state.together_api_key = together_api_key + + # Initialize Together AI client only after confirming API key exists + try: + client = Together(api_key=together_api_key) + + # Preprocess and save the uploaded file + temp_path, columns, df = preprocess_and_save(uploaded_file) + if temp_path: + # Rest of your code for user query handling + user_query = st.text_input("Ask a query about the data:") + if st.button("Submit Query"): + response_message, base64_image = chat_with_llm(user_query, temp_path, columns, df) + + # Display AI's response + st.write("AI's Response:") + st.write(response_message) + + # Display the plot if generated + if base64_image: + st.image(base64.b64decode(base64_image), use_container_width=True) + else: + st.write("No plot generated.") + else: + st.error("Failed to preprocess and save the data.") + except Exception as e: + st.error(f"Error initializing Together AI client: {str(e)}") else: - st.warning("Please provide API keys and upload a file.") \ No newline at end of file + st.warning("Please upload a file.") \ No newline at end of file diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt b/ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt index 9740fa5..99a78ca 100644 --- a/ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt +++ b/ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt @@ -1,3 +1,4 @@ -e2b-code-interpreter==1.0.3 -togetherai==1.3.10 -streamlit==1.41.1 \ No newline at end of file +python-dotenv +together +e2b +streamlit \ No newline at end of file diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/test.py b/ai_agent_tutorials/ai_data_visualisation_agent/test.py new file mode 100644 index 0000000..dca9a6c --- /dev/null +++ b/ai_agent_tutorials/ai_data_visualisation_agent/test.py @@ -0,0 +1,202 @@ +import streamlit as st +import os +from dotenv import load_dotenv +import json +import re +from together import Together +from e2b_code_interpreter import Sandbox +from typing import Optional, List, Any +import tempfile + +# Load environment variables +load_dotenv() + +# Get API keys from environment variables +TOGETHER_API_KEY = os.getenv("TOGETHER_API_KEY") +E2B_API_KEY = os.getenv("E2B_API_KEY") + +# Define the Together AI model to use +MODEL_NAME = "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo" + +# System prompt for the LLM +SYSTEM_PROMPT = """You are a highly skilled Python data scientist. Your task is to analyze datasets and generate Python code to solve data-related problems. Follow these guidelines: + +1. **Data Preprocessing**: + - Always check for missing or invalid values in the dataset. + - Handle missing values by either removing rows/columns or imputing them appropriately. + - Convert columns to the correct data types (e.g., numeric, datetime). + - Filter out rows with invalid or inconsistent data. + +2. **Data Analysis**: + - Perform exploratory data analysis (EDA) to understand the dataset. + - Use statistical methods to analyze relationships between variables. + - If the task involves machine learning (e.g., linear regression), ensure the data is properly prepared (e.g., feature scaling, train-test split). + +3. **Visualization**: + - Use libraries like `matplotlib` or `seaborn` for creating visualizations. + - Ensure plots are clear, labeled, and informative (e.g., include titles, axis labels, legends). + - Save plots as images (e.g., PNG) and return them as base64-encoded strings. + +4. **Code Quality**: + - Write clean, modular, and well-commented Python code. + - Handle potential errors gracefully (e.g., invalid data, missing columns). + - Include necessary imports (e.g., `pandas`, `numpy`, `matplotlib`, `seaborn`). + +5. **Output**: + - Always return the Python code to solve the task. + - If the task involves visualization, include the code to generate and save the plot.""" + +# Function to execute code in the E2B Sandbox +def code_interpret(e2b_code_interpreter, code): + print("Running code interpreter...") + exec = e2b_code_interpreter.run_code( + code, + on_stderr=lambda stderr: print("[Code Interpreter]", stderr), + on_stdout=lambda stdout: print("[Code Interpreter]", stdout), + ) + + if exec.error: + print("[Code Interpreter ERROR]", exec.error) + else: + return exec.results + +# Initialize Together AI client +client = Together(api_key=TOGETHER_API_KEY) + +# Regex pattern to extract Python code blocks from LLM responses +pattern = re.compile(r"```python\n(.*?)\n```", re.DOTALL) + +# Function to extract Python code from LLM responses +def match_code_blocks(llm_response): + match = pattern.search(llm_response) + if match: + code = match.group(1) + print("Extracted Python code:") + print(code) + return code + return "" + +# Function to interact with the LLM and execute code in the sandbox +def chat_with_llm(e2b_code_interpreter, user_message): + """ + Interact with LLM and execute code in sandbox. + + Args: + e2b_code_interpreter: The E2B Sandbox instance + user_message: User's query string + + Returns: + List of results from code execution + """ + print(f"\n{'='*50}\nUser message: {user_message}\n{'='*50}") + + # Add file path information to the user message + enhanced_message = f""" +The dataset is located at '/data.csv' in the current directory. +User query: {user_message} +Important: Always use '/data.csv' as the path when reading the dataset. +""" + + # Prepare messages for the LLM + messages = [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": enhanced_message}, + ] + + # Get response from Together AI + response = client.chat.completions.create( + model=MODEL_NAME, + messages=messages, + ) + + # Extract the response message + response_message = response.choices[0].message.content + print("LLM Response:") + print(response_message) + + # Extract Python code from the response + python_code = match_code_blocks(response_message) + if python_code: + # Execute the code in the sandbox + code_interpreter_results = code_interpret(e2b_code_interpreter, python_code) + return code_interpreter_results + else: + print(f"Failed to match any Python code in model's response: {response_message}") + return [] + +# Function to upload a dataset to the E2B Sandbox +def upload_dataset(code_interpreter: Sandbox, uploaded_file: Any) -> str: + """ + Upload a dataset to the E2B Sandbox from Streamlit's uploaded file. + + Args: + code_interpreter: The E2B Sandbox instance + uploaded_file: Streamlit's UploadedFile object + + Returns: + str: Path to the uploaded dataset in the sandbox + """ + print("Uploading dataset to Code Interpreter sandbox...") + + try: + # Create a temporary file to store the uploaded content + with tempfile.NamedTemporaryFile(delete=False, suffix='.csv') as tmp_file: + tmp_file.write(uploaded_file.getvalue()) + dataset_path = tmp_file.name + + # Upload the dataset to the sandbox + with open(dataset_path, "rb") as f: + code_interpreter.files.write("/data.csv", f) + + # Clean up the temporary file + os.unlink(dataset_path) + + print("Dataset uploaded to: /data.csv") + return "/data.csv" + except Exception as error: + print("Error during file upload:", error) + raise error + +def main(): + """Main function to run the Streamlit application.""" + st.title("AI Data Visualization Agent") + st.write("Upload your dataset and ask questions about it!") + + # File uploader + uploaded_file = st.file_uploader("Choose a CSV file", type="csv") + + # Text input for the query + user_query = st.text_input("Enter your visualization query:") + + # Process button + if st.button("Generate Visualization") and uploaded_file is not None and user_query: + try: + with Sandbox(api_key=E2B_API_KEY) as code_interpreter: + # Upload the dataset + upload_dataset(code_interpreter, uploaded_file) + + # Get and execute the visualization code + with st.spinner("Generating visualization..."): + code_results = chat_with_llm(code_interpreter, user_query) + + # Display results + if code_results: + first_result = code_results[0] + + # If there's an image output + if hasattr(first_result, "png"): + st.image(first_result.png, caption="Generated Visualization") + else: + st.write("Results:", first_result) + else: + st.error("No results generated") + + except Exception as e: + st.error(f"An error occurred: {e}") + elif not uploaded_file: + st.warning("Please upload a dataset first") + elif not user_query: + st.warning("Please enter a query") + +if __name__ == "__main__": + main() \ No newline at end of file From 340aa821fb8afa26a3f8e9e636ea3e6c1b1b6281 Mon Sep 17 00:00:00 2001 From: Madhu Date: Sat, 4 Jan 2025 23:58:25 +0530 Subject: [PATCH 3/7] new issues-2 --- .../ai_data_visualisation_agent/test.py | 31 ++++++++++++------- 1 file changed, 19 insertions(+), 12 deletions(-) diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/test.py b/ai_agent_tutorials/ai_data_visualisation_agent/test.py index dca9a6c..d4472e0 100644 --- a/ai_agent_tutorials/ai_data_visualisation_agent/test.py +++ b/ai_agent_tutorials/ai_data_visualisation_agent/test.py @@ -7,6 +7,9 @@ from together import Together from e2b_code_interpreter import Sandbox from typing import Optional, List, Any import tempfile +import base64 +from io import BytesIO +from PIL import Image # Load environment variables load_dotenv() @@ -86,7 +89,7 @@ def chat_with_llm(e2b_code_interpreter, user_message): user_message: User's query string Returns: - List of results from code execution + Base64-encoded image data or None if no image is generated """ print(f"\n{'='*50}\nUser message: {user_message}\n{'='*50}") @@ -119,10 +122,15 @@ Important: Always use '/data.csv' as the path when reading the dataset. if python_code: # Execute the code in the sandbox code_interpreter_results = code_interpret(e2b_code_interpreter, python_code) - return code_interpreter_results + + # Return the base64-encoded image data + if code_interpreter_results and hasattr(code_interpreter_results[0], "png"): + return code_interpreter_results[0].png + else: + return None else: print(f"Failed to match any Python code in model's response: {response_message}") - return [] + return None # Function to upload a dataset to the E2B Sandbox def upload_dataset(code_interpreter: Sandbox, uploaded_file: Any) -> str: @@ -177,19 +185,18 @@ def main(): # Get and execute the visualization code with st.spinner("Generating visualization..."): - code_results = chat_with_llm(code_interpreter, user_query) + image_data = chat_with_llm(code_interpreter, user_query) # Display results - if code_results: - first_result = code_results[0] + if image_data: + # Decode the base64-encoded image data + image_bytes = base64.b64decode(image_data) + image = Image.open(BytesIO(image_bytes)) - # If there's an image output - if hasattr(first_result, "png"): - st.image(first_result.png, caption="Generated Visualization") - else: - st.write("Results:", first_result) + # Display the image in Streamlit + st.image(image, caption="Generated Visualization") else: - st.error("No results generated") + st.error("No visualization generated") except Exception as e: st.error(f"An error occurred: {e}") From 8cdb9bd6634f15da6b579b6d63311bcdc75b2813 Mon Sep 17 00:00:00 2001 From: Madhu Date: Sun, 5 Jan 2025 16:55:18 +0530 Subject: [PATCH 4/7] successful working of code --- .../ai_data_visualisation_agent/.gitignore | 1 - .../ai_data_visualisation_agent.py | 11 +- .../requirements.txt | 11 +- .../ai_data_visualisation_agent/test.py | 286 ++++++++---------- 4 files changed, 138 insertions(+), 171 deletions(-) delete mode 100644 ai_agent_tutorials/ai_data_visualisation_agent/.gitignore diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/.gitignore b/ai_agent_tutorials/ai_data_visualisation_agent/.gitignore deleted file mode 100644 index 2eea525..0000000 --- a/ai_agent_tutorials/ai_data_visualisation_agent/.gitignore +++ /dev/null @@ -1 +0,0 @@ -.env \ No newline at end of file diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py b/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py index adf1b53..429b09f 100644 --- a/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py +++ b/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py @@ -25,9 +25,6 @@ def preprocess_and_save(file): st.error("Unsupported file format. Please upload a CSV or Excel file.") return None, None, None - # Log the data types of columns before preprocessing - st.write("Data types before preprocessing:") - st.write(df.dtypes) # Ensure string columns are properly quoted for col in df.select_dtypes(include=['object']): @@ -53,10 +50,6 @@ def preprocess_and_save(file): # Drop rows with all NaN values df.dropna(how='all', inplace=True) - # Log the data types of columns after preprocessing - st.write("Data types after preprocessing:") - st.write(df.dtypes) - # Create a temporary file to save the preprocessed data with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as temp_file: temp_path = temp_file.name @@ -137,7 +130,7 @@ def chat_with_llm(user_message, file_path, columns, df): return response_message, None # Set up Streamlit app -st.title("AI Data Scientist") +st.title("AI Data Visualisation Agent") # Sidebar for API keys and file upload st.sidebar.header("API Keys") @@ -151,7 +144,7 @@ uploaded_file = st.sidebar.file_uploader("Upload CSV or Excel File", type=['csv' # System prompt (dynamic based on the uploaded file) SYSTEM_PROMPT = """ -You are a Python data scientist. You have access to a CSV file located at '{file_path}'. +You are a Python data scientist and Visualisation expert. You have access to a CSV file located at '{file_path}'. The dataset has the following columns: {columns}. You can read this file into a DataFrame using `df = pd.read_csv('{file_path}')` and perform data analysis tasks based on user queries. Make sure to handle missing values and data type inconsistencies. When generating plots, diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt b/ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt index 99a78ca..8019c78 100644 --- a/ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt +++ b/ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt @@ -1,4 +1,9 @@ +together>=0.2.8 +e2b>=0.12.0 python-dotenv -together -e2b -streamlit \ No newline at end of file +Pillow +streamlit +pandas +matplotlib +plotly +seaborn>=0.12.0 \ No newline at end of file diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/test.py b/ai_agent_tutorials/ai_data_visualisation_agent/test.py index d4472e0..3ca61ca 100644 --- a/ai_agent_tutorials/ai_data_visualisation_agent/test.py +++ b/ai_agent_tutorials/ai_data_visualisation_agent/test.py @@ -1,209 +1,179 @@ -import streamlit as st import os -from dotenv import load_dotenv import json import re -from together import Together -from e2b_code_interpreter import Sandbox -from typing import Optional, List, Any -import tempfile +from typing import Optional, List, Any, Tuple +from dotenv import load_dotenv +from PIL import Image +import io +import streamlit as st +import pandas as pd import base64 from io import BytesIO from PIL import Image +from together import Together +from e2b_code_interpreter import Sandbox # Load environment variables load_dotenv() -# Get API keys from environment variables -TOGETHER_API_KEY = os.getenv("TOGETHER_API_KEY") -E2B_API_KEY = os.getenv("E2B_API_KEY") - -# Define the Together AI model to use -MODEL_NAME = "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo" - -# System prompt for the LLM -SYSTEM_PROMPT = """You are a highly skilled Python data scientist. Your task is to analyze datasets and generate Python code to solve data-related problems. Follow these guidelines: - -1. **Data Preprocessing**: - - Always check for missing or invalid values in the dataset. - - Handle missing values by either removing rows/columns or imputing them appropriately. - - Convert columns to the correct data types (e.g., numeric, datetime). - - Filter out rows with invalid or inconsistent data. - -2. **Data Analysis**: - - Perform exploratory data analysis (EDA) to understand the dataset. - - Use statistical methods to analyze relationships between variables. - - If the task involves machine learning (e.g., linear regression), ensure the data is properly prepared (e.g., feature scaling, train-test split). - -3. **Visualization**: - - Use libraries like `matplotlib` or `seaborn` for creating visualizations. - - Ensure plots are clear, labeled, and informative (e.g., include titles, axis labels, legends). - - Save plots as images (e.g., PNG) and return them as base64-encoded strings. - -4. **Code Quality**: - - Write clean, modular, and well-commented Python code. - - Handle potential errors gracefully (e.g., invalid data, missing columns). - - Include necessary imports (e.g., `pandas`, `numpy`, `matplotlib`, `seaborn`). - -5. **Output**: - - Always return the Python code to solve the task. - - If the task involves visualization, include the code to generate and save the plot.""" - -# Function to execute code in the E2B Sandbox -def code_interpret(e2b_code_interpreter, code): - print("Running code interpreter...") - exec = e2b_code_interpreter.run_code( - code, - on_stderr=lambda stderr: print("[Code Interpreter]", stderr), - on_stdout=lambda stdout: print("[Code Interpreter]", stdout), - ) - - if exec.error: - print("[Code Interpreter ERROR]", exec.error) - else: - return exec.results - -# Initialize Together AI client -client = Together(api_key=TOGETHER_API_KEY) - -# Regex pattern to extract Python code blocks from LLM responses +# Regex pattern to extract code from LLM response pattern = re.compile(r"```python\n(.*?)\n```", re.DOTALL) -# Function to extract Python code from LLM responses -def match_code_blocks(llm_response): +def code_interpret(e2b_code_interpreter: Sandbox, code: str) -> Optional[List[Any]]: + """ + Runs the given Python code in the E2B sandbox. + + Args: + e2b_code_interpreter: The E2B sandbox instance + code: Python code to execute + + Returns: + Optional[List[Any]]: Results from code execution + """ + with st.spinner('Executing code in E2B sandbox...'): + exec = e2b_code_interpreter.run_code(code, + on_stderr=lambda stderr: st.error(f"[Code Interpreter] {stderr}"), + on_stdout=lambda stdout: st.info(f"[Code Interpreter] {stdout}")) + + if exec.error: + st.error(f"[Code Interpreter ERROR] {exec.error}") + return None + return exec.results + +def match_code_blocks(llm_response: str) -> str: + """ + Extracts Python code blocks from the LLM response. + + Args: + llm_response: The response from the LLM + + Returns: + str: Extracted Python code or empty string + """ match = pattern.search(llm_response) if match: code = match.group(1) - print("Extracted Python code:") - print(code) return code return "" -# Function to interact with the LLM and execute code in the sandbox -def chat_with_llm(e2b_code_interpreter, user_message): +def chat_with_llm(e2b_code_interpreter: Sandbox, user_message: str, dataset_path: str) -> Tuple[Optional[List[Any]], str]: """ - Interact with LLM and execute code in sandbox. + Sends the user message to the LLM and executes the generated code. Args: - e2b_code_interpreter: The E2B Sandbox instance - user_message: User's query string - + e2b_code_interpreter: The E2B sandbox instance + user_message: User's query message + dataset_path: Path to the uploaded dataset + Returns: - Base64-encoded image data or None if no image is generated + Tuple[Optional[List[Any]], str]: Code execution results and LLM response """ - print(f"\n{'='*50}\nUser message: {user_message}\n{'='*50}") + # Update system prompt to include dataset path information + system_prompt = f"""You're a Python data scientist and data visualization expert. You are given a dataset at path '{dataset_path}' and also the user's query. +You need to analyze the dataset and answer the user's query with a response and you run Python code to solve them. +IMPORTANT: Always use the dataset path variable '{dataset_path}' in your code when reading the CSV file.""" - # Add file path information to the user message - enhanced_message = f""" -The dataset is located at '/data.csv' in the current directory. -User query: {user_message} -Important: Always use '/data.csv' as the path when reading the dataset. -""" - - # Prepare messages for the LLM messages = [ - {"role": "system", "content": SYSTEM_PROMPT}, - {"role": "user", "content": enhanced_message}, + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_message}, ] - # Get response from Together AI - response = client.chat.completions.create( - model=MODEL_NAME, - messages=messages, - ) + with st.spinner('Getting response from together AI...'): + client = Together(api_key=st.session_state.together_api_key) + response = client.chat.completions.create( + model=st.session_state.model_name, + messages=messages, + ) - # Extract the response message - response_message = response.choices[0].message.content - print("LLM Response:") - print(response_message) - - # Extract Python code from the response - python_code = match_code_blocks(response_message) - if python_code: - # Execute the code in the sandbox - code_interpreter_results = code_interpret(e2b_code_interpreter, python_code) + response_message = response.choices[0].message + python_code = match_code_blocks(response_message.content) - # Return the base64-encoded image data - if code_interpreter_results and hasattr(code_interpreter_results[0], "png"): - return code_interpreter_results[0].png + if python_code: + code_interpreter_results = code_interpret(e2b_code_interpreter, python_code) + return code_interpreter_results, response_message.content else: - return None - else: - print(f"Failed to match any Python code in model's response: {response_message}") - return None + st.warning(f"Failed to match any Python code in model's response") + return None, response_message.content -# Function to upload a dataset to the E2B Sandbox -def upload_dataset(code_interpreter: Sandbox, uploaded_file: Any) -> str: +def upload_dataset(code_interpreter: Sandbox, uploaded_file) -> str: """ - Upload a dataset to the E2B Sandbox from Streamlit's uploaded file. + Uploads the dataset to the E2B sandbox. Args: - code_interpreter: The E2B Sandbox instance - uploaded_file: Streamlit's UploadedFile object - + code_interpreter: The E2B sandbox instance + uploaded_file: Streamlit uploaded file + Returns: - str: Path to the uploaded dataset in the sandbox + str: Path where file was uploaded """ - print("Uploading dataset to Code Interpreter sandbox...") + dataset_path = f"./{uploaded_file.name}" try: - # Create a temporary file to store the uploaded content - with tempfile.NamedTemporaryFile(delete=False, suffix='.csv') as tmp_file: - tmp_file.write(uploaded_file.getvalue()) - dataset_path = tmp_file.name - - # Upload the dataset to the sandbox - with open(dataset_path, "rb") as f: - code_interpreter.files.write("/data.csv", f) - - # Clean up the temporary file - os.unlink(dataset_path) - - print("Dataset uploaded to: /data.csv") - return "/data.csv" + code_interpreter.files.write(dataset_path, uploaded_file) + return dataset_path except Exception as error: - print("Error during file upload:", error) + st.error(f"Error during file upload: {error}") raise error + def main(): - """Main function to run the Streamlit application.""" - st.title("AI Data Visualization Agent") + """Main Streamlit application.""" + st.title("AI Data Visualization Assistant") st.write("Upload your dataset and ask questions about it!") - # File uploader + # Sidebar for API keys and model name + with st.sidebar: + st.header("API Keys and Model Configuration") + st.session_state.together_api_key = st.text_input("Enter Together API Key", type="password") + st.session_state.e2b_api_key = st.text_input("Enter E2B API Key", type="password") + st.session_state.model_name = st.text_input("Enter Model Name", value="meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo") + uploaded_file = st.file_uploader("Choose a CSV file", type="csv") - # Text input for the query - user_query = st.text_input("Enter your visualization query:") - - # Process button - if st.button("Generate Visualization") and uploaded_file is not None and user_query: - try: - with Sandbox(api_key=E2B_API_KEY) as code_interpreter: - # Upload the dataset - upload_dataset(code_interpreter, uploaded_file) - - # Get and execute the visualization code - with st.spinner("Generating visualization..."): - image_data = chat_with_llm(code_interpreter, user_query) - - # Display results - if image_data: - # Decode the base64-encoded image data - image_bytes = base64.b64decode(image_data) - image = Image.open(BytesIO(image_bytes)) + if uploaded_file is not None: + # Display dataset preview + df = pd.read_csv(uploaded_file) + st.write("Dataset Preview:") + st.dataframe(df.head()) + + # Query input + query = st.text_area("What would you like to know about your data?", + "Can you compare the average cost for two people between different categories?") + + if st.button("Analyze"): + if not st.session_state.together_api_key or not st.session_state.e2b_api_key: + st.error("Please enter both API keys in the sidebar.") + else: + with Sandbox(api_key=st.session_state.e2b_api_key) as code_interpreter: + # Upload the dataset + dataset_path = upload_dataset(code_interpreter, uploaded_file) - # Display the image in Streamlit - st.image(image, caption="Generated Visualization") - else: - st.error("No visualization generated") + # Pass dataset_path to chat_with_llm + code_results, llm_response = chat_with_llm(code_interpreter, query, dataset_path) - except Exception as e: - st.error(f"An error occurred: {e}") - elif not uploaded_file: - st.warning("Please upload a dataset first") - elif not user_query: - st.warning("Please enter a query") + # Display LLM's text response + st.write("AI Response:") + st.write(llm_response) + + # Display results/visualizations + if code_results: + for result in code_results: + if hasattr(result, 'png') and result.png: # Check if PNG data is available + # Decode the base64-encoded PNG data + png_data = base64.b64decode(result.png) + + # Convert PNG data to an image and display it + image = Image.open(BytesIO(png_data)) + st.image(image, caption="Generated Visualization", use_container_width=False) + elif hasattr(result, 'figure'): # For matplotlib figures + fig = result.figure # Extract the matplotlib figure + st.pyplot(fig) # Display using st.pyplot + elif hasattr(result, 'show'): # For plotly figures + st.plotly_chart(result) + elif isinstance(result, (pd.DataFrame, pd.Series)): + st.dataframe(result) + else: + st.write(result) if __name__ == "__main__": main() \ No newline at end of file From 9fd044b81edc63cdbc3739116ddae045a2c21d02 Mon Sep 17 00:00:00 2001 From: Madhu Date: Sun, 5 Jan 2025 17:22:09 +0530 Subject: [PATCH 5/7] code is perfect --- .../ai_data_visualisation_agent.py | 378 +++++++++--------- .../ai_data_visualisation_agent/test.py | 179 --------- 2 files changed, 195 insertions(+), 362 deletions(-) delete mode 100644 ai_agent_tutorials/ai_data_visualisation_agent/test.py diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py b/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py index 429b09f..682cf09 100644 --- a/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py +++ b/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py @@ -1,202 +1,214 @@ +import os +import json +import re +import sys +import io +import contextlib +import warnings +from typing import Optional, List, Any, Tuple +from dotenv import load_dotenv +from PIL import Image import streamlit as st import pandas as pd -import tempfile -import re -from together import Together -import csv -from dotenv import load_dotenv import base64 -import matplotlib.pyplot as plt -import io -import seaborn as sns +from io import BytesIO +from together import Together +from e2b_code_interpreter import Sandbox -# Load environment variables -load_dotenv() +# Suppress Pydantic warnings globally +warnings.filterwarnings("ignore", category=UserWarning, module="pydantic") -# Function to preprocess and save the uploaded file to a temporary file -def preprocess_and_save(file): - try: - # Read the uploaded file into a DataFrame - if file.name.endswith('.csv'): - df = pd.read_csv(file, encoding='utf-8', na_values=['NA', 'N/A', 'missing']) - elif file.name.endswith('.xlsx'): - df = pd.read_excel(file, na_values=['NA', 'N/A', 'missing']) - else: - st.error("Unsupported file format. Please upload a CSV or Excel file.") - return None, None, None - - - # Ensure string columns are properly quoted - for col in df.select_dtypes(include=['object']): - df[col] = df[col].astype(str).replace({r'"': '""'}, regex=True) - - # Parse dates and numeric columns - for col in df.columns: - if 'date' in col.lower(): - df[col] = pd.to_datetime(df[col], errors='coerce') - elif df[col].dtype == 'object': - try: - # Handle columns with values like "4.1/5" - if df[col].str.contains('/').any(): - # Split the values and take the first part (e.g., "4.1/5" -> 4.1) - df[col] = df[col].str.split('/').str[0] - # Convert to numeric, coerce errors to NaN - df[col] = pd.to_numeric(df[col], errors='coerce') - except (ValueError, TypeError): - # Keep as is if conversion fails - st.warning(f"Could not convert column '{col}' to numeric. Keeping as string.") - pass - - # Drop rows with all NaN values - df.dropna(how='all', inplace=True) - - # Create a temporary file to save the preprocessed data - with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as temp_file: - temp_path = temp_file.name - # Save the DataFrame to the temporary CSV file - df.to_csv(temp_path, index=False, quoting=csv.QUOTE_ALL) - - return temp_path, df.columns.tolist(), df # Return the DataFrame as well - except Exception as e: - st.error(f"Error processing file: {e}") - return None, None, None +# Regex pattern to extract code from LLM response +pattern = re.compile(r"```python\n(.*?)\n```", re.DOTALL) -# Function to execute Python code and generate plots -def execute_code(code: str, df): - try: - # Define locals with necessary imports and the DataFrame - local_env = { - 'pd': pd, - 'df': df, - 'plt': plt, - 'sns': sns # if seaborn is needed - } - # Execute the code in the local environment - exec(code, globals(), local_env) +def code_interpret(e2b_code_interpreter: Sandbox, code: str) -> Optional[List[Any]]: + """ + Runs the given Python code in the E2B sandbox. + + Args: + e2b_code_interpreter: The E2B sandbox instance + code: Python code to execute - # Check if a plot was generated - if 'plt' in local_env: - # Save the plot to a BytesIO object - buf = io.BytesIO() - plt.savefig(buf, format='png') - plt.close() - buf.seek(0) - # Encode the plot as base64 - base64_image = base64.b64encode(buf.read()).decode('utf-8') - return base64_image - else: - st.warning("No plot generated. Ensure the data being plotted is numeric.") + Returns: + Optional[List[Any]]: Results from code execution + """ + with st.spinner('Executing code in E2B sandbox...'): + # Capture stdout and stderr + stdout_capture = io.StringIO() + stderr_capture = io.StringIO() + + with contextlib.redirect_stdout(stdout_capture), contextlib.redirect_stderr(stderr_capture): + # Suppress warnings during code execution + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + exec = e2b_code_interpreter.run_code(code) + + # Log stderr (warnings and errors) to the terminal + if stderr_capture.getvalue(): + print("[Code Interpreter Warnings/Errors]", file=sys.stderr) + print(stderr_capture.getvalue(), file=sys.stderr) + + # Log stdout (normal output) to the terminal + if stdout_capture.getvalue(): + print("[Code Interpreter Output]", file=sys.stdout) + print(stdout_capture.getvalue(), file=sys.stdout) + + if exec.error: + print(f"[Code Interpreter ERROR] {exec.error}", file=sys.stderr) return None - except Exception as e: - st.error(f"Error executing code: {e}") - return None + return exec.results -# Function to communicate with Together AI -def chat_with_llm(user_message, file_path, columns, df): - print(f"\n{'='*50}\nUser message: {user_message}\n{'='*50}") +def match_code_blocks(llm_response: str) -> str: + """ + Extracts Python code blocks from the LLM response. + + Args: + llm_response: The response from the LLM + + Returns: + str: Extracted Python code or empty string + """ + match = pattern.search(llm_response) + if match: + code = match.group(1) + return code + return "" - # Update the system prompt with the file path, columns, and plot path - system_prompt = SYSTEM_PROMPT.format( - file_path=file_path, - columns=columns, - ) - - # Add a hint to include a plot if the user asks for visualization - if "plot" in user_message.lower(): - system_prompt += " Include a plot in your response and output the base64 string of the plot image." +def chat_with_llm(e2b_code_interpreter: Sandbox, user_message: str, dataset_path: str) -> Tuple[Optional[List[Any]], str]: + """ + Sends the user message to the LLM and executes the generated code. + + Args: + e2b_code_interpreter: The E2B sandbox instance + user_message: User's query message + dataset_path: Path to the uploaded dataset + + Returns: + Tuple[Optional[List[Any]], str]: Code execution results and LLM response + """ + # Update system prompt to include dataset path information + system_prompt = f"""You're a Python data scientist and data visualization expert. You are given a dataset at path '{dataset_path}' and also the user's query. +You need to analyze the dataset and answer the user's query with a response and you run Python code to solve them. +IMPORTANT: Always use the dataset path variable '{dataset_path}' in your code when reading the CSV file.""" messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_message}, ] - # Use the Together API key from session state - response = client.chat.completions.create( - model="meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo", - messages=messages, - ) + with st.spinner('Getting response from Together AI LLM model...'): + client = Together(api_key=st.session_state.together_api_key) + response = client.chat.completions.create( + model=st.session_state.model_name, + messages=messages, + ) - response_message = response.choices[0].message.content - print("LLM Response:", response_message) # Debug: Print the LLM's response - python_code = match_code_blocks(response_message) - print("Extracted Python Code:", python_code) # Debug: Print the extracted code - - if python_code: - # Execute the code and generate the plot - base64_image = execute_code(python_code, df) - return response_message, base64_image - else: - print(f"Failed to match any Python code in model's response {response_message}") - return response_message, None - -# Set up Streamlit app -st.title("AI Data Visualisation Agent") - -# Sidebar for API keys and file upload -st.sidebar.header("API Keys") -together_api_key = st.sidebar.text_input("Together AI API Key", type="password") - -# Store API key in session state -if 'together_api_key' not in st.session_state: - st.session_state.together_api_key = None - -uploaded_file = st.sidebar.file_uploader("Upload CSV or Excel File", type=['csv', 'xlsx']) - -# System prompt (dynamic based on the uploaded file) -SYSTEM_PROMPT = """ -You are a Python data scientist and Visualisation expert. You have access to a CSV file located at '{file_path}'. -The dataset has the following columns: {columns}. -You can read this file into a DataFrame using `df = pd.read_csv('{file_path}')` and perform data analysis tasks based on user queries. -Make sure to handle missing values and data type inconsistencies. When generating plots, -use matplotlib or seaborn and output the plot as a base64 string. -Always check if the data being plotted is numeric. If the data is not numeric, preprocess it to convert it to numeric values. -Always respond with the Python code to answer the user's query, and include visualizations only if explicitly requested. -""" - -# Function to match Python code blocks -pattern = re.compile(r"```python\n(.*?)\n```", re.DOTALL) - -def match_code_blocks(llm_response): - match = pattern.search(llm_response) - if match: - code = match.group(1) - # Remove comments and extra text - code = "\n".join([line for line in code.split("\n") if not line.strip().startswith("#")]) - return code - return "" - -# Main app logic -if uploaded_file: - if not together_api_key: - st.warning("Please provide the Together AI API key.") - else: - # Update session state with API key - st.session_state.together_api_key = together_api_key + response_message = response.choices[0].message + python_code = match_code_blocks(response_message.content) - # Initialize Together AI client only after confirming API key exists - try: - client = Together(api_key=together_api_key) - - # Preprocess and save the uploaded file - temp_path, columns, df = preprocess_and_save(uploaded_file) - if temp_path: - # Rest of your code for user query handling - user_query = st.text_input("Ask a query about the data:") - if st.button("Submit Query"): - response_message, base64_image = chat_with_llm(user_query, temp_path, columns, df) - - # Display AI's response - st.write("AI's Response:") - st.write(response_message) - - # Display the plot if generated - if base64_image: - st.image(base64.b64decode(base64_image), use_container_width=True) - else: - st.write("No plot generated.") + if python_code: + code_interpreter_results = code_interpret(e2b_code_interpreter, python_code) + return code_interpreter_results, response_message.content + else: + st.warning(f"Failed to match any Python code in model's response") + return None, response_message.content + +def upload_dataset(code_interpreter: Sandbox, uploaded_file) -> str: + """ + Uploads the dataset to the E2B sandbox. + + Args: + code_interpreter: The E2B sandbox instance + uploaded_file: Streamlit uploaded file + + Returns: + str: Path where file was uploaded + """ + dataset_path = f"./{uploaded_file.name}" + + try: + code_interpreter.files.write(dataset_path, uploaded_file) + return dataset_path + except Exception as error: + st.error(f"Error during file upload: {error}") + raise error + + +def main(): + """Main Streamlit application.""" + st.title("AI Data Visualization Agent") + st.write("Upload your dataset and ask questions about it!") + + # Sidebar for API keys and model selection + with st.sidebar: + st.header("API Keys and Model Configuration") + st.session_state.together_api_key = st.text_input("Enter Together API Key", type="password") + st.session_state.e2b_api_key = st.text_input("Enter E2B API Key", type="password") + + # Add model selection dropdown + model_options = { + "Meta-Llama 3.1 405B": "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo", + "DeepSeek V3": "deepseek-ai/DeepSeek-V3", + "Qwen 2.5 7B": "Qwen/Qwen2.5-7B-Instruct-Turbo", + "Meta-Llama 3.3 70B": "meta-llama/Llama-3.3-70B-Instruct-Turbo" + } + selected_model = st.selectbox( + "Select Model", + options=list(model_options.keys()), + index=0 # Default to first option + ) + st.session_state.model_name = model_options[selected_model] + + uploaded_file = st.file_uploader("Choose a CSV file", type="csv") + + if uploaded_file is not None: + # Display dataset with toggle + df = pd.read_csv(uploaded_file) + st.write("Dataset:") + show_full = st.checkbox("Show full dataset") + if show_full: + st.dataframe(df) + else: + st.write("Preview (first 5 rows):") + st.dataframe(df.head()) + # Query input + query = st.text_area("What would you like to know about your data?", + "Can you compare the average cost for two people between different categories?") + + if st.button("Analyze"): + if not st.session_state.together_api_key or not st.session_state.e2b_api_key: + st.error("Please enter both API keys in the sidebar.") else: - st.error("Failed to preprocess and save the data.") - except Exception as e: - st.error(f"Error initializing Together AI client: {str(e)}") -else: - st.warning("Please upload a file.") \ No newline at end of file + with Sandbox(api_key=st.session_state.e2b_api_key) as code_interpreter: + # Upload the dataset + dataset_path = upload_dataset(code_interpreter, uploaded_file) + + # Pass dataset_path to chat_with_llm + code_results, llm_response = chat_with_llm(code_interpreter, query, dataset_path) + + # Display LLM's text response + st.write("AI Response:") + st.write(llm_response) + + # Display results/visualizations + if code_results: + for result in code_results: + if hasattr(result, 'png') and result.png: # Check if PNG data is available + # Decode the base64-encoded PNG data + png_data = base64.b64decode(result.png) + + # Convert PNG data to an image and display it + image = Image.open(BytesIO(png_data)) + st.image(image, caption="Generated Visualization", use_container_width=False) + elif hasattr(result, 'figure'): # For matplotlib figures + fig = result.figure # Extract the matplotlib figure + st.pyplot(fig) # Display using st.pyplot + elif hasattr(result, 'show'): # For plotly figures + st.plotly_chart(result) + elif isinstance(result, (pd.DataFrame, pd.Series)): + st.dataframe(result) + else: + st.write(result) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/test.py b/ai_agent_tutorials/ai_data_visualisation_agent/test.py deleted file mode 100644 index 3ca61ca..0000000 --- a/ai_agent_tutorials/ai_data_visualisation_agent/test.py +++ /dev/null @@ -1,179 +0,0 @@ -import os -import json -import re -from typing import Optional, List, Any, Tuple -from dotenv import load_dotenv -from PIL import Image -import io -import streamlit as st -import pandas as pd -import base64 -from io import BytesIO -from PIL import Image -from together import Together -from e2b_code_interpreter import Sandbox - -# Load environment variables -load_dotenv() - -# Regex pattern to extract code from LLM response -pattern = re.compile(r"```python\n(.*?)\n```", re.DOTALL) - -def code_interpret(e2b_code_interpreter: Sandbox, code: str) -> Optional[List[Any]]: - """ - Runs the given Python code in the E2B sandbox. - - Args: - e2b_code_interpreter: The E2B sandbox instance - code: Python code to execute - - Returns: - Optional[List[Any]]: Results from code execution - """ - with st.spinner('Executing code in E2B sandbox...'): - exec = e2b_code_interpreter.run_code(code, - on_stderr=lambda stderr: st.error(f"[Code Interpreter] {stderr}"), - on_stdout=lambda stdout: st.info(f"[Code Interpreter] {stdout}")) - - if exec.error: - st.error(f"[Code Interpreter ERROR] {exec.error}") - return None - return exec.results - -def match_code_blocks(llm_response: str) -> str: - """ - Extracts Python code blocks from the LLM response. - - Args: - llm_response: The response from the LLM - - Returns: - str: Extracted Python code or empty string - """ - match = pattern.search(llm_response) - if match: - code = match.group(1) - return code - return "" - -def chat_with_llm(e2b_code_interpreter: Sandbox, user_message: str, dataset_path: str) -> Tuple[Optional[List[Any]], str]: - """ - Sends the user message to the LLM and executes the generated code. - - Args: - e2b_code_interpreter: The E2B sandbox instance - user_message: User's query message - dataset_path: Path to the uploaded dataset - - Returns: - Tuple[Optional[List[Any]], str]: Code execution results and LLM response - """ - # Update system prompt to include dataset path information - system_prompt = f"""You're a Python data scientist and data visualization expert. You are given a dataset at path '{dataset_path}' and also the user's query. -You need to analyze the dataset and answer the user's query with a response and you run Python code to solve them. -IMPORTANT: Always use the dataset path variable '{dataset_path}' in your code when reading the CSV file.""" - - messages = [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": user_message}, - ] - - with st.spinner('Getting response from together AI...'): - client = Together(api_key=st.session_state.together_api_key) - response = client.chat.completions.create( - model=st.session_state.model_name, - messages=messages, - ) - - response_message = response.choices[0].message - python_code = match_code_blocks(response_message.content) - - if python_code: - code_interpreter_results = code_interpret(e2b_code_interpreter, python_code) - return code_interpreter_results, response_message.content - else: - st.warning(f"Failed to match any Python code in model's response") - return None, response_message.content - -def upload_dataset(code_interpreter: Sandbox, uploaded_file) -> str: - """ - Uploads the dataset to the E2B sandbox. - - Args: - code_interpreter: The E2B sandbox instance - uploaded_file: Streamlit uploaded file - - Returns: - str: Path where file was uploaded - """ - dataset_path = f"./{uploaded_file.name}" - - try: - code_interpreter.files.write(dataset_path, uploaded_file) - return dataset_path - except Exception as error: - st.error(f"Error during file upload: {error}") - raise error - - -def main(): - """Main Streamlit application.""" - st.title("AI Data Visualization Assistant") - st.write("Upload your dataset and ask questions about it!") - - # Sidebar for API keys and model name - with st.sidebar: - st.header("API Keys and Model Configuration") - st.session_state.together_api_key = st.text_input("Enter Together API Key", type="password") - st.session_state.e2b_api_key = st.text_input("Enter E2B API Key", type="password") - st.session_state.model_name = st.text_input("Enter Model Name", value="meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo") - - uploaded_file = st.file_uploader("Choose a CSV file", type="csv") - - if uploaded_file is not None: - # Display dataset preview - df = pd.read_csv(uploaded_file) - st.write("Dataset Preview:") - st.dataframe(df.head()) - - # Query input - query = st.text_area("What would you like to know about your data?", - "Can you compare the average cost for two people between different categories?") - - if st.button("Analyze"): - if not st.session_state.together_api_key or not st.session_state.e2b_api_key: - st.error("Please enter both API keys in the sidebar.") - else: - with Sandbox(api_key=st.session_state.e2b_api_key) as code_interpreter: - # Upload the dataset - dataset_path = upload_dataset(code_interpreter, uploaded_file) - - # Pass dataset_path to chat_with_llm - code_results, llm_response = chat_with_llm(code_interpreter, query, dataset_path) - - # Display LLM's text response - st.write("AI Response:") - st.write(llm_response) - - # Display results/visualizations - if code_results: - for result in code_results: - if hasattr(result, 'png') and result.png: # Check if PNG data is available - # Decode the base64-encoded PNG data - png_data = base64.b64decode(result.png) - - # Convert PNG data to an image and display it - image = Image.open(BytesIO(png_data)) - st.image(image, caption="Generated Visualization", use_container_width=False) - elif hasattr(result, 'figure'): # For matplotlib figures - fig = result.figure # Extract the matplotlib figure - st.pyplot(fig) # Display using st.pyplot - elif hasattr(result, 'show'): # For plotly figures - st.plotly_chart(result) - elif isinstance(result, (pd.DataFrame, pd.Series)): - st.dataframe(result) - else: - st.write(result) - -if __name__ == "__main__": - main() \ No newline at end of file From e2eec7e80f5a93053c3e97a29c933e40b262eb07 Mon Sep 17 00:00:00 2001 From: Madhu Date: Sun, 5 Jan 2025 19:31:14 +0530 Subject: [PATCH 6/7] completed all files --- .../ai_data_visualisation_agent/README.md | 39 +++++++++++ .../ai_data_visualisation_agent.py | 68 +++++-------------- .../requirements.txt | 10 ++- 3 files changed, 59 insertions(+), 58 deletions(-) diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/README.md b/ai_agent_tutorials/ai_data_visualisation_agent/README.md index e69de29..dc69328 100644 --- a/ai_agent_tutorials/ai_data_visualisation_agent/README.md +++ b/ai_agent_tutorials/ai_data_visualisation_agent/README.md @@ -0,0 +1,39 @@ +# AI Data Visualization Agent + +This Assistant is designed to help anyone create and visualize data using natural language commands, and it is built using Together AI and E2B Code Interpreter. User gets to upload a dataset and ask questions to the LLM to get the data visualized. This demo can be considered as a demo for the E2B Code Interpreter and Together AI, for anyone who's getting started with these libraries! + +## Features + +- 🎨 Natural language-driven visualization creation +- 📊 Support for multiple chart types (line, bar, scatter, pie, bubble) +- 📈 Automatic data preprocessing and cleaning +- 🎯 Available Models: + - Meta-Llama 3.1 405B + - DeepSeek V3 + - Qwen 2.5 7B + - Meta-Llama 3.3 70B +- 📱 The Code runs in the E2B Sandbox environment, so it is secure and fast +- Streamlit for clear and interactive user interface + +## How to Run + +Follow the steps below to set up and run the application: +Before anything else, Please get a free Together AI API Key here: https://api.together.ai/signin +Get a free E2B API Key here: https://e2b.dev/ ; https://e2b.dev/docs/legacy/getting-started/api-key + +1. **Clone the Repository**: + ```bash + git clone https://github.com/Shubhamsaboo/awesome-llm-apps.git + cd ai_agent_tutorials/ai_data_visualisation_agent + ``` + +2. **Install the dependencies** + ```bash + pip install -r requirements.txt + ``` + +3. **Run the Streamlit app** + ```bash + streamlit run ai_data_visualisation_agent.py + ``` + diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py b/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py index 682cf09..5260b4e 100644 --- a/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py +++ b/ai_agent_tutorials/ai_data_visualisation_agent/ai_data_visualisation_agent.py @@ -6,7 +6,6 @@ import io import contextlib import warnings from typing import Optional, List, Any, Tuple -from dotenv import load_dotenv from PIL import Image import streamlit as st import pandas as pd @@ -15,40 +14,24 @@ from io import BytesIO from together import Together from e2b_code_interpreter import Sandbox -# Suppress Pydantic warnings globally warnings.filterwarnings("ignore", category=UserWarning, module="pydantic") -# Regex pattern to extract code from LLM response pattern = re.compile(r"```python\n(.*?)\n```", re.DOTALL) def code_interpret(e2b_code_interpreter: Sandbox, code: str) -> Optional[List[Any]]: - """ - Runs the given Python code in the E2B sandbox. - - Args: - e2b_code_interpreter: The E2B sandbox instance - code: Python code to execute - - Returns: - Optional[List[Any]]: Results from code execution - """ with st.spinner('Executing code in E2B sandbox...'): - # Capture stdout and stderr stdout_capture = io.StringIO() stderr_capture = io.StringIO() with contextlib.redirect_stdout(stdout_capture), contextlib.redirect_stderr(stderr_capture): - # Suppress warnings during code execution with warnings.catch_warnings(): warnings.simplefilter("ignore") exec = e2b_code_interpreter.run_code(code) - # Log stderr (warnings and errors) to the terminal if stderr_capture.getvalue(): print("[Code Interpreter Warnings/Errors]", file=sys.stderr) print(stderr_capture.getvalue(), file=sys.stderr) - # Log stdout (normal output) to the terminal if stdout_capture.getvalue(): print("[Code Interpreter Output]", file=sys.stdout) print(stdout_capture.getvalue(), file=sys.stdout) @@ -59,15 +42,6 @@ def code_interpret(e2b_code_interpreter: Sandbox, code: str) -> Optional[List[An return exec.results def match_code_blocks(llm_response: str) -> str: - """ - Extracts Python code blocks from the LLM response. - - Args: - llm_response: The response from the LLM - - Returns: - str: Extracted Python code or empty string - """ match = pattern.search(llm_response) if match: code = match.group(1) @@ -75,17 +49,6 @@ def match_code_blocks(llm_response: str) -> str: return "" def chat_with_llm(e2b_code_interpreter: Sandbox, user_message: str, dataset_path: str) -> Tuple[Optional[List[Any]], str]: - """ - Sends the user message to the LLM and executes the generated code. - - Args: - e2b_code_interpreter: The E2B sandbox instance - user_message: User's query message - dataset_path: Path to the uploaded dataset - - Returns: - Tuple[Optional[List[Any]], str]: Code execution results and LLM response - """ # Update system prompt to include dataset path information system_prompt = f"""You're a Python data scientist and data visualization expert. You are given a dataset at path '{dataset_path}' and also the user's query. You need to analyze the dataset and answer the user's query with a response and you run Python code to solve them. @@ -114,16 +77,6 @@ IMPORTANT: Always use the dataset path variable '{dataset_path}' in your code wh return None, response_message.content def upload_dataset(code_interpreter: Sandbox, uploaded_file) -> str: - """ - Uploads the dataset to the E2B sandbox. - - Args: - code_interpreter: The E2B sandbox instance - uploaded_file: Streamlit uploaded file - - Returns: - str: Path where file was uploaded - """ dataset_path = f"./{uploaded_file.name}" try: @@ -139,11 +92,22 @@ def main(): st.title("AI Data Visualization Agent") st.write("Upload your dataset and ask questions about it!") - # Sidebar for API keys and model selection + # Initialize session state variables + if 'together_api_key' not in st.session_state: + st.session_state.together_api_key = '' + if 'e2b_api_key' not in st.session_state: + st.session_state.e2b_api_key = '' + if 'model_name' not in st.session_state: + st.session_state.model_name = '' + with st.sidebar: st.header("API Keys and Model Configuration") - st.session_state.together_api_key = st.text_input("Enter Together API Key", type="password") - st.session_state.e2b_api_key = st.text_input("Enter E2B API Key", type="password") + st.session_state.together_api_key = st.sidebar.text_input("Together AI API Key", type="password") + st.sidebar.info("💡 Everyone gets a free $1 credit by Together AI - AI Acceleration Cloud platform") + st.sidebar.markdown("[Get Together AI API Key](https://api.together.ai/signin)") + + st.session_state.e2b_api_key = st.sidebar.text_input("Enter E2B API Key", type="password") + st.sidebar.markdown("[Get E2B API Key](https://e2b.dev/docs/legacy/getting-started/api-key)") # Add model selection dropdown model_options = { @@ -152,12 +116,12 @@ def main(): "Qwen 2.5 7B": "Qwen/Qwen2.5-7B-Instruct-Turbo", "Meta-Llama 3.3 70B": "meta-llama/Llama-3.3-70B-Instruct-Turbo" } - selected_model = st.selectbox( + st.session_state.model_name = st.selectbox( "Select Model", options=list(model_options.keys()), index=0 # Default to first option ) - st.session_state.model_name = model_options[selected_model] + st.session_state.model_name = model_options[st.session_state.model_name] uploaded_file = st.file_uploader("Choose a CSV file", type="csv") diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt b/ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt index 8019c78..2ec4fbe 100644 --- a/ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt +++ b/ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt @@ -1,9 +1,7 @@ -together>=0.2.8 -e2b>=0.12.0 -python-dotenv -Pillow +together==1.3.10 +e2b-code-interpreter==1.0.3 +e2b==1.0.5 +Pillow==10.4.0 streamlit pandas matplotlib -plotly -seaborn>=0.12.0 \ No newline at end of file From 744780314b4cf4c3956a748aa5a396641160425f Mon Sep 17 00:00:00 2001 From: Madhu Shantan Date: Sun, 5 Jan 2025 19:50:57 +0530 Subject: [PATCH 7/7] added demoo --- ai_agent_tutorials/ai_data_visualisation_agent/README.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/ai_agent_tutorials/ai_data_visualisation_agent/README.md b/ai_agent_tutorials/ai_data_visualisation_agent/README.md index dc69328..802d724 100644 --- a/ai_agent_tutorials/ai_data_visualisation_agent/README.md +++ b/ai_agent_tutorials/ai_data_visualisation_agent/README.md @@ -2,6 +2,10 @@ This Assistant is designed to help anyone create and visualize data using natural language commands, and it is built using Together AI and E2B Code Interpreter. User gets to upload a dataset and ask questions to the LLM to get the data visualized. This demo can be considered as a demo for the E2B Code Interpreter and Together AI, for anyone who's getting started with these libraries! +## Demo + +https://github.com/user-attachments/assets/d8414c37-5edd-4e4d-a7b1-b9ab500bd8cd + ## Features - 🎨 Natural language-driven visualization creation