Displaying plots in custom AI Agents

raresaxpo — Mon, 04 Aug 2025 08:44:51 GMT

I am trying to build an agent that uses custom functions running python code. However, I can't find a way to make the agent display plots or images. How can I tackle this?

Re: Displaying plots in custom AI Agents

Khaja_Zaffer — Mon, 04 Aug 2025 09:15:00 GMT

Hello @raresaxpo

Good day!

The content was asked from GPT to make a structured response but the whole process is not from GPT.

Step-by-Step Guide: Building an AI Data Visualization Agent

This guide provides a detailed walkthrough to help you build a fully functional AI data visualization agent using Together AI, E2B, and Streamlit. This agent will allow you to interact with your data using natural language and automatically generate various types of plots and charts.

What You Will Build

A Streamlit application that serves as an interactive data visualization assistant with the following features:

Natural Language Interface: Ask questions about your data in plain English.
Multiple Visualization Types: Automatically generates line, bar, scatter, pie, and bubble charts.
Secure Code Execution: Uses E2B's sandboxed environment to safely run AI-generated code.
Real-time Display: Visualizations are generated and displayed on the fly.
Interactive UI: A simple Streamlit interface for uploading data and interacting with the agent.

Prerequisites

Before you begin, ensure you have the following installed and configured:

Python: Version 3.10 or higher.
Code Editor: We recommend VS Code or PyCharm.
Together AI API Key: You can obtain a free key by signing up on the Together AI website.
E2B Code Interpreting API Key: You can get a free key by signing up on the E2B website.
Git: Installed on your machine to clone the repository.

Step 1: Setting Up the Environment

First, you need to set up your project environment by cloning the source code repository and installing the necessary libraries.

Open your terminal or command prompt.
Clone the GitHub repository: This command will download all the project files into a new folder named awesome-llm-apps.
Bash
```
git clone https://github.com/Shubhamsaboo/awesome-llm-apps.git
```

Navigate to the project directory:

Bash

cd starter_ai_agents/ai_data_visualisation_agent

Install the required Python dependencies: The requirements.txt file lists all the libraries needed for the application.
Bash
```
pip install -r requirements.txt
```

Step 2: Creating the Streamlit Application

Now you will create the main Python file for the application. The following code will handle all the logic, from user input to code execution and visualization display.

Create a new file named ai_data_visualisation_agent.py in your project directory.

Add the following code to the file. This code is broken down into sections for clarity.

Imports:

Python

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

Code Interpretation Function: This function executes the AI-generated Python code in the E2B sandbox.

Python

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

LLM Interaction Function: This function sends the user's query and the dataset context to the LLM (Large Language Model) to get a Python code response.

Python

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}
        ]
    )

Dataset Handling Function: This function uploads the user's file to the E2B sandbox.

Python

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

Streamlit main Function: This is the core of the application, which sets up the user interface, handles input, and calls the other functions.

Python

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

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

    # File Upload
    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())

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

    # Visualization Display and Error Handling
    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)

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

if __name__ == "__main__":
    main()

Step 3: Running the Application

Once the code is in place, you can launch the application from your terminal.

Open your terminal and make sure you are in the project directory (starter_ai_agents/ai_data_visualisation_agent).
Run the Streamlit command:
Bash
```
streamlit run ai_data_visualisation_agent.py
```
Streamlit will start a local web server and open the application in your default web browser. The URL is typically http://localhost:8501.

Step 4: Using the Application

Enter your API Keys: In the sidebar of the Streamlit application, enter your API keys for Together AI and E2B.
Upload a CSV file: Use the file uploader to select a CSV dataset from your computer.
Ask a question: Type a natural language query in the text area, such as "Show me a bar chart of sales by product category" or "What is the correlation between price and quantity?".
Click "Analyze": The agent will process your request, generate and execute the Python code, and display the resulting visualization directly in the application.

This process provides a comprehensive solution for creating an AI agent that can generate and display data visualizations, offering a powerful tool for data exploration and analysis.

topic Displaying plots in custom AI Agents in Generative AI

Displaying plots in custom AI Agents

Re: Displaying plots in custom AI Agents

Step-by-Step Guide: Building an AI Data Visualization Agent

What You Will Build

Prerequisites

Step 1: Setting Up the Environment

Step 2: Creating the Streamlit Application

Step 3: Running the Application

Step 4: Using the Application