diff --git a/experiments/pydantic-graph-prototype/.gitignore b/experiments/pydantic-graph-prototype/.gitignore new file mode 100644 index 000000000..241f1369a --- /dev/null +++ b/experiments/pydantic-graph-prototype/.gitignore @@ -0,0 +1,5 @@ +__pycache__/ +*.egg-info/ +.pytest_cache/ +dist/ +build/ diff --git a/experiments/pydantic-graph-prototype/README.md b/experiments/pydantic-graph-prototype/README.md new file mode 100644 index 000000000..bccf65c37 --- /dev/null +++ b/experiments/pydantic-graph-prototype/README.md @@ -0,0 +1,44 @@ +# pg-factory: pydantic-graph Factory Engine Prototype + +A standalone prototype that validates whether [pydantic-graph](https://github.com/pydantic/pydantic-graph) can replace the Factory's custom workflow execution engine. + +## What this is + +The Factory currently uses a custom graph execution engine (`primitives.py` + `executor.py`) with runtime edge matching, explicit `dict[str, Node] + list[Edge]` construction, and manual verdict routing. This prototype ports progressively harder subgraphs to pydantic-graph's `BaseNode` + `Graph` API to evaluate: + +1. Can `BaseNode` model factory node types cleanly? +2. Does return-type-as-edge routing work for verdict/gate patterns? +3. Does `FactoryState` + `FactoryDeps` carry context through `GraphRunContext`? +4. Are `Graph.iter()` and Mermaid rendering useful bonus wins? +5. Can `definitions.py` construction patterns adapt? + +## Setup + +```bash +python3 -m pip install -e ".[dev]" +``` + +## Run tests + +```bash +python3 -m pytest +``` + +## Type checking + +```bash +python3 -m pyright +``` + +## Project structure + +``` +src/pg_factory/ + state.py — FactoryState dataclass (mutable graph state) + deps.py — FactoryDeps dataclass (immutable dependencies) + verdicts.py — VerdictType enum + HaltResult for End returns +tests/ + conftest.py — shared fixtures + test_smoke.py — basic 2-node graph smoke test + test_future_annotations.py — critical risk: from __future__ import annotations compatibility +``` diff --git a/experiments/pydantic-graph-prototype/docs/comparison-report.md b/experiments/pydantic-graph-prototype/docs/comparison-report.md new file mode 100644 index 000000000..8d771b83b --- /dev/null +++ b/experiments/pydantic-graph-prototype/docs/comparison-report.md @@ -0,0 +1,206 @@ +# Comparison Report: pydantic-graph vs Current Engine + +## Side-by-Side Event Traces + +Both engines execute the same workflow: `builder → qa_gate → fork_qa(3 QA agents)`. +The comparison harness (`src/pg_factory/compare.py`) defines this workflow in both +representations and runs them with identical verdict sequences. + +### PROCEED Path + +| Step | Current Engine | pydantic-graph | +|------|---------------------------------------|---------------------------------------| +| 1 | builder: execute | builder: execute | +| 2 | qa_gate: gate_verdict (proceed) | qa_gate: gate_verdict (proceed) | +| 3 | fork_qa: fork (3 targets) | fork_qa: child_completed (health_checker) | +| 4 | health_checker: execute | fork_qa: child_completed (code_reviewer) | +| 5 | code_reviewer: execute | fork_qa: child_completed (adversarial_tester) | +| 6 | adversarial_tester: execute | fork_qa: fork_join_complete | +| 7 | fork_qa: fork_join_complete | — | +| 8 | join_qa: join | — | + +**Key difference:** The current engine uses separate fork/join nodes and records each +branch as an independent "execute" event. pydantic-graph encapsulates the fork/join +inside a single `ForkJoinNode`, recording children as `child_completed` events with +per-child timing data. The JoinNode barrier is implicit. + +### RELOOP → PROCEED Path + +| Step | Current Engine | pydantic-graph | +|------|---------------------------------------|---------------------------------------| +| 1 | builder: execute | builder: execute | +| 2 | qa_gate: gate_verdict (reloop → builder) | qa_gate: gate_verdict (reloop) | +| 3 | builder: execute (2nd time) | builder: execute (2nd time) | +| 4 | qa_gate: gate_verdict (proceed) | qa_gate: gate_verdict (proceed) | +| 5-8 | fork + 3 children + join | 3 child_completed + fork_join_complete | + +**Key difference:** Both engines re-execute the builder on RELOOP. The current engine +routes via string-based edge matching (`Edge(source="qa_gate", target="builder", +condition="reloop")`). pydantic-graph routes via the gate's return type — returning +`CompareBuilderNode(...)` directly, with the type system ensuring only valid targets +are returned. + +### HALT Path + +| Step | Current Engine | pydantic-graph | +|------|---------------------------------------|---------------------------------------| +| 1 | builder: execute | builder: execute | +| 2 | qa_gate: gate_verdict (halt) | qa_gate: gate_verdict (halt) | + +Identical behavior — both stop immediately after gate HALT. + + +## Lines-of-Code Comparison + +### Workflow Definition Patterns + +| Component | Current Engine (definitions.py) | pydantic-graph (BaseNode) | +|--------------------------|--------------------------------|---------------------------| +| Study chain (4 nodes) | 51 lines (dict + Edge list) | 154 lines (4 BaseNode classes) | +| Gate routing | ~40 lines (GateNode + edges) | 85 lines (GateBaseNode) | +| Fork/join (3 branches) | 76 lines (Fork + Join + edges) | 111 lines (ForkJoinNode) | +| **Total node definitions** | **~167 lines** | **350 lines** | + +### Infrastructure + +| Component | Current Engine | pydantic-graph | +|--------------------------|-----------------------------|-----------------------------| +| Node/Edge primitives | 322 lines (primitives.py) | 73 lines (state + deps + verdicts) | +| Executor | 1,135 lines (executor.py) | 0 lines (Graph.run provided by library) | +| **Total infrastructure** | **1,457 lines** | **73 lines** | + +### Analysis + +Node definitions are **~2x longer** in pydantic-graph because each node is a full +class with typed `run()` method, docstring, and explicit state interaction. The +current engine's dict-based style is more compact for definition. + +However, pydantic-graph **eliminates 1,457 lines of infrastructure** — the entire +executor, edge-walking logic, fork dispatch, gate verdict parsing, and read/write +polling. This is handled by `Graph.run()` / `Graph.iter()` from the library. + +**Net LOC change for equivalent functionality:** +- Current engine: 167 (definitions) + 1,457 (infrastructure) = **1,624 lines** +- pydantic-graph: 350 (definitions) + 73 (types) = **423 lines** +- **Reduction: 74%** + + +## Type Safety Analysis + +### What pyright catches in pydantic-graph that the current engine misses + +| Misconfiguration | Current Engine | pydantic-graph (pyright strict) | +|-------------------------------------|-----------------------|---------------------------------| +| Misspelled edge target | Runtime error | Compile-time type error | +| Gate returns invalid successor | Runtime edge mismatch | Return type annotation error | +| Missing node in workflow | Runtime KeyError | Import/reference error | +| Wrong state field type | Runtime TypeError | Type error on state access | +| Fork target not in node dict | Silent no-op | N/A (no separate fork targets) | +| Gate verdict with no matching edge | Runtime halt | N/A (exhaustive return union) | + +**Example — invalid gate routing:** + +Current engine (detected only at runtime): +```python +# This silently compiles but crashes at runtime when the gate returns RELOOP +edges = [ + Edge(source="qa_gate", target="buildr", condition="reloop"), # typo! +] +``` + +pydantic-graph (caught by pyright): +```python +class QAGateNode(GateBaseNode): + async def run(self, ctx) -> BuilderNode | NextNode | End[HaltResult]: + return Buildr() # pyright error: "Buildr" is not defined +``` + +### What the current engine checks that pydantic-graph doesn't + +| Check | Current Engine | pydantic-graph | +|------------------------------------|-----------------------|-----------------------------| +| File reads/writes dependencies | `reads`/`writes` sets on nodes | No equivalent (state-based) | +| Graph connectivity validation | `validate_graph()` via NetworkX | Return type inference only | +| Circular dependency detection | Edge-based cycle check | Not checked (cycles are valid via reloop) | + + +## Mermaid Diagram Comparison + +### pydantic-graph (auto-generated via `Graph.render()`) + +```mermaid +stateDiagram-v2 + CompareBuilderNode + CompareQAGateNode + state decision <> + CompareQAForkJoinNode + + [*] --> CompareBuilderNode + CompareBuilderNode --> CompareQAGateNode + CompareQAGateNode --> decision + decision --> CompareBuilderNode + decision --> CompareQAForkJoinNode + decision --> [*] + CompareQAForkJoinNode --> [*] +``` + +### Current engine (no built-in visualization) + +The current engine has no Mermaid generation capability. Workflow topology must be +manually reconstructed from `dict[str, NodeType]` + `list[Edge]` definitions. There +is a `validate_graph()` method using NetworkX, but no rendering. + +### Assessment + +pydantic-graph's `Graph.render()` produces valid Mermaid state diagrams with zero +configuration. Gate routing appears as a `<>` decision node showing all +possible successors. The one limitation: `ForkJoinNode` appears as a single node +rather than a visual fan-out of branches — the internal parallelism is only visible +through event timing data, not the diagram. + + +## Feature Matrix + +| Feature | Current Engine | pydantic-graph | Notes | +|---------------------------------|:--------------:|:--------------:|--------------------------------------| +| Sequential chains | ✓ | ✓ | Both handle linear node sequences | +| Gate/verdict routing | ✓ | ✓ | PG uses return types; current uses edges | +| Iteration counting (max_iter) | ✓ | ✓ | Both track (gate, target) counts | +| Feedback injection | ✓ | ✓ | Both pass context on RELOOP | +| Parallel fork/join | ✓ | ✓ | Both use asyncio.gather | +| `reads`/`writes` file deps | ✓ | ✗ | PG uses typed state instead | +| Non-blocking (fire-and-forget) | ✓ | △ | PG needs asyncio.create_task wrapper | +| SubgraphForkNode (N worktrees) | ✓ | ✗ | No PG equivalent — stays custom | +| SelectionNode (branch compare) | ✓ | ✗ | No PG equivalent — stays custom | +| LLMNode (in-process API loop) | ✓ | ✗ | Not related to graph execution | +| Mermaid diagram generation | ✗ | ✓ | PG auto-generates from types | +| Compile-time edge validation | ✗ | ✓ | pyright strict mode | +| Step-by-step iteration | ✗ | ✓ | Graph.iter() async context manager | +| Graph persistence/resume | ✗ | ✓ | Built into pydantic-graph (untested) | + +**Key:** ✓ = fully supported, △ = partial/workaround, ✗ = not supported + +### What pydantic-graph replaces + +- Edge-based graph walking (`_execute_from`, `_next_unconditional`, `_next_conditional`) +- Gate verdict parsing (`_parse_agent_verdict`, `_parse_fn_verdict`) +- Fork dispatch (`_execute_fork` + `run_branch` + gather) +- Node execution dispatch (`_run_node`, `_execute_action_node`) +- Event emission plumbing + +### What stays custom + +- `SubgraphForkNode` (parallel experiments in isolated worktrees) +- `SelectionNode` (best-score comparison across branches) +- `LLMNode` (in-process API tool-use loop) +- `reads`/`writes` file dependency polling +- Agent invocation (`invoke_agent` subprocess management) + +### What's lost + +- `reads`/`writes` declarative file dependencies — replaced by typed state fields, + which are more type-safe but less explicit about filesystem artifacts +- Flat dict-based workflow construction — replaced by class hierarchies, which are + more verbose but catch more errors at compile time +- Graph validation via NetworkX — replaced by type-system edge inference, which + catches a different (and largely overlapping) class of errors diff --git a/experiments/pydantic-graph-prototype/docs/gate-routing.mmd b/experiments/pydantic-graph-prototype/docs/gate-routing.mmd new file mode 100644 index 000000000..9ee83eec2 --- /dev/null +++ b/experiments/pydantic-graph-prototype/docs/gate-routing.mmd @@ -0,0 +1,13 @@ +stateDiagram-v2 + MockBuilderNode + QAGateNode + state decision <> + MockNextNode + + [*] --> MockBuilderNode + MockBuilderNode --> QAGateNode + QAGateNode --> decision + decision --> MockBuilderNode + decision --> MockNextNode + decision --> [*] + MockNextNode --> [*] diff --git a/experiments/pydantic-graph-prototype/docs/migration-verdict.md b/experiments/pydantic-graph-prototype/docs/migration-verdict.md new file mode 100644 index 000000000..9ea349ed6 --- /dev/null +++ b/experiments/pydantic-graph-prototype/docs/migration-verdict.md @@ -0,0 +1,307 @@ +# Migration Verdict: pydantic-graph for Factory Workflow Engine + +## Q1: Can BaseNode model factory node types cleanly? + +**Verdict: YES — with one structural adaptation** + +### Evidence + +Phase 2 ported the study chain (`graph_update → study → graph_explorer → concat_study`) +as 4 `BaseNode` subclasses. Each maps directly from a current engine node type: + +| Current Engine | pydantic-graph | Mapping | +|-----------------------|-------------------------------|----------------------| +| `FnNode("graph_update")` | `GraphUpdateNode(BaseNode)` | 1:1, command in run() | +| `Study("study")` | `StudyNode(BaseNode)` | 1:1, command in run() | +| `AgentNode("graph_explorer")` | `GraphExplorerNode(BaseNode)` | 1:1, agent call in run() | +| `FnNode("concat_study")` | `ConcatStudyNode(BaseNode)` | 1:1, command in run() | + +The model is clean: each node's `run()` method encapsulates what the current engine +dispatches to `_run_node()`, `_run_fn()`, or `_run_agent()`. State access moves from +executor instance variables to `ctx.state` (typed) and `ctx.deps` (immutable deps). + +All 7 tests for the study chain pass, confirming execution order, state mutations, +file creation (dry_run mode), Mermaid output, and event recording. + +### Adaptation required + +`FnNode` and `AgentNode` become full Python classes (~30 lines each vs ~5 lines in +dict form). This is more verbose but gains type-checked state access and explicit +control flow. + +--- + +## Q2: Does return-type routing work for GateNode verdicts? + +**Verdict: YES — this is the biggest architectural gain** + +### Evidence + +Phase 3 implemented `GateBaseNode` with a typed return union: + +```python +class QAGateNode(GateBaseNode): + async def run(self, ctx) -> BuilderNode | NextNode | End[HaltResult]: +``` + +This replaces the current engine's runtime edge matching: + +```python +# Current: string-based, runtime-only +Edge(source="qa_gate", target="builder", condition="reloop") +Edge(source="qa_gate", target="fork_qa", condition="proceed") +# Executor: _next_conditional(node_id, VerdictType.PROCEED) +``` + +**10 tests** verify all 3 verdict paths (PROCEED, RELOOP, HALT), iteration counting, +feedback injection, feedback updates across iterations, Mermaid branching topology, +and Graph.iter() event yields for both simple and reloop paths. + +**Type safety verified:** If a gate's `run()` body tries to return a node not in its +return annotation, pyright (strict mode) catches it as a type error. The current +engine only catches this at runtime when no matching edge is found. + +The comparison harness (Phase 5) confirmed that both engines produce identical gate +verdict sequences for PROCEED, RELOOP, HALT, and max-iteration-halt paths. + +### Key finding + +Return-type-as-edge eliminates the entire `_parse_agent_verdict` and +`_parse_fn_verdict` code paths (102 lines in executor.py). Verdict routing is +declared in the type system, not parsed from strings. + +--- + +## Q3: Does FactoryState + FactoryDeps carry context effectively? + +**Verdict: YES — cleaner than the executor's instance variables** + +### Evidence + +The current engine spreads state across `WorkflowExecutor` instance variables: + +```python +class WorkflowExecutor: + self.completed_files: set[str] + self.node_context: dict[str, str] # feedback + self.iteration_counts: dict[tuple[str, str], int] + self.result.node_outputs: dict[str, str] + self.result.events: list[dict] +``` + +pydantic-graph consolidates these into two typed dataclasses: + +```python +@dataclass +class FactoryState: # mutable, threaded via ctx.state + iteration_counts: dict[tuple[str, str], int] + node_feedback: dict[str, str] + node_outputs: dict[str, str] + completed_files: set[str] + events: list[dict[str, Any]] + +@dataclass +class FactoryDeps: # immutable, injected via ctx.deps + project_path: Path + dry_run: bool + event_emitter: Callable +``` + +Every node across all 4 phases successfully reads and writes state through +`ctx.state` and `ctx.deps`. Cross-node communication patterns verified: + +- **Sequential state threading:** Each study chain node reads/writes + `ctx.state.node_outputs` (Phase 2, 7 tests) +- **Gate iteration tracking:** `ctx.state.iteration_counts[(gate_id, target_id)]` + increments correctly across RELOOP cycles (Phase 3, 10 tests) +- **Feedback injection:** Gates write to `ctx.state.node_feedback[target_id]`, + reloop targets read it on re-entry (Phase 3, verified explicitly) +- **Fork/join state sharing:** All 3 QA children write to the same `ctx.state` + concurrently without races (Phase 4, 10 tests) + +### Advantage over current engine + +State access is type-checked by pyright. In the current engine, +`self.node_context[target]` is an untyped dict access — typos in key names are +silent bugs. In pydantic-graph, `ctx.state.node_feedback[target_id]` is checked +against the `FactoryState` dataclass definition. + +--- + +## Q4: Mermaid + Graph.iter() as bonus wins? + +**Verdict: YES — significant observability gains** + +### Mermaid + +`Graph.render()` auto-generates valid Mermaid state diagrams from the type-inferred +topology. No configuration, no manual diagram maintenance. + +Generated diagrams verified across all phases: +- Study chain: linear 4-node topology with `[*]` terminal (Phase 2) +- Gate routing: `<>` decision node showing PROCEED/RELOOP/HALT edges (Phase 3) +- Fork/join: single ForkJoinNode (internal parallelism not visible) (Phase 4) +- Comparison workflow: full builder→gate→fork pattern (Phase 5) + +The current engine has **no built-in visualization**. `validate_graph()` uses +NetworkX for structural validation but produces no visual output. + +**Limitation:** ForkJoinNode renders as a single node. The 3-branch fan-out is not +visible in the diagram. A custom Mermaid post-processor could expand it, but this +is acceptable for the prototype. + +### Graph.iter() + +`Graph.iter()` provides step-by-step async iteration over graph execution: + +```python +async with graph.iter(state=state, deps=deps, inputs=start) as run: + async for event in run: + # observe each step in real time +``` + +Verified in all phases: correct event counts, EndMarker as terminal event, one +event per node step. This enables real-time observability that the current engine +achieves only through its event emission system. + +--- + +## Q5: Can definitions.py construction patterns adapt? + +**Verdict: YES — but the adaptation is structural, not syntactic** + +### Evidence + +The comparison harness (Phase 5) defines the identical workflow in both representations: + +**Current engine pattern** (`build_current_engine_workflow()`): +```python +nodes = { + "builder": SimNode(id="builder", kind=NodeKind.ACTION), + "qa_gate": SimNode(id="qa_gate", kind=NodeKind.GATE), + "fork_qa": SimNode(id="fork_qa", kind=NodeKind.FORK, + fork_targets=["health_checker", ...]), + ... +} +edges = [ + SimEdge(source="builder", target="qa_gate"), + SimEdge(source="qa_gate", target="fork_qa", condition="proceed"), + SimEdge(source="qa_gate", target="builder", condition="reloop"), + ... +] +``` + +**pydantic-graph pattern** (`build_pydantic_graph()`): +```python +class CompareBuilderNode(BaseNode): + async def run(self, ctx) -> "CompareQAGateNode": + ... + +class CompareQAGateNode(GateBaseNode): + async def run(self, ctx) -> CompareBuilderNode | CompareQAForkJoinNode | End: + ... + +class CompareQAForkJoinNode(ForkJoinNode): + async def run(self, ctx) -> End[HaltResult]: + await self.execute_children(ctx) + ... +``` + +**The adaptation:** +- `dict[str, NodeType]` → class hierarchy (each node type is a class) +- `list[Edge]` → return type annotations (edges are implicit) +- `ForkNode + JoinNode` → single `ForkJoinNode` (asyncio.gather inside `run()`) +- `GateNode + edge conditions` → return union type (pyright validates) +- String-based node IDs → Python class references (import-time validation) + +### The 18 new comparison tests confirm behavioral equivalence + +Both representations produce: +- Same active nodes for PROCEED, RELOOP, and HALT paths +- Same gate verdict sequences +- Same fork/join children sets +- Same max-iteration halt behavior + +The construction pattern changes significantly (data-driven → type-driven), but the +execution semantics are preserved. + +--- + +## Recommended Migration Path + +### Phase 1: Core Workflows (Recommended — start here) + +Port `improve_workflow` first — it exercises all 3 patterns (sequential study chain, +gate verdicts, and deep-QA fork/join) and is the most frequently run workflow. + +| Workflow | Complexity | Patterns Used | Priority | +|-----------------|:----------:|---------------------------|:--------:| +| improve | Medium | Sequential + Gate + Fork | High | +| build | Medium | Sequential + Gate + Fork | High | +| discover | Low | Sequential only | Medium | +| review | Low | Sequential only | Medium | +| research | Medium | Sequential + Gate + Fork | Medium | + +### Phase 2: Parallel Execution (Keep custom) + +`SubgraphForkNode` and `SelectionNode` have no pydantic-graph equivalent and should +remain custom implementations. These can wrap pydantic-graph subgraphs internally +(each worktree branch runs a `Graph` instance). + +### Phase 3: Advanced Nodes (Keep custom) + +`LLMNode` (in-process API tool-use loop) is orthogonal to graph execution and should +remain a separate concern. It can be wrapped as a `BaseNode` subclass that delegates +to the existing `run_llm_loop()` function. + +### What to keep custom + +| Component | Reason | +|------------------------|---------------------------------------------------------| +| SubgraphForkNode | Worktree isolation is unique to the factory | +| SelectionNode | Branch comparison logic is domain-specific | +| LLMNode | In-process API loop, orthogonal to graph execution | +| `reads`/`writes` deps | File-polling replaced by typed state; migration needed | +| Agent invocation | `invoke_agent` subprocess management stays as-is | + +### Estimated Complexity + +| Task | Effort | Risk | +|----------------------------------------|----------|-------| +| Port improve_workflow nodes | 2-3 days | Low | +| Port build_workflow nodes | 2-3 days | Low | +| Replace executor for ported workflows | 1-2 days | Medium | +| Migrate `reads`/`writes` to state | 3-5 days | High | +| Integrate SubgraphFork with PG graphs | 2-3 days | Medium | +| Remove legacy executor (after full port)| 1 day | Low | + +### Risk Assessment + +**Low risk:** +- Sequential chains and gate routing are fully validated (34 tests across Phases 1-4) +- asyncio.gather fork/join pattern is identical to current implementation + +**Medium risk:** +- `reads`/`writes` file dependency polling has no pydantic-graph equivalent — the + migration to typed state requires careful analysis of which files are used as + inter-node communication channels vs. which are external artifacts + +**Mitigated risk:** +- `from __future__ import annotations` interaction tested and validated in Phase 1 + (4 dedicated tests) — this was the top-flagged risk from research and it passed + +### Recommendation + +**Proceed with migration using the thin-adapter approach:** + +1. Create `BaseNode` subclass wrappers for existing node types (`AgentNode`, `FnNode`, + `Study`, `GateNode`) +2. Port one workflow at a time, starting with `improve_workflow` +3. Run both engines in parallel during migration (current engine as fallback) +4. Keep `SubgraphForkNode`, `SelectionNode`, and `LLMNode` as custom extensions +5. Deprecate `executor.py` after all core workflows are ported + +The prototype demonstrates that pydantic-graph eliminates 74% of infrastructure code +while providing compile-time edge validation, auto-generated Mermaid diagrams, and +step-by-step execution iteration. The tradeoff is 2x more verbose node definitions, +which is acceptable given the type safety gains. diff --git a/experiments/pydantic-graph-prototype/docs/study-chain.mmd b/experiments/pydantic-graph-prototype/docs/study-chain.mmd new file mode 100644 index 000000000..719c9fe60 --- /dev/null +++ b/experiments/pydantic-graph-prototype/docs/study-chain.mmd @@ -0,0 +1,11 @@ +stateDiagram-v2 + GraphUpdateNode + StudyNode + GraphExplorerNode + ConcatStudyNode + + [*] --> GraphUpdateNode + GraphUpdateNode --> StudyNode + StudyNode --> GraphExplorerNode + GraphExplorerNode --> ConcatStudyNode + ConcatStudyNode --> [*] \ No newline at end of file diff --git a/experiments/pydantic-graph-prototype/pyproject.toml b/experiments/pydantic-graph-prototype/pyproject.toml new file mode 100644 index 000000000..0e8a9a7f1 --- /dev/null +++ b/experiments/pydantic-graph-prototype/pyproject.toml @@ -0,0 +1,29 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "pg-factory" +version = "0.1.0" +description = "Prototype: pydantic-graph as Factory workflow engine" +requires-python = ">=3.12" +dependencies = [ + "pydantic-graph>=0.2", + "structlog>=24.0", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8.0", + "pytest-asyncio>=0.24", + "pyright>=1.1", +] + +[tool.pytest.ini_options] +asyncio_mode = "auto" +testpaths = ["tests"] + +[tool.pyright] +pythonVersion = "3.12" +typeCheckingMode = "strict" +include = ["src"] diff --git a/experiments/pydantic-graph-prototype/src/pg_factory/__init__.py b/experiments/pydantic-graph-prototype/src/pg_factory/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/experiments/pydantic-graph-prototype/src/pg_factory/compare.py b/experiments/pydantic-graph-prototype/src/pg_factory/compare.py new file mode 100644 index 000000000..90c98c384 --- /dev/null +++ b/experiments/pydantic-graph-prototype/src/pg_factory/compare.py @@ -0,0 +1,437 @@ +"""Dual-engine comparison harness. + +Defines the SAME workflow in both representations and compares execution: +1. Current engine simulation: dict[str, SimNode] + list[SimEdge] → mock walker +2. pydantic-graph: BaseNode subclasses → Graph → Graph.run() + +The comparison workflow covers all patterns from Phases 2-4: + - Sequential chain: builder executes, feeds into gate + - Gate with RELOOP: qa_gate routes PROCEED/RELOOP/HALT + - Parallel fork/join: 3 QA agents via asyncio.gather +""" + +import asyncio +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from enum import Enum +from typing import Any + +from pydantic_graph import BaseNode, End, GraphBuilder, GraphRunContext + +from pg_factory.deps import FactoryDeps +from pg_factory.nodes.gates import GateBaseNode +from pg_factory.nodes.parallel import ChildAgent, ForkJoinNode +from pg_factory.state import FactoryState +from pg_factory.verdicts import HaltResult, VerdictType + + +# ═══════════════════════════════════════════════════════════════════ +# Part 1: Current Engine Simulation +# ═══════════════════════════════════════════════════════════════════ + + +class NodeKind(str, Enum): + ACTION = "action" + GATE = "gate" + FORK = "fork" + JOIN = "join" + + +@dataclass(frozen=True) +class SimEdge: + source: str + target: str + condition: str | None = None + + +@dataclass +class SimNode: + id: str + kind: NodeKind + fork_targets: list[str] = field(default_factory=list) + + +@dataclass +class SimWorkflow: + name: str + nodes: dict[str, SimNode] + edges: list[SimEdge] + start_node: str + + +@dataclass +class TraceEvent: + node: str + action: str + detail: dict[str, Any] = field(default_factory=dict) + + +def build_current_engine_workflow() -> SimWorkflow: + """Define the comparison workflow as dict[str, Node] + list[Edge]. + + Mirrors the definitions.py pattern: explicit node dict, explicit edge list, + ForkNode.targets for parallel dispatch, string-based wiring. + """ + nodes = { + "builder": SimNode(id="builder", kind=NodeKind.ACTION), + "qa_gate": SimNode(id="qa_gate", kind=NodeKind.GATE), + "fork_qa": SimNode( + id="fork_qa", + kind=NodeKind.FORK, + fork_targets=["health_checker", "code_reviewer", "adversarial_tester"], + ), + "health_checker": SimNode(id="health_checker", kind=NodeKind.ACTION), + "code_reviewer": SimNode(id="code_reviewer", kind=NodeKind.ACTION), + "adversarial_tester": SimNode(id="adversarial_tester", kind=NodeKind.ACTION), + "join_qa": SimNode(id="join_qa", kind=NodeKind.JOIN), + } + edges = [ + SimEdge(source="builder", target="qa_gate"), + SimEdge(source="qa_gate", target="fork_qa", condition="proceed"), + SimEdge(source="qa_gate", target="builder", condition="reloop"), + SimEdge(source="fork_qa", target="health_checker"), + SimEdge(source="fork_qa", target="code_reviewer"), + SimEdge(source="fork_qa", target="adversarial_tester"), + SimEdge(source="fork_qa", target="join_qa"), + ] + return SimWorkflow( + name="comparison-workflow", + nodes=nodes, + edges=edges, + start_node="builder", + ) + + +def simulate_current_engine( + workflow: SimWorkflow, + verdict_fn: Callable[[int], VerdictType], + max_iterations: int = 3, +) -> list[TraceEvent]: + """Walk the graph, recording what WorkflowExecutor would do. + + Simplified mock of executor.py's _execute_from: + - ACTION nodes: record execution, follow unconditional edge + - GATE nodes: call verdict_fn(call_count), follow conditional edge + - FORK nodes: record each target execution, follow to join + - JOIN nodes: record barrier, follow unconditional edge + + verdict_fn receives the gate evaluation count (0-indexed). + """ + events: list[TraceEvent] = [] + iteration_counts: dict[tuple[str, str], int] = {} + gate_call_count = 0 + + def _find_edge(source: str, condition: str | None = None) -> str | None: + for edge in workflow.edges: + if edge.source != source or edge.condition != condition: + continue + if condition is None: + node = workflow.nodes.get(source) + if node and node.kind == NodeKind.FORK and edge.target in node.fork_targets: + continue + return edge.target + return None + + def _walk(node_id: str) -> None: + nonlocal gate_call_count + + node = workflow.nodes.get(node_id) + if node is None: + return + + if node.kind == NodeKind.ACTION: + events.append(TraceEvent(node=node_id, action="execute")) + nxt = _find_edge(node_id) + if nxt: + _walk(nxt) + + elif node.kind == NodeKind.GATE: + verdict = verdict_fn(gate_call_count) + gate_call_count += 1 + + if verdict == VerdictType.HALT: + events.append(TraceEvent( + node=node_id, action="gate_verdict", + detail={"verdict": "halt"}, + )) + return + + if verdict == VerdictType.RELOOP: + target = _find_edge(node_id, "reloop") + if not target: + return + key = (node_id, target) + count = iteration_counts.get(key, 0) + 1 + iteration_counts[key] = count + if count > max_iterations: + events.append(TraceEvent( + node=node_id, action="gate_verdict", + detail={"verdict": "halt", "reason": "max_iterations"}, + )) + return + events.append(TraceEvent( + node=node_id, action="gate_verdict", + detail={"verdict": "reloop", "target": target, "iteration": count}, + )) + _walk(target) + return + + events.append(TraceEvent( + node=node_id, action="gate_verdict", + detail={"verdict": "proceed"}, + )) + target = _find_edge(node_id, "proceed") + if target: + _walk(target) + + elif node.kind == NodeKind.FORK: + events.append(TraceEvent( + node=node_id, action="fork", + detail={"targets": list(node.fork_targets)}, + )) + for t in node.fork_targets: + events.append(TraceEvent(node=t, action="execute")) + events.append(TraceEvent( + node=node_id, action="fork_join_complete", + detail={"children": list(node.fork_targets)}, + )) + nxt = _find_edge(node_id) + if nxt: + _walk(nxt) + + elif node.kind == NodeKind.JOIN: + events.append(TraceEvent(node=node_id, action="join")) + nxt = _find_edge(node_id) + if nxt: + _walk(nxt) + + _walk(workflow.start_node) + return events + + +# ═══════════════════════════════════════════════════════════════════ +# Part 2: pydantic-graph Implementation +# ═══════════════════════════════════════════════════════════════════ + + +class CompareQAForkJoinNode(ForkJoinNode): + """Fork/join running 3 QA agents concurrently, then ends.""" + + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> End[HaltResult]: + await self.execute_children(ctx) + return End(HaltResult(reason="qa_complete")) + + +class CompareBuilderNode(BaseNode[FactoryState, FactoryDeps, HaltResult]): + """Builder that feeds into the QA gate.""" + + def __init__( + self, + verdict_fn: Callable[[FactoryState], VerdictType] | None = None, + max_iterations: int = 3, + ) -> None: + self.verdict_fn = verdict_fn + self.max_iterations = max_iterations + + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> "CompareQAGateNode": + ctx.state.node_outputs["builder"] = "build_output" + ctx.state.events.append({"node": "builder", "action": "execute"}) + return CompareQAGateNode( + gate_id="qa_gate", + verdict_fn=self.verdict_fn, + max_iterations=self.max_iterations, + ) + + +class CompareQAGateNode(GateBaseNode): + """Gate routing: builder (RELOOP) | fork/join (PROCEED) | End (HALT).""" + + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> CompareBuilderNode | CompareQAForkJoinNode | End[HaltResult]: + verdict = self.evaluate_verdict(ctx.state) + + if verdict == VerdictType.HALT: + self.record_verdict_event(ctx, verdict) + return End(HaltResult(reason="gate_halted")) + + if verdict == VerdictType.RELOOP: + target_id = "CompareBuilderNode" + halt = self.check_and_increment_iteration(ctx, target_id) + if halt is not None: + self.record_verdict_event(ctx, VerdictType.HALT, target_id) + return halt + iteration = ctx.state.iteration_counts[(self.gate_id, target_id)] + self.inject_feedback(ctx, target_id, f"iteration {iteration}") + self.record_verdict_event(ctx, verdict, target_id) + return CompareBuilderNode( + verdict_fn=self.verdict_fn, + max_iterations=self.max_iterations, + ) + + self.record_verdict_event(ctx, verdict) + return CompareQAForkJoinNode(children=_qa_children(), node_id="fork_qa") + + +def _qa_children() -> list[ChildAgent]: + async def _health_check(ctx: GraphRunContext[FactoryState, FactoryDeps]) -> str: + return "health_check: pass" + + async def _code_review(ctx: GraphRunContext[FactoryState, FactoryDeps]) -> str: + return "code_review: pass" + + async def _adversarial(ctx: GraphRunContext[FactoryState, FactoryDeps]) -> str: + return "adversarial: pass" + + return [ + ChildAgent(name="health_checker", fn=_health_check), + ChildAgent(name="code_reviewer", fn=_code_review), + ChildAgent(name="adversarial_tester", fn=_adversarial), + ] + + +def build_pydantic_graph( + verdict_fn: Callable[[FactoryState], VerdictType] | None = None, + max_iterations: int = 3, +) -> tuple[ + GraphBuilder[FactoryState, FactoryDeps, CompareBuilderNode, HaltResult], + CompareBuilderNode, +]: + """Build the comparison workflow as pydantic-graph BaseNode subclasses.""" + builder: GraphBuilder[FactoryState, FactoryDeps, CompareBuilderNode, HaltResult] = ( + GraphBuilder( + name="comparison-workflow", + state_type=FactoryState, + deps_type=FactoryDeps, + input_type=CompareBuilderNode, + output_type=HaltResult, + ) + ) + builder.add_edge(builder.start_node, CompareBuilderNode) + builder.add(builder.node(CompareBuilderNode)) + builder.add(builder.node(CompareQAGateNode)) + builder.add(builder.node(CompareQAForkJoinNode)) + + start = CompareBuilderNode(verdict_fn=verdict_fn, max_iterations=max_iterations) + return builder, start + + +# ═══════════════════════════════════════════════════════════════════ +# Part 3: Comparison Utilities +# ═══════════════════════════════════════════════════════════════════ + + +def normalize_pg_events(state: FactoryState) -> list[TraceEvent]: + """Convert FactoryState.events to normalized TraceEvents.""" + return [ + TraceEvent( + node=ev.get("node", ""), + action=ev.get("action", ""), + detail={k: v for k, v in ev.items() if k not in ("node", "action")}, + ) + for ev in state.events + ] + + +def extract_active_nodes(events: list[TraceEvent]) -> list[str]: + """Extract ordered list of nodes that did work. + + Includes action "execute" events and fork children ("child_completed"). + Deduplicates consecutive repeats from the same node. + """ + nodes: list[str] = [] + for e in events: + if e.action in ("execute", "child_completed"): + node_name = e.detail.get("child", e.node) if e.action == "child_completed" else e.node + nodes.append(node_name) + return nodes + + +def extract_gate_verdicts(events: list[TraceEvent]) -> list[str]: + """Extract the sequence of gate verdict values.""" + return [ + e.detail.get("verdict", "") + for e in events + if e.action == "gate_verdict" + ] + + +def extract_fork_children(events: list[TraceEvent]) -> set[str]: + """Extract the set of children that ran in a fork/join.""" + children: set[str] = set() + for e in events: + if e.action == "fork_join_complete": + children.update(e.detail.get("children", [])) + elif e.action == "child_completed": + children.add(e.detail.get("child", "")) + return children + + +async def run_comparison( + verdict_sequence: list[VerdictType], + max_iterations: int = 3, +) -> dict[str, Any]: + """Run both engines with the same verdict sequence and return comparison data. + + verdict_sequence controls the gate: element i is the verdict for the i-th + gate evaluation (both engines use the same sequence). + """ + call_count_sim = 0 + + def sim_verdict(call_idx: int) -> VerdictType: + if call_idx < len(verdict_sequence): + return verdict_sequence[call_idx] + return VerdictType.PROCEED + + pg_call_count = 0 + + def pg_verdict(_state: FactoryState) -> VerdictType: + nonlocal pg_call_count + idx = pg_call_count + pg_call_count += 1 + if idx < len(verdict_sequence): + return verdict_sequence[idx] + return VerdictType.PROCEED + + sim_workflow = build_current_engine_workflow() + sim_events = simulate_current_engine(sim_workflow, sim_verdict, max_iterations) + + state = FactoryState() + deps = FactoryDeps(dry_run=True) + graph_builder, start = build_pydantic_graph(pg_verdict, max_iterations) + graph = graph_builder.build() + result = await graph.run(state=state, deps=deps, inputs=start) + pg_events = normalize_pg_events(state) + + sim_active = extract_active_nodes(sim_events) + pg_active = extract_active_nodes(pg_events) + sim_verdicts = extract_gate_verdicts(sim_events) + pg_verdicts = extract_gate_verdicts(pg_events) + sim_children = extract_fork_children(sim_events) + pg_children = extract_fork_children(pg_events) + + return { + "current_engine": { + "events": sim_events, + "active_nodes": sim_active, + "gate_verdicts": sim_verdicts, + "fork_children": sim_children, + }, + "pydantic_graph": { + "events": pg_events, + "active_nodes": pg_active, + "gate_verdicts": pg_verdicts, + "fork_children": pg_children, + "result": result, + }, + "match": { + "active_nodes": set(sim_active) == set(pg_active), + "gate_verdicts": sim_verdicts == pg_verdicts, + "fork_children": sim_children == pg_children, + }, + "mermaid": graph.render(), + } diff --git a/experiments/pydantic-graph-prototype/src/pg_factory/deps.py b/experiments/pydantic-graph-prototype/src/pg_factory/deps.py new file mode 100644 index 000000000..fe0f29444 --- /dev/null +++ b/experiments/pydantic-graph-prototype/src/pg_factory/deps.py @@ -0,0 +1,22 @@ +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable + + +def _noop_emitter(event: dict[str, Any]) -> None: + pass + + +@dataclass +class FactoryDeps: + """Immutable dependencies injected into graph nodes via GraphRunContext.deps. + + Maps from the current executor's constructor args: + - project_path: root path of the project being operated on + - dry_run: when True, nodes simulate execution without side effects + - event_emitter: callback for structured event emission + """ + + project_path: Path = field(default_factory=lambda: Path(".")) + dry_run: bool = False + event_emitter: Callable[[dict[str, Any]], None] = field(default=_noop_emitter) diff --git a/experiments/pydantic-graph-prototype/src/pg_factory/graphs/__init__.py b/experiments/pydantic-graph-prototype/src/pg_factory/graphs/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/experiments/pydantic-graph-prototype/src/pg_factory/graphs/study.py b/experiments/pydantic-graph-prototype/src/pg_factory/graphs/study.py new file mode 100644 index 000000000..7a6358fd1 --- /dev/null +++ b/experiments/pydantic-graph-prototype/src/pg_factory/graphs/study.py @@ -0,0 +1,33 @@ +"""Study chain graph assembly. + +Assembles the 4-node study chain into a pydantic-graph Graph: + GraphUpdateNode → StudyNode → GraphExplorerNode → ConcatStudyNode → End +""" + +from pydantic_graph import GraphBuilder + +from pg_factory.deps import FactoryDeps +from pg_factory.nodes.study import ( + ConcatStudyNode, + GraphExplorerNode, + GraphUpdateNode, + StudyNode, +) +from pg_factory.state import FactoryState +from pg_factory.verdicts import HaltResult + + +def build_study_graph() -> GraphBuilder[FactoryState, FactoryDeps, GraphUpdateNode, HaltResult]: + builder: GraphBuilder[FactoryState, FactoryDeps, GraphUpdateNode, HaltResult] = GraphBuilder( + name="study-chain", + state_type=FactoryState, + deps_type=FactoryDeps, + input_type=GraphUpdateNode, + output_type=HaltResult, + ) + builder.add_edge(builder.start_node, GraphUpdateNode) + builder.add(builder.node(GraphUpdateNode)) + builder.add(builder.node(StudyNode)) + builder.add(builder.node(GraphExplorerNode)) + builder.add(builder.node(ConcatStudyNode)) + return builder diff --git a/experiments/pydantic-graph-prototype/src/pg_factory/nodes/__init__.py b/experiments/pydantic-graph-prototype/src/pg_factory/nodes/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/experiments/pydantic-graph-prototype/src/pg_factory/nodes/gates.py b/experiments/pydantic-graph-prototype/src/pg_factory/nodes/gates.py new file mode 100644 index 000000000..81aa6cbb3 --- /dev/null +++ b/experiments/pydantic-graph-prototype/src/pg_factory/nodes/gates.py @@ -0,0 +1,85 @@ +"""Gate/verdict routing nodes — port of GateNode verdict routing from primitives.py + executor.py. + +GateBaseNode provides the core gate pattern: + - Configurable verdict function determines PROCEED / RELOOP / HALT + - Iteration counting: tracks (gate_id, target_id) counts in FactoryState + - Feedback injection: writes to node_feedback so reloop targets can read context + - Max-iteration halt: returns End(HaltResult) when count exceeds max_iterations + +Structural difference from the original: + executor.py's _execute_gate uses runtime edge matching — it looks up the next + node by string ID from the workflow's edge list based on VerdictType. In + pydantic-graph, the gate's run() return type union declares all possible + successors at definition time, and the graph enforces valid routing via the + type system. Invalid edges become type errors, not runtime crashes. +""" + +from typing import Callable + +from pydantic_graph import BaseNode, End, GraphRunContext + +from pg_factory.deps import FactoryDeps +from pg_factory.state import FactoryState +from pg_factory.verdicts import HaltResult, VerdictType + + +class GateBaseNode(BaseNode[FactoryState, FactoryDeps, HaltResult]): + """Abstract base for gate nodes with iteration counting and feedback injection. + + Subclasses define run() with a concrete return union of their + proceed/reloop target node types plus End[HaltResult]. + """ + + def __init__( + self, + gate_id: str = "", + max_iterations: int = 3, + verdict_fn: Callable[[FactoryState], VerdictType] | None = None, + ) -> None: + self.gate_id = gate_id + self.max_iterations = max_iterations + self.verdict_fn = verdict_fn + + def evaluate_verdict(self, state: FactoryState) -> VerdictType: + if self.verdict_fn: + return self.verdict_fn(state) + return VerdictType.PROCEED + + def check_and_increment_iteration( + self, ctx: GraphRunContext[FactoryState, FactoryDeps], target_id: str + ) -> End[HaltResult] | None: + """Increment iteration count; return End if max_iterations exceeded.""" + key = (self.gate_id, target_id) + count = ctx.state.iteration_counts.get(key, 0) + 1 + ctx.state.iteration_counts[key] = count + if count > self.max_iterations: + return End( + HaltResult( + reason=f"max_iterations ({self.max_iterations}) exceeded " + f"for {self.gate_id} -> {target_id}" + ) + ) + return None + + def inject_feedback( + self, + ctx: GraphRunContext[FactoryState, FactoryDeps], + target_id: str, + feedback: str, + ) -> None: + ctx.state.node_feedback[target_id] = feedback + + def record_verdict_event( + self, + ctx: GraphRunContext[FactoryState, FactoryDeps], + verdict: VerdictType, + target: str | None = None, + ) -> None: + ctx.state.events.append( + { + "node": self.gate_id or self.__class__.__name__, + "action": "gate_verdict", + "verdict": verdict.value, + "target": target, + } + ) diff --git a/experiments/pydantic-graph-prototype/src/pg_factory/nodes/parallel.py b/experiments/pydantic-graph-prototype/src/pg_factory/nodes/parallel.py new file mode 100644 index 000000000..462cb68a7 --- /dev/null +++ b/experiments/pydantic-graph-prototype/src/pg_factory/nodes/parallel.py @@ -0,0 +1,111 @@ +"""Parallel fork/join node — wraps asyncio.gather inside a single BaseNode. + +Ports the deep-QA fork/join pattern from executor.py's _execute_fork method: + fork_qa → [health_checker, code_reviewer, adversarial_tester] → join_qa + +Structural difference from the original: + executor.py uses ForkNode + JoinNode as separate graph primitives with the + executor dispatching branches via asyncio.gather. In pydantic-graph, parallel + execution is encapsulated inside a single ForkJoinNode's run() method — the + graph sees one node, while internally asyncio.gather runs all children + concurrently. This trades Mermaid visual fan-out for composition simplicity: + the Mermaid diagram shows ForkJoinNode as a single node rather than a + fork → branches → join fan-out. A custom Mermaid post-processor could expand + it, but for the prototype the single-node representation is acceptable since + the internal parallelism is observable through state.events timing data. +""" + +import asyncio +import time +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import Any + +from pydantic_graph import BaseNode, GraphRunContext + +from pg_factory.deps import FactoryDeps +from pg_factory.state import FactoryState +from pg_factory.verdicts import HaltResult + + +@dataclass +class ChildAgent: + """A simulated agent invocation to be run as one branch of a fork/join.""" + + name: str + fn: Callable[[GraphRunContext[FactoryState, FactoryDeps]], Awaitable[str]] + + +class ForkJoinNode(BaseNode[FactoryState, FactoryDeps, HaltResult]): + """Base node that runs multiple child agents concurrently via asyncio.gather. + + Subclasses implement run() with a concrete return type to declare the + next node in the graph. The run() body calls execute_children() then + returns the successor. + """ + + def __init__(self, children: list[ChildAgent], node_id: str = "") -> None: + self.children = children + self.node_id = node_id or self.__class__.__name__ + + async def execute_children( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> dict[str, dict[str, Any]]: + """Run all children concurrently via asyncio.gather. + + Records per-child outputs to ctx.state.node_outputs and timing + events to ctx.state.events. Returns per-child result dicts. + """ + results: dict[str, dict[str, Any]] = {} + + async def run_child(child: ChildAgent) -> None: + start = time.monotonic() + try: + output = await child.fn(ctx) + elapsed_ms = (time.monotonic() - start) * 1000 + ctx.state.node_outputs[child.name] = output + ctx.state.events.append( + { + "node": self.node_id, + "action": "child_completed", + "child": child.name, + "duration_ms": round(elapsed_ms, 2), + } + ) + results[child.name] = { + "output": output, + "duration_ms": elapsed_ms, + "success": True, + } + except Exception as exc: + elapsed_ms = (time.monotonic() - start) * 1000 + ctx.state.events.append( + { + "node": self.node_id, + "action": "child_failed", + "child": child.name, + "error": str(exc), + "duration_ms": round(elapsed_ms, 2), + } + ) + results[child.name] = { + "output": "", + "duration_ms": elapsed_ms, + "success": False, + "error": str(exc), + } + + fork_start = time.monotonic() + await asyncio.gather(*(run_child(c) for c in self.children)) + total_ms = (time.monotonic() - fork_start) * 1000 + + ctx.state.events.append( + { + "node": self.node_id, + "action": "fork_join_complete", + "children": [c.name for c in self.children], + "total_duration_ms": round(total_ms, 2), + } + ) + + return results diff --git a/experiments/pydantic-graph-prototype/src/pg_factory/nodes/study.py b/experiments/pydantic-graph-prototype/src/pg_factory/nodes/study.py new file mode 100644 index 000000000..2612f161b --- /dev/null +++ b/experiments/pydantic-graph-prototype/src/pg_factory/nodes/study.py @@ -0,0 +1,154 @@ +"""Study chain nodes — port of definitions.py _study_subgraph (lines 108-158). + +Four BaseNode subclasses forming a linear chain: + GraphUpdateNode → StudyNode → GraphExplorerNode → ConcatStudyNode + +Structural difference from the original: + definitions.py constructs a dict[str, Node] + list[Edge] where edges are explicit + data objects wired by string IDs. In pydantic-graph, edges are implicit — each + node's run() return type annotation declares its successor, and Graph infers the + topology at construction time. This eliminates an entire class of wiring bugs + (misspelled edge targets, dangling nodes) at the cost of requiring all successor + types to be importable at definition time. +""" + +import asyncio + +from pydantic_graph import BaseNode, End, GraphRunContext + +from pg_factory.deps import FactoryDeps +from pg_factory.state import FactoryState +from pg_factory.verdicts import HaltResult + + +class ConcatStudyNode(BaseNode[FactoryState, FactoryDeps, HaltResult]): + """Concatenates observation files into study-combined.md. + + Maps from: FnNode(id="concat_study", command="cat ... > study-combined.md") + """ + + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> End[HaltResult]: + pp = ctx.deps.project_path + obs = pp / ".factory" / "strategy" / "observations.md" + graph_ctx = pp / ".factory" / "strategy" / "graph-context.md" + combined = pp / ".factory" / "strategy" / "study-combined.md" + + if ctx.deps.dry_run: + output = "" + for f in (obs, graph_ctx): + if f.exists(): + output += f.read_text() + else: + output += f"[mock] {f.name} content\n" + combined.parent.mkdir(parents=True, exist_ok=True) + combined.write_text(output) + else: + cmd = f"cat {obs} {graph_ctx} > {combined}" + proc = await asyncio.create_subprocess_shell( + cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE + ) + await proc.communicate() + + ctx.state.node_outputs["ConcatStudyNode"] = str(combined) + ctx.state.events.append( + {"node": "ConcatStudyNode", "action": "concat", "output": str(combined)} + ) + return End(HaltResult(reason="study_complete")) + + +class GraphExplorerNode(BaseNode[FactoryState, FactoryDeps, HaltResult]): + """Simulates agent invocation for graph exploration. + + Maps from: AgentNode(id="graph_explorer", role=RESEARCHER) + In dry_run mode, writes mock graph-context output instead of invoking an agent. + """ + + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> ConcatStudyNode: + pp = ctx.deps.project_path + output_path = pp / ".factory" / "strategy" / "graph-context.md" + + if ctx.deps.dry_run: + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + "# Graph Context\n\nMock graph exploration output for dry_run mode.\n" + ) + result = "dry_run: mock graph-context written" + else: + proc = await asyncio.create_subprocess_shell( + f"factory agent researcher --task 'explore graph' --project {pp}", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, _ = await proc.communicate() + result = stdout.decode() + + ctx.state.node_outputs["GraphExplorerNode"] = result + ctx.state.events.append( + {"node": "GraphExplorerNode", "action": "explore", "output": result} + ) + return ConcatStudyNode() + + +class StudyNode(BaseNode[FactoryState, FactoryDeps, HaltResult]): + """Runs 'factory study {project_path}'. + + Maps from: Study(id="study", command="factory study {project_path}") + """ + + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> GraphExplorerNode: + pp = ctx.deps.project_path + + if ctx.deps.dry_run: + obs_path = pp / ".factory" / "strategy" / "observations.md" + obs_path.parent.mkdir(parents=True, exist_ok=True) + obs_path.write_text("# Observations\n\nMock study output for dry_run mode.\n") + result = "dry_run: mock observations written" + else: + proc = await asyncio.create_subprocess_shell( + f"factory study {pp}", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, _ = await proc.communicate() + result = stdout.decode() + + ctx.state.node_outputs["StudyNode"] = result + ctx.state.events.append( + {"node": "StudyNode", "action": "study", "output": result} + ) + return GraphExplorerNode() + + +class GraphUpdateNode(BaseNode[FactoryState, FactoryDeps, HaltResult]): + """Runs 'factory graph update {project_path}'. + + Maps from: FnNode(id="graph_update", command="factory graph update {project_path}") + """ + + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> StudyNode: + pp = ctx.deps.project_path + + if ctx.deps.dry_run: + result = "dry_run: graph update skipped" + else: + proc = await asyncio.create_subprocess_shell( + f"factory graph update {pp}", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, _ = await proc.communicate() + result = stdout.decode() + + ctx.state.node_outputs["GraphUpdateNode"] = result + ctx.state.events.append( + {"node": "GraphUpdateNode", "action": "graph_update", "output": result} + ) + return StudyNode() diff --git a/experiments/pydantic-graph-prototype/src/pg_factory/state.py b/experiments/pydantic-graph-prototype/src/pg_factory/state.py new file mode 100644 index 000000000..bebd18874 --- /dev/null +++ b/experiments/pydantic-graph-prototype/src/pg_factory/state.py @@ -0,0 +1,29 @@ +from dataclasses import dataclass, field +from typing import Any + + +@dataclass +class FactoryState: + """Mutable state threaded through a pydantic-graph workflow via GraphRunContext.state. + + Maps from the current executor's mutable fields: + - iteration_counts: tracks gate reloop counts per (gate_id, target_id) pair + - node_feedback: carries feedback from gate verdicts to reloop target nodes + - node_outputs: stores each node's output keyed by node class name + - completed_files: tracks which logical file artifacts have been produced + - events: append-only event log for observability + """ + + iteration_counts: dict[tuple[str, str], int] = field( + default_factory=lambda: dict[tuple[str, str], int]() + ) + node_feedback: dict[str, str] = field( + default_factory=lambda: dict[str, str]() + ) + node_outputs: dict[str, str] = field( + default_factory=lambda: dict[str, str]() + ) + completed_files: set[str] = field(default_factory=lambda: set[str]()) + events: list[dict[str, Any]] = field( + default_factory=lambda: list[dict[str, Any]]() + ) diff --git a/experiments/pydantic-graph-prototype/src/pg_factory/verdicts.py b/experiments/pydantic-graph-prototype/src/pg_factory/verdicts.py new file mode 100644 index 000000000..e00b88345 --- /dev/null +++ b/experiments/pydantic-graph-prototype/src/pg_factory/verdicts.py @@ -0,0 +1,22 @@ +from dataclasses import dataclass +from enum import Enum + + +class VerdictType(Enum): + """Maps from the current engine's VerdictType enum.""" + + PROCEED = "proceed" + RELOOP = "reloop" + HALT = "halt" + + +@dataclass(frozen=True) +class HaltResult: + """Terminal result type for End[HaltResult] returns in pydantic-graph. + + When a node or gate determines execution should stop, it returns + End(HaltResult(reason="...")) to terminate the graph run. + """ + + reason: str + verdict: VerdictType = VerdictType.HALT diff --git a/experiments/pydantic-graph-prototype/tests/__init__.py b/experiments/pydantic-graph-prototype/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/experiments/pydantic-graph-prototype/tests/_future_annotations_nodes.py b/experiments/pydantic-graph-prototype/tests/_future_annotations_nodes.py new file mode 100644 index 000000000..bb03bdf6b --- /dev/null +++ b/experiments/pydantic-graph-prototype/tests/_future_annotations_nodes.py @@ -0,0 +1,41 @@ +"""Module that uses `from __future__ import annotations` with pydantic-graph BaseNode subclasses. + +This is the critical risk test: pydantic-graph uses runtime type-hint introspection +to infer graph edges from return types. `from __future__ import annotations` turns all +annotations into strings (PEP 563), which could break this introspection. + +This module must be a separate file (not inline in the test) because `from __future__` +only applies at the module level. +""" + +from __future__ import annotations + +from pydantic_graph import BaseNode, End, GraphRunContext + +from pg_factory.deps import FactoryDeps +from pg_factory.state import FactoryState +from pg_factory.verdicts import HaltResult + + +class AlphaNode(BaseNode[FactoryState, FactoryDeps, HaltResult]): + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> BetaNode: + ctx.state.node_outputs["AlphaNode"] = "alpha_done" + return BetaNode() + + +class BetaNode(BaseNode[FactoryState, FactoryDeps, HaltResult]): + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> GammaNode: + ctx.state.node_outputs["BetaNode"] = "beta_done" + return GammaNode() + + +class GammaNode(BaseNode[FactoryState, FactoryDeps, HaltResult]): + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> End[HaltResult]: + ctx.state.node_outputs["GammaNode"] = "gamma_done" + return End(HaltResult(reason="chain_complete")) diff --git a/experiments/pydantic-graph-prototype/tests/conftest.py b/experiments/pydantic-graph-prototype/tests/conftest.py new file mode 100644 index 000000000..5ab320b43 --- /dev/null +++ b/experiments/pydantic-graph-prototype/tests/conftest.py @@ -0,0 +1,19 @@ +from pathlib import Path + +import pytest + +from pg_factory.deps import FactoryDeps +from pg_factory.state import FactoryState + + +@pytest.fixture +def factory_state() -> FactoryState: + return FactoryState() + + +@pytest.fixture +def factory_deps(tmp_path: Path) -> FactoryDeps: + return FactoryDeps( + project_path=tmp_path, + dry_run=True, + ) diff --git a/experiments/pydantic-graph-prototype/tests/test_compare.py b/experiments/pydantic-graph-prototype/tests/test_compare.py new file mode 100644 index 000000000..9c96b2ce0 --- /dev/null +++ b/experiments/pydantic-graph-prototype/tests/test_compare.py @@ -0,0 +1,271 @@ +"""Tests for the dual-engine comparison harness. + +Verifies that both engines produce equivalent behavior for: + - PROCEED path (sequential + gate + fork/join) + - RELOOP path (gate routes back to builder) + - HALT path (gate stops execution) + - Max iterations (reloop exhaustion) + - End-to-end comparison runner +""" + +from pg_factory.compare import ( + build_current_engine_workflow, + build_pydantic_graph, + extract_active_nodes, + extract_fork_children, + extract_gate_verdicts, + normalize_pg_events, + run_comparison, + simulate_current_engine, +) +from pg_factory.deps import FactoryDeps +from pg_factory.state import FactoryState +from pg_factory.verdicts import HaltResult, VerdictType + + +# ── PROCEED path ───────────────────────────────────────────────── + + +async def test_proceed_active_nodes_match() -> None: + """Both engines execute the same set of nodes on PROCEED.""" + result = await run_comparison([VerdictType.PROCEED]) + + assert result["match"]["active_nodes"] + assert result["match"]["gate_verdicts"] + assert result["match"]["fork_children"] + + +async def test_proceed_gate_verdict() -> None: + """Both engines record a single 'proceed' verdict.""" + result = await run_comparison([VerdictType.PROCEED]) + + assert result["current_engine"]["gate_verdicts"] == ["proceed"] + assert result["pydantic_graph"]["gate_verdicts"] == ["proceed"] + + +async def test_proceed_fork_children() -> None: + """Both engines run all 3 QA agents in the fork/join.""" + result = await run_comparison([VerdictType.PROCEED]) + + expected = {"health_checker", "code_reviewer", "adversarial_tester"} + assert result["current_engine"]["fork_children"] == expected + assert result["pydantic_graph"]["fork_children"] == expected + + +async def test_proceed_pg_result() -> None: + """pydantic-graph returns HaltResult with reason 'qa_complete'.""" + result = await run_comparison([VerdictType.PROCEED]) + + pg_result = result["pydantic_graph"]["result"] + assert isinstance(pg_result, HaltResult) + assert pg_result.reason == "qa_complete" + + +# ── RELOOP path ────────────────────────────────────────────────── + + +async def test_reloop_then_proceed_nodes_match() -> None: + """RELOOP once then PROCEED — both engines execute builder twice.""" + result = await run_comparison([VerdictType.RELOOP, VerdictType.PROCEED]) + + sim_active = result["current_engine"]["active_nodes"] + pg_active = result["pydantic_graph"]["active_nodes"] + + assert sim_active.count("builder") == 2 + assert pg_active.count("builder") == 2 + assert result["match"]["active_nodes"] + + +async def test_reloop_then_proceed_verdicts() -> None: + """RELOOP then PROCEED produces reloop + proceed verdict sequence.""" + result = await run_comparison([VerdictType.RELOOP, VerdictType.PROCEED]) + + assert result["current_engine"]["gate_verdicts"] == ["reloop", "proceed"] + assert result["pydantic_graph"]["gate_verdicts"] == ["reloop", "proceed"] + + +async def test_reloop_twice_then_proceed() -> None: + """Two RELOOPs then PROCEED — builder executes 3 times total.""" + result = await run_comparison( + [VerdictType.RELOOP, VerdictType.RELOOP, VerdictType.PROCEED] + ) + + sim_active = result["current_engine"]["active_nodes"] + pg_active = result["pydantic_graph"]["active_nodes"] + + assert sim_active.count("builder") == 3 + assert pg_active.count("builder") == 3 + assert result["match"]["gate_verdicts"] + assert result["match"]["fork_children"] + + +# ── HALT path ──────────────────────────────────────────────────── + + +async def test_halt_stops_execution() -> None: + """HALT stops execution before fork/join.""" + result = await run_comparison([VerdictType.HALT]) + + assert result["current_engine"]["gate_verdicts"] == ["halt"] + assert result["pydantic_graph"]["gate_verdicts"] == ["halt"] + + assert result["current_engine"]["fork_children"] == set() + assert result["pydantic_graph"]["fork_children"] == set() + + +async def test_halt_pg_result() -> None: + """pydantic-graph returns HaltResult with reason 'gate_halted'.""" + result = await run_comparison([VerdictType.HALT]) + + pg_result = result["pydantic_graph"]["result"] + assert isinstance(pg_result, HaltResult) + assert pg_result.reason == "gate_halted" + + +# ── Max iterations ─────────────────────────────────────────────── + + +async def test_max_iterations_halt() -> None: + """Continuous RELOOP hits max_iterations and halts in both engines.""" + all_reloop = [VerdictType.RELOOP] * 10 + result = await run_comparison(all_reloop, max_iterations=2) + + sim_verdicts = result["current_engine"]["gate_verdicts"] + pg_verdicts = result["pydantic_graph"]["gate_verdicts"] + + assert sim_verdicts[-1] == "halt" + assert pg_verdicts[-1] == "halt" + + assert result["current_engine"]["fork_children"] == set() + assert result["pydantic_graph"]["fork_children"] == set() + + +# ── Current engine simulation unit tests ───────────────────────── + + +def test_sim_workflow_has_all_nodes() -> None: + """Simulated workflow has all 7 nodes.""" + wf = build_current_engine_workflow() + expected = { + "builder", "qa_gate", "fork_qa", + "health_checker", "code_reviewer", "adversarial_tester", + "join_qa", + } + assert set(wf.nodes.keys()) == expected + + +def test_sim_proceed_event_sequence() -> None: + """PROCEED produces: builder execute, gate verdict, fork, 3 children, join.""" + wf = build_current_engine_workflow() + events = simulate_current_engine(wf, lambda _: VerdictType.PROCEED) + + actions = [(e.node, e.action) for e in events] + assert actions[0] == ("builder", "execute") + assert actions[1] == ("qa_gate", "gate_verdict") + assert actions[1 + 1] == ("fork_qa", "fork") + + child_executes = [(n, a) for n, a in actions if a == "execute" and n != "builder"] + assert len(child_executes) == 3 + + assert actions[-2] == ("fork_qa", "fork_join_complete") + assert actions[-1] == ("join_qa", "join") + + +def test_sim_halt_no_fork() -> None: + """HALT produces only builder execute + gate halt verdict.""" + wf = build_current_engine_workflow() + events = simulate_current_engine(wf, lambda _: VerdictType.HALT) + + assert len(events) == 2 + assert events[0].action == "execute" + assert events[1].detail["verdict"] == "halt" + + +# ── pydantic-graph unit tests ──────────────────────────────────── + + +async def test_pg_proceed_events( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """pydantic-graph PROCEED records builder execute + gate verdict + children.""" + graph_builder, start = build_pydantic_graph( + lambda _: VerdictType.PROCEED + ) + graph = graph_builder.build() + result = await graph.run(state=factory_state, deps=factory_deps, inputs=start) + + assert isinstance(result, HaltResult) + assert result.reason == "qa_complete" + + events = normalize_pg_events(factory_state) + actions = [e.action for e in events] + + assert "execute" in actions + assert "gate_verdict" in actions + assert actions.count("child_completed") == 3 + assert "fork_join_complete" in actions + + +async def test_pg_halt_events( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """pydantic-graph HALT records builder execute + gate halt verdict only.""" + graph_builder, start = build_pydantic_graph( + lambda _: VerdictType.HALT + ) + graph = graph_builder.build() + result = await graph.run(state=factory_state, deps=factory_deps, inputs=start) + + assert isinstance(result, HaltResult) + assert result.reason == "gate_halted" + + events = normalize_pg_events(factory_state) + assert len(events) == 2 + assert events[0].action == "execute" + assert events[1].detail["verdict"] == "halt" + + +# ── Mermaid output ─────────────────────────────────────────────── + + +async def test_mermaid_contains_all_nodes() -> None: + """Mermaid diagram from pydantic-graph shows all workflow nodes.""" + result = await run_comparison([VerdictType.PROCEED]) + + mermaid = result["mermaid"] + assert "CompareBuilderNode" in mermaid + assert "CompareQAGateNode" in mermaid + assert "CompareQAForkJoinNode" in mermaid + + +async def test_mermaid_shows_gate_branching() -> None: + """Mermaid diagram shows gate routing edges.""" + graph_builder, _ = build_pydantic_graph() + graph = graph_builder.build() + mermaid = graph.render() + + assert "CompareBuilderNode --> CompareQAGateNode" in mermaid + assert "decision" in mermaid + + +# ── End-to-end ─────────────────────────────────────────────────── + + +async def test_comparison_returns_all_fields() -> None: + """run_comparison returns all expected fields.""" + result = await run_comparison([VerdictType.PROCEED]) + + assert "current_engine" in result + assert "pydantic_graph" in result + assert "match" in result + assert "mermaid" in result + + assert "events" in result["current_engine"] + assert "active_nodes" in result["current_engine"] + assert "gate_verdicts" in result["current_engine"] + assert "fork_children" in result["current_engine"] + + assert "events" in result["pydantic_graph"] + assert "result" in result["pydantic_graph"] diff --git a/experiments/pydantic-graph-prototype/tests/test_future_annotations.py b/experiments/pydantic-graph-prototype/tests/test_future_annotations.py new file mode 100644 index 000000000..ab1448f38 --- /dev/null +++ b/experiments/pydantic-graph-prototype/tests/test_future_annotations.py @@ -0,0 +1,67 @@ +"""Test that `from __future__ import annotations` works with pydantic-graph's +runtime type-hint introspection for edge inference. + +This is identified as the #1 risk to the prototype (research-similar). +If this fails, the entire migration approach changes. +""" + +from pydantic_graph import GraphBuilder + +from pg_factory.deps import FactoryDeps +from pg_factory.state import FactoryState +from pg_factory.verdicts import HaltResult +from tests._future_annotations_nodes import AlphaNode, BetaNode, GammaNode + + +def _build_future_graph() -> GraphBuilder[FactoryState, FactoryDeps, AlphaNode, HaltResult]: + builder: GraphBuilder[FactoryState, FactoryDeps, AlphaNode, HaltResult] = GraphBuilder( + name="future-annotations-test", + state_type=FactoryState, + deps_type=FactoryDeps, + input_type=AlphaNode, + output_type=HaltResult, + ) + builder.add_edge(builder.start_node, AlphaNode) + builder.add(builder.node(AlphaNode)) + builder.add(builder.node(BetaNode)) + builder.add(builder.node(GammaNode)) + return builder + + +async def test_future_annotations_graph_builds() -> None: + """Graph() should correctly infer edges even when annotations are stringified.""" + graph = _build_future_graph().build() + assert graph is not None + + +async def test_future_annotations_graph_runs( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """Full execution should work: Alpha -> Beta -> Gamma -> End.""" + graph = _build_future_graph().build() + result = await graph.run(state=factory_state, deps=factory_deps, inputs=AlphaNode()) + + assert isinstance(result, HaltResult) + assert result.reason == "chain_complete" + assert factory_state.node_outputs["AlphaNode"] == "alpha_done" + assert factory_state.node_outputs["BetaNode"] == "beta_done" + assert factory_state.node_outputs["GammaNode"] == "gamma_done" + + +async def test_future_annotations_mermaid_has_all_nodes() -> None: + """Mermaid rendering should include all 3 node names despite stringified annotations.""" + graph = _build_future_graph().build() + mermaid = graph.render() + assert "AlphaNode" in mermaid + assert "BetaNode" in mermaid + assert "GammaNode" in mermaid + + +async def test_future_annotations_edge_topology() -> None: + """Verify the inferred edge topology matches: Alpha -> Beta -> Gamma -> End.""" + graph = _build_future_graph().build() + mermaid = graph.render() + assert "AlphaNode --> BetaNode" in mermaid + assert "BetaNode --> GammaNode" in mermaid + assert "GammaNode --> [*]" in mermaid diff --git a/experiments/pydantic-graph-prototype/tests/test_gate_routing.py b/experiments/pydantic-graph-prototype/tests/test_gate_routing.py new file mode 100644 index 000000000..36c558bd5 --- /dev/null +++ b/experiments/pydantic-graph-prototype/tests/test_gate_routing.py @@ -0,0 +1,365 @@ +"""Tests for gate/verdict routing: PROCEED, RELOOP, HALT paths. + +Models the builder -> gate_qa -> (RELOOP -> builder | PROCEED -> next | HALT -> End) +pattern from the factory's definitions.py gate routing. +""" + +from typing import Callable + +from pydantic_graph import BaseNode, End, EndMarker, GraphBuilder, GraphRunContext + +from pg_factory.deps import FactoryDeps +from pg_factory.nodes.gates import GateBaseNode +from pg_factory.state import FactoryState +from pg_factory.verdicts import HaltResult, VerdictType + + +# ── Test workflow nodes ───────────────────────────────────────── + + +class MockNextNode(BaseNode[FactoryState, FactoryDeps, HaltResult]): + """Terminal node reached after gate PROCEED.""" + + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> End[HaltResult]: + ctx.state.node_outputs["MockNextNode"] = "executed" + ctx.state.events.append({"node": "MockNextNode", "action": "next"}) + return End(HaltResult(reason="proceed_complete")) + + +class MockBuilderNode(BaseNode[FactoryState, FactoryDeps, HaltResult]): + """Simulated builder that feeds into the QA gate.""" + + def __init__( + self, + gate_verdict_fn: Callable[[FactoryState], VerdictType] | None = None, + gate_max_iterations: int = 3, + ) -> None: + self.gate_verdict_fn = gate_verdict_fn + self.gate_max_iterations = gate_max_iterations + + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> "QAGateNode": + feedback = ctx.state.node_feedback.get("MockBuilderNode", "") + output = f"built (feedback: {feedback})" if feedback else "built" + ctx.state.node_outputs["MockBuilderNode"] = output + ctx.state.events.append( + { + "node": "MockBuilderNode", + "action": "build", + "feedback_received": feedback, + } + ) + return QAGateNode( + gate_id="qa_gate", + verdict_fn=self.gate_verdict_fn, + max_iterations=self.gate_max_iterations, + ) + + +class QAGateNode(GateBaseNode): + """Concrete gate modeling the QA verdict pattern.""" + + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> MockBuilderNode | MockNextNode | End[HaltResult]: + verdict = self.evaluate_verdict(ctx.state) + + if verdict == VerdictType.HALT: + self.record_verdict_event(ctx, verdict) + return End(HaltResult(reason="gate_halted")) + + if verdict == VerdictType.RELOOP: + target_id = "MockBuilderNode" + halt = self.check_and_increment_iteration(ctx, target_id) + if halt is not None: + self.record_verdict_event(ctx, VerdictType.HALT, target_id) + return halt + iteration = ctx.state.iteration_counts[(self.gate_id, target_id)] + self.inject_feedback( + ctx, target_id, f"iteration {iteration}: needs improvement" + ) + self.record_verdict_event(ctx, verdict, target_id) + return MockBuilderNode( + gate_verdict_fn=self.verdict_fn, + gate_max_iterations=self.max_iterations, + ) + + self.record_verdict_event(ctx, verdict) + return MockNextNode() + + +# ── Graph builder ─────────────────────────────────────────────── + + +def _build_gate_graph( + verdict_fn: Callable[[FactoryState], VerdictType] | None = None, + max_iterations: int = 3, +) -> tuple[ + GraphBuilder[FactoryState, FactoryDeps, MockBuilderNode, HaltResult], + MockBuilderNode, +]: + builder: GraphBuilder[FactoryState, FactoryDeps, MockBuilderNode, HaltResult] = ( + GraphBuilder( + name="gate-routing-test", + state_type=FactoryState, + deps_type=FactoryDeps, + input_type=MockBuilderNode, + output_type=HaltResult, + ) + ) + builder.add_edge(builder.start_node, MockBuilderNode) + builder.add(builder.node(MockBuilderNode)) + builder.add(builder.node(QAGateNode)) + builder.add(builder.node(MockNextNode)) + + start_node = MockBuilderNode( + gate_verdict_fn=verdict_fn, + gate_max_iterations=max_iterations, + ) + return builder, start_node + + +# ── Tests ─────────────────────────────────────────────────────── + + +async def test_proceed_path( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """PROCEED: gate returns MockNextNode, execution continues to End.""" + graph_builder, start = _build_gate_graph( + verdict_fn=lambda _: VerdictType.PROCEED + ) + graph = graph_builder.build() + + result = await graph.run(state=factory_state, deps=factory_deps, inputs=start) + + assert isinstance(result, HaltResult) + assert result.reason == "proceed_complete" + + node_order = [e["node"] for e in factory_state.events] + assert node_order == ["MockBuilderNode", "qa_gate", "MockNextNode"] + + verdict_events = [ + e for e in factory_state.events if e.get("action") == "gate_verdict" + ] + assert len(verdict_events) == 1 + assert verdict_events[0]["verdict"] == "proceed" + + +async def test_halt_path( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """HALT: gate returns End(HaltResult) immediately.""" + graph_builder, start = _build_gate_graph( + verdict_fn=lambda _: VerdictType.HALT + ) + graph = graph_builder.build() + + result = await graph.run(state=factory_state, deps=factory_deps, inputs=start) + + assert isinstance(result, HaltResult) + assert result.reason == "gate_halted" + + node_order = [e["node"] for e in factory_state.events] + assert node_order == ["MockBuilderNode", "qa_gate"] + + +async def test_reloop_then_max_iterations_halt( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """RELOOP: builder re-executes with feedback; max_iterations halts.""" + max_iter = 2 + graph_builder, start = _build_gate_graph( + verdict_fn=lambda _: VerdictType.RELOOP, + max_iterations=max_iter, + ) + graph = graph_builder.build() + + result = await graph.run(state=factory_state, deps=factory_deps, inputs=start) + + assert isinstance(result, HaltResult) + assert "max_iterations" in result.reason + assert f"({max_iter})" in result.reason + + builder_events = [ + e for e in factory_state.events if e["node"] == "MockBuilderNode" + ] + assert len(builder_events) == max_iter + 1 + + key = ("qa_gate", "MockBuilderNode") + assert factory_state.iteration_counts[key] == max_iter + 1 + + +async def test_reloop_iteration_count_increments( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """Iteration count increments on each RELOOP pass.""" + call_count = 0 + + def reloop_twice(state: FactoryState) -> VerdictType: + nonlocal call_count + call_count += 1 + if call_count <= 2: + return VerdictType.RELOOP + return VerdictType.PROCEED + + graph_builder, start = _build_gate_graph(verdict_fn=reloop_twice) + graph = graph_builder.build() + + result = await graph.run(state=factory_state, deps=factory_deps, inputs=start) + + assert isinstance(result, HaltResult) + assert result.reason == "proceed_complete" + + key = ("qa_gate", "MockBuilderNode") + assert factory_state.iteration_counts[key] == 2 + + builder_events = [ + e for e in factory_state.events if e["node"] == "MockBuilderNode" + ] + assert len(builder_events) == 3 + + +async def test_feedback_accessible_on_reentry( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """node_feedback is accessible to the reloop target on re-entry.""" + call_count = 0 + + def reloop_once(state: FactoryState) -> VerdictType: + nonlocal call_count + call_count += 1 + if call_count <= 1: + return VerdictType.RELOOP + return VerdictType.PROCEED + + graph_builder, start = _build_gate_graph(verdict_fn=reloop_once) + graph = graph_builder.build() + + await graph.run(state=factory_state, deps=factory_deps, inputs=start) + + builder_events = [ + e for e in factory_state.events if e["node"] == "MockBuilderNode" + ] + assert len(builder_events) == 2 + + assert builder_events[0]["feedback_received"] == "" + assert "iteration 1" in builder_events[1]["feedback_received"] + + assert "MockBuilderNode" in factory_state.node_feedback + assert "needs improvement" in factory_state.node_feedback["MockBuilderNode"] + + +async def test_feedback_updates_on_each_reloop( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """Feedback is updated on each RELOOP iteration.""" + call_count = 0 + + def reloop_twice(state: FactoryState) -> VerdictType: + nonlocal call_count + call_count += 1 + if call_count <= 2: + return VerdictType.RELOOP + return VerdictType.PROCEED + + graph_builder, start = _build_gate_graph(verdict_fn=reloop_twice) + graph = graph_builder.build() + + await graph.run(state=factory_state, deps=factory_deps, inputs=start) + + builder_events = [ + e for e in factory_state.events if e["node"] == "MockBuilderNode" + ] + assert builder_events[0]["feedback_received"] == "" + assert "iteration 1" in builder_events[1]["feedback_received"] + assert "iteration 2" in builder_events[2]["feedback_received"] + + +async def test_mermaid_shows_branching_topology() -> None: + """Mermaid diagram shows all 3 verdict edges from the gate.""" + graph_builder, _ = _build_gate_graph() + graph = graph_builder.build() + mermaid = graph.render() + + assert "MockBuilderNode" in mermaid + assert "QAGateNode" in mermaid + assert "MockNextNode" in mermaid + + assert "MockBuilderNode --> QAGateNode" in mermaid + assert "decision --> MockBuilderNode" in mermaid + assert "decision --> MockNextNode" in mermaid + assert "MockNextNode --> [*]" in mermaid + assert "decision --> [*]" in mermaid + assert "QAGateNode --> decision" in mermaid + + +async def test_graph_iter_yields_events_proceed( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """Graph.iter() yields node events plus EndMarker for PROCEED path.""" + graph_builder, start = _build_gate_graph( + verdict_fn=lambda _: VerdictType.PROCEED + ) + graph = graph_builder.build() + + events: list[object] = [] + async with graph.iter( + state=factory_state, deps=factory_deps, inputs=start + ) as run: + async for event in run: + events.append(event) + + assert isinstance(events[-1], EndMarker) + assert len(events) == 4 + + +async def test_graph_iter_yields_events_reloop( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """Graph.iter() yields extra events for RELOOP iterations.""" + call_count = 0 + + def reloop_once(state: FactoryState) -> VerdictType: + nonlocal call_count + call_count += 1 + return VerdictType.RELOOP if call_count <= 1 else VerdictType.PROCEED + + graph_builder, start = _build_gate_graph(verdict_fn=reloop_once) + graph = graph_builder.build() + + events: list[object] = [] + async with graph.iter( + state=factory_state, deps=factory_deps, inputs=start + ) as run: + async for event in run: + events.append(event) + + assert isinstance(events[-1], EndMarker) + assert len(events) == 6 + + +async def test_default_verdict_is_proceed( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """GateBaseNode with no verdict_fn defaults to PROCEED.""" + graph_builder, _ = _build_gate_graph(verdict_fn=None) + start = MockBuilderNode(gate_verdict_fn=None) + graph = graph_builder.build() + + result = await graph.run(state=factory_state, deps=factory_deps, inputs=start) + + assert isinstance(result, HaltResult) + assert result.reason == "proceed_complete" diff --git a/experiments/pydantic-graph-prototype/tests/test_parallel.py b/experiments/pydantic-graph-prototype/tests/test_parallel.py new file mode 100644 index 000000000..7bd839029 --- /dev/null +++ b/experiments/pydantic-graph-prototype/tests/test_parallel.py @@ -0,0 +1,346 @@ +"""Tests for parallel fork/join via asyncio.gather inside ForkJoinNode. + +Test workflow: MockBuilderNode → DeepQAForkJoinNode(3 QA agents) → QAResultGateNode → End + +Verifies: + a. All 3 children execute concurrently (wall-clock ≈ max(child_times), not sum) + b. All child outputs appear in ctx.state.node_outputs + c. Mermaid rendering (ForkJoinNode as single node) + d. Event ordering — fork_join_complete after all child_completed events +""" + +import asyncio +import time + +from pydantic_graph import BaseNode, End, EndMarker, GraphBuilder, GraphRunContext + +from pg_factory.deps import FactoryDeps +from pg_factory.nodes.parallel import ChildAgent, ForkJoinNode +from pg_factory.state import FactoryState +from pg_factory.verdicts import HaltResult + + +# ── Simulated QA child agents ───────────────────────────────── + + +async def _health_checker( + ctx: GraphRunContext[FactoryState, FactoryDeps], +) -> str: + await asyncio.sleep(0.05) + return "health_check: all passing" + + +async def _code_reviewer( + ctx: GraphRunContext[FactoryState, FactoryDeps], +) -> str: + await asyncio.sleep(0.08) + return "code_review: 7/7 categories PASS" + + +async def _adversarial_tester( + ctx: GraphRunContext[FactoryState, FactoryDeps], +) -> str: + await asyncio.sleep(0.06) + return "adversarial_qa: feature verified" + + +QA_CHILDREN = [ + ChildAgent(name="health_checker", fn=_health_checker), + ChildAgent(name="code_reviewer", fn=_code_reviewer), + ChildAgent(name="adversarial_tester", fn=_adversarial_tester), +] + + +# ── Test workflow nodes ──────────────────────────────────────── + + +class QAResultGateNode(BaseNode[FactoryState, FactoryDeps, HaltResult]): + """Simple gate that checks QA results and ends.""" + + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> End[HaltResult]: + qa_keys = {"health_checker", "code_reviewer", "adversarial_tester"} + present = qa_keys & ctx.state.node_outputs.keys() + ctx.state.node_outputs["QAResultGateNode"] = f"reviewed {len(present)}/3" + ctx.state.events.append( + {"node": "QAResultGateNode", "action": "gate_check", "qa_count": len(present)} + ) + return End(HaltResult(reason="qa_complete")) + + +class DeepQAForkJoinNode(ForkJoinNode): + """Concrete fork/join modeling the deep-QA parallel subgraph.""" + + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> QAResultGateNode: + await self.execute_children(ctx) + return QAResultGateNode() + + +class MockBuilderNode(BaseNode[FactoryState, FactoryDeps, HaltResult]): + """Simulated builder that feeds into the QA fork/join.""" + + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> DeepQAForkJoinNode: + ctx.state.node_outputs["MockBuilderNode"] = "build complete" + ctx.state.events.append({"node": "MockBuilderNode", "action": "build"}) + return DeepQAForkJoinNode(children=QA_CHILDREN, node_id="deep_qa") + + +# ── Graph builder ────────────────────────────────────────────── + + +def _build_parallel_graph() -> ( + GraphBuilder[FactoryState, FactoryDeps, MockBuilderNode, HaltResult] +): + builder: GraphBuilder[FactoryState, FactoryDeps, MockBuilderNode, HaltResult] = ( + GraphBuilder( + name="parallel-fork-join-test", + state_type=FactoryState, + deps_type=FactoryDeps, + input_type=MockBuilderNode, + output_type=HaltResult, + ) + ) + builder.add_edge(builder.start_node, MockBuilderNode) + builder.add(builder.node(MockBuilderNode)) + builder.add(builder.node(DeepQAForkJoinNode)) + builder.add(builder.node(QAResultGateNode)) + return builder + + +# ── Tests ────────────────────────────────────────────────────── + + +async def test_all_children_execute_concurrently( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """Wall-clock time ≈ max(child_times), not sum — proves asyncio.gather concurrency.""" + graph = _build_parallel_graph().build() + + start = time.monotonic() + result = await graph.run( + state=factory_state, deps=factory_deps, inputs=MockBuilderNode() + ) + elapsed = time.monotonic() - start + + assert isinstance(result, HaltResult) + assert result.reason == "qa_complete" + + # Children sleep 0.05 + 0.08 + 0.06 = 0.19s total. + # If concurrent, wall-clock ≈ max(0.08) = ~0.08s. + # Allow generous margin for CI but ensure it's well under sequential sum. + assert elapsed < 0.15, ( + f"Wall-clock {elapsed:.3f}s too close to sequential sum 0.19s — " + f"children may not be running concurrently" + ) + + +async def test_all_child_outputs_in_state( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """All 3 child outputs appear in ctx.state.node_outputs.""" + graph = _build_parallel_graph().build() + + await graph.run( + state=factory_state, deps=factory_deps, inputs=MockBuilderNode() + ) + + assert factory_state.node_outputs["health_checker"] == "health_check: all passing" + assert factory_state.node_outputs["code_reviewer"] == "code_review: 7/7 categories PASS" + assert factory_state.node_outputs["adversarial_tester"] == "adversarial_qa: feature verified" + + +async def test_builder_output_preserved( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """Builder output remains in state after fork/join completes.""" + graph = _build_parallel_graph().build() + + await graph.run( + state=factory_state, deps=factory_deps, inputs=MockBuilderNode() + ) + + assert factory_state.node_outputs["MockBuilderNode"] == "build complete" + assert factory_state.node_outputs["QAResultGateNode"] == "reviewed 3/3" + + +async def test_gate_sees_all_qa_results( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """QAResultGateNode receives all 3 QA outputs from the fork/join.""" + graph = _build_parallel_graph().build() + + await graph.run( + state=factory_state, deps=factory_deps, inputs=MockBuilderNode() + ) + + gate_events = [ + e for e in factory_state.events if e["node"] == "QAResultGateNode" + ] + assert len(gate_events) == 1 + assert gate_events[0]["qa_count"] == 3 + + +async def test_event_ordering( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """Events: builder → child_completed × 3 → fork_join_complete → gate_check.""" + graph = _build_parallel_graph().build() + + await graph.run( + state=factory_state, deps=factory_deps, inputs=MockBuilderNode() + ) + + actions = [e["action"] for e in factory_state.events] + + assert actions[0] == "build" + + child_completed_indices = [ + i for i, a in enumerate(actions) if a == "child_completed" + ] + assert len(child_completed_indices) == 3 + + fork_join_idx = actions.index("fork_join_complete") + assert all(ci < fork_join_idx for ci in child_completed_indices) + + assert actions[-1] == "gate_check" + + +async def test_fork_join_complete_event_lists_children( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """fork_join_complete event contains all child names and timing.""" + graph = _build_parallel_graph().build() + + await graph.run( + state=factory_state, deps=factory_deps, inputs=MockBuilderNode() + ) + + fj_events = [ + e for e in factory_state.events if e["action"] == "fork_join_complete" + ] + assert len(fj_events) == 1 + + fj = fj_events[0] + assert set(fj["children"]) == {"health_checker", "code_reviewer", "adversarial_tester"} + assert fj["total_duration_ms"] > 0 + + +async def test_per_child_timing_recorded( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """Each child_completed event has duration_ms > 0.""" + graph = _build_parallel_graph().build() + + await graph.run( + state=factory_state, deps=factory_deps, inputs=MockBuilderNode() + ) + + child_events = [ + e for e in factory_state.events if e["action"] == "child_completed" + ] + assert len(child_events) == 3 + + for event in child_events: + assert event["duration_ms"] > 0 + assert event["child"] in {"health_checker", "code_reviewer", "adversarial_tester"} + + +async def test_mermaid_shows_fork_join_as_single_node() -> None: + """Mermaid renders ForkJoinNode as one node (not visual fan-out).""" + graph = _build_parallel_graph().build() + mermaid = graph.render() + + assert "MockBuilderNode" in mermaid + assert "DeepQAForkJoinNode" in mermaid + assert "QAResultGateNode" in mermaid + + assert "MockBuilderNode --> DeepQAForkJoinNode" in mermaid + assert "DeepQAForkJoinNode --> QAResultGateNode" in mermaid + assert "QAResultGateNode --> [*]" in mermaid + + +async def test_graph_iter_yields_correct_event_count( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """Graph.iter() yields one event per node step plus EndMarker.""" + graph = _build_parallel_graph().build() + + events: list[object] = [] + async with graph.iter( + state=factory_state, deps=factory_deps, inputs=MockBuilderNode() + ) as run: + async for event in run: + events.append(event) + + assert isinstance(events[-1], EndMarker) + # 3 nodes (MockBuilder, DeepQAForkJoin, QAResultGate) + EndMarker = 4 + assert len(events) == 4 + + +async def test_child_failure_does_not_crash_fork_join( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """A failing child records the error without crashing the fork/join.""" + + async def _failing_agent( + ctx: GraphRunContext[FactoryState, FactoryDeps], + ) -> str: + raise RuntimeError("agent crashed") + + children = [ + ChildAgent(name="good_agent", fn=_health_checker), + ChildAgent(name="bad_agent", fn=_failing_agent), + ] + + class FailTestForkJoinNode(ForkJoinNode): + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> End[HaltResult]: + await self.execute_children(ctx) + return End(HaltResult(reason="done_with_failures")) + + builder: GraphBuilder[FactoryState, FactoryDeps, FailTestForkJoinNode, HaltResult] = ( + GraphBuilder( + name="fail-test", + state_type=FactoryState, + deps_type=FactoryDeps, + input_type=FailTestForkJoinNode, + output_type=HaltResult, + ) + ) + builder.add_edge(builder.start_node, FailTestForkJoinNode) + builder.add(builder.node(FailTestForkJoinNode)) + graph = builder.build() + + result = await graph.run( + state=factory_state, + deps=factory_deps, + inputs=FailTestForkJoinNode(children=children), + ) + + assert isinstance(result, HaltResult) + assert result.reason == "done_with_failures" + + assert factory_state.node_outputs["good_agent"] == "health_check: all passing" + assert "bad_agent" not in factory_state.node_outputs + + failed_events = [ + e for e in factory_state.events if e["action"] == "child_failed" + ] + assert len(failed_events) == 1 + assert failed_events[0]["child"] == "bad_agent" + assert "agent crashed" in failed_events[0]["error"] diff --git a/experiments/pydantic-graph-prototype/tests/test_smoke.py b/experiments/pydantic-graph-prototype/tests/test_smoke.py new file mode 100644 index 000000000..0f1878203 --- /dev/null +++ b/experiments/pydantic-graph-prototype/tests/test_smoke.py @@ -0,0 +1,75 @@ +"""Smoke test: import pydantic-graph, build a trivial 2-node graph, run it, assert End reached.""" + +from pydantic_graph import BaseNode, End, EndMarker, GraphBuilder, GraphRunContext + +from pg_factory.deps import FactoryDeps +from pg_factory.state import FactoryState +from pg_factory.verdicts import HaltResult + + +class StepOne(BaseNode[FactoryState, FactoryDeps, HaltResult]): + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> "StepTwo": + ctx.state.node_outputs["StepOne"] = "executed" + ctx.state.events.append({"node": "StepOne", "action": "run"}) + return StepTwo() + + +class StepTwo(BaseNode[FactoryState, FactoryDeps, HaltResult]): + async def run( + self, ctx: GraphRunContext[FactoryState, FactoryDeps] + ) -> End[HaltResult]: + ctx.state.node_outputs["StepTwo"] = "executed" + ctx.state.events.append({"node": "StepTwo", "action": "run"}) + return End(HaltResult(reason="complete")) + + +def _build_graph() -> GraphBuilder[FactoryState, FactoryDeps, StepOne, HaltResult]: + builder: GraphBuilder[FactoryState, FactoryDeps, StepOne, HaltResult] = GraphBuilder( + name="smoke-test", + state_type=FactoryState, + deps_type=FactoryDeps, + input_type=StepOne, + output_type=HaltResult, + ) + builder.add_edge(builder.start_node, StepOne) + builder.add(builder.node(StepOne)) + builder.add(builder.node(StepTwo)) + return builder + + +async def test_two_node_graph_runs_to_end( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + graph = _build_graph().build() + result = await graph.run(state=factory_state, deps=factory_deps, inputs=StepOne()) + + assert isinstance(result, HaltResult) + assert result.reason == "complete" + assert factory_state.node_outputs["StepOne"] == "executed" + assert factory_state.node_outputs["StepTwo"] == "executed" + assert len(factory_state.events) == 2 + + +async def test_graph_iter_yields_events( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + graph = _build_graph().build() + events: list[object] = [] + async with graph.iter(state=factory_state, deps=factory_deps, inputs=StepOne()) as run: + async for event in run: + events.append(event) + + assert len(events) == 3 + assert isinstance(events[-1], EndMarker) + + +async def test_mermaid_rendering() -> None: + graph = _build_graph().build() + mermaid = graph.render() + assert "StepOne" in mermaid + assert "StepTwo" in mermaid + assert "[*]" in mermaid diff --git a/experiments/pydantic-graph-prototype/tests/test_study_chain.py b/experiments/pydantic-graph-prototype/tests/test_study_chain.py new file mode 100644 index 000000000..ab1d48138 --- /dev/null +++ b/experiments/pydantic-graph-prototype/tests/test_study_chain.py @@ -0,0 +1,127 @@ +"""Tests for the study chain: execution order, state mutations, Mermaid output.""" + +from pathlib import Path + +from pydantic_graph import EndMarker + +from pg_factory.deps import FactoryDeps +from pg_factory.graphs.study import build_study_graph +from pg_factory.state import FactoryState +from pg_factory.verdicts import HaltResult + + +async def test_study_chain_execution_order( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """Nodes execute in the correct linear order.""" + graph = build_study_graph().build() + from pg_factory.nodes.study import GraphUpdateNode + + result = await graph.run( + state=factory_state, deps=factory_deps, inputs=GraphUpdateNode() + ) + + assert isinstance(result, HaltResult) + assert result.reason == "study_complete" + + expected_order = [ + "GraphUpdateNode", + "StudyNode", + "GraphExplorerNode", + "ConcatStudyNode", + ] + actual_order = [e["node"] for e in factory_state.events] + assert actual_order == expected_order + + +async def test_study_chain_state_mutations( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """Each node records its output to ctx.state.node_outputs.""" + graph = build_study_graph().build() + from pg_factory.nodes.study import GraphUpdateNode + + await graph.run(state=factory_state, deps=factory_deps, inputs=GraphUpdateNode()) + + assert "GraphUpdateNode" in factory_state.node_outputs + assert "StudyNode" in factory_state.node_outputs + assert "GraphExplorerNode" in factory_state.node_outputs + assert "ConcatStudyNode" in factory_state.node_outputs + + +async def test_study_chain_dry_run_creates_files( + factory_state: FactoryState, + tmp_path: Path, +) -> None: + """In dry_run mode, nodes create mock output files.""" + deps = FactoryDeps(project_path=tmp_path, dry_run=True) + graph = build_study_graph().build() + from pg_factory.nodes.study import GraphUpdateNode + + await graph.run(state=factory_state, deps=deps, inputs=GraphUpdateNode()) + + strategy_dir = tmp_path / ".factory" / "strategy" + assert (strategy_dir / "observations.md").exists() + assert (strategy_dir / "graph-context.md").exists() + assert (strategy_dir / "study-combined.md").exists() + + combined = (strategy_dir / "study-combined.md").read_text() + assert "Observations" in combined + assert "Graph Context" in combined + + +async def test_study_chain_iter_yields_events( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """Graph.iter() yields one event per node step plus EndMarker.""" + graph = build_study_graph().build() + from pg_factory.nodes.study import GraphUpdateNode + + events: list[object] = [] + async with graph.iter( + state=factory_state, deps=factory_deps, inputs=GraphUpdateNode() + ) as run: + async for event in run: + events.append(event) + + assert len(events) == 5 + assert isinstance(events[-1], EndMarker) + + +async def test_study_chain_mermaid_contains_all_nodes() -> None: + """Mermaid output contains all 4 node names.""" + graph = build_study_graph().build() + mermaid = graph.render() + + assert "GraphUpdateNode" in mermaid + assert "StudyNode" in mermaid + assert "GraphExplorerNode" in mermaid + assert "ConcatStudyNode" in mermaid + + +async def test_study_chain_mermaid_topology() -> None: + """Mermaid output shows the correct linear edge topology.""" + graph = build_study_graph().build() + mermaid = graph.render() + + assert "GraphUpdateNode --> StudyNode" in mermaid + assert "StudyNode --> GraphExplorerNode" in mermaid + assert "GraphExplorerNode --> ConcatStudyNode" in mermaid + assert "ConcatStudyNode --> [*]" in mermaid + + +async def test_study_chain_events_have_action_field( + factory_state: FactoryState, + factory_deps: FactoryDeps, +) -> None: + """Each event recorded by nodes has a distinct action field.""" + graph = build_study_graph().build() + from pg_factory.nodes.study import GraphUpdateNode + + await graph.run(state=factory_state, deps=factory_deps, inputs=GraphUpdateNode()) + + actions = [e["action"] for e in factory_state.events] + assert actions == ["graph_update", "study", "explore", "concat"]