new issues
This commit is contained in:
parent
47abbdf46d
commit
d3cfeaab2e
4 changed files with 296 additions and 103 deletions
1
ai_agent_tutorials/ai_data_visualisation_agent/.gitignore
vendored
Normal file
1
ai_agent_tutorials/ai_data_visualisation_agent/.gitignore
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
.env
|
||||
|
|
@ -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.")
|
||||
|
|
@ -1,3 +1,4 @@
|
|||
e2b-code-interpreter==1.0.3
|
||||
togetherai==1.3.10
|
||||
streamlit==1.41.1
|
||||
python-dotenv
|
||||
together
|
||||
e2b
|
||||
streamlit
|
||||
202
ai_agent_tutorials/ai_data_visualisation_agent/test.py
Normal file
202
ai_agent_tutorials/ai_data_visualisation_agent/test.py
Normal 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()
|
||||
Loading…
Reference in a new issue