new issues

This commit is contained in:
Madhu 2025-01-04 23:47:56 +05:30
parent 47abbdf46d
commit d3cfeaab2e
4 changed files with 296 additions and 103 deletions

View file

@ -0,0 +1 @@
.env

View file

@ -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.")
st.warning("Please upload a file.")

View file

@ -1,3 +1,4 @@
e2b-code-interpreter==1.0.3
togetherai==1.3.10
streamlit==1.41.1
python-dotenv
together
e2b
streamlit

View file

@ -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()