defmodule Planck.AI.Stream do
@moduledoc """
Streaming event types and normalization from `req_llm` chunks.
`req_llm` emits `%ReqLLM.StreamChunk{}` structs with a `:type` field.
`from_req_llm/1` maps each chunk type to the canonical event tuples consumed
by `Planck.Agent` and callers of `Planck.AI.stream/3`.
Tool call arguments arrive as JSON fragments spread across multiple `:meta`
chunks. This module buffers the fragments and emits a single assembled
`{:tool_call_complete, ...}` per tool call when the stream finishes.
Exceptions raised during stream consumption (e.g. `ReqLLM.Error.API.Stream`
on HTTP errors) are caught and emitted as `{:error, exception}` events,
halting the stream gracefully.
## Event types
- `{:text_delta, text}` — a chunk of assistant text
- `{:thinking_delta, text}` — a chunk of extended thinking / reasoning
- `{:tool_call_complete, %{id, name, args}}` — a fully assembled tool call
- `{:done, %{stop_reason, usage}}` — stream finished; includes token usage
- `{:error, reason}` — an error occurred during streaming
"""
@type tool_call :: %{
id: String.t(),
name: String.t(),
args: map()
}
@type usage :: %{
input_tokens: non_neg_integer(),
output_tokens: non_neg_integer()
}
@type done :: %{
stop_reason: atom(),
usage: usage()
}
@type t ::
{:text_delta, String.t()}
| {:thinking_delta, String.t()}
| {:tool_call_complete, tool_call()}
| {:done, done()}
| {:error, term()}
@doc """
Converts a stream of `ReqLLM.StreamChunk` structs into event tuples.
## Examples
iex> chunks = [
...> %{type: :content, text: "Hello"},
...> %{type: :meta, metadata: %{finish_reason: :stop, usage: %{input_tokens: 10, output_tokens: 5}}}
...> ]
iex> chunks |> Planck.AI.Stream.from_req_llm() |> Enum.to_list()
[{:text_delta, "Hello"}, {:done, %{stop_reason: :stop, usage: %{input_tokens: 10, output_tokens: 5}}}]
"""
@spec from_req_llm(Enumerable.t()) :: Enumerable.t(t())
def from_req_llm(chunks) do
Stream.resource(
fn -> {pull(chunks), %{calls: %{}, fragments: %{}}} end,
&emit/1,
fn _ -> :ok end
)
end
# --- Private ---
#
# How one-at-a-time pulling works
# --------------------------------
# Stream.resource drives everything. Its `next_fun` (`emit/1`) is called once
# per consumer demand and must return early — if we let Enumerable.reduce run
# to completion over the whole upstream, we'd lose all incrementality.
#
# The trick is `suspend_reducer/2`, which immediately returns `{:suspend, el}`
# for every element. This causes Enumerable.reduce to pause after the first
# element and return:
#
# {:suspended, element, continuation}
#
# where `continuation` is a closure over the remaining elements. `step_result`
# converts that into `{:item, element, continuation}` so `emit` can process
# one chunk and store the continuation as the next state.
#
# The next call to `emit` resumes by calling `pull(continuation)`, which
# advances by one more element. When the upstream is exhausted, reduce returns
# `{:done, _}` instead of `{:suspended, ...}`, which becomes `:done` —
# signalling `emit` to halt the stream.
#
# Importantly, `pull` is synchronous: if the upstream is an HTTP stream, it
# blocks until the next chunk arrives over the network. So even if the
# consumer is `Enum.to_list` (which is eager), it can't race ahead — it waits
# on each `pull` just like a lazy consumer would.
#
# Continuation arity
# ------------------
# List/Stream enumerables produce arity-1 continuations (the reducer is
# already captured in the closure). Stream.resource-based enumerables produce
# arity-2 continuations and require the reducer to be passed back in on each
# call. The two `pull/1` function clauses handle each case.
defp emit({:done, _tool_acc}), do: {:halt, :done}
defp emit({{:error, reason}, _tool_acc}), do: {[{:error, reason}], {:done, %{}}}
defp emit({{:item, chunk, cont}, tool_acc}) do
{events, new_acc} = process_chunk(chunk, tool_acc)
{events, {pull(cont), new_acc}}
end
# Tool call identity chunk: buffer name/id by index, emit nothing yet.
defp process_chunk(%{type: :tool_call, name: name} = chunk, acc) do
index = (chunk.metadata && chunk.metadata[:index]) || 0
id = (chunk.metadata && chunk.metadata[:id]) || ""
{[], %{acc | calls: Map.put(acc.calls, index, {id, name})}}
end
# Argument fragment: accumulate JSON string by index, emit nothing yet.
defp process_chunk(
%{type: :meta, metadata: %{tool_call_args: %{index: index, fragment: frag}}},
acc
) do
existing = Map.get(acc.fragments, index, [])
{[], %{acc | fragments: Map.put(acc.fragments, index, [existing | [frag]])}}
end
# Final meta chunk: emit assembled tool calls then done.
defp process_chunk(%{type: :meta, metadata: metadata}, acc) do
tool_events = assemble_tool_calls(acc)
stop_reason = Map.get(metadata, :finish_reason, :stop)
usage = Map.get(metadata, :usage, %{input_tokens: 0, output_tokens: 0})
{tool_events ++ [{:done, %{stop_reason: stop_reason, usage: usage}}],
%{calls: %{}, fragments: %{}}}
end
# All other chunk types (text, thinking, unknown).
defp process_chunk(chunk, acc) do
{[normalize_chunk(chunk)], acc}
end
defp assemble_tool_calls(%{calls: calls}) when calls == %{}, do: []
defp assemble_tool_calls(%{calls: calls, fragments: fragments}) do
calls
|> Enum.sort_by(fn {index, _} -> index end)
|> Enum.map(fn {index, {id, name}} ->
{:tool_call_complete, %{id: id, name: name, args: decode_args(fragments, index)}}
end)
end
defp decode_args(fragments, index) do
case Map.get(fragments, index) do
nil -> %{}
frags -> frags |> IO.iodata_to_binary() |> ensure_object() |> parse_json()
end
end
defp ensure_object("{" <> _ = json), do: json
defp ensure_object(json), do: "{" <> json
defp parse_json(json) do
case Jason.decode(json) do
{:ok, decoded} when is_map(decoded) -> decoded
_ -> %{}
end
end
# Arity-1 continuations: returned by list and %Stream{} enumerables.
defp pull(cont) when is_function(cont, 1) do
cont.({:cont, nil}) |> step_result()
rescue
e -> {:error, e}
catch
:exit, reason -> {:error, {:exit, reason}}
end
# Arity-2 continuations: returned by Stream.resource, which passes both the
# accumulator AND the reducer through the continuation so it can resume
# enumeration with the same reducer on the next call.
defp pull(cont) when is_function(cont, 2) do
cont.({:cont, nil}, &suspend_reducer/2) |> step_result()
rescue
e -> {:error, e}
catch
:exit, reason -> {:error, {:exit, reason}}
end
# Initial call: start reducing the upstream enumerable one element at a time.
defp pull(enum) do
Enumerable.reduce(enum, {:cont, nil}, &suspend_reducer/2) |> step_result()
rescue
e -> {:error, e}
catch
:exit, reason -> {:error, {:exit, reason}}
end
defp step_result({:suspended, chunk, cont}), do: {:item, chunk, cont}
defp step_result({:done, _}), do: :done
defp step_result({:halted, _}), do: :done
defp suspend_reducer(el, _acc), do: {:suspend, el}
defp normalize_chunk(%{type: :content, text: text}), do: {:text_delta, text}
defp normalize_chunk(%{type: :thinking, text: text}), do: {:thinking_delta, text}
defp normalize_chunk(chunk), do: {:error, {:unknown_chunk, chunk}}
end
