Crate graphflow_stream

Expand description

§graphflow-stream

The astream_events LangGraph gives Python, for graph-flow in Rust.

§Install

cargo add graphflow-stream

§Usage

use graph_flow::{Context, NextAction, Task, TaskResult, error::Result};
use graphflow_stream::{emit_token, spawn_task};

struct MyLlmTask;

#[async_trait::async_trait]
impl Task for MyLlmTask {
    fn id(&self) -> &str { "my_llm_task" }

    async fn run(&self, _context: Context) -> Result<TaskResult> {
        for delta in ["Hel", "lo", "!"] {           // e.g. deltas from Rig
            emit_token("my_llm_task", delta).await;
        }
        Ok(TaskResult::new(Some("Hello!".into()), NextAction::Continue))
    }
}

let (mut rx, handle) = spawn_task(Arc::new(MyLlmTask), Context::new(), 32);
while let Some(event) = rx.recv().await {
    // forward over SSE / WebSocket / stdout as it arrives
}
let result = handle.await??; // same TaskResult you'd get from task.run()

Just want the full streamed text back, no manual loop? One line:

let text = graphflow_stream::collect_text(Arc::new(MyLlmTask), Context::new(), 32).await?;

emit_token/emit_started/emit_finished/emit_failed are ambient — call them from anywhere inside Task::run, no new trait to implement, no-op if nothing is listening. spawn_graph(flow_runner, session_id, buffer) does the same for a whole FlowRunner run; SubgraphTask wraps a nested Graph as one Task and streams through automatically.

Replay debugging: record(rx).await turns a run into a Recording (serializable, so it can be saved to disk), and recording.replay(buffer) plays it back on a fresh channel with the original timing — inspect a past run, or demo a UI without hitting an LLM again.

§Orchestration: map over a runtime list, vote across runs

graph_flow’s built-in FanOutTask runs a fixed set of children decided at construction time. DynamicMapTask covers what LangGraph’s Send API covers in Python — fan out over however many items context holds this run (one child per retrieved document, one per subtask an LLM just planned):

use graphflow_stream::DynamicMapTask;

let map_task = DynamicMapTask::new("summarize_retrieved", |ctx: &Context| {
    let docs: Vec<String> = ctx.get("retrieved_docs").unwrap_or_default();
    docs.into_iter()
        .map(|doc_id| Arc::new(SummarizeDoc { doc_id }) as Arc<dyn Task>)
        .collect()
}).with_prefix("summaries");

map_task.run(context).await?; // writes summaries.<doc_id>.response for each doc

EnsembleTask runs the same task several times concurrently and reduces the responses — self-consistency prompting, sample an LLM call a few times and combine instead of trusting one draw:

use graphflow_stream::{EnsembleTask, majority_vote};

let ensemble = EnsembleTask::new("classify_intent", ClassifyIntent, 5, majority_vote);
let result = ensemble.run(context).await?; // most common of 5 concurrent runs

majority_vote ships built in; pass any Fn(Vec<String>) -> String for a custom reducer (join, longest, an LLM-as-judge pick).

More examples (full_graph, sse_axum, websocket_axum, replay, map_and_ensemble) in examples/.

§Benchmarks

cargo bench (criterion, benches/overhead.rs):

Scenario	Time
`task.run()` direct — no `graphflow-stream` involved	~1.0 µs
`emit_token()` with nobody listening (ambient no-op)	~30 ns / call
`spawn_task()` streaming 100 tokens to a draining receiver	~2.0 µs / token
`record()` capturing 100 streamed tokens	~1.4 µs / token
`DynamicMapTask::run`, 10 items	~239 µs
`EnsembleTask::run`, 5 runs	~215 µs

§License

MIT

Structs§

DynamicMapTask: Fans out over a list of child tasks built from context at run time — the same shape as graph_flow::fanout::FanOutTask, except the number and identity of children isn’t fixed until the task actually runs. This is what LangGraph’s Send API covers in Python: “run this step once per item in a list I only know at runtime” (once per retrieved document, once per subtask an LLM just planned out), which a construction-time Vec of children can’t express.
EnsembleTask: Runs the same Task runs times concurrently and reduces the responses into one — self-consistency prompting: sample an LLM call several times (usually with temperature > 0) and combine the answers instead of trusting a single draw.
RecordedEvent: A StreamEvent plus its offset from the start of the recording.
Recording: A recorded run: every StreamEvent a StreamReceiver produced, with timing.
SubgraphTask: A Task that drives a whole inner Graph, so a graph can be a node in another graph.

Enums§

StreamEvent: An incremental event emitted while a task or graph run is in flight.

Functions§

collect_text: Run a Task and join its streamed StreamEvent::Token deltas into one String.
emit: Send an event on the ambient stream, if one is active; a no-op otherwise.
emit_failed: Emit a StreamEvent::TaskFailed on the ambient stream.
emit_finished: Emit a StreamEvent::TaskFinished on the ambient stream.
emit_started: Emit a StreamEvent::TaskStarted on the ambient stream.
emit_token: Emit a StreamEvent::Token on the ambient stream.
forward_text_stream: Drain any Stream<Item = String> (e.g. a mapped LLM completion) into emit_token calls.
majority_vote: A ready-made EnsembleTask reducer: the most common response wins, ties broken by whichever came first. Fits self-consistency prompting, where every run should converge on the same discrete answer.
record: Drain a StreamReceiver into a Recording, timestamping each event from the start.
run_streaming: Run fut with an ambient stream scope active; the primitive spawn_task/spawn_graph use.
spawn_graph: Drive a FlowRunner to completion, streaming every task’s emit_* calls out.
spawn_task: Run a single Task, streaming its emit_* calls out through the returned receiver.

Type Aliases§

StreamReceiver: Receiving half of a stream channel; drain this to observe a run in progress.
StreamSender: Sending half of a stream channel; a running task’s emit_* calls write here.