successful working of code

This commit is contained in:
Madhu 2025-01-05 16:55:18 +05:30
parent 340aa821fb
commit 8cdb9bd663
4 changed files with 138 additions and 171 deletions

View file

@ -1 +0,0 @@
.env

View file

@ -25,9 +25,6 @@ def preprocess_and_save(file):
st.error("Unsupported file format. Please upload a CSV or Excel file.") st.error("Unsupported file format. Please upload a CSV or Excel file.")
return None, None, None 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 # Ensure string columns are properly quoted
for col in df.select_dtypes(include=['object']): for col in df.select_dtypes(include=['object']):
@ -53,10 +50,6 @@ def preprocess_and_save(file):
# Drop rows with all NaN values # Drop rows with all NaN values
df.dropna(how='all', inplace=True) 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 # Create a temporary file to save the preprocessed data
with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as temp_file: with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as temp_file:
temp_path = temp_file.name temp_path = temp_file.name
@ -137,7 +130,7 @@ def chat_with_llm(user_message, file_path, columns, df):
return response_message, None return response_message, None
# Set up Streamlit app # Set up Streamlit app
st.title("AI Data Scientist") st.title("AI Data Visualisation Agent")
# Sidebar for API keys and file upload # Sidebar for API keys and file upload
st.sidebar.header("API Keys") 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 (dynamic based on the uploaded file)
SYSTEM_PROMPT = """ 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}. 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. 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, Make sure to handle missing values and data type inconsistencies. When generating plots,

View file

@ -1,4 +1,9 @@
together>=0.2.8
e2b>=0.12.0
python-dotenv python-dotenv
together Pillow
e2b streamlit
streamlit pandas
matplotlib
plotly
seaborn>=0.12.0

View file

@ -1,209 +1,179 @@
import streamlit as st
import os import os
from dotenv import load_dotenv
import json import json
import re import re
from together import Together from typing import Optional, List, Any, Tuple
from e2b_code_interpreter import Sandbox from dotenv import load_dotenv
from typing import Optional, List, Any from PIL import Image
import tempfile import io
import streamlit as st
import pandas as pd
import base64 import base64
from io import BytesIO from io import BytesIO
from PIL import Image from PIL import Image
from together import Together
from e2b_code_interpreter import Sandbox
# Load environment variables # Load environment variables
load_dotenv() load_dotenv()
# Get API keys from environment variables # Regex pattern to extract code from LLM response
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) pattern = re.compile(r"```python\n(.*?)\n```", re.DOTALL)
# Function to extract Python code from LLM responses def code_interpret(e2b_code_interpreter: Sandbox, code: str) -> Optional[List[Any]]:
def match_code_blocks(llm_response): """
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) match = pattern.search(llm_response)
if match: if match:
code = match.group(1) code = match.group(1)
print("Extracted Python code:")
print(code)
return code return code
return "" return ""
# Function to interact with the LLM and execute code in the sandbox def chat_with_llm(e2b_code_interpreter: Sandbox, user_message: str, dataset_path: str) -> Tuple[Optional[List[Any]], str]:
def chat_with_llm(e2b_code_interpreter, user_message):
""" """
Interact with LLM and execute code in sandbox. Sends the user message to the LLM and executes the generated code.
Args: Args:
e2b_code_interpreter: The E2B Sandbox instance e2b_code_interpreter: The E2B sandbox instance
user_message: User's query string user_message: User's query message
dataset_path: Path to the uploaded dataset
Returns: 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 = [ messages = [
{"role": "system", "content": SYSTEM_PROMPT}, {"role": "system", "content": system_prompt},
{"role": "user", "content": enhanced_message}, {"role": "user", "content": user_message},
] ]
# Get response from Together AI with st.spinner('Getting response from together AI...'):
response = client.chat.completions.create( client = Together(api_key=st.session_state.together_api_key)
model=MODEL_NAME, response = client.chat.completions.create(
messages=messages, model=st.session_state.model_name,
) messages=messages,
)
# Extract the response message response_message = response.choices[0].message
response_message = response.choices[0].message.content python_code = match_code_blocks(response_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 the base64-encoded image data if python_code:
if code_interpreter_results and hasattr(code_interpreter_results[0], "png"): code_interpreter_results = code_interpret(e2b_code_interpreter, python_code)
return code_interpreter_results[0].png return code_interpreter_results, response_message.content
else: else:
return None st.warning(f"Failed to match any Python code in model's response")
else: return None, response_message.content
print(f"Failed to match any Python code in model's response: {response_message}")
return None
# Function to upload a dataset to the E2B Sandbox def upload_dataset(code_interpreter: Sandbox, uploaded_file) -> str:
def upload_dataset(code_interpreter: Sandbox, uploaded_file: Any) -> str:
""" """
Upload a dataset to the E2B Sandbox from Streamlit's uploaded file. Uploads the dataset to the E2B sandbox.
Args: Args:
code_interpreter: The E2B Sandbox instance code_interpreter: The E2B sandbox instance
uploaded_file: Streamlit's UploadedFile object uploaded_file: Streamlit uploaded file
Returns: 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: try:
# Create a temporary file to store the uploaded content code_interpreter.files.write(dataset_path, uploaded_file)
with tempfile.NamedTemporaryFile(delete=False, suffix='.csv') as tmp_file: return dataset_path
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: except Exception as error:
print("Error during file upload:", error) st.error(f"Error during file upload: {error}")
raise error raise error
def main(): def main():
"""Main function to run the Streamlit application.""" """Main Streamlit application."""
st.title("AI Data Visualization Agent") st.title("AI Data Visualization Assistant")
st.write("Upload your dataset and ask questions about it!") 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") uploaded_file = st.file_uploader("Choose a CSV file", type="csv")
# Text input for the query if uploaded_file is not None:
user_query = st.text_input("Enter your visualization query:") # Display dataset preview
df = pd.read_csv(uploaded_file)
# Process button st.write("Dataset Preview:")
if st.button("Generate Visualization") and uploaded_file is not None and user_query: st.dataframe(df.head())
try:
with Sandbox(api_key=E2B_API_KEY) as code_interpreter: # Query input
# Upload the dataset query = st.text_area("What would you like to know about your data?",
upload_dataset(code_interpreter, uploaded_file) "Can you compare the average cost for two people between different categories?")
# Get and execute the visualization code if st.button("Analyze"):
with st.spinner("Generating visualization..."): if not st.session_state.together_api_key or not st.session_state.e2b_api_key:
image_data = chat_with_llm(code_interpreter, user_query) st.error("Please enter both API keys in the sidebar.")
else:
# Display results with Sandbox(api_key=st.session_state.e2b_api_key) as code_interpreter:
if image_data: # Upload the dataset
# Decode the base64-encoded image data dataset_path = upload_dataset(code_interpreter, uploaded_file)
image_bytes = base64.b64decode(image_data)
image = Image.open(BytesIO(image_bytes))
# Display the image in Streamlit # Pass dataset_path to chat_with_llm
st.image(image, caption="Generated Visualization") code_results, llm_response = chat_with_llm(code_interpreter, query, dataset_path)
else:
st.error("No visualization generated")
except Exception as e: # Display LLM's text response
st.error(f"An error occurred: {e}") st.write("AI Response:")
elif not uploaded_file: st.write(llm_response)
st.warning("Please upload a dataset first")
elif not user_query: # Display results/visualizations
st.warning("Please enter a query") 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__": if __name__ == "__main__":
main() main()