From 382d9d49187720c18da68c134a6cada7166b39da Mon Sep 17 00:00:00 2001 From: Andrew Tritt Date: Tue, 1 Sep 2026 17:17:15 -0700 Subject: [PATCH] Add JSON output for reconstruction pipeline provenance --- src/dsagt/mcp/registry_tools.py | 2 +- src/dsagt/provenance.py | 41 ++++++++++- tests/test_pipeline.py | 116 ++++++++++++++++++++++++++++++++ 3 files changed, 156 insertions(+), 3 deletions(-) diff --git a/src/dsagt/mcp/registry_tools.py b/src/dsagt/mcp/registry_tools.py index baff6617..b7e7820b 100644 --- a/src/dsagt/mcp/registry_tools.py +++ b/src/dsagt/mcp/registry_tools.py @@ -524,7 +524,7 @@ def _registry_tools_and_handlers( "properties": { "format": { "type": "string", - "enum": ["bash", "snakemake"], + "enum": ["bash", "snakemake", "json"], "default": "bash", }, }, diff --git a/src/dsagt/provenance.py b/src/dsagt/provenance.py index 079658a9..6df8860c 100644 --- a/src/dsagt/provenance.py +++ b/src/dsagt/provenance.py @@ -577,6 +577,37 @@ def render_snakemake(records: list[dict], deps: dict[int, list[int]]) -> str: return "\n".join(lines) +def compute_terminal_outputs(records: list[dict]) -> list[str]: + """Derive terminal outputs: output files consumed as input by no record. + + Order follows first appearance across ``records`` (already sorted by + execution start time), with duplicates removed. + """ + all_inputs: set[str] = set() + for record in records: + all_inputs.update(record["execution"].get("input_files", [])) + + terminal: list[str] = [] + seen: set[str] = set() + for record in records: + for f in record["execution"].get("output_files", []): + if f not in all_inputs and f not in seen: + terminal.append(f) + seen.add(f) + + return terminal + + +def render_json(records: list[dict], deps: dict[int, list[int]]) -> str: + """Render the pipeline as structured JSON: records, dependency graph, and terminal outputs.""" + payload = { + "records": records, + "dependency_graph": deps, + "terminal_outputs": compute_terminal_outputs(records), + } + return json.dumps(payload, indent=2) + + def _shell_quote(s: str) -> str: """Quote a string for shell if it contains special characters.""" if not s: @@ -601,19 +632,25 @@ def reconstruct_pipeline( session_id : str, optional Filter records to a specific session. fmt : str - Output format: "bash" or "snakemake". + Output format: "bash", "snakemake", or "json". Returns ------- str - The rendered pipeline script. + The rendered pipeline script, or (for "json") a JSON document. """ records = load_pipeline_records(trace_dir, session_id) if not records: + if fmt == "json": + return json.dumps( + {"records": [], "dependency_graph": {}, "terminal_outputs": []} + ) return f"# No execution records found{' for session ' + session_id if session_id else ''}\n" deps = build_dependency_graph(records) if fmt == "snakemake": return render_snakemake(records, deps) + if fmt == "json": + return render_json(records, deps) return render_bash(records, deps) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 79aaee8d..b9e9ded8 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -8,9 +8,11 @@ from dsagt.provenance import ( build_dependency_graph, + compute_terminal_outputs, load_pipeline_records, reconstruct_pipeline, render_bash, + render_json, render_snakemake, ) @@ -269,6 +271,97 @@ def test_multi_step(self): assert '"contigs.fa"' in workflow +# --------------------------------------------------------------------------- +# compute_terminal_outputs +# --------------------------------------------------------------------------- + + +class TestComputeTerminalOutputs: + + def test_diamond_dependency(self): + """A produces two files; B and C each consume one; D combines them into + a final output that nothing downstream consumes.""" + records = [ + _make_record("a", ["a"], output_files=["x.fq", "y.fq"]), + _make_record("b", ["b"], input_files=["x.fq"], output_files=["bx.txt"]), + _make_record("c", ["c"], input_files=["y.fq"], output_files=["cy.txt"]), + _make_record( + "d", + ["d"], + input_files=["bx.txt", "cy.txt"], + output_files=["final.txt"], + ), + ] + assert compute_terminal_outputs(records) == ["final.txt"] + + def test_independent_leaves(self): + """Several unrelated single-step branches: every output is terminal.""" + records = [ + _make_record("a", ["a"], output_files=["a.out"]), + _make_record("b", ["b"], output_files=["b.out"]), + _make_record("c", ["c"], output_files=["c.out"]), + ] + assert compute_terminal_outputs(records) == ["a.out", "b.out", "c.out"] + + def test_linear_pipeline_only_last_output_terminal(self): + records = [ + _make_record("a", ["a"], output_files=["x.txt"]), + _make_record("b", ["b"], input_files=["x.txt"], output_files=["y.txt"]), + ] + assert compute_terminal_outputs(records) == ["y.txt"] + + def test_no_outputs(self): + records = [_make_record("a", ["a"])] + assert compute_terminal_outputs(records) == [] + + def test_duplicate_output_deduplicated(self): + """The same terminal file produced twice only appears once.""" + records = [ + _make_record("a", ["a"], output_files=["x.txt"], record_id="r1"), + _make_record("b", ["b"], output_files=["x.txt"], record_id="r2"), + ] + assert compute_terminal_outputs(records) == ["x.txt"] + + +# --------------------------------------------------------------------------- +# render_json +# --------------------------------------------------------------------------- + + +class TestRenderJson: + + def test_structure(self): + records = [ + _make_record( + "fastp", ["fastp"], input_files=["raw.fq"], output_files=["clean.fq"] + ), + ] + deps = build_dependency_graph(records) + result = json.loads(render_json(records, deps)) + + assert result["records"] == records + assert result["dependency_graph"] == {"0": []} + assert result["terminal_outputs"] == ["clean.fq"] + + def test_diamond_terminal_outputs(self): + records = [ + _make_record("a", ["a"], output_files=["x.fq", "y.fq"]), + _make_record("b", ["b"], input_files=["x.fq"], output_files=["bx.txt"]), + _make_record("c", ["c"], input_files=["y.fq"], output_files=["cy.txt"]), + _make_record( + "d", + ["d"], + input_files=["bx.txt", "cy.txt"], + output_files=["final.txt"], + ), + ] + deps = build_dependency_graph(records) + result = json.loads(render_json(records, deps)) + + assert result["dependency_graph"] == {"0": [], "1": [0], "2": [0], "3": [1, 2]} + assert result["terminal_outputs"] == ["final.txt"] + + # --------------------------------------------------------------------------- # reconstruct_pipeline (end-to-end) # --------------------------------------------------------------------------- @@ -300,11 +393,34 @@ def test_snakemake_format(self, tmp_path): assert "rule fastp_1:" in workflow + def test_json_format(self, tmp_path): + _write_record( + tmp_path, + _make_record( + "fastp", + ["fastp"], + input_files=["raw.fq"], + output_files=["clean.fq"], + record_id="r1", + ), + ) + result = json.loads(reconstruct_pipeline(tmp_path, fmt="json")) + + assert len(result["records"]) == 1 + assert result["records"][0]["code_name"] == "fastp" + assert result["dependency_graph"] == {"0": []} + assert result["terminal_outputs"] == ["clean.fq"] + def test_empty_returns_comment(self, tmp_path): tmp_path.mkdir(exist_ok=True) result = reconstruct_pipeline(tmp_path) assert "No execution records found" in result + def test_empty_json_format(self, tmp_path): + tmp_path.mkdir(exist_ok=True) + result = json.loads(reconstruct_pipeline(tmp_path, fmt="json")) + assert result == {"records": [], "dependency_graph": {}, "terminal_outputs": []} + def test_session_filter(self, tmp_path): _write_record( tmp_path, _make_record("a", ["a"], session_id="s1", record_id="r1")