Using LLM for Data Visualization

We are building a lightweight data visualization agent that turns raw CSV strings into rendered charts by generating matplotlib code through an LLM. It saves data analysts and backend engineers from writing repetitive plotting boilerplate for every ad-hoc request. I route the generation through Oxlo.ai because its flat per-request pricing keeps costs predictable even when I send long system prompts or iterate on phrasing.

What you'll need

This project is entirely Python. You do not need Jupyter, Node, or a frontend framework. You will send CSV text and natural language instructions to an Oxlo.ai model, then run the returned Python code locally. The OpenAI SDK is the only client library required because Oxlo.ai is fully compatible with the OpenAI API spec. You will not need to learn a new SDK or rewrite existing code.

• Python 3.10 or newer
• An Oxlo.ai API key from https://portal.oxlo.ai
• The OpenAI SDK: pip install openai
• matplotlib and pandas: pip install matplotlib pandas

I also recommend setting your API key as an environment variable so you do not hardcode secrets in the script.

Step 1: Configure the Oxlo.ai client

I initialize the OpenAI-compatible client with Oxlo.ai's base URL. Keeping the key in OXLO_API_KEY lets me share the script without leaking credentials.

import os
from openai import OpenAI

client = OpenAI(
    base_url="https://api.oxlo.ai/v1",
    api_key=os.getenv("OXLO_API_KEY")
)

Step 2: Define the system prompt

The system prompt constrains the model to return only executable Python inside a markdown fence. I also tell it that the CSV is already available as a StringIO object named data so the generated code does not try to open a file.

SYSTEM_PROMPT = (
    "You are a data visualization assistant. The user provides CSV text and a request.\n"
    "Write Python code that:\n"
    "1. Reads the CSV from a StringIO object named 'data' into a pandas DataFrame.\n"
    "2. Uses matplotlib to create the requested chart.\n"
    "3. Calls plt.savefig('output.png', bbox_inches='tight') and plt.close().\n"
    "Return ONLY the code inside a single ```python block. Do not add explanations."
)

Step 3: Extract code from the LLM response

Even with clear instructions, the model sometimes adds a brief intro sentence before the code block. A small regex helper grabs the first fenced Python snippet so execution never fails on markdown formatting.

import re

def extract_code(text: str) -> str:
    pattern = r"```python\s*(.*?)```"
    match = re.search(pattern, text, re.DOTALL)
    if match:
        return match.group(1).strip()
    return text.strip()

Step 4: Build the visualization generator

This helper formats the user request, sends it to Oxlo.ai using the deepseek-v3.2 model, and returns cleaned code. I use deepseek-v3.2 because it handles coding and reasoning tasks efficiently, and it is available on Oxlo.ai's free tier if I want to experiment before upgrading.

def generate_plot_code(csv_text: str, instruction: str) -> str:
    user_content = f"CSV data:\n{csv_text}\n\nRequest: {instruction}"

    response = client.chat.completions.create(
        model="deepseek-v3.2",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_content},
        ],
    )

    raw = response.choices[0].message.content
    return extract_code(raw)

Step 5: Execute the generated code safely

I run the generated code in a restricted namespace preloaded with pandas, matplotlib, and StringIO. This avoids polluting the global scope and prevents the script from accessing unintended variables.

import io
import pandas as pd
import matplotlib.pyplot as plt

def run_viz_code(code: str, csv_text: str):
    namespace = {
        "pd": pd,
        "plt": plt,
        "io": io,
        "data": io.StringIO(csv_text),
    }
    exec(code, namespace)
    print("Chart saved to output.png")

Step 6: Wire the agent together

The visualize entrypoint chains generation and execution. The main guard at the bottom provides a sample dataset so you can test the agent immediately after filling in your API key.

def visualize(csv_text: str, instruction: str):
    code = generate_plot_code(csv_text, instruction)
    print("Generated code:\n", code, "\n")
    run_viz_code(code, csv_text)

if __name__ == "__main__":
    sample_csv = """Date,Product,Revenue
2024-01-01,Widget,1200
2024-01-01,Gadget,950
2024-02-01,Widget,1500
2024-02-01,Gadget,1100
2024-03-01,Widget,1350
2024-03-01,Gadget,1600"""

    visualize(
        sample_csv,
        "Create a grouped bar chart of revenue by product for each month."
    )

Run it

Save all the previous blocks into a single file named viz_agent.py, set your OXLO_API_KEY, and run it from the terminal. The agent sends the sample CSV and request to Oxlo.ai, receives Python code, and executes it to produce output.png.

$ export OXLO_API_KEY="sk-..."
$ python viz_agent.py
Generated code:
 import pandas as pd
 import matplotlib.pyplot as plt
 
 df = pd.read_csv(data)
 monthly = df.groupby(['Date', 'Product'])['Revenue'].sum().unstack()
 monthly.plot(kind='bar')
 plt.title('Revenue by Product')
 plt.tight_layout()
 plt.savefig('output.png', bbox_inches='tight')
 plt.close()
 
Chart saved to output.png

Opening output.png shows a grouped bar chart with months on the x-axis and revenue on the y-axis, one series per product. The exact code returned by the model may vary slightly between runs. If the chart is missing a legend, you can add 'include a legend' to the instruction string and rerun. Because Oxlo.ai charges a flat rate per request, iterating on the prompt or sending larger CSV snippets does not change the cost.

Next steps

Two concrete ways to extend this agent. First, add a lightweight web interface with Streamlit so stakeholders can paste CSVs and type requests without touching Python. Second, store successful generations in a local SQLite cache keyed by a hash of the CSV schema and instruction. When a similar request arrives, you return the cached code instantly and skip the API call entirely.

For production use, you should also wrap the exec call in a try-except block and validate that no unsafe builtins are injected. The restricted namespace already limits exposure, but defense in depth is worth the extra lines. Both extensions sit cleanly on top of the Oxlo.ai backend. If you are currently using a token-based provider for code generation, moving this workload to Oxlo.ai can cut costs significantly for long-context prompts. You can compare plans at https://oxlo.ai/pricing.