Add image arg to up command (#4385)

Using this argument, you can get more customization since you can do
`langgraph build` or directly `docker build` your image and then re-use
the `langgraph up --image my-image` and have it also spin up redis &
postgres for you.

Easier then writing your own compose file

libs/cli/langgraph_cli/cli.py

@@ -168,6 +168,13 @@ def cli():
 @OPT_DEBUGGER_BASE_URL
 @OPT_WATCH
 @OPT_POSTGRES_URI
+@click.option(
+    "--image",
+    type=str,
+    default=None,
+    help="Docker image to use for the langgraph-api service. If specified, skips building and uses this image directly."
+    " Useful if you want to test against an image already built using `langgraph build`.",
+)
 @click.option(
     "--wait",
     is_flag=True,
@@ -187,6 +194,7 @@ def up(
     debugger_port: Optional[int],
     debugger_base_url: Optional[str],
     postgres_uri: Optional[str],
+    image: Optional[str],
 ):
     click.secho("Starting LangGraph API server...", fg="green")
     click.secho(
@@ -207,6 +215,7 @@ For production use, requires a license key in env var LANGGRAPH_CLOUD_LICENSE_KE
             debugger_port=debugger_port,
             debugger_base_url=debugger_base_url,
             postgres_uri=postgres_uri,
+            image=image,
         )
         # add up + options
         args.extend(["up", "--remove-orphans"])
@@ -692,6 +701,7 @@ def prepare_args_and_stdin(
     debugger_port: Optional[int] = None,
     debugger_base_url: Optional[str] = None,
     postgres_uri: Optional[str] = None,
+    image: Optional[str] = None,
 ) -> Tuple[List[str], str]:
     assert config_path.exists(), f"Config file not found: {config_path}"
     # prepare args
@@ -701,6 +711,7 @@ def prepare_args_and_stdin(
         debugger_port=debugger_port,
         debugger_base_url=debugger_base_url,
         postgres_uri=postgres_uri,
+        image=image,  # Pass image to compose YAML generator
     )
     args = [
         "--project-directory",
@@ -716,6 +727,7 @@ def prepare_args_and_stdin(
         config,
         watch=watch,
         base_image=langgraph_cli.config.default_base_image(config),
+        image=image,
     )
     return args, stdin
 
@@ -733,6 +745,7 @@ def prepare(
     debugger_port: Optional[int] = None,
     debugger_base_url: Optional[str] = None,
     postgres_uri: Optional[str] = None,
+    image: Optional[str] = None,
 ) -> Tuple[List[str], str]:
     """Prepare the arguments and stdin for running the LangGraph API server."""
     config_json = langgraph_cli.config.validate_config_file(config_path)
@@ -757,5 +770,6 @@ def prepare(
         debugger_port=debugger_port,
         debugger_base_url=debugger_base_url or f"http://127.0.0.1:{port}",
         postgres_uri=postgres_uri,
+        image=image,
     )
     return args, stdin

libs/cli/langgraph_cli/config.py

@@ -1288,6 +1288,7 @@ def config_to_compose(
     config_path: pathlib.Path,
     config: Config,
     base_image: Optional[str] = None,
+    image: Optional[str] = None,
     watch: bool = False,
 ) -> str:
     base_image = base_image or default_base_image(config)
@@ -1314,19 +1315,28 @@ def config_to_compose(
 """
     else:
         watch_str = ""
+    if image:
+        return f"""
+{textwrap.indent(env_vars_str, "            ")}
+        {env_file_str}
+        {watch_str}
+"""
 
-    dockerfile, additional_contexts = config_to_docker(config_path, config, base_image)
+    else:
+        dockerfile, additional_contexts = config_to_docker(
+            config_path, config, base_image
+        )
 
-    additional_contexts_str = "\n".join(
-        f"                - {name}: {path}"
-        for name, path in additional_contexts.items()
-    )
-    if additional_contexts_str:
-        additional_contexts_str = f"""
+        additional_contexts_str = "\n".join(
+            f"                - {name}: {path}"
+            for name, path in additional_contexts.items()
+        )
+        if additional_contexts_str:
+            additional_contexts_str = f"""
             additional_contexts:
 {additional_contexts_str}"""
 
-    return f"""
+        return f"""
 {textwrap.indent(env_vars_str, "            ")}
         {env_file_str}
         pull_policy: build

libs/cli/langgraph_cli/docker.py

@@ -143,6 +143,8 @@ def compose_as_dict(
     debugger_base_url: Optional[str] = None,
     # postgres://user:password@host:port/database?option=value
     postgres_uri: Optional[str] = None,
+    # If you are running against an already-built image, you can pass it here
+    image: Optional[str] = None,
 ) -> dict:
     """Create a docker compose file as a dictionary in YML style."""
     if postgres_uri is None:
@@ -211,6 +213,8 @@ def compose_as_dict(
             "POSTGRES_URI": postgres_uri,
         },
     }
+    if image:
+        services["langgraph-api"]["image"] = image
 
     # If Postgres is included, add it to the dependencies of langgraph-api
     if include_db:
@@ -244,6 +248,7 @@ def compose(
     debugger_base_url: Optional[str] = None,
     # postgres://user:password@host:port/database?option=value
     postgres_uri: Optional[str] = None,
+    image: Optional[str] = None,
 ) -> str:
     """Create a docker compose file as a string."""
     compose_content = compose_as_dict(
@@ -252,6 +257,7 @@ def compose(
         debugger_port=debugger_port,
         debugger_base_url=debugger_base_url,
         postgres_uri=postgres_uri,
+        image=image,
     )
     compose_str = dict_to_yaml(compose_content)
     return compose_str

libs/cli/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-cli"
-version = "0.2.6"
+version = "0.2.7"
 description = "CLI for interacting with LangGraph API"
 authors = []
 license = "MIT"

libs/cli/tests/unit_tests/cli/test_cli.py

@@ -160,6 +160,110 @@ services:
     assert clean_empty_lines(actual_stdin) == expected_stdin
 
 
+def test_prepare_args_and_stdin_with_image() -> None:
+    # this basically serves as an end-to-end test for using config and docker helpers
+    config_path = pathlib.Path(__file__).parent / "langgraph.json"
+    config = validate_config(
+        Config(dependencies=[".", "../../.."], graphs={"agent": "agent.py:graph"})
+    )
+    port = 8000
+    debugger_port = 8001
+    debugger_graph_url = f"http://127.0.0.1:{port}"
+
+    actual_args, actual_stdin = prepare_args_and_stdin(
+        capabilities=DEFAULT_DOCKER_CAPABILITIES,
+        config_path=config_path,
+        config=config,
+        docker_compose=pathlib.Path("custom-docker-compose.yml"),
+        port=port,
+        debugger_port=debugger_port,
+        debugger_base_url=debugger_graph_url,
+        watch=True,
+        image="my-cool-image",
+    )
+
+    expected_args = [
+        "--project-directory",
+        str(pathlib.Path(__file__).parent.absolute()),
+        "-f",
+        "custom-docker-compose.yml",
+        "-f",
+        "-",
+    ]
+    expected_stdin = f"""volumes:
+    langgraph-data:
+        driver: local
+services:
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
+    langgraph-postgres:
+        image: pgvector/pgvector:pg16
+        ports:
+            - "5433:5432"
+        environment:
+            POSTGRES_DB: postgres
+            POSTGRES_USER: postgres
+            POSTGRES_PASSWORD: postgres
+        command:
+            - postgres
+            - -c
+            - shared_preload_libraries=vector
+        volumes:
+            - langgraph-data:/var/lib/postgresql/data
+        healthcheck:
+            test: pg_isready -U postgres
+            start_period: 10s
+            timeout: 1s
+            retries: 5
+            interval: 60s
+            start_interval: 1s
+    langgraph-debugger:
+        image: langchain/langgraph-debugger
+        restart: on-failure
+        depends_on:
+            langgraph-postgres:
+                condition: service_healthy
+        ports:
+            - "{debugger_port}:3968"
+        environment:
+            VITE_STUDIO_LOCAL_GRAPH_URL: {debugger_graph_url}
+    langgraph-api:
+        ports:
+            - "8000:8000"
+        depends_on:
+            langgraph-redis:
+                condition: service_healthy
+            langgraph-postgres:
+                condition: service_healthy
+        environment:
+            REDIS_URI: redis://langgraph-redis:6379
+            POSTGRES_URI: {DEFAULT_POSTGRES_URI}
+        image: my-cool-image
+        healthcheck:
+            test: python /api/healthcheck.py
+            interval: 60s
+            start_interval: 1s
+            start_period: 10s
+        
+        
+        develop:
+            watch:
+                - path: langgraph.json
+                  action: rebuild
+                - path: .
+                  action: rebuild
+                - path: ../../..
+                  action: rebuild\
+"""
+    assert actual_args == expected_args
+    assert clean_empty_lines(actual_stdin) == expected_stdin
+
+
 def test_version_option() -> None:
     """Test the --version option of the CLI."""
     runner = CliRunner()