Fix failing tests and add helper functions for docker image

libs/langchain/langchain/utilities/docker_containers.py

@@ -52,11 +52,16 @@ class DockerImage:
     def __init__(self, name: str):
         """Note that it does not pull the image from the internet.
         It only represents a tag so it must exist on your system.
+        It throws ValueError if docker image by that name does not exist locally.
         """
-        self.name = name
-        # check if image exists
-        docker_client = get_docker_client()
-        if len(docker_client.images.list(name=name)) < 1:
+        splitted_name = name.split(":")
+        if len(splitted_name) == 1:
+            # by default, image has latest tag.
+            self.name = name + ":latest"
+        else:
+            self.name = name
+
+        if not self.exists(name):
             raise ValueError(
                 f"Invalid value: name={name} does not exist on your system."
                 "Use DockerImage.from_tag() to pull it."
@@ -65,6 +70,21 @@ def __init__(self, name: str):
     def __repr__(self) -> str:
         return f"DockerImage(name={self.name})"
 
+    @classmethod
+    def exists(cls, name: str) -> bool:
+        """Checks if the docker image exists"""
+        docker_client = get_docker_client()
+        return len(docker_client.images.list(name=name)) > 0
+
+    @classmethod
+    def remove(cls, name: str) -> None:
+        """WARNING: Removes image from the system, be cautious with this function.
+        It is irreversible operation!.
+        """
+        if cls.exists(name):
+            docker_client = get_docker_client()
+            docker_client.images.remove(name)
+
     @classmethod
     def from_tag(
         cls,
@@ -78,10 +98,13 @@ def from_tag(
         Example: repository = "python" tag = "3.9-slim"
         """
         docker_client = get_docker_client()
+        name = f"{repository}:{tag}"
+        if len(docker_client.images.list(name=name)) > 0:
+            return cls(name=name)
         docker_client.images.pull(
             repository=repository, tag=tag, auth_config=auth_config
         )
-        return cls(name=f"{repository}:{tag}")
+        return cls(name=name)
 
     @classmethod
     def from_dockerfile(
@@ -146,6 +169,13 @@ def __enter__(self) -> "DockerContainer":
         """Enters container context. It means that container is started and you can
         execute commands inside it.
         """
+        self.unsafe_start()
+        return self
+
+    def unsafe_start(self) -> None:
+        """Starts container without entering it.
+        Please prefer to use with DockerContainer statement.
+        """
         assert self._container is None, "You cannot re-entry container"
         # tty=True is required to keep container alive
         self._container = self._client.containers.run(
@@ -154,7 +184,6 @@ def __enter__(self) -> "DockerContainer":
             tty=True,
             **self._run_kwargs,
         )
-        return self
 
     def __exit__(
         self,
@@ -163,6 +192,7 @@ def __exit__(
         traceback: Optional[TracebackType],
     ) -> bool:
         """Cleanup container on exit."""
+        assert self._container is not None, "You cannot exit unstarted container."
         if exc_type is not None:
             # re-throw exception. try to stop container and remove it
             try:
@@ -171,10 +201,16 @@ def __exit__(
                 print("Failed to stop and remove container to cleanup exception.", e)
             return False
         else:
-            self._cleanup()
-            self._container = None
+            self.unsafe_exit()
             return True
 
+    def unsafe_exit(self):
+        """Cleanup container on exit. Please prefer to use `with` statement."""
+        if self._container is None:
+            return
+        self._cleanup()
+        self._container = None
+
     def spawn_run(
         self, command: Union[str, List[str]], **kwargs: Any
     ) -> Tuple[int, bytes]:

libs/langchain/tests/integration_tests/utilities/docker_test_data/Dockerfile

@@ -1,7 +1,5 @@
 # This is a test dockerfile that will be used to test the docker_containers.
-FROM python:3.11-slim
-RUN pip install cowsay
-# This runs cowsay with moo so that we can test that the image works
-# and additionally every other run command within container 
-# will execute but with moo prefix for the user input.
-ENTRYPOINT ["python3", "-m", "cowsay", "moo"]
+FROM python:3.11-alpine
+RUN pip install --no-cache-dir cowsay==6.0
+# This runs cowsay and it requires arguments like -t "hello world".
+ENTRYPOINT ["python3", "-m", "cowsay"]

libs/langchain/tests/integration_tests/utilities/test_docker_containers.py

@@ -2,6 +2,7 @@
 from typing import cast
 
 import pytest
+
 from langchain.utilities.docker_containers import (
     DockerContainer,
     DockerImage,
@@ -31,14 +32,11 @@ def run_container_cowsay(image: DockerImage) -> None:
     # by ENTRYPOINT defined in dockerfile.
     try:
         container = DockerContainer(image)
-        ret_code, log = container.spawn_run("I like langchain!")
+        ret_code, log = container.spawn_run('-t "I like langchain!"')
         assert ret_code == 0
-        assert (
-            log.find(b"moo I like langchain") >= 0
-        ), "Cowsay should say same words with moo"
+        assert log.find(b"I like langchain") >= 0, "Cowsay should say same words"
     finally:
-        docker_client = get_docker_client()
-        docker_client.images.remove(image.name)
+        DockerImage.remove(image)
 
 
 @pytest.mark.requires("docker")