From 436902e5a34246d9c74b808d5195e1d0e3d989b0 Mon Sep 17 00:00:00 2001
From: Tat Dat Duong <david@duong.cz>
Date: Fri, 11 Apr 2025 01:04:26 +0200
Subject: [PATCH] Consolidate node_version and python_version fix

---
 libs/cli/langgraph_cli/cli.py    | 19 +++----------------
 libs/cli/langgraph_cli/config.py | 24 ++++++++++++++++++------
 2 files changed, 21 insertions(+), 22 deletions(-)

diff --git a/libs/cli/langgraph_cli/cli.py b/libs/cli/langgraph_cli/cli.py
index 140def614..07a71c8b9 100644
--- a/libs/cli/langgraph_cli/cli.py
+++ b/libs/cli/langgraph_cli/cli.py
@@ -280,11 +280,7 @@ def _build(
             subp_exec(
                 "docker",
                 "pull",
-                (
-                    f"{base_image}:{config_json['node_version']}"
-                    if config_json.get("node_version")
-                    else f"{base_image}:{config_json['python_version']}"
-                ),
+                langgraph_cli.config.docker_tag(config_json, base_image),
                 verbose=True,
             )
         )
@@ -709,11 +705,7 @@ def prepare_args_and_stdin(
         config_path,
         config,
         watch=watch,
-        base_image=(
-            "langchain/langgraphjs-api"
-            if config.get("node_version") and not config.get("python_version")
-            else "langchain/langgraph-api"
-        ),
+        base_image=langgraph_cli.config.default_base_image(config),
     )
     return args, stdin
 
@@ -740,12 +732,7 @@ def prepare(
             subp_exec(
                 "docker",
                 "pull",
-                (
-                    f"langchain/langgraphjs-api:{config_json['node_version']}"
-                    if config_json.get("node_version")
-                    and not config_json.get("python_version")
-                    else f"langchain/langgraph-api:{config_json['python_version']}"
-                ),
+                langgraph_cli.config.docker_tag(config_json),
                 verbose=verbose,
             )
         )
diff --git a/libs/cli/langgraph_cli/config.py b/libs/cli/langgraph_cli/config.py
index 878254fd5..d26065011 100644
--- a/libs/cli/langgraph_cli/config.py
+++ b/libs/cli/langgraph_cli/config.py
@@ -1234,14 +1234,24 @@ def node_config_to_docker(
     return os.linesep.join(docker_file_contents), {}
 
 
+def default_base_image(config: Config) -> str:
+    if config.get("node_version") and not config.get("python_version"):
+        return "langchain/langgraphjs-api"
+    return "langchain/langgraph-api"
+
+
+def docker_tag(config: Config, base_image: Optional[str] = None) -> str:
+    base_image = base_image or default_base_image(config)
+
+    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']}"
+
+
 def config_to_docker(
     config_path: pathlib.Path, config: Config, base_image: Optional[str] = None
 ) -> tuple[str, dict[str, str]]:
-    base_image = base_image or (
-        "langchain/langgraphjs-api"
-        if config.get("node_version") and not config.get("python_version")
-        else "langchain/langgraph-api"
-    )
+    base_image = base_image or default_base_image(config)
 
     if config.get("node_version") and not config.get("python_version"):
         return node_config_to_docker(config_path, config, base_image)
@@ -1252,9 +1262,11 @@ def config_to_docker(
 def config_to_compose(
     config_path: pathlib.Path,
     config: Config,
-    base_image: str,
+    base_image: Optional[str] = None,
     watch: bool = False,
 ) -> str:
+    base_image = base_image or default_base_image(config)
+
     env_vars = config["env"].items() if isinstance(config["env"], dict) else {}
     env_vars_str = "\n".join(f'            {k}: "{v}"' for k, v in env_vars)
     env_file_str = (