Mejorar pipelines / DAG de los agentes

[miguelgfierro]

Target Module

genai

Problem Statement

Más declarativo

Proposed Solution

TODO

Alternatives Considered

No response

Additional Context

No response

[miguelgfierro]

Plan: LangGraph-shaped StateGraph API for Firefly pipelines (Issue #147)

Context

GitHub issue #147 — "Mejorar pipelines / DAG de los agentes" — asks for a more declarative pipeline API. The issue itself is a stub ("Más declarativo / TODO") with no design.

Firefly's current PipelineBuilder is a string-keyed, port-based DAG (closer to Haystack than LangGraph). It is mature on infrastructure (async, retries, OTel, failure strategies, eager scheduling) but lacks the conceptual model that has become the agentic-framework standard:

• Branching is split across two mechanisms (BranchStep + per-node condition lambdas) and is not inspectable as data.
• Cycles are forbidden (Kahn's topological sort), so ReAct loops, retry-with-critique, and multi-turn agent reflection can't be expressed.
• There is no typed shared state; data flows through stringly-typed output_key → input_key ports.
• The DAG cannot be rendered or serialised.

Industry research (LangGraph, LlamaIndex Workflows, Burr, pydantic-graph, AutoGen GraphFlow, Hamilton) shows the dominant agentic pattern is LangGraph's StateGraph with add_conditional_edges. This plan adopts that pattern, additively, while reusing Firefly's existing execution infrastructure.

Goal

Introduce a StateGraph API shaped exactly like LangGraph's, living alongside the current PipelineBuilder. Old pipelines keep working; new agentic workflows use StateGraph.

Decision points (defaults chosen; flag any pushback)

Decision	Default	Rationale
Scope	Additive — coexist with PipelineBuilder	Zero breakage; lets us migrate internal callers gradually. PipelineBuilder stays the right tool for pure ETL-shaped DAGs.
Cycles in v1	Yes — frontier executor with recursion_limit=25	Loops are LangGraph's defining feature for agents. Skipping them gives us the syntax but not the semantics.
Send (runtime fan-out) in v1	Yes	Replaces FanOutStep. Small addition once the frontier executor exists.
Mermaid + JSON export in v1	Yes	Cheap (one nested loop), high value for docs / debugging / Studio.
Checkpointer / HITL interrupt	Deferred to v2	Real Postgres/Redis checkpointers are a project on their own.
Streaming (app.stream())	Deferred to v2	Requires generator-based engine refactor.
Subgraphs	Deferred to v2	Composing one StateGraph as a node inside another.

If any of these defaults is wrong, say so before implementation starts.

Target API (canonical example)

from typing import TypedDict, Annotated
from fireflyframework_agentic.pipeline import StateGraph, START, END, Send
from fireflyframework_agentic.pipeline.reducers import append

class AgentState(TypedDict):
    messages: Annotated[list[str], append]
    intent: str | None
    answer: str | None

async def classify(state: AgentState) -> dict:
    return {"intent": await llm_classify(state["messages"])}

async def answer(state: AgentState) -> dict:
    return {"answer": await llm_answer(state["messages"])}

async def escalate(state: AgentState) -> dict:
    return {"answer": "Escalated to human."}

def route(state: AgentState) -> str:
    return "escalate" if state["intent"] == "complaint" else "answer"

g = StateGraph(AgentState)
g.add_node("classify", classify)
g.add_node("answer", answer)
g.add_node("escalate", escalate)
g.add_edge(START, "classify")
g.add_conditional_edges("classify", route, {"answer": "answer", "escalate": "escalate"})
g.add_edge("answer", END)
g.add_edge("escalate", END)

app = g.compile()
result = await app.invoke({"messages": ["I want a refund."]})
print(app.get_graph().to_mermaid())

Implementation outline

All paths relative to /home/u/signature/fireflyframework-agentic/fireflyframework_agentic/pipeline/.

New files

• state_graph.py — StateGraph[State], START, END sentinels, Send dataclass.
  • StateGraph(state_schema) — accepts a TypedDict subclass or Pydantic model.
  • add_node(name, fn) — fn is async (state) -> dict.
  • add_edge(source, target) — source/target may be node name, START, or END.
  • add_conditional_edges(source, router, mapping) — mapping is {label: target}; router returns a label, a Send, or a list of Sends.
  • compile() -> CompiledStateGraph — validates connectivity (every node reachable from START, every node reaches END or is part of a cycle that does).
• compiled_graph.py — CompiledStateGraph.
  • async invoke(state) -> State — frontier executor (see below).
  • get_graph() -> GraphView — .to_mermaid(), .to_json(), .to_ascii().
• reducers.py — built-in reducers: append, extend, merge_dict, replace (default).
  • Reducers are looked up via typing.get_type_hints(state_schema, include_extras=True) for Annotated[T, reducer_fn]. Fields without annotation get replace.
• frontier_executor.py — frontier-based BFS executor supporting cycles.
  • Maintains a queue of (node_name, state_slice) tasks.
  • Per-step: pop frontier, run all current tasks concurrently with asyncio.gather, merge each return value into shared state via the relevant reducer, evaluate outgoing edges (incl. conditional + Send), enqueue next frontier.
  • Tracks per-node visit count; aborts when any node's visits exceed recursion_limit (default 25, override per-compile).
  • Reuses DAGNode.retry_max, timeout_seconds, failure_strategy, OTel span emission, and PipelineEventHandler callbacks. The retry / timeout / failure-strategy code lives today in pipeline/engine.py; lift the relevant helpers into a shared _node_runtime.py so both PipelineEngine and frontier_executor use them — do not duplicate.

Modified files

• __init__.py — add exports: StateGraph, START, END, Send, append, extend, merge_dict, replace.
• engine.py — extract the per-node execution helper (retry loop + timeout + OTel + event hooks) into a new private module _node_runtime.py that both PipelineEngine and the new frontier_executor import. No behavioural change to PipelineEngine.

Tests

Create tests/pipeline/test_state_graph.py with plain pytest functions (no classes, per project rules). Cover:

1. Linear three-node graph START → a → b → c → END returns expected state.
2. Reducers: append accumulates, default replace overwrites, custom reducer via Annotated.
3. Conditional edges: router returns label; correct branch runs; non-matching label raises.
4. Cycle: simple two-node loop with a router that exits after N iterations; recursion_limit triggers when the router never exits.
5. Send: router returns [Send("worker", payload_a), Send("worker", payload_b)]; both worker invocations run concurrently; merge reducer combines results.
6. Mermaid export: golden-file test for a small graph.
7. Reuse of DAGNode retry/timeout: a flaky node retries and eventually succeeds; a slow node hits its timeout; failure strategies behave as in PipelineEngine.
8. Connectivity validation: compile() rejects unreachable nodes and dead-end nodes (with explanatory error).

Run with pytest tests/pipeline/test_state_graph.py -v from the repo root inside the ~/.venvs/firefly venv (Python 3.13).

Docs

• Add docs/pipeline/state_graph.md with the canonical example, a section on reducers, a section on conditional edges, a section on loops with recursion_limit, and a "when to use StateGraph vs PipelineBuilder" table.
• Update top-level README.md to mention both APIs in the pipeline section.
• Do not mark PipelineBuilder as deprecated yet — that decision waits until at least one production pipeline has migrated.

Reused existing code (do not reinvent)

• pipeline/engine.py — retry-with-backoff loop, timeout enforcement, OTel span emission, event-handler dispatch. Lift into _node_runtime.py and call from both executors.
• pipeline/dag.py — DAGNode already carries retry_max, timeout_seconds, failure_strategy, condition, backoff_factor. Reuse for StateGraph nodes too; the conditional-edge mechanism layers on top, it does not replace per-node settings.
• pipeline/context.py — PipelineContext keeps correlation_id, memory, metadata. The CompiledStateGraph carries an internal PipelineContext for OTel/memory parity with PipelineEngine. The user-facing state is the typed state dict; PipelineContext is plumbing.
• pipeline/result.py — NodeResult, ExecutionTraceEntry, UsageSummary. CompiledStateGraph.invoke() returns the final state, but emits the same per-node tracing into a PipelineResult-shaped object available as app.last_result for parity.
• pipeline/steps.py — AgentStep, ReasoningStep, CallableStep, EmbeddingStep, RetrievalStep, BatchLLMStep. These are still useful inside StateGraph nodes; provide a thin as_node_fn(step) adapter that turns a StepExecutor into the (state) -> dict signature, so users can drop existing steps into a StateGraph without rewriting.

Verification

End-to-end checks after implementation:

1. pytest tests/pipeline/ -v — full pipeline suite green (existing + new StateGraph tests).
2. pytest tests/ -v — full repo suite green; no regression in PipelineEngine callers.
3. Run the canonical example in docs/pipeline/state_graph.md as a smoke script; print the Mermaid; eyeball that it matches the topology.
4. Run one existing pipeline (pick any from the repo's integration tests) ported to StateGraph side by side with the original; outputs should match.
5. ruff check . and mypy fireflyframework_agentic/pipeline/ clean.

What is explicitly NOT in v1

• Checkpointer / HITL interrupt+resume.
• app.stream() partial-state streaming.
• Subgraphs (one StateGraph as a node inside another).
• Migration of existing internal callers off PipelineBuilder — separate follow-up issue.
• Deprecation of BranchStep / FanOutStep — same.

Estimated size

One PR, ~800–1200 lines of new code (state_graph + compiled_graph + frontier_executor + reducers + adapter) plus ~400 lines of tests plus docs. The _node_runtime.py extraction is the only touch to existing code and should be a pure refactor with no behavioural change.

[miguelgfierro]

Alternative plan: pydantic-graph as the model (Issue #147)

This is an alternative to the LangGraph-shaped StateGraph plan posted above. It targets the same goal — a more declarative pipeline API for Firefly — but using pydantic-graph (from pydantic-ai) as the model instead of LangGraph.

The two plans are mutually exclusive. Pick one before implementation begins.

Why consider pydantic-graph at all

The LangGraph proposal is more LangGraph-shaped, but LangGraph itself isn't truly declarative — it's a builder pattern (g.add_node(...).add_edge(...)) that produces a data-shaped output. If issue #147's actual ask ("Más declarativo") means genuinely declarative, pydantic-graph delivers it; LangGraph does not.

The single design idea that separates pydantic-graph from everything else: edges live in the type system. Each node's run method has a return type, and that return type is the set of outgoing edges. mypy/pyright statically verify that nodes only transition to declared successors.

Comparison vs. the LangGraph-shaped plan

Dimension	LangGraph-shaped StateGraph (other plan)	pydantic-graph (this plan)
Edge declaration	g.add_edge("a", "b") + add_conditional_edges(...) — method calls	Return-type union of BaseNode.run — types
Static edge checking	None — node names are strings	Full — mypy errors on undeclared transitions
Typed state	TypedDict with reducers (Annotated[list, append])	@dataclass or pydantic model, mutated via ctx.state.field
Branching	Router fn returns label → mapping dict picks target	if cond: return A() else: return B() inside run
Cycles	Yes, with recursion_limit=25	Yes, no built-in limit (manual counter)
Fan-out (Send)	First-class	Weak — asyncio.gather inside a node, no per-branch graph nodes
Inspection	to_mermaid(), to_json()	mermaid_code(), snapshots
HITL / persistence	Deferred to v2	Built-in: FileStatePersistence + iter_from_persistence
Ecosystem coupling	Drags langchain-core	Pure Python, MIT, no LLM deps
House-style fit	Introduces TypedDict-with-reducers model	Pydantic everywhere — same model as NodeResult, DAGEdge, etc.
Industry mindshare	Dominant LLM-orchestration	Smaller, but rising with pydantic-ai

Goal

Adopt pydantic-graph as Firefly's declarative pipeline API. Live alongside the current PipelineBuilder (additive); migrate internal callers gradually.

Decision points (defaults chosen; flag any pushback)

Decision	Default	Rationale
Scope	Additive — coexist with PipelineBuilder	Zero breakage. PipelineBuilder stays the right tool for pure ETL-shaped DAGs.
Direct use vs. thin wrapper	Thin wrapper (firefly.pipeline.graph) re-exporting BaseNode, End, Graph, GraphRunContext plus Firefly extensions	Keeps users on pydantic-graph's API directly; the wrapper adds retries/timeouts/OTel parity.
Cycles in v1	Yes, with optional recursion_limit enforced by wrapper	Defining feature for agentic loops. pydantic-graph leaves the limit to the caller; we add it.
HITL / persistence in v1	Yes — re-export FileStatePersistence	It's free; comes with the library.
Fan-out story	Defer first-class Send equivalent; document asyncio.gather-inside-a-node pattern	The weak spot. Address in v2 when GraphBuilder parallel-step support in pydantic-graph stabilises, or build a Firefly-specific primitive then.
Streaming	Deferred to v2	Same as the other plan.

Target API (canonical example — same workflow as the other plan)

from dataclasses import dataclass, field
from pydantic_graph import BaseNode, End, Graph, GraphRunContext
from fireflyframework_agentic.pipeline.graph import run_with_observability  # Firefly wrapper

@dataclass
class AgentState:
    messages: list[str] = field(default_factory=list)
    intent: str | None = None

@dataclass
class Classify(BaseNode[AgentState]):
    async def run(self, ctx: GraphRunContext[AgentState]) -> "Answer | Escalate":
        ctx.state.intent = await llm_classify(ctx.state.messages)
        return Escalate() if ctx.state.intent == "complaint" else Answer()

@dataclass
class Answer(BaseNode[AgentState, None, str]):
    async def run(self, ctx: GraphRunContext[AgentState]) -> End[str]:
        return End(await llm_answer(ctx.state.messages))

@dataclass
class Escalate(BaseNode[AgentState, None, str]):
    async def run(self, ctx: GraphRunContext[AgentState]) -> End[str]:
        return End("Escalated to human.")

graph = Graph(nodes=[Classify, Answer, Escalate])

# Firefly wrapper adds OTel, retries, recursion_limit
final = await run_with_observability(graph, start=Classify(), state=AgentState(messages=["I want a refund."]))
print(graph.mermaid_code())

No add_node. No add_edge. No router fn. No mapping dict. Topology = types.

Implementation outline

All paths relative to /home/u/signature/fireflyframework-agentic/fireflyframework_agentic/pipeline/.

New files

• graph/__init__.py — re-exports BaseNode, End, Graph, GraphRunContext, FileStatePersistence, FullStatePersistence from pydantic_graph for the canonical import path.
• graph/runtime.py — Firefly's observability + retries + recursion-limit wrapper:
  • async run_with_observability(graph, *, start, state, deps=None, persistence=None, recursion_limit=25, retry_policy=None, event_handler=None) -> T.
  • Wraps graph.iter() to inject:
    • Per-node retry-with-backoff (reuse _node_runtime.py extracted from engine.py).
    • OTel span per node.
    • PipelineEventHandler callbacks (on_node_start, on_node_complete, etc.).
    • Recursion-limit enforcement (counts node visits; raises RecursionLimitExceeded over threshold).
• graph/adapters.py — bridge to existing Firefly steps:
  • agent_node(agent, prompt_key="input") -> type[BaseNode] — factory that produces a BaseNode subclass wrapping a FireflyAgent.
  • Same for reasoning_node(pattern), embedding_node(embedder), retrieval_node(store), batch_llm_node(...).
  • Each factory generates a @dataclass class _GeneratedNode(BaseNode[...]): async def run(...): ... so users can drop existing components into a pydantic-graph workflow without rewriting.

Modified files

• pipeline/__init__.py — add re-export: from .graph import *.
• pipeline/engine.py — extract per-node retry/timeout/OTel/event-hook helper into _node_runtime.py (same refactor as the other plan; pure no-op extraction).

Dependencies

• Add pydantic-graph>=1.100 to pyproject.toml (pure-Python, MIT, no langchain coupling).

Tests

tests/pipeline/test_graph.py, plain pytest functions. Cover:

1. Linear three-node graph (Classify → Answer → End) returns expected output.
2. Branching: Classify returns Answer | Escalate; correct branch runs based on state.
3. Cycle: a counter node that loops back to itself N times then Ends; mypy passes; runtime exit works.
4. recursion_limit enforced by the Firefly wrapper triggers RecursionLimitExceeded when a cycle never exits.
5. Retries: a flaky BaseNode retries via the wrapper's retry policy and eventually succeeds.
6. Timeout: a slow node hits the wrapper's timeout and surfaces a typed error.
7. OTel spans emitted per node visit (snapshot the span hierarchy).
8. FileStatePersistence: a graph paused mid-run and resumed from disk produces identical final state.
9. Adapter: an existing FireflyAgent wrapped via agent_node(agent) runs inside a graph.
10. mypy: a mypy invocation on a fixture file with a wrong return type fails (run as a subprocess in the test).

Run with pytest tests/pipeline/test_graph.py -v from the repo root inside ~/.venvs/firefly.

Docs

• docs/pipeline/graph.md — canonical example, branching, cycles, persistence/HITL, adapter usage, "when to use Graph vs PipelineBuilder" table, "what about fan-out" honest note.
• Update top-level README.md to mention both APIs.

Reused existing code (do not reinvent)

• pipeline/engine.py retry/timeout/OTel helpers → extracted into _node_runtime.py, called by graph/runtime.py.
• pipeline/context.py PipelineContext is not reused at the user-facing layer (state is now ctx.state from pydantic-graph), but the wrapper carries an internal PipelineContext for correlation IDs and memory injection.
• pipeline/result.py NodeResult, ExecutionTraceEntry, UsageSummary — emitted by the wrapper into an optional last_result for parity with PipelineEngine.
• pipeline/steps.py — wrapped by the adapters in graph/adapters.py. No rewrite needed.

Verification

1. pytest tests/pipeline/ -v — full pipeline suite green.
2. pytest tests/ -v — full repo suite green; PipelineEngine regression-free.
3. mypy fireflyframework_agentic/pipeline/ clean. Also run mypy on tests/pipeline/test_graph.py to confirm static checks fire on the negative fixture.
4. Run the canonical example in docs/pipeline/graph.md as a smoke script; print Mermaid; check topology.
5. Pick one existing pipeline; port it to a pydantic-graph version; outputs should match the original.
6. ruff check . clean.

What is explicitly NOT in v1

• First-class Send-equivalent for runtime fan-out.
• Streaming partial state.
• Subgraphs (composing one Graph as a node inside another — pydantic-graph supports this natively in recent versions; just not in our v1 scope).
• Migration of existing internal callers off PipelineBuilder — separate follow-up.
• Deprecation of BranchStep / FanOutStep — separate follow-up.

When to pick this plan over the LangGraph one

Pick pydantic-graph (this plan) if:

• Static type-safety on edges is a top priority.
• The team values "topology in types" over industry mindshare.
• Pydantic-everywhere house style matters more than LangGraph familiarity.
• Heavy LLM fan-out / map-reduce is not a near-term workload.

Pick LangGraph-shaped StateGraph (other plan) if:

• Industry familiarity / hiring story matters.
• Dynamic fan-out with Send is a near-term workload.
• LangSmith integration is wanted out of the box.

Estimated size

Smaller than the LangGraph plan: ~400–600 LOC of wrapper + adapters + tests (we are not reimplementing the graph executor — pydantic-graph provides it). Plus the same _node_runtime.py extraction (pure refactor, no behaviour change).

[miguelgfierro]

Pipeline API comparison

Same workflow, three APIs. Read top-to-bottom to compare ergonomics for issue #147.

The workflow

A customer-support agent. Input: a list of user messages. Step 1 classifies the user's intent via an LLM (returns "complaint" or any other label). Step 2 branches: if the intent is "complaint", escalate to a human and return "Escalated to human."; otherwise generate an answer via an LLM. The final output is a single string. The LLM calls (llm_classify, llm_answer) are stubbed — the point is the wiring, not the model.

A. Current Firefly PipelineBuilder

Port-based DAG. Nodes are added by string id, edges connect string ids with output_key / input_key ports. Branching is split between a BranchStep (which writes a string label as its output) and per-downstream-node condition lambdas that inspect that label via the PipelineContext. The final result is read out of the context by node id. Uses the real API from fireflyframework_agentic/pipeline/.

from typing import Any

from fireflyframework_agentic.pipeline.builder import PipelineBuilder
from fireflyframework_agentic.pipeline.context import PipelineContext
from fireflyframework_agentic.pipeline.steps import BranchStep, CallableStep


async def llm_classify(messages: list[str]) -> str:
    return "complaint"


async def llm_answer(messages: list[str]) -> str:
    return "Here is your answer."


async def classify_step(ctx: PipelineContext, inputs: dict[str, Any]) -> str:
    messages = inputs.get("input", ctx.inputs)
    return await llm_classify(messages)


async def answer_step(ctx: PipelineContext, inputs: dict[str, Any]) -> str:
    messages = ctx.inputs
    return await llm_answer(messages)


async def escalate_step(ctx: PipelineContext, inputs: dict[str, Any]) -> str:
    return "Escalated to human."


def route(inputs: dict[str, Any]) -> str:
    intent = inputs.get("input", "")
    return "escalate" if intent == "complaint" else "answer"


def is_branch(label: str) -> Any:
    def _cond(ctx: PipelineContext) -> bool:
        result = ctx.get_node_result("route")
        return getattr(result, "output", None) == label

    return _cond


pipeline = (
    PipelineBuilder("support-agent")
    .add_node("classify", CallableStep(classify_step))
    .add_node("route", BranchStep(router=route))
    .add_node("answer", CallableStep(answer_step), condition=is_branch("answer"))
    .add_node("escalate", CallableStep(escalate_step), condition=is_branch("escalate"))
    .add_edge("classify", "route")
    .add_edge("route", "answer")
    .add_edge("route", "escalate")
    .build()
)


async def run(messages: list[str]) -> str:
    result = await pipeline.execute(messages)
    answer = result.context.get_node_output("answer")
    escalated = result.context.get_node_output("escalate")
    return answer if answer is not None else escalated

Observations:

• Branching is bifurcated: BranchStep writes a label, then each downstream node carries its own condition lambda that reaches back into the context to read that label. Two coupled mechanisms for one concept.
• Node ids are stringly typed everywhere. Typos in "route" / "answer" / "escalate" are runtime failures.
• There is no shared state object. Data flows by reading upstream node outputs from the context by id, or via the port keys on edges. Reading the final answer requires the caller to know which branch ran and probe both.
• Every node signature is async (PipelineContext, dict) -> Any. The inputs dict is keyed by edge input_key and is not statically typed.
• No cycles allowed (Kahn topological sort in dag.py). ReAct-style loops cannot be expressed.
• Strengths: mature runtime — retries, timeouts, failure strategies, OTel spans, event hooks are already wired into PipelineEngine.
• LOC (excl. blanks and imports): 41.

B. Proposed StateGraph (LangGraph-shaped)

Aspirational — this API does not exist in Firefly yet. It is the canonical example from plan i-would-like-you-calm-kite.md. Nodes are async functions over a typed TypedDict state; each returns a partial state dict that is merged via per-field reducers. Branching is one call: add_conditional_edges(source, router_fn, {label: target}). Final result is the full state dict.

from typing import Annotated, TypedDict

from fireflyframework_agentic.pipeline import END, START, StateGraph
from fireflyframework_agentic.pipeline.reducers import append


class AgentState(TypedDict):
    messages: Annotated[list[str], append]
    intent: str | None
    answer: str | None


async def llm_classify(messages: list[str]) -> str:
    return "complaint"


async def llm_answer(messages: list[str]) -> str:
    return "Here is your answer."


async def classify(state: AgentState) -> dict:
    return {"intent": await llm_classify(state["messages"])}


async def answer(state: AgentState) -> dict:
    return {"answer": await llm_answer(state["messages"])}


async def escalate(state: AgentState) -> dict:
    return {"answer": "Escalated to human."}


def route(state: AgentState) -> str:
    return "escalate" if state["intent"] == "complaint" else "answer"


g = StateGraph(AgentState)
g.add_node("classify", classify)
g.add_node("answer", answer)
g.add_node("escalate", escalate)
g.add_edge(START, "classify")
g.add_conditional_edges("classify", route, {"answer": "answer", "escalate": "escalate"})
g.add_edge("answer", END)
g.add_edge("escalate", END)

app = g.compile()


async def run(messages: list[str]) -> str:
    result = await app.invoke({"messages": messages})
    return result["answer"]

Observations:

• One typed shared-state object (AgentState). Node bodies read by key; static checkers can flag state["intnet"] typos on a TypedDict.
• Branching is a single call. The {label: target} mapping is data — the graph can be rendered and inspected.
• Node functions are plain async (State) -> dict. No PipelineContext, no inputs dict, no output_key/input_key plumbing.
• Construction is still imperative — g.add_node(...) calls in sequence, names repeated as strings in add_edge / add_conditional_edges. String-id typos remain a runtime failure (mitigated by compile()-time connectivity validation, not by mypy).
• Reducers (Annotated[list[str], append]) move merge semantics into the type. Powerful, but invisible to a reader who is not looking at the schema.
• Supports cycles (per plan). ReAct loops, retry-with-critique, multi-turn reflection are expressible.
• LOC (excl. blanks and imports): 32.

C. pydantic-graph alternative

Real library (pydantic-ai's pydantic_graph). Each node is a @dataclass BaseNode subclass; its async run(ctx) returns either the next node instance or End(value). Edges are encoded in the return-type union of run. There is no builder, no add_node / add_edge — the graph topology is discovered by inspecting the return-type annotations of every node.

from __future__ import annotations

from dataclasses import dataclass, field

from pydantic_graph import BaseNode, End, Graph, GraphRunContext


async def llm_classify(messages: list[str]) -> str:
    return "complaint"


async def llm_answer(messages: list[str]) -> str:
    return "Here is your answer."


@dataclass
class State:
    messages: list[str] = field(default_factory=list)
    intent: str | None = None


@dataclass
class Classify(BaseNode[State, None, str]):
    async def run(self, ctx: GraphRunContext[State]) -> Answer | Escalate:
        ctx.state.intent = await llm_classify(ctx.state.messages)
        return Escalate() if ctx.state.intent == "complaint" else Answer()


@dataclass
class Answer(BaseNode[State, None, str]):
    async def run(self, ctx: GraphRunContext[State]) -> End[str]:
        return End(await llm_answer(ctx.state.messages))


@dataclass
class Escalate(BaseNode[State, None, str]):
    async def run(self, ctx: GraphRunContext[State]) -> End[str]:
        return End("Escalated to human.")


graph = Graph(nodes=[Classify, Answer, Escalate])


async def run(messages: list[str]) -> str:
    result = await graph.run(Classify(), state=State(messages=messages))
    return result.output

Observations:

• Edges live in types. Classify.run returning Answer | Escalate is the edge declaration; mypy/pyright will flag if a run method returns a node that is not in its declared union.
• No string ids anywhere. Renaming a node is a refactor across class names, caught by tooling.
• The graph topology is inspectable by reading return-type annotations — Graph.mermaid_code() is built in.
• Cost: one class per node. Trivial workflows feel ceremonious. State is a separate dataclass that nodes mutate in ctx.state, which is less functional than StateGraph's "return partial state, framework reduces" model.
• Branching is return Escalate() if ... else Answer(). No router function, no label mapping — just Python.
• Cycles are expressed naturally — a node can return itself or any previously-defined node, as long as it is in the union.
• LOC (excl. blanks and imports): 27.

Side-by-side summary

Aspect	A. PipelineBuilder	B. StateGraph	C. pydantic-graph
LOC (this workflow)	41	32	27
Branching	BranchStep + per-node condition lambdas (two mechanisms)	add_conditional_edges(src, router, {label: target}) (one call)	Return-type union on run
Where state lives	Per-node outputs in PipelineContext._results, keyed by node id	Typed TypedDict state, merged by reducers	Dataclass State mutated via ctx.state
Type safety on state access	None — inputs.get("input") and ctx.get_node_output(id) are Any	TypedDict — static checkers catch typos in state["..."]	Full — State is a dataclass with typed fields
Type safety on graph edges	None — strings	Strings, but compile() validates connectivity	Full — edges are types; mypy/pyright catch invalid transitions
Graph inspectability	DAG inspectable in memory; no exporter	Plan promises Mermaid + JSON export	Graph.mermaid_code() built in
mypy catches bad transitions	No	No (string ids)	Yes
Cycles / loops	Forbidden	Yes (recursion_limit)	Yes (natural via return types)
Fan-out / parallel branches	FanOutStep + multiple downstream edges	Send(node, payload) from a router	Not first-class; manual concurrency in a node
Ecosystem / familiarity	Internal only	Mirrors LangGraph — current de-facto standard	Small but growing (Pydantic AI ecosystem)
Runtime maturity	High — retries, timeouts, OTel, failure strategies in place	Inherits the same runtime via _node_runtime.py extraction (per plan)	Independent runtime; less battle-tested

Verdict

PipelineBuilder is the most verbose by every measure and bifurcates branching across two coupled mechanisms — fine for linear ETL DAGs, awkward for agent control flow. StateGraph halves the ceremony and matches the industry-standard mental model, but it remains a string-id imperative builder: mypy cannot catch a misspelled node name, only compile() can. pydantic-graph is the only one of the three where the graph topology is type-checked by mypy and the edges are inspectable without a builder — but it pays for that with a class per node, which dominates the LOC win on small graphs and would feel heavier on large ones. For issue #147 the StateGraph shape is the right move: it is the convention readers will recognise, and the type-safety gap versus pydantic-graph is small relative to the readability gain over the current API.

[miguelgfierro]

Alternative plan: evolve PipelineBuilder — best of both, no new dependency (Issue #147)

This is a third alternative to the two plans above. Same goal — a more declarative pipeline API for Firefly — but instead of introducing a parallel StateGraph class (plan A) or adopting pydantic-graph (plan B), it evolves the existing PipelineBuilder to absorb the best ideas of each.

The three plans are mutually exclusive. Pick one before implementation begins.

Motivation

Plan A is implemented natively (no langgraph dependency), but it introduces a second pipeline API that lives alongside PipelineBuilder. Two APIs for "build a DAG" is a long-term tax: docs split, the team has to learn both, internal callers fragment.

Plan B is genuinely declarative (topology in types) but adds a real third-party dependency, has a weak fan-out story, and forces a BaseNode subclass-per-step pattern that does not match how Firefly steps are written today.

"Declarative" actually pulls in two directions:

• Declarative-as-data (LangGraph): topology ends up as inspectable data — renderable as Mermaid, serialisable as JSON, branching collapsed into one call.
• Declarative-as-types (pydantic-graph): edges live in run's return annotations, mypy enforces them.

This plan delivers both — additively — on the existing PipelineBuilder, without a parallel class and without a new dependency.

Goal

Keep one pipeline API. Extend PipelineBuilder so it can express what plan A's StateGraph and plan B's Graph express, while staying source-compatible with every pipeline written today.

Decision points (defaults chosen; flag any pushback)

Decision	Default	Rationale
Scope	Evolve PipelineBuilder in place; everything additive and opt-in	Avoid two parallel APIs. Existing pipelines unchanged.
Typed shared state	Opt-in via PipelineBuilder("name", state=AgentState)	Powerful, but the port-based model still works for ETL-shaped pipelines that don't need shared state.
State schema type	Pydantic model OR TypedDict	Pydantic-everywhere house style, but TypedDict accepted for users who prefer LangGraph idioms.
Reducers	Annotated[T, reducer_fn] on fields	Same mechanism in both plan A and pydantic-graph idioms; familiar.
Unified branching	New .branch(source, router, {label: target})	Collapses BranchStep + per-node condition lambdas into one inspectable call. BranchStep and condition= keep working (soft-deprecated).
Optional typed edges	If node return annotation is Literal[...] or a node-ref union, .build() validates declared edges against it	Static-ish edge checking without forcing a BaseNode subclass pattern. Fully optional — annotate or don't.
Cycles	Yes, with recursion_limit=25, enabled only when the DAG declares itself cyclic	Pure DAGs keep the current Kahn path (zero behavioural change); cyclic graphs route to a frontier executor.
Runtime fan-out (Send-equivalent)	Yes — router may return a list of Send(target, payload)	Replaces FanOutStep. Needed for agentic map-reduce.
Mermaid + JSON export	Yes, on the DAG itself (dag.to_mermaid(), dag.to_json())	Cheap, high value, useful to Studio.
Checkpointer / HITL / streaming	Deferred to v2	Same as plans A and B.

Target API (canonical example — same workflow as plans A and B)

from typing import Annotated
from pydantic import BaseModel
from fireflyframework_agentic.pipeline import PipelineBuilder, Send
from fireflyframework_agentic.pipeline.reducers import append


class AgentState(BaseModel):
    messages: Annotated[list[str], append] = []
    intent: str | None = None
    answer: str | None = None


async def classify(state: AgentState) -> dict:
    return {"intent": await llm_classify(state.messages)}


async def answer(state: AgentState) -> dict:
    return {"answer": await llm_answer(state.messages)}


async def escalate(state: AgentState) -> dict:
    return {"answer": "Escalated to human."}


def route(state: AgentState) -> str:
    return "escalate" if state.intent == "complaint" else "answer"


pipeline = (
    PipelineBuilder("support-agent", state=AgentState)
    .add_node("classify", classify)
    .add_node("answer", answer)
    .add_node("escalate", escalate)
    .add_edge("classify")  # source-only: implicit START -> classify when only one entry
    .branch("classify", route, {"answer": "answer", "escalate": "escalate"})
    .build()
)


async def run(messages: list[str]) -> str:
    result = await pipeline.invoke(AgentState(messages=messages))
    return result.answer

Compared to today's PipelineBuilder:

• One state object, not port-based plumbing. Nodes are async (State) -> dict, return value merged via reducers.
• One branching call, not BranchStep + per-node condition lambdas reading back through the context.
• Typed result: pipeline.invoke(state) returns the final AgentState, not a PipelineResult you have to probe by node id.

Compared to plan A's StateGraph: same shape, no second API class — same PipelineBuilder, just richer.

Compared to plan B's pydantic-graph: same pydantic-everywhere feel, no new dependency, fan-out is first-class, but topology is data not types (unless you opt into typed-edge annotations — see below).

Optional: typed edges (pydantic-graph influence)

Nodes can declare their successors in the return annotation; .build() validates the declared edges match:

from typing import Literal

async def classify(state: AgentState) -> tuple[dict, Literal["answer", "escalate"]]:
    intent = await llm_classify(state.messages)
    target = "escalate" if intent == "complaint" else "answer"
    return {"intent": intent}, target

If .branch("classify", ...) declares a target the annotation doesn't allow, .build() raises. This is fully optional — annotate or don't.

Implementation outline

All paths relative to /home/u/signature/fireflyframework-agentic/fireflyframework_agentic/pipeline/.

Modified files

• builder.py
  • New optional state kwarg on __init__.
  • add_node(name, fn) accepts async (state) -> dict when state= is set, alongside today's (ctx, inputs) -> Any form. Type of the callable is sniffed by signature, no flag needed.
  • New .branch(source, router, mapping) method.
  • New .invoke(state) shortcut on the built pipeline that returns the final state object (calls engine.execute under the hood, then projects state out of the context).
• dag.py
  • DAG learns to carry a state_schema and a reducers map derived from typing.get_type_hints(state_schema, include_extras=True).
  • DAG.is_cyclic() — used by the engine to pick the executor.
  • DAG.to_mermaid() and DAG.to_json().
• engine.py
  • Per-node retry/timeout/OTel/event-hook helper extracted into _node_runtime.py (pure refactor — same as both other plans propose).
  • PipelineEngine.execute dispatches: cyclic DAG → frontier_executor, acyclic DAG → existing Kahn path. No change to the acyclic path.

New files

• _node_runtime.py — shared retry/timeout/OTel/event-hook helper. Pure refactor lift from engine.py. Both executors call it.
• frontier_executor.py — frontier-based BFS executor with recursion_limit. Used for cyclic DAGs and for Send fan-out. Same shape as plan A's, but invoked by PipelineEngine, not by a new CompiledStateGraph class.
• reducers.py — append, extend, merge_dict, replace (default).
• state_adapter.py — internal: project state from PipelineContext and back. Lets state-based nodes coexist with (ctx, inputs) nodes in the same DAG during migration.

Soft-deprecation (not removal)

• BranchStep keeps working but emits a DeprecationWarning pointing users to .branch(...).
• FanOutStep keeps working but emits a DeprecationWarning pointing to Send.
• Per-node condition= keeps working unchanged.

Nothing in existing code breaks. Removal is a separate decision tracked in a follow-up issue, gated on internal callers migrating.

Tests

tests/pipeline/test_builder_evolved.py, plain pytest functions. Cover:

1. State-based node: a single async (state) -> dict node mutates state correctly.
2. Reducers: append accumulates across nodes; default replace overwrites; custom reducer via Annotated.
3. .branch(...): router returns label, correct branch runs; unknown label raises a clear error at runtime.
4. Cycle: two-node loop with router exit; recursion_limit triggers when never exits.
5. Send fan-out: router returns [Send("worker", a), Send("worker", b)]; both run concurrently; merge reducer combines.
6. Mermaid export: golden-file test.
7. Mixed-mode pipeline: one (ctx, inputs) node and one (state) -> dict node in the same DAG run correctly, with state adapter bridging.
8. Typed edges: a node annotated with Literal["a", "b"] rejects .branch(..., {"c": ...}) at .build().
9. Retry/timeout/failure-strategy: a flaky state-based node retries via the lifted _node_runtime; behaves identically to the same scenario on a (ctx, inputs) node.
10. Regression: every existing test in tests/pipeline/ still passes unchanged.

Run with pytest tests/pipeline/ -v from the repo root inside ~/.venvs/firefly.

Docs

• Update docs/pipeline/ with: the new state-based form, reducers, .branch, cycles, Send, Mermaid export, "when do I need a typed state" table.
• Update top-level README.md pipeline section.
• Do not mark the port-based form as deprecated — both forms remain first-class. Port-based is the right tool for pure ETL.

Reused existing code (do not reinvent)

• pipeline/engine.py — retry-with-backoff, timeout, OTel, event-hook dispatch lift into _node_runtime.py.
• pipeline/dag.py — DAGNode retains retry_max, timeout_seconds, failure_strategy. Reused by both executors.
• pipeline/context.py — PipelineContext remains the internal plumbing. State-based nodes never see it; port-based nodes get it unchanged.
• pipeline/result.py — NodeResult, ExecutionTraceEntry, UsageSummary continue to be emitted for both node forms. pipeline.invoke(state) returns the final state but pipeline.execute(state_or_inputs) is still available for callers that want the full PipelineResult.
• pipeline/steps.py — AgentStep, ReasoningStep, etc. continue to work in port-based mode unchanged. Add a tiny as_state_node(step) helper so they can also be dropped into a state-based pipeline.

Verification

1. pytest tests/pipeline/ -v — full pipeline suite green, including all preexisting tests.
2. pytest tests/ -v — full repo suite green.
3. ruff check . and mypy fireflyframework_agentic/pipeline/ clean.
4. Run the canonical example as a smoke script; print Mermaid; eyeball topology.
5. Pick one existing pipeline; rewrite it in state-based form side by side with the original; outputs should match.
6. Verify deprecation warnings fire for BranchStep and FanOutStep and not otherwise.

What is explicitly NOT in v1

• Checkpointer / HITL interrupt+resume.
• pipeline.stream() partial-state streaming.
• Subgraphs.
• Migration of internal callers off the port-based form — separate follow-up.
• Removal of BranchStep / FanOutStep — separate follow-up, gated on migration.

When to pick this plan over A or B

Pick this plan if:

• One pipeline API is worth more than perfect alignment with LangGraph or pydantic-graph idioms.
• Adding a third-party dependency (plan B) is unacceptable.
• "Declarative-as-data" is enough; full type-system-encoded edges (plan B) are nice-to-have, not load-bearing.
• Source compatibility with every existing pipeline matters.

Pick plan A if:

• A clean, separate API for agentic workflows is worth the cost of maintaining two builders.
• Matching LangGraph idioms exactly (so users coming from LangGraph feel at home) is a hiring/onboarding priority.

Pick plan B if:

• Static type-safety on edges is the top priority.
• Heavy LLM fan-out is not a near-term workload.
• Adding pydantic-graph as a dependency is acceptable.

Estimated size

~600–900 LOC of new + modified code (builder extension + frontier executor + reducers + Mermaid/JSON + state adapter), ~400 LOC of tests, plus docs. The _node_runtime.py extraction is the only behaviour-preserving refactor of existing code.