Add support for specifying a custom base image in docker commands (#4500)

build & dockerfile commands

can be specified via CLI  > langgraph.json

libs/cli/langgraph_cli/cli.py

@@ -438,8 +438,17 @@ tests
     ),
     is_flag=True,
 )
+@click.option(
+    "--base-image",
+    help="Base image to use for the LangGraph API server. Defaults to langchain/langgraph-api or langchain/langgraphjs-api",
+)
 @log_command
-def dockerfile(save_path: str, config: pathlib.Path, add_docker_compose: bool) -> None:
+def dockerfile(
+    save_path: str,
+    config: pathlib.Path,
+    add_docker_compose: bool,
+    base_image: Optional[str] = None,
+) -> None:
     save_path = pathlib.Path(save_path).absolute()
     secho(f"🔍 Validating configuration at path: {config}", fg="yellow")
     config_json = langgraph_cli.config.validate_config_file(config)
@@ -449,7 +458,7 @@ def dockerfile(save_path: str, config: pathlib.Path, add_docker_compose: bool) -
     dockerfile, additional_contexts = langgraph_cli.config.config_to_docker(
         config,
         config_json,
-        None,
+        base_image=base_image,
     )
     with open(str(save_path), "w", encoding="utf-8") as f:
         f.write(dockerfile)
@@ -709,7 +718,10 @@ def prepare_args_and_stdin(
     debugger_port: Optional[int] = None,
     debugger_base_url: Optional[str] = None,
     postgres_uri: Optional[str] = None,
+    # Like "my-tag" (if you already built it locally)
     image: Optional[str] = None,
+    # Like "langchain/langgraphjs-api" or "langchain/langgraph-api
+    base_image: Optional[str] = None,
 ) -> Tuple[List[str], str]:
     assert config_path.exists(), f"Config file not found: {config_path}"
     # prepare args

libs/cli/langgraph_cli/config.py

@@ -362,6 +362,11 @@ class Config(TypedDict, total=False):
     """Optional. Internal use only.
     """
 
+    base_image: Optional[str]
+    """Optional. Base image to use for the LangGraph API server.
+    
+    Defaults to langchain/langgraph-api or langchain/langgraphjs-api."""
+
     pip_config_file: Optional[str]
     """Optional. Path to a pip config file (e.g., "/etc/pip.conf" or "pip.ini") for controlling
     package installation (custom indices, credentials, etc.).
@@ -517,6 +522,7 @@ def validate_config(config: Config) -> Config:
         "python_version": python_version,
         "pip_config_file": config.get("pip_config_file"),
         "_INTERNAL_docker_tag": config.get("_INTERNAL_docker_tag"),
+        "base_image": config.get("base_image"),
         "dependencies": config.get("dependencies", []),
         "dockerfile_lines": config.get("dockerfile_lines", []),
         "graphs": config.get("graphs", {}),
@@ -1199,9 +1205,12 @@ ADD {relpath} /deps/{name}
                 "# -- End of JS dependencies install --",
             ]
         )
-
+    if "/langgraph-server" in base_image:
+        image_str = f"{base_image}-py{docker_tag}"
+    else:
+        image_str = f"{base_image}:{docker_tag}"
     docker_file_contents = [
-        f"FROM {base_image}:{docker_tag}",
+        f"FROM {image_str}",
         "",
         os.linesep.join(config["dockerfile_lines"]),
         "",
@@ -1285,6 +1294,8 @@ def node_config_to_docker(
 
 
 def default_base_image(config: Config) -> str:
+    if config.get("base_image"):
+        return config["base_image"]
     if config.get("node_version") and not config.get("python_version"):
         return "langchain/langgraphjs-api"
     return "langchain/langgraph-api"
@@ -1298,6 +1309,9 @@ def docker_tag(
     if config.get("_INTERNAL_docker_tag"):
         return f"{base_image}:{config['_INTERNAL_docker_tag']}"
 
+    if "/langgraph-server" in base_image:
+        return f"{base_image}-py{config['python_version']}"
+
     if config.get("node_version") and not config.get("python_version"):
         return f"{base_image}:{config['node_version']}"
     return f"{base_image}:{config['python_version']}"

libs/cli/pyproject.toml

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

libs/cli/schemas/schema.json

@@ -51,6 +51,17 @@
               ],
               "description": "Optional. Custom authentication config, including the path to your Python auth logic and\nthe OpenAPI security definitions it uses.\n"
             },
+            "base_image": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "Optional. Base image to use for the LangGraph API server.\n\nDefaults to langchain/langgraph-api or langchain/langgraphjs-api.\n"
+            },
             "checkpointer": {
               "anyOf": [
                 {
@@ -178,6 +189,17 @@
               ],
               "description": "Optional. Custom authentication config, including the path to your Python auth logic and\nthe OpenAPI security definitions it uses.\n"
             },
+            "base_image": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "Optional. Base image to use for the LangGraph API server.\n\nDefaults to langchain/langgraph-api or langchain/langgraphjs-api.\n"
+            },
             "checkpointer": {
               "anyOf": [
                 {

libs/cli/schemas/schema.v0.json

@@ -51,6 +51,17 @@
               ],
               "description": "Optional. Custom authentication config, including the path to your Python auth logic and\nthe OpenAPI security definitions it uses.\n"
             },
+            "base_image": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "Optional. Base image to use for the LangGraph API server.\n\nDefaults to langchain/langgraph-api or langchain/langgraphjs-api.\n"
+            },
             "checkpointer": {
               "anyOf": [
                 {
@@ -178,6 +189,17 @@
               ],
               "description": "Optional. Custom authentication config, including the path to your Python auth logic and\nthe OpenAPI security definitions it uses.\n"
             },
+            "base_image": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "Optional. Base image to use for the LangGraph API server.\n\nDefaults to langchain/langgraph-api or langchain/langgraphjs-api.\n"
+            },
             "checkpointer": {
               "anyOf": [
                 {

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

@@ -1,5 +1,6 @@
 import json
 import pathlib
+import re
 import shutil
 import tempfile
 import textwrap
@@ -347,6 +348,35 @@ def test_dockerfile_command_new_style_config() -> None:
         assert save_path.exists()
 
 
+def test_dockerfile_command_with_base_image() -> None:
+    """Test the 'dockerfile' command with a base image."""
+    runner = CliRunner()
+    config_content = {
+        "python_version": "3.11",
+        "graphs": {"agent": "agent.py:graph"},
+        "dependencies": ["."],
+        "base_image": "langchain/langgraph-server:0.2",
+    }
+    with temporary_config_folder(config_content) as temp_dir:
+        save_path = temp_dir / "Dockerfile"
+        agent_path = temp_dir / "agent.py"
+        agent_path.parent.mkdir(parents=True, exist_ok=True)
+        agent_path.touch()
+
+        result = runner.invoke(
+            cli,
+            ["dockerfile", str(save_path), "--config", str(temp_dir / "config.json")],
+        )
+
+        assert result.exit_code == 0, result.output
+        assert "✅ Created: Dockerfile" in result.output
+
+        assert save_path.exists()
+        with open(save_path) as f:
+            dockerfile = f.read()
+            assert re.match("FROM langchain/langgraph-server:0.2-py3.*", dockerfile)
+
+
 def test_dockerfile_command_with_docker_compose() -> None:
     """Test the 'dockerfile' command with Docker Compose configuration."""
     runner = CliRunner()

libs/cli/tests/unit_tests/test_config.py

@@ -30,6 +30,7 @@ def test_validate_config():
     actual_config = validate_config(expected_config)
     expected_config = {
         "_INTERNAL_docker_tag": None,
+        "base_image": None,
         "python_version": "3.11",
         "node_version": None,
         "pip_config_file": None,
@@ -49,6 +50,7 @@ def test_validate_config():
     env = ".env"
     expected_config = {
         "_INTERNAL_docker_tag": None,
+        "base_image": None,
         "python_version": "3.12",
         "node_version": None,
         "pip_config_file": "pipconfig.txt",