new issues

ai_agent_tutorials/ai_data_visualisation_agent/.gitignore

@@ -0,0 +1 @@
+.env

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
+import base64
-from typing import Optional, Union, List
+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
+# Function to execute Python code and generate plots
-def code_interpret(code: str) -> str:
+def execute_code(code: str, df):
-    """
-    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")
+        # Define locals with necessary imports and the DataFrame
-        # Convert list output to string if needed
+        local_env = {
-        stdout = execution.logs.stdout
+            'pd': pd,
-        if isinstance(stdout, list):
+            'df': df,
-            return '\n'.join(map(str, stdout))
+            'plt': plt,
-        return stdout if stdout else ""
+            'sns': sns  # if seaborn is needed
-    except Exception as e:
+        }
-        return f"Error executing code: {str(e)}"
+        # Execute the code in the local environment
+        exec(code, globals(), local_env)
         
-# Function to communicate with LLM
+        # Check if a plot was generated
-def chat_with_llm(user_message, file_path, columns):
+        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:
+        st.error(f"Error executing code: {e}")
+        return None
+
+# 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
+        # Execute the code and generate the plot
-        python_code = python_code.replace(
+        base64_image = execute_code(python_code, df)
-            "df['approx_cost(for two people)'] = df['approx_cost(for two people)'].str.replace(',', '')",
+        return response_message, base64_image
-            "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
+        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:
+if uploaded_file:
-    # Preprocess and save the uploaded file to a temporary file
+    if not together_api_key:
-    temp_path, columns, df = preprocess_and_save(uploaded_file)
+        st.warning("Please provide the Together AI API key.")
-    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.")
+        # 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.")

ai_agent_tutorials/ai_data_visualisation_agent/requirements.txt

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

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