docker wrapper tool for untrusted execution

langchain/utilities/docker.py

@@ -0,0 +1,61 @@
+"""Wrapper for untrusted code exectuion on docker."""
+# TODO: Validation:
+#       - verify gVisor runtime (runsc) if available
+#       - pass arbitrary image names
+
+import docker
+from docker.client import DockerClient  # type: ignore
+from docker.errors import APIError, ContainerError
+
+from typing import Any, Dict
+from typing import Optional
+from pydantic import BaseModel, PrivateAttr, Extra, root_validator, validator
+
+
+class DockerWrapper(BaseModel, extra=Extra.forbid):
+    """Executes arbitrary payloads and returns the output."""
+
+    _docker_client: DockerClient = PrivateAttr()
+    image: Optional[str] = "alpine"
+
+    # use env by default when create docker client
+    from_env: Optional[bool] = True
+
+    def __init__(self, **kwargs):
+        """Initialize docker client."""
+        super().__init__(**kwargs)
+
+        if self.from_env:
+            self._docker_client = docker.from_env()
+
+    @property
+    def client(self) -> DockerClient:
+        """Docker client."""
+        return self._docker_client
+
+    @property
+    def info(self) -> Any:
+        """Prints docker `info`."""
+        return self._docker_client.info()
+
+    @root_validator()
+    def validate_all(cls, values: Dict) -> Dict:
+        """Validate environment."""
+        # print("root validator")
+        return values
+
+    def run(self, query: str, **kwargs: Any) -> str:
+        """Run arbitrary shell command inside a container.
+
+        Args:
+            **kwargs: Pass extra parameters to DockerClient.container.run.
+
+        """
+        try:
+            image = getattr(kwargs, "image", self.image)
+            return self._docker_client.containers.run(image,
+                                                      query,
+                                                      remove=True)
+        except ContainerError as e:
+            return f"STDERR: {e}"
+        # TODO: handle docker APIError ?

poetry.lock

@@ -1140,6 +1140,28 @@ idna = ["idna (>=2.1,<4.0)"]
 trio = ["trio (>=0.14,<0.23)"]
 wmi = ["wmi (>=1.5.1,<2.0.0)"]
 
+[[package]]
+name = "docker"
+version = "6.0.1"
+description = "A Python library for the Docker Engine API."
+category = "main"
+optional = true
+python-versions = ">=3.7"
+files = [
+    {file = "docker-6.0.1-py3-none-any.whl", hash = "sha256:dbcb3bd2fa80dca0788ed908218bf43972772009b881ed1e20dfc29a65e49782"},
+    {file = "docker-6.0.1.tar.gz", hash = "sha256:896c4282e5c7af5c45e8b683b0b0c33932974fe6e50fc6906a0a83616ab3da97"},
+]
+
+[package.dependencies]
+packaging = ">=14.0"
+pywin32 = {version = ">=304", markers = "sys_platform == \"win32\""}
+requests = ">=2.26.0"
+urllib3 = ">=1.26.0"
+websocket-client = ">=0.32.0"
+
+[package.extras]
+ssh = ["paramiko (>=2.4.3)"]
+
 [[package]]
 name = "docutils"
 version = "0.17.1"
@@ -4851,7 +4873,7 @@ files = [
 name = "pywin32"
 version = "305"
 description = "Python for Window Extensions"
-category = "dev"
+category = "main"
 optional = false
 python-versions = "*"
 files = [
@@ -6059,7 +6081,7 @@ files = [
 ]
 
 [package.dependencies]
-greenlet = {version = "!=0.4.17", markers = "python_version >= \"3\" and (platform_machine == \"aarch64\" or platform_machine == \"ppc64le\" or platform_machine == \"x86_64\" or platform_machine == \"amd64\" or platform_machine == \"AMD64\" or platform_machine == \"win32\" or platform_machine == \"WIN32\")"}
+greenlet = {version = "!=0.4.17", markers = "python_version >= \"3\" and platform_machine == \"aarch64\" or python_version >= \"3\" and platform_machine == \"ppc64le\" or python_version >= \"3\" and platform_machine == \"x86_64\" or python_version >= \"3\" and platform_machine == \"amd64\" or python_version >= \"3\" and platform_machine == \"AMD64\" or python_version >= \"3\" and platform_machine == \"win32\" or python_version >= \"3\" and platform_machine == \"WIN32\""}
 
 [package.extras]
 aiomysql = ["aiomysql", "greenlet (!=0.4.17)"]
@@ -7183,7 +7205,7 @@ files = [
 name = "websocket-client"
 version = "1.5.1"
 description = "WebSocket client for Python with low level API options"
-category = "dev"
+category = "main"
 optional = false
 python-versions = ">=3.7"
 files = [
@@ -7498,10 +7520,11 @@ docs = ["furo", "jaraco.packaging (>=9)", "jaraco.tidelift (>=1.4)", "rst.linker
 testing = ["flake8 (<5)", "func-timeout", "jaraco.functools", "jaraco.itertools", "more-itertools", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=1.3)", "pytest-flake8", "pytest-mypy (>=0.9.1)"]
 
 [extras]
-all = ["anthropic", "cohere", "openai", "nlpcloud", "huggingface_hub", "manifest-ml", "elasticsearch", "opensearch-py", "google-search-results", "faiss-cpu", "sentence-transformers", "transformers", "spacy", "nltk", "wikipedia", "beautifulsoup4", "tiktoken", "torch", "jinja2", "pinecone-client", "weaviate-client", "redis", "google-api-python-client", "wolframalpha", "qdrant-client", "tensorflow-text", "pypdf", "networkx", "nomic"]
+all = ["anthropic", "cohere", "openai", "nlpcloud", "huggingface_hub", "manifest-ml", "elasticsearch", "opensearch-py", "google-search-results", "faiss-cpu", "sentence-transformers", "transformers", "spacy", "nltk", "wikipedia", "beautifulsoup4", "tiktoken", "torch", "jinja2", "pinecone-client", "weaviate-client", "redis", "google-api-python-client", "wolframalpha", "qdrant-client", "tensorflow-text", "pypdf", "networkx", "nomic", "docker"]
+docker = ["docker"]
 llms = ["anthropic", "cohere", "openai", "nlpcloud", "huggingface_hub", "manifest-ml", "torch", "transformers"]
 
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.8.1,<4.0"
-content-hash = "449d9958004f9b0af5667b02f866313913f9bd9c939870898873c0e3198a9cb4"
+content-hash = "e817b6b0f985c4178f4cd1bc5bea92130e79092e5fff0d41a03a5dbcfe1047cd"

pyproject.toml

@@ -51,6 +51,7 @@ pypdf = {version = "^3.4.0", optional = true}
 networkx = {version="^2.6.3", optional = true}
 aleph-alpha-client = {version="^2.15.0", optional = true}
 deeplake = {version = "^3.2.9", optional = true}
+docker = {version = "^6.0.1", optional = true}
 
 [tool.poetry.group.docs.dependencies]
 autodoc_pydantic = "^1.8.0"
@@ -95,8 +96,9 @@ jupyter = "^1.0.0"
 playwright = "^1.28.0"
 
 [tool.poetry.extras]
+docker = ["docker"]
 llms = ["anthropic", "cohere", "openai", "nlpcloud", "huggingface_hub", "manifest-ml", "torch", "transformers"]
-all = ["anthropic", "cohere", "openai", "nlpcloud", "huggingface_hub", "manifest-ml", "elasticsearch", "opensearch-py", "google-search-results", "faiss-cpu", "sentence_transformers", "transformers", "spacy", "nltk", "wikipedia", "beautifulsoup4", "tiktoken", "torch", "jinja2", "pinecone-client", "weaviate-client", "redis", "google-api-python-client", "wolframalpha", "qdrant-client", "tensorflow-text", "pypdf", "networkx", "nomic"]
+all = ["anthropic", "cohere", "openai", "nlpcloud", "huggingface_hub", "manifest-ml", "elasticsearch", "opensearch-py", "google-search-results", "faiss-cpu", "sentence_transformers", "transformers", "spacy", "nltk", "wikipedia", "beautifulsoup4", "tiktoken", "torch", "jinja2", "pinecone-client", "weaviate-client", "redis", "google-api-python-client", "wolframalpha", "qdrant-client", "tensorflow-text", "pypdf", "networkx", "nomic", "docker"]
 
 [tool.ruff]
 select = [

tests/unit_tests/test_docker.py

@@ -0,0 +1,25 @@
+"""Test the docker wrapper utility."""
+
+import pytest
+from langchain.utilities.docker import DockerWrapper
+
+
+def test_command_default_image() -> None:
+    """Test running a command with the default alpine image."""
+    docker = DockerWrapper()
+    output = docker.run("cat /etc/os-release")
+    assert output.find(b"alpine")
+
+def test_inner_failing_command() -> None:
+    """Test inner command with non zero exit"""
+    docker = DockerWrapper()
+    output = docker.run("ls /inner-failing-command")
+    assert str(output).startswith("STDERR")
+
+def test_entrypoint_failure() -> None:
+    """Test inner command with non zero exit"""
+    docker = DockerWrapper()
+    output = docker.run("todo handle APIError")
+
+
+