Build an AI Data Visualization Agent

The demand for AI-powered data visualization tools is surging as businesses seek faster, more intuitive ways to understand their data. We can tap into this growing market by building our own AI-powered visualization tools that integrate seamlessly with existing data workflows.

In this tutorial, we'll build an AI Data Visualization Agent using Together AI's powerful language models and E2B's secure code execution environment. This agent will understand natural language queries about your data and automatically generate appropriate visualizations, making data exploration intuitive and efficient.

E2B is an open-source infrastructure that provides secure sandboxed environments for running AI-generated code. Using E2B's Python SDK, we can safely execute code generated by language models, making it perfect for creating an AI-powered data visualization tool that can handle various data formats and create multiple types of charts.

This tutorial can also be considered as a demo for the E2B Code Interpreter and Together AI, for anyone who's getting started with these libraries!

Don’t forget to share this tutorial on your social channels and tag Unwind AI (X, LinkedIn, Threads, Facebook) to support us!

What We’re Building

This Streamlit application creates an interactive data visualization assistant that can understand natural language queries and generate appropriate visualizations using LLMs.

Features:

1. Natural language query interface for data visualization

2. Support for multiple visualization types (line, bar, scatter, pie, bubble charts)

3. Automatic data preprocessing and cleaning

4. Secure code execution in E2B's sandboxed environment

5. Interactive Streamlit interface for easy data upload and visualization

6. Real-time visualization generation and display

7. Available Models:

   • Meta-Llama 3.1 405B

   • DeepSeek V3

   • Qwen 2.5 7B

   • Meta-Llama 3.3 70B

Prerequisites

Before we begin, make sure you have the following:

1. Python installed on your machine (version 3.10 or higher is recommended)

2. Your Together AI and E2B Code Interpreting API Key

3. A code editor of your choice (we recommend VS Code or PyCharm for their excellent Python support)

4. Basic familiarity with Python programming

Step-by-Step Instructions

Setting Up the Environment

First, let's get our development environment ready:

1. Clone the GitHub repository:

git clone https://github.com/Shubhamsaboo/awesome-llm-apps.git

🌟 Don't forget to star the opensource repo to show your support.

2. Go to the ai_data_visualisation_agent.py folder:

cd ai_agent_tutorials/ai_data_visualisation_agent

3. Install the required dependencies:

pip install -r requirements.txt

4. Get your API Key: You'll need to obtain Together AI and E2B Code Interpreting free API keys.

Creating the Streamlit App

Let’s create our app. Create a new file ai_data_visualisation_agent.py and add the following code:

1. Let's set up our imports:

import os
import json
import re
import sys
import io
import contextlib
import warnings
from typing import Optional, List, Any, Tuple
from PIL import Image
import streamlit as st
import pandas as pd
import base64
from io import BytesIO
from together import Together
from e2b_code_interpreter import Sandbox

2. Create code interpretation function:

def code_interpret(e2b_code_interpreter: Sandbox, code: str) -> Optional[List[Any]]:
    with st.spinner('Executing code in E2B sandbox...'):
        stdout_capture = io.StringIO()
        stderr_capture = io.StringIO()
        
        with contextlib.redirect_stdout(stdout_capture):
            exec = e2b_code_interpreter.run_code(code)
            
        return exec.results

3. Implement LLM interaction:

def chat_with_llm(e2b_code_interpreter: Sandbox, user_message: str, dataset_path: str):
    system_prompt = f"""You're a Python data scientist and visualization expert.
    Dataset at path '{dataset_path}'
    Analyze and answer with Python code."""
    
    client = Together(api_key=st.session_state.together_api_key)
    response = client.chat.completions.create(
        model=st.session_state.model_name,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message}
        ]
    )

4. Add dataset handling:

def upload_dataset(code_interpreter: Sandbox, uploaded_file) -> str:
    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

5. Set up Streamlit interface:

def main():
    st.title("AI Data Visualization Agent")
    
    with st.sidebar:
        st.header("API Keys and Model Configuration")
        st.session_state.together_api_key = st.sidebar.text_input(
            "Together AI API Key", 
            type="password"
        )
        st.session_state.e2b_api_key = st.sidebar.text_input(
            "E2B API Key", 
            type="password"
        )

6. Add model selection:

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"
}
st.session_state.model_name = st.selectbox(
    "Select Model",
    options=list(model_options.keys())
)

7. Implement file upload and preview:

uploaded_file = st.file_uploader("Choose a CSV file", type="csv")
if uploaded_file:
    df = pd.read_csv(uploaded_file)
    show_full = st.checkbox("Show full dataset")
    if show_full:
        st.dataframe(df)
    else:
        st.dataframe(df.head())

8. Add query processing:

query = st.text_area(
    "What would you like to know about your data?",
    "Can you compare the average cost between categories?"
)

if st.button("Analyze"):
    with Sandbox(api_key=st.session_state.e2b_api_key) as code_interpreter:
        dataset_path = upload_dataset(code_interpreter, uploaded_file)
        code_results, llm_response = chat_with_llm(
            code_interpreter, 
            query, 
            dataset_path
        )

9. Handle visualization display:

if code_results:
    for result in code_results:
        if hasattr(result, 'png'):
            png_data = base64.b64decode(result.png)
            image = Image.open(BytesIO(png_data))
            st.image(image)
        elif hasattr(result, 'figure'):
            st.pyplot(result.figure)
        elif hasattr(result, 'show'):
            st.plotly_chart(result)

10. Error handling and validation:

try:
    code_interpreter_results = code_interpret(
        e2b_code_interpreter, 
        python_code
    )
except Exception as error:
    st.error(f"Error executing code: {error}")
    return None

Running the App

With our code in place, it's time to launch the app.

• In your terminal, navigate to the project folder, and run the following command

streamlit run ai_data_visualisation_agent.py

• Streamlit will provide a local URL (typically http://localhost:8501).

Working Application Demo

Conclusion

Your AI data visualization agent is now ready. Put in your API keys, upload a dataset, and ask questions to get the data visualized.

To enhance the application further, consider:

• Add support for more file formats (Excel, JSON, SQL databases)

• Implement visualization templates for common query types

• Add the ability to export visualizations in various formats

• Add memory to remember user preferences and previous visualizations

Keep experimenting and refining to build smarter AI solutions!

We share hands-on tutorials like this 2-3 times a week, to help you stay ahead in the world of AI. If you're serious about leveling up your AI skills and staying ahead of the curve, subscribe now and be the first to access our latest tutorials.

Don’t forget to share this tutorial on your social channels and tag Unwind AI (X, LinkedIn, Threads, Facebook) to support us!