feat(sandbox): adding docker sandbox (#1517)

* feat(sandbox): adding docker sandbox * feat(sandbox): add serializer and deserializer * feat(dockerSandbox): complete implementation of docker sandbox * fix(sandbox): spell mistakes * fix(sandbox): fix readme command * fix(langchan): poetry lock file update * Update pandasai/sandbox/sandbox.py Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com> * fix(docker): make docker not use network * fix(ruff): errors in code formatting * feat(sandbox): add notebook for docker sandbox * added notebook and documentation for sandbox * feat(Sandbox): typo and update documentation * fix(docker): create docker from command line * Update extensions/sandbox/docker/pandasai_docker/docker_sandbox.py Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com> * Update extensions/sandbox/docker/pandasai_docker/docker_sandbox.py Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com> --------- Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com> Co-authored-by: gdcsinaptik <[email protected]>
sinaptik-ai · Jan 28, 2025 · d4a1d0e · d4a1d0e
1 parent 56bdd8c
commit d4a1d0e
Show file tree

Hide file tree

Showing 17 changed files with 2,980 additions and 7 deletions.
diff --git a/docs/mint.json b/docs/mint.json
@@ -3,7 +3,7 @@
     "logo": {
       "light": "/logo/logo.png",
       "dark": "/logo/logo.png",
-      "href": "https://pandas-ai.com"
+      "href": "https://getpanda.ai"
     },
     "favicon": "/favicon.svg",
     "colors": {
@@ -72,9 +72,14 @@
         "pages": ["v3/ai-dashboards", "v3/share-dataframes", "v3/permission-management"],
         "version": "v3"
       },
+      {
+        "group": "Advanced Usage",
+        "pages": ["v3/agent"],
+        "version": "v3"
+      },
       {
         "group": "Backwards Compatibility",
-        "pages": ["v3/agent", "v3/smart-dataframes", "v3/smart-datalakes"],
+        "pages": ["v3/smart-dataframes", "v3/smart-datalakes"],
         "version": "v3"
       },
       {

diff --git a/docs/v3/agent.mdx b/docs/v3/agent.mdx
@@ -1,18 +1,18 @@
 ---
-title: 'Agent'
-description: 'Add few-shot learning to your PandaAI agent'
+title: "Agent"
+description: "Add few-shot learning to your PandaAI agent"
 ---
 
 <Note title="Beta Notice">
-Release v3 is currently in beta. This documentation reflects the features and functionality in progress and may change before the final release.
+  Release v3 is currently in beta. This documentation reflects the features and
+  functionality in progress and may change before the final release.
 </Note>
 
 You can train PandaAI to understand your data better and to improve its performance. Training is as easy as calling the `train` method on the `Agent`.
 
-
 ## Prerequisites
 
-Before you start training PandaAI, you need to set your PandaAI API key. 
+Before you start training PandaAI, you need to set your PandaAI API key.
 You can generate your API key by signing up at [https://app.pandabi.ai](https://app.pandabi.ai).
 
 ```python
@@ -134,6 +134,60 @@ print(response)
 # The model will use the information provided in the training to generate a response
 ```
 
+## Using the Sandbox Environment
+
+To enhance security and protect against malicious code through prompt injection, PandaAI provides a sandbox environment for code execution. The sandbox runs your code in an isolated Docker container, ensuring that potentially harmful operations are contained.
+
+### Installation
+
+Before using the sandbox, you need to install Docker on your machine and ensure it is running.
+
+First, install the sandbox package:
+
+```bash
+pip install pandasai-docker
+```
+
+### Basic Usage
+
+Here's how to use the sandbox with your PandaAI agent:
+
+```python
+from pandasai import Agent
+from pandasai_docker import DockerSandbox
+
+# Initialize the sandbox
+sandbox = DockerSandbox()
+sandbox.start()
+
+# Create an agent with the sandbox
+df = pai.read_csv("data.csv")
+agent = Agent([df], sandbox=sandbox)
+
+# Chat with the agent - code will run in the sandbox
+response = agent.chat("Calculate the average sales")
+
+# Don't forget to stop the sandbox when done
+sandbox.stop()
+```
+
+### Customizing the Sandbox
+
+You can customize the sandbox environment by specifying a custom name and Dockerfile:
+
+```python
+sandbox = DockerSandbox(
+    "custom-sandbox-name",
+    "/path/to/custom/Dockerfile"
+)
+```
+
+<Note>
+  The sandbox works offline and provides an additional layer of security for
+  code execution. It's particularly useful when working with untrusted data or
+  when you need to ensure that code execution is isolated from your main system.
+</Note>
+
 ## Troubleshooting
 
 In some cases, you might get an error like this: `No vector store provided. Please provide a vector store to train the agent`. It means no API key has been generated to use the `BambooVectorStore`.
@@ -153,9 +207,11 @@ vector_store = BambooVectorStor(api_key="YOUR_PANDABI_API_KEY")
 # Instantiate the agent with the custom vector store
 agent = Agent(connector, config={...} vectorstore=vector_store)
 ```
+
 ## Custom Head
 
 In some cases, you might want to provide custom data samples to the conversational agent to improve its understanding and responses. For example, you might want to:
+
 - Provide better examples that represent your data patterns
 - Avoid sharing sensitive information
 - Guide the agent with specific data scenarios

diff --git a/examples/docker_sandbox.ipynb b/examples/docker_sandbox.ipynb
@@ -0,0 +1,123 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Execute code in a sandbox\n",
+    "\n",
+    "To enhance security and protect yourself from malicious code through prompt injection, \n",
+    "we make it possible to run code in a sandbox environment.\n",
+    "This notebook explains how to do it."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Install the package\n",
+    "\n",
+    "First of all you need to install the python package. You can use pip to install it"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install pandasai-docker"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Execute the code in the sandbox with the agent\n",
+    "\n",
+    "Please keep in mind the sandbox works offline. \n",
+    "Once you have installed the package, you can start the sandbox with the following code.\n",
+    "For the purpose of this example, we are going to use bambooLLM as the LLM chosen for the agent."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandasai as pai\n",
+    "from pandasai import Agent\n",
+    "from pandasai_docker import DockerSandbox\n",
+    "\n",
+    "pai.set_api_key(\"YOUR_API_KEY\")\n",
+    "\n",
+    "# initialize the sandbox\n",
+    "sandbox = DockerSandbox()\n",
+    "sandbox.start()\n",
+    "\n",
+    "# read a csv as df\n",
+    "df = pai.read_csv(\"./data/heart.csv\")\n",
+    "\n",
+    "# pass the csv and the sandbox to the agent\n",
+    "agent = Agent([df], memory_size=10, sandbox=sandbox)\n",
+    "\n",
+    "# Chat with the Agent\n",
+    "response = agent.chat(\"plot top five artists streams\")\n",
+    "\n",
+    "# stop the sandbox (docker container)\n",
+    "sandbox.stop()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Customize the sandbox\n",
+    "\n",
+    "You can decide the name and path of your sandbox by passing them as positional arguments in the DockerSandbox()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sandbox = DockerSandbox(\"pandaai-sandbox\", \"/path/to/Dockerfile\")\n",
+    "\n",
+    "# read a csv as df\n",
+    "df = pai.read_csv(\"./data/heart.csv\")\n",
+    "\n",
+    "# pass the csv and the sandbox to the agent\n",
+    "agent = Agent([df], memory_size=10, sandbox=sandbox)\n",
+    "\n",
+    "# Chat with the Agent\n",
+    "response = agent.chat(\"plot top five artists streams\")\n",
+    "\n",
+    "sandbox.stop()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}
diff --git a/extensions/sandbox/docker/README.md b/extensions/sandbox/docker/README.md
@@ -0,0 +1,9 @@
+# Docker Sandbox Extension for PandasAI
+
+## Installation
+
+You can install this extension using poetry:
+
+```bash
+poetry add pandasai-docker
+```
diff --git a/extensions/sandbox/docker/pandasai_docker/Dockerfile b/extensions/sandbox/docker/pandasai_docker/Dockerfile
@@ -0,0 +1,12 @@
+FROM python:3.9
+
+LABEL image_name="pandasai-sandbox"
+
+# Install required Python packages
+RUN pip install pandas numpy matplotlib
+
+# Set the working directory inside the container
+WORKDIR /app
+
+# Default command keeps the container running (useful for testing or debugging)
+CMD ["sleep", "infinity"]
diff --git a/extensions/sandbox/docker/pandasai_docker/__init__.py b/extensions/sandbox/docker/pandasai_docker/__init__.py
@@ -0,0 +1,3 @@
+from .docker_sandbox import DockerSandbox
+
+__all__ = ["DockerSandbox"]
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		from .docker_sandbox import DockerSandbox

		__all__ = ["DockerSandbox"]