forked from Archives/langchain
docker wrapper tool for untrusted execution
parent
f61858163d
commit
895f862662
@ -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 ?
|
@ -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")
|
||||
|
||||
|
||||
|
Loading…
Reference in New Issue