diff --git a/docs-agent-mcp/pipelines/code-pipeline.py b/docs-agent-mcp/pipelines/code-pipeline.py index c6c6e05..df5c3de 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, diff --git a/docs-agent-mcp/pipelines/utils/code_utils.py b/docs-agent-mcp/pipelines/utils/code_utils.py index a29b061..4a0d8f9 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 5e10e0f..142a710 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.""" @@ -203,6 +246,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"