Add dag_construction example module

Reference implementation of the partition -> build_dag -> merge_view
pipeline for turning reasoning traces into DAGs, plus per-step
customization notes and CI.

Closes #1

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: zzzhr97

.github/workflows/dag-construction-ci.yml

@@ -0,0 +1,28 @@
+name: dag-construction-ci
+
+on:
+  push:
+    paths:
+      - "dag_construction/**"
+      - ".github/workflows/dag-construction-ci.yml"
+  pull_request:
+    paths:
+      - "dag_construction/**"
+      - ".github/workflows/dag-construction-ci.yml"
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: dag_construction
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - run: python -m pip install --upgrade pip
+      - run: pip install -e ".[dev]"
+      - run: ruff check .
+      - run: mypy trm_dag
+      - run: pytest -q

README.md

@@ -20,6 +20,8 @@
 
 This repository provides the TRM-Preference dataset, TRM weights, and implementations for TRM training and TRM-guided policy optimization.
 
+For modular reasoning-trace DAG construction, see [`dag_construction/`](dag_construction/).
+
 The **Thinking Reward Model (TRM)** evaluates the quality of *reasoning traces* rather than final answers. We study how to optimize reasoning itself: instead of only asking “Is the answer correct?”, we ask:
 
 > **Is this a good way to think?**

dag_construction/.env.example

@@ -0,0 +1,6 @@
+# Whatever your OpenAI-compatible endpoint expects.
+# This may be a full URL (e.g. https://api.example.com/v1) or a proxy token.
+# The value is passed to the OpenAI client unchanged.
+DAG_CONSTRUCTION_BASE_URL="..."
+DAG_CONSTRUCTION_MODEL="deepseek-v3.2"
+DAG_CONSTRUCTION_API_KEY="..."

dag_construction/README.md

@@ -0,0 +1,117 @@
+# TRM DAG Construction
+
+Build DAGs over reasoning traces.
+
+> Example / reference implementation. Defaults reflect one specific setup — swap pieces out for your own data and models.
+
+## Pipeline
+
+Three steps, in order:
+
+```text
+(1) partition  ->  (2) build_dag  ->  (3) merge_view
+```
+
+1. **`partition`** — split a raw `trace` into an ordered list of steps.
+2. **`build_dag`** — for each new step, ask an LLM to classify how it attaches to earlier steps as `continue` / `backtrack` / `merge`, then validate and majority-vote into parents.
+3. **`merge_view`** — collapse linear `continue` chains into super-nodes.
+
+## Quickstart
+
+```bash
+cd dag_construction
+pip install -e ".[dev]"
+cp .env.example .env
+# Fill DAG_CONSTRUCTION_API_KEY, DAG_CONSTRUCTION_MODEL, optional DAG_CONSTRUCTION_BASE_URL.
+python examples/run_single.py
+```
+
+Batch JSONL:
+
+```bash
+python -m trm_dag build --input items.jsonl --output out.jsonl --num-threads 8 --resume
+```
+
+## Customizing each step
+
+### 1. `partition`
+
+- Different keywords: `partition_keyword(..., keywords=("Step", "Stage", ...))` or CLI `--keywords ...`.
+- Skip it: pass rows with `"steps": [...]` directly to `build_dag_batch`.
+- Custom splitter: write your own `(steps, fillers)` function and feed `"steps"` in. See `trm_dag/partition.py` for the contract.
+
+### 2. `build_dag`
+
+- Prompt: edit `trm_dag/prompts/dag_construction.txt` (keep the `$current_id`, `$input_steps`, `$available_steps`, `$leaf_steps`, `$last_previous_step_id` placeholders).
+- LLM backend: any OpenAI-compatible endpoint via `DAG_CONSTRUCTION_BASE_URL` / `DAG_CONSTRUCTION_MODEL`. For non-OpenAI APIs, implement the `ChatClient` protocol in `trm_dag/client.py` and pass `client=` to `build_dag` / `build_dag_batch`.
+- Sampling: `OpenAIConfig(temperature, top_p, n, max_tokens)` and `DagParams(regen_limit)`, or matching CLI flags.
+- Attachment pool size: `DagParams(main_path_cap=..., other_leaf_cap=...)`.
+- Action set / voting: edit `parse_action_previous`, `_validate_candidate`, `majority_vote` in `trm_dag/components.py` (and the prompt).
+- Determinism: `DagParams(random_seed=0)`.
+
+### 3. `merge_view`
+
+- Different merge rule: edit `collapse_continue_chains` in `trm_dag/components.py` (`is_continue_edge`, `is_chain_head`).
+- Skip it: just read `result["dag_graph_raw"]`.
+- Per-super-node summarization: post-process `dag_graph_merged["merged_nodes"]` with your own summarizer.
+
+## Python API
+
+```python
+from trm_dag import DagParams, OpenAIConfig, build_dag
+
+result = build_dag(
+    prompt,
+    steps,
+    openai_config=OpenAIConfig.from_env(),
+    dag_params=DagParams(random_seed=0),
+)
+
+raw = result["dag_graph_raw"]
+merged = result["dag_graph_merged"]
+```
+
+Batch:
+
+```python
+from trm_dag import DagParams, OpenAIConfig, build_dag_batch
+
+rows = [
+    {"prompt": "...", "steps": ["Step one.", "Then two."]},
+    {"prompt": "...", "trace": "Step one.\n\nThen two."},
+]
+
+out_rows = build_dag_batch(
+    rows,
+    openai_config=OpenAIConfig.from_env(),
+    dag_params=DagParams(regen_limit=5, random_seed=0),
+    output_path="items_dag.jsonl",
+    resume=True,
+)
+```
+
+## CLI
+
+```bash
+python -m trm_dag build --input items.jsonl --output items_dag.jsonl
+```
+
+Build flags: `--input`, `--output`, `--num-threads`, `--resume` / `--no-resume` / `--rebuild`, `--keywords`, `--regen-limit`, `--main-path-cap`, `--other-leaf-cap`, `--random-seed`.
+
+LLM flags: `--api-key`, `--base-url`, `--model`, `--temperature`, `--top-p`, `--max-tokens`, `--n`, `--timeout`.
+
+## Input And Output
+
+Input rows can contain pre-split steps:
+
+```json
+{"prompt": "...", "steps": ["Step one.", "Then two."]}
+```
+
+or a raw trace:
+
+```json
+{"prompt": "...", "trace": "Step one.\\n\\nThen two."}
+```
+
+Output rows add `dag_graph_raw`, `dag_graph_merged`, `dag_parse_errors`, `dag_usage`. When input contains `trace`, `build_dag_batch` also keeps generated `steps` and `fillers` on the output row.

dag_construction/examples/data.json

@@ -0,0 +1,18 @@
+[
+  {
+    "prompt": "A train trip has a fixed distance and the same short stop each time. The first plan takes 5 hours at speed v. A faster plan takes 4 hours at speed v+15. Find the total time for a third plan at speed v+10 with the same stop.",
+    "steps": [
+      "Let the distance be d miles and the stop time be h hours.",
+      "From the first plan, d/v + h = 5.",
+      "From the faster plan, d/(v+15) + h = 4.",
+      "One tempting branch is to treat the one-hour saving as if it came directly from adding 15 mph.",
+      "On that shortcut, adding 10 mph would save about two thirds of an hour, giving a rough answer near 4 hours 20 minutes.",
+      "But that branch is shaky: the trip has a fixed distance, so the time saved cannot be scaled directly from the speed increase.",
+      "I should go back to the two trip equations and compare the travel-time parts instead.",
+      "Subtracting the two equations gives d/v - d/(v+15) = 1.",
+      "Solving the equations gives v = 30, d = 90, and h = 2 hours.",
+      "For the third plan, the travel time is 90/(30+10) = 2.25 hours.",
+      "Therefore the total third-plan time is 2.25 + 2 = 4.25 hours, or 255 minutes."
+    ]
+  }
+]

dag_construction/examples/run_batch.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+from trm_dag import DagParams, OpenAIConfig, build_dag_batch
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input", default=str(Path(__file__).with_name("data.json")))
+    parser.add_argument("--output", default="out.jsonl")
+    parser.add_argument("--num-threads", type=int, default=4)
+    args = parser.parse_args()
+
+    rows = json.loads(Path(args.input).read_text(encoding="utf-8"))
+    build_dag_batch(
+        rows,
+        openai_config=OpenAIConfig.from_env(max_tokens=1024),
+        dag_params=DagParams(random_seed=0),
+        num_threads=args.num_threads,
+        output_path=args.output,
+        resume=True,
+    )
+    print(f"wrote {args.output}")
+
+
+if __name__ == "__main__":
+    main()

dag_construction/examples/run_single.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from trm_dag import DagParams, OpenAIConfig, build_dag
+
+
+def main() -> None:
+    item = json.loads((Path(__file__).with_name("data.json")).read_text(encoding="utf-8"))[0]
+    result = build_dag(
+        item["prompt"],
+        item["steps"],
+        openai_config=OpenAIConfig.from_env(max_tokens=1024),
+        dag_params=DagParams(random_seed=0),
+    )
+    print(
+        json.dumps(
+            {
+                "dag_graph_raw": result["dag_graph_raw"],
+                "dag_graph_merged": result["dag_graph_merged"],
+            },
+            ensure_ascii=False,
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

dag_construction/pyproject.toml

@@ -0,0 +1,53 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "trm-dag"
+version = "0.1.0"
+description = "A small, hackable DAG builder for TRM reasoning traces."
+readme = "README.md"
+requires-python = ">=3.10"
+license = {text = "MIT"}
+authors = [{name = "Haoran Zhang"}]
+dependencies = [
+  "openai>=1.0",
+  "python-dotenv>=1.0",
+  "rich>=13.0",
+]
+
+[project.optional-dependencies]
+dev = [
+  "pytest>=8.0",
+  "ruff>=0.5",
+  "mypy>=1.10",
+]
+
+[project.scripts]
+trm-dag = "trm_dag.cli:main"
+
+[tool.setuptools.packages.find]
+include = ["trm_dag*"]
+exclude = ["scripts*", "tests*"]
+
+[tool.setuptools.package-data]
+trm_dag = ["prompts/*.txt"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP", "B", "C4", "SIM"]
+ignore = ["E501"]
+
+[tool.mypy]
+python_version = "3.10"
+strict = true
+warn_unused_ignores = true
+warn_return_any = true
+disallow_any_generics = false
+
+[[tool.mypy.overrides]]
+module = ["openai", "dotenv", "rich.*"]
+ignore_missing_imports = true

dag_construction/tests/conftest.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from trm_dag.config import OpenAIConfig  # noqa: E402
+
+
+@dataclass
+class MockOpenAIClient:
+    responses: list[list[str]]
+    calls: int = 0
+    prompts: list[str] = field(default_factory=list)
+
+    def complete(self, prompt: str, config: OpenAIConfig, *, n: int) -> tuple[list[str], dict[str, int]]:
+        self.prompts.append(prompt)
+        idx = self.calls
+        self.calls += 1
+        if idx >= len(self.responses):
+            batch = self.responses[-1] if self.responses else ["ok\n<|action|>continue"]
+        else:
+            batch = self.responses[idx]
+        return batch[:n], {
+            "prompt_tokens": 1,
+            "completion_tokens": 1,
+            "total_tokens": 2,
+        }
+
+
+def openai_config() -> OpenAIConfig:
+    return OpenAIConfig(api_key="test", base_url=None, model="mock", n=1)

dag_construction/tests/test_dag_validity.py

@@ -0,0 +1,60 @@
+from collections import deque
+
+from conftest import MockOpenAIClient, openai_config
+
+from trm_dag import DagParams, build_dag
+
+
+def assert_valid_dag(dag: dict) -> None:
+    parents = {int(k): [int(x) for x in v] for k, v in dag["parents"].items()}
+    children = {int(k): [int(x) for x in v] for k, v in dag["children"].items()}
+    actions = {int(k): v for k, v in dag["actions"].items()}
+    assert actions[0] == "root"
+    assert all(action in {"root", "continue", "backtrack", "merge"} for action in actions.values())
+    for node, parent_list in parents.items():
+        if node != 0:
+            assert len(parent_list) >= 1
+
+    seen = {0}
+    queue = deque([0])
+    while queue:
+        node = queue.popleft()
+        for child in children.get(node, []):
+            if child not in seen:
+                seen.add(child)
+                queue.append(child)
+    assert seen == set(parents)
+
+    visiting: set[int] = set()
+    done: set[int] = set()
+
+    def dfs(node: int) -> None:
+        assert node not in visiting
+        if node in done:
+            return
+        visiting.add(node)
+        for child in children.get(node, []):
+            dfs(child)
+        visiting.remove(node)
+        done.add(node)
+
+    dfs(0)
+
+
+def test_mock_pipeline_produces_valid_dag() -> None:
+    client = MockOpenAIClient(
+        [
+            ["s1 follows\n<|action|>continue"],
+            ["restart\n<|action|>backtrack\n<|previous|>0"],
+            ["join\n<|action|>merge\n<|previous|>1,2"],
+        ]
+    )
+    result = build_dag(
+        "p",
+        ["root", "a", "b", "c"],
+        openai_config=openai_config(),
+        dag_params=DagParams(regen_limit=1, random_seed=0),
+        client=client,
+    )
+    assert_valid_dag(result["dag_graph_raw"])
+    assert result["dag_graph_raw"]["leaves"] == [3]