From aed881c58f9f0eb4f74f648be3f7f3557950fe27 Mon Sep 17 00:00:00 2001 From: Sayak Satpathi Date: Fri, 30 Oct 2026 12:40:32 +0530 Subject: [PATCH 1/4] fix(pipelines): don't duplicate the first node's decorator in the module header parse_python_ast set the module_header boundary at the first node's def/class line (nodes[0].lineno). When that node is decorated, its own chunk starts at the decorator line, so the decorator was captured in both the header chunk and the function chunk -- common for Kubeflow code using @dsl.component / @dsl.pipeline. Use the decorator line as the header boundary when the first node has decorators. Adds a regression test. Signed-off-by: Sayak Satpathi --- docs-agent-mcp/pipelines/utils/code_utils.py | 12 +++++++---- tests/test_code_utils.py | 22 ++++++++++++++++++++ 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/docs-agent-mcp/pipelines/utils/code_utils.py b/docs-agent-mcp/pipelines/utils/code_utils.py index a29b0610..4a0d8f98 100644 --- a/docs-agent-mcp/pipelines/utils/code_utils.py +++ b/docs-agent-mcp/pipelines/utils/code_utils.py @@ -123,10 +123,14 @@ def parse_python_ast(content: str, file_path: str = "") -> list[dict]: "file_type": "python", }] - # Extract module header (everything before first function/class) - first_node_line = nodes[0].lineno # 1-indexed - if first_node_line > 1: - header = "\n".join(lines[:first_node_line - 1]).strip() + # Extract module header (everything before the first function/class). Use + # the decorator line as the boundary when the first node is decorated, so + # the decorator is not duplicated in both the header and the function chunk. + first_start_line = nodes[0].lineno # 1-indexed + if nodes[0].decorator_list: + first_start_line = nodes[0].decorator_list[0].lineno + if first_start_line > 1: + header = "\n".join(lines[:first_start_line - 1]).strip() if header and len(header) >= 10: chunks.append({ "content": header, diff --git a/tests/test_code_utils.py b/tests/test_code_utils.py index 5e10e0f1..513669ae 100644 --- a/tests/test_code_utils.py +++ b/tests/test_code_utils.py @@ -203,6 +203,28 @@ async def run_pipeline(): assert chunks[1]["content"].startswith("@decorator") assert chunks[1]["file_type"] == "python" + def test_decorated_first_node_not_duplicated_in_header(self): + # Regression: the first node's decorator must belong to the function + # chunk only. Previously the module header (which stops at the def line, + # not the decorator line) also captured the decorator, duplicating it -- + # common for Kubeflow code using @dsl.component / @dsl.pipeline. + content = ( + "import os\n" + "import sys\n" + "\n" + "@dsl.component\n" + "def my_comp(a):\n" + " return a + 1\n" + ) + + chunks = parse_python_ast(content, "pipeline.py") + + header = next(c for c in chunks if c["resource_kind"] == "module_header") + func = next(c for c in chunks if c["resource_name"] == "my_comp") + assert "@dsl.component" not in header["content"] + assert header["content"] == "import os\nimport sys" + assert func["content"].startswith("@dsl.component") + def test_returns_module_chunk_when_no_top_level_defs(self): content = "PIPELINE_ROOT = '/tmp/pipeline'\nDEFAULT_TIMEOUT = 30\n" From f1595c8802580413c8d07a2e78a7ef05a9073e32 Mon Sep 17 00:00:00 2001 From: Sayak Satpathi Date: Thu, 5 Nov 2026 06:41:00 +0530 Subject: [PATCH 2/4] fix(pipelines): handle decorators in inline module header --- docs-agent-mcp/pipelines/code-pipeline.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/docs-agent-mcp/pipelines/code-pipeline.py b/docs-agent-mcp/pipelines/code-pipeline.py index c6c6e05d..df5c3de8 100644 --- a/docs-agent-mcp/pipelines/code-pipeline.py +++ b/docs-agent-mcp/pipelines/code-pipeline.py @@ -248,9 +248,11 @@ def parse_python_ast(content, file_path=""): "file_type": "python", }] - first_node_line = nodes[0].lineno - if first_node_line > 1: - header = "\n".join(lines[:first_node_line - 1]).strip() + first_start_line = nodes[0].lineno + if nodes[0].decorator_list: + first_start_line = nodes[0].decorator_list[0].lineno + if first_start_line > 1: + header = "\n".join(lines[:first_start_line - 1]).strip() if header and len(header) >= 10: chunks.append({ "content": header, From fc800b3e75ab61a7c99dd1cf519b5caf92c68701 Mon Sep 17 00:00:00 2001 From: Sayak Satpathi Date: Thu, 5 Nov 2026 06:41:05 +0530 Subject: [PATCH 3/4] test(pipelines): cover inline decorator header boundary --- tests/test_code_utils.py | 43 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/tests/test_code_utils.py b/tests/test_code_utils.py index 513669ae..997bc1dc 100644 --- a/tests/test_code_utils.py +++ b/tests/test_code_utils.py @@ -115,6 +115,49 @@ def test_compiled_code_pipeline_injects_required_secrets(tmp_path): assert "github-pat" in compiled +def test_inline_code_pipeline_does_not_duplicate_first_decorator(monkeypatch, tmp_path): + """The KFP component's inline parser must keep decorators out of the header.""" + module = load_code_pipeline_module() + source = ( + "import os\\n" + "import sys\\n" + "\\n" + "@dsl.component\\n" + "def my_comp(a):\\n" + " return a + 1\\n" + ) + input_path = tmp_path / "code.jsonl" + input_path.write_text( + json.dumps({ + "content": source, + "path": "pipeline.py", + "repo": "kubeflow/pipelines", + "file_name": "pipeline.py", + }) + "\\n" + ) + monkeypatch.setattr( + "requests.post", + lambda *args, **kwargs: FakeResponse([[0.1], [0.2]]), + ) + output_path = tmp_path / "embedded.jsonl" + + module.chunk_and_embed_code.python_func( + code_data=SimpleNamespace(path=str(input_path)), + chunk_size=1000, + chunk_overlap=100, + embeddings_service_url="http://embeddings.test/embed", + embedding_batch_size=8, + embedded_data=SimpleNamespace(path=str(output_path)), + ) + + chunks = [json.loads(line) for line in output_path.read_text().splitlines()] + header = next(chunk for chunk in chunks if chunk["resource_kind"] == "module_header") + function = next(chunk for chunk in chunks if chunk["resource_name"] == "my_comp") + assert header["content_text"] == "import os\\nimport sys" + assert "@dsl.component" not in header["content_text"] + assert function["content_text"].startswith("@dsl.component") + + class TestParseYamlDocuments: """Tests for Kubernetes YAML-aware chunking.""" From 77208fc6f26ccab79a8e62f8063d2024202a6706 Mon Sep 17 00:00:00 2001 From: Sayak Satpathi Date: Thu, 5 Nov 2026 06:41:19 +0530 Subject: [PATCH 4/4] test(pipelines): fix inline decorator regression fixture --- tests/test_code_utils.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/test_code_utils.py b/tests/test_code_utils.py index 997bc1dc..142a710e 100644 --- a/tests/test_code_utils.py +++ b/tests/test_code_utils.py @@ -119,12 +119,12 @@ def test_inline_code_pipeline_does_not_duplicate_first_decorator(monkeypatch, tm """The KFP component's inline parser must keep decorators out of the header.""" module = load_code_pipeline_module() source = ( - "import os\\n" - "import sys\\n" - "\\n" - "@dsl.component\\n" - "def my_comp(a):\\n" - " return a + 1\\n" + "import os\n" + "import sys\n" + "\n" + "@dsl.component\n" + "def my_comp(a):\n" + " return a + 1\n" ) input_path = tmp_path / "code.jsonl" input_path.write_text( @@ -133,7 +133,7 @@ def test_inline_code_pipeline_does_not_duplicate_first_decorator(monkeypatch, tm "path": "pipeline.py", "repo": "kubeflow/pipelines", "file_name": "pipeline.py", - }) + "\\n" + }) + "\n" ) monkeypatch.setattr( "requests.post", @@ -153,7 +153,7 @@ def test_inline_code_pipeline_does_not_duplicate_first_decorator(monkeypatch, tm chunks = [json.loads(line) for line in output_path.read_text().splitlines()] header = next(chunk for chunk in chunks if chunk["resource_kind"] == "module_header") function = next(chunk for chunk in chunks if chunk["resource_name"] == "my_comp") - assert header["content_text"] == "import os\\nimport sys" + assert header["content_text"] == "import os\nimport sys" assert "@dsl.component" not in header["content_text"] assert function["content_text"].startswith("@dsl.component")