cli: Make base image configurable in build command (#831)

* cli: Make base image configurable in build command

* Lint

* Fix

libs/cli/langgraph_cli/cli.py

@@ -15,7 +15,6 @@ from langgraph_cli.constants import DEFAULT_CONFIG, DEFAULT_PORT
 from langgraph_cli.docker import DockerCapabilities
 from langgraph_cli.exec import Runner, subp_exec
 from langgraph_cli.progress import Progress
-from langgraph_cli.util import clean_empty_lines
 
 OPT_DOCKER_COMPOSE = click.option(
     "--docker-compose",
@@ -351,14 +350,20 @@ def logs(
     \b
     """,
 )
+@click.option(
+    "--base-image",
+    hidden=True,
+)
 @cli.command(help="Build langgraph API server docker image")
 @log_command
 def build(
     config: pathlib.Path,
     platform: Optional[str],
+    base_image: Optional[str],
     pull: bool,
     tag: str,
 ):
+    base_image = base_image or "langchain/langgraph-api"
     with open(config) as f:
         config_json = langgraph_cli.config.validate_config(json.load(f))
     with Runner() as runner:
@@ -370,7 +375,7 @@ def build(
                 subp_exec(
                     "docker",
                     "pull",
-                    f"langchain/langgraph-api:{config_json['python_version']}",
+                    f"{base_image}:{config_json['python_version']}",
                 )
             )
         # apply options
@@ -383,7 +388,7 @@ def build(
         if platform:
             args.extend(["--platform", platform])
         # apply config
-        stdin = langgraph_cli.config.config_to_docker(config, config_json)
+        stdin = langgraph_cli.config.config_to_docker(config, config_json, base_image)
         # run docker build
         runner.run(
             subp_exec(
@@ -392,109 +397,6 @@ def build(
         )
 
 
-@cli.group(help="Export langgraph compose files")
-def export():
-    pass
-
-
-@click.option(
-    "--output",
-    "-o",
-    help="Output path to write the docker compose file to",
-    type=click.Path(
-        exists=False,
-        file_okay=True,
-        dir_okay=False,
-        resolve_path=True,
-        path_type=pathlib.Path,
-    ),
-    required=True,
-)
-@OPT_CONFIG
-@OPT_PORT
-@OPT_WATCH
-@OPT_LANGGRAPH_API_PATH
-@export.command(name="compose", help="Export docker compose file")
-@log_command
-def export_compose(
-    output: pathlib.Path,
-    config: pathlib.Path,
-    port: int,
-    watch: bool,
-    langgraph_api_path: Optional[pathlib.Path],
-):
-    with Runner() as runner:
-        capabilities = langgraph_cli.docker.check_capabilities(runner)
-        _, stdin = prepare(
-            runner,
-            capabilities=capabilities,
-            config_path=config,
-            docker_compose=None,
-            pull=False,
-            watch=watch,
-            langgraph_api_path=langgraph_api_path,
-            port=port,
-            verbose=False,
-        )
-
-    with open(output, "w") as f:
-        f.write(clean_empty_lines(stdin))
-
-
-@click.option(
-    "--output",
-    "-o",
-    help="Output path (directory) to write the helm chart to",
-    type=click.Path(
-        exists=False,
-        file_okay=False,
-        dir_okay=True,
-        resolve_path=True,
-        path_type=pathlib.Path,
-    ),
-    required=True,
-)
-@OPT_PORT
-@OPT_DOCKER_COMPOSE
-@OPT_CONFIG
-@export.command(
-    name="helm",
-    help="Build and export a helm chart to deploy to a Kubernetes cluster",
-    hidden=True,
-)
-@log_command
-def export_helm(
-    output: pathlib.Path,
-    config: pathlib.Path,
-    docker_compose: Optional[pathlib.Path],
-    port: int,
-):
-    with open(config) as f:
-        config_json = langgraph_cli.config.validate_config(json.load(f))
-
-    with Runner() as runner:
-        # check docker available
-        capabilities = langgraph_cli.docker.check_capabilities(runner)
-        # prepare args
-        stdin = langgraph_cli.docker.compose(capabilities, port=port)
-        args = [
-            "convert",
-            "--chart",
-            "-o",
-            str(output),
-            "-v",
-        ]
-        # apply options
-        if docker_compose:
-            args.extend(["-f", str(docker_compose)])
-
-        args.extend(["-f", "-"])  # stdin
-        # apply config
-        stdin += langgraph_cli.config.config_to_compose(config, config_json)
-        # run kompose convert
-        runner.run(subp_exec("kompose", *args, input=stdin))
-
-
 def prepare_args_and_stdin(
     *,
     capabilities: DockerCapabilities,
@@ -524,7 +426,11 @@ def prepare_args_and_stdin(
     args.extend(["-f", "-"])  # stdin
     # apply config
     stdin += langgraph_cli.config.config_to_compose(
-        config_path, config, watch=watch, langgraph_api_path=langgraph_api_path
+        config_path,
+        config,
+        watch=watch,
+        langgraph_api_path=langgraph_api_path,
+        base_image="langchain/langgraph-api",
     )
     return args, stdin
 

libs/cli/langgraph_cli/config.py

@@ -191,7 +191,7 @@ def _update_graph_paths(
             config["graphs"][graph_id] = f"{module_str}:{attr_str}"
 
 
-def config_to_docker(config_path: pathlib.Path, config: Config):
+def config_to_docker(config_path: pathlib.Path, config: Config, base_image: str):
     # configure pip
     pip_install = "pip install -c /api/constraints.txt"
     if config.get("pip_config_file"):
@@ -237,7 +237,7 @@ EOF"""
         for fullpath, relpath in local_deps.real_pkgs.items()
     )
 
-    return f"""FROM langchain/langgraph-api:{config['python_version']}
+    return f"""FROM {base_image}:{config['python_version']}
 
 {os.linesep.join(config["dockerfile_lines"])}
 
@@ -261,6 +261,7 @@ ENV LANGSERVE_GRAPHS='{json.dumps(config["graphs"])}'
 def config_to_compose(
     config_path: pathlib.Path,
     config: Config,
+    base_image: str,
     watch: bool = False,
     langgraph_api_path: Optional[pathlib.Path] = None,
 ):
@@ -301,6 +302,6 @@ def config_to_compose(
         build:
             context: .
             dockerfile_inline: |
-{textwrap.indent(config_to_docker(config_path, config), "                ")}
+{textwrap.indent(config_to_docker(config_path, config, base_image), "                ")}
         {watch_str}
 """

libs/cli/tests/unit_tests/test_config.py

@@ -66,7 +66,9 @@ def test_validate_config():
 def test_config_to_docker_simple():
     graphs = {"agent": "./agent.py:graph"}
     actual_docker_stdin = config_to_docker(
-        PATH_TO_CONFIG, validate_config({"dependencies": ["."], "graphs": graphs})
+        PATH_TO_CONFIG,
+        validate_config({"dependencies": ["."], "graphs": graphs}),
+        "langchain/langgraph-api",
     )
     expected_docker_stdin = """\
 FROM langchain/langgraph-api:3.11
@@ -96,6 +98,7 @@ def test_config_to_docker_pipconfig():
                 "pip_config_file": "pipconfig.txt",
             }
         ),
+        "langchain/langgraph-api",
     )
     expected_docker_stdin = """\
 FROM langchain/langgraph-api:3.11
@@ -122,13 +125,16 @@ def test_config_to_docker_invalid_inputs():
         config_to_docker(
             PATH_TO_CONFIG,
             validate_config({"dependencies": ["./missing"], "graphs": graphs}),
+            "langchain/langgraph-api",
         )
 
     # test missing local module
     with pytest.raises(FileNotFoundError):
         graphs = {"agent": "./missing_agent.py:graph"}
         config_to_docker(
-            PATH_TO_CONFIG, validate_config({"dependencies": ["."], "graphs": graphs})
+            PATH_TO_CONFIG,
+            validate_config({"dependencies": ["."], "graphs": graphs}),
+            "langchain/langgraph-api",
         )
 
 
@@ -142,9 +148,10 @@ def test_config_to_docker_local_deps():
                 "graphs": graphs,
             }
         ),
+        "langchain/langgraph-api-custom",
     )
     expected_docker_stdin = """\
-FROM langchain/langgraph-api:3.11
+FROM langchain/langgraph-api-custom:3.11
 ADD ./graphs /deps/__outer_graphs/src
 COPY <<EOF /deps/__outer_graphs/pyproject.toml
 [project]
@@ -177,6 +184,7 @@ dependencies = ["langchain"]"""
                 "graphs": graphs,
             }
         ),
+        "langchain/langgraph-api",
     )
     os.remove(pyproject_path)
     expected_docker_stdin = """FROM langchain/langgraph-api:3.11
@@ -200,6 +208,7 @@ def test_config_to_docker_end_to_end():
                 "dockerfile_lines": ["ARG meow", "ARG foo"],
             }
         ),
+        "langchain/langgraph-api",
     )
     expected_docker_stdin = """FROM langchain/langgraph-api:3.12
 ARG meow
@@ -242,7 +251,9 @@ def test_config_to_compose_simple_config():
                 WORKDIR /deps/__outer_unit_tests/unit_tests
         """
     actual_compose_stdin = config_to_compose(
-        PATH_TO_CONFIG, validate_config({"dependencies": ["."], "graphs": graphs})
+        PATH_TO_CONFIG,
+        validate_config({"dependencies": ["."], "graphs": graphs}),
+        "langchain/langgraph-api",
     )
     assert clean_empty_lines(actual_compose_stdin) == expected_compose_stdin
 
@@ -255,7 +266,7 @@ def test_config_to_compose_env_vars():
         build:
             context: .
             dockerfile_inline: |
-                FROM langchain/langgraph-api:3.11
+                FROM langchain/langgraph-api-custom:3.11
                 ADD . /deps/__outer_unit_tests/unit_tests
                 COPY <<EOF /deps/__outer_unit_tests/pyproject.toml
                 [project]
@@ -278,6 +289,7 @@ def test_config_to_compose_env_vars():
                 "env": {"OPENAI_API_KEY": openai_api_key},
             }
         ),
+        "langchain/langgraph-api-custom",
     )
     assert clean_empty_lines(actual_compose_stdin) == expected_compose_stdin
 
@@ -306,6 +318,7 @@ def test_config_to_compose_env_file():
     actual_compose_stdin = config_to_compose(
         PATH_TO_CONFIG,
         validate_config({"dependencies": ["."], "graphs": graphs, "env": ".env"}),
+        "langchain/langgraph-api",
     )
     assert clean_empty_lines(actual_compose_stdin) == expected_compose_stdin
 
@@ -345,6 +358,7 @@ def test_config_to_compose_watch():
     actual_compose_stdin = config_to_compose(
         PATH_TO_CONFIG,
         validate_config({"dependencies": ["."], "graphs": graphs}),
+        "langchain/langgraph-api",
         watch=True,
     )
     assert clean_empty_lines(actual_compose_stdin) == expected_compose_stdin
@@ -389,6 +403,7 @@ def test_config_to_compose_end_to_end():
     actual_compose_stdin = config_to_compose(
         PATH_TO_CONFIG,
         validate_config({"dependencies": ["."], "graphs": graphs, "env": ".env"}),
+        "langchain/langgraph-api",
         watch=True,
         langgraph_api_path="path/to/langgraph/api",
     )