Is there a way for me to enable my LLM agent to write and execute python code in a sandboxed environment?

[tituslhy]

Question Validation

• [X] I have searched both the documentation and discord for an answer.

Question

I built a ReAct SQL agent but want to confer it the ability to plot the results from its SQL statement. For now, this has to be defined as a custom tool which means that the code to plot has to be specified in detail (for e.g. https://www.linkedin.com/feed/update/urn:li:activity:7186333474256035840/). I was wondering if there's a way to allow for the LLM to write its own python code to plot the data table in anyway it likes (so it doesn't have to be a line graph), execute it and then show the results. Kind of like Vanna AI - I know that there's a Vanna AI integration with LlamaIndex but the team designing the integration specifically muted the plotting capability.

[dosubot[bot]]

Yes, you can enable your LLM agent to write and execute Python code in a sandboxed environment to plot the results from its SQL statements using the CodeInterpreterToolSpec class from LlamaIndex. This class allows the execution of arbitrary Python code, including importing libraries and running SQL queries, and can be used to generate plots.

LlamaIndex provides a secure environment by restricting imports and built-in functions to ensure safe execution. Here is an example of how you might use this class to execute Python code that includes plotting:

from llama_index.tools.code_interpreter.base import CodeInterpreterToolSpec

# Initialize the code interpreter tool
code_interpreter = CodeInterpreterToolSpec()

# Example code to execute: running a SQL query and plotting the results
code = """
import sqlite3
import matplotlib.pyplot as plt

# Connect to the database (replace with your database connection details)
conn = sqlite3.connect('example.db')
cursor = conn.cursor()

# Execute a SQL query
cursor.execute('SELECT column1, column2 FROM example_table')
data = cursor.fetchall()

# Close the connection
conn.close()

# Extract data for plotting
x, y = zip(*data)

# Plot the results
plt.plot(x, y)
plt.xlabel('Column 1')
plt.ylabel('Column 2')
plt.title('SQL Query Results')
plt.savefig('plot.png')
"""

# Execute the code
output = code_interpreter.code_interpreter(code)
print(output)

This code snippet connects to an SQLite database, executes a SQL query, fetches the results, and plots them using Matplotlib. The plot is saved as an image file (plot.png). The code_interpreter method will execute this code and return the standard output and error messages.

LlamaIndex ensures the safety of the executed code by restricting imports to a predefined set of allowed libraries (e.g., math, time, datetime, pandas, scipy, numpy, matplotlib, plotly, seaborn) and by verifying that the code does not contain any references to private or dunder methods, disallowed built-ins, or any imports outside the allowed set ^[1][2].

To continue talking to Dosu, mention @dosu.

[logan-markewich]

@tituslhy we recently released azure support for this https://llamahub.ai/l/tools/llama-index-tools-azure-code-interpreter?from=

been meaning to make an integration with e2b as well.

[tituslhy]

@logan-markewich Sorry can I check if there is a way to return the plotly figure in the AgentChatReponse class so that I can pass it to a frontend tool like streamlit/chainlit?

I've been trying some of @dosu's suggestions in this thread https://github.com/run-llama/llama_index/discussions/14262#discussioncomment-9835901. The suggestions don't work because the default response class is AgentChatResponse, but I see it's trying to do by editing the taskoutput and agent runner class.

[dosubot[bot]]

To modify the AgentChatResponse class to return a Plotly figure so that you can pass it to a frontend tool like Streamlit or Chainlit, you can add a method that generates and returns a Plotly figure. Here is an example of how you can do this:

import plotly.graph_objects as go
from typing import List, Dict, Any, Optional
from dataclasses import dataclass, field

@dataclass
class AgentChatResponse:
    """Agent chat response."""

    response: str = ""
    sources: List[ToolOutput] = field(default_factory=list)
    source_nodes: List[NodeWithScore] = field(default_factory=list)
    is_dummy_stream: bool = False
    metadata: Optional[Dict[str, Any]] = None

    def __post_init__(self) -> None:
        if self.sources and not self.source_nodes:
            for tool_output in self.sources:
                if isinstance(tool_output.raw_output, (Response, StreamingResponse)):
                    self.source_nodes.extend(tool_output.raw_output.source_nodes)

    def __str__(self) -> str:
        return self.response

    @property
    def response_gen(self) -> Generator[str, None, None]:
        """Used for fake streaming, i.e. with tool outputs."""
        if not self.is_dummy_stream:
            raise ValueError(
                "response_gen is only available for streaming responses. "
                "Set is_dummy_stream=True if you still want a generator."
            )

        for token in self.response.split(" "):
            yield token + " "
            time.sleep(0.1)

    async def async_response_gen(self) -> AsyncGenerator[str, None]:
        """Used for fake streaming, i.e. with tool outputs."""
        if not self.is_dummy_stream:
            raise ValueError(
                "response_gen is only available for streaming responses. "
                "Set is_dummy_stream=True if you still want a generator."
            )

        for token in self.response.split(" "):
            yield token + " "
            await asyncio.sleep(0.1)

    def generate_plotly_figure(self) -> go.Figure:
        """Generate a Plotly figure."""
        fig = go.Figure(data=go.Bar(y=[2, 3, 1]))
        fig.update_layout(title="Sample Plotly Figure")
        return fig

With this method, you can now generate a Plotly figure and pass it to a frontend tool like Streamlit or Chainlit. Here is an example of how you might use this in a Streamlit application:

import streamlit as st
from your_module import AgentChatResponse

# Create an instance of AgentChatResponse
response = AgentChatResponse(response="Sample response")

# Generate the Plotly figure
fig = response.generate_plotly_figure()

# Display the figure in Streamlit
st.plotly_chart(fig)

This will display the Plotly figure in your Streamlit application ^[1].

To continue talking to Dosu, mention @dosu.