diff --git a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/langgraph_workflow.py b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/langgraph_workflow.py index 83bd043fed..657bccadb7 100644 --- a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/langgraph_workflow.py +++ b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/langgraph_workflow.py @@ -193,6 +193,16 @@ def convert_to_chat_response_chunk(value: LanggraphWrapperOutput) -> ChatRespons return ChatResponseChunk.from_string(text) +def split_graph_path(graph: str) -> tuple[str, str]: + """Split a graph reference into its module path and graph name.""" + module_path, separator, name = graph.rpartition(":") + if not separator or not module_path or not name: + raise ValueError(f"Graph definition path '{graph}' must contain a non-empty module path and graph name " + "separated by a colon (e.g., '/path/to/module.py:graph_name').") + + return module_path, name + + @register_function(config_type=LanggraphWrapperConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN]) async def register(config: LanggraphWrapperConfig, b: Builder): @@ -227,14 +237,7 @@ async def register(config: LanggraphWrapperConfig, b: Builder): f"Env '{config.env}' is not a valid type. At the moment, we only support strings and dictionaries.") # Now process the graph. - # Check that config.graph contains exactly one colon - if config.graph.count(":") != 1: - raise ValueError( - f"Graph definition path '{config.graph}' must contain exactly one colon to split module and name " - f"(e.g., '/path/to/module.py:graph_name'). Found {config.graph.count(':')}.") - - # Split the graph path into module and name - module_path, name = config.graph.rsplit(":", 1) + module_path, name = split_graph_path(config.graph) unique_module_name = f"langgraph_workflow_{uuid.uuid4().hex[:8]}" diff --git a/packages/nvidia_nat_langchain/tests/test_langgraph_workflow.py b/packages/nvidia_nat_langchain/tests/test_langgraph_workflow.py index 4e8b8642d9..758c5cb835 100644 --- a/packages/nvidia_nat_langchain/tests/test_langgraph_workflow.py +++ b/packages/nvidia_nat_langchain/tests/test_langgraph_workflow.py @@ -28,6 +28,27 @@ from nat.plugins.langchain.langgraph_workflow import LanggraphWrapperFunction from nat.plugins.langchain.langgraph_workflow import LanggraphWrapperInput from nat.plugins.langchain.langgraph_workflow import LanggraphWrapperOutput +from nat.plugins.langchain.langgraph_workflow import split_graph_path + + +@pytest.mark.parametrize( + ("graph", "expected"), + [ + ("/path/to/module.py:graph_name", ("/path/to/module.py", "graph_name")), + (r"C:\\path\\to\\module.py:graph", (r"C:\\path\\to\\module.py", "graph")), + ("relative.py:graph", ("relative.py", "graph")), + ], +) +def test_split_graph_path(graph, expected): + """Graph references split at the final colon, including Windows drive paths.""" + assert split_graph_path(graph) == expected + + +@pytest.mark.parametrize("graph", ["module.py", ":graph", "module.py:", ""]) +def test_split_graph_path_rejects_missing_parts(graph): + """Graph references require both a module path and graph name.""" + with pytest.raises(ValueError, match="non-empty module path"): + split_graph_path(graph) class TestConvertChatRequest: