From f2fef44fef8e7572eab34d8cd5763d59d8ff7142 Mon Sep 17 00:00:00 2001 From: blob42 Date: Fri, 24 Feb 2023 00:43:05 +0100 Subject: [PATCH] docker wrapper tool for untrusted execution --- langchain/utilities/docker.py | 61 +++++++++++++++++++++++++++++++++ poetry.lock | 33 +++++++++++++++--- pyproject.toml | 4 ++- tests/unit_tests/test_docker.py | 25 ++++++++++++++ 4 files changed, 117 insertions(+), 6 deletions(-) create mode 100644 langchain/utilities/docker.py create mode 100644 tests/unit_tests/test_docker.py diff --git a/langchain/utilities/docker.py b/langchain/utilities/docker.py new file mode 100644 index 00000000..9d44416b --- /dev/null +++ b/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 ? diff --git a/poetry.lock b/poetry.lock index 18aa9dd0..700c7cfb 100644 --- a/poetry.lock +++ b/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" diff --git a/pyproject.toml b/pyproject.toml index 1387ef07..f2abcc7c 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/tests/unit_tests/test_docker.py b/tests/unit_tests/test_docker.py new file mode 100644 index 00000000..bd64a9ab --- /dev/null +++ b/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") + + +