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']}"